poimen-workflows/action/memory.go

package action

import (
	"context"
	"fmt"
	"os"

	"github.com/rockliang/poimen/workflows/internal/memory"
)

// RetrieveMemoryInput input for RetrieveMemoryActivity
type RetrieveMemoryInput struct {
	// Query semantic search query
	Query string `json:"query"`

	// Project memory project (default: "poimen")
	Project string `json:"project,omitempty"`

	// Scope retrieval scope: "skills", "lessons", "all" (default: "all")
	Scope string `json:"scope,omitempty"`

	// Limit max results (default: 10)
	Limit int `json:"limit,omitempty"`

	// LevelFilter filter by level: L1, L2, R (reference)
	LevelFilter []string `json:"levelFilter,omitempty"`

	// Tool tool context for skill matching
	Tool string `json:"tool,omitempty"`

	// Task task description for context retrieval
	Task string `json:"task,omitempty"`
}

// RetrieveMemoryOutput output from RetrieveMemoryActivity
type RetrieveMemoryOutput struct {
	// Skills relevant skills found
	Skills []MemorySkill `json:"skills"`

	// Lessons relevant lessons/knowledge found
	Lessons []MemoryLesson `json:"lessons"`

	// References reference documents found
	References []MemoryReference `json:"references"`

	// TotalResults total results found
	TotalResults int `json:"totalResults"`

	// Budget token budget info
	Budget MemoryBudget `json:"budget"`
}

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

// MemoryLesson lesson from memory
type MemoryLesson struct {
	ID         string  `json:"id"`
	Text       string  `json:"text"`
	Level      string  `json:"level"`
	Score      float32 `json:"score"`
	Breadcrumb string  `json:"breadcrumb,omitempty"`
}

// MemoryReference reference document from memory
type MemoryReference struct {
	ID         string  `json:"id"`
	Text       string  `json:"text"`
	Score      float32 `json:"score"`
	Breadcrumb string  `json:"breadcrumb,omitempty"`
}

// MemoryBudget token budget tracking
type MemoryBudget struct {
	Requested int `json:"requested"`
	Used      int `json:"used"`
}

// RetrieveMemoryActivity retrieves relevant knowledge from poimen-memory
func RetrieveMemoryActivity(ctx context.Context, in RetrieveMemoryInput) (RetrieveMemoryOutput, error) {
	logger := newActivityLogger(ctx)
	logger.Info("RetrieveMemoryActivity started", "query", in.Query, "scope", in.Scope)

	output := RetrieveMemoryOutput{
		Skills:     []MemorySkill{},
		Lessons:    []MemoryLesson{},
		References: []MemoryReference{},
	}

	// Get memory service URL and token
	baseURL := os.Getenv("POIMEN_MEMORY_URL")
	if baseURL == "" {
		baseURL = "http://poimen-memory.poimen.svc.cluster.local:8080"
	}

	token := os.Getenv("POIMEN_MEMORY_TOKEN")
	// Token optional for internal cluster access

	// Set defaults
	project := in.Project
	if project == "" {
		project = "poimen"
	}

	scope := in.Scope
	if scope == "" {
		scope = "all"
	}

	limit := in.Limit
	if limit == 0 {
		limit = 10
	}

	client := memory.NewClient(baseURL, token)

	// If tool/task provided, use Context API for skill matching
	if in.Tool != "" || in.Task != "" {
		contextResp, err := client.Context(ctx, &memory.ContextRequest{
			Project:         project,
			Tool:            in.Tool,
			Task:            in.Task,
			SignatureSource: in.Query,
			Scope:           "tool_context",
			Budget:          8192,
		})
		if err != nil {
			logger.Warn("Context retrieval failed, falling back to query", "error", err)
		} else {
			// Extract skills
			for _, skill := range contextResp.Skills {
				output.Skills = append(output.Skills, MemorySkill{
					Name: skill.Name,
					Why:  skill.Why,
				})
			}

			// Extract lessons
			for i, lesson := range contextResp.Lessons {
				output.Lessons = append(output.Lessons, MemoryLesson{
					ID:         fmt.Sprintf("ctx-%d", i),
					Text:       lesson.Text,
					Level:      lesson.Level,
					Score:      lesson.Score,
					Breadcrumb: "",
				})
			}

			output.Budget = MemoryBudget{
				Requested: contextResp.Budget.Requested,
				Used:      contextResp.Budget.Used,
			}
			output.TotalResults = len(output.Skills) + len(output.Lessons)
		}
	}

	// Also do semantic query for additional context
	if scope == "all" || scope == "lessons" || scope == "references" {
		levelFilter := in.LevelFilter
		if len(levelFilter) == 0 {
			levelFilter = []string{"L1", "L2"}
		}

		queryResp, err := client.Query(ctx, &memory.QueryRequest{
			Project:     project,
			Query:       in.Query,
			LevelFilter: levelFilter,
			Limit:       limit,
			Scope:       "all",
		})
		if err != nil {
			logger.Warn("Query failed", "error", err)
		} else {
			for _, result := range queryResp.Results {
				if result.Level == "R" {
					output.References = append(output.References, MemoryReference{
						ID:         result.ID,
						Text:       result.Text,
						Score:      result.Score,
						Breadcrumb: result.Breadcrumb,
					})
				} else {
					// Avoid duplicates from Context call
					found := false
					for _, existing := range output.Lessons {
						if existing.ID == result.ID {
							found = true
							break
						}
					}
					if !found {
						output.Lessons = append(output.Lessons, MemoryLesson{
							ID:         result.ID,
							Text:       result.Text,
							Level:      result.Level,
							Score:      result.Score,
							Breadcrumb: result.Breadcrumb,
						})
					}
				}
			}
			output.TotalResults = len(output.Skills) + len(output.Lessons) + len(output.References)
		}
	}

	logger.Info("RetrieveMemoryActivity completed",
		"skills", len(output.Skills),
		"lessons", len(output.Lessons),
		"references", len(output.References))

	return output, nil
}

// FormatMemoryForPrompt formats memory output for LLM prompt injection
func FormatMemoryForPrompt(mem RetrieveMemoryOutput) string {
	if mem.TotalResults == 0 {
		return ""
	}

	var result string

	if len(mem.Skills) > 0 {
		result += "\n## Relevant Skills\n"
		for _, skill := range mem.Skills {
			result += fmt.Sprintf("- **%s**: %s\n", skill.Name, skill.Why)
		}
	}

	if len(mem.Lessons) > 0 {
		result += "\n## Relevant Knowledge\n"
		for _, lesson := range mem.Lessons {
			result += fmt.Sprintf("- [%s] %s\n", lesson.Level, truncate(lesson.Text, 200))
		}
	}

	if len(mem.References) > 0 {
		result += "\n## Reference Documents\n"
		for _, ref := range mem.References {
			result += fmt.Sprintf("- %s\n", truncate(ref.Text, 200))
		}
	}

	return result
}

func truncate(s string, maxLen int) string {
	if len(s) <= maxLen {
		return s
	}
	return s[:maxLen] + "..."
}