poimen-workflows/internal/routing/llm_router.go

package routing

import (
	"context"
	"encoding/json"
	"fmt"
	"regexp"
	"strings"
)

// LLMRouterInput is input to the llm-router activity
type LLMRouterInput struct {
	Message       string                 `json:"message"`
	Context       map[string]interface{} `json:"context,omitempty"`       // Optional context (repo, branch, etc)
	MemoryContext *MemoryContext         `json:"memoryContext,omitempty"` // Optional memory retrieval results
	UseMemory     bool                   `json:"useMemory,omitempty"`     // Enable memory retrieval (default: false)
}

// MemoryContext holds retrieved memory for prompt injection
type MemoryContext struct {
	Skills     []MemorySkill     `json:"skills"`
	Lessons    []MemoryLesson    `json:"lessons"`
	References []MemoryReference `json:"references"`
}

// MemorySkill from memory service
type MemorySkill struct {
	Name        string `json:"name"`
	Description string `json:"description"`
	Why         string `json:"why,omitempty"`
}

// MemoryLesson from memory service
type MemoryLesson struct {
	ID    string `json:"id"`
	Text  string `json:"text"`
	Level string `json:"level"`
}

// MemoryReference from memory service
type MemoryReference struct {
	ID   string `json:"id"`
	Text string `json:"text"`
}

// LLMRouterOutput is output from the llm-router activity
type LLMRouterOutput struct {
	Spec     *WorkflowSpec     `json:"spec,omitempty"`
	CronSpec *CronWorkflowSpec `json:"cronSpec,omitempty"`
	IsCron   bool              `json:"isCron"`
	Error    string            `json:"error,omitempty"`
}

// LLMRouter orchestrates intent analysis and spec generation
type LLMRouter struct {
	client        *LLMClient
	knowledgeBase *KnowledgeBase
}

// NewLLMRouter creates a new LLM router
func NewLLMRouter(kb *KnowledgeBase) (*LLMRouter, error) {
	return &LLMRouter{
		client:        NewLLMClient(),
		knowledgeBase: kb,
	}, nil
}

// Route analyzes user message and generates appropriate workflow spec
func (r *LLMRouter) Route(ctx context.Context, input LLMRouterInput) (*LLMRouterOutput, error) {
	// 1. Analyze intent using LLM
	intent, err := r.analyzeIntent(ctx, input)
	if err != nil {
		return nil, fmt.Errorf("intent analysis failed: %w", err)
	}

	// 2. Build workflow spec based on intent
	if intent.IsCron {
		cronSpec, err := r.buildCronSpec(intent, input)
		if err != nil {
			return nil, fmt.Errorf("cron spec build failed: %w", err)
		}
		return &LLMRouterOutput{
			CronSpec: cronSpec,
			IsCron:   true,
		}, nil
	}

	spec, err := r.buildSpec(intent, input)
	if err != nil {
		return nil, fmt.Errorf("spec build failed: %w", err)
	}

	return &LLMRouterOutput{
		Spec:   spec,
		IsCron: false,
	}, nil
}

// Intent represents analyzed user intent
type Intent struct {
	Activities     []string               `json:"activities"`      // Selected activity names
	Parameters     map[string]interface{} `json:"parameters"`      // Extracted parameters
	IsCron         bool                   `json:"isCron"`          // Is scheduled workflow?
	CronSchedule   string                 `json:"cronSchedule"`    // Cron expression if scheduled
	CronTimezone   string                 `json:"cronTimezone"`    // Timezone for cron
	WorkflowName   string                 `json:"workflowName"`    // Generated workflow name
	ErrorHandling  string                 `json:"errorHandling"`   // "retry", "fail-fast", "continue"
}

// analyzeIntent uses LLM to understand user request
func (r *LLMRouter) analyzeIntent(ctx context.Context, input LLMRouterInput) (*Intent, error) {
	// Build prompt with knowledge base context
	prompt := r.buildIntentPrompt(input)

	// Call LLM
	response, err := r.client.Chat(ctx, intentSystemPrompt, prompt)
	if err != nil {
		return nil, fmt.Errorf("LLM call failed: %w", err)
	}

	// Parse LLM response
	intent, err := parseIntentResponse(response)
	if err != nil {
		return nil, fmt.Errorf("failed to parse intent: %w", err)
	}

	// Validate activities exist
	for _, actName := range intent.Activities {
		if !r.knowledgeBase.HasActivity(actName) {
			return nil, fmt.Errorf("unknown activity: %s", actName)
		}
	}

	return intent, nil
}

// buildIntentPrompt creates the prompt for intent analysis
func (r *LLMRouter) buildIntentPrompt(input LLMRouterInput) string {
	// Get activity summaries
	var activityList strings.Builder
	for _, act := range r.knowledgeBase.Activities {
		activityList.WriteString(fmt.Sprintf("- %s: %s (category: %s, timeout: %s, flaky: %v)\n",
			act.Name, act.Description, act.Category,
			act.Constraints.DefaultTimeout, act.Constraints.IsFlaky))
	}

	// Build context string
	contextStr := ""
	if len(input.Context) > 0 {
		ctxBytes, _ := json.Marshal(input.Context)
		contextStr = fmt.Sprintf("\nProvided context: %s", string(ctxBytes))
	}

	// Build memory context string
	memoryStr := r.formatMemoryContext(input.MemoryContext)

	return fmt.Sprintf(`User request: %s
%s%s
Available activities:
%s
Analyze the request and output JSON with:
- activities: ordered list of activity names to execute
- parameters: extracted parameters from request (repo URL, branch, etc)
- isCron: true if user wants scheduled/recurring execution
- cronSchedule: cron expression if scheduled (e.g., "0 2 * * *" for 2 AM daily)
- cronTimezone: timezone (default "UTC")
- workflowName: short descriptive name
- errorHandling: "retry" (default), "fail-fast", or "continue"

Output ONLY valid JSON, no explanation.`, input.Message, contextStr, memoryStr, activityList.String())
}

// formatMemoryContext formats memory context for prompt injection
func (r *LLMRouter) formatMemoryContext(mem *MemoryContext) string {
	if mem == nil {
		return ""
	}

	var sb strings.Builder

	if len(mem.Skills) > 0 {
		sb.WriteString("\n\nRelevant skills from memory:\n")
		for _, skill := range mem.Skills {
			if skill.Why != "" {
				sb.WriteString(fmt.Sprintf("- %s: %s (reason: %s)\n", skill.Name, skill.Description, skill.Why))
			} else {
				sb.WriteString(fmt.Sprintf("- %s: %s\n", skill.Name, skill.Description))
			}
		}
	}

	if len(mem.Lessons) > 0 {
		sb.WriteString("\nRelevant knowledge from memory:\n")
		for _, lesson := range mem.Lessons {
			text := lesson.Text
			if len(text) > 300 {
				text = text[:300] + "..."
			}
			sb.WriteString(fmt.Sprintf("- [%s] %s\n", lesson.Level, text))
		}
	}

	if len(mem.References) > 0 {
		sb.WriteString("\nReference documents:\n")
		for _, ref := range mem.References {
			text := ref.Text
			if len(text) > 200 {
				text = text[:200] + "..."
			}
			sb.WriteString(fmt.Sprintf("- %s\n", text))
		}
	}

	return sb.String()
}

// parseIntentResponse extracts Intent from LLM response
func parseIntentResponse(response string) (*Intent, error) {
	// Try to extract JSON from response
	response = strings.TrimSpace(response)
	
	// Handle markdown code blocks
	if strings.HasPrefix(response, "```") {
		re := regexp.MustCompile("```(?:json)?\\s*([\\s\\S]*?)```")
		matches := re.FindStringSubmatch(response)
		if len(matches) > 1 {
			response = strings.TrimSpace(matches[1])
		}
	}

	var intent Intent
	if err := json.Unmarshal([]byte(response), &intent); err != nil {
		return nil, fmt.Errorf("invalid JSON from LLM: %w\nResponse: %s", err, response)
	}

	// Set defaults
	if intent.CronTimezone == "" {
		intent.CronTimezone = "UTC"
	}
	if intent.ErrorHandling == "" {
		intent.ErrorHandling = "retry"
	}
	if intent.WorkflowName == "" {
		intent.WorkflowName = "generated-workflow"
	}

	return &intent, nil
}

// buildSpec creates WorkflowSpec from intent
func (r *LLMRouter) buildSpec(intent *Intent, input LLMRouterInput) (*WorkflowSpec, error) {
	if len(intent.Activities) == 0 {
		return nil, fmt.Errorf("no activities selected")
	}

	states := make([]State, 0, len(intent.Activities)+1)

	// Build states for each activity
	for i, actName := range intent.Activities {
		act := r.knowledgeBase.GetActivity(actName)
		
		state := State{
			Name:     actName,
			Type:     StateTypeTask,
			Resource: actName,
			Parameters: r.buildParameters(act, intent, i),
			Timeout:  act.Constraints.DefaultTimeout,
			Retry:    r.buildRetryPolicy(act, intent),
		}

		// Set next state or end
		if i < len(intent.Activities)-1 {
			state.Next = intent.Activities[i+1]
		} else {
			state.End = true
		}

		// Add catch clause for flaky activities
		if act.Constraints.IsFlaky && intent.ErrorHandling != "fail-fast" {
			state.Catch = []CatchClause{
				{
					ErrorEquals: []string{"ActivityError", "TimeoutError"},
					Next:        "HandleError",
				},
			}
		}

		states = append(states, state)
	}

	// Add error handler if needed
	hasFlaky := false
	for _, actName := range intent.Activities {
		act := r.knowledgeBase.GetActivity(actName)
		if act != nil && act.Constraints.IsFlaky {
			hasFlaky = true
			break
		}
	}
	if hasFlaky && intent.ErrorHandling != "fail-fast" {
		states = append(states, State{
			Name:  "HandleError",
			Type:  StateTypeFail,
			Error: "WorkflowError",
			Cause: "Activity failed after retries",
		})
	}

	// Build input map
	inputMap := make(map[string]interface{})
	for k, v := range intent.Parameters {
		inputMap[k] = v
	}
	for k, v := range input.Context {
		if _, exists := inputMap[k]; !exists {
			inputMap[k] = v
		}
	}

	return &WorkflowSpec{
		Name:   intent.WorkflowName,
		Input:  inputMap,
		States: states,
	}, nil
}

// buildCronSpec creates CronWorkflowSpec from intent
func (r *LLMRouter) buildCronSpec(intent *Intent, input LLMRouterInput) (*CronWorkflowSpec, error) {
	// First build regular spec
	spec, err := r.buildSpec(intent, input)
	if err != nil {
		return nil, err
	}

	// Get schedule - check both intent and parameters (LLM sometimes puts it in parameters)
	schedule := intent.CronSchedule
	if schedule == "" {
		if sched, ok := intent.Parameters["cronSchedule"].(string); ok {
			schedule = sched
		}
	}
	if schedule == "" {
		if sched, ok := spec.Input["cronSchedule"].(string); ok {
			schedule = sched
			delete(spec.Input, "cronSchedule") // Remove from input
		}
	}

	// Get timezone
	timezone := intent.CronTimezone
	if timezone == "" {
		if tz, ok := intent.Parameters["cronTimezone"].(string); ok {
			timezone = tz
		}
	}
	if timezone == "" {
		if tz, ok := spec.Input["cronTimezone"].(string); ok {
			timezone = tz
			delete(spec.Input, "cronTimezone") // Remove from input
		}
	}
	if timezone == "" {
		timezone = "UTC"
	}

	return &CronWorkflowSpec{
		Name:          spec.Name,
		Type:          "CronWorkflow",
		Schedule:      schedule,
		Timezone:      timezone,
		Input:         spec.Input,
		States:        spec.States,
		MaxConcurrent: 1,
		Timeout:       "1h",
		EnableHistory: true,
	}, nil
}

// buildParameters creates parameter map for activity
func (r *LLMRouter) buildParameters(act *ActivityMetadata, intent *Intent, stateIndex int) map[string]interface{} {
	params := make(map[string]interface{})

	for inputName, inputDef := range act.Inputs {
		// Check if parameter was extracted from intent
		if val, ok := intent.Parameters[inputName]; ok {
			params[inputName] = val
			continue
		}

		// Check for JSONPath reference from previous state
		if stateIndex > 0 {
			prevAct := intent.Activities[stateIndex-1]
			prevActDef := r.knowledgeBase.GetActivity(prevAct)
			
			// Look for matching output from previous activity
			for outName := range prevActDef.Outputs {
				if outName == inputName || strings.EqualFold(outName, inputName) {
					params[inputName] = fmt.Sprintf("${%s.output.%s}", prevAct, outName)
					break
				}
			}
		}

		// Use default if available
		if params[inputName] == nil && inputDef.Default != nil {
			params[inputName] = inputDef.Default
		}

		// Use input reference for common fields
		if params[inputName] == nil {
			if inputName == "repo" || inputName == "path" || inputName == "branch" {
				params[inputName] = fmt.Sprintf("${input.%s}", inputName)
			}
		}
	}

	return params
}

// buildRetryPolicy creates retry policy based on activity constraints
func (r *LLMRouter) buildRetryPolicy(act *ActivityMetadata, intent *Intent) *RetryPolicy {
	if intent.ErrorHandling == "fail-fast" {
		return &RetryPolicy{
			MaxAttempts:     1,
			BackoffRate:     1.0,
			InitialInterval: "1s",
		}
	}

	return &RetryPolicy{
		MaxAttempts:     int32(act.Constraints.RecommendedRetries),
		BackoffRate:     act.Constraints.RetryBackoff,
		InitialInterval: "1s",
		MaxInterval:     "30s",
	}
}

const intentSystemPrompt = `You are an intelligent workflow router. Your job is to:
1. Understand what the user wants to accomplish
2. Select the appropriate activities from the available list
3. Order them correctly based on dependencies
4. Extract any parameters mentioned (URLs, branches, etc)
5. Detect if user wants scheduled/recurring execution
6. Use any relevant knowledge from memory to inform your decisions

Rules:
- Always include CloneRepoActivity first if any analysis activity is needed
- Order activities respecting dependencies
- If user mentions "daily", "every hour", "weekly", etc → set isCron=true and cronSchedule
- Common cron patterns: "0 2 * * *" (2 AM daily), "0 * * * *" (hourly), "0 0 * * 0" (weekly Sunday)
- Extract repo URLs, branch names, severity levels from the message
- workflowName should be short and descriptive (kebab-case)
- If memory context includes relevant skills or lessons, incorporate that knowledge
- Skills from memory may suggest specific activity parameters or ordering

Output ONLY valid JSON.`