From a38fb588ec02ad72c99d5848f8687aadea0da94f Mon Sep 17 00:00:00 2001 From: Kunal Kushwaha Date: Sun, 8 Feb 2026 16:55:18 +0900 Subject: [PATCH 1/2] mermaid diagrams fixed --- go.mod | 1 + go.sum | 2 + internal/audit/mermaid.go | 324 ++++++++++++++++++++++++++++---------- 3 files changed, 246 insertions(+), 81 deletions(-) diff --git a/go.mod b/go.mod index 6b13879..5981ea4 100644 --- a/go.mod +++ b/go.mod @@ -5,6 +5,7 @@ go 1.24.1 require ( github.com/BurntSushi/toml v1.5.0 github.com/Masterminds/sprig/v3 v3.3.0 + github.com/TyphonHill/go-mermaid v1.0.0 github.com/agenticgokit/agenticgokit v0.5.5 github.com/charmbracelet/bubbles v0.21.0 github.com/charmbracelet/bubbletea v1.3.10 diff --git a/go.sum b/go.sum index 1a99102..306af57 100644 --- a/go.sum +++ b/go.sum @@ -17,6 +17,8 @@ github.com/Microsoft/go-winio v0.6.2 h1:F2VQgta7ecxGYO8k3ZZz3RS8fVIXVxONVUPlNERo github.com/Microsoft/go-winio v0.6.2/go.mod h1:yd8OoFMLzJbo9gZq8j5qaps8bJ9aShtEA8Ipt1oGCvU= github.com/ProtonMail/go-crypto v1.1.6 h1:ZcV+Ropw6Qn0AX9brlQLAUXfqLBc7Bl+f/DmNxpLfdw= github.com/ProtonMail/go-crypto v1.1.6/go.mod h1:rA3QumHc/FZ8pAHreoekgiAbzpNsfQAosU5td4SnOrE= +github.com/TyphonHill/go-mermaid v1.0.0 h1:VtmgQwgZA+KNHJvG/O591ibBVuDkGhg2+F/olVXnXAs= +github.com/TyphonHill/go-mermaid v1.0.0/go.mod h1:BqMEbKnr2HHpZ4lJJvGjL47v6rZAUpJcOaE/db1Ppwc= github.com/agenticgokit/agenticgokit v0.5.5 h1:f/+2EbiIImlUsK8RP23V3W1D5pFtS+EgH/vCAqzPEF4= github.com/agenticgokit/agenticgokit v0.5.5/go.mod h1:0EwU951CZIGYwEOLnC5hJbC9lhNvM85FhrL6NTTDIZo= github.com/anmitsu/go-shlex v0.0.0-20200514113438-38f4b401e2be h1:9AeTilPcZAjCFIImctFaOjnTIavg87rW78vTPkQqLI8= diff --git a/internal/audit/mermaid.go b/internal/audit/mermaid.go index 5659a5d..92ef9b1 100644 --- a/internal/audit/mermaid.go +++ b/internal/audit/mermaid.go @@ -2,104 +2,207 @@ package audit import ( "fmt" + "sort" + "strconv" "strings" + + "github.com/TyphonHill/go-mermaid/diagrams/flowchart" ) // GenerateMermaid creates a Mermaid flowchart from a TraceObject func GenerateMermaid(obj *TraceObject) string { - var b strings.Builder - - b.WriteString("```mermaid\n") - b.WriteString("flowchart TD\n") - - // Create nodes for each event - for i, event := range obj.Events { - nodeID := fmt.Sprintf("N%d", i) - label := formatNodeLabel(event) - shape := getNodeShape(event.Type) - - b.WriteString(fmt.Sprintf(" %s%s\n", nodeID, shape(label))) - } - - b.WriteString("\n") - - // Create edges between consecutive events - for i := 0; i < len(obj.Events)-1; i++ { - b.WriteString(fmt.Sprintf(" N%d --> N%d\n", i, i+1)) - } - - // Add styling - b.WriteString("\n") - b.WriteString(" %% Styling\n") - for i, event := range obj.Events { - style := getNodeStyle(event.Type) - if style != "" { - b.WriteString(fmt.Sprintf(" style N%d %s\n", i, style)) - } - } - - b.WriteString("```\n") - - return b.String() + return GenerateMermaidWithHierarchy(obj) } // GenerateMermaidWithHierarchy creates a Mermaid diagram respecting parent-child relationships func GenerateMermaidWithHierarchy(obj *TraceObject) string { - var b strings.Builder - // Build parent map parentMap := make(map[string][]int) spanIDToIndex := make(map[string]int) + childrenBySpan := make(map[string][]string) + spanByIndex := make([]string, len(obj.Events)) for i, event := range obj.Events { spanIDToIndex[event.SpanID] = i + spanByIndex[i] = event.SpanID if event.ParentID != "" && event.ParentID != "0000000000000000" { parentMap[event.ParentID] = append(parentMap[event.ParentID], i) + childrenBySpan[event.ParentID] = append(childrenBySpan[event.ParentID], event.SpanID) } } - b.WriteString("```mermaid\n") - b.WriteString("flowchart TD\n") + diagram := flowchart.NewFlowchart() + diagram.EnableMarkdownFence() + diagram.SetDirection(flowchart.FlowchartDirectionTopDown) + diagram.Config.SetHtmlLabels(true) - // Create nodes + nodes := make([]*flowchart.Node, len(obj.Events)) for i, event := range obj.Events { - nodeID := fmt.Sprintf("N%d", i) label := formatNodeLabel(event) - shape := getNodeShape(event.Type) - b.WriteString(fmt.Sprintf(" %s%s\n", nodeID, shape(label))) + node := diagram.AddNode(label) + applyFlowchartShape(node, event.Type) + if style := getFlowchartStyle(event.Type); style != nil { + node.SetStyle(style) + } + nodes[i] = node } - b.WriteString("\n") - - // Create edges based on parent-child relationships - for parentSpanID, children := range parentMap { + parentIndices := make([]int, 0, len(parentMap)) + for parentSpanID := range parentMap { if parentIdx, ok := spanIDToIndex[parentSpanID]; ok { + parentIndices = append(parentIndices, parentIdx) + } + } + sort.Ints(parentIndices) + + addedLinks := make(map[string]bool) + addLink := func(fromIdx, toIdx int) { + key := fmt.Sprintf("%d->%d", fromIdx, toIdx) + if addedLinks[key] { + return + } + addedLinks[key] = true + diagram.AddLink(nodes[fromIdx], nodes[toIdx]) + } + + // Special handling for sequential workflows: chain steps and nest descendants + sequentialParents := make([]int, 0) + for _, parentIdx := range parentIndices { + if isWorkflowSequential(obj.Events[parentIdx]) { + sequentialParents = append(sequentialParents, parentIdx) + } + } + + if len(sequentialParents) > 0 { + for _, parentIdx := range sequentialParents { + parentSpanID := obj.Events[parentIdx].SpanID + children := parentMap[parentSpanID] + stepChildren := make([]int, 0) for _, childIdx := range children { - b.WriteString(fmt.Sprintf(" N%d --> N%d\n", parentIdx, childIdx)) + if isWorkflowStep(obj.Events[childIdx]) { + stepChildren = append(stepChildren, childIdx) + } + } + + sort.Slice(stepChildren, func(i, j int) bool { + idxI := stepChildren[i] + idxJ := stepChildren[j] + stepI, okI := getStepIndex(obj.Events[idxI]) + stepJ, okJ := getStepIndex(obj.Events[idxJ]) + if okI && okJ { + return stepI < stepJ + } + if okI != okJ { + return okI + } + return obj.Events[idxI].Timestamp.Before(obj.Events[idxJ].Timestamp) + }) + + if len(stepChildren) > 0 { + addLink(parentIdx, stepChildren[0]) + for i := 0; i < len(stepChildren)-1; i++ { + addLink(stepChildren[i], stepChildren[i+1]) + } + } + + for _, stepIdx := range stepChildren { + descendants := collectDescendantIndices(spanByIndex[stepIdx], spanIDToIndex, childrenBySpan, obj) + if len(descendants) == 0 { + continue + } + sort.Slice(descendants, func(i, j int) bool { + return obj.Events[descendants[i]].Timestamp.Before(obj.Events[descendants[j]].Timestamp) + }) + addLink(stepIdx, descendants[0]) + for i := 0; i < len(descendants)-1; i++ { + addLink(descendants[i], descendants[i+1]) + } } } + + return diagram.String() } - // For orphan nodes (no parent in trace), connect sequentially - hasParent := make(map[int]bool) - for _, children := range parentMap { - for _, idx := range children { - hasParent[idx] = true + for _, parentIdx := range parentIndices { + parentEvent := obj.Events[parentIdx] + parentSpanID := parentEvent.SpanID + children := parentMap[parentSpanID] + + if isWorkflowSequential(parentEvent) { + stepChildren := make([]int, 0) + otherChildren := make([]int, 0) + for _, childIdx := range children { + if isWorkflowStep(obj.Events[childIdx]) { + stepChildren = append(stepChildren, childIdx) + } else { + otherChildren = append(otherChildren, childIdx) + } + } + + if len(stepChildren) > 0 { + sort.Slice(stepChildren, func(i, j int) bool { + idxI := stepChildren[i] + idxJ := stepChildren[j] + stepI, okI := getStepIndex(obj.Events[idxI]) + stepJ, okJ := getStepIndex(obj.Events[idxJ]) + if okI && okJ { + return stepI < stepJ + } + if okI != okJ { + return okI + } + return obj.Events[idxI].Timestamp.Before(obj.Events[idxJ].Timestamp) + }) + + addLink(parentIdx, stepChildren[0]) + for i := 0; i < len(stepChildren)-1; i++ { + addLink(stepChildren[i], stepChildren[i+1]) + } + } + + sort.Ints(otherChildren) + for _, childIdx := range otherChildren { + addLink(parentIdx, childIdx) + } + continue } - } - // Add styling - b.WriteString("\n") - for i, event := range obj.Events { - style := getNodeStyle(event.Type) - if style != "" { - b.WriteString(fmt.Sprintf(" style N%d %s\n", i, style)) + sort.Ints(children) + for _, childIdx := range children { + addLink(parentIdx, childIdx) } } - b.WriteString("```\n") + return diagram.String() +} + +func collectDescendantIndices(rootSpanID string, spanIDToIndex map[string]int, childrenBySpan map[string][]string, obj *TraceObject) []int { + var result []int + queue := []string{rootSpanID} + visited := make(map[string]bool) + + for len(queue) > 0 { + current := queue[0] + queue = queue[1:] + if visited[current] { + continue + } + visited[current] = true + children := childrenBySpan[current] + for _, childSpan := range children { + idx, ok := spanIDToIndex[childSpan] + if ok { + if isWorkflowStep(obj.Events[idx]) || isWorkflowSequential(obj.Events[idx]) { + // Skip other workflow step nodes to avoid cross-linking + } else { + result = append(result, idx) + } + } + queue = append(queue, childSpan) + } + } - return b.String() + return result } // formatNodeLabel creates a concise label for the node @@ -109,19 +212,68 @@ func formatNodeLabel(event TraceEvent) string { // Get a short description desc := event.SpanName - if len(desc) > 25 { - desc = desc[:22] + "..." + if stepName, ok := event.Metadata["agk.workflow.step_name"].(string); ok && stepName != "" { + desc = "step:" + stepName + } + if agentName, ok := event.Metadata["agk.agent.name"].(string); ok && agentName != "" { + desc = fmt.Sprintf("%s @%s", desc, agentName) + } + if len(desc) > 60 { + desc = desc[:57] + "..." } - // Add duration if significant + // Add duration on new line duration := "" - if event.DurationMs > 100 { - duration = fmt.Sprintf(" (%dms)", event.DurationMs) + if event.DurationMs > 0 { + duration = fmt.Sprintf("
%dms", event.DurationMs) } return fmt.Sprintf("%s %s%s", icon, desc, duration) } +func isWorkflowSequential(event TraceEvent) bool { + name := strings.ToLower(event.SpanName) + return strings.Contains(name, "workflow.sequential") +} + +func isWorkflowStep(event TraceEvent) bool { + name := strings.ToLower(event.SpanName) + if strings.Contains(name, "workflow.step") { + return true + } + if stepName, ok := event.Metadata["agk.workflow.step_name"].(string); ok && stepName != "" { + return true + } + return false +} + +func getStepIndex(event TraceEvent) (int, bool) { + if raw, ok := event.Metadata["agk.workflow.step_index"]; ok { + switch v := raw.(type) { + case int: + return v, true + case int64: + return int(v), true + case float64: + return int(v), true + case float32: + return int(v), true + case string: + if parsed, err := parseInt(v); err == nil { + return parsed, true + } + } + } + if !event.Timestamp.IsZero() { + return int(event.Timestamp.UnixMilli()), false + } + return 0, false +} + +func parseInt(value string) (int, error) { + return strconv.Atoi(value) +} + // getEventIcon returns an emoji for the event type func getEventIcon(eventType EventType) string { switch eventType { @@ -141,37 +293,47 @@ func getEventIcon(eventType EventType) string { } // getNodeShape returns a function that wraps the label in the appropriate shape -func getNodeShape(eventType EventType) func(string) string { +func applyFlowchartShape(node *flowchart.Node, eventType EventType) { switch eventType { case EventTypeThought: - return func(label string) string { return fmt.Sprintf("([%s])", label) } // Stadium + node.SetShape(flowchart.NodeShapeTerminal) case EventTypeToolCall: - return func(label string) string { return fmt.Sprintf("[[%s]]", label) } // Subroutine + node.SetShape(flowchart.NodeShapeSubprocess) case EventTypeObservation: - return func(label string) string { return fmt.Sprintf("[/%s/]", label) } // Parallelogram + node.SetShape(flowchart.NodeShapeInputOutput) case EventTypeLLMCall: - return func(label string) string { return fmt.Sprintf("{%s}", label) } // Rhombus + node.SetShape(flowchart.NodeShapeDecision) case EventTypeDecision: - return func(label string) string { return fmt.Sprintf("{{%s}}", label) } // Hexagon + node.SetShape(flowchart.NodeShapePrepare) default: - return func(label string) string { return fmt.Sprintf("[%s]", label) } + node.SetShape(flowchart.NodeShapeProcess) } } -// getNodeStyle returns Mermaid styling for the event type -func getNodeStyle(eventType EventType) string { +// getFlowchartStyle returns Mermaid styling for the event type +func getFlowchartStyle(eventType EventType) *flowchart.NodeStyle { + style := flowchart.NewNodeStyle() + style.StrokeWidth = 1 + switch eventType { case EventTypeThought: - return "fill:#e1f5fe,stroke:#01579b" + style.Fill = "#e1f5fe" + style.Stroke = "#01579b" case EventTypeToolCall: - return "fill:#e8f5e9,stroke:#1b5e20" + style.Fill = "#e8f5e9" + style.Stroke = "#1b5e20" case EventTypeObservation: - return "fill:#fff3e0,stroke:#e65100" + style.Fill = "#fff3e0" + style.Stroke = "#e65100" case EventTypeLLMCall: - return "fill:#f3e5f5,stroke:#4a148c" + style.Fill = "#f3e5f5" + style.Stroke = "#4a148c" case EventTypeDecision: - return "fill:#fce4ec,stroke:#880e4f" + style.Fill = "#fce4ec" + style.Stroke = "#880e4f" default: - return "" + return nil } + + return style } From 7f08d0466542c7cb7c58e49117fcc00fe5bb0ee2 Mon Sep 17 00:00:00 2001 From: Kunal Kushwaha Date: Sun, 8 Feb 2026 17:27:27 +0900 Subject: [PATCH 2/2] fixed templates --- cmd/init.go | 63 ++++++++++++++++++++++++++++++++++------------------- 1 file changed, 40 insertions(+), 23 deletions(-) diff --git a/cmd/init.go b/cmd/init.go index c5abe18..5fcbfcc 100644 --- a/cmd/init.go +++ b/cmd/init.go @@ -4,6 +4,8 @@ import ( "fmt" "os" "path/filepath" + "sort" + "strings" "github.com/agenticgokit/agenticgokit/observability" "github.com/fatih/color" @@ -48,10 +50,10 @@ Examples: agk init my-project # Initialize with built-in template - agk init my-project --template single-agent + agk init my-project --template quickstart # Initialize from a community template (registry) - agk init my-project --template rag-agent + agk init my-project --template # Initialize from a GitHub repository agk init my-project --template github.com/username/my-template @@ -59,13 +61,13 @@ Examples: # Initialize from a specific version agk init my-project --template github.com/username/my-template@v1.0.0 - # Non-interactive initialization - agk init my-project --template single-agent --llm openai --agent-type single --force + # Non-interactive initialization + agk init my-project --template quickstart --llm openai --force # Initialize in specific directory agk init my-project --output ./projects - # List available templates + # List available templates agk init --list`, Args: func(cmd *cobra.Command, args []string) error { // Allow zero args only when listing templates @@ -226,33 +228,48 @@ func listTemplates() { color.Cyan("\nšŸ“‹ Available AgenticGoKit Templates\n") color.Cyan("ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•\n\n") + // Built-in templates templates := scaffold.GetAllTemplates() + color.Cyan("Built-in:\n") for i, tmpl := range templates { - // Template name and complexity color.Green("%d. %s %s\n", i+1, tmpl.Name, tmpl.Complexity) - - // Description fmt.Printf(" %s\n", color.YellowString(tmpl.Description)) - - // Features if len(tmpl.Features) > 0 { fmt.Printf(" Features: %v\n", color.CyanString("%v", tmpl.Features)) } - - // File count fmt.Printf(" Files: %s\n", color.MagentaString("%d", tmpl.FileCount)) - - // Usage example - templateID := "" - switch tmpl.Name { - case "Quickstart": - templateID = "quickstart" - case "Workflow": - templateID = "workflow" + fmt.Printf(" Usage: %s\n", color.HiBlackString("agk init my-project --template %s", strings.ToLower(tmpl.Name))) + if i < len(templates)-1 { + fmt.Println() } - fmt.Printf(" Usage: %s\n", color.HiBlackString("agk init my-project --template %s", templateID)) + } - if i < len(templates)-1 { + // Registry templates + color.Cyan("\nRegistry:\n") + index, err := registry.FetchIndex(registry.DefaultRegistryURL) + if err != nil { + fmt.Printf(" %s\n", color.YellowString("Unable to fetch registry templates: %v", err)) + fmt.Println() + return + } + + if len(index.Templates) == 0 { + fmt.Printf(" %s\n", color.YellowString("No templates found in registry.")) + fmt.Println() + return + } + + registryNames := make([]string, 0, len(index.Templates)) + for name := range index.Templates { + registryNames = append(registryNames, name) + } + sort.Strings(registryNames) + for i, name := range registryNames { + source := index.Templates[name] + color.Green("%d. %s\n", i+1, name) + fmt.Printf(" Source: %s\n", color.HiBlackString(source)) + fmt.Printf(" Usage: %s\n", color.HiBlackString("agk init my-project --template %s", name)) + if i < len(registryNames)-1 { fmt.Println() } } @@ -334,7 +351,7 @@ func init() { // Define flags initCmd.Flags().BoolVar(&initListTemplates, "list", false, "List available templates") initCmd.Flags().StringVarP(&initTemplate, "template", "t", "quickstart", - "Template type: quickstart, workflow") + "Template name (built-in: quickstart, workflow; or a registry template)") initCmd.Flags().StringVarP(&initOutputDir, "output", "o", ".", "Output directory for the project") initCmd.Flags().BoolVarP(&initInteractive, "interactive", "i", false, "Enable interactive prompts") initCmd.Flags().BoolVarP(&initForce, "force", "f", false, "Force overwrite existing files")