poimen-workflows/action/canvas_reasoner.go

package action

import (
	"context"
	"encoding/json"
	"fmt"

	"github.com/rockliang/poimen/workflows/action/llm"
	"github.com/rockliang/poimen/workflows/pkg/db"
	"github.com/rockliang/poimen/workflows/statemachine"
)

// CanvasReasonerInput infers connections between nodes using LLM reasoning
type CanvasReasonerInput struct {
	Nodes []db.WorkflowNode `json:"nodes"`       // Canvas nodes
	Edges []db.WorkflowEdge `json:"edges"`       // Existing edges
	// If true, only suggest new edges; if false, redesign entire canvas
	PreserveExisting bool `json:"preserve_existing,omitempty"`
	AuthToken        string `json:"auth_token,omitempty"` // JWT for LLM calls
}

// RelationWording describes semantic meaning of an edge
type RelationWording struct {
	Verb               string `json:"verb"`                 // outputs, inputs, depends-on, etc
	SourceOutput       string `json:"source_output"`        // What source produces
	TargetInput        string `json:"target_input"`         // What target requires
	ConnectionType     string `json:"connection_type"`      // direct-map, requires-transformer, conditional
	Confidence         float64 `json:"confidence"`          // 0.0-1.0
	SemanticMatch      string `json:"semantic_match"`       // Human-readable explanation
	TransformerNeeded  string `json:"transformer_needed,omitempty"` // If transformation required
}

// EdgeWithWording pairs an edge with its semantic description
type EdgeWithWording struct {
	Source           string          `json:"source"`
	Target           string          `json:"target"`
	RelationType     string          `json:"relation_type"`      // data-flow, dependency, conditional, parallel
	RelationLabel    string          `json:"relation_label"`     // e.g., "CloneRepo outputs path → AnalyzeCode requires path"
	RelationWording  RelationWording `json:"relation_wording"`
}

// CanvasReasonerOutput returns suggested edges and reasoning
type CanvasReasonerOutput struct {
	SuggestedEdges        []EdgeWithWording             `json:"suggested_edges"`         // Edges with wording
	RemovedEdges          []db.WorkflowEdge             `json:"removed_edges,omitempty"` // Edges to remove
	Reasoning             string                        `json:"reasoning"`               // LLM explanation
	Confidence            float64                       `json:"confidence"`              // 0.0-1.0
	IncompatibleEdges     []IncompatibilityWarning      `json:"incompatible_edges,omitempty"` // Can't connect
	DisconnectedNodes     []string                      `json:"disconnected_nodes,omitempty"` // No connections
	UserAlerts            []string                      `json:"user_alerts,omitempty"`        // Human-readable warnings
}

// CanvasReasonerActivity uses LLM to infer connections between workflow activities
func CanvasReasonerActivity(ctx context.Context, in CanvasReasonerInput) (CanvasReasonerOutput, error) {
	logger := newActivityLogger(ctx)

	output := CanvasReasonerOutput{
		SuggestedEdges: []db.WorkflowEdge{},
	}

	if len(in.Nodes) == 0 {
		return output, fmt.Errorf("no nodes provided")
	}

	logger.logf("info", "Analyzing canvas with %d nodes, %d edges", len(in.Nodes), len(in.Edges))

	// Build activity descriptions for LLM context
	nodeDesc := buildNodeDescriptions(in.Nodes)
	edgeDesc := buildEdgeDescriptions(in.Edges)

	// Create prompt for LLM reasoning with relation wording
	systemPrompt := `You are a workflow automation expert. Analyze activities and suggest logical connections with semantic descriptions.

CRITICAL RULES:
1. Only suggest edges where outputs→inputs match
2. Provide relation wording: verb, source_output, target_input
3. Assess connection confidence (0.0-1.0)
4. Flag type mismatches that need transformers

Respond with JSON:
{
  "edges": [
    {
      "source": "node-1",
      "target": "node-2",
      "relation_type": "data-flow|dependency|conditional|parallel",
      "relation_label": "Node1 outputs X → Node2 requires X",
      "relation_wording": {
        "verb": "outputs|depends-on|triggers|etc",
        "source_output": "field_name (type): description",
        "target_input": "field_name (type, required?): description",
        "connection_type": "direct-map|requires-transformer|conditional",
        "confidence": 0.95,
        "semantic_match": "Explanation of why this makes sense"
      }
    }
  ],
  "reasoning": "Overall workflow structure explanation",
  "confidence": 0.85
}`

	userPrompt := fmt.Sprintf(`Canvas Analysis:

Nodes (including inputs/outputs):
%s

Current Edges:
%s

Task: %s

KEY RULES:
- Preserve existing edges and suggest only NEW edges to add
- SKIP any connections where input/output types don't match
- If an activity has no outputs, it cannot be a source
- If an activity has no inputs, it cannot be a target
- Note any activities that are hard to connect (terminal activities, generators, etc)

Return ONLY valid JSON, no markdown code blocks.`, nodeDesc, edgeDesc, getReasoningTask(in.PreserveExisting))

	logger.logf("info", "Calling LLM reasoning (preserve_existing=%v)", in.PreserveExisting)

	// Call LLM
	client, err := llm.NewClient()
	if err != nil {
		return output, fmt.Errorf("failed to create LLM client: %w", err)
	}

	response, err := client.CreateMessage(ctx, llm.MessageInput{
		Model: statemachine.ModelSpec{
			ModelID: "reasoning", // Use reasoning model for complex analysis
		},
		SystemPrompt: systemPrompt,
		Messages: []llm.MessageParam{
			{
				Role:    "user",
				Content: userPrompt,
			},
		},
		AuthToken: in.AuthToken,
	})

	if err != nil {
		return output, fmt.Errorf("LLM reasoning failed: %w", err)
	}

	// Parse LLM response
	var reasonerResp struct {
		Edges      []db.WorkflowEdge `json:"edges"`
		Reasoning  string            `json:"reasoning"`
		Confidence float64           `json:"confidence"`
	}

	if err := json.Unmarshal([]byte(response), &reasonerResp); err != nil {
		logger.logf("warn", "Failed to parse LLM response as JSON: %v", err)
		// Try to extract from response text
		output.Reasoning = response
		output.Confidence = 0.5
		return output, fmt.Errorf("failed to parse LLM response: %w", err)
	}

	// Validate suggested edges
	nodeMap := make(map[string]bool)
	for _, n := range in.Nodes {
		nodeMap[n.ID] = true
	}

	validEdges := []db.WorkflowEdge{}
	for _, edge := range reasonerResp.Edges {
		if !nodeMap[edge.Source] {
			logger.logf("warn", "Suggested edge references unknown source: %s", edge.Source)
			continue
		}
		if !nodeMap[edge.Target] {
			logger.logf("warn", "Suggested edge references unknown target: %s", edge.Target)
			continue
		}
		// Don't suggest self-loops
		if edge.Source == edge.Target {
			logger.logf("warn", "Skipping self-loop: %s", edge.Source)
			continue
		}
		validEdges = append(validEdges, edge)
	}

	output.SuggestedEdges = validEdges
	output.Reasoning = reasonerResp.Reasoning
	output.Confidence = reasonerResp.Confidence

	// Check compatibility of suggested edges
	incompatibilities := CheckCanvasConnectivity(in.Nodes, validEdges)
	if len(incompatibilities) > 0 {
		output.IncompatibleEdges = incompatibilities
		logger.logf("warn", "Found %d incompatible edge connections", len(incompatibilities))

		// Generate user-friendly alerts
		for i, incompat := range incompatibilities {
			if i < 5 { // Limit to 5 alerts to avoid spam
				alert := fmt.Sprintf(
					"⚠️  %s → %s: %s. %s",
					incompat.Source, incompat.Target, incompat.Reason, incompat.Suggestion,
				)
				output.UserAlerts = append(output.UserAlerts, alert)
			}
		}
	}

	// Identify disconnected nodes
	disconnected := IdentifyDisconnectedNodes(in.Nodes, validEdges)
	if len(disconnected) > 0 {
		output.DisconnectedNodes = disconnected
		logger.logf("warn", "Found %d disconnected nodes", len(disconnected))

		for _, nodeID := range disconnected {
			var label string
			for _, node := range in.Nodes {
				if node.ID == nodeID {
					label = node.Label
					break
				}
			}
			alert := fmt.Sprintf(
				"🔌 Node '%s' has no connections. Consider adding edges or removing it.",
				label,
			)
			output.UserAlerts = append(output.UserAlerts, alert)
		}
	}

	logger.logf("info", "LLM suggested %d edges with confidence %.2f | %d incompatibilities | %d disconnected",
		len(validEdges), output.Confidence, len(incompatibilities), len(disconnected))

	return output, nil
}

// buildNodeDescriptions creates readable node descriptions for LLM (including schemas)
func buildNodeDescriptions(nodes []db.WorkflowNode) string {
	var desc string
	for i, node := range nodes {
		desc += fmt.Sprintf("%d. [%s] %s (type: %s)\n", i+1, node.ID, node.Label, node.Type)

		// Add input/output schema info
		if schema, err := getActivitySchema(node.Type); err == nil {
			if len(schema.Inputs) > 0 {
				desc += fmt.Sprintf("   INPUTS: %v\n", getInputNames(schema.Inputs))
			} else {
				desc += fmt.Sprintf("   INPUTS: none (generator/trigger)\n")
			}
			if len(schema.Outputs) > 0 {
				desc += fmt.Sprintf("   OUTPUTS: %v\n", getOutputNames(schema.Outputs))
			} else {
				desc += fmt.Sprintf("   OUTPUTS: none (terminal/sink)\n")
			}
		}

		if node.Data != nil {
			if b, err := json.MarshalIndent(node.Data, "   ", "  "); err == nil {
				desc += fmt.Sprintf("   CONFIG: %s\n", string(b))
			}
		}
	}
	return desc
}

// buildEdgeDescriptions creates readable edge descriptions for LLM
func buildEdgeDescriptions(edges []db.WorkflowEdge) string {
	if len(edges) == 0 {
		return "None"
	}
	var desc string
	for i, edge := range edges {
		desc += fmt.Sprintf("%d. %s → %s\n", i+1, edge.Source, edge.Target)
	}
	return desc
}

// getReasoningTask returns task description based on preservation mode
func getReasoningTask(preserveExisting bool) string {
	if preserveExisting {
		return "Keep all existing edges and suggest ONLY NEW edges to improve workflow"
	}
	return "Design optimal workflow by suggesting all connections and noting any redundant edges"
}