CHORUS/pkg/slurp/context/lightrag.go

package context

import (
	"context"
	"fmt"
	"strings"

	"chorus/pkg/mcp"
	"chorus/pkg/ucxl"
)

// LightRAGEnricher enriches context nodes with RAG-retrieved information
type LightRAGEnricher struct {
	client      *mcp.LightRAGClient
	defaultMode mcp.QueryMode
	enabled     bool
}

// NewLightRAGEnricher creates a new LightRAG context enricher
func NewLightRAGEnricher(client *mcp.LightRAGClient, defaultMode string) *LightRAGEnricher {
	if client == nil {
		return &LightRAGEnricher{enabled: false}
	}

	mode := mcp.QueryModeHybrid // Default to hybrid
	switch defaultMode {
	case "naive":
		mode = mcp.QueryModeNaive
	case "local":
		mode = mcp.QueryModeLocal
	case "global":
		mode = mcp.QueryModeGlobal
	case "hybrid":
		mode = mcp.QueryModeHybrid
	}

	return &LightRAGEnricher{
		client:      client,
		defaultMode: mode,
		enabled:     true,
	}
}

// EnrichContextNode enriches a ContextNode with LightRAG data
// This queries LightRAG for relevant information and adds it to the node's insights
func (e *LightRAGEnricher) EnrichContextNode(ctx context.Context, node *ContextNode) error {
	if !e.enabled || e.client == nil {
		return nil // No-op if not enabled
	}

	// Build query from node information
	query := e.buildQuery(node)
	if query == "" {
		return nil // Nothing to query
	}

	// Query LightRAG for context
	ragContext, err := e.client.GetContext(ctx, query, e.defaultMode)
	if err != nil {
		// Non-fatal - just log and continue
		return fmt.Errorf("lightrag query failed (non-fatal): %w", err)
	}

	// Add RAG context to insights if we got meaningful data
	if strings.TrimSpace(ragContext) != "" {
		insight := fmt.Sprintf("RAG Context: %s", strings.TrimSpace(ragContext))
		node.Insights = append(node.Insights, insight)

		// Update RAG confidence based on response quality
		// This is a simple heuristic - could be more sophisticated
		if len(ragContext) > 100 {
			node.RAGConfidence = 0.8
		} else if len(ragContext) > 50 {
			node.RAGConfidence = 0.6
		} else {
			node.RAGConfidence = 0.4
		}
	}

	return nil
}

// EnrichResolvedContext enriches a ResolvedContext with LightRAG data
// This is called after context resolution to add additional RAG-retrieved insights
func (e *LightRAGEnricher) EnrichResolvedContext(ctx context.Context, resolved *ResolvedContext) error {
	if !e.enabled || e.client == nil {
		return nil // No-op if not enabled
	}

	// Build query from resolved context
	query := fmt.Sprintf("Purpose: %s\nSummary: %s\nTechnologies: %s",
		resolved.Purpose,
		resolved.Summary,
		strings.Join(resolved.Technologies, ", "))

	// Query LightRAG
	ragContext, err := e.client.GetContext(ctx, query, e.defaultMode)
	if err != nil {
		return fmt.Errorf("lightrag query failed (non-fatal): %w", err)
	}

	// Add to insights if meaningful
	if strings.TrimSpace(ragContext) != "" {
		insight := fmt.Sprintf("RAG Enhancement: %s", strings.TrimSpace(ragContext))
		resolved.Insights = append(resolved.Insights, insight)

		// Boost confidence slightly if RAG provided good context
		if len(ragContext) > 100 {
			resolved.ResolutionConfidence = min(1.0, resolved.ResolutionConfidence*1.1)
		}
	}

	return nil
}

// EnrichBatchResolution enriches a batch resolution with LightRAG data
// Efficiently processes multiple addresses by batching queries where possible
func (e *LightRAGEnricher) EnrichBatchResolution(ctx context.Context, batch *BatchResolutionResult) error {
	if !e.enabled || e.client == nil {
		return nil // No-op if not enabled
	}

	// Enrich each resolved context
	for _, resolved := range batch.Results {
		if err := e.EnrichResolvedContext(ctx, resolved); err != nil {
			// Log error but continue with other contexts
			// Errors are non-fatal for enrichment
			continue
		}
	}

	return nil
}

// InsertContextNode inserts a context node into LightRAG for future retrieval
// This builds the knowledge base over time as contexts are created
func (e *LightRAGEnricher) InsertContextNode(ctx context.Context, node *ContextNode) error {
	if !e.enabled || e.client == nil {
		return nil // No-op if not enabled
	}

	// Build text representation of the context node
	text := e.buildTextRepresentation(node)
	description := fmt.Sprintf("Context for %s: %s", node.Path, node.Summary)

	// Insert into LightRAG
	if err := e.client.Insert(ctx, text, description); err != nil {
		return fmt.Errorf("failed to insert context into lightrag: %w", err)
	}

	return nil
}

// IsEnabled returns whether LightRAG enrichment is enabled
func (e *LightRAGEnricher) IsEnabled() bool {
	return e.enabled
}

// buildQuery constructs a search query from a ContextNode
func (e *LightRAGEnricher) buildQuery(node *ContextNode) string {
	var parts []string

	if node.Purpose != "" {
		parts = append(parts, node.Purpose)
	}

	if node.Summary != "" {
		parts = append(parts, node.Summary)
	}

	if len(node.Technologies) > 0 {
		parts = append(parts, strings.Join(node.Technologies, " "))
	}

	if len(node.Tags) > 0 {
		parts = append(parts, strings.Join(node.Tags, " "))
	}

	return strings.Join(parts, " ")
}

// buildTextRepresentation builds a text representation for storage in LightRAG
func (e *LightRAGEnricher) buildTextRepresentation(node *ContextNode) string {
	var builder strings.Builder

	builder.WriteString(fmt.Sprintf("Path: %s\n", node.Path))
	builder.WriteString(fmt.Sprintf("UCXL Address: %s\n", node.UCXLAddress.String()))
	builder.WriteString(fmt.Sprintf("Summary: %s\n", node.Summary))
	builder.WriteString(fmt.Sprintf("Purpose: %s\n", node.Purpose))

	if len(node.Technologies) > 0 {
		builder.WriteString(fmt.Sprintf("Technologies: %s\n", strings.Join(node.Technologies, ", ")))
	}

	if len(node.Tags) > 0 {
		builder.WriteString(fmt.Sprintf("Tags: %s\n", strings.Join(node.Tags, ", ")))
	}

	if len(node.Insights) > 0 {
		builder.WriteString("Insights:\n")
		for _, insight := range node.Insights {
			builder.WriteString(fmt.Sprintf("  - %s\n", insight))
		}
	}

	if node.Language != nil {
		builder.WriteString(fmt.Sprintf("Language: %s\n", *node.Language))
	}

	return builder.String()
}

func min(a, b float64) float64 {
	if a < b {
		return a
	}
	return b
}