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
}

// CanvasReasonerOutput returns suggested edges and reasoning
type CanvasReasonerOutput struct {
	SuggestedEdges []db.WorkflowEdge `json:"suggested_edges"`  // New edges to add
	RemovedEdges   []db.WorkflowEdge `json:"removed_edges,omitempty"` // Edges to remove (if redesign)
	Reasoning      string            `json:"reasoning"`        // LLM explanation
	Confidence     float64           `json:"confidence"`       // 0.0-1.0
}

// 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
	systemPrompt := `You are a workflow automation expert. Analyze the following activities and suggest logical connections (edges) between them based on:
1. Activity input/output compatibility
2. Logical execution order
3. Data flow requirements
4. Common workflow patterns

Respond with JSON containing:
{
  "edges": [{"source": "node-1", "target": "node-2"}, ...],
  "reasoning": "explanation of why these connections make sense",
  "confidence": 0.85
}`

	userPrompt := fmt.Sprintf(`Canvas Analysis:

Nodes:
%s

Current Edges:
%s

Task: %s

Preserve existing edges and suggest only NEW edges to add.
If any existing edges don't make sense, note them but keep them unless explicitly wrong.

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

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

	return output, nil
}

// buildNodeDescriptions creates readable node descriptions for LLM
func buildNodeDescriptions(nodes []db.WorkflowNode) string {
	var desc string
	for i, node := range nodes {
		desc += fmt.Sprintf("%d. %s (type: %s)\n", i+1, node.ID, node.Type)
		desc += fmt.Sprintf("   Label: %s\n", node.Label)
		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"
}