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
}

// getStringFromMap safely extracts a string from a map
func getStringFromMap(m map[string]interface{}, key string) string {
	if m == nil {
		return ""
	}
	if v, ok := m[key].(string); ok {
		return v
	}
	return ""
}

// firstNonEmpty returns the first non-empty string from the list
func firstNonEmpty(values ...string) string {
	for _, v := range values {
		if v != "" {
			return v
		}
	}
	return ""
}

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

	// Extract schedule from multiple sources
	schedule := firstNonEmpty(
		intent.CronSchedule,
		getStringFromMap(intent.Parameters, "cronSchedule"),
		getStringFromMap(spec.Input, "cronSchedule"),
	)

	// Extract timezone from multiple sources, default to UTC
	timezone := firstNonEmpty(
		intent.CronTimezone,
		getStringFromMap(intent.Parameters, "cronTimezone"),
		getStringFromMap(spec.Input, "cronTimezone"),
		"UTC",
	)

	// Clean cron fields from input
	delete(spec.Input, "cronSchedule")
	delete(spec.Input, "cronTimezone")

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

// isCommonInputField checks if field name is a common workflow input
func isCommonInputField(name string) bool {
	switch name {
	case "repo", "path", "branch":
		return true
	}
	return false
}

// paramResolver resolves activity parameters from multiple sources
type paramResolver struct {
	intent    *Intent
	kb        *KnowledgeBase
	prevState string
}

// resolve finds parameter value from intent, previous output, default, or input ref
func (r *paramResolver) resolve(inputName string, inputDef InputField) interface{} {
	// 1. From intent parameters
	if val, ok := r.intent.Parameters[inputName]; ok {
		return val
	}

	// 2. From previous state output
	if val := r.fromPrevOutput(inputName); val != nil {
		return val
	}

	// 3. Default value
	if inputDef.Default != nil {
		return inputDef.Default
	}

	// 4. Input reference for common fields
	if isCommonInputField(inputName) {
		return fmt.Sprintf("${input.%s}", inputName)
	}

	return nil
}

// fromPrevOutput checks if previous activity has matching output
func (r *paramResolver) fromPrevOutput(inputName string) interface{} {
	if r.prevState == "" {
		return nil
	}
	prevActDef := r.kb.GetActivity(r.prevState)
	if prevActDef == nil {
		return nil
	}
	for outName := range prevActDef.Outputs {
		if outName == inputName || strings.EqualFold(outName, inputName) {
			return fmt.Sprintf("${%s.output.%s}", r.prevState, outName)
		}
	}
	return nil
}

// buildParameters creates parameter map for activity
func (r *LLMRouter) buildParameters(act *ActivityMetadata, intent *Intent, stateIndex int) map[string]interface{} {
	var prevState string
	if stateIndex > 0 {
		prevState = intent.Activities[stateIndex-1]
	}

	resolver := &paramResolver{
		intent:    intent,
		kb:        r.knowledgeBase,
		prevState: prevState,
	}

	params := make(map[string]interface{})
	for name, def := range act.Inputs {
		if val := resolver.resolve(name, def); val != nil {
			params[name] = val
		}
	}
	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.`