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 }