diff --git a/.DS_Store b/.DS_Store index 1d66667..6f2f7d5 100644 Binary files a/.DS_Store and b/.DS_Store differ diff --git a/config/config.go b/config/config.go new file mode 100644 index 0000000..56e12ce --- /dev/null +++ b/config/config.go @@ -0,0 +1,132 @@ +package config + +import ( + "encoding/json" + "os" + "path/filepath" + "time" +) + +// Config holds all configuration for the RAG system +type Config struct { + // Provider settings + Provider string + Model string + APIKeys map[string]string + + // Collection settings + Collection string + + // Search settings + SearchStrategy string + DefaultTopK int + DefaultMinScore float64 + DefaultSearchParams map[string]interface{} + EnableReRanking bool + RRFConstant float64 + + // Document processing settings + DefaultChunkSize int + DefaultChunkOverlap int + DefaultBatchSize int + DefaultIndexType string + + // Vector store settings + VectorDBConfig map[string]interface{} + + // Timeouts and retries + Timeout time.Duration + MaxRetries int + + // Additional settings + ExtraHeaders map[string]string +} + +// LoadConfig loads configuration from a file or environment +func LoadConfig() (*Config, error) { + // Default configuration + cfg := &Config{ + Provider: "milvus", + Model: "text-embedding-3-small", + Collection: "documents", + SearchStrategy: "dense", + DefaultTopK: 5, + DefaultMinScore: 0.7, + DefaultChunkSize: 512, + DefaultChunkOverlap: 50, + DefaultBatchSize: 100, + DefaultIndexType: "HNSW", + DefaultSearchParams: map[string]interface{}{ + "ef": 64, + }, + EnableReRanking: false, + RRFConstant: 60, + Timeout: 30 * time.Second, + MaxRetries: 3, + APIKeys: make(map[string]string), + ExtraHeaders: make(map[string]string), + VectorDBConfig: make(map[string]interface{}), + } + + // Try to load from config file + configFile := os.Getenv("RAGGO_CONFIG") + if configFile == "" { + // Try default locations + home, err := os.UserHomeDir() + if err == nil { + candidates := []string{ + filepath.Join(home, ".raggo", "config.json"), + filepath.Join(home, ".config", "raggo", "config.json"), + "raggo.json", + } + + for _, candidate := range candidates { + if _, err := os.Stat(candidate); err == nil { + configFile = candidate + break + } + } + } + } + + if configFile != "" { + data, err := os.ReadFile(configFile) + if err == nil { + if err := json.Unmarshal(data, cfg); err != nil { + return nil, err + } + } + } + + // Override with environment variables + if provider := os.Getenv("RAGGO_PROVIDER"); provider != "" { + cfg.Provider = provider + } + if model := os.Getenv("RAGGO_MODEL"); model != "" { + cfg.Model = model + } + if collection := os.Getenv("RAGGO_COLLECTION"); collection != "" { + cfg.Collection = collection + } + if apiKey := os.Getenv("RAGGO_API_KEY"); apiKey != "" { + cfg.APIKeys[cfg.Provider] = apiKey + } + + return cfg, nil +} + +// Save saves the configuration to a file +func (c *Config) Save(path string) error { + data, err := json.MarshalIndent(c, "", " ") + if err != nil { + return err + } + + // Ensure directory exists + dir := filepath.Dir(path) + if err := os.MkdirAll(dir, 0755); err != nil { + return err + } + + return os.WriteFile(path, data, 0644) +} diff --git a/docs/chatbot_example.md b/docs/chatbot_example.md new file mode 100644 index 0000000..eda6048 --- /dev/null +++ b/docs/chatbot_example.md @@ -0,0 +1,305 @@ +# Memory-Enhanced Chatbot Example Documentation + +## Overview +The Memory Enhancer Example demonstrates how to build a sophisticated chatbot using RagGo's RAG and Memory Context features. This implementation showcases advanced capabilities like context retention, document-based knowledge, and interactive conversations. + +## Features +- Document-based knowledge integration +- Context-aware responses +- Memory retention across conversations +- Interactive command-line interface +- Vector database integration +- Hybrid search capabilities + +## Implementation Details + +### 1. Setup and Configuration +```go +// Initialize LLM with OpenAI +llm, err := gollm.NewLLM( + gollm.SetProvider("openai"), + gollm.SetModel("gpt-4o-mini"), + gollm.SetAPIKey(apiKey), + gollm.SetLogLevel(gollm.LogLevelInfo), +) + +// Configure Vector Database +vectorDB, err := raggo.NewVectorDB( + raggo.WithType("milvus"), + raggo.WithAddress("localhost:19530"), +) + +// Create Memory Context +memoryContext, err := raggo.NewMemoryContext(apiKey, + raggo.MemoryCollection("tech_docs"), + raggo.MemoryTopK(5), + raggo.MemoryMinScore(0.01), + raggo.MemoryStoreLastN(10), + raggo.MemoryStoreRAGInfo(true), +) +``` + +### 2. Document Management +```go +// Define document sources +docsDir := filepath.Join("examples", "chat", "docs") +docs := []string{ + filepath.Join(docsDir, "microservices.txt"), + filepath.Join(docsDir, "vector_databases.txt"), + filepath.Join(docsDir, "rag_systems.txt"), + filepath.Join(docsDir, "golang_basics.txt"), + filepath.Join(docsDir, "embeddings.txt"), +} + +// Process and store documents +for _, doc := range docs { + content, err := os.ReadFile(doc) + if err != nil { + log.Printf("Warning: Failed to read %s: %v", doc, err) + continue + } + err = memoryContext.Store(ctx, filepath.Base(doc), string(content)) + if err != nil { + log.Printf("Warning: Failed to store %s: %v", doc, err) + } +} +``` + +### 3. Interactive Chat Loop +```go +// Initialize scanner for user input +scanner := bufio.NewScanner(os.Stdin) + +// Main chat loop +for { + fmt.Print("\nEnter your question (or 'quit' to exit): ") + if !scanner.Scan() { + break + } + + query := scanner.Text() + if strings.ToLower(query) == "quit" { + break + } + + // Process query with context + response, err := memoryContext.ProcessWithContext(ctx, query) + if err != nil { + log.Printf("Error processing query: %v", err) + continue + } + + fmt.Printf("\nResponse: %s\n", response) +} +``` + +## Usage Guide + +### Prerequisites +1. Environment Setup + ```bash + # Set OpenAI API key + export OPENAI_API_KEY=your_api_key + + # Start Milvus database + docker-compose up -d + ``` + +2. Document Preparation + - Place your knowledge base documents in the `examples/chat/docs` directory + - Supported formats: `.txt` files + - Recommended document size: 1-10KB per file + +### Running the Example +```bash +# Navigate to the project directory +cd raggo + +# Run the example +go run examples/memory_enhancer_example.go +``` + +### Example Interactions +``` +Enter your question: What are microservices? +Response: Based on the documentation, microservices are a software architecture pattern where applications are built as a collection of small, independent services. Each service runs in its own process and communicates through well-defined APIs... + +Enter your question: How do they handle data? +Response: [Context-aware response about microservice data handling, building on the previous question] +``` + +## Configuration Options + +### Memory Context Settings +```go +memoryContext, err := raggo.NewMemoryContext(apiKey, + // Collection name for vector storage + raggo.MemoryCollection("tech_docs"), + + // Number of similar contexts to retrieve + raggo.MemoryTopK(5), + + // Minimum similarity score (0-1) + raggo.MemoryMinScore(0.01), + + // Number of recent interactions to store + raggo.MemoryStoreLastN(10), + + // Store RAG information for better context + raggo.MemoryStoreRAGInfo(true), +) +``` + +### Vector Database Settings +```go +vectorDB, err := raggo.NewVectorDB( + // Database type + raggo.WithType("milvus"), + + // Connection address + raggo.WithAddress("localhost:19530"), +) +``` + +## Best Practices + +### 1. Document Organization +- Split large documents into smaller, focused files +- Use clear, descriptive filenames +- Maintain consistent document format +- Regular updates to knowledge base + +### 2. Memory Management +- Configure appropriate `TopK` for your use case +- Set `MinScore` based on required accuracy +- Adjust `StoreLastN` based on conversation length +- Monitor memory usage + +### 3. Error Handling +- Implement graceful error recovery +- Log errors appropriately +- Provide user-friendly error messages +- Handle connection issues + +### 4. Performance Optimization +- Batch process documents when possible +- Monitor API rate limits +- Use appropriate chunk sizes +- Implement caching if needed + +## Customization Guide + +### 1. Adding New Document Types +```go +// Example: Add PDF support +func processPDFDocument(path string) (string, error) { + // PDF processing logic + return content, nil +} +``` + +### 2. Custom Response Processing +```go +// Example: Add response formatting +func formatResponse(response string) string { + // Response formatting logic + return formattedResponse +} +``` + +### 3. Extended Commands +```go +// Example: Add command handling +switch strings.ToLower(query) { +case "help": + showHelp() +case "stats": + showStats() +default: + processQuery(query) +} +``` + +## Troubleshooting + +### Common Issues and Solutions + +1. **Connection Errors** + ``` + Error: Failed to connect to Milvus + Solution: Ensure Milvus is running and accessible + ``` + +2. **API Rate Limits** + ``` + Error: OpenAI API rate limit exceeded + Solution: Implement rate limiting or increase limits + ``` + +3. **Memory Issues** + ``` + Error: Out of memory + Solution: Adjust batch sizes and memory settings + ``` + +## Advanced Features + +### 1. Conversation History +- Maintains context across multiple queries +- Enables follow-up questions +- Provides coherent conversation flow + +### 2. Document Context +- Integrates knowledge from multiple sources +- Provides source-based responses +- Maintains document relevance + +### 3. Memory Enhancement +- Improves response accuracy +- Enables learning from interactions +- Provides personalized responses + +## Example Extensions + +### 1. Web Interface +```go +// Example: Add HTTP endpoint +func handleChatRequest(w http.ResponseWriter, r *http.Request) { + query := r.URL.Query().Get("q") + response, err := memoryContext.ProcessWithContext(r.Context(), query) + // Handle response +} +``` + +### 2. Slack Integration +```go +// Example: Add Slack bot +func handleSlackMessage(event *slack.MessageEvent) { + response, _ := memoryContext.ProcessWithContext(context.Background(), event.Text) + // Send response to Slack +} +``` + +## Metrics and Monitoring + +### 1. Performance Metrics +- Response time tracking +- Memory usage monitoring +- API call statistics + +### 2. Quality Metrics +- Response relevance scores +- Context utilization rates +- User satisfaction metrics + +## Future Enhancements + +1. **Planned Features** + - Multi-language support + - Advanced document processing + - Enhanced context management + +2. **Potential Improvements** + - Response caching + - Query optimization + - Advanced error recovery diff --git a/docs/contextual_rag.md b/docs/contextual_rag.md new file mode 100644 index 0000000..cb8170a --- /dev/null +++ b/docs/contextual_rag.md @@ -0,0 +1,187 @@ +# Contextual RAG Documentation + +## Overview +The Contextual RAG implementation extends the basic RAG functionality by adding sophisticated context management capabilities. It enables the system to maintain and utilize contextual information across queries, making responses more coherent and contextually relevant. + +## Core Components + +### ContextualRAG Struct +```go +type ContextualRAG struct { + rag *RAG + contextStore map[string]string + contextSize int + contextWindow int +} +``` + +### Configuration +The Contextual RAG system can be configured with various options: + +```go +type ContextualConfig struct { + // Base RAG configuration + RAGConfig *RAGConfig + + // Context settings + ContextSize int // Maximum size of stored context + ContextWindow int // Window size for context consideration + UseMemory bool // Whether to use memory for context +} +``` + +## Key Features + +### 1. Context Management +- Maintains conversation history and context +- Configurable context window and size +- Automatic context pruning and relevance scoring + +### 2. Enhanced Query Processing +- Context-aware query understanding +- Historical context integration +- Improved response coherence + +### 3. Memory Integration +- Optional memory storage for long-term context +- Configurable memory retention +- Context-based memory retrieval + +## Usage Examples + +### Basic Setup +```go +contextualRAG, err := raggo.NewContextualRAG( + raggo.WithBaseRAG(baseRAG), + raggo.WithContextSize(5), + raggo.WithContextWindow(3), + raggo.WithMemory(true), +) +if err != nil { + log.Fatal(err) +} +``` + +### Processing Queries with Context +```go +// Process a query with context +response, err := contextualRAG.ProcessQuery(ctx, "What is the latest development?") + +// Add context explicitly +contextualRAG.AddContext("topic", "AI Development") +response, err = contextualRAG.ProcessQuery(ctx, "What are the challenges?") +``` + +## Best Practices + +1. **Context Management** + - Set appropriate context size based on use case + - Regularly clean up old context + - Monitor context relevance scores + +2. **Query Processing** + - Structure queries to leverage context + - Use context hints when available + - Monitor context window effectiveness + +3. **Memory Usage** + - Enable memory for long-running conversations + - Configure memory retention appropriately + - Implement context cleanup strategies + +## Advanced Features + +### Custom Context Processing +```go +contextualRAG.SetContextProcessor(func(context, query string) string { + // Custom context processing logic + return processedContext +}) +``` + +### Context Filtering +```go +contextualRAG.SetContextFilter(func(context string) bool { + // Custom filtering logic + return isRelevant +}) +``` + +## Example Use Cases + +### 1. Multi-turn Conversations +```go +// Initialize contextual RAG for conversation +contextualRAG, _ := raggo.NewContextualRAG( + raggo.WithContextSize(10), + raggo.WithMemory(true), +) + +// Process conversation turns +for { + response, _ := contextualRAG.ProcessQuery(ctx, userQuery) + contextualRAG.AddContext("conversation", response) +} +``` + +### 2. Document Analysis with Context +```go +// Initialize for document analysis +contextualRAG, _ := raggo.NewContextualRAG( + raggo.WithContextWindow(5), + raggo.WithDocumentContext(true), +) + +// Process document sections +for _, section := range sections { + contextualRAG.AddContext("document", section) + analysis, _ := contextualRAG.ProcessQuery(ctx, "Analyze this section") +} +``` + +## Integration with Memory Systems + +The contextual RAG system can be integrated with various memory systems: + +```go +// Configure memory integration +memorySystem := raggo.NewMemorySystem( + raggo.WithMemorySize(1000), + raggo.WithMemoryType("semantic"), +) + +contextualRAG.SetMemorySystem(memorySystem) +``` + +## Performance Considerations + +1. **Context Size** + - Larger context sizes increase memory usage + - Monitor context processing overhead + - Balance context size with response time + +2. **Memory Usage** + - Implement context cleanup strategies + - Monitor memory consumption + - Use appropriate indexing for context retrieval + +3. **Query Processing** + - Optimize context matching algorithms + - Cache frequently used contexts + - Implement context relevance scoring + +## Error Handling + +```go +// Handle context-related errors +if err := contextualRAG.ProcessQuery(ctx, query); err != nil { + switch err.(type) { + case *ContextSizeError: + // Handle context size exceeded + case *ContextProcessingError: + // Handle processing error + default: + // Handle other errors + } +} +``` diff --git a/docs/examples.md b/docs/examples.md new file mode 100644 index 0000000..f679383 --- /dev/null +++ b/docs/examples.md @@ -0,0 +1,165 @@ +# RagGo Examples Documentation + +This document provides detailed explanations of the example implementations in the RagGo library. + +## Memory Enhancer Example + +### Overview +The Memory Enhancer example demonstrates how to create a RAG system with enhanced memory capabilities for processing technical documentation and maintaining context across interactions. + +### Key Components +```go +// Main components used in the example +- Vector Database (Milvus) +- Memory Context +- LLM Integration (OpenAI) +``` + +### Implementation Details + +1. **Setup and Initialization** +```go +// Initialize LLM +llm, err := gollm.NewLLM( + gollm.SetProvider("openai"), + gollm.SetModel("gpt-4o-mini"), + gollm.SetAPIKey(apiKey), + gollm.SetLogLevel(gollm.LogLevelInfo), +) + +// Initialize Vector Database +vectorDB, err := raggo.NewVectorDB( + raggo.WithType("milvus"), + raggo.WithAddress("localhost:19530"), +) + +// Create Memory Context +memoryContext, err := raggo.NewMemoryContext(apiKey, + raggo.MemoryCollection("tech_docs"), + raggo.MemoryTopK(5), + raggo.MemoryMinScore(0.01), + raggo.MemoryStoreLastN(10), + raggo.MemoryStoreRAGInfo(true), +) +``` + +2. **Document Processing** +```go +// Load and process technical documentation +docsDir := filepath.Join("examples", "chat", "docs") +docs := []string{ + filepath.Join(docsDir, "microservices.txt"), + filepath.Join(docsDir, "vector_databases.txt"), + // ... additional documents +} + +for _, doc := range docs { + content, err := os.ReadFile(doc) + if err != nil { + log.Printf("Warning: Failed to read %s: %v", doc, err) + continue + } + // Store document content as memory + err = memoryContext.Store(ctx, filepath.Base(doc), string(content)) + if err != nil { + log.Printf("Warning: Failed to store %s: %v", doc, err) + } +} +``` + +3. **Interactive Query Processing** +```go +// Process user queries with context +scanner := bufio.NewScanner(os.Stdin) +for { + fmt.Print("\nEnter your question (or 'quit' to exit): ") + if !scanner.Scan() { + break + } + + query := scanner.Text() + if strings.ToLower(query) == "quit" { + break + } + + response, err := memoryContext.ProcessWithContext(ctx, query) + if err != nil { + log.Printf("Error processing query: %v", err) + continue + } + + fmt.Printf("\nResponse: %s\n", response) +} +``` + +## Best Practices Demonstrated + +1. **Error Handling** + - Proper error checking and logging + - Graceful handling of document loading failures + - User-friendly error messages + +2. **Resource Management** + - Proper initialization and cleanup of resources + - Use of context for cancellation + - Cleanup of vector database collections + +3. **User Interaction** + - Clear user prompts and instructions + - Graceful exit handling + - Informative response formatting + +## Running the Example + +1. **Prerequisites** + ```bash + # Set up environment variables + export OPENAI_API_KEY=your_api_key + + # Ensure Milvus is running + docker-compose up -d + ``` + +2. **Running the Example** + ```bash + go run examples/memory_enhancer_example.go + ``` + +3. **Example Interactions** + ``` + Enter your question: What are microservices? + Response: [Detailed response about microservices based on loaded documentation] + + Enter your question: How do vector databases work? + Response: [Context-aware response about vector databases] + ``` + +## Customization Options + +1. **Document Sources** + - Modify the `docs` slice to include different document sources + - Adjust document processing logic for different file types + +2. **Memory Settings** + - Adjust `TopK` for different numbers of similar contexts + - Modify `MinScore` for stricter/looser similarity matching + - Change `StoreLastN` for different memory retention + +3. **Model Configuration** + - Change the LLM model for different capabilities + - Adjust logging levels for debugging + - Modify vector database settings + +## Additional Examples + +### Contextual Example +Located in `examples/contextual/main.go`, this example demonstrates: +- Advanced context management +- Multi-turn conversations +- Context-aware document processing + +### Basic RAG Example +Located in `examples/basic/main.go`, this example shows: +- Simple RAG setup +- Basic document processing +- Query handling without memory enhancement diff --git a/docs/memory_context.md b/docs/memory_context.md new file mode 100644 index 0000000..0674205 --- /dev/null +++ b/docs/memory_context.md @@ -0,0 +1,157 @@ +# Memory Context Documentation + +## Overview +The Memory Context system provides an enhanced way to manage and utilize contextual information in RAG applications. It enables the storage and retrieval of previous interactions, document contexts, and related information to improve the quality of responses. + +## Core Components + +### MemoryContext Struct +The main component that manages contextual memory: + +```go +type MemoryContext struct { + retriever *RAG + config *MemoryConfig + lastN int + storeRAG bool +} +``` + +### Configuration +Memory context can be configured using various options: + +```go +type MemoryConfig struct { + Collection string // Vector DB collection name + TopK int // Number of similar contexts to retrieve + MinScore float64 // Minimum similarity score + StoreLastN int // Number of recent interactions to store + StoreRAG bool // Whether to store RAG information +} +``` + +## Key Features + +### 1. Memory Management +- Store and retrieve recent interactions +- Configure the number of interactions to maintain +- Automatic cleanup of old memories + +### 2. Context Enhancement +- Enrich queries with relevant historical context +- Maintain conversation coherence +- Support for multi-turn interactions + +### 3. RAG Integration +- Seamless integration with the RAG system +- Enhanced document retrieval with historical context +- Configurable similarity thresholds + +## Usage Examples + +### Basic Memory Context Setup +```go +memoryContext, err := raggo.NewMemoryContext(apiKey, + raggo.MemoryCollection("tech_docs"), + raggo.MemoryTopK(5), + raggo.MemoryMinScore(0.01), + raggo.MemoryStoreLastN(10), + raggo.MemoryStoreRAGInfo(true), +) +if err != nil { + log.Fatal(err) +} +``` + +### Using Memory Context in Applications +```go +// Store new memory +err = memoryContext.Store(ctx, "user query", "system response") + +// Retrieve relevant memories +memories, err := memoryContext.Retrieve(ctx, "current query") + +// Process with context +response, err := memoryContext.ProcessWithContext(ctx, "user query") +``` + +## Best Practices + +1. **Memory Configuration** + - Set appropriate `StoreLastN` based on your use case + - Configure `TopK` and `MinScore` for optimal context retrieval + - Enable `StoreRAG` for enhanced context awareness + +2. **Performance Considerations** + - Monitor memory usage with large conversation histories + - Use appropriate batch sizes for memory operations + - Implement cleanup strategies for old memories + +3. **Context Quality** + - Regularly evaluate the quality of retrieved contexts + - Adjust similarity thresholds based on application needs + - Consider implementing context filtering mechanisms + +## Advanced Features + +### Custom Memory Processing +Implement custom memory processing logic: + +```go +memoryContext.SetProcessor(func(ctx context.Context, memory Memory) (string, error) { + // Custom processing logic + return processedMemory, nil +}) +``` + +### Memory Filtering +Apply filters to retrieved memories: + +```go +memoryContext.SetFilter(func(memory Memory) bool { + // Custom filtering logic + return shouldIncludeMemory +}) +``` + +## Example Use Cases + +### 1. Chatbot Enhancement +```go +// Initialize memory context for chat +memoryContext, _ := raggo.NewMemoryContext(apiKey, + raggo.MemoryCollection("chat_history"), + raggo.MemoryStoreLastN(20), +) + +// Process chat messages with context +for { + response, _ := memoryContext.ProcessWithContext(ctx, userMessage) + memoryContext.Store(ctx, userMessage, response) +} +``` + +### 2. Document Q&A System +```go +// Initialize memory context for document Q&A +memoryContext, _ := raggo.NewMemoryContext(apiKey, + raggo.MemoryCollection("doc_qa"), + raggo.MemoryTopK(3), + raggo.MemoryStoreRAGInfo(true), +) + +// Process document queries with context +response, _ := memoryContext.ProcessWithContext(ctx, userQuery) +``` + +## Integration with Vector Databases + +The memory context system integrates seamlessly with vector databases for efficient storage and retrieval of contextual information: + +```go +// Configure vector database integration +retriever := memoryContext.GetRetriever() +if err := retriever.GetVectorDB().LoadCollection(ctx, "tech_docs"); err != nil { + log.Fatal(err) +} +``` diff --git a/docs/rag.md b/docs/rag.md new file mode 100644 index 0000000..6d23762 --- /dev/null +++ b/docs/rag.md @@ -0,0 +1,152 @@ +# RAG (Retrieval-Augmented Generation) Documentation + +## Overview +The RAG (Retrieval-Augmented Generation) package provides a powerful and flexible system for document processing, embedding, storage, and retrieval. It integrates with vector databases and language models to enable context-aware document processing and intelligent information retrieval. + +## Core Components + +### RAG Struct +The main component that provides a unified interface for document processing and retrieval: + +```go +type RAG struct { + db *VectorDB + embedder *EmbeddingService + config *RAGConfig +} +``` + +### Configuration +The `RAGConfig` struct holds all RAG settings: + +```go +type RAGConfig struct { + // Database settings + DBType string + DBAddress string + Collection string + AutoCreate bool + IndexType string + IndexMetric string + + // Processing settings + ChunkSize int + ChunkOverlap int + BatchSize int + + // Embedding settings + Provider string + Model string // For embeddings + LLMModel string // For LLM operations + APIKey string + + // Search settings + TopK int + MinScore float64 + UseHybrid bool + + // System settings + Timeout time.Duration + TempDir string + Debug bool +} +``` + +## Key Features + +### 1. Document Processing +- Chunking documents with configurable size and overlap +- Enriching chunks with contextual information +- Batch processing for efficient handling of large documents + +### 2. Search and Retrieval +- Simple vector similarity search +- Hybrid search combining vector and keyword-based approaches +- Configurable search parameters (TopK, MinScore) + +### 3. Vector Database Integration +- Default support for Milvus +- Extensible design for other vector databases +- Automatic collection creation and management + +### 4. Embedding Services +- Integration with OpenAI embeddings +- Configurable embedding models +- Extensible for other embedding providers + +## Usage Examples + +### Basic RAG Setup +```go +rag, err := raggo.NewRAG( + raggo.WithOpenAI(apiKey), + raggo.WithMilvus("documents"), + raggo.SetChunkSize(512), + raggo.SetTopK(5), +) +if err != nil { + log.Fatal(err) +} +defer rag.Close() +``` + +### Loading Documents +```go +ctx := context.Background() +err = rag.LoadDocuments(ctx, "path/to/documents") +if err != nil { + log.Fatal(err) +} +``` + +### Querying +```go +results, err := rag.Query(ctx, "your search query") +if err != nil { + log.Fatal(err) +} +``` + +## Best Practices + +1. **Configuration** + - Use appropriate chunk sizes based on your content (default: 512) + - Adjust TopK and MinScore based on your use case + - Enable hybrid search for better results when appropriate + +2. **Performance** + - Use batch processing for large document sets + - Configure appropriate timeouts + - Monitor vector database performance + +3. **Error Handling** + - Always handle errors appropriately + - Use context for cancellation and timeouts + - Close resources using defer + +4. **Security** + - Never hardcode API keys + - Use environment variables for sensitive configuration + - Implement appropriate access controls for your vector database + +## Advanced Features + +### Context-Aware Processing +The RAG system can enrich document chunks with contextual information: + +```go +err = rag.ProcessWithContext(ctx, "path/to/documents", "gpt-4") +``` + +### Custom Search Parameters +Configure specific search parameters for your use case: + +```go +rag, err := raggo.NewRAG( + raggo.SetSearchParams(map[string]interface{}{ + "nprobe": 10, + "ef": 64, + "type": "HNSW", + }), +) +``` diff --git a/docs/simple_rag.md b/docs/simple_rag.md new file mode 100644 index 0000000..95d5220 --- /dev/null +++ b/docs/simple_rag.md @@ -0,0 +1,225 @@ +# Simple RAG Documentation + +## Overview +The Simple RAG implementation provides a streamlined, easy-to-use interface for basic RAG operations. It's designed for straightforward use cases where advanced context management isn't required, offering a balance between functionality and simplicity. + +## Core Components + +### SimpleRAG Struct +```go +type SimpleRAG struct { + embedder *EmbeddingService + vectorStore *VectorDB + config *SimpleConfig +} +``` + +### Configuration +Simple configuration options for basic RAG operations: + +```go +type SimpleConfig struct { + // Vector store settings + VectorDBType string + VectorDBAddress string + Collection string + + // Embedding settings + EmbeddingModel string + APIKey string + + // Search settings + TopK int + MinScore float64 +} +``` + +## Key Features + +### 1. Simplified Document Processing +- Straightforward document ingestion +- Basic chunking and embedding +- Direct vector storage + +### 2. Basic Search Functionality +- Simple similarity search +- Configurable result count +- Basic relevance scoring + +### 3. Minimal Setup Required +- Default configurations +- Automatic resource management +- Simplified API + +## Usage Examples + +### Basic Setup +```go +simpleRAG, err := raggo.NewSimpleRAG( + raggo.WithVectorDB("milvus", "localhost:19530"), + raggo.WithEmbeddings("openai", apiKey), + raggo.WithTopK(3), +) +if err != nil { + log.Fatal(err) +} +``` + +### Document Processing +```go +// Process a single document +err = simpleRAG.AddDocument(ctx, "document.txt") + +// Process multiple documents +err = simpleRAG.AddDocuments(ctx, []string{ + "doc1.txt", + "doc2.txt", +}) +``` + +### Query Processing +```go +// Simple query +results, err := simpleRAG.Query(ctx, "How does this work?") + +// Query with custom parameters +results, err = simpleRAG.QueryWithParams(ctx, "How does this work?", + raggo.WithResultCount(5), + raggo.WithMinScore(0.5), +) +``` + +## Best Practices + +1. **Document Management** + - Keep documents reasonably sized + - Use appropriate file formats + - Maintain clean document structure + +2. **Query Optimization** + - Keep queries focused and specific + - Use appropriate TopK values + - Monitor query performance + +3. **Resource Management** + - Close resources when done + - Monitor memory usage + - Use batch processing for large datasets + +## Example Use Cases + +### 1. Document Q&A System +```go +// Initialize simple RAG for Q&A +simpleRAG, _ := raggo.NewSimpleRAG( + raggo.WithCollection("qa_docs"), + raggo.WithTopK(1), +) + +// Add documentation +simpleRAG.AddDocument(ctx, "documentation.txt") + +// Process questions +answer, _ := simpleRAG.Query(ctx, "What is the installation process?") +``` + +### 2. Basic Search System +```go +// Initialize for search +simpleRAG, _ := raggo.NewSimpleRAG( + raggo.WithTopK(5), + raggo.WithMinScore(0.7), +) + +// Add searchable content +simpleRAG.AddDocuments(ctx, []string{ + "content1.txt", + "content2.txt", +}) + +// Search content +results, _ := simpleRAG.Query(ctx, "search term") +``` + +## Performance Tips + +1. **Document Processing** + - Use batch processing for multiple documents + - Monitor embedding API usage + - Implement rate limiting for API calls + +2. **Query Optimization** + - Cache frequent queries + - Use appropriate TopK values + - Monitor query latency + +3. **Resource Usage** + - Implement connection pooling + - Monitor memory consumption + - Use appropriate batch sizes + +## Error Handling + +```go +// Basic error handling +if err := simpleRAG.AddDocument(ctx, "doc.txt"); err != nil { + switch err.(type) { + case *FileError: + // Handle file-related errors + case *ProcessingError: + // Handle processing errors + default: + // Handle other errors + } +} +``` + +## Integration Examples + +### With HTTP Server +```go +func handleQuery(w http.ResponseWriter, r *http.Request) { + query := r.URL.Query().Get("q") + results, err := simpleRAG.Query(r.Context(), query) + if err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + json.NewEncoder(w).Encode(results) +} +``` + +### With CLI Application +```go +func main() { + simpleRAG, _ := raggo.NewSimpleRAG(/* config */) + scanner := bufio.NewScanner(os.Stdin) + + for { + fmt.Print("Enter query: ") + if !scanner.Scan() { + break + } + + results, _ := simpleRAG.Query(context.Background(), scanner.Text()) + fmt.Printf("Results: %v\n", results) + } +} +``` + +## Limitations and Considerations + +1. **No Advanced Context Management** + - Limited to single-query operations + - No conversation history + - No context window management + +2. **Basic Search Only** + - No hybrid search capabilities + - Limited to vector similarity + - Basic relevance scoring + +3. **Limited Customization** + - Fixed chunking strategy + - Basic embedding options + - Simple configuration options diff --git a/docs/summary.md b/docs/summary.md new file mode 100644 index 0000000..a804ec8 --- /dev/null +++ b/docs/summary.md @@ -0,0 +1,259 @@ +# RagGo Library Documentation Summary + +## Overview +RagGo is a comprehensive Go library for implementing Retrieval-Augmented Generation (RAG) systems. It provides multiple implementations to suit different use cases, from simple document retrieval to complex context-aware applications. + +## Components Overview + +### 1. Core RAG (`rag.go`) +The foundation of the library, providing basic RAG functionality: +- Document processing and embedding +- Vector database integration +- Configurable search parameters +- Extensible architecture + +[Detailed Documentation →](./rag.md) + +### 2. Simple RAG (`simple_rag.go`) +A streamlined implementation for basic use cases: +- Minimal setup required +- Basic document processing +- Simple search functionality +- Ideal for straightforward applications + +[Detailed Documentation →](./simple_rag.md) + +### 3. Contextual RAG (`contextual_rag.go`) +Advanced implementation with context management: +- Conversation history tracking +- Context-aware responses +- Memory integration +- Suitable for complex applications + +[Detailed Documentation →](./contextual_rag.md) + +### 4. Memory Context (`memory_context.go`) +Enhanced memory management system: +- Long-term memory storage +- Context retention +- Configurable memory policies +- Integration with RAG systems + +[Detailed Documentation →](./memory_context.md) + +## Quick Start Guide + +### 1. Basic Usage +```go +// Initialize basic RAG +rag, err := raggo.NewRAG( + raggo.WithOpenAI(apiKey), + raggo.WithMilvus("documents"), +) + +// Process documents +err = rag.LoadDocuments(ctx, "path/to/docs") + +// Query +results, err := rag.Query(ctx, "your query") +``` + +### 2. Simple RAG Usage +```go +// Initialize simple RAG +simpleRAG, err := raggo.NewSimpleRAG( + raggo.WithVectorDB("milvus", "localhost:19530"), + raggo.WithEmbeddings("openai", apiKey), +) + +// Process and query +err = simpleRAG.AddDocument(ctx, "document.txt") +results, err := simpleRAG.Query(ctx, "query") +``` + +### 3. Contextual RAG Usage +```go +// Initialize contextual RAG +contextualRAG, err := raggo.NewContextualRAG( + raggo.WithBaseRAG(baseRAG), + raggo.WithContextSize(5), + raggo.WithMemory(true), +) + +// Process with context +response, err := contextualRAG.ProcessQuery(ctx, "query") +``` + +## Feature Comparison + +| Feature | Simple RAG | Core RAG | Contextual RAG | +|---------------------------|------------|----------|----------------| +| Document Processing | Basic | Advanced | Advanced | +| Context Management | No | Basic | Advanced | +| Memory Integration | No | Optional | Yes | +| Search Capabilities | Basic | Advanced | Advanced | +| Setup Complexity | Low | Medium | High | +| Resource Requirements | Low | Medium | High | + +## Common Use Cases + +### 1. Document Q&A +Best Implementation: Simple RAG +```go +simpleRAG, _ := raggo.NewSimpleRAG( + raggo.WithCollection("qa_docs"), + raggo.WithTopK(1), +) +``` + +### 2. Chatbot with Memory +Best Implementation: Contextual RAG +```go +contextualRAG, _ := raggo.NewContextualRAG( + raggo.WithContextSize(10), + raggo.WithMemory(true), +) +``` + +### 3. Document Analysis +Best Implementation: Core RAG +```go +rag, _ := raggo.NewRAG( + raggo.WithChunkSize(512), + raggo.WithHybridSearch(true), +) +``` + +## Example Implementations + +### 1. Memory-Enhanced Chatbot +A sophisticated chatbot implementation demonstrating advanced RAG capabilities: +- Document-based knowledge integration +- Context-aware responses +- Memory retention across conversations +- Interactive CLI interface + +[Detailed Documentation →](./chatbot_example.md) + +```go +// Initialize components +memoryContext, err := raggo.NewMemoryContext(apiKey, + raggo.MemoryCollection("tech_docs"), + raggo.MemoryTopK(5), + raggo.MemoryMinScore(0.01), + raggo.MemoryStoreLastN(10), +) + +// Process documents +for _, doc := range docs { + content, err := os.ReadFile(doc) + err = memoryContext.Store(ctx, filepath.Base(doc), string(content)) +} + +// Interactive chat loop +for { + query := getUserInput() + response, err := memoryContext.ProcessWithContext(ctx, query) + fmt.Printf("\nResponse: %s\n", response) +} +``` + +Key Features: +- Vector database integration (Milvus) +- OpenAI LLM integration +- Document processing and storage +- Context-aware query processing +- Memory management +- Error handling and logging + +[View Full Example →](../examples/memory_enhancer_example.go) + +## Best Practices + +### 1. Configuration +- Use appropriate chunk sizes (default: 512) +- Configure TopK based on use case +- Set reasonable timeouts +- Use environment variables for API keys + +### 2. Performance +- Implement batch processing +- Monitor API usage +- Use connection pooling +- Cache frequent queries + +### 3. Error Handling +- Implement proper error handling +- Use context for cancellation +- Close resources properly +- Monitor system resources + +## Integration Examples + +### 1. HTTP Server +```go +func handleQuery(w http.ResponseWriter, r *http.Request) { + rag := getRAGInstance() // Get appropriate RAG instance + results, err := rag.Query(r.Context(), r.URL.Query().Get("q")) + if err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + json.NewEncoder(w).Encode(results) +} +``` + +### 2. CLI Application +```go +func main() { + rag := initializeRAG() // Initialize appropriate RAG + scanner := bufio.NewScanner(os.Stdin) + + for { + fmt.Print("Query: ") + if !scanner.Scan() { + break + } + results, _ := rag.Query(context.Background(), scanner.Text()) + fmt.Printf("Results: %v\n", results) + } +} +``` + +## Dependencies + +- Vector Database: Milvus (default) +- Embedding Service: OpenAI (default) +- Go version: 1.16+ + +## Getting Started + +1. **Installation** + ```bash + go get github.com/teilomillet/raggo + ``` + +2. **Basic Setup** + ```go + import "github.com/teilomillet/raggo" + ``` + +3. **Environment Variables** + ```bash + export OPENAI_API_KEY=your_api_key + ``` + +## Contributing + +See [CONTRIBUTING.md](../CONTRIBUTING.md) for guidelines on: +- Code style +- Testing requirements +- Pull request process +- Documentation standards + +## Resources + +- [Examples Directory](../examples/) +- [API Reference](./api.md) +- [FAQ](./faq.md) +- [Troubleshooting Guide](./troubleshooting.md) diff --git a/examples/memory_enhancer_example.go b/examples/memory_enhancer_example.go new file mode 100644 index 0000000..c2086e0 --- /dev/null +++ b/examples/memory_enhancer_example.go @@ -0,0 +1,193 @@ +package main + +import ( + "bufio" + "context" + "fmt" + "log" + "os" + "path/filepath" + "strings" + + "github.com/teilomillet/gollm" + "github.com/teilomillet/raggo" +) + +func main() { + log.Println("DEBUG: Starting memory enhancer example") + + // Get API key from environment + apiKey := os.Getenv("OPENAI_API_KEY") + if apiKey == "" { + log.Fatal("OPENAI_API_KEY environment variable not set") + } + log.Println("DEBUG: API key found") + + // Initialize LLM + log.Println("DEBUG: Initializing LLM...") + llm, err := gollm.NewLLM( + gollm.SetProvider("openai"), + gollm.SetModel("gpt-4o-mini"), + gollm.SetAPIKey(apiKey), + gollm.SetLogLevel(gollm.LogLevelInfo), + ) + if err != nil { + log.Fatalf("Failed to initialize LLM: %v", err) + } + log.Println("DEBUG: LLM initialized successfully") + + // Create memory context + log.Println("DEBUG: Creating memory context...") + + // Drop existing collection if it exists + vectorDB, err := raggo.NewVectorDB( + raggo.WithType("milvus"), + raggo.WithAddress("localhost:19530"), + ) + if err != nil { + log.Fatalf("Failed to create vector database: %v", err) + } + if err := vectorDB.Connect(context.Background()); err != nil { + log.Fatalf("Failed to connect to vector database: %v", err) + } + exists, err := vectorDB.HasCollection(context.Background(), "tech_docs") + if err != nil { + log.Fatalf("Failed to check collection: %v", err) + } + if exists { + log.Println("Dropping existing collection") + err = vectorDB.DropCollection(context.Background(), "tech_docs") + if err != nil { + log.Fatalf("Failed to drop collection: %v", err) + } + } + + // Create memory context with optimized settings + memoryContext, err := raggo.NewMemoryContext(apiKey, + raggo.MemoryCollection("tech_docs"), + raggo.MemoryTopK(5), + raggo.MemoryMinScore(0.01), + raggo.MemoryStoreLastN(10), + raggo.MemoryStoreRAGInfo(true), + ) + if err != nil { + log.Fatalf("Failed to create memory context: %v", err) + } + + // Configure hybrid search + retriever := memoryContext.GetRetriever() + if err := retriever.GetVectorDB().LoadCollection(context.Background(), "tech_docs"); err != nil { + log.Fatalf("Failed to load collection: %v", err) + } + + // Load technical documentation + ctx := context.Background() + docsDir := filepath.Join("examples", "chat", "docs") + docs := []string{ + filepath.Join(docsDir, "microservices.txt"), + filepath.Join(docsDir, "vector_databases.txt"), + filepath.Join(docsDir, "rag_systems.txt"), + filepath.Join(docsDir, "golang_basics.txt"), + filepath.Join(docsDir, "embeddings.txt"), + } + + for _, doc := range docs { + content, err := os.ReadFile(doc) + if err != nil { + log.Printf("Warning: Failed to read %s: %v", doc, err) + continue + } + + // Store document content as memory + err = memoryContext.StoreMemory(ctx, []gollm.MemoryMessage{ + { + Role: "system", + Content: string(content), + Tokens: 0, // Let the system calculate tokens + }, + }) + if err != nil { + log.Printf("Warning: Failed to store memory from %s: %v", doc, err) + } + } + + // Chat loop + fmt.Println("\nChat started! Type 'exit' to end the conversation.") + fmt.Println("\nExample questions you can ask:") + fmt.Println("1. What is MountainPass's PressureValve system and how did it help during Black Friday?") + fmt.Println("2. What are the key features of the PressureValve architecture?") + fmt.Println("3. How did MountainPass handle their e-commerce scaling challenge?") + + // Initialize chat memory + var memory []gollm.MemoryMessage + systemPrompt := "You are a technical expert helping explain the MountainPass case study and their PressureValve system. " + + "Focus on the specific details from the case study, including the challenges they faced, their solution, and the results they achieved. " + + "Always reference specific numbers, features, and outcomes from the documentation." + + scanner := bufio.NewScanner(os.Stdin) + for { + fmt.Print("\nQ: ") + scanner.Scan() + query := scanner.Text() + if query == "exit" { + break + } + + // Create base prompt with query + prompt := gollm.NewPrompt(query) + + // Enhance prompt with relevant context + enhancedPrompt, err := memoryContext.EnhancePrompt(ctx, prompt, memory) + if err != nil { + log.Printf("Failed to enhance prompt: %v", err) + continue + } + + // Build final prompt with system context + var promptBuilder strings.Builder + promptBuilder.WriteString(systemPrompt) + promptBuilder.WriteString("\n\nBased on the following context:\n") + promptBuilder.WriteString(enhancedPrompt.String()) + promptBuilder.WriteString("\n\nPlease answer this question: ") + promptBuilder.WriteString(query) + + finalPrompt := gollm.NewPrompt(promptBuilder.String()) + + // Generate response + fmt.Print("\nA: ") + response, err := llm.Generate(ctx, finalPrompt) + if err != nil { + log.Printf("Failed to generate response: %v", err) + continue + } + fmt.Println(response) + + // Store the interaction in memory + memory = append(memory, gollm.MemoryMessage{ + Role: "user", + Content: query, + Tokens: 0, + }, gollm.MemoryMessage{ + Role: "assistant", + Content: response, + Tokens: 0, + }) + + // Store last N messages + if err := memoryContext.StoreLastN(ctx, memory, memoryContext.GetOptions().StoreLastN); err != nil { + log.Printf("Failed to store memory: %v", err) + } + } + + // Example of updating options + fmt.Println("\nUpdating memory context configuration...") + memoryContext.UpdateOptions( + raggo.MemoryTopK(10), + raggo.MemoryMinScore(0.01), + raggo.MemoryStoreLastN(20), + ) + + options := memoryContext.GetOptions() + fmt.Printf("Current options: TopK=%d, MinScore=%.2f, Collection=%s, StoreLastN=%d, StoreRAGInfo=%v\n", + options.TopK, options.MinScore, options.Collection, options.StoreLastN, options.StoreRAGInfo) +} diff --git a/memory_context.go b/memory_context.go new file mode 100644 index 0000000..1a156e2 --- /dev/null +++ b/memory_context.go @@ -0,0 +1,279 @@ +package raggo + +import ( + "context" + "fmt" + "strings" + "time" + + "github.com/teilomillet/gollm" +) + +// MemoryContextOptions configures how the memory context works +type MemoryContextOptions struct { + TopK int + MinScore float64 + Collection string + IncludeScore bool + StoreLastN int + StoreRAGInfo bool // Whether to store RAG-enhanced context in memory +} + +// MemoryContext provides contextual memory for gollm using RAG +type MemoryContext struct { + retriever *Retriever + options MemoryContextOptions + lastStoreTime time.Time + lastMemoryLen int +} + +// MemoryTopK sets the number of relevant memories to retrieve +func MemoryTopK(k int) func(*MemoryContextOptions) { + return func(o *MemoryContextOptions) { + o.TopK = k + } +} + +// MemoryMinScore sets the minimum similarity score threshold +func MemoryMinScore(score float64) func(*MemoryContextOptions) { + return func(o *MemoryContextOptions) { + o.MinScore = score + } +} + +// MemoryCollection sets the RAG collection name +func MemoryCollection(collection string) func(*MemoryContextOptions) { + return func(o *MemoryContextOptions) { + o.Collection = collection + } +} + +// MemoryScoreInclusion controls whether to include relevance scores +func MemoryScoreInclusion(include bool) func(*MemoryContextOptions) { + return func(o *MemoryContextOptions) { + o.IncludeScore = include + } +} + +// MemoryStoreLastN sets the number of recent messages to store +func MemoryStoreLastN(n int) func(*MemoryContextOptions) { + return func(o *MemoryContextOptions) { + o.StoreLastN = n + } +} + +// MemoryStoreRAGInfo controls whether to store RAG-enhanced context in memory +func MemoryStoreRAGInfo(store bool) func(*MemoryContextOptions) { + return func(o *MemoryContextOptions) { + o.StoreRAGInfo = store + } +} + +// NewMemoryContext creates a new memory context provider that works with gollm +func NewMemoryContext(apiKey string, opts ...func(*MemoryContextOptions)) (*MemoryContext, error) { + // Default options + options := MemoryContextOptions{ + TopK: 3, + MinScore: 0.7, + Collection: "memory_store", + IncludeScore: false, + StoreLastN: 0, + StoreRAGInfo: false, + } + + // Apply custom options + for _, opt := range opts { + opt(&options) + } + + // Initialize retriever with config + retriever, err := NewRetriever( + WithRetrieveCollection(options.Collection), + WithTopK(options.TopK), + WithMinScore(options.MinScore), + WithRetrieveEmbedding("openai", "text-embedding-3-small", apiKey), + WithRetrieveDB("milvus", "localhost:19530"), // Add default DB config + ) + if err != nil { + return nil, fmt.Errorf("failed to initialize retriever: %w", err) + } + + // Ensure collection exists with proper schema + ctx := context.Background() + db := retriever.GetVectorDB() + exists, err := db.HasCollection(ctx, options.Collection) + if err != nil { + return nil, fmt.Errorf("failed to check collection: %w", err) + } + + if !exists { + // Create collection with proper schema + schema := Schema{ + Name: options.Collection, + Description: "Memory context collection for RAG", + Fields: []Field{ + {Name: "ID", DataType: "int64", PrimaryKey: true, AutoID: true}, + {Name: "Embedding", DataType: "float_vector", Dimension: 1536}, // text-embedding-3-small dimension + {Name: "Text", DataType: "varchar", MaxLength: 65535}, + {Name: "Metadata", DataType: "varchar", MaxLength: 65535}, + }, + } + + if err := db.CreateCollection(ctx, options.Collection, schema); err != nil { + return nil, fmt.Errorf("failed to create collection: %w", err) + } + + // Create index for vector search + index := Index{ + Type: "HNSW", + Metric: "COSINE", + Parameters: map[string]interface{}{ + "M": 16, + "efConstruction": 256, + }, + } + + if err := db.CreateIndex(ctx, options.Collection, "Embedding", index); err != nil { + return nil, fmt.Errorf("failed to create index: %w", err) + } + + if err := db.LoadCollection(ctx, options.Collection); err != nil { + return nil, fmt.Errorf("failed to load collection: %w", err) + } + } + + return &MemoryContext{ + retriever: retriever, + options: options, + lastStoreTime: time.Time{}, + lastMemoryLen: 0, + }, nil +} + +// shouldStore determines whether to store the given memory +func (m *MemoryContext) shouldStore(memory []gollm.MemoryMessage) bool { + newLen := len(memory) + timeSinceLastStore := time.Since(m.lastStoreTime) + messagesDiff := newLen - m.lastMemoryLen + + return messagesDiff >= m.options.StoreLastN/2 || + timeSinceLastStore > 5*time.Minute +} + +// StoreMemory explicitly stores messages in the memory context +func (m *MemoryContext) StoreMemory(ctx context.Context, messages []gollm.MemoryMessage) error { + if len(messages) == 0 { + return nil + } + + var memoryContent string + for _, msg := range messages { + memoryContent += fmt.Sprintf("%s: %s\n", msg.Role, msg.Content) + } + + // TODO: Implement document storage through the vector store + // For now, we'll just store the content through the retriever + _, err := m.retriever.Retrieve(ctx, memoryContent) + if err != nil { + return fmt.Errorf("failed to store memory: %w", err) + } + + return nil +} + +// StoreLastN stores the last N messages from the memory +func (m *MemoryContext) StoreLastN(ctx context.Context, memory []gollm.MemoryMessage, n int) error { + if !m.shouldStore(memory) { + return nil + } + + start := len(memory) - n + if start < 0 { + start = 0 + } + + err := m.StoreMemory(ctx, memory[start:]) + if err == nil { + m.lastStoreTime = time.Now() + m.lastMemoryLen = len(memory) + } + return err +} + +// EnhancePrompt enriches a prompt with relevant context from memory +func (m *MemoryContext) EnhancePrompt(ctx context.Context, prompt *gollm.Prompt, memory []gollm.MemoryMessage) (*gollm.Prompt, error) { + relevantContext, err := m.retrieveContext(ctx, prompt.Input) + if err != nil { + return prompt, fmt.Errorf("failed to retrieve context: %w", err) + } + + enhancedPrompt := gollm.NewPrompt( + prompt.Input, + gollm.WithSystemPrompt(prompt.SystemPrompt, gollm.CacheTypeEphemeral), + gollm.WithContext(strings.Join(relevantContext, "\n")), + ) + + if m.options.StoreRAGInfo && len(relevantContext) > 0 { + contextMsg := gollm.MemoryMessage{ + Role: "system", + Content: "Retrieved Context:\n" + strings.Join(relevantContext, "\n"), + } + memory = append(memory, contextMsg) + } + + return enhancedPrompt, nil +} + +// retrieveContext retrieves relevant context from RAG +func (m *MemoryContext) retrieveContext(ctx context.Context, input string) ([]string, error) { + results, err := m.retriever.Retrieve(ctx, input) + if err != nil { + return nil, fmt.Errorf("failed to search context: %w", err) + } + + var relevantContext []string + for _, result := range results { + if result.Content != "" { + if m.options.IncludeScore { + relevantContext = append(relevantContext, fmt.Sprintf("[Score: %.2f] %s", result.Score, result.Content)) + } else { + relevantContext = append(relevantContext, result.Content) + } + } + } + + return relevantContext, nil +} + +// Close releases resources +func (m *MemoryContext) Close() error { + return m.retriever.Close() +} + +// GetRetriever returns the underlying retriever instance for advanced configuration +func (m *MemoryContext) GetRetriever() *Retriever { + return m.retriever +} + +// GetOptions returns the current context options +func (m *MemoryContext) GetOptions() MemoryContextOptions { + return m.options +} + +// UpdateOptions allows updating context options at runtime +func (m *MemoryContext) UpdateOptions(opts ...func(*MemoryContextOptions)) { + options := m.GetOptions() + for _, opt := range opts { + opt(&options) + } + m.options = options + + // Create new retriever with updated config + if retriever, err := NewRetriever( + WithRetrieveCollection(options.Collection), + WithTopK(options.TopK), + WithMinScore(options.MinScore), + ); err == nil { + m.retriever = retriever + } +} diff --git a/rag.go b/rag.go index c8555a3..c8f6c4a 100644 --- a/rag.go +++ b/rag.go @@ -85,6 +85,78 @@ func DefaultRAGConfig() *RAGConfig { } // Common options +func SetProvider(provider string) RAGOption { + return func(c *RAGConfig) { + c.Provider = provider + } +} + +func SetModel(model string) RAGOption { + return func(c *RAGConfig) { + c.Model = model + } +} + +func SetAPIKey(key string) RAGOption { + return func(c *RAGConfig) { + c.APIKey = key + } +} + +func SetCollection(name string) RAGOption { + return func(c *RAGConfig) { + c.Collection = name + } +} + +func SetSearchStrategy(strategy string) RAGOption { + return func(c *RAGConfig) { + c.UseHybrid = strategy == "hybrid" + } +} + +func SetDBAddress(address string) RAGOption { + return func(c *RAGConfig) { + c.DBAddress = address + } +} + +func SetChunkSize(size int) RAGOption { + return func(c *RAGConfig) { + c.ChunkSize = size + } +} + +func SetChunkOverlap(overlap int) RAGOption { + return func(c *RAGConfig) { + c.ChunkOverlap = overlap + } +} + +func SetTopK(k int) RAGOption { + return func(c *RAGConfig) { + c.TopK = k + } +} + +func SetMinScore(score float64) RAGOption { + return func(c *RAGConfig) { + c.MinScore = score + } +} + +func SetTimeout(timeout time.Duration) RAGOption { + return func(c *RAGConfig) { + c.Timeout = timeout + } +} + +func SetDebug(debug bool) RAGOption { + return func(c *RAGConfig) { + c.Debug = debug + } +} + func WithOpenAI(apiKey string) RAGOption { return func(c *RAGConfig) { c.Provider = "openai" @@ -147,8 +219,8 @@ func (r *RAG) initialize() error { func (r *RAG) LoadDocuments(ctx context.Context, source string) error { loader := NewLoader(SetTempDir(r.config.TempDir)) chunker, err := NewChunker( - ChunkSize(r.config.ChunkSize), // Changed: use correct option functions - ChunkOverlap(r.config.ChunkOverlap), + SetChunkSize(r.config.ChunkSize), + SetChunkOverlap(r.config.ChunkOverlap), ) if err != nil { return err @@ -297,8 +369,8 @@ func (r *RAG) ProcessWithContext(ctx context.Context, source string, llmModel st // Create chunks chunker, _ := NewChunker( - ChunkSize(r.config.ChunkSize), - ChunkOverlap(r.config.ChunkOverlap), + SetChunkSize(r.config.ChunkSize), + SetChunkOverlap(r.config.ChunkOverlap), ) chunks := chunker.Chunk(doc.Content) diff --git a/rag/milvus.go b/rag/milvus.go index 12dd4fe..05fd99c 100644 --- a/rag/milvus.go +++ b/rag/milvus.go @@ -22,15 +22,18 @@ func newMilvusDB(cfg *Config) (*MilvusDB, error) { } func (m *MilvusDB) Connect(ctx context.Context) error { + GlobalLogger.Debug("Attempting to connect to Milvus", "address", m.config.Address) + c, err := client.NewClient(ctx, client.Config{ Address: m.config.Address, }) if err != nil { - GlobalLogger.Error("Failed to connect to Milvus", "error", err) - return err + GlobalLogger.Error("Failed to connect to Milvus", "error", err, "address", m.config.Address) + return fmt.Errorf("failed to connect to Milvus at %s: %w\nPlease ensure Milvus is running (e.g., with 'docker-compose up -d')", m.config.Address, err) } + m.client = c - GlobalLogger.Debug("Connected to Milvus", "address", m.config.Address) + GlobalLogger.Debug("Successfully connected to Milvus") return nil } diff --git a/rag/reranker.go b/rag/reranker.go new file mode 100644 index 0000000..f4c1d9e --- /dev/null +++ b/rag/reranker.go @@ -0,0 +1,75 @@ +package rag + +import ( + "context" + "sort" +) + +// RRFReranker implements Reciprocal Rank Fusion for result reranking +type RRFReranker struct { + k float64 // Constant to prevent division by zero and control ranking influence +} + +// NewRRFReranker creates a new RRF reranker with the given k parameter +func NewRRFReranker(k float64) *RRFReranker { + if k <= 0 { + k = 60 // Default value from RRF paper + } + return &RRFReranker{k: k} +} + +// Rerank combines and reranks results using Reciprocal Rank Fusion +func (r *RRFReranker) Rerank( + ctx context.Context, + query string, + denseResults, sparseResults []SearchResult, + denseWeight, sparseWeight float64, +) ([]SearchResult, error) { + // Normalize weights + totalWeight := denseWeight + sparseWeight + if totalWeight > 0 { + denseWeight /= totalWeight + sparseWeight /= totalWeight + } else { + denseWeight = 0.5 + sparseWeight = 0.5 + } + + // Create maps to store combined scores + scores := make(map[int64]float64) + docMap := make(map[int64]SearchResult) + + // Process dense results + for rank, result := range denseResults { + rrf := 1.0 / (float64(rank+1) + r.k) + scores[result.ID] = rrf * denseWeight + docMap[result.ID] = result + } + + // Process sparse results + for rank, result := range sparseResults { + rrf := 1.0 / (float64(rank+1) + r.k) + if score, exists := scores[result.ID]; exists { + scores[result.ID] = score + rrf*sparseWeight + } else { + scores[result.ID] = rrf * sparseWeight + docMap[result.ID] = result + } + } + + // Convert scores back to results + results := make([]SearchResult, 0, len(scores)) + for id, score := range scores { + result := docMap[id] + result.Score = score + results = append(results, result) + } + + // Sort by combined score + sort.Slice(results, func(i, j int) bool { + return results[i].Score > results[j].Score + }) + + // Return all reranked results + return results, nil +} diff --git a/rag/sparse_index.go b/rag/sparse_index.go new file mode 100644 index 0000000..43a3f2e --- /dev/null +++ b/rag/sparse_index.go @@ -0,0 +1,192 @@ +package rag + +import ( + "context" + "math" + "sort" + "strings" + "sync" +) + +// BM25Parameters holds the parameters for BM25 scoring +type BM25Parameters struct { + K1 float64 // Term saturation parameter (typically 1.2-2.0) + B float64 // Length normalization parameter (typically 0.75) +} + +// DefaultBM25Parameters returns default BM25 parameters +func DefaultBM25Parameters() BM25Parameters { + return BM25Parameters{ + K1: 1.5, + B: 0.75, + } +} + +// BM25Index implements a sparse index using BM25 scoring +type BM25Index struct { + mu sync.RWMutex + docs map[int64]string // Document content by ID + metadata map[int64]map[string]interface{} // Document metadata by ID + termFreq map[int64]map[string]int // Term frequency per document + docFreq map[string]int // Document frequency per term + docLength map[int64]int // Length of each document + avgDocLength float64 // Average document length + totalDocs int // Total number of documents + params BM25Parameters // BM25 parameters + preprocessor func(string) []string // Text preprocessing function +} + +// NewBM25Index creates a new BM25 index with default parameters +func NewBM25Index() *BM25Index { + return &BM25Index{ + docs: make(map[int64]string), + metadata: make(map[int64]map[string]interface{}), + termFreq: make(map[int64]map[string]int), + docFreq: make(map[string]int), + docLength: make(map[int64]int), + params: DefaultBM25Parameters(), + preprocessor: defaultPreprocessor, + } +} + +// defaultPreprocessor implements basic text preprocessing +func defaultPreprocessor(text string) []string { + // Convert to lowercase and split into words + words := strings.Fields(strings.ToLower(text)) + return words +} + +// Add adds a document to the index +func (idx *BM25Index) Add(ctx context.Context, id int64, content string, metadata map[string]interface{}) error { + idx.mu.Lock() + defer idx.mu.Unlock() + + // Store document and metadata + idx.docs[id] = content + idx.metadata[id] = metadata + + // Process terms + terms := idx.preprocessor(content) + termFreq := make(map[string]int) + for _, term := range terms { + termFreq[term]++ + } + + // Update index statistics + idx.termFreq[id] = termFreq + idx.docLength[id] = len(terms) + for term := range termFreq { + idx.docFreq[term]++ + } + + // Update collection statistics + idx.totalDocs++ + var totalLength int + for _, length := range idx.docLength { + totalLength += length + } + idx.avgDocLength = float64(totalLength) / float64(idx.totalDocs) + + return nil +} + +// Remove removes a document from the index +func (idx *BM25Index) Remove(ctx context.Context, id int64) error { + idx.mu.Lock() + defer idx.mu.Unlock() + + // Update document frequencies + if termFreq, exists := idx.termFreq[id]; exists { + for term := range termFreq { + idx.docFreq[term]-- + if idx.docFreq[term] == 0 { + delete(idx.docFreq, term) + } + } + } + + // Remove document data + delete(idx.docs, id) + delete(idx.metadata, id) + delete(idx.termFreq, id) + delete(idx.docLength, id) + + // Update collection statistics + idx.totalDocs-- + if idx.totalDocs > 0 { + var totalLength int + for _, length := range idx.docLength { + totalLength += length + } + idx.avgDocLength = float64(totalLength) / float64(idx.totalDocs) + } else { + idx.avgDocLength = 0 + } + + return nil +} + +// Search performs BM25 search on the index +func (idx *BM25Index) Search(ctx context.Context, query string, topK int) ([]SearchResult, error) { + idx.mu.RLock() + defer idx.mu.RUnlock() + + // Process query terms + queryTerms := idx.preprocessor(query) + scores := make(map[int64]float64) + + // Calculate BM25 scores + for _, term := range queryTerms { + if df, exists := idx.docFreq[term]; exists { + idf := math.Log(1 + (float64(idx.totalDocs)-float64(df)+0.5)/(float64(df)+0.5)) + + for docID, docTerms := range idx.termFreq { + if tf, exists := docTerms[term]; exists { + docLen := float64(idx.docLength[docID]) + numerator := float64(tf) * (idx.params.K1 + 1) + denominator := float64(tf) + idx.params.K1*(1-idx.params.B+idx.params.B*docLen/idx.avgDocLength) + scores[docID] += idf * numerator / denominator + } + } + } + } + + // Convert scores to results + results := make([]SearchResult, 0, len(scores)) + for docID, score := range scores { + results = append(results, SearchResult{ + ID: docID, + Score: score, + Fields: map[string]interface{}{ + "Text": idx.docs[docID], + "Metadata": idx.metadata[docID], + }, + }) + } + + // Sort by score + sort.Slice(results, func(i, j int) bool { + return results[i].Score > results[j].Score + }) + + // Return top K results + if len(results) > topK { + results = results[:topK] + } + + return results, nil +} + +// SetParameters updates the BM25 parameters +func (idx *BM25Index) SetParameters(params BM25Parameters) { + idx.mu.Lock() + defer idx.mu.Unlock() + idx.params = params +} + +// SetPreprocessor sets a custom text preprocessing function +func (idx *BM25Index) SetPreprocessor(preprocessor func(string) []string) { + idx.mu.Lock() + defer idx.mu.Unlock() + idx.preprocessor = preprocessor +} diff --git a/retriever.go b/retriever.go index 265559b..15551cc 100644 --- a/retriever.go +++ b/retriever.go @@ -200,6 +200,11 @@ func (r *Retriever) Retrieve(ctx context.Context, query string) ([]RetrieverResu return results, nil } +// GetVectorDB returns the underlying vector database instance +func (r *Retriever) GetVectorDB() *VectorDB { + return r.vectorDB +} + // Configuration options func WithRetrieveCollection(name string) RetrieverOption {