├── .gitignore ├── LICENSE ├── README.md ├── chunker.go ├── concurrentloader.go ├── config └── config.go ├── contextual.go ├── contextual_rag.go ├── data ├── leaves.txt └── sky.txt ├── docs ├── chatbot_example.md ├── contextual_rag.md ├── examples.md ├── memory_context.md ├── rag.md ├── simple_rag.md └── summary.md ├── embedder.go ├── examples ├── chat │ ├── docs │ │ ├── embeddings.txt │ │ ├── golang_basics.txt │ │ ├── microservices.txt │ │ ├── rag_systems.txt │ │ ├── sample.txt │ │ └── vector_databases.txt │ ├── main.go │ ├── v2.go │ └── v3.go ├── chromem │ ├── README.md │ └── main.go ├── chunker_examples.go ├── concurrent_loader_example.go ├── contextual │ ├── custom_llm.go │ └── main.go ├── embedding_example.go ├── full_process.go ├── loader_example.go ├── memory_enhancer_example.go ├── milvus_example.go ├── parser_example.go ├── process_embedding_benchmark.go ├── processing_benchmark.go ├── recruit_example.go ├── simple │ └── main.go ├── tsne_example.go └── vectordb_example.go ├── go.mod ├── go.sum ├── loader.go ├── logger.go ├── memory_context.go ├── parser.go ├── rag.go ├── rag ├── chromem.go ├── chunk.go ├── embed.go ├── example_vectordb.go ├── load.go ├── log.go ├── memory.go ├── milvus.go ├── parse.go ├── providers │ ├── example_provider.go │ ├── openai.go │ └── register.go ├── reranker.go ├── sparse_index.go └── vector_interface.go ├── register.go ├── retriever.go ├── simple_rag.go └── vectordb.go /.gitignore: -------------------------------------------------------------------------------- 1 | # If you prefer the allow list template instead of the deny list, see community template: 2 | # https://github.com/github/gitignore/blob/main/community/Golang/Go.AllowList.gitignore 3 | # 4 | # Binaries for programs and plugins 5 | *.exe 6 | *.exe~ 7 | *.dll 8 | *.so 9 | *.dylib 10 | 11 | # Test binary, built with `go test -c` 12 | *.test 13 | 14 | # Output of the go coverage tool, specifically when used with LiteIDE 15 | *.out 16 | 17 | # Dependency directories (remove the comment below to include it) 18 | vendor/ 19 | testdata/ 20 | 21 | # Go workspace file 22 | go.work 23 | go.work.sum 24 | data/chromem.db/ 25 | 26 | .DS_Store 27 | 28 | *.json -------------------------------------------------------------------------------- /chunker.go: -------------------------------------------------------------------------------- 1 | // Package raggo provides a high-level interface for text chunking and token management, 2 | // designed for use in retrieval-augmented generation (RAG) applications. 3 | package raggo 4 | 5 | import ( 6 | "github.com/teilomillet/raggo/rag" 7 | ) 8 | 9 | // Chunk represents a piece of text with associated metadata including its content, 10 | // token count, and position within the original document. It tracks: 11 | // - The actual text content 12 | // - Number of tokens in the chunk 13 | // - Starting and ending sentence indices 14 | type Chunk = rag.Chunk 15 | 16 | // Chunker defines the interface for text chunking implementations. 17 | // Implementations of this interface provide strategies for splitting text 18 | // into semantically meaningful chunks while preserving context. 19 | type Chunker interface { 20 | // Chunk splits the input text into a slice of Chunks according to the 21 | // implementation's strategy. 22 | Chunk(text string) []Chunk 23 | } 24 | 25 | // TokenCounter defines the interface for counting tokens in text. 26 | // Different implementations can provide various tokenization strategies, 27 | // from simple word-based counting to model-specific subword tokenization. 28 | type TokenCounter interface { 29 | // Count returns the number of tokens in the given text according to 30 | // the implementation's tokenization strategy. 31 | Count(text string) int 32 | } 33 | 34 | // ChunkerOption is a function type for configuring Chunker instances. 35 | // It follows the functional options pattern for clean and flexible configuration. 36 | type ChunkerOption = rag.TextChunkerOption 37 | 38 | // NewChunker creates a new Chunker with the given options. 39 | // By default, it creates a TextChunker with: 40 | // - Chunk size: 200 tokens 41 | // - Chunk overlap: 50 tokens 42 | // - Default word-based token counter 43 | // - Basic sentence splitter 44 | // 45 | // Use the provided option functions to customize these settings. 46 | func NewChunker(options ...ChunkerOption) (Chunker, error) { 47 | return rag.NewTextChunker(options...) 48 | } 49 | 50 | // ChunkSize sets the target size of each chunk in tokens. 51 | // This determines how much text will be included in each chunk 52 | // before starting a new one. 53 | func ChunkSize(size int) ChunkerOption { 54 | return func(tc *rag.TextChunker) { 55 | tc.ChunkSize = size 56 | } 57 | } 58 | 59 | // ChunkOverlap sets the number of tokens that should overlap between 60 | // adjacent chunks. This helps maintain context across chunk boundaries 61 | // and improves retrieval quality. 62 | func ChunkOverlap(overlap int) ChunkerOption { 63 | return func(tc *rag.TextChunker) { 64 | tc.ChunkOverlap = overlap 65 | } 66 | } 67 | 68 | // WithTokenCounter sets a custom token counter implementation. 69 | // This allows you to use different tokenization strategies, such as: 70 | // - Word-based counting (DefaultTokenCounter) 71 | // - Model-specific tokenization (TikTokenCounter) 72 | // - Custom tokenization schemes 73 | func WithTokenCounter(counter TokenCounter) ChunkerOption { 74 | return func(tc *rag.TextChunker) { 75 | tc.TokenCounter = counter 76 | } 77 | } 78 | 79 | // WithSentenceSplitter sets a custom sentence splitter function. 80 | // The function should take a string and return a slice of strings, 81 | // where each string is a sentence. This allows for: 82 | // - Custom sentence boundary detection 83 | // - Language-specific splitting rules 84 | // - Special handling of abbreviations or formatting 85 | func WithSentenceSplitter(splitter func(string) []string) ChunkerOption { 86 | return func(tc *rag.TextChunker) { 87 | tc.SentenceSplitter = splitter 88 | } 89 | } 90 | 91 | // DefaultSentenceSplitter returns the basic sentence splitter function 92 | // that splits text on common punctuation marks (., !, ?). 93 | // Suitable for simple English text without complex formatting. 94 | func DefaultSentenceSplitter() func(string) []string { 95 | return rag.DefaultSentenceSplitter 96 | } 97 | 98 | // SmartSentenceSplitter returns an advanced sentence splitter that handles: 99 | // - Multiple punctuation marks 100 | // - Quoted sentences 101 | // - Parenthetical content 102 | // - Lists and enumerations 103 | // 104 | // Recommended for complex text with varied formatting and structure. 105 | func SmartSentenceSplitter() func(string) []string { 106 | return rag.SmartSentenceSplitter 107 | } 108 | 109 | // NewDefaultTokenCounter creates a simple word-based token counter 110 | // that splits text on whitespace. Suitable for basic use cases 111 | // where exact token counts aren't critical. 112 | func NewDefaultTokenCounter() TokenCounter { 113 | return &rag.DefaultTokenCounter{} 114 | } 115 | 116 | // NewTikTokenCounter creates a token counter using the tiktoken library, 117 | // which implements the same tokenization used by OpenAI models. 118 | // The encoding parameter specifies which tokenization model to use 119 | // (e.g., "cl100k_base" for GPT-4, "p50k_base" for GPT-3). 120 | func NewTikTokenCounter(encoding string) (TokenCounter, error) { 121 | return rag.NewTikTokenCounter(encoding) 122 | } 123 | -------------------------------------------------------------------------------- /concurrentloader.go: -------------------------------------------------------------------------------- 1 | // Package raggo provides utilities for concurrent document loading and processing. 2 | package raggo 3 | 4 | import ( 5 | "context" 6 | "fmt" 7 | "io" 8 | "math" 9 | "os" 10 | "path/filepath" 11 | "strings" 12 | "sync" 13 | 14 | "github.com/teilomillet/raggo/rag" 15 | ) 16 | 17 | // ConcurrentPDFLoader extends the basic Loader interface with concurrent PDF processing 18 | // capabilities. It provides efficient handling of multiple PDF files by: 19 | // - Loading files in parallel using goroutines 20 | // - Managing concurrent file operations safely 21 | // - Handling file duplication when needed 22 | // - Providing progress tracking and error handling 23 | type ConcurrentPDFLoader interface { 24 | // Embeds the basic Loader interface 25 | Loader 26 | 27 | // LoadPDFsConcurrent loads a specified number of PDF files concurrently from a source directory. 28 | // If the source directory contains fewer files than the requested count, it automatically 29 | // duplicates existing PDFs to reach the desired number. 30 | // 31 | // Parameters: 32 | // - ctx: Context for cancellation and timeout 33 | // - sourceDir: Directory containing source PDF files 34 | // - targetDir: Directory where duplicated PDFs will be stored 35 | // - count: Desired number of PDF files to load 36 | // 37 | // Returns: 38 | // - []string: Paths to all successfully loaded files 39 | // - error: Any error encountered during the process 40 | // 41 | // Example usage: 42 | // loader := raggo.NewConcurrentPDFLoader(raggo.SetTimeout(1*time.Minute)) 43 | // files, err := loader.LoadPDFsConcurrent(ctx, "source", "target", 10) 44 | LoadPDFsConcurrent(ctx context.Context, sourceDir string, targetDir string, count int) ([]string, error) 45 | } 46 | 47 | // concurrentPDFLoaderWrapper wraps the internal loader and adds concurrent PDF loading capability. 48 | // It implements thread-safe operations and efficient resource management. 49 | type concurrentPDFLoaderWrapper struct { 50 | internal *rag.Loader 51 | } 52 | 53 | // NewConcurrentPDFLoader creates a new ConcurrentPDFLoader with the given options. 54 | // It supports all standard loader options plus concurrent processing capabilities. 55 | // 56 | // Options can include: 57 | // - SetTimeout: Maximum time for loading operations 58 | // - SetTempDir: Directory for temporary files 59 | // - SetRetryCount: Number of retries for failed operations 60 | // 61 | // Example: 62 | // loader := raggo.NewConcurrentPDFLoader( 63 | // raggo.SetTimeout(1*time.Minute), 64 | // raggo.SetTempDir(os.TempDir()), 65 | // ) 66 | func NewConcurrentPDFLoader(opts ...LoaderOption) ConcurrentPDFLoader { 67 | return &concurrentPDFLoaderWrapper{internal: rag.NewLoader(opts...)} 68 | } 69 | 70 | // LoadPDFsConcurrent implements the concurrent PDF loading strategy. 71 | // It performs the following steps: 72 | // 1. Lists all PDF files in the source directory 73 | // 2. Creates the target directory if it doesn't exist 74 | // 3. Duplicates PDFs if necessary to reach the desired count 75 | // 4. Loads files concurrently using goroutines 76 | // 5. Collects results and errors from concurrent operations 77 | // 78 | // The function uses channels for thread-safe communication and a WaitGroup 79 | // to ensure all operations complete before returning. 80 | func (clw *concurrentPDFLoaderWrapper) LoadPDFsConcurrent(ctx context.Context, sourceDir string, targetDir string, count int) ([]string, error) { 81 | pdfs, err := listPDFFiles(sourceDir) 82 | if err != nil { 83 | return nil, fmt.Errorf("failed to list PDF files in directory: %w", err) 84 | } 85 | 86 | if len(pdfs) == 0 { 87 | return nil, fmt.Errorf("no PDF files found in directory: %s", sourceDir) 88 | } 89 | 90 | // Create target directory if it doesn't exist 91 | if err := os.MkdirAll(targetDir, 0755); err != nil { 92 | return nil, fmt.Errorf("failed to create target directory: %w", err) 93 | } 94 | 95 | // Duplicate PDFs if necessary 96 | duplicatedPDFs, err := duplicatePDFs(pdfs, targetDir, count) 97 | if err != nil { 98 | return nil, fmt.Errorf("failed to duplicate PDFs: %w", err) 99 | } 100 | 101 | var wg sync.WaitGroup 102 | results := make(chan string, len(duplicatedPDFs)) 103 | errors := make(chan error, len(duplicatedPDFs)) 104 | 105 | for _, pdf := range duplicatedPDFs { 106 | wg.Add(1) 107 | go func(pdfPath string) { 108 | defer wg.Done() 109 | loadedPath, err := clw.internal.LoadFile(ctx, pdfPath) 110 | if err != nil { 111 | errors <- err 112 | return 113 | } 114 | results <- loadedPath 115 | }(pdf) 116 | } 117 | 118 | go func() { 119 | wg.Wait() 120 | close(results) 121 | close(errors) 122 | }() 123 | 124 | var loadedFiles []string 125 | var loadErrors []error 126 | 127 | for i := 0; i < len(duplicatedPDFs); i++ { 128 | select { 129 | case result := <-results: 130 | loadedFiles = append(loadedFiles, result) 131 | case err := <-errors: 132 | loadErrors = append(loadErrors, err) 133 | } 134 | } 135 | 136 | if len(loadErrors) > 0 { 137 | return loadedFiles, fmt.Errorf("encountered %d errors during loading", len(loadErrors)) 138 | } 139 | 140 | return loadedFiles, nil 141 | } 142 | 143 | // LoadURL implements the Loader interface by loading a document from a URL. 144 | // This method is inherited from the basic Loader interface. 145 | func (clw *concurrentPDFLoaderWrapper) LoadURL(ctx context.Context, url string) (string, error) { 146 | return clw.internal.LoadURL(ctx, url) 147 | } 148 | 149 | // LoadFile implements the Loader interface by loading a single file. 150 | // This method is inherited from the basic Loader interface. 151 | func (clw *concurrentPDFLoaderWrapper) LoadFile(ctx context.Context, path string) (string, error) { 152 | return clw.internal.LoadFile(ctx, path) 153 | } 154 | 155 | // LoadDir implements the Loader interface by loading all files in a directory. 156 | // This method is inherited from the basic Loader interface. 157 | func (clw *concurrentPDFLoaderWrapper) LoadDir(ctx context.Context, dir string) ([]string, error) { 158 | return clw.internal.LoadDir(ctx, dir) 159 | } 160 | 161 | // listPDFFiles returns a list of all PDF files in the given directory. 162 | // It recursively walks through the directory tree and identifies files 163 | // with a .pdf extension (case-insensitive). 164 | func listPDFFiles(dir string) ([]string, error) { 165 | var pdfs []string 166 | err := filepath.Walk(dir, func(path string, info os.FileInfo, err error) error { 167 | if err != nil { 168 | return err 169 | } 170 | if !info.IsDir() && strings.ToLower(filepath.Ext(path)) == ".pdf" { 171 | pdfs = append(pdfs, path) 172 | } 173 | return nil 174 | }) 175 | return pdfs, err 176 | } 177 | 178 | // duplicatePDFs duplicates the given PDF files to reach the desired count. 179 | // If the number of source PDFs is less than the desired count, it creates 180 | // copies with unique names by appending a counter to the original filename. 181 | // 182 | // The function ensures that: 183 | // - Each copy has a unique name 184 | // - The total number of files matches the desired count 185 | // - File copying is performed safely 186 | func duplicatePDFs(pdfs []string, targetDir string, desiredCount int) ([]string, error) { 187 | var duplicatedPDFs []string 188 | numOriginalPDFs := len(pdfs) 189 | 190 | if numOriginalPDFs >= desiredCount { 191 | return pdfs[:desiredCount], nil 192 | } 193 | 194 | duplicationsNeeded := int(math.Ceil(float64(desiredCount) / float64(numOriginalPDFs))) 195 | 196 | for i := 0; i < duplicationsNeeded; i++ { 197 | for _, pdf := range pdfs { 198 | if len(duplicatedPDFs) >= desiredCount { 199 | break 200 | } 201 | 202 | newFileName := fmt.Sprintf("%s_copy%d%s", strings.TrimSuffix(filepath.Base(pdf), ".pdf"), i, ".pdf") 203 | newFilePath := filepath.Join(targetDir, newFileName) 204 | 205 | if err := copyFile(pdf, newFilePath); err != nil { 206 | return nil, fmt.Errorf("failed to copy file %s: %w", pdf, err) 207 | } 208 | 209 | duplicatedPDFs = append(duplicatedPDFs, newFilePath) 210 | } 211 | } 212 | 213 | return duplicatedPDFs, nil 214 | } 215 | 216 | // copyFile performs a safe copy of a file from src to dst. 217 | // It handles: 218 | // - Opening source and destination files 219 | // - Proper resource cleanup with defer 220 | // - Efficient copying with io.Copy 221 | func copyFile(src, dst string) error { 222 | sourceFile, err := os.Open(src) 223 | if err != nil { 224 | return err 225 | } 226 | defer sourceFile.Close() 227 | 228 | destFile, err := os.Create(dst) 229 | if err != nil { 230 | return err 231 | } 232 | defer destFile.Close() 233 | 234 | _, err = io.Copy(destFile, sourceFile) 235 | return err 236 | } 237 | -------------------------------------------------------------------------------- /config/config.go: -------------------------------------------------------------------------------- 1 | // Package config provides a flexible configuration management system for the Raggo 2 | // Retrieval-Augmented Generation (RAG) framework. It handles configuration loading, 3 | // validation, and persistence with support for multiple sources: 4 | // - Configuration files (JSON) 5 | // - Environment variables 6 | // - Programmatic defaults 7 | // 8 | // The package implements a hierarchical configuration system where settings can be 9 | // overridden in the following order (highest to lowest precedence): 10 | // 1. Environment variables 11 | // 2. Configuration file 12 | // 3. Default values 13 | package config 14 | 15 | import ( 16 | "encoding/json" 17 | "os" 18 | "path/filepath" 19 | "time" 20 | ) 21 | 22 | // Config holds all configuration for the RAG system. It provides a centralized 23 | // way to manage settings across different components of the system. 24 | // 25 | // Configuration categories: 26 | // - Provider settings: Embedding and service providers 27 | // - Collection settings: Vector database collections 28 | // - Search settings: Retrieval and ranking parameters 29 | // - Document processing: Text chunking and batching 30 | // - Vector store: Database-specific configuration 31 | // - System settings: Timeouts, retries, and headers 32 | type Config struct { 33 | // Provider settings configure the embedding and service providers 34 | Provider string // Service provider (e.g., "milvus", "openai") 35 | Model string // Model identifier for embeddings 36 | APIKeys map[string]string // API keys for different providers 37 | 38 | // Collection settings define the vector database structure 39 | Collection string // Name of the vector collection 40 | 41 | // Search settings control retrieval behavior and ranking 42 | SearchStrategy string // Search method (e.g., "dense", "hybrid") 43 | DefaultTopK int // Default number of results to return 44 | DefaultMinScore float64 // Minimum similarity score threshold 45 | DefaultSearchParams map[string]interface{} // Additional search parameters 46 | EnableReRanking bool // Enable result re-ranking 47 | RRFConstant float64 // Reciprocal Rank Fusion constant 48 | 49 | // Document processing settings for text handling 50 | DefaultChunkSize int // Size of text chunks 51 | DefaultChunkOverlap int // Overlap between consecutive chunks 52 | DefaultBatchSize int // Number of items per processing batch 53 | DefaultIndexType string // Type of vector index (e.g., "HNSW") 54 | 55 | // Vector store settings for database configuration 56 | VectorDBConfig map[string]interface{} // Database-specific settings 57 | 58 | // Timeouts and retries for system operations 59 | Timeout time.Duration // Operation timeout 60 | MaxRetries int // Maximum retry attempts 61 | 62 | // Additional settings for extended functionality 63 | ExtraHeaders map[string]string // Additional HTTP headers 64 | } 65 | 66 | // LoadConfig loads configuration from multiple sources, combining them according 67 | // to the precedence rules. It automatically searches for configuration files in 68 | // standard locations and applies environment variable overrides. 69 | // 70 | // Configuration file search paths: 71 | // 1. $RAGGO_CONFIG environment variable 72 | // 2. ~/.raggo/config.json 73 | // 3. ~/.config/raggo/config.json 74 | // 4. ./raggo.json 75 | // 76 | // Environment variable overrides: 77 | // - RAGGO_PROVIDER: Service provider 78 | // - RAGGO_MODEL: Model identifier 79 | // - RAGGO_COLLECTION: Collection name 80 | // - RAGGO_API_KEY: Default API key 81 | // 82 | // Example usage: 83 | // 84 | // cfg, err := config.LoadConfig() 85 | // if err != nil { 86 | // log.Fatal(err) 87 | // } 88 | // fmt.Printf("Using provider: %s\n", cfg.Provider) 89 | func LoadConfig() (*Config, error) { 90 | // Default configuration with production-ready settings 91 | cfg := &Config{ 92 | Provider: "milvus", // Fast, open-source vector database 93 | Model: "text-embedding-3-small", // Latest OpenAI embedding model 94 | Collection: "documents", // Default collection name 95 | SearchStrategy: "dense", // Pure vector similarity search 96 | DefaultTopK: 5, // Conservative number of results 97 | DefaultMinScore: 0.7, // High confidence threshold 98 | DefaultChunkSize: 512, // Balanced chunk size 99 | DefaultChunkOverlap: 50, // Moderate overlap 100 | DefaultBatchSize: 100, // Efficient batch size 101 | DefaultIndexType: "HNSW", // Fast approximate search 102 | DefaultSearchParams: map[string]interface{}{ 103 | "ef": 64, // HNSW search depth 104 | }, 105 | EnableReRanking: false, // Disabled by default 106 | RRFConstant: 60, // Standard RRF constant 107 | Timeout: 30 * time.Second, // Conservative timeout 108 | MaxRetries: 3, // Reasonable retry count 109 | APIKeys: make(map[string]string), 110 | ExtraHeaders: make(map[string]string), 111 | VectorDBConfig: make(map[string]interface{}), 112 | } 113 | 114 | // Try to load from config file 115 | configFile := os.Getenv("RAGGO_CONFIG") 116 | if configFile == "" { 117 | // Try default locations 118 | home, err := os.UserHomeDir() 119 | if err == nil { 120 | candidates := []string{ 121 | filepath.Join(home, ".raggo", "config.json"), 122 | filepath.Join(home, ".config", "raggo", "config.json"), 123 | "raggo.json", 124 | } 125 | 126 | for _, candidate := range candidates { 127 | if _, err := os.Stat(candidate); err == nil { 128 | configFile = candidate 129 | break 130 | } 131 | } 132 | } 133 | } 134 | 135 | if configFile != "" { 136 | data, err := os.ReadFile(configFile) 137 | if err == nil { 138 | if err := json.Unmarshal(data, cfg); err != nil { 139 | return nil, err 140 | } 141 | } 142 | } 143 | 144 | // Override with environment variables 145 | if provider := os.Getenv("RAGGO_PROVIDER"); provider != "" { 146 | cfg.Provider = provider 147 | } 148 | if model := os.Getenv("RAGGO_MODEL"); model != "" { 149 | cfg.Model = model 150 | } 151 | if collection := os.Getenv("RAGGO_COLLECTION"); collection != "" { 152 | cfg.Collection = collection 153 | } 154 | if apiKey := os.Getenv("RAGGO_API_KEY"); apiKey != "" { 155 | cfg.APIKeys[cfg.Provider] = apiKey 156 | } 157 | 158 | return cfg, nil 159 | } 160 | 161 | // Save persists the configuration to a JSON file at the specified path. 162 | // It creates any necessary parent directories and sets appropriate file 163 | // permissions. 164 | // 165 | // Example usage: 166 | // 167 | // cfg := &Config{ 168 | // Provider: "milvus", 169 | // Model: "text-embedding-3-small", 170 | // } 171 | // err := cfg.Save("~/.raggo/config.json") 172 | // if err != nil { 173 | // log.Fatal(err) 174 | // } 175 | func (c *Config) Save(path string) error { 176 | data, err := json.MarshalIndent(c, "", " ") 177 | if err != nil { 178 | return err 179 | } 180 | 181 | // Ensure directory exists 182 | dir := filepath.Dir(path) 183 | if err := os.MkdirAll(dir, 0755); err != nil { 184 | return err 185 | } 186 | 187 | return os.WriteFile(path, data, 0644) 188 | } 189 | -------------------------------------------------------------------------------- /contextual.go: -------------------------------------------------------------------------------- 1 | // Package raggo provides advanced Retrieval-Augmented Generation (RAG) capabilities 2 | // with contextual awareness and memory management. 3 | package raggo 4 | 5 | import ( 6 | "context" 7 | "fmt" 8 | ) 9 | 10 | // ContextualStoreOptions configures how documents are processed and stored with 11 | // contextual information. It provides settings for: 12 | // - Vector database collection management 13 | // - Document chunking and processing 14 | // - Embedding model configuration 15 | // - Batch processing controls 16 | // 17 | // This configuration is designed to optimize the balance between processing 18 | // efficiency and context preservation. 19 | type ContextualStoreOptions struct { 20 | // Collection specifies the vector database collection name 21 | Collection string 22 | 23 | // APIKey is the authentication key for the embedding provider 24 | APIKey string 25 | 26 | // ChunkSize determines the size of text chunks in tokens 27 | // Larger chunks preserve more context but use more memory 28 | ChunkSize int 29 | 30 | // ChunkOverlap controls how much text overlaps between chunks 31 | // More overlap helps maintain context across chunk boundaries 32 | ChunkOverlap int 33 | 34 | // BatchSize sets how many documents to process simultaneously 35 | // Higher values increase throughput but use more memory 36 | BatchSize int 37 | 38 | // ModelName specifies which language model to use for context generation 39 | // This model enriches chunks with additional contextual information 40 | ModelName string 41 | } 42 | 43 | // StoreWithContext processes documents and stores them with enhanced contextual information. 44 | // It uses a combination of: 45 | // - Semantic chunking for document segmentation 46 | // - Language model enrichment for context generation 47 | // - Vector embedding for efficient retrieval 48 | // - Batch processing for performance 49 | // 50 | // The function automatically handles: 51 | // - Default configuration values 52 | // - Resource management 53 | // - Error handling and reporting 54 | // - Context-aware processing 55 | // 56 | // Example usage: 57 | // 58 | // opts := raggo.ContextualStoreOptions{ 59 | // Collection: "my_docs", 60 | // APIKey: os.Getenv("OPENAI_API_KEY"), 61 | // ChunkSize: 512, 62 | // BatchSize: 100, 63 | // } 64 | // 65 | // err := raggo.StoreWithContext(ctx, "path/to/docs", opts) 66 | func StoreWithContext(ctx context.Context, source string, opts ContextualStoreOptions) error { 67 | // Use default values if not specified 68 | if opts.ChunkSize == 0 { 69 | opts.ChunkSize = 512 70 | } 71 | if opts.ChunkOverlap == 0 { 72 | opts.ChunkOverlap = 64 73 | } 74 | if opts.BatchSize == 0 { 75 | opts.BatchSize = 100 76 | } 77 | if opts.ModelName == "" { 78 | opts.ModelName = "gpt-4o-mini" 79 | } 80 | 81 | // Initialize RAG with context-aware configuration 82 | rag, err := NewRAG( 83 | WithMilvus(opts.Collection), 84 | WithOpenAI(opts.APIKey), 85 | func(c *RAGConfig) { 86 | c.ChunkSize = opts.ChunkSize 87 | c.ChunkOverlap = opts.ChunkOverlap 88 | c.BatchSize = opts.BatchSize 89 | }, 90 | ) 91 | if err != nil { 92 | return fmt.Errorf("failed to initialize RAG: %w", err) 93 | } 94 | defer rag.Close() 95 | 96 | // Process and store documents with enhanced context 97 | return rag.ProcessWithContext(ctx, source, opts.ModelName) 98 | } 99 | -------------------------------------------------------------------------------- /data/leaves.txt: -------------------------------------------------------------------------------- 1 | Leaves are green because chlorophyll absorbs red and blue light. -------------------------------------------------------------------------------- /data/sky.txt: -------------------------------------------------------------------------------- 1 | The sky is blue because of Rayleigh scattering. -------------------------------------------------------------------------------- /docs/contextual_rag.md: -------------------------------------------------------------------------------- 1 | # Contextual RAG Documentation 2 | 3 | ## Overview 4 | 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. 5 | 6 | ## Core Components 7 | 8 | ### ContextualRAG Struct 9 | ```go 10 | type ContextualRAG struct { 11 | rag *RAG 12 | contextStore map[string]string 13 | contextSize int 14 | contextWindow int 15 | } 16 | ``` 17 | 18 | ### Configuration 19 | The Contextual RAG system can be configured with various options: 20 | 21 | ```go 22 | type ContextualConfig struct { 23 | // Base RAG configuration 24 | RAGConfig *RAGConfig 25 | 26 | // Context settings 27 | ContextSize int // Maximum size of stored context 28 | ContextWindow int // Window size for context consideration 29 | UseMemory bool // Whether to use memory for context 30 | } 31 | ``` 32 | 33 | ## Key Features 34 | 35 | ### 1. Context Management 36 | - Maintains conversation history and context 37 | - Configurable context window and size 38 | - Automatic context pruning and relevance scoring 39 | 40 | ### 2. Enhanced Query Processing 41 | - Context-aware query understanding 42 | - Historical context integration 43 | - Improved response coherence 44 | 45 | ### 3. Memory Integration 46 | - Optional memory storage for long-term context 47 | - Configurable memory retention 48 | - Context-based memory retrieval 49 | 50 | ## Usage Examples 51 | 52 | ### Basic Setup 53 | ```go 54 | contextualRAG, err := raggo.NewContextualRAG( 55 | raggo.WithBaseRAG(baseRAG), 56 | raggo.WithContextSize(5), 57 | raggo.WithContextWindow(3), 58 | raggo.WithMemory(true), 59 | ) 60 | if err != nil { 61 | log.Fatal(err) 62 | } 63 | ``` 64 | 65 | ### Processing Queries with Context 66 | ```go 67 | // Process a query with context 68 | response, err := contextualRAG.ProcessQuery(ctx, "What is the latest development?") 69 | 70 | // Add context explicitly 71 | contextualRAG.AddContext("topic", "AI Development") 72 | response, err = contextualRAG.ProcessQuery(ctx, "What are the challenges?") 73 | ``` 74 | 75 | ## Best Practices 76 | 77 | 1. **Context Management** 78 | - Set appropriate context size based on use case 79 | - Regularly clean up old context 80 | - Monitor context relevance scores 81 | 82 | 2. **Query Processing** 83 | - Structure queries to leverage context 84 | - Use context hints when available 85 | - Monitor context window effectiveness 86 | 87 | 3. **Memory Usage** 88 | - Enable memory for long-running conversations 89 | - Configure memory retention appropriately 90 | - Implement context cleanup strategies 91 | 92 | ## Advanced Features 93 | 94 | ### Custom Context Processing 95 | ```go 96 | contextualRAG.SetContextProcessor(func(context, query string) string { 97 | // Custom context processing logic 98 | return processedContext 99 | }) 100 | ``` 101 | 102 | ### Context Filtering 103 | ```go 104 | contextualRAG.SetContextFilter(func(context string) bool { 105 | // Custom filtering logic 106 | return isRelevant 107 | }) 108 | ``` 109 | 110 | ## Example Use Cases 111 | 112 | ### 1. Multi-turn Conversations 113 | ```go 114 | // Initialize contextual RAG for conversation 115 | contextualRAG, _ := raggo.NewContextualRAG( 116 | raggo.WithContextSize(10), 117 | raggo.WithMemory(true), 118 | ) 119 | 120 | // Process conversation turns 121 | for { 122 | response, _ := contextualRAG.ProcessQuery(ctx, userQuery) 123 | contextualRAG.AddContext("conversation", response) 124 | } 125 | ``` 126 | 127 | ### 2. Document Analysis with Context 128 | ```go 129 | // Initialize for document analysis 130 | contextualRAG, _ := raggo.NewContextualRAG( 131 | raggo.WithContextWindow(5), 132 | raggo.WithDocumentContext(true), 133 | ) 134 | 135 | // Process document sections 136 | for _, section := range sections { 137 | contextualRAG.AddContext("document", section) 138 | analysis, _ := contextualRAG.ProcessQuery(ctx, "Analyze this section") 139 | } 140 | ``` 141 | 142 | ## Integration with Memory Systems 143 | 144 | The contextual RAG system can be integrated with various memory systems: 145 | 146 | ```go 147 | // Configure memory integration 148 | memorySystem := raggo.NewMemorySystem( 149 | raggo.WithMemorySize(1000), 150 | raggo.WithMemoryType("semantic"), 151 | ) 152 | 153 | contextualRAG.SetMemorySystem(memorySystem) 154 | ``` 155 | 156 | ## Performance Considerations 157 | 158 | 1. **Context Size** 159 | - Larger context sizes increase memory usage 160 | - Monitor context processing overhead 161 | - Balance context size with response time 162 | 163 | 2. **Memory Usage** 164 | - Implement context cleanup strategies 165 | - Monitor memory consumption 166 | - Use appropriate indexing for context retrieval 167 | 168 | 3. **Query Processing** 169 | - Optimize context matching algorithms 170 | - Cache frequently used contexts 171 | - Implement context relevance scoring 172 | 173 | ## Error Handling 174 | 175 | ```go 176 | // Handle context-related errors 177 | if err := contextualRAG.ProcessQuery(ctx, query); err != nil { 178 | switch err.(type) { 179 | case *ContextSizeError: 180 | // Handle context size exceeded 181 | case *ContextProcessingError: 182 | // Handle processing error 183 | default: 184 | // Handle other errors 185 | } 186 | } 187 | ``` 188 | -------------------------------------------------------------------------------- /docs/examples.md: -------------------------------------------------------------------------------- 1 | # RagGo Examples Documentation 2 | 3 | This document provides detailed explanations of the example implementations in the RagGo library. 4 | 5 | ## Memory Enhancer Example 6 | 7 | ### Overview 8 | The Memory Enhancer example demonstrates how to create a RAG system with enhanced memory capabilities for processing technical documentation and maintaining context across interactions. 9 | 10 | ### Key Components 11 | ```go 12 | // Main components used in the example 13 | - Vector Database (Milvus) 14 | - Memory Context 15 | - LLM Integration (OpenAI) 16 | ``` 17 | 18 | ### Implementation Details 19 | 20 | 1. **Setup and Initialization** 21 | ```go 22 | // Initialize LLM 23 | llm, err := gollm.NewLLM( 24 | gollm.SetProvider("openai"), 25 | gollm.SetModel("gpt-4o-mini"), 26 | gollm.SetAPIKey(apiKey), 27 | gollm.SetLogLevel(gollm.LogLevelInfo), 28 | ) 29 | 30 | // Initialize Vector Database 31 | vectorDB, err := raggo.NewVectorDB( 32 | raggo.WithType("milvus"), 33 | raggo.WithAddress("localhost:19530"), 34 | ) 35 | 36 | // Create Memory Context 37 | memoryContext, err := raggo.NewMemoryContext(apiKey, 38 | raggo.MemoryCollection("tech_docs"), 39 | raggo.MemoryTopK(5), 40 | raggo.MemoryMinScore(0.01), 41 | raggo.MemoryStoreLastN(10), 42 | raggo.MemoryStoreRAGInfo(true), 43 | ) 44 | ``` 45 | 46 | 2. **Document Processing** 47 | ```go 48 | // Load and process technical documentation 49 | docsDir := filepath.Join("examples", "chat", "docs") 50 | docs := []string{ 51 | filepath.Join(docsDir, "microservices.txt"), 52 | filepath.Join(docsDir, "vector_databases.txt"), 53 | // ... additional documents 54 | } 55 | 56 | for _, doc := range docs { 57 | content, err := os.ReadFile(doc) 58 | if err != nil { 59 | log.Printf("Warning: Failed to read %s: %v", doc, err) 60 | continue 61 | } 62 | // Store document content as memory 63 | err = memoryContext.Store(ctx, filepath.Base(doc), string(content)) 64 | if err != nil { 65 | log.Printf("Warning: Failed to store %s: %v", doc, err) 66 | } 67 | } 68 | ``` 69 | 70 | 3. **Interactive Query Processing** 71 | ```go 72 | // Process user queries with context 73 | scanner := bufio.NewScanner(os.Stdin) 74 | for { 75 | fmt.Print("\nEnter your question (or 'quit' to exit): ") 76 | if !scanner.Scan() { 77 | break 78 | } 79 | 80 | query := scanner.Text() 81 | if strings.ToLower(query) == "quit" { 82 | break 83 | } 84 | 85 | response, err := memoryContext.ProcessWithContext(ctx, query) 86 | if err != nil { 87 | log.Printf("Error processing query: %v", err) 88 | continue 89 | } 90 | 91 | fmt.Printf("\nResponse: %s\n", response) 92 | } 93 | ``` 94 | 95 | ## Best Practices Demonstrated 96 | 97 | 1. **Error Handling** 98 | - Proper error checking and logging 99 | - Graceful handling of document loading failures 100 | - User-friendly error messages 101 | 102 | 2. **Resource Management** 103 | - Proper initialization and cleanup of resources 104 | - Use of context for cancellation 105 | - Cleanup of vector database collections 106 | 107 | 3. **User Interaction** 108 | - Clear user prompts and instructions 109 | - Graceful exit handling 110 | - Informative response formatting 111 | 112 | ## Running the Example 113 | 114 | 1. **Prerequisites** 115 | ```bash 116 | # Set up environment variables 117 | export OPENAI_API_KEY=your_api_key 118 | 119 | # Ensure Milvus is running 120 | docker-compose up -d 121 | ``` 122 | 123 | 2. **Running the Example** 124 | ```bash 125 | go run examples/memory_enhancer_example.go 126 | ``` 127 | 128 | 3. **Example Interactions** 129 | ``` 130 | Enter your question: What are microservices? 131 | Response: [Detailed response about microservices based on loaded documentation] 132 | 133 | Enter your question: How do vector databases work? 134 | Response: [Context-aware response about vector databases] 135 | ``` 136 | 137 | ## Customization Options 138 | 139 | 1. **Document Sources** 140 | - Modify the `docs` slice to include different document sources 141 | - Adjust document processing logic for different file types 142 | 143 | 2. **Memory Settings** 144 | - Adjust `TopK` for different numbers of similar contexts 145 | - Modify `MinScore` for stricter/looser similarity matching 146 | - Change `StoreLastN` for different memory retention 147 | 148 | 3. **Model Configuration** 149 | - Change the LLM model for different capabilities 150 | - Adjust logging levels for debugging 151 | - Modify vector database settings 152 | 153 | ## Additional Examples 154 | 155 | ### Contextual Example 156 | Located in `examples/contextual/main.go`, this example demonstrates: 157 | - Advanced context management 158 | - Multi-turn conversations 159 | - Context-aware document processing 160 | 161 | ### Basic RAG Example 162 | Located in `examples/basic/main.go`, this example shows: 163 | - Simple RAG setup 164 | - Basic document processing 165 | - Query handling without memory enhancement 166 | -------------------------------------------------------------------------------- /docs/memory_context.md: -------------------------------------------------------------------------------- 1 | # Memory Context Documentation 2 | 3 | ## Overview 4 | 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. 5 | 6 | ## Core Components 7 | 8 | ### MemoryContext Struct 9 | The main component that manages contextual memory: 10 | 11 | ```go 12 | type MemoryContext struct { 13 | retriever *RAG 14 | config *MemoryConfig 15 | lastN int 16 | storeRAG bool 17 | } 18 | ``` 19 | 20 | ### Configuration 21 | Memory context can be configured using various options: 22 | 23 | ```go 24 | type MemoryConfig struct { 25 | Collection string // Vector DB collection name 26 | TopK int // Number of similar contexts to retrieve 27 | MinScore float64 // Minimum similarity score 28 | StoreLastN int // Number of recent interactions to store 29 | StoreRAG bool // Whether to store RAG information 30 | } 31 | ``` 32 | 33 | ## Key Features 34 | 35 | ### 1. Memory Management 36 | - Store and retrieve recent interactions 37 | - Configure the number of interactions to maintain 38 | - Automatic cleanup of old memories 39 | 40 | ### 2. Context Enhancement 41 | - Enrich queries with relevant historical context 42 | - Maintain conversation coherence 43 | - Support for multi-turn interactions 44 | 45 | ### 3. RAG Integration 46 | - Seamless integration with the RAG system 47 | - Enhanced document retrieval with historical context 48 | - Configurable similarity thresholds 49 | 50 | ## Usage Examples 51 | 52 | ### Basic Memory Context Setup 53 | ```go 54 | memoryContext, err := raggo.NewMemoryContext(apiKey, 55 | raggo.MemoryCollection("tech_docs"), 56 | raggo.MemoryTopK(5), 57 | raggo.MemoryMinScore(0.01), 58 | raggo.MemoryStoreLastN(10), 59 | raggo.MemoryStoreRAGInfo(true), 60 | ) 61 | if err != nil { 62 | log.Fatal(err) 63 | } 64 | ``` 65 | 66 | ### Using Memory Context in Applications 67 | ```go 68 | // Store new memory 69 | err = memoryContext.Store(ctx, "user query", "system response") 70 | 71 | // Retrieve relevant memories 72 | memories, err := memoryContext.Retrieve(ctx, "current query") 73 | 74 | // Process with context 75 | response, err := memoryContext.ProcessWithContext(ctx, "user query") 76 | ``` 77 | 78 | ## Best Practices 79 | 80 | 1. **Memory Configuration** 81 | - Set appropriate `StoreLastN` based on your use case 82 | - Configure `TopK` and `MinScore` for optimal context retrieval 83 | - Enable `StoreRAG` for enhanced context awareness 84 | 85 | 2. **Performance Considerations** 86 | - Monitor memory usage with large conversation histories 87 | - Use appropriate batch sizes for memory operations 88 | - Implement cleanup strategies for old memories 89 | 90 | 3. **Context Quality** 91 | - Regularly evaluate the quality of retrieved contexts 92 | - Adjust similarity thresholds based on application needs 93 | - Consider implementing context filtering mechanisms 94 | 95 | ## Advanced Features 96 | 97 | ### Custom Memory Processing 98 | Implement custom memory processing logic: 99 | 100 | ```go 101 | memoryContext.SetProcessor(func(ctx context.Context, memory Memory) (string, error) { 102 | // Custom processing logic 103 | return processedMemory, nil 104 | }) 105 | ``` 106 | 107 | ### Memory Filtering 108 | Apply filters to retrieved memories: 109 | 110 | ```go 111 | memoryContext.SetFilter(func(memory Memory) bool { 112 | // Custom filtering logic 113 | return shouldIncludeMemory 114 | }) 115 | ``` 116 | 117 | ## Example Use Cases 118 | 119 | ### 1. Chatbot Enhancement 120 | ```go 121 | // Initialize memory context for chat 122 | memoryContext, _ := raggo.NewMemoryContext(apiKey, 123 | raggo.MemoryCollection("chat_history"), 124 | raggo.MemoryStoreLastN(20), 125 | ) 126 | 127 | // Process chat messages with context 128 | for { 129 | response, _ := memoryContext.ProcessWithContext(ctx, userMessage) 130 | memoryContext.Store(ctx, userMessage, response) 131 | } 132 | ``` 133 | 134 | ### 2. Document Q&A System 135 | ```go 136 | // Initialize memory context for document Q&A 137 | memoryContext, _ := raggo.NewMemoryContext(apiKey, 138 | raggo.MemoryCollection("doc_qa"), 139 | raggo.MemoryTopK(3), 140 | raggo.MemoryStoreRAGInfo(true), 141 | ) 142 | 143 | // Process document queries with context 144 | response, _ := memoryContext.ProcessWithContext(ctx, userQuery) 145 | ``` 146 | 147 | ## Integration with Vector Databases 148 | 149 | The memory context system integrates seamlessly with vector databases for efficient storage and retrieval of contextual information: 150 | 151 | ```go 152 | // Configure vector database integration 153 | retriever := memoryContext.GetRetriever() 154 | if err := retriever.GetVectorDB().LoadCollection(ctx, "tech_docs"); err != nil { 155 | log.Fatal(err) 156 | } 157 | ``` 158 | -------------------------------------------------------------------------------- /docs/rag.md: -------------------------------------------------------------------------------- 1 | # RAG (Retrieval-Augmented Generation) Documentation 2 | 3 | ## Overview 4 | 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. 5 | 6 | ## Core Components 7 | 8 | ### RAG Struct 9 | The main component that provides a unified interface for document processing and retrieval: 10 | 11 | ```go 12 | type RAG struct { 13 | db *VectorDB 14 | embedder *EmbeddingService 15 | config *RAGConfig 16 | } 17 | ``` 18 | 19 | ### Configuration 20 | The `RAGConfig` struct holds all RAG settings: 21 | 22 | ```go 23 | type RAGConfig struct { 24 | // Database settings 25 | DBType string 26 | DBAddress string 27 | Collection string 28 | AutoCreate bool 29 | IndexType string 30 | IndexMetric string 31 | 32 | // Processing settings 33 | ChunkSize int 34 | ChunkOverlap int 35 | BatchSize int 36 | 37 | // Embedding settings 38 | Provider string 39 | Model string // For embeddings 40 | LLMModel string // For LLM operations 41 | APIKey string 42 | 43 | // Search settings 44 | TopK int 45 | MinScore float64 46 | UseHybrid bool 47 | 48 | // System settings 49 | Timeout time.Duration 50 | TempDir string 51 | Debug bool 52 | } 53 | ``` 54 | 55 | ## Key Features 56 | 57 | ### 1. Document Processing 58 | - Chunking documents with configurable size and overlap 59 | - Enriching chunks with contextual information 60 | - Batch processing for efficient handling of large documents 61 | 62 | ### 2. Search and Retrieval 63 | - Simple vector similarity search 64 | - Hybrid search combining vector and keyword-based approaches 65 | - Configurable search parameters (TopK, MinScore) 66 | 67 | ### 3. Vector Database Integration 68 | - Default support for Milvus 69 | - Extensible design for other vector databases 70 | - Automatic collection creation and management 71 | 72 | ### 4. Embedding Services 73 | - Integration with OpenAI embeddings 74 | - Configurable embedding models 75 | - Extensible for other embedding providers 76 | 77 | ## Usage Examples 78 | 79 | ### Basic RAG Setup 80 | ```go 81 | rag, err := raggo.NewRAG( 82 | raggo.WithOpenAI(apiKey), 83 | raggo.WithMilvus("documents"), 84 | raggo.SetChunkSize(512), 85 | raggo.SetTopK(5), 86 | ) 87 | if err != nil { 88 | log.Fatal(err) 89 | } 90 | defer rag.Close() 91 | ``` 92 | 93 | ### Loading Documents 94 | ```go 95 | ctx := context.Background() 96 | err = rag.LoadDocuments(ctx, "path/to/documents") 97 | if err != nil { 98 | log.Fatal(err) 99 | } 100 | ``` 101 | 102 | ### Querying 103 | ```go 104 | results, err := rag.Query(ctx, "your search query") 105 | if err != nil { 106 | log.Fatal(err) 107 | } 108 | ``` 109 | 110 | ## Best Practices 111 | 112 | 1. **Configuration** 113 | - Use appropriate chunk sizes based on your content (default: 512) 114 | - Adjust TopK and MinScore based on your use case 115 | - Enable hybrid search for better results when appropriate 116 | 117 | 2. **Performance** 118 | - Use batch processing for large document sets 119 | - Configure appropriate timeouts 120 | - Monitor vector database performance 121 | 122 | 3. **Error Handling** 123 | - Always handle errors appropriately 124 | - Use context for cancellation and timeouts 125 | - Close resources using defer 126 | 127 | 4. **Security** 128 | - Never hardcode API keys 129 | - Use environment variables for sensitive configuration 130 | - Implement appropriate access controls for your vector database 131 | 132 | ## Advanced Features 133 | 134 | ### Context-Aware Processing 135 | The RAG system can enrich document chunks with contextual information: 136 | 137 | ```go 138 | err = rag.ProcessWithContext(ctx, "path/to/documents", "gpt-4") 139 | ``` 140 | 141 | ### Custom Search Parameters 142 | Configure specific search parameters for your use case: 143 | 144 | ```go 145 | rag, err := raggo.NewRAG( 146 | raggo.SetSearchParams(map[string]interface{}{ 147 | "nprobe": 10, 148 | "ef": 64, 149 | "type": "HNSW", 150 | }), 151 | ) 152 | ``` 153 | -------------------------------------------------------------------------------- /docs/simple_rag.md: -------------------------------------------------------------------------------- 1 | # Simple RAG Documentation 2 | 3 | ## Overview 4 | 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. 5 | 6 | ## Core Components 7 | 8 | ### SimpleRAG Struct 9 | ```go 10 | type SimpleRAG struct { 11 | embedder *EmbeddingService 12 | vectorStore *VectorDB 13 | config *SimpleConfig 14 | } 15 | ``` 16 | 17 | ### Configuration 18 | Simple configuration options for basic RAG operations: 19 | 20 | ```go 21 | type SimpleConfig struct { 22 | // Vector store settings 23 | VectorDBType string 24 | VectorDBAddress string 25 | Collection string 26 | 27 | // Embedding settings 28 | EmbeddingModel string 29 | APIKey string 30 | 31 | // Search settings 32 | TopK int 33 | MinScore float64 34 | } 35 | ``` 36 | 37 | ## Key Features 38 | 39 | ### 1. Simplified Document Processing 40 | - Straightforward document ingestion 41 | - Basic chunking and embedding 42 | - Direct vector storage 43 | 44 | ### 2. Basic Search Functionality 45 | - Simple similarity search 46 | - Configurable result count 47 | - Basic relevance scoring 48 | 49 | ### 3. Minimal Setup Required 50 | - Default configurations 51 | - Automatic resource management 52 | - Simplified API 53 | 54 | ## Usage Examples 55 | 56 | ### Basic Setup 57 | ```go 58 | simpleRAG, err := raggo.NewSimpleRAG( 59 | raggo.WithVectorDB("milvus", "localhost:19530"), 60 | raggo.WithEmbeddings("openai", apiKey), 61 | raggo.WithTopK(3), 62 | ) 63 | if err != nil { 64 | log.Fatal(err) 65 | } 66 | ``` 67 | 68 | ### Document Processing 69 | ```go 70 | // Process a single document 71 | err = simpleRAG.AddDocument(ctx, "document.txt") 72 | 73 | // Process multiple documents 74 | err = simpleRAG.AddDocuments(ctx, []string{ 75 | "doc1.txt", 76 | "doc2.txt", 77 | }) 78 | ``` 79 | 80 | ### Query Processing 81 | ```go 82 | // Simple query 83 | results, err := simpleRAG.Query(ctx, "How does this work?") 84 | 85 | // Query with custom parameters 86 | results, err = simpleRAG.QueryWithParams(ctx, "How does this work?", 87 | raggo.WithResultCount(5), 88 | raggo.WithMinScore(0.5), 89 | ) 90 | ``` 91 | 92 | ## Best Practices 93 | 94 | 1. **Document Management** 95 | - Keep documents reasonably sized 96 | - Use appropriate file formats 97 | - Maintain clean document structure 98 | 99 | 2. **Query Optimization** 100 | - Keep queries focused and specific 101 | - Use appropriate TopK values 102 | - Monitor query performance 103 | 104 | 3. **Resource Management** 105 | - Close resources when done 106 | - Monitor memory usage 107 | - Use batch processing for large datasets 108 | 109 | ## Example Use Cases 110 | 111 | ### 1. Document Q&A System 112 | ```go 113 | // Initialize simple RAG for Q&A 114 | simpleRAG, _ := raggo.NewSimpleRAG( 115 | raggo.WithCollection("qa_docs"), 116 | raggo.WithTopK(1), 117 | ) 118 | 119 | // Add documentation 120 | simpleRAG.AddDocument(ctx, "documentation.txt") 121 | 122 | // Process questions 123 | answer, _ := simpleRAG.Query(ctx, "What is the installation process?") 124 | ``` 125 | 126 | ### 2. Basic Search System 127 | ```go 128 | // Initialize for search 129 | simpleRAG, _ := raggo.NewSimpleRAG( 130 | raggo.WithTopK(5), 131 | raggo.WithMinScore(0.7), 132 | ) 133 | 134 | // Add searchable content 135 | simpleRAG.AddDocuments(ctx, []string{ 136 | "content1.txt", 137 | "content2.txt", 138 | }) 139 | 140 | // Search content 141 | results, _ := simpleRAG.Query(ctx, "search term") 142 | ``` 143 | 144 | ## Performance Tips 145 | 146 | 1. **Document Processing** 147 | - Use batch processing for multiple documents 148 | - Monitor embedding API usage 149 | - Implement rate limiting for API calls 150 | 151 | 2. **Query Optimization** 152 | - Cache frequent queries 153 | - Use appropriate TopK values 154 | - Monitor query latency 155 | 156 | 3. **Resource Usage** 157 | - Implement connection pooling 158 | - Monitor memory consumption 159 | - Use appropriate batch sizes 160 | 161 | ## Error Handling 162 | 163 | ```go 164 | // Basic error handling 165 | if err := simpleRAG.AddDocument(ctx, "doc.txt"); err != nil { 166 | switch err.(type) { 167 | case *FileError: 168 | // Handle file-related errors 169 | case *ProcessingError: 170 | // Handle processing errors 171 | default: 172 | // Handle other errors 173 | } 174 | } 175 | ``` 176 | 177 | ## Integration Examples 178 | 179 | ### With HTTP Server 180 | ```go 181 | func handleQuery(w http.ResponseWriter, r *http.Request) { 182 | query := r.URL.Query().Get("q") 183 | results, err := simpleRAG.Query(r.Context(), query) 184 | if err != nil { 185 | http.Error(w, err.Error(), http.StatusInternalServerError) 186 | return 187 | } 188 | json.NewEncoder(w).Encode(results) 189 | } 190 | ``` 191 | 192 | ### With CLI Application 193 | ```go 194 | func main() { 195 | simpleRAG, _ := raggo.NewSimpleRAG(/* config */) 196 | scanner := bufio.NewScanner(os.Stdin) 197 | 198 | for { 199 | fmt.Print("Enter query: ") 200 | if !scanner.Scan() { 201 | break 202 | } 203 | 204 | results, _ := simpleRAG.Query(context.Background(), scanner.Text()) 205 | fmt.Printf("Results: %v\n", results) 206 | } 207 | } 208 | ``` 209 | 210 | ## Limitations and Considerations 211 | 212 | 1. **No Advanced Context Management** 213 | - Limited to single-query operations 214 | - No conversation history 215 | - No context window management 216 | 217 | 2. **Basic Search Only** 218 | - No hybrid search capabilities 219 | - Limited to vector similarity 220 | - Basic relevance scoring 221 | 222 | 3. **Limited Customization** 223 | - Fixed chunking strategy 224 | - Basic embedding options 225 | - Simple configuration options 226 | -------------------------------------------------------------------------------- /docs/summary.md: -------------------------------------------------------------------------------- 1 | # RagGo Library Documentation Summary 2 | 3 | ## Overview 4 | 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. 5 | 6 | ## Components Overview 7 | 8 | ### 1. Core RAG (`rag.go`) 9 | The foundation of the library, providing basic RAG functionality: 10 | - Document processing and embedding 11 | - Vector database integration 12 | - Configurable search parameters 13 | - Extensible architecture 14 | 15 | [Detailed Documentation →](./rag.md) 16 | 17 | ### 2. Simple RAG (`simple_rag.go`) 18 | A streamlined implementation for basic use cases: 19 | - Minimal setup required 20 | - Basic document processing 21 | - Simple search functionality 22 | - Ideal for straightforward applications 23 | 24 | [Detailed Documentation →](./simple_rag.md) 25 | 26 | ### 3. Contextual RAG (`contextual_rag.go`) 27 | Advanced implementation with context management: 28 | - Conversation history tracking 29 | - Context-aware responses 30 | - Memory integration 31 | - Suitable for complex applications 32 | 33 | [Detailed Documentation →](./contextual_rag.md) 34 | 35 | ### 4. Memory Context (`memory_context.go`) 36 | Enhanced memory management system: 37 | - Long-term memory storage 38 | - Context retention 39 | - Configurable memory policies 40 | - Integration with RAG systems 41 | 42 | [Detailed Documentation →](./memory_context.md) 43 | 44 | ## Quick Start Guide 45 | 46 | ### 1. Basic Usage 47 | ```go 48 | // Initialize basic RAG 49 | rag, err := raggo.NewRAG( 50 | raggo.WithOpenAI(apiKey), 51 | raggo.WithMilvus("documents"), 52 | ) 53 | 54 | // Process documents 55 | err = rag.LoadDocuments(ctx, "path/to/docs") 56 | 57 | // Query 58 | results, err := rag.Query(ctx, "your query") 59 | ``` 60 | 61 | ### 2. Simple RAG Usage 62 | ```go 63 | // Initialize simple RAG 64 | simpleRAG, err := raggo.NewSimpleRAG( 65 | raggo.WithVectorDB("milvus", "localhost:19530"), 66 | raggo.WithEmbeddings("openai", apiKey), 67 | ) 68 | 69 | // Process and query 70 | err = simpleRAG.AddDocument(ctx, "document.txt") 71 | results, err := simpleRAG.Query(ctx, "query") 72 | ``` 73 | 74 | ### 3. Contextual RAG Usage 75 | ```go 76 | // Initialize contextual RAG 77 | contextualRAG, err := raggo.NewContextualRAG( 78 | raggo.WithBaseRAG(baseRAG), 79 | raggo.WithContextSize(5), 80 | raggo.WithMemory(true), 81 | ) 82 | 83 | // Process with context 84 | response, err := contextualRAG.ProcessQuery(ctx, "query") 85 | ``` 86 | 87 | ## Feature Comparison 88 | 89 | | Feature | Simple RAG | Core RAG | Contextual RAG | 90 | |---------------------------|------------|----------|----------------| 91 | | Document Processing | Basic | Advanced | Advanced | 92 | | Context Management | No | Basic | Advanced | 93 | | Memory Integration | No | Optional | Yes | 94 | | Search Capabilities | Basic | Advanced | Advanced | 95 | | Setup Complexity | Low | Medium | High | 96 | | Resource Requirements | Low | Medium | High | 97 | 98 | ## Common Use Cases 99 | 100 | ### 1. Document Q&A 101 | Best Implementation: Simple RAG 102 | ```go 103 | simpleRAG, _ := raggo.NewSimpleRAG( 104 | raggo.WithCollection("qa_docs"), 105 | raggo.WithTopK(1), 106 | ) 107 | ``` 108 | 109 | ### 2. Chatbot with Memory 110 | Best Implementation: Contextual RAG 111 | ```go 112 | contextualRAG, _ := raggo.NewContextualRAG( 113 | raggo.WithContextSize(10), 114 | raggo.WithMemory(true), 115 | ) 116 | ``` 117 | 118 | ### 3. Document Analysis 119 | Best Implementation: Core RAG 120 | ```go 121 | rag, _ := raggo.NewRAG( 122 | raggo.WithChunkSize(512), 123 | raggo.WithHybridSearch(true), 124 | ) 125 | ``` 126 | 127 | ## Example Implementations 128 | 129 | ### 1. Memory-Enhanced Chatbot 130 | A sophisticated chatbot implementation demonstrating advanced RAG capabilities: 131 | - Document-based knowledge integration 132 | - Context-aware responses 133 | - Memory retention across conversations 134 | - Interactive CLI interface 135 | 136 | [Detailed Documentation →](./chatbot_example.md) 137 | 138 | ```go 139 | // Initialize components 140 | memoryContext, err := raggo.NewMemoryContext(apiKey, 141 | raggo.MemoryCollection("tech_docs"), 142 | raggo.MemoryTopK(5), 143 | raggo.MemoryMinScore(0.01), 144 | raggo.MemoryStoreLastN(10), 145 | ) 146 | 147 | // Process documents 148 | for _, doc := range docs { 149 | content, err := os.ReadFile(doc) 150 | err = memoryContext.Store(ctx, filepath.Base(doc), string(content)) 151 | } 152 | 153 | // Interactive chat loop 154 | for { 155 | query := getUserInput() 156 | response, err := memoryContext.ProcessWithContext(ctx, query) 157 | fmt.Printf("\nResponse: %s\n", response) 158 | } 159 | ``` 160 | 161 | Key Features: 162 | - Vector database integration (Milvus) 163 | - OpenAI LLM integration 164 | - Document processing and storage 165 | - Context-aware query processing 166 | - Memory management 167 | - Error handling and logging 168 | 169 | [View Full Example →](../examples/memory_enhancer_example.go) 170 | 171 | ## Best Practices 172 | 173 | ### 1. Configuration 174 | - Use appropriate chunk sizes (default: 512) 175 | - Configure TopK based on use case 176 | - Set reasonable timeouts 177 | - Use environment variables for API keys 178 | 179 | ### 2. Performance 180 | - Implement batch processing 181 | - Monitor API usage 182 | - Use connection pooling 183 | - Cache frequent queries 184 | 185 | ### 3. Error Handling 186 | - Implement proper error handling 187 | - Use context for cancellation 188 | - Close resources properly 189 | - Monitor system resources 190 | 191 | ## Integration Examples 192 | 193 | ### 1. HTTP Server 194 | ```go 195 | func handleQuery(w http.ResponseWriter, r *http.Request) { 196 | rag := getRAGInstance() // Get appropriate RAG instance 197 | results, err := rag.Query(r.Context(), r.URL.Query().Get("q")) 198 | if err != nil { 199 | http.Error(w, err.Error(), http.StatusInternalServerError) 200 | return 201 | } 202 | json.NewEncoder(w).Encode(results) 203 | } 204 | ``` 205 | 206 | ### 2. CLI Application 207 | ```go 208 | func main() { 209 | rag := initializeRAG() // Initialize appropriate RAG 210 | scanner := bufio.NewScanner(os.Stdin) 211 | 212 | for { 213 | fmt.Print("Query: ") 214 | if !scanner.Scan() { 215 | break 216 | } 217 | results, _ := rag.Query(context.Background(), scanner.Text()) 218 | fmt.Printf("Results: %v\n", results) 219 | } 220 | } 221 | ``` 222 | 223 | ## Dependencies 224 | 225 | - Vector Database: Milvus (default) 226 | - Embedding Service: OpenAI (default) 227 | - Go version: 1.16+ 228 | 229 | ## Getting Started 230 | 231 | 1. **Installation** 232 | ```bash 233 | go get github.com/teilomillet/raggo 234 | ``` 235 | 236 | 2. **Basic Setup** 237 | ```go 238 | import "github.com/teilomillet/raggo" 239 | ``` 240 | 241 | 3. **Environment Variables** 242 | ```bash 243 | export OPENAI_API_KEY=your_api_key 244 | ``` 245 | 246 | ## Contributing 247 | 248 | See [CONTRIBUTING.md](../CONTRIBUTING.md) for guidelines on: 249 | - Code style 250 | - Testing requirements 251 | - Pull request process 252 | - Documentation standards 253 | 254 | ## Resources 255 | 256 | - [Examples Directory](../examples/) 257 | - [API Reference](./api.md) 258 | - [FAQ](./faq.md) 259 | - [Troubleshooting Guide](./troubleshooting.md) 260 | -------------------------------------------------------------------------------- /embedder.go: -------------------------------------------------------------------------------- 1 | // Package raggo provides a high-level interface for text embedding and retrieval 2 | // operations in RAG (Retrieval-Augmented Generation) systems. It simplifies the 3 | // process of converting text into vector embeddings using various providers. 4 | package raggo 5 | 6 | import ( 7 | "context" 8 | "fmt" 9 | 10 | "github.com/teilomillet/raggo/rag" 11 | "github.com/teilomillet/raggo/rag/providers" 12 | ) 13 | 14 | // EmbeddedChunk represents a chunk of text with its embeddings and metadata. 15 | // It serves as the core data structure for storing and retrieving embedded content 16 | // in the RAG system. 17 | // 18 | // Structure: 19 | // - Text: The original text content that was embedded 20 | // - Embeddings: Vector representations from different models/providers 21 | // - Metadata: Additional context and information about the chunk 22 | // 23 | // Example: 24 | // 25 | // chunk := EmbeddedChunk{ 26 | // Text: "Sample text content", 27 | // Embeddings: map[string][]float64{ 28 | // "default": []float64{0.1, 0.2, 0.3}, 29 | // }, 30 | // Metadata: map[string]interface{}{ 31 | // "source": "document1.txt", 32 | // "timestamp": time.Now(), 33 | // }, 34 | // } 35 | type EmbeddedChunk = rag.EmbeddedChunk 36 | 37 | // EmbedderOption is a function type for configuring the Embedder. 38 | // It follows the functional options pattern to provide a clean and 39 | // flexible configuration API. 40 | // 41 | // Common options include: 42 | // - SetEmbedderProvider: Choose the embedding service provider 43 | // - SetEmbedderModel: Select the specific embedding model 44 | // - SetEmbedderAPIKey: Configure authentication 45 | // - SetOption: Set custom provider-specific options 46 | type EmbedderOption = rag.EmbedderOption 47 | 48 | // SetEmbedderProvider sets the provider for the Embedder. 49 | // Supported providers include: 50 | // - "openai": OpenAI's text-embedding-ada-002 and other models 51 | // - "cohere": Cohere's embedding models 52 | // - "local": Local embedding models (if configured) 53 | // 54 | // Example: 55 | // 56 | // embedder, err := NewEmbedder( 57 | // SetEmbedderProvider("openai"), 58 | // SetEmbedderModel("text-embedding-ada-002"), 59 | // ) 60 | func SetEmbedderProvider(provider string) EmbedderOption { 61 | return rag.SetProvider(provider) 62 | } 63 | 64 | // SetEmbedderModel sets the specific model to use for embedding. 65 | // Available models depend on the chosen provider: 66 | // - OpenAI: "text-embedding-ada-002" (recommended) 67 | // - Cohere: "embed-multilingual-v2.0" 68 | // - Local: Depends on configured models 69 | // 70 | // Example: 71 | // 72 | // embedder, err := NewEmbedder( 73 | // SetEmbedderProvider("openai"), 74 | // SetEmbedderModel("text-embedding-ada-002"), 75 | // ) 76 | func SetEmbedderModel(model string) EmbedderOption { 77 | return rag.SetModel(model) 78 | } 79 | 80 | // SetEmbedderAPIKey sets the authentication key for the embedding service. 81 | // This is required for most cloud-based embedding providers. 82 | // 83 | // Security Note: Store API keys securely and never commit them to version control. 84 | // Consider using environment variables or secure key management systems. 85 | // 86 | // Example: 87 | // 88 | // embedder, err := NewEmbedder( 89 | // SetEmbedderProvider("openai"), 90 | // SetEmbedderAPIKey(os.Getenv("OPENAI_API_KEY")), 91 | // ) 92 | func SetEmbedderAPIKey(apiKey string) EmbedderOption { 93 | return rag.SetAPIKey(apiKey) 94 | } 95 | 96 | // SetOption sets a custom option for the Embedder. 97 | // This allows for provider-specific configuration that isn't covered 98 | // by the standard options. 99 | // 100 | // Example: 101 | // 102 | // embedder, err := NewEmbedder( 103 | // SetEmbedderProvider("openai"), 104 | // SetOption("timeout", 30*time.Second), 105 | // SetOption("max_retries", 3), 106 | // ) 107 | func SetOption(key string, value interface{}) EmbedderOption { 108 | return rag.SetOption(key, value) 109 | } 110 | 111 | // Embedder interface defines the contract for embedding implementations. 112 | // This allows for different embedding providers to be used interchangeably. 113 | type Embedder = providers.Embedder 114 | 115 | // NewEmbedder creates a new Embedder instance based on the provided options. 116 | // It handles provider selection and configuration, returning a ready-to-use 117 | // embedding interface. 118 | // 119 | // Returns an error if: 120 | // - No provider is specified 121 | // - The provider is not supported 122 | // - Configuration is invalid 123 | // - Authentication fails 124 | // 125 | // Example: 126 | // 127 | // embedder, err := NewEmbedder( 128 | // SetEmbedderProvider("openai"), 129 | // SetEmbedderModel("text-embedding-ada-002"), 130 | // SetEmbedderAPIKey(os.Getenv("OPENAI_API_KEY")), 131 | // ) 132 | // if err != nil { 133 | // log.Fatal(err) 134 | // } 135 | func NewEmbedder(opts ...EmbedderOption) (Embedder, error) { 136 | return rag.NewEmbedder(opts...) 137 | } 138 | 139 | // EmbeddingService handles the embedding process for text content. 140 | // It supports multiple embedders for different fields or purposes, 141 | // allowing for flexible embedding strategies. 142 | type EmbeddingService struct { 143 | embedders map[string]Embedder 144 | } 145 | 146 | // NewEmbeddingService creates a new embedding service with the specified embedder 147 | // as the default embedding provider. 148 | // 149 | // Example: 150 | // 151 | // embedder, _ := NewEmbedder(SetEmbedderProvider("openai")) 152 | // service := NewEmbeddingService(embedder) 153 | func NewEmbeddingService(embedder Embedder) *EmbeddingService { 154 | return &EmbeddingService{ 155 | embedders: map[string]Embedder{"default": embedder}, 156 | } 157 | } 158 | 159 | // EmbedChunks processes a slice of text chunks and generates embeddings for each one. 160 | // It supports multiple embedding fields per chunk, using different embedders 161 | // for each field if configured. 162 | // 163 | // The function: 164 | // 1. Processes each chunk through configured embedders 165 | // 2. Combines embeddings from all fields 166 | // 3. Preserves chunk metadata 167 | // 4. Handles errors for individual chunks 168 | // 169 | // Example: 170 | // 171 | // chunks := []rag.Chunk{ 172 | // {Text: "First chunk", TokenSize: 10}, 173 | // {Text: "Second chunk", TokenSize: 12}, 174 | // } 175 | // embedded, err := service.EmbedChunks(ctx, chunks) 176 | func (s *EmbeddingService) EmbedChunks(ctx context.Context, chunks []rag.Chunk) ([]rag.EmbeddedChunk, error) { 177 | embeddedChunks := make([]rag.EmbeddedChunk, 0, len(chunks)) 178 | for _, chunk := range chunks { 179 | embeddings := make(map[string][]float64) 180 | for field, embedder := range s.embedders { 181 | embedding, err := embedder.Embed(ctx, chunk.Text) 182 | if err != nil { 183 | return nil, fmt.Errorf("error embedding chunk for field %s: %w", field, err) 184 | } 185 | embeddings[field] = embedding 186 | } 187 | embeddedChunk := rag.EmbeddedChunk{ 188 | Text: chunk.Text, 189 | Embeddings: embeddings, 190 | Metadata: map[string]interface{}{ 191 | "token_size": chunk.TokenSize, 192 | "start_sentence": chunk.StartSentence, 193 | "end_sentence": chunk.EndSentence, 194 | }, 195 | } 196 | embeddedChunks = append(embeddedChunks, embeddedChunk) 197 | } 198 | return embeddedChunks, nil 199 | } 200 | 201 | // Embed generates embeddings for a single text string using the default embedder. 202 | // This is a convenience method for simple embedding operations. 203 | // 204 | // Example: 205 | // 206 | // text := "Sample text to embed" 207 | // embedding, err := service.Embed(ctx, text) 208 | // if err != nil { 209 | // log.Fatal(err) 210 | // } 211 | func (s *EmbeddingService) Embed(ctx context.Context, text string) ([]float64, error) { 212 | // Get the default embedder 213 | embedder, ok := s.embedders["default"] 214 | if !ok { 215 | return nil, fmt.Errorf("no default embedder found") 216 | } 217 | 218 | // Get embedding using the default embedder 219 | embedding, err := embedder.Embed(ctx, text) 220 | if err != nil { 221 | return nil, fmt.Errorf("error embedding text: %w", err) 222 | } 223 | 224 | return embedding, nil 225 | } 226 | -------------------------------------------------------------------------------- /examples/chat/docs/embeddings.txt: -------------------------------------------------------------------------------- 1 | Understanding Vector Embeddings in Machine Learning 2 | 3 | Vector embeddings are numerical representations of data in a high-dimensional space. They capture semantic relationships and enable efficient similarity comparisons. 4 | 5 | Types of Embeddings: 6 | 7 | 1. Text Embeddings 8 | - Word embeddings (Word2Vec, GloVe) 9 | - Sentence embeddings (USE, SBERT) 10 | - Document embeddings (OpenAI text-embedding-ada-002) 11 | 12 | 2. Image Embeddings 13 | - CNN feature vectors 14 | - CLIP embeddings 15 | - ResNet features 16 | 17 | Properties of Good Embeddings: 18 | 1. Similar items have similar vectors 19 | 2. Relationships are preserved 20 | 3. Meaningful vector operations 21 | 4. Consistent dimensionality 22 | 23 | Common Use Cases: 24 | - Semantic search 25 | - Document classification 26 | - Recommendation systems 27 | - Clustering similar items 28 | 29 | Popular Embedding Models: 30 | - OpenAI Embeddings 31 | - Cohere Embeddings 32 | - Hugging Face Models 33 | - Custom trained models 34 | 35 | Best Practices: 36 | 1. Choose appropriate dimension size 37 | 2. Normalize vectors when needed 38 | 3. Use domain-specific models 39 | 4. Regular model updates 40 | -------------------------------------------------------------------------------- /examples/chat/docs/golang_basics.txt: -------------------------------------------------------------------------------- 1 | Go (also known as Golang) is a statically typed, compiled programming language designed at Google. Here are some key features: 2 | 3 | 1. Simplicity and Readability 4 | Go emphasizes simplicity in its design. The language has a clean syntax and a small set of keywords, making it easy to learn and read. 5 | 6 | 2. Built-in Concurrency Support 7 | Go provides goroutines for concurrent execution and channels for communication between goroutines. A goroutine is a lightweight thread managed by the Go runtime. 8 | 9 | Example of concurrent programming in Go: 10 | ```go 11 | go func() { 12 | // This runs concurrently 13 | fmt.Println("Hello from goroutine!") 14 | }() 15 | ``` 16 | 17 | 3. Package System 18 | Go has a powerful package system for code organization. The standard library provides many essential packages for common programming tasks. 19 | 20 | 4. Error Handling 21 | Go uses explicit error handling instead of exceptions. Functions often return an error value that must be checked: 22 | ```go 23 | result, err := someFunction() 24 | if err != nil { 25 | // Handle error 26 | } 27 | ``` 28 | 29 | 5. Fast Compilation 30 | Go compiles very quickly compared to other compiled languages, making the development cycle more efficient. 31 | -------------------------------------------------------------------------------- /examples/chat/docs/microservices.txt: -------------------------------------------------------------------------------- 1 | Microservices Architecture Guide 2 | 3 | Microservices architecture is an approach to building applications as a collection of small, independent services. Each service runs in its own process and communicates through well-defined APIs. 4 | 5 | Key Principles: 6 | 7 | 1. Service Independence 8 | - Each service is independently deployable 9 | - Services can be written in different languages 10 | - Separate databases per service 11 | - Loose coupling between services 12 | 13 | 2. Communication Patterns 14 | - REST APIs 15 | - gRPC 16 | - Message queues (RabbitMQ, Kafka) 17 | - Event-driven architecture 18 | 19 | 3. Service Discovery 20 | - Service registry 21 | - Load balancing 22 | - Health checking 23 | - API gateway 24 | 25 | 4. Data Management 26 | - Database per service 27 | - Event sourcing 28 | - CQRS pattern 29 | - Data consistency patterns 30 | 31 | 5. Deployment 32 | - Containerization (Docker) 33 | - Orchestration (Kubernetes) 34 | - CI/CD pipelines 35 | - Blue-green deployment 36 | 37 | Real-World Example: The MountainPass Story 38 | 39 | In late 2022, a startup called MountainPass faced a critical challenge with their e-commerce platform. During Black Friday, their monolithic system crashed under heavy traffic, resulting in $2.8M in lost sales. Despite having 40 engineers and robust infrastructure, they couldn't handle the peak load. 40 | 41 | Their lead architect, Sarah Chen, proposed a groundbreaking solution called "PressureValve" - a dynamic microservices architecture that would automatically adapt to traffic patterns. The PressureValve system worked by: 42 | - Monitoring real-time traffic and system loads 43 | - Automatically scaling individual services based on demand 44 | - Managing inventory synchronization for their mountain gear rentals 45 | - Implementing smart circuit breakers at the service level 46 | 47 | Despite initial skepticism from the engineering team, Sarah's implementation proved successful. During the next Black Friday, with 20x normal traffic levels, the system handled 145% more orders than the previous year with zero downtime. 48 | 49 | Key Features of PressureValve: 50 | 1. Adaptive service mesh 51 | 2. Real-time load monitoring 52 | 3. Intelligent auto-scaling 53 | 4. Custom inventory synchronization 54 | 5. Automatic circuit breaking 55 | 56 | Common Challenges: 57 | - Distributed system complexity 58 | - Service boundary definition 59 | - Data consistency 60 | - Monitoring and tracing 61 | - Network reliability 62 | 63 | Best Practices: 64 | 1. Start small 65 | 2. Focus on business domains 66 | 3. Implement robust monitoring 67 | 4. Plan for failure 68 | 5. Automate everything possible 69 | -------------------------------------------------------------------------------- /examples/chat/docs/rag_systems.txt: -------------------------------------------------------------------------------- 1 | Understanding RAG (Retrieval Augmented Generation) Systems 2 | 3 | RAG systems combine information retrieval with language model generation to produce more accurate and contextual responses. Here's how they work: 4 | 5 | 1. Document Processing 6 | - Documents are split into chunks 7 | - Each chunk is converted into a vector embedding 8 | - Embeddings are stored in a vector database 9 | 10 | 2. Retrieval Process 11 | When a query is received: 12 | - Convert query to vector embedding 13 | - Search vector database for similar content 14 | - Retrieve most relevant chunks 15 | 16 | 3. Generation 17 | - Retrieved context is combined with the original query 18 | - Large Language Model (LLM) generates response using this context 19 | - Results are more factual and grounded in source documents 20 | 21 | 4. Benefits 22 | - Reduced hallucination in LLM responses 23 | - Up-to-date information from your document base 24 | - Source attribution for generated content 25 | - Lower costs through efficient prompting 26 | 27 | 5. Common Challenges 28 | - Choosing optimal chunk sizes 29 | - Balancing retrieval speed vs accuracy 30 | - Managing context window limitations 31 | - Handling document updates efficiently 32 | 33 | Best Practices: 34 | - Regular document updates 35 | - Careful chunk size selection 36 | - Quality embeddings generation 37 | - Proper prompt engineering 38 | -------------------------------------------------------------------------------- /examples/chat/docs/sample.txt: -------------------------------------------------------------------------------- 1 | MountainPass's PressureValve system is an innovative load balancing solution that helped the company handle unprecedented traffic during Black Friday. The system automatically scales resources based on incoming traffic patterns and distributes load across multiple servers. During Black Friday, when traffic spiked by 300%, PressureValve successfully maintained system stability and ensured zero downtime by dynamically allocating resources and routing requests efficiently. -------------------------------------------------------------------------------- /examples/chat/docs/vector_databases.txt: -------------------------------------------------------------------------------- 1 | Vector Databases: A Comprehensive Overview 2 | 3 | Vector databases are specialized database systems designed to store and query high-dimensional vectors efficiently. They are crucial for modern machine learning applications, especially in similarity search and recommendation systems. 4 | 5 | Key Concepts: 6 | 7 | 1. Vector Embeddings 8 | Vectors represent data points in a high-dimensional space. Text, images, or any other data can be converted into vector embeddings using machine learning models. 9 | 10 | 2. Similarity Search 11 | Vector databases excel at finding similar items using distance metrics like: 12 | - Euclidean distance (L2) 13 | - Cosine similarity 14 | - Inner product (IP) 15 | 16 | 3. Indexing Methods 17 | Common indexing techniques include: 18 | - HNSW (Hierarchical Navigable Small World) 19 | - IVF (Inverted File Index) 20 | - PQ (Product Quantization) 21 | 22 | These methods enable fast approximate nearest neighbor (ANN) search in high-dimensional spaces. 23 | 24 | 4. Popular Vector Database Systems 25 | - Milvus: Open-source vector database with high performance 26 | - Pinecone: Cloud-native vector database service 27 | - Weaviate: Vector search engine with GraphQL API 28 | - FAISS: Facebook AI Similarity Search library 29 | 30 | Use Cases: 31 | - Semantic search 32 | - Image similarity 33 | - Recommendation engines 34 | - Face recognition 35 | - Document deduplication 36 | -------------------------------------------------------------------------------- /examples/chat/main.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "context" 5 | "fmt" 6 | "log" 7 | "os" 8 | "path/filepath" 9 | "strings" 10 | "time" 11 | 12 | "github.com/teilomillet/gofh" 13 | "github.com/teilomillet/gollm" 14 | "github.com/teilomillet/raggo" 15 | ) 16 | 17 | var ( 18 | llm gollm.LLM 19 | retriever *raggo.Retriever 20 | ) 21 | 22 | func main() { 23 | // Initialize LLM 24 | var err error 25 | llm, err = gollm.NewLLM( 26 | gollm.SetProvider("openai"), 27 | gollm.SetModel("gpt-4o-mini"), 28 | gollm.SetAPIKey(os.Getenv("OPENAI_API_KEY")), 29 | ) 30 | if err != nil { 31 | log.Fatal(err) 32 | } 33 | 34 | // Get the docs directory 35 | wd, err := os.Getwd() 36 | if err != nil { 37 | log.Fatal(err) 38 | } 39 | // Fix the path construction - remove the duplicate "examples/chat" 40 | docsDir := filepath.Join(wd, "docs") 41 | 42 | // Print the path for debugging 43 | log.Printf("Loading documents from: %s", docsDir) 44 | 45 | // Verify directory exists 46 | if _, err := os.Stat(docsDir); os.IsNotExist(err) { 47 | log.Fatalf("Documents directory does not exist: %s", docsDir) 48 | } 49 | 50 | // Create context with timeout 51 | ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute) 52 | defer cancel() 53 | 54 | // First, let's clean up any existing collection 55 | vectorDB, err := raggo.NewVectorDB( 56 | raggo.WithType("milvus"), 57 | raggo.WithAddress("localhost:19530"), 58 | raggo.WithTimeout(5*time.Minute), 59 | ) 60 | if err != nil { 61 | log.Fatal(err) 62 | } 63 | 64 | // Connect to the database 65 | err = vectorDB.Connect(ctx) 66 | if err != nil { 67 | log.Fatal(err) 68 | } 69 | defer vectorDB.Close() 70 | 71 | // Check and drop existing collection 72 | exists, err := vectorDB.HasCollection(ctx, "chat_docs") 73 | if err != nil { 74 | log.Fatal(err) 75 | } 76 | if exists { 77 | log.Println("Dropping existing collection") 78 | err = vectorDB.DropCollection(ctx, "chat_docs") 79 | if err != nil { 80 | log.Fatal(err) 81 | } 82 | } 83 | 84 | // Register documents with explicit chunking and debug output 85 | log.Println("Registering documents with debug settings...") 86 | err = raggo.Register(ctx, docsDir, 87 | raggo.WithVectorDB("milvus", map[string]string{"address": "localhost:19530"}), 88 | raggo.WithCollection("chat_docs", true), 89 | raggo.WithChunking(200, 50), // Adjusted chunk size and overlap 90 | raggo.WithEmbedding( 91 | "openai", 92 | "text-embedding-3-small", 93 | os.Getenv("OPENAI_API_KEY"), 94 | ), 95 | ) 96 | if err != nil { 97 | log.Fatal(err) 98 | } 99 | 100 | // Create and load index 101 | log.Println("Creating and loading index...") 102 | err = vectorDB.CreateIndex(ctx, "chat_docs", "Embedding", raggo.Index{ 103 | Type: "HNSW", 104 | Metric: "L2", 105 | Parameters: map[string]interface{}{ 106 | "M": 16, 107 | "efConstruction": 256, 108 | }, 109 | }) 110 | if err != nil { 111 | log.Fatal(err) 112 | } 113 | 114 | // Load the collection 115 | err = vectorDB.LoadCollection(ctx, "chat_docs") 116 | if err != nil { 117 | log.Fatal(err) 118 | } 119 | 120 | // Initialize retriever with debug settings 121 | log.Println("Initializing retriever with debug settings...") 122 | retriever, err = raggo.NewRetriever( 123 | raggo.WithRetrieveDB("milvus", "localhost:19530"), 124 | raggo.WithRetrieveCollection("chat_docs"), 125 | raggo.WithTopK(10), // More results 126 | raggo.WithMinScore(0.1), // Lower threshold 127 | raggo.WithHybrid(false), // Simple search first 128 | raggo.WithRetrieveEmbedding( 129 | "openai", 130 | "text-embedding-3-small", 131 | os.Getenv("OPENAI_API_KEY"), 132 | ), 133 | ) 134 | if err != nil { 135 | log.Fatal(err) 136 | } 137 | defer retriever.Close() 138 | 139 | // Create web app 140 | app := gofh.New() 141 | 142 | // Main page 143 | app.Get("/").Handle(func(c *gofh.Context) gofh.Element { 144 | return gofh.Div( 145 | gofh.El("style", ` 146 | body { max-width: 800px; margin: 0 auto; padding: 20px; font-family: system-ui; } 147 | #chat { margin-bottom: 20px; } 148 | .message { padding: 10px; margin: 5px 0; border-radius: 5px; } 149 | .user { background: #e3f2fd; } 150 | .ai { background: #f5f5f5; } 151 | form { display: flex; gap: 10px; } 152 | input { flex: 1; padding: 8px; border: 1px solid #ddd; border-radius: 4px; } 153 | button { padding: 8px 16px; background: #1976d2; color: white; border: none; border-radius: 4px; cursor: pointer; } 154 | .sources { font-size: 0.8em; color: #666; margin-top: 5px; } 155 | `), 156 | gofh.H1("RAG Chat Demo"), 157 | gofh.Div().ID("chat"), 158 | gofh.Form( 159 | gofh.Input("text", "msg"). 160 | Attr("placeholder", "Type your question..."). 161 | Attr("class", "w-full"), 162 | gofh.Button("Send"), 163 | ). 164 | Attr("hx-post", "/chat"). 165 | Attr("hx-target", "#chat"). 166 | Attr("hx-swap", "beforeend"), 167 | ) 168 | }) 169 | 170 | // Chat endpoint 171 | app.Post("/chat").Handle(func(c *gofh.Context) gofh.Element { 172 | msg := c.GetFormValue("msg") 173 | 174 | log.Printf("\n=== New Search Query ===") 175 | log.Printf("Query: %q", msg) 176 | 177 | results, err := retriever.Retrieve(c.Request.Context(), msg) 178 | if err != nil { 179 | log.Printf("❌ Retrieval error: %v", err) 180 | } 181 | 182 | log.Printf("\n=== Search Results ===") 183 | log.Printf("Found %d results", len(results)) 184 | 185 | var contexts []string 186 | var sources []string 187 | 188 | for i, result := range results { 189 | log.Printf("\nResult %d:", i+1) 190 | log.Printf("Score: %.4f", result.Score) 191 | log.Printf("Source: %s", result.Source) 192 | log.Printf("Content length: %d", len(result.Content)) 193 | log.Printf("Content: %s", truncateString(result.Content, 200)) 194 | if result.Metadata != nil { 195 | log.Printf("Metadata: %+v", result.Metadata) 196 | } 197 | 198 | contexts = append(contexts, result.Content) 199 | if result.Source != "" { 200 | shortPath := filepath.Base(result.Source) 201 | sources = append(sources, fmt.Sprintf("%s (%.2f)", shortPath, result.Score)) 202 | } 203 | } 204 | 205 | // Generate response with improved prompt 206 | prompt := fmt.Sprintf(`Here are some relevant sections from our documentation: 207 | 208 | %s 209 | 210 | Based on this information, please answer the following question: %s 211 | 212 | If the information isn't found in the provided context, please say so clearly.`, 213 | strings.Join(contexts, "\n\n---\n\n"), 214 | msg, 215 | ) 216 | 217 | resp, err := llm.Generate(c.Request.Context(), gollm.NewPrompt(prompt)) 218 | if err != nil { 219 | resp = "Error: " + err.Error() 220 | } 221 | 222 | userMsg := gofh.Div( 223 | gofh.P("You: "+msg), 224 | ).Attr("class", "message user") 225 | 226 | aiMsg := gofh.Div( 227 | gofh.P("AI: "+resp), 228 | gofh.P("Sources: "+strings.Join(sources, ", ")).Attr("class", "sources"), 229 | ).Attr("class", "message ai") 230 | 231 | return gofh.Div(userMsg, aiMsg) 232 | }) 233 | 234 | log.Println("Chat server starting on http://localhost:8080") 235 | log.Fatal(app.Serve()) 236 | } 237 | 238 | func truncateString(s string, n int) string { 239 | if len(s) <= n { 240 | return s 241 | } 242 | return s[:n] + "..." 243 | } 244 | -------------------------------------------------------------------------------- /examples/chat/v2.go: -------------------------------------------------------------------------------- 1 | // examples/chat/v2.go 2 | package main 3 | 4 | import ( 5 | "context" 6 | "fmt" 7 | "log" 8 | "os" 9 | "path/filepath" 10 | "strings" 11 | 12 | "github.com/teilomillet/gofh" 13 | "github.com/teilomillet/gollm" 14 | "github.com/teilomillet/raggo" 15 | ) 16 | 17 | type ChatApp struct { 18 | rag *raggo.RAG 19 | llm gollm.LLM 20 | } 21 | 22 | func NewChatApp() (*ChatApp, error) { 23 | // Initialize LLM 24 | llm, err := gollm.NewLLM( 25 | gollm.SetProvider("openai"), 26 | gollm.SetModel("gpt-4o-mini"), 27 | gollm.SetAPIKey(os.Getenv("OPENAI_API_KEY")), 28 | ) 29 | if err != nil { 30 | return nil, fmt.Errorf("failed to initialize LLM: %w", err) 31 | } 32 | 33 | // Initialize RAG with OpenAI and local Milvus 34 | rag, err := raggo.NewRAG( 35 | raggo.WithOpenAI(os.Getenv("OPENAI_API_KEY")), 36 | raggo.WithMilvus("chat_docs"), 37 | // Optional: Add custom settings 38 | func(c *raggo.RAGConfig) { 39 | c.ChunkSize = 200 40 | c.ChunkOverlap = 50 41 | c.TopK = 10 42 | c.MinScore = 0.1 43 | c.UseHybrid = false 44 | c.SearchParams = map[string]interface{}{ // Add search params 45 | "nprobe": 10, 46 | "ef": 64, 47 | "type": "HNSW", 48 | } 49 | }, 50 | ) 51 | if err != nil { 52 | return nil, fmt.Errorf("failed to initialize RAG: %w", err) 53 | } 54 | 55 | return &ChatApp{ 56 | rag: rag, 57 | llm: llm, 58 | }, nil 59 | } 60 | 61 | func (app *ChatApp) LoadDocuments(ctx context.Context, docsDir string) error { 62 | log.Printf("Loading documents from: %s", docsDir) 63 | return app.rag.LoadDocuments(ctx, docsDir) 64 | } 65 | 66 | func (app *ChatApp) Close() error { 67 | return app.rag.Close() 68 | } 69 | 70 | func main() { 71 | // Create new chat app 72 | app, err := NewChatApp() 73 | if err != nil { 74 | log.Fatal(err) 75 | } 76 | defer app.Close() 77 | 78 | // Get the current working directory 79 | wd, err := os.Getwd() 80 | if err != nil { 81 | log.Fatal(err) 82 | } 83 | 84 | // Construct the correct path to docs directory 85 | // Changed: Remove duplicate "examples/chat" from path 86 | docsDir := filepath.Join(wd, "docs") 87 | 88 | // Debug: Print the path we're trying to use 89 | log.Printf("Looking for docs in: %s", docsDir) 90 | 91 | // Verify directory exists 92 | if _, err := os.Stat(docsDir); os.IsNotExist(err) { 93 | log.Fatalf("Documents directory does not exist: %s", docsDir) 94 | } 95 | 96 | // Load documents 97 | err = app.LoadDocuments(context.Background(), docsDir) 98 | if err != nil { 99 | log.Fatal(err) 100 | } 101 | 102 | // Create web app 103 | webApp := gofh.New() 104 | 105 | // Main page 106 | webApp.Get("/").Handle(func(c *gofh.Context) gofh.Element { 107 | return gofh.Div( 108 | gofh.El("style", ` 109 | body { max-width: 800px; margin: 0 auto; padding: 20px; font-family: system-ui; } 110 | #chat { margin-bottom: 20px; } 111 | .message { padding: 10px; margin: 5px 0; border-radius: 5px; } 112 | .user { background: #e3f2fd; } 113 | .ai { background: #f5f5f5; } 114 | form { display: flex; gap: 10px; } 115 | input { flex: 1; padding: 8px; border: 1px solid #ddd; border-radius: 4px; } 116 | button { padding: 8px 16px; background: #1976d2; color: white; border: none; border-radius: 4px; cursor: pointer; } 117 | .sources { font-size: 0.8em; color: #666; margin-top: 5px; } 118 | `), 119 | gofh.H1("RAG Chat Demo"), 120 | gofh.Div().ID("chat"), 121 | gofh.Form( 122 | gofh.Input("text", "msg"). 123 | Attr("placeholder", "Type your question..."). 124 | Attr("class", "w-full"), 125 | gofh.Button("Send"), 126 | ). 127 | Attr("hx-post", "/chat"). 128 | Attr("hx-target", "#chat"). 129 | Attr("hx-swap", "beforeend"), 130 | ) 131 | }) 132 | 133 | // Chat endpoint 134 | webApp.Post("/chat").Handle(func(c *gofh.Context) gofh.Element { 135 | msg := c.GetFormValue("msg") 136 | if msg == "" { 137 | return renderError(fmt.Errorf("empty message")) 138 | } 139 | 140 | log.Printf("\n=== New Search Query ===\nQuery: %q", msg) 141 | 142 | // Get relevant documents 143 | results, err := app.rag.Query(c.Request.Context(), msg) 144 | if err != nil { 145 | log.Printf("❌ Retrieval error: %v", err) 146 | return renderError(err) 147 | } 148 | 149 | log.Printf("Found %d results", len(results)) 150 | 151 | // Prepare context and sources 152 | var contexts []string 153 | var sources []string 154 | for _, result := range results { 155 | contexts = append(contexts, result.Content) 156 | shortPath := filepath.Base(result.Source) 157 | sources = append(sources, fmt.Sprintf("%s (%.2f)", shortPath, result.Score)) 158 | } 159 | 160 | // Generate response 161 | prompt := fmt.Sprintf(`Here are some relevant sections from our documentation: 162 | 163 | %s 164 | 165 | Based on this information, please answer the following question: %s 166 | 167 | If the information isn't found in the provided context, please say so clearly.`, 168 | strings.Join(contexts, "\n\n---\n\n"), 169 | msg, 170 | ) 171 | 172 | resp, err := app.llm.Generate(c.Request.Context(), gollm.NewPrompt(prompt)) 173 | if err != nil { 174 | return renderError(err) 175 | } 176 | 177 | return renderChat(msg, resp, sources) 178 | }) 179 | 180 | log.Println("Chat server starting on http://localhost:8080") 181 | log.Fatal(webApp.Serve()) 182 | } 183 | 184 | // UI helpers 185 | const defaultStyles = ` 186 | body { max-width: 800px; margin: 0 auto; padding: 20px; font-family: system-ui; } 187 | #chat { margin-bottom: 20px; } 188 | .message { padding: 10px; margin: 5px 0; border-radius: 5px; } 189 | .user { background: #e3f2fd; } 190 | .ai { background: #f5f5f5; } 191 | .error { background: #ffebee; } 192 | form { display: flex; gap: 10px; } 193 | input { flex: 1; padding: 8px; border: 1px solid #ddd; border-radius: 4px; } 194 | button { padding: 8px 16px; background: #1976d2; color: white; border: none; border-radius: 4px; cursor: pointer; } 195 | .sources { font-size: 0.8em; color: #666; margin-top: 5px; } 196 | ` 197 | 198 | func renderChat(query, response string, sources []string) gofh.Element { 199 | return gofh.Div( 200 | gofh.Div( 201 | gofh.P("You: "+query), 202 | ).Attr("class", "message user"), 203 | gofh.Div( 204 | gofh.P("AI: "+response), 205 | gofh.P("Sources: "+strings.Join(sources, ", ")).Attr("class", "sources"), 206 | ).Attr("class", "message ai"), 207 | ) 208 | } 209 | 210 | func renderError(err error) gofh.Element { 211 | return gofh.Div( 212 | gofh.P("Error: "+err.Error()), 213 | ).Attr("class", "message error") 214 | } 215 | -------------------------------------------------------------------------------- /examples/chat/v3.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "context" 5 | "fmt" 6 | "log" 7 | "os" 8 | "path/filepath" 9 | "strings" 10 | 11 | "github.com/teilomillet/gofh" 12 | "github.com/teilomillet/gollm" 13 | "github.com/teilomillet/raggo" 14 | ) 15 | 16 | type ChatApp struct { 17 | rag *raggo.RAG 18 | llm gollm.LLM 19 | } 20 | 21 | func NewChatApp() (*ChatApp, error) { 22 | // Initialize LLM 23 | llm, err := gollm.NewLLM( 24 | gollm.SetProvider("openai"), 25 | gollm.SetModel("gpt-4o-mini"), 26 | gollm.SetAPIKey(os.Getenv("OPENAI_API_KEY")), 27 | ) 28 | if err != nil { 29 | return nil, fmt.Errorf("failed to initialize LLM: %w", err) 30 | } 31 | 32 | // Initialize RAG with context processing 33 | rag, err := raggo.NewRAG( 34 | raggo.WithOpenAI(os.Getenv("OPENAI_API_KEY")), 35 | raggo.WithMilvus("chat_docs"), 36 | func(c *raggo.RAGConfig) { 37 | c.ChunkSize = 100 // Increased from 200 38 | c.ChunkOverlap = 50 // Increased from 50 39 | c.BatchSize = 10 40 | c.TopK = 5 // Increased from 5 41 | c.MinScore = 0.1 // Decreased from 0.1 42 | }, 43 | ) 44 | if err != nil { 45 | return nil, fmt.Errorf("failed to initialize RAG: %w", err) 46 | } 47 | 48 | return &ChatApp{ 49 | rag: rag, 50 | llm: llm, 51 | }, nil 52 | } 53 | 54 | func (app *ChatApp) LoadDocuments(ctx context.Context, docsDir string) error { 55 | log.Printf("Loading documents from: %s", docsDir) 56 | return app.rag.ProcessWithContext(ctx, docsDir, "gpt-4o-mini") 57 | } 58 | 59 | func (app *ChatApp) Close() error { 60 | return app.rag.Close() 61 | } 62 | 63 | func main() { 64 | // Enable debug logging 65 | raggo.SetLogLevel(raggo.LogLevelDebug) 66 | 67 | // Create new chat app 68 | app, err := NewChatApp() 69 | if err != nil { 70 | log.Fatal(err) 71 | } 72 | defer app.Close() 73 | 74 | // Get the current working directory 75 | wd, err := os.Getwd() 76 | if err != nil { 77 | log.Fatal(err) 78 | } 79 | 80 | // Construct the path to docs directory 81 | docsDir := filepath.Join(wd, "docs") 82 | 83 | // Debug: Print the path we're trying to use 84 | log.Printf("Looking for docs in: %s", docsDir) 85 | 86 | // Verify directory exists 87 | if _, err := os.Stat(docsDir); os.IsNotExist(err) { 88 | log.Fatalf("Documents directory does not exist: %s", docsDir) 89 | } 90 | 91 | // Process each file in the directory 92 | files, err := os.ReadDir(docsDir) 93 | if err != nil { 94 | log.Fatalf("Failed to read directory: %v", err) 95 | } 96 | 97 | for _, file := range files { 98 | if file.IsDir() { 99 | continue 100 | } 101 | filePath := filepath.Join(docsDir, file.Name()) 102 | log.Printf("Processing file: %s", filePath) 103 | err = app.LoadDocuments(context.Background(), filePath) 104 | if err != nil { 105 | log.Printf("Error processing file %s: %v", filePath, err) 106 | // Continue with the next file instead of stopping 107 | continue 108 | } 109 | } 110 | 111 | // Create web app 112 | webApp := gofh.New() 113 | 114 | // Main page 115 | webApp.Get("/").Handle(func(c *gofh.Context) gofh.Element { 116 | return gofh.Div( 117 | gofh.El("style", ` 118 | body { max-width: 800px; margin: 0 auto; padding: 20px; font-family: system-ui; } 119 | #chat { margin-bottom: 20px; } 120 | .message { padding: 10px; margin: 5px 0; border-radius: 5px; } 121 | .user { background: #e3f2fd; } 122 | .ai { background: #f5f5f5; } 123 | form { display: flex; gap: 10px; } 124 | input { flex: 1; padding: 8px; border: 1px solid #ddd; border-radius: 4px; } 125 | button { padding: 8px 16px; background: #1976d2; color: white; border: none; border-radius: 4px; cursor: pointer; } 126 | .sources { font-size: 0.8em; color: #666; margin-top: 5px; } 127 | `), 128 | gofh.H1("Context-Enhanced RAG Chat Demo"), 129 | gofh.Div().ID("chat"), 130 | gofh.Form( 131 | gofh.Input("text", "msg"). 132 | Attr("placeholder", "Type your question..."). 133 | Attr("class", "w-full"), 134 | gofh.Button("Send"), 135 | ). 136 | Attr("hx-post", "/chat"). 137 | Attr("hx-target", "#chat"). 138 | Attr("hx-swap", "beforeend"), 139 | ) 140 | }) 141 | 142 | // Chat endpoint 143 | webApp.Post("/chat").Handle(func(c *gofh.Context) gofh.Element { 144 | msg := c.GetFormValue("msg") 145 | 146 | log.Printf("\n=== New Search Query ===") 147 | log.Printf("Query: %q", msg) 148 | 149 | results, err := app.rag.Query(c.Request.Context(), msg) 150 | if err != nil { 151 | log.Printf("❌ Retrieval error: %v", err) 152 | return renderError(err) 153 | } 154 | 155 | log.Printf("\n=== Search Results ===") 156 | log.Printf("Found %d results", len(results)) 157 | 158 | var contexts []string 159 | var sources []string 160 | 161 | for i, result := range results { 162 | log.Printf("\nResult %d:", i+1) 163 | log.Printf("Score: %.4f", result.Score) 164 | log.Printf("Source: %s", result.Source) 165 | log.Printf("Content length: %d", len(result.Content)) 166 | log.Printf("Content: %s", truncateString(result.Content, 500)) // Increased from 200 167 | if result.Metadata != nil { 168 | log.Printf("Metadata: %+v", result.Metadata) 169 | } 170 | 171 | contexts = append(contexts, result.Content) 172 | if result.Source != "" { 173 | shortPath := filepath.Base(result.Source) 174 | sources = append(sources, fmt.Sprintf("%s (%.2f)", shortPath, result.Score)) 175 | } 176 | } 177 | 178 | // Generate response with improved prompt 179 | prompt := fmt.Sprintf(`Here are some relevant sections from our documentation: 180 | 181 | %s 182 | 183 | Based on this information, please answer the following question: %s 184 | 185 | If the information isn't found in the provided context, please say so clearly.`, 186 | strings.Join(contexts, "\n\n---\n\n"), 187 | msg, 188 | ) 189 | 190 | resp, err := app.llm.Generate(c.Request.Context(), gollm.NewPrompt(prompt)) 191 | if err != nil { 192 | return renderError(err) 193 | } 194 | 195 | return renderChat(msg, resp, sources) 196 | }) 197 | 198 | log.Println("Context-enhanced chat server starting on http://localhost:8080") 199 | log.Fatal(webApp.Serve()) 200 | } 201 | 202 | // UI helpers 203 | func renderChat(query, response string, sources []string) gofh.Element { 204 | return gofh.Div( 205 | gofh.Div( 206 | gofh.P("You: "+query), 207 | ).Attr("class", "message user"), 208 | gofh.Div( 209 | gofh.P("AI: "+response), 210 | gofh.P("Sources: "+strings.Join(sources, ", ")).Attr("class", "sources"), 211 | ).Attr("class", "message ai"), 212 | ) 213 | } 214 | 215 | func renderError(err error) gofh.Element { 216 | return gofh.Div( 217 | gofh.P("Error: "+err.Error()), 218 | ).Attr("class", "message error") 219 | } 220 | 221 | func truncateString(s string, maxLen int) string { 222 | if len(s) <= maxLen { 223 | return s 224 | } 225 | return s[:maxLen] + "..." 226 | } 227 | -------------------------------------------------------------------------------- /examples/chromem/README.md: -------------------------------------------------------------------------------- 1 | # Chromem Example 2 | 3 | This example demonstrates how to use the Chromem vector database with Raggo's SimpleRAG interface. 4 | 5 | ## Prerequisites 6 | 7 | 1. Go 1.16 or later 8 | 2. OpenAI API key (set as environment variable `OPENAI_API_KEY`) 9 | 10 | ## Running the Example 11 | 12 | 1. Set your OpenAI API key: 13 | ```bash 14 | export OPENAI_API_KEY='your-api-key' 15 | ``` 16 | 17 | 2. Run the example: 18 | ```bash 19 | go run main.go 20 | ``` 21 | 22 | ## What it Does 23 | 24 | 1. Creates a new SimpleRAG instance with Chromem as the vector database 25 | 2. Creates sample documents about natural phenomena 26 | 3. Adds the documents to the database 27 | 4. Performs a semantic search using the query "Why is the sky blue?" 28 | 5. Prints the response based on the relevant documents found 29 | 30 | ## Expected Output 31 | 32 | ``` 33 | Question: Why is the sky blue? 34 | 35 | Answer: The sky appears blue because of a phenomenon called Rayleigh scattering. When sunlight travels through Earth's atmosphere, it collides with gas molecules. These molecules scatter blue wavelengths of light more strongly than red wavelengths, which is why we see the sky as blue. 36 | ``` 37 | 38 | ## Configuration 39 | 40 | The example uses the following configuration: 41 | - Vector Database: Chromem (persistent mode) 42 | - Collection Name: knowledge-base 43 | - Embedding Model: text-embedding-3-small 44 | - Chunk Size: 200 characters 45 | - Chunk Overlap: 50 characters 46 | - Top K Results: 1 47 | - Minimum Score: 0.1 48 | 49 | ## Notes 50 | 51 | - The database is stored in `./data/chromem.db` 52 | - Sample documents are created in the `./data` directory 53 | - The example uses persistent storage mode for Chromem 54 | -------------------------------------------------------------------------------- /examples/chromem/main.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "context" 5 | "fmt" 6 | "os" 7 | "path/filepath" 8 | 9 | "github.com/teilomillet/raggo" 10 | ) 11 | 12 | func main() { 13 | // Enable debug logging 14 | raggo.SetLogLevel(raggo.LogLevelDebug) 15 | 16 | // Create a temporary directory for our documents 17 | tmpDir := "./data" 18 | err := os.MkdirAll(tmpDir, 0755) 19 | if err != nil { 20 | fmt.Printf("Error creating temp directory: %v\n", err) 21 | os.Exit(1) 22 | } 23 | 24 | // Create sample documents 25 | docs := map[string]string{ 26 | "sky.txt": "The sky is blue because of Rayleigh scattering.", 27 | "leaves.txt": "Leaves are green because chlorophyll absorbs red and blue light.", 28 | } 29 | 30 | for filename, content := range docs { 31 | err := os.WriteFile(filepath.Join(tmpDir, filename), []byte(content), 0644) 32 | if err != nil { 33 | fmt.Printf("Error writing file %s: %v\n", filename, err) 34 | os.Exit(1) 35 | } 36 | } 37 | 38 | // Initialize RAG with Chromem 39 | config := raggo.SimpleRAGConfig{ 40 | Collection: "knowledge-base", 41 | DBType: "chromem", 42 | DBAddress: "./data/chromem.db", 43 | Model: "text-embedding-3-small", // OpenAI embedding model 44 | APIKey: os.Getenv("OPENAI_API_KEY"), 45 | Dimension: 1536, // text-embedding-3-small dimension 46 | // TopK is determined dynamically by the number of documents 47 | } 48 | 49 | raggo.Debug("Creating SimpleRAG with config", "config", config) 50 | 51 | rag, err := raggo.NewSimpleRAG(config) 52 | if err != nil { 53 | fmt.Printf("Error creating SimpleRAG: %v\n", err) 54 | os.Exit(1) 55 | } 56 | defer rag.Close() 57 | 58 | ctx := context.Background() 59 | 60 | // Add documents from the directory 61 | raggo.Debug("Adding documents from directory", "dir", tmpDir) 62 | err = rag.AddDocuments(ctx, tmpDir) 63 | if err != nil { 64 | fmt.Printf("Error adding documents: %v\n", err) 65 | os.Exit(1) 66 | } 67 | 68 | // Search for documents 69 | raggo.Debug("Searching for documents", "query", "Why is the sky blue?") 70 | response, err := rag.Search(ctx, "Why is the sky blue?") 71 | if err != nil { 72 | fmt.Printf("Error searching: %v\n", err) 73 | os.Exit(1) 74 | } 75 | 76 | fmt.Printf("Response: %s\n", response) 77 | } 78 | -------------------------------------------------------------------------------- /examples/chunker_examples.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "fmt" 5 | "log" 6 | "os" 7 | "path/filepath" 8 | 9 | "github.com/teilomillet/raggo" 10 | ) 11 | 12 | func main() { 13 | // Set log level to Info 14 | raggo.SetLogLevel(raggo.LogLevelInfo) 15 | 16 | // Create a new Parser 17 | parser := raggo.NewParser() 18 | 19 | // Create a new Chunker 20 | chunker, err := raggo.NewChunker( 21 | raggo.ChunkSize(100), // Set chunk size to 100 tokens 22 | raggo.ChunkOverlap(20), // Set chunk overlap to 20 tokens 23 | raggo.WithSentenceSplitter(raggo.SmartSentenceSplitter()), 24 | raggo.WithTokenCounter(raggo.NewDefaultTokenCounter()), 25 | ) 26 | if err != nil { 27 | log.Fatalf("Failed to create chunker: %v", err) 28 | } 29 | 30 | // Get the path to a PDF file in the testdata directory 31 | wd, err := os.Getwd() 32 | if err != nil { 33 | log.Fatalf("Failed to get working directory: %v", err) 34 | } 35 | pdfPath := filepath.Join(wd, "testdata", "CV.pdf") // Make sure this file exists 36 | 37 | // Parse the PDF file 38 | doc, err := parser.Parse(pdfPath) 39 | if err != nil { 40 | log.Fatalf("Failed to parse PDF: %v", err) 41 | } 42 | 43 | fmt.Printf("Parsed document with %d characters\n", len(doc.Content)) 44 | 45 | // Chunk the parsed content 46 | chunks := chunker.Chunk(doc.Content) 47 | 48 | fmt.Printf("Created %d chunks from the document\n", len(chunks)) 49 | 50 | // Print details of the first few chunks 51 | for i, chunk := range chunks { 52 | if i >= 5 { 53 | break // Only print the first 5 chunks 54 | } 55 | fmt.Printf("Chunk %d:\n", i+1) 56 | fmt.Printf(" Token Size: %d\n", chunk.TokenSize) 57 | fmt.Printf(" Start Sentence: %d\n", chunk.StartSentence) 58 | fmt.Printf(" End Sentence: %d\n", chunk.EndSentence) 59 | fmt.Printf(" Preview: %s\n", truncateString(chunk.Text, 100)) 60 | fmt.Println() 61 | } 62 | } 63 | 64 | // truncateString truncates a string to a specified length, adding an ellipsis if truncated 65 | func truncateString(s string, length int) string { 66 | if len(s) <= length { 67 | return s 68 | } 69 | return s[:length-3] + "..." 70 | } 71 | -------------------------------------------------------------------------------- /examples/concurrent_loader_example.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "context" 5 | "fmt" 6 | "log" 7 | "os" 8 | "path/filepath" 9 | "time" 10 | 11 | "github.com/teilomillet/raggo" 12 | ) 13 | 14 | func main() { 15 | // Set log level to Warn to reduce output noise 16 | raggo.SetLogLevel(raggo.LogLevelWarn) 17 | 18 | // Create a new ConcurrentPDFLoader with custom options 19 | loader := raggo.NewConcurrentPDFLoader( 20 | raggo.SetLoaderTimeout(1*time.Minute), 21 | raggo.SetTempDir(os.TempDir()), 22 | ) 23 | 24 | // Get the current working directory 25 | wd, err := os.Getwd() 26 | if err != nil { 27 | log.Fatalf("Failed to get working directory: %v", err) 28 | } 29 | 30 | // Specify the source directory containing PDF files 31 | sourceDir := filepath.Join(wd, "testdata") 32 | 33 | // Create a temporary directory for duplicated PDFs 34 | targetDir, err := os.MkdirTemp("", "raggo-pdf-test") 35 | if err != nil { 36 | log.Fatalf("Failed to create temp directory: %v", err) 37 | } 38 | defer os.RemoveAll(targetDir) 39 | 40 | // Use the concurrent loader to load 10 PDF files (some may be duplicates) 41 | desiredCount := 1000 42 | start := time.Now() 43 | loadedFiles, err := loader.LoadPDFsConcurrent(context.Background(), sourceDir, targetDir, desiredCount) 44 | if err != nil { 45 | log.Fatalf("Error loading PDF files concurrently: %v", err) 46 | } 47 | elapsed := time.Since(start) 48 | 49 | fmt.Printf("Loaded %d PDF files concurrently in %v\n", len(loadedFiles), elapsed) 50 | for i, file := range loadedFiles { 51 | fmt.Printf("%d. %s\n", i+1, file) 52 | } 53 | } 54 | -------------------------------------------------------------------------------- /examples/contextual/custom_llm.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "context" 5 | "fmt" 6 | "os" 7 | "path/filepath" 8 | "time" 9 | 10 | "github.com/teilomillet/gollm" 11 | "github.com/teilomillet/raggo" 12 | ) 13 | 14 | func main() { 15 | // Create a custom LLM instance with specific configuration 16 | llm, err := gollm.NewLLM( 17 | gollm.SetProvider("openai"), 18 | gollm.SetModel("gpt-4o-mini"), 19 | gollm.SetAPIKey(os.Getenv("OPENAI_API_KEY")), 20 | gollm.SetMaxTokens(500), 21 | gollm.SetMaxRetries(5), 22 | gollm.SetRetryDelay(time.Second*3), 23 | gollm.SetLogLevel(gollm.LogLevelInfo), 24 | ) 25 | if err != nil { 26 | fmt.Printf("Failed to create LLM: %v\n", err) 27 | os.Exit(1) 28 | } 29 | 30 | // Create RAG with custom LLM 31 | config := &raggo.ContextualRAGConfig{ 32 | Collection: "my_contextual_docs", 33 | LLM: llm, // Use our custom LLM instance 34 | } 35 | 36 | rag, err := raggo.NewContextualRAG(config) 37 | if err != nil { 38 | fmt.Printf("Failed to initialize RAG: %v\n", err) 39 | os.Exit(1) 40 | } 41 | defer rag.Close() 42 | 43 | // Add documents 44 | docsPath := filepath.Join(os.Getenv("HOME"), "Desktop", "raggo", "examples", "chat", "docs") 45 | if err := rag.AddDocuments(context.Background(), docsPath); err != nil { 46 | fmt.Printf("Failed to add documents: %v\n", err) 47 | os.Exit(1) 48 | } 49 | 50 | // Wait a moment for the collection to be fully loaded 51 | time.Sleep(time.Second * 2) 52 | 53 | // Example query 54 | query := "What is MountainPass's PressureValve system and how did it help during Black Friday?" 55 | response, err := rag.Search(context.Background(), query) 56 | if err != nil { 57 | fmt.Printf("Failed to search: %v\n", err) 58 | os.Exit(1) 59 | } 60 | 61 | fmt.Println(response) 62 | } 63 | -------------------------------------------------------------------------------- /examples/contextual/main.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "context" 5 | "fmt" 6 | "os" 7 | "path/filepath" 8 | 9 | "github.com/teilomillet/raggo" 10 | ) 11 | 12 | func main() { 13 | // Example 1: Basic Contextual RAG Usage 14 | // This demonstrates the simplest way to use contextual RAG with default settings 15 | basicExample() 16 | 17 | // Example 2: Advanced Configuration 18 | // This shows how to customize the contextual RAG behavior 19 | advancedExample() 20 | } 21 | 22 | func basicExample() { 23 | fmt.Println("\n=== Basic Contextual RAG Example ===") 24 | 25 | // Initialize RAG with default settings 26 | // This automatically: 27 | // - Uses OpenAI's text-embedding-3-small model for embeddings 28 | // - Sets optimal chunk size and overlap for context preservation 29 | // - Configures reasonable TopK and similarity thresholds 30 | rag, err := raggo.NewDefaultContextualRAG("basic_contextual_docs") 31 | if err != nil { 32 | fmt.Printf("Failed to initialize RAG: %v\n", err) 33 | os.Exit(1) 34 | } 35 | defer rag.Close() 36 | 37 | // Add documents - the system will automatically: 38 | // - Split documents into semantic chunks 39 | // - Generate rich context for each chunk 40 | // - Store embeddings with contextual information 41 | docsPath := filepath.Join("examples", "docs") 42 | if err := rag.AddDocuments(context.Background(), docsPath); err != nil { 43 | fmt.Printf("Failed to add documents: %v\n", err) 44 | os.Exit(1) 45 | } 46 | 47 | // Simple search with automatic context enhancement 48 | query := "What are the key features of the product?" 49 | response, err := rag.Search(context.Background(), query) 50 | if err != nil { 51 | fmt.Printf("Failed to search: %v\n", err) 52 | os.Exit(1) 53 | } 54 | 55 | fmt.Printf("\nQuery: %s\nResponse: %s\n", query, response) 56 | } 57 | 58 | func advancedExample() { 59 | fmt.Println("\n=== Advanced Contextual RAG Example ===") 60 | 61 | // Create a custom configuration 62 | config := &raggo.ContextualRAGConfig{ 63 | Collection: "advanced_contextual_docs", 64 | Model: "text-embedding-3-small", // Embedding model 65 | LLMModel: "gpt-4o-mini", // Model for context generation 66 | ChunkSize: 300, // Larger chunks for more context 67 | ChunkOverlap: 75, // 25% overlap for better continuity 68 | TopK: 5, // Number of similar chunks to retrieve 69 | MinScore: 0.7, // Higher threshold for better relevance 70 | } 71 | 72 | // Initialize RAG with custom configuration 73 | rag, err := raggo.NewContextualRAG(config) 74 | if err != nil { 75 | fmt.Printf("Failed to initialize RAG: %v\n", err) 76 | os.Exit(1) 77 | } 78 | defer rag.Close() 79 | 80 | // Add documents with enhanced context 81 | docsPath := filepath.Join("examples", "docs") 82 | if err := rag.AddDocuments(context.Background(), docsPath); err != nil { 83 | fmt.Printf("Failed to add documents: %v\n", err) 84 | os.Exit(1) 85 | } 86 | 87 | // Complex query demonstrating context-aware search 88 | query := "How does the system handle high load scenarios and what optimizations are in place?" 89 | response, err := rag.Search(context.Background(), query) 90 | if err != nil { 91 | fmt.Printf("Failed to search: %v\n", err) 92 | os.Exit(1) 93 | } 94 | 95 | fmt.Printf("\nQuery: %s\nResponse: %s\n", query, response) 96 | } 97 | -------------------------------------------------------------------------------- /examples/embedding_example.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "context" 5 | "fmt" 6 | "log" 7 | "os" 8 | "path/filepath" 9 | 10 | "github.com/teilomillet/raggo" 11 | ) 12 | 13 | func main() { 14 | // Set log level to Info 15 | raggo.SetLogLevel(raggo.LogLevelInfo) 16 | 17 | // Create a new Parser 18 | parser := raggo.NewParser() 19 | 20 | // Create a new Chunker 21 | chunker, err := raggo.NewChunker( 22 | raggo.ChunkSize(100), // Set chunk size to 100 tokens 23 | raggo.ChunkOverlap(20), // Set chunk overlap to 20 tokens 24 | raggo.WithSentenceSplitter(raggo.SmartSentenceSplitter()), 25 | raggo.WithTokenCounter(raggo.NewDefaultTokenCounter()), 26 | ) 27 | if err != nil { 28 | log.Fatalf("Failed to create chunker: %v", err) 29 | } 30 | 31 | // Create a new Embedder 32 | embedder, err := raggo.NewEmbedder( 33 | raggo.SetEmbedderProvider("openai"), 34 | raggo.SetEmbedderAPIKey(os.Getenv("OPENAI_API_KEY")), // Make sure to set this environment variable 35 | raggo.SetEmbedderModel("text-embedding-3-small"), 36 | ) 37 | if err != nil { 38 | log.Fatalf("Failed to create embedder: %v", err) 39 | } 40 | 41 | // Create an EmbeddingService 42 | embeddingService := raggo.NewEmbeddingService(embedder) 43 | 44 | // Get the path to a PDF file in the testdata directory 45 | wd, err := os.Getwd() 46 | if err != nil { 47 | log.Fatalf("Failed to get working directory: %v", err) 48 | } 49 | pdfPath := filepath.Join(wd, "testdata", "CV.pdf") // Make sure this file exists 50 | 51 | // Parse the PDF file 52 | doc, err := parser.Parse(pdfPath) 53 | if err != nil { 54 | log.Fatalf("Failed to parse PDF: %v", err) 55 | } 56 | 57 | fmt.Printf("Parsed document with %d characters\n", len(doc.Content)) 58 | 59 | // Chunk the parsed content 60 | chunks := chunker.Chunk(doc.Content) 61 | 62 | fmt.Printf("Created %d chunks from the document\n", len(chunks)) 63 | 64 | // Embed the chunks 65 | embeddedChunks, err := embeddingService.EmbedChunks(context.Background(), chunks) 66 | if err != nil { 67 | log.Fatalf("Failed to embed chunks: %v", err) 68 | } 69 | 70 | fmt.Printf("Successfully embedded %d chunks\n", len(embeddedChunks)) 71 | 72 | // Print details of the first few embedded chunks 73 | for i, chunk := range embeddedChunks { 74 | if i >= 3 { 75 | break // Only print the first 3 chunks 76 | } 77 | fmt.Printf("Embedded Chunk %d:\n", i+1) 78 | fmt.Printf(" Text: %s\n", truncateString(chunk.Text, 50)) 79 | fmt.Printf(" Embedding Vector Length: %d\n", len(chunk.Embeddings["default"])) 80 | fmt.Printf(" Metadata: %v\n", chunk.Metadata) 81 | fmt.Println() 82 | } 83 | } 84 | 85 | // truncateString truncates a string to a specified length, adding an ellipsis if truncated 86 | func truncateString(s string, length int) string { 87 | if len(s) <= length { 88 | return s 89 | } 90 | return s[:length-3] + "..." 91 | } 92 | -------------------------------------------------------------------------------- /examples/loader_example.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "context" 5 | "fmt" 6 | "log" 7 | "os" 8 | "path/filepath" 9 | "time" 10 | 11 | "github.com/teilomillet/raggo" 12 | ) 13 | 14 | func main() { 15 | // Set log level to Debug for more verbose output 16 | raggo.SetLogLevel(raggo.LogLevelWarn) 17 | 18 | // Create a new Loader with custom options 19 | loader := raggo.NewLoader( 20 | raggo.SetLoaderTimeout(1*time.Minute), 21 | raggo.SetTempDir(os.TempDir()), 22 | ) 23 | 24 | // Example 1: Load from URL 25 | urlExample(loader) 26 | 27 | // Example 2: Load single file 28 | fileExample(loader) 29 | 30 | // Example 3: Load directory 31 | dirExample(loader) 32 | } 33 | 34 | func urlExample(loader raggo.Loader) { 35 | fmt.Println("Example 1: Loading from URL") 36 | url := "https://www.w3.org/WAI/ER/tests/xhtml/testfiles/resources/pdf/dummy.pdf" 37 | urlPath, err := loader.LoadURL(context.Background(), url) 38 | if err != nil { 39 | log.Printf("Error loading URL: %v\n", err) 40 | } else { 41 | fmt.Printf("Successfully loaded URL to: %s\n", urlPath) 42 | } 43 | fmt.Println() 44 | } 45 | 46 | func fileExample(loader raggo.Loader) { 47 | fmt.Println("Example 2: Loading single file") 48 | // Create a temporary file for testing 49 | tempFile, err := os.CreateTemp("", "raggo-test-*.txt") 50 | if err != nil { 51 | log.Printf("Error creating temp file: %v\n", err) 52 | return 53 | } 54 | defer os.Remove(tempFile.Name()) 55 | 56 | _, err = tempFile.WriteString("This is a test file for raggo loader.") 57 | if err != nil { 58 | log.Printf("Error writing to temp file: %v\n", err) 59 | return 60 | } 61 | tempFile.Close() 62 | 63 | filePath, err := loader.LoadFile(context.Background(), tempFile.Name()) 64 | if err != nil { 65 | log.Printf("Error loading file: %v\n", err) 66 | } else { 67 | fmt.Printf("Successfully loaded file to: %s\n", filePath) 68 | } 69 | fmt.Println() 70 | } 71 | 72 | func dirExample(loader raggo.Loader) { 73 | fmt.Println("Example 3: Loading directory") 74 | // Create a temporary directory with some files for testing 75 | tempDir, err := os.MkdirTemp("", "raggo-test-dir") 76 | if err != nil { 77 | log.Printf("Error creating temp directory: %v\n", err) 78 | return 79 | } 80 | defer os.RemoveAll(tempDir) 81 | 82 | // Create a few test files in the directory 83 | for i := 1; i <= 3; i++ { 84 | fileName := filepath.Join(tempDir, fmt.Sprintf("test-file-%d.txt", i)) 85 | err := os.WriteFile(fileName, []byte(fmt.Sprintf("This is test file %d", i)), 0644) 86 | if err != nil { 87 | log.Printf("Error creating test file: %v\n", err) 88 | return 89 | } 90 | } 91 | 92 | dirPaths, err := loader.LoadDir(context.Background(), tempDir) 93 | if err != nil { 94 | log.Printf("Error loading directory: %v\n", err) 95 | } else { 96 | fmt.Printf("Successfully loaded %d files from directory:\n", len(dirPaths)) 97 | for _, path := range dirPaths { 98 | fmt.Printf("- %s\n", path) 99 | } 100 | } 101 | } 102 | -------------------------------------------------------------------------------- /examples/milvus_example.go: -------------------------------------------------------------------------------- 1 | // File: milvus_example.go 2 | 3 | package main 4 | 5 | import ( 6 | "context" 7 | "log" 8 | "math/rand" 9 | "time" 10 | 11 | "github.com/teilomillet/raggo" 12 | ) 13 | 14 | const ( 15 | nEntities, dim = 10000, 128 16 | collectionName = "hello_multi_vectors" 17 | idCol, keyCol, embeddingCol1, embeddingCol2 = "ID", "key", "vector1", "vector2" 18 | topK = 3 19 | ) 20 | 21 | func main() { 22 | ctx := context.Background() 23 | 24 | log.Println("start connecting to Milvus") 25 | db, err := raggo.NewVectorDB( 26 | raggo.WithType("milvus"), 27 | raggo.WithAddress("localhost:19530"), 28 | ) 29 | if err != nil { 30 | log.Fatalf("failed to create Milvus client, err: %v", err) 31 | } 32 | if err := db.Connect(ctx); err != nil { 33 | log.Fatalf("failed to connect to Milvus, err: %v", err) 34 | } 35 | defer db.Close() 36 | 37 | // delete collection if exists 38 | has, err := db.HasCollection(ctx, collectionName) 39 | if err != nil { 40 | log.Fatalf("failed to check collection exists, err: %v", err) 41 | } 42 | if has { 43 | db.DropCollection(ctx, collectionName) 44 | } 45 | 46 | // create collection 47 | log.Printf("create collection `%s`\n", collectionName) 48 | schema := raggo.Schema{ 49 | Name: collectionName, 50 | Description: "hello_multi_vectors is a demo collection with multiple vector fields", 51 | Fields: []raggo.Field{ 52 | {Name: idCol, DataType: "int64", PrimaryKey: true, AutoID: true}, 53 | {Name: keyCol, DataType: "int64"}, 54 | {Name: embeddingCol1, DataType: "float_vector", Dimension: dim}, 55 | {Name: embeddingCol2, DataType: "float_vector", Dimension: dim}, 56 | }, 57 | } 58 | 59 | if err := db.CreateCollection(ctx, collectionName, schema); err != nil { 60 | log.Fatalf("create collection failed, err: %v", err) 61 | } 62 | 63 | // Generate and insert data 64 | var records []raggo.Record 65 | for i := 0; i < nEntities; i++ { 66 | vec1 := make([]float64, dim) 67 | vec2 := make([]float64, dim) 68 | for j := 0; j < dim; j++ { 69 | vec1[j] = rand.Float64() 70 | vec2[j] = rand.Float64() 71 | } 72 | records = append(records, raggo.Record{ 73 | Fields: map[string]interface{}{ 74 | keyCol: rand.Int63() % 512, 75 | embeddingCol1: vec1, 76 | embeddingCol2: vec2, 77 | }, 78 | }) 79 | } 80 | 81 | log.Println("start to insert data into collection") 82 | if err := db.Insert(ctx, collectionName, records); err != nil { 83 | log.Fatalf("failed to insert random data into `%s`, err: %v", collectionName, err) 84 | } 85 | 86 | log.Println("insert data done, start to flush") 87 | if err := db.Flush(ctx, collectionName); err != nil { 88 | log.Fatalf("failed to flush data, err: %v", err) 89 | } 90 | log.Println("flush data done") 91 | 92 | // build index 93 | log.Println("start creating index HNSW") 94 | index := raggo.Index{ 95 | Type: "HNSW", 96 | Metric: "L2", 97 | Parameters: map[string]interface{}{ 98 | "M": 16, 99 | "efConstruction": 256, 100 | }, 101 | } 102 | if err := db.CreateIndex(ctx, collectionName, embeddingCol1, index); err != nil { 103 | log.Fatalf("failed to create index for %s, err: %v", embeddingCol1, err) 104 | } 105 | if err := db.CreateIndex(ctx, collectionName, embeddingCol2, index); err != nil { 106 | log.Fatalf("failed to create index for %s, err: %v", embeddingCol2, err) 107 | } 108 | 109 | log.Printf("build HNSW index done for collection `%s`\n", collectionName) 110 | log.Printf("start to load collection `%s`\n", collectionName) 111 | 112 | // load collection 113 | if err := db.LoadCollection(ctx, collectionName); err != nil { 114 | log.Fatalf("failed to load collection, err: %v", err) 115 | } 116 | 117 | log.Println("load collection done") 118 | // Prepare vectors for hybrid search 119 | vec2search1 := records[len(records)-2].Fields[embeddingCol1].([]float64) 120 | vec2search2 := records[len(records)-1].Fields[embeddingCol2].([]float64) 121 | 122 | begin := time.Now() 123 | log.Println("start to execute hybrid search") 124 | searchVectors := map[string]raggo.Vector{ 125 | embeddingCol1: vec2search1, 126 | embeddingCol2: vec2search2, 127 | } 128 | 129 | db.SetColumnNames([]string{keyCol, embeddingCol1, embeddingCol2}) 130 | 131 | result, err := db.HybridSearch(ctx, collectionName, searchVectors, topK, "L2", map[string]interface{}{ 132 | "type": "HNSW", 133 | "ef": 100, 134 | }, nil) 135 | if err != nil { 136 | log.Fatalf("failed to perform hybrid search, err: %v", err) 137 | } 138 | log.Printf("hybrid search `%s` done, latency %v\n", collectionName, time.Since(begin)) 139 | for _, rs := range result { 140 | log.Printf("ID: %d, score %f, embedding1: %v, embedding2: %v\n", 141 | rs.ID, rs.Score, 142 | rs.Fields[embeddingCol1], rs.Fields[embeddingCol2]) 143 | } 144 | 145 | db.DropCollection(ctx, collectionName) 146 | } 147 | 148 | -------------------------------------------------------------------------------- /examples/parser_example.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "context" 5 | "fmt" 6 | "log" 7 | "os" 8 | "path/filepath" 9 | 10 | "github.com/teilomillet/raggo" 11 | ) 12 | 13 | func main() { 14 | // Set log level to Info by default 15 | raggo.SetLogLevel(raggo.LogLevelInfo) 16 | 17 | parser := raggo.NewParser() 18 | loader := raggo.NewLoader() 19 | 20 | fmt.Println("Running examples with INFO level logging:") 21 | runExamples(parser, loader) 22 | 23 | } 24 | 25 | func runExamples(parser raggo.Parser, loader raggo.Loader) { 26 | // Example 1: Parse PDF file 27 | pdfExample(parser) 28 | 29 | // Example 2: Parse text file 30 | textExample(parser) 31 | 32 | // Example 3: Parse directory 33 | dirExample(loader) 34 | } 35 | 36 | func pdfExample(parser raggo.Parser) { 37 | fmt.Println("Example 1: Parsing PDF file") 38 | wd, err := os.Getwd() 39 | if err != nil { 40 | log.Fatalf("Failed to get working directory: %v", err) 41 | } 42 | pdfPath := filepath.Join(wd, "testdata", "CV.pdf") 43 | 44 | doc, err := parser.Parse(pdfPath) 45 | if err != nil { 46 | log.Printf("Error parsing PDF: %v\n", err) 47 | return 48 | } 49 | fmt.Printf("PDF parsed. Content length: %d\n", len(doc.Content)) 50 | } 51 | 52 | func textExample(parser raggo.Parser) { 53 | fmt.Println("Example 2: Parsing text file") 54 | tempFile, err := os.CreateTemp("", "raggo-test-*.txt") 55 | if err != nil { 56 | log.Printf("Error creating temp file: %v\n", err) 57 | return 58 | } 59 | defer os.Remove(tempFile.Name()) 60 | 61 | content := "This is a test file for raggo parser." 62 | if _, err := tempFile.WriteString(content); err != nil { 63 | log.Printf("Error writing to temp file: %v\n", err) 64 | return 65 | } 66 | tempFile.Close() 67 | 68 | doc, err := parser.Parse(tempFile.Name()) 69 | if err != nil { 70 | log.Printf("Error parsing text file: %v\n", err) 71 | return 72 | } 73 | fmt.Printf("Text file parsed. Content length: %d\n", len(doc.Content)) 74 | } 75 | 76 | func dirExample(loader raggo.Loader) { 77 | fmt.Println("Example 3: Loading directory") 78 | wd, err := os.Getwd() 79 | if err != nil { 80 | log.Fatalf("Failed to get working directory: %v", err) 81 | } 82 | testDataDir := filepath.Join(wd, "testdata") 83 | 84 | paths, err := loader.LoadDir(context.Background(), testDataDir) 85 | if err != nil { 86 | log.Printf("Error loading directory: %v\n", err) 87 | } else { 88 | fmt.Printf("Loaded %d files from directory\n", len(paths)) 89 | for i, path := range paths { 90 | fmt.Printf("%d: %s\n", i+1, path) 91 | } 92 | } 93 | } 94 | -------------------------------------------------------------------------------- /examples/process_embedding_benchmark.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "context" 5 | "fmt" 6 | "log" 7 | "os" 8 | "path/filepath" 9 | "sync" 10 | "time" 11 | 12 | "github.com/teilomillet/gollm" 13 | "github.com/teilomillet/raggo" 14 | ) 15 | 16 | const EXPECTED_FILE_COUNT = 1000 17 | 18 | func main() { 19 | raggo.SetLogLevel(raggo.LogLevelDebug) 20 | 21 | wd, err := os.Getwd() 22 | if err != nil { 23 | log.Fatalf("Failed to get working directory: %v", err) 24 | } 25 | 26 | targetDir := filepath.Join(wd, "benchmark_data") 27 | 28 | parser := raggo.NewParser() 29 | chunker, err := raggo.NewChunker( 30 | raggo.ChunkSize(512), 31 | raggo.ChunkOverlap(64), 32 | ) 33 | if err != nil { 34 | log.Fatalf("Failed to create chunker: %v", err) 35 | } 36 | 37 | embedder, err := raggo.NewEmbedder( 38 | raggo.SetEmbedderProvider("openai"), 39 | raggo.SetEmbedderAPIKey(os.Getenv("OPENAI_API_KEY")), 40 | raggo.SetEmbedderModel("text-embedding-3-small"), 41 | ) 42 | if err != nil { 43 | log.Fatalf("Failed to create embedder: %v", err) 44 | } 45 | 46 | llm, err := gollm.NewLLM( 47 | gollm.SetProvider("openai"), 48 | gollm.SetModel("gpt-4o-mini"), 49 | gollm.SetAPIKey(os.Getenv("OPENAI_API_KEY")), 50 | gollm.SetMaxTokens(2048), 51 | ) 52 | if err != nil { 53 | log.Fatalf("Failed to create LLM: %v", err) 54 | } 55 | 56 | // Create VectorDB instance 57 | vectorDB, err := raggo.NewVectorDB( 58 | raggo.WithType("milvus"), 59 | raggo.WithAddress("localhost:19530"), 60 | raggo.WithTimeout(30*time.Second), 61 | ) 62 | if err != nil { 63 | log.Fatalf("Failed to create vector database: %v", err) 64 | } 65 | defer vectorDB.Close() 66 | 67 | // Connect to the database 68 | ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) 69 | defer cancel() 70 | 71 | if err := vectorDB.Connect(ctx); err != nil { 72 | log.Fatalf("Failed to connect to vector database: %v", err) 73 | } 74 | 75 | collectionName := "benchmark_docs" 76 | 77 | // Create collection with schema if it doesn't exist 78 | exists, err := vectorDB.HasCollection(ctx, collectionName) 79 | if err != nil { 80 | log.Fatalf("Failed to check collection existence: %v", err) 81 | } 82 | 83 | if exists { 84 | err = vectorDB.DropCollection(ctx, collectionName) 85 | if err != nil { 86 | log.Fatalf("Failed to drop existing collection: %v", err) 87 | } 88 | } 89 | 90 | schema := raggo.Schema{ 91 | Name: collectionName, 92 | Fields: []raggo.Field{ 93 | {Name: "ID", DataType: "int64", PrimaryKey: true, AutoID: false}, 94 | {Name: "Embedding", DataType: "float_vector", Dimension: 1536}, 95 | {Name: "Text", DataType: "varchar", MaxLength: 65535}, 96 | {Name: "Metadata", DataType: "json", MaxLength: 65535}, 97 | }, 98 | } 99 | 100 | err = vectorDB.CreateCollection(ctx, collectionName, schema) 101 | if err != nil { 102 | log.Fatalf("Failed to create collection: %v", err) 103 | } 104 | 105 | // Create index for vector search 106 | err = vectorDB.CreateIndex(ctx, collectionName, "Embedding", raggo.Index{ 107 | Type: "HNSW", 108 | Metric: "L2", 109 | Parameters: map[string]interface{}{ 110 | "M": 16, 111 | "efConstruction": 256, 112 | }, 113 | }) 114 | if err != nil { 115 | log.Fatalf("Failed to create index: %v", err) 116 | } 117 | 118 | // Load the collection 119 | err = vectorDB.LoadCollection(ctx, collectionName) 120 | if err != nil { 121 | log.Fatalf("Failed to load collection: %v", err) 122 | } 123 | 124 | benchmarkPDFProcessing(parser, chunker, embedder, llm, vectorDB, targetDir, collectionName) 125 | } 126 | 127 | func benchmarkPDFProcessing(parser raggo.Parser, chunker raggo.Chunker, embedder raggo.Embedder, llm gollm.LLM, vectorDB *raggo.VectorDB, targetDir, collectionName string) { 128 | files, err := filepath.Glob(filepath.Join(targetDir, "*.pdf")) 129 | if err != nil { 130 | log.Fatalf("Failed to list PDF files: %v", err) 131 | } 132 | 133 | if len(files) != EXPECTED_FILE_COUNT { 134 | log.Fatalf("Expected %d files, but found %d. Please run the bash script to prepare the correct number of files.", EXPECTED_FILE_COUNT, len(files)) 135 | } 136 | 137 | fmt.Printf("Starting PDF processing benchmark with %d files...\n", len(files)) 138 | 139 | start := time.Now() 140 | 141 | var wg sync.WaitGroup 142 | var mu sync.Mutex 143 | errorCount := 0 144 | successCount := 0 145 | totalTokens := 0 146 | embedCount := 0 147 | summaryCount := 0 148 | 149 | for _, file := range files { 150 | wg.Add(1) 151 | go func(filePath string) { 152 | defer wg.Done() 153 | tokens, embeds, summaries, err := processAndEmbedPDF(parser, chunker, embedder, llm, vectorDB, filePath, collectionName) 154 | mu.Lock() 155 | defer mu.Unlock() 156 | if err != nil { 157 | log.Printf("Error processing %s: %v", filePath, err) 158 | errorCount++ 159 | } else { 160 | successCount++ 161 | totalTokens += tokens 162 | embedCount += embeds 163 | summaryCount += summaries 164 | } 165 | }(file) 166 | } 167 | 168 | wg.Wait() 169 | duration := time.Since(start) 170 | 171 | fmt.Printf("\nBenchmark Results:\n") 172 | fmt.Printf("Total files: %d\n", len(files)) 173 | fmt.Printf("Files successfully processed: %d\n", successCount) 174 | fmt.Printf("Files with processing errors: %d\n", errorCount) 175 | fmt.Printf("Total processing time: %v\n", duration) 176 | fmt.Printf("Total tokens processed: %d\n", totalTokens) 177 | fmt.Printf("Total embeddings created: %d\n", embedCount) 178 | fmt.Printf("Total summaries generated: %d\n", summaryCount) 179 | if successCount > 0 { 180 | fmt.Printf("Average time per successfully processed file: %v\n", duration/time.Duration(successCount)) 181 | } 182 | fmt.Printf("Average tokens per second: %.2f\n", float64(totalTokens)/duration.Seconds()) 183 | fmt.Printf("Average embeddings per second: %.2f\n", float64(embedCount)/duration.Seconds()) 184 | fmt.Printf("Average summaries per second: %.2f\n", float64(summaryCount)/duration.Seconds()) 185 | } 186 | 187 | func processAndEmbedPDF(parser raggo.Parser, chunker raggo.Chunker, embedder raggo.Embedder, llm gollm.LLM, vectorDB *raggo.VectorDB, filePath, collectionName string) (int, int, int, error) { 188 | log.Printf("Processing file: %s", filePath) 189 | 190 | doc, err := parser.Parse(filePath) 191 | if err != nil { 192 | return 0, 0, 0, fmt.Errorf("error parsing PDF: %w", err) 193 | } 194 | 195 | if len(doc.Content) == 0 { 196 | return 0, 0, 0, fmt.Errorf("parsed PDF content is empty") 197 | } 198 | 199 | log.Printf("Successfully parsed PDF: %s, Content length: %d", filePath, len(doc.Content)) 200 | 201 | summaryPrompt := gollm.NewPrompt(fmt.Sprintf("Summarize the following text in 2-3 sentences:\n\n%s", doc.Content)) 202 | summary, err := llm.Generate(context.Background(), summaryPrompt) 203 | if err != nil { 204 | return 0, 0, 0, fmt.Errorf("error generating summary: %w", err) 205 | } 206 | 207 | log.Printf("Generated summary for %s, Summary length: %d", filePath, len(summary)) 208 | 209 | chunks := chunker.Chunk(doc.Content) 210 | totalTokens := 0 211 | embedCount := 0 212 | 213 | // Create records for batch insertion 214 | var records []raggo.Record 215 | for i, chunk := range chunks { 216 | totalTokens += len(chunk.Text) // Simple approximation 217 | 218 | embedding, err := embedder.Embed(context.Background(), chunk.Text) 219 | if err != nil { 220 | return totalTokens, embedCount, 1, fmt.Errorf("error embedding chunk: %w", err) 221 | } 222 | embedCount++ 223 | 224 | // Create record with metadata 225 | records = append(records, raggo.Record{ 226 | Fields: map[string]interface{}{ 227 | "ID": int64(i), 228 | "Embedding": embedding, 229 | "Text": chunk.Text, 230 | "Metadata": map[string]interface{}{ 231 | "source": filePath, 232 | "chunk": i, 233 | "summary": summary, 234 | "timestamp": time.Now().Unix(), 235 | }, 236 | }, 237 | }) 238 | } 239 | 240 | // Batch insert records 241 | err = vectorDB.Insert(context.Background(), collectionName, records) 242 | if err != nil { 243 | return totalTokens, embedCount, 1, fmt.Errorf("error inserting into vector database: %w", err) 244 | } 245 | 246 | log.Printf("Successfully processed %s: %d tokens, %d embeddings", filePath, totalTokens, embedCount) 247 | 248 | return totalTokens, embedCount, 1, nil 249 | } 250 | -------------------------------------------------------------------------------- /examples/processing_benchmark.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "context" 5 | "fmt" 6 | "log" 7 | "os" 8 | "path/filepath" 9 | "sync" 10 | "sync/atomic" 11 | "time" 12 | 13 | "github.com/pkoukk/tiktoken-go" 14 | "github.com/teilomillet/gollm" 15 | "github.com/teilomillet/raggo" 16 | "golang.org/x/time/rate" 17 | ) 18 | 19 | const ( 20 | embeddingTPM = 5_000_000 // Tokens per minute for text-embedding-3-small 21 | embeddingRPM = 5_000 // Requests per minute for text-embedding-3-small 22 | defaultTimeout = 5 * time.Minute 23 | defaultBatchSize = 10 24 | defaultCollectionName = "pdf_embeddings" // Add default collection name 25 | ) 26 | 27 | type RateLimitedEmbedder struct { 28 | embedder raggo.Embedder 29 | tokenLimiter *rate.Limiter 30 | requestLimiter *rate.Limiter 31 | tokenCounter *TikTokenCounter 32 | } 33 | 34 | func NewRateLimitedEmbedder(embedder raggo.Embedder) (*RateLimitedEmbedder, error) { 35 | tokenCounter, err := NewTikTokenCounter("cl100k_base") // Use the appropriate encoding for your model 36 | if err != nil { 37 | return nil, fmt.Errorf("failed to create TikTokenCounter: %w", err) 38 | } 39 | 40 | return &RateLimitedEmbedder{ 41 | embedder: embedder, 42 | tokenLimiter: rate.NewLimiter(rate.Limit(embeddingTPM/60), embeddingTPM), // Convert TPM to tokens per second 43 | requestLimiter: rate.NewLimiter(rate.Limit(embeddingRPM/60), embeddingRPM), // Convert RPM to requests per second 44 | tokenCounter: tokenCounter, 45 | }, nil 46 | } 47 | 48 | func (rle *RateLimitedEmbedder) Embed(ctx context.Context, text string) ([]float64, error) { 49 | tokens := rle.tokenCounter.Count(text) 50 | 51 | if err := rle.tokenLimiter.WaitN(ctx, tokens); err != nil { 52 | return nil, fmt.Errorf("rate limit exceeded for tokens: %w", err) 53 | } 54 | 55 | if err := rle.requestLimiter.Wait(ctx); err != nil { 56 | return nil, fmt.Errorf("rate limit exceeded for requests: %w", err) 57 | } 58 | 59 | return rle.embedder.Embed(ctx, text) 60 | } 61 | 62 | type TikTokenCounter struct { 63 | tke *tiktoken.Tiktoken 64 | } 65 | 66 | func NewTikTokenCounter(encoding string) (*TikTokenCounter, error) { 67 | tke, err := tiktoken.GetEncoding(encoding) 68 | if err != nil { 69 | return nil, fmt.Errorf("failed to get encoding: %w", err) 70 | } 71 | return &TikTokenCounter{tke: tke}, nil 72 | } 73 | 74 | func (ttc *TikTokenCounter) Count(text string) int { 75 | return len(ttc.tke.Encode(text, nil, nil)) 76 | } 77 | 78 | func main() { 79 | // Set log level to Info for more detailed output 80 | raggo.SetLogLevel(raggo.LogLevelInfo) 81 | 82 | // Get the current working directory 83 | wd, err := os.Getwd() 84 | if err != nil { 85 | log.Fatalf("Failed to get working directory: %v", err) 86 | } 87 | 88 | // Construct the path to the testdata directory 89 | sourceDir := filepath.Join(wd, "testdata") 90 | targetDir := filepath.Join(wd, "benchmark_data") 91 | 92 | // Create the target directory if it doesn't exist 93 | if err := os.MkdirAll(targetDir, 0755); err != nil { 94 | log.Fatalf("Failed to create target directory: %v", err) 95 | } 96 | 97 | // Initialize components 98 | parser := raggo.PDFParser() // Initialize the PDF parser 99 | chunker, err := raggo.NewChunker( 100 | raggo.ChunkSize(512), 101 | raggo.ChunkOverlap(64), 102 | ) 103 | if err != nil { 104 | log.Fatalf("Failed to create chunker: %v", err) 105 | } 106 | 107 | embedder, err := raggo.NewEmbedder( 108 | raggo.SetEmbedderProvider("openai"), 109 | raggo.SetEmbedderAPIKey(os.Getenv("OPENAI_API_KEY")), 110 | raggo.SetEmbedderModel("text-embedding-3-small"), 111 | ) 112 | if err != nil { 113 | log.Fatalf("Failed to create embedder: %v", err) 114 | } 115 | 116 | rateLimitedEmbedder, err := NewRateLimitedEmbedder(embedder) 117 | if err != nil { 118 | log.Fatalf("Failed to create rate-limited embedder: %v", err) 119 | } 120 | 121 | // Create LLM and VectorDB 122 | llm, err := gollm.NewLLM( /* LLM configuration */ ) 123 | if err != nil { 124 | log.Fatalf("Failed to create LLM: %v", err) 125 | } 126 | 127 | vectorDB, err := raggo.NewVectorDB( /* VectorDB configuration */ ) 128 | if err != nil { 129 | log.Fatalf("Failed to create VectorDB: %v", err) 130 | } 131 | 132 | // Run the benchmark 133 | benchmarkPDFProcessing( 134 | parser, // parser 135 | chunker, // chunker 136 | rateLimitedEmbedder.embedder, // embedder 137 | llm, // llm 138 | vectorDB, // vectorDB 139 | sourceDir, // sourceDir 140 | targetDir, // targetDir 141 | ) 142 | } 143 | 144 | func benchmarkPDFProcessing( 145 | parser raggo.Parser, 146 | chunker raggo.Chunker, 147 | embedder raggo.Embedder, 148 | llm gollm.LLM, 149 | vectorDB *raggo.VectorDB, 150 | sourceDir string, 151 | targetDir string, 152 | ) { 153 | fmt.Printf("Starting PDF processing benchmark...\n") 154 | 155 | start := time.Now() 156 | 157 | // Load PDFs 158 | loadStart := time.Now() 159 | loader := raggo.NewConcurrentPDFLoader( 160 | raggo.SetLoaderTimeout(5*time.Minute), // Increased timeout 161 | raggo.SetTempDir(targetDir), 162 | ) 163 | loadedFiles, err := loader.LoadPDFsConcurrent(context.Background(), sourceDir, targetDir, 100) 164 | loadDuration := time.Since(loadStart) 165 | 166 | if err != nil { 167 | log.Printf("Warning: Encountered errors while loading PDF files: %v", err) 168 | } 169 | 170 | fmt.Printf("Attempted to load PDF files, successfully loaded %d in %v\n", len(loadedFiles), loadDuration) 171 | 172 | if len(loadedFiles) == 0 { 173 | log.Fatalf("No files were successfully loaded. Cannot continue benchmark.") 174 | } 175 | 176 | // Process and embed PDFs 177 | var wg sync.WaitGroup 178 | processStart := time.Now() 179 | errorCount := 0 180 | successCount := 0 181 | var totalTokens int64 182 | var embedCount int64 183 | var totalChunks int64 184 | 185 | for _, file := range loadedFiles { 186 | wg.Add(1) 187 | go func(filePath string) { 188 | defer wg.Done() 189 | tokens, embeds, chunks, err := processAndEmbedPDF( 190 | parser, 191 | chunker, 192 | embedder, 193 | llm, 194 | vectorDB, 195 | defaultCollectionName, // Add collection name 196 | filePath, 197 | ) 198 | if err != nil { 199 | log.Printf("Error processing %s: %v", filePath, err) 200 | errorCount++ 201 | } else { 202 | successCount++ 203 | atomic.AddInt64(&totalTokens, int64(tokens)) 204 | atomic.AddInt64(&embedCount, int64(embeds)) 205 | atomic.AddInt64(&totalChunks, int64(chunks)) 206 | } 207 | }(file) 208 | } 209 | 210 | wg.Wait() 211 | processDuration := time.Since(processStart) 212 | 213 | totalDuration := time.Since(start) 214 | 215 | // Print benchmark results 216 | fmt.Printf("\nBenchmark Results:\n") 217 | fmt.Printf("Total files attempted: %d\n", len(loadedFiles)) 218 | fmt.Printf("Files successfully loaded: %d\n", len(loadedFiles)) 219 | fmt.Printf("Files successfully processed: %d\n", successCount) 220 | fmt.Printf("Files with processing errors: %d\n", errorCount) 221 | fmt.Printf("Loading time: %v\n", loadDuration) 222 | fmt.Printf("Processing time: %v\n", processDuration) 223 | fmt.Printf("Total time: %v\n", totalDuration) 224 | fmt.Printf("Total tokens processed: %d\n", totalTokens) 225 | fmt.Printf("Total embeddings created: %d\n", embedCount) 226 | fmt.Printf("Total chunks processed: %d\n", totalChunks) 227 | if successCount > 0 { 228 | fmt.Printf("Average time per successfully processed file: %v\n", totalDuration/time.Duration(successCount)) 229 | } 230 | fmt.Printf("Average tokens per second: %.2f\n", float64(totalTokens)/processDuration.Seconds()) 231 | fmt.Printf("Average embeddings per second: %.2f\n", float64(embedCount)/processDuration.Seconds()) 232 | } 233 | 234 | func processAndEmbedPDF( 235 | parser raggo.Parser, 236 | chunker raggo.Chunker, 237 | embedder raggo.Embedder, 238 | llm gollm.LLM, 239 | vectorDB *raggo.VectorDB, 240 | collectionName string, 241 | filePath string, 242 | ) (int, int, int, error) { 243 | // Parse the PDF 244 | doc, err := parser.Parse(filePath) 245 | if err != nil { 246 | return 0, 0, 0, fmt.Errorf("error parsing PDF: %w", err) 247 | } 248 | 249 | // Chunk the content 250 | chunks := chunker.Chunk(doc.Content) 251 | totalTokens := 0 252 | embedCount := 0 253 | 254 | // Embed each chunk 255 | for _, chunk := range chunks { 256 | tokens := len(chunk.Text) 257 | 258 | _, err := embedder.Embed(context.Background(), chunk.Text) 259 | if err != nil { 260 | return totalTokens, embedCount, len(chunks), fmt.Errorf("error embedding chunk: %w", err) 261 | } 262 | 263 | embedCount++ 264 | totalTokens += tokens 265 | } 266 | 267 | log.Printf("Processed and embedded %s: %d characters, %d chunks, %d tokens", filePath, len(doc.Content), len(chunks), totalTokens) 268 | return totalTokens, embedCount, len(chunks), nil 269 | } 270 | -------------------------------------------------------------------------------- /examples/simple/main.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "context" 5 | "fmt" 6 | "log" 7 | "path/filepath" 8 | 9 | "github.com/teilomillet/raggo" 10 | ) 11 | 12 | func main() { 13 | // Create a new SimpleRAG instance with default configuration 14 | config := raggo.DefaultConfig() 15 | config.Collection = "my_documents" 16 | 17 | rag, err := raggo.NewSimpleRAG(config) 18 | if err != nil { 19 | log.Fatalf("Failed to create SimpleRAG: %v", err) 20 | } 21 | defer rag.Close() 22 | 23 | // Get absolute path to docs directory 24 | docsPath := filepath.Join("examples", "chat", "docs") 25 | 26 | // Add all documents from the directory 27 | ctx := context.Background() 28 | err = rag.AddDocuments(ctx, docsPath) 29 | if err != nil { 30 | log.Fatalf("Failed to add documents: %v", err) 31 | } 32 | 33 | // Ask a question about PressureValve 34 | query := "What is MountainPass's PressureValve system and how did it help during Black Friday?" 35 | response, err := rag.Search(ctx, query) 36 | if err != nil { 37 | log.Fatalf("Failed to search: %v", err) 38 | } 39 | 40 | // Print response 41 | fmt.Printf("\nQuestion: %s\n", query) 42 | fmt.Printf("\nAnswer: %s\n", response) 43 | } 44 | -------------------------------------------------------------------------------- /go.mod: -------------------------------------------------------------------------------- 1 | module github.com/teilomillet/raggo 2 | 3 | go 1.23.1 4 | 5 | require ( 6 | github.com/danaugrs/go-tsne/tsne v0.0.0-20220306155740-2250969e057f 7 | github.com/ledongthuc/pdf v0.0.0-20240201131950-da5b75280b06 8 | github.com/milvus-io/milvus-sdk-go/v2 v2.4.2 9 | github.com/philippgille/chromem-go v0.7.0 10 | github.com/pkoukk/tiktoken-go v0.1.7 11 | github.com/teilomillet/gofh v0.0.0-20240802075906-9ed4e405f11a 12 | github.com/teilomillet/gollm v0.1.1 13 | golang.org/x/time v0.8.0 14 | gonum.org/v1/gonum v0.15.1 15 | ) 16 | 17 | require ( 18 | github.com/bahlo/generic-list-go v0.2.0 // indirect 19 | github.com/buger/jsonparser v1.1.1 // indirect 20 | github.com/caarlos0/env/v11 v11.2.2 // indirect 21 | github.com/cockroachdb/errors v1.11.3 // indirect 22 | github.com/cockroachdb/logtags v0.0.0-20230118201751-21c54148d20b // indirect 23 | github.com/cockroachdb/redact v1.1.5 // indirect 24 | github.com/davecgh/go-spew v1.1.1 // indirect 25 | github.com/dlclark/regexp2 v1.11.4 // indirect 26 | github.com/gabriel-vasile/mimetype v1.4.7 // indirect 27 | github.com/getsentry/sentry-go v0.29.1 // indirect 28 | github.com/go-playground/locales v0.14.1 // indirect 29 | github.com/go-playground/universal-translator v0.18.1 // indirect 30 | github.com/go-playground/validator/v10 v10.23.0 // indirect 31 | github.com/gogo/protobuf v1.3.2 // indirect 32 | github.com/golang/protobuf v1.5.4 // indirect 33 | github.com/google/uuid v1.6.0 // indirect 34 | github.com/grpc-ecosystem/go-grpc-middleware v1.4.0 // indirect 35 | github.com/invopop/jsonschema v0.12.0 // indirect 36 | github.com/kr/pretty v0.3.1 // indirect 37 | github.com/kr/text v0.2.0 // indirect 38 | github.com/leodido/go-urn v1.4.0 // indirect 39 | github.com/mailru/easyjson v0.7.7 // indirect 40 | github.com/milvus-io/milvus-proto/go-api/v2 v2.4.16 // indirect 41 | github.com/pkg/errors v0.9.1 // indirect 42 | github.com/pmezard/go-difflib v1.0.0 // indirect 43 | github.com/rogpeppe/go-internal v1.13.1 // indirect 44 | github.com/stretchr/objx v0.5.2 // indirect 45 | github.com/stretchr/testify v1.9.0 // indirect 46 | github.com/tidwall/gjson v1.18.0 // indirect 47 | github.com/tidwall/match v1.1.1 // indirect 48 | github.com/tidwall/pretty v1.2.1 // indirect 49 | github.com/wk8/go-ordered-map/v2 v2.1.8 // indirect 50 | golang.org/x/crypto v0.35.0 // indirect 51 | golang.org/x/net v0.31.0 // indirect 52 | golang.org/x/sync v0.11.0 // indirect 53 | golang.org/x/sys v0.30.0 // indirect 54 | golang.org/x/text v0.22.0 // indirect 55 | google.golang.org/genproto/googleapis/rpc v0.0.0-20241118233622-e639e219e697 // indirect 56 | google.golang.org/grpc v1.68.0 // indirect 57 | google.golang.org/protobuf v1.35.2 // indirect 58 | gopkg.in/yaml.v3 v3.0.1 // indirect 59 | ) 60 | -------------------------------------------------------------------------------- /loader.go: -------------------------------------------------------------------------------- 1 | // Package raggo provides a high-level interface for document loading and processing 2 | // in RAG (Retrieval-Augmented Generation) systems. The loader component handles 3 | // various input sources with support for concurrent operations and configurable 4 | // behaviors. 5 | package raggo 6 | 7 | import ( 8 | "context" 9 | "net/http" 10 | "time" 11 | 12 | "github.com/teilomillet/raggo/rag" 13 | ) 14 | 15 | // Loader represents the main interface for loading documents from various sources. 16 | // It provides a unified API for handling: 17 | // - URLs: Download and process remote documents 18 | // - Files: Load and process local files 19 | // - Directories: Recursively process directory contents 20 | // 21 | // The interface is designed to be thread-safe and supports concurrent operations 22 | // with configurable timeouts and error handling. 23 | type Loader interface { 24 | // LoadURL downloads and processes a document from the given URL. 25 | // The function handles: 26 | // - HTTP/HTTPS downloads 27 | // - Timeout management 28 | // - Temporary file storage 29 | // 30 | // Returns the processed content and any errors encountered. 31 | LoadURL(ctx context.Context, url string) (string, error) 32 | 33 | // LoadFile processes a local file at the given path. 34 | // The function: 35 | // - Verifies file existence 36 | // - Handles file reading 37 | // - Manages temporary storage 38 | // 39 | // Returns the processed content and any errors encountered. 40 | LoadFile(ctx context.Context, path string) (string, error) 41 | 42 | // LoadDir recursively processes all files in a directory. 43 | // The function: 44 | // - Walks the directory tree 45 | // - Processes each file 46 | // - Handles errors gracefully 47 | // 48 | // Returns paths to all processed files and any errors encountered. 49 | LoadDir(ctx context.Context, dir string) ([]string, error) 50 | } 51 | 52 | // loaderWrapper encapsulates the internal loader implementation 53 | // providing a clean interface while maintaining all functionality. 54 | type loaderWrapper struct { 55 | internal *rag.Loader 56 | } 57 | 58 | // LoaderOption is a functional option for configuring a Loader. 59 | // It follows the functional options pattern to provide a clean 60 | // and extensible configuration API. 61 | // 62 | // Common options include: 63 | // - WithHTTPClient: Custom HTTP client configuration 64 | // - SetLoaderTimeout: Operation timeout settings 65 | // - SetTempDir: Temporary storage location 66 | type LoaderOption = rag.LoaderOption 67 | 68 | // WithHTTPClient sets a custom HTTP client for the Loader. 69 | // This enables customization of: 70 | // - Transport settings 71 | // - Proxy configuration 72 | // - Authentication mechanisms 73 | // - Connection pooling 74 | // 75 | // Example: 76 | // 77 | // client := &http.Client{ 78 | // Timeout: 60 * time.Second, 79 | // Transport: &http.Transport{ 80 | // MaxIdleConns: 10, 81 | // IdleConnTimeout: 30 * time.Second, 82 | // }, 83 | // } 84 | // loader := NewLoader(WithHTTPClient(client)) 85 | func WithHTTPClient(client *http.Client) LoaderOption { 86 | return rag.WithHTTPClient(client) 87 | } 88 | 89 | // SetLoaderTimeout sets a custom timeout for all loader operations. 90 | // The timeout applies to: 91 | // - URL downloads 92 | // - File operations 93 | // - Directory traversal 94 | // 95 | // Example: 96 | // 97 | // // Set a 2-minute timeout for all operations 98 | // loader := NewLoader(SetLoaderTimeout(2 * time.Minute)) 99 | func SetLoaderTimeout(timeout time.Duration) LoaderOption { 100 | return rag.WithTimeout(timeout) 101 | } 102 | 103 | // SetTempDir sets the temporary directory for file operations. 104 | // This directory is used for: 105 | // - Storing downloaded files 106 | // - Creating temporary copies 107 | // - Processing large documents 108 | // 109 | // Example: 110 | // 111 | // // Use a custom temporary directory 112 | // loader := NewLoader(SetTempDir("/path/to/temp")) 113 | func SetTempDir(dir string) LoaderOption { 114 | return rag.WithTempDir(dir) 115 | } 116 | 117 | // NewLoader creates a new Loader with the specified options. 118 | // It initializes a loader with sensible defaults and applies 119 | // any provided configuration options. 120 | // 121 | // Default settings: 122 | // - Standard HTTP client 123 | // - 30-second timeout 124 | // - System temporary directory 125 | // 126 | // Example: 127 | // 128 | // loader := NewLoader( 129 | // WithHTTPClient(customClient), 130 | // SetLoaderTimeout(time.Minute), 131 | // SetTempDir("/custom/temp"), 132 | // ) 133 | func NewLoader(opts ...LoaderOption) Loader { 134 | return &loaderWrapper{internal: rag.NewLoader(opts...)} 135 | } 136 | 137 | // LoadURL downloads and processes a document from the given URL. 138 | // The function handles the entire download process including: 139 | // - Context and timeout management 140 | // - HTTP request execution 141 | // - Response processing 142 | // - Temporary file management 143 | func (lw *loaderWrapper) LoadURL(ctx context.Context, url string) (string, error) { 144 | return lw.internal.LoadURL(ctx, url) 145 | } 146 | 147 | // LoadFile processes a local file at the given path. 148 | // The function ensures safe file handling by: 149 | // - Verifying file existence 150 | // - Creating temporary copies 151 | // - Managing file resources 152 | // - Handling processing errors 153 | func (lw *loaderWrapper) LoadFile(ctx context.Context, path string) (string, error) { 154 | return lw.internal.LoadFile(ctx, path) 155 | } 156 | 157 | // LoadDir recursively processes all files in a directory. 158 | // The function provides robust directory handling: 159 | // - Recursive traversal 160 | // - Error tolerance (continues on file errors) 161 | // - Progress tracking 162 | // - Resource cleanup 163 | func (lw *loaderWrapper) LoadDir(ctx context.Context, dir string) ([]string, error) { 164 | return lw.internal.LoadDir(ctx, dir) 165 | } 166 | -------------------------------------------------------------------------------- /logger.go: -------------------------------------------------------------------------------- 1 | // Package raggo provides a high-level logging interface for the Raggo framework, 2 | // built on top of the core rag package logging system. It offers: 3 | // - Multiple severity levels (Debug, Info, Warn, Error) 4 | // - Structured logging with key-value pairs 5 | // - Global log level control 6 | // - Consistent logging across the framework 7 | package raggo 8 | 9 | import ( 10 | "github.com/teilomillet/raggo/rag" 11 | ) 12 | 13 | // LogLevel represents the severity of a log message. 14 | // It is used to control which messages are output and 15 | // to indicate the importance of logged information. 16 | // 17 | // Available levels (from least to most severe): 18 | // - LogLevelDebug: Detailed information for debugging 19 | // - LogLevelInfo: General operational messages 20 | // - LogLevelWarn: Warning conditions 21 | // - LogLevelError: Error conditions 22 | // - LogLevelOff: Disable all logging 23 | type LogLevel = rag.LogLevel 24 | 25 | // Log levels define the available logging severities. 26 | // Higher levels include messages from all lower levels. 27 | const ( 28 | // LogLevelOff disables all logging output 29 | LogLevelOff = rag.LogLevelOff 30 | 31 | // LogLevelError enables only error messages 32 | // Use for conditions that prevent normal operation 33 | LogLevelError = rag.LogLevelError 34 | 35 | // LogLevelWarn enables warning and error messages 36 | // Use for potentially harmful situations 37 | LogLevelWarn = rag.LogLevelWarn 38 | 39 | // LogLevelInfo enables info, warning, and error messages 40 | // Use for general operational information 41 | LogLevelInfo = rag.LogLevelInfo 42 | 43 | // LogLevelDebug enables all message types 44 | // Use for detailed debugging information 45 | LogLevelDebug = rag.LogLevelDebug 46 | ) 47 | 48 | // Logger interface defines the logging operations available. 49 | // It supports structured logging with key-value pairs for 50 | // better log aggregation and analysis. 51 | // 52 | // Example usage: 53 | // 54 | // logger.Debug("Processing document", 55 | // "filename", "example.pdf", 56 | // "size", 1024, 57 | // "chunks", 5) 58 | type Logger = rag.Logger 59 | 60 | // SetLogLevel sets the global log level for the raggo package. 61 | // Messages below this level will not be logged. 62 | // 63 | // Example usage: 64 | // 65 | // // Enable all logging including debug 66 | // raggo.SetLogLevel(raggo.LogLevelDebug) 67 | // 68 | // // Only log errors 69 | // raggo.SetLogLevel(raggo.LogLevelError) 70 | // 71 | // // Disable all logging 72 | // raggo.SetLogLevel(raggo.LogLevelOff) 73 | func SetLogLevel(level LogLevel) { 74 | rag.SetGlobalLogLevel(level) 75 | } 76 | 77 | // Debug logs a message at debug level with optional key-value pairs. 78 | // Debug messages provide detailed information for troubleshooting. 79 | // 80 | // Example usage: 81 | // 82 | // raggo.Debug("Processing chunk", 83 | // "id", chunkID, 84 | // "size", chunkSize, 85 | // "overlap", overlap) 86 | func Debug(msg string, keysAndValues ...interface{}) { 87 | rag.GlobalLogger.Debug(msg, keysAndValues...) 88 | } 89 | 90 | // Info logs a message at info level with optional key-value pairs. 91 | // Info messages provide general operational information. 92 | // 93 | // Example usage: 94 | // 95 | // raggo.Info("Document added successfully", 96 | // "collection", collectionName, 97 | // "documentID", docID) 98 | func Info(msg string, keysAndValues ...interface{}) { 99 | rag.GlobalLogger.Info(msg, keysAndValues...) 100 | } 101 | 102 | // Warn logs a message at warning level with optional key-value pairs. 103 | // Warning messages indicate potential issues that don't prevent operation. 104 | // 105 | // Example usage: 106 | // 107 | // raggo.Warn("High memory usage detected", 108 | // "usedMB", memoryUsed, 109 | // "threshold", threshold) 110 | func Warn(msg string, keysAndValues ...interface{}) { 111 | rag.GlobalLogger.Warn(msg, keysAndValues...) 112 | } 113 | 114 | // Error logs a message at error level with optional key-value pairs. 115 | // Error messages indicate serious problems that affect normal operation. 116 | // 117 | // Example usage: 118 | // 119 | // raggo.Error("Failed to connect to vector store", 120 | // "error", err, 121 | // "retries", retryCount) 122 | func Error(msg string, keysAndValues ...interface{}) { 123 | rag.GlobalLogger.Error(msg, keysAndValues...) 124 | } 125 | -------------------------------------------------------------------------------- /rag/chunk.go: -------------------------------------------------------------------------------- 1 | // Package rag provides text chunking capabilities for processing documents into 2 | // manageable pieces suitable for vector embedding and retrieval. 3 | package rag 4 | 5 | import ( 6 | "fmt" 7 | "strings" 8 | 9 | "github.com/pkoukk/tiktoken-go" 10 | ) 11 | 12 | // Chunk represents a piece of text with associated metadata for tracking its position 13 | // and size within the original document. 14 | type Chunk struct { 15 | // Text contains the actual content of the chunk 16 | Text string 17 | // TokenSize represents the number of tokens in this chunk 18 | TokenSize int 19 | // StartSentence is the index of the first sentence in this chunk 20 | StartSentence int 21 | // EndSentence is the index of the last sentence in this chunk (exclusive) 22 | EndSentence int 23 | } 24 | 25 | // Chunker defines the interface for text chunking implementations. 26 | // Different implementations can provide various strategies for splitting text 27 | // while maintaining context and semantic meaning. 28 | type Chunker interface { 29 | // Chunk splits the input text into a slice of Chunks according to the 30 | // implementation's strategy. 31 | Chunk(text string) []Chunk 32 | } 33 | 34 | // TokenCounter defines the interface for counting tokens in a string. 35 | // This abstraction allows for different tokenization strategies (e.g., words, subwords). 36 | type TokenCounter interface { 37 | // Count returns the number of tokens in the given text according to the 38 | // implementation's tokenization strategy. 39 | Count(text string) int 40 | } 41 | 42 | // TextChunker provides an advanced implementation of the Chunker interface 43 | // with support for overlapping chunks and custom tokenization. 44 | type TextChunker struct { 45 | // ChunkSize is the target size of each chunk in tokens 46 | ChunkSize int 47 | // ChunkOverlap is the number of tokens that should overlap between adjacent chunks 48 | ChunkOverlap int 49 | // TokenCounter is used to count tokens in text segments 50 | TokenCounter TokenCounter 51 | // SentenceSplitter is a function that splits text into sentences 52 | SentenceSplitter func(string) []string 53 | } 54 | 55 | // NewTextChunker creates a new TextChunker with the given options. 56 | // It uses sensible defaults if no options are provided: 57 | // - ChunkSize: 200 tokens 58 | // - ChunkOverlap: 50 tokens 59 | // - TokenCounter: DefaultTokenCounter 60 | // - SentenceSplitter: DefaultSentenceSplitter 61 | func NewTextChunker(options ...TextChunkerOption) (*TextChunker, error) { 62 | tc := &TextChunker{ 63 | ChunkSize: 200, 64 | ChunkOverlap: 50, 65 | TokenCounter: &DefaultTokenCounter{}, 66 | SentenceSplitter: DefaultSentenceSplitter, 67 | } 68 | 69 | for _, option := range options { 70 | option(tc) 71 | } 72 | 73 | return tc, nil 74 | } 75 | 76 | // TextChunkerOption is a function type for configuring TextChunker instances. 77 | // This follows the functional options pattern for clean and flexible configuration. 78 | type TextChunkerOption func(*TextChunker) 79 | 80 | // Chunk splits the input text into chunks while preserving sentence boundaries 81 | // and maintaining the specified overlap between chunks. The algorithm: 82 | // 1. Splits the text into sentences 83 | // 2. Builds chunks by adding sentences until the chunk size limit is reached 84 | // 3. Creates overlap with previous chunk when starting a new chunk 85 | // 4. Tracks token counts and sentence indices for each chunk 86 | func (tc *TextChunker) Chunk(text string) []Chunk { 87 | sentences := tc.SentenceSplitter(text) 88 | var chunks []Chunk 89 | var currentChunk Chunk 90 | currentTokenCount := 0 91 | 92 | for i, sentence := range sentences { 93 | sentenceTokenCount := tc.TokenCounter.Count(sentence) 94 | 95 | if currentTokenCount+sentenceTokenCount > tc.ChunkSize && currentTokenCount > 0 { 96 | chunks = append(chunks, currentChunk) 97 | 98 | overlapStart := max(currentChunk.StartSentence, currentChunk.EndSentence-tc.estimateOverlapSentences(sentences, currentChunk.EndSentence, tc.ChunkOverlap)) 99 | currentChunk = Chunk{ 100 | Text: strings.Join(sentences[overlapStart:i+1], " "), 101 | TokenSize: 0, 102 | StartSentence: overlapStart, 103 | EndSentence: i + 1, 104 | } 105 | currentTokenCount = 0 106 | for j := overlapStart; j <= i; j++ { 107 | currentTokenCount += tc.TokenCounter.Count(sentences[j]) 108 | } 109 | } else { 110 | if currentTokenCount == 0 { 111 | currentChunk.StartSentence = i 112 | } 113 | currentChunk.Text += sentence + " " 114 | currentChunk.EndSentence = i + 1 115 | currentTokenCount += sentenceTokenCount 116 | } 117 | currentChunk.TokenSize = currentTokenCount 118 | } 119 | 120 | if currentChunk.TokenSize > 0 { 121 | chunks = append(chunks, currentChunk) 122 | } 123 | 124 | return chunks 125 | } 126 | 127 | // estimateOverlapSentences calculates how many sentences from the end of the 128 | // previous chunk should be included in the next chunk to achieve the desired 129 | // token overlap. 130 | func (tc *TextChunker) estimateOverlapSentences(sentences []string, endSentence, desiredOverlap int) int { 131 | overlapTokens := 0 132 | overlapSentences := 0 133 | for i := endSentence - 1; i >= 0 && overlapTokens < desiredOverlap; i-- { 134 | overlapTokens += tc.TokenCounter.Count(sentences[i]) 135 | overlapSentences++ 136 | } 137 | return overlapSentences 138 | } 139 | 140 | // DefaultSentenceSplitter provides a basic implementation for splitting text into sentences. 141 | // It uses common punctuation marks (., !, ?) as sentence boundaries. 142 | func DefaultSentenceSplitter(text string) []string { 143 | return strings.FieldsFunc(text, func(r rune) bool { 144 | return r == '.' || r == '!' || r == '?' 145 | }) 146 | } 147 | 148 | // SmartSentenceSplitter provides an advanced sentence splitting implementation that handles: 149 | // - Multiple punctuation marks (., !, ?) 150 | // - Common abbreviations 151 | // - Quoted sentences 152 | // - Parenthetical sentences 153 | // - Lists and enumerations 154 | func SmartSentenceSplitter(text string) []string { 155 | var sentences []string 156 | var currentSentence strings.Builder 157 | inQuote := false 158 | 159 | for _, r := range text { 160 | currentSentence.WriteRune(r) 161 | 162 | if r == '"' { 163 | inQuote = !inQuote 164 | } 165 | 166 | if (r == '.' || r == '!' || r == '?') && !inQuote { 167 | // Check if it's really the end of a sentence 168 | if len(sentences) > 0 || currentSentence.Len() > 1 { 169 | sentences = append(sentences, strings.TrimSpace(currentSentence.String())) 170 | currentSentence.Reset() 171 | } 172 | } 173 | } 174 | 175 | // Add any remaining text as a sentence 176 | if currentSentence.Len() > 0 { 177 | sentences = append(sentences, strings.TrimSpace(currentSentence.String())) 178 | } 179 | 180 | return sentences 181 | } 182 | 183 | // DefaultTokenCounter provides a simple word-based token counting implementation. 184 | // It splits text on whitespace to approximate token counts. This is suitable 185 | // for basic use cases but may not accurately reflect subword tokenization 186 | // used by language models. 187 | type DefaultTokenCounter struct{} 188 | 189 | // Count returns the number of words in the text, using whitespace as a delimiter. 190 | func (dtc *DefaultTokenCounter) Count(text string) int { 191 | return len(strings.Fields(text)) 192 | } 193 | 194 | // TikTokenCounter provides accurate token counting using the tiktoken library, 195 | // which implements the tokenization schemes used by OpenAI models. 196 | type TikTokenCounter struct { 197 | tke *tiktoken.Tiktoken 198 | } 199 | 200 | // NewTikTokenCounter creates a new TikTokenCounter using the specified encoding. 201 | // Common encodings include: 202 | // - "cl100k_base" (GPT-4, ChatGPT) 203 | // - "p50k_base" (GPT-3) 204 | // - "r50k_base" (Codex) 205 | func NewTikTokenCounter(encoding string) (*TikTokenCounter, error) { 206 | tke, err := tiktoken.GetEncoding(encoding) 207 | if err != nil { 208 | return nil, fmt.Errorf("failed to get encoding: %w", err) 209 | } 210 | return &TikTokenCounter{tke: tke}, nil 211 | } 212 | 213 | // Count returns the exact number of tokens in the text according to the 214 | // specified tiktoken encoding. 215 | func (ttc *TikTokenCounter) Count(text string) int { 216 | return len(ttc.tke.Encode(text, nil, nil)) 217 | } 218 | 219 | // max returns the larger of two integers. 220 | func max(a, b int) int { 221 | if a > b { 222 | return a 223 | } 224 | return b 225 | } 226 | -------------------------------------------------------------------------------- /rag/embed.go: -------------------------------------------------------------------------------- 1 | // Package rag provides functionality for converting text into vector embeddings 2 | // using various embedding providers (e.g., OpenAI, Cohere, local models). 3 | package rag 4 | 5 | import ( 6 | "context" 7 | "fmt" 8 | 9 | "github.com/teilomillet/raggo/rag/providers" 10 | ) 11 | 12 | // EmbedderConfig holds the configuration for creating an Embedder instance. 13 | // It supports multiple embedding providers and their specific options. 14 | type EmbedderConfig struct { 15 | // Provider specifies the embedding service to use (e.g., "openai", "cohere") 16 | Provider string 17 | // Options contains provider-specific configuration parameters 18 | Options map[string]interface{} 19 | } 20 | 21 | // EmbedderOption is a function type for configuring the EmbedderConfig. 22 | // It follows the functional options pattern for clean and flexible configuration. 23 | type EmbedderOption func(*EmbedderConfig) 24 | 25 | // SetProvider sets the provider for the Embedder. 26 | // Common providers include: 27 | // - "openai": OpenAI's text-embedding-ada-002 and other models 28 | // - "cohere": Cohere's embedding models 29 | // - "local": Local embedding models 30 | func SetProvider(provider string) EmbedderOption { 31 | return func(c *EmbedderConfig) { 32 | c.Provider = provider 33 | } 34 | } 35 | 36 | // SetModel sets the specific model to use for embedding. 37 | // The available models depend on the chosen provider. 38 | // Examples: 39 | // - OpenAI: "text-embedding-ada-002" 40 | // - Cohere: "embed-multilingual-v2.0" 41 | func SetModel(model string) EmbedderOption { 42 | return func(c *EmbedderConfig) { 43 | c.Options["model"] = model 44 | } 45 | } 46 | 47 | // SetAPIKey sets the authentication key for the embedding service. 48 | // This is required for most cloud-based embedding providers. 49 | func SetAPIKey(apiKey string) EmbedderOption { 50 | return func(c *EmbedderConfig) { 51 | c.Options["api_key"] = apiKey 52 | } 53 | } 54 | 55 | // SetOption sets a custom option for the Embedder. 56 | // This allows for provider-specific configuration options 57 | // that aren't covered by the standard options. 58 | func SetOption(key string, value interface{}) EmbedderOption { 59 | return func(c *EmbedderConfig) { 60 | c.Options[key] = value 61 | } 62 | } 63 | 64 | // NewEmbedder creates a new Embedder instance based on the provided options. 65 | // It uses the provider factory system to instantiate the appropriate embedder 66 | // implementation. Returns an error if: 67 | // - No provider is specified 68 | // - The specified provider is not registered 69 | // - The provider factory fails to create an embedder 70 | func NewEmbedder(opts ...EmbedderOption) (providers.Embedder, error) { 71 | config := &EmbedderConfig{ 72 | Options: make(map[string]interface{}), 73 | } 74 | for _, opt := range opts { 75 | opt(config) 76 | } 77 | if config.Provider == "" { 78 | return nil, fmt.Errorf("provider must be specified") 79 | } 80 | factory, err := providers.GetEmbedderFactory(config.Provider) 81 | if err != nil { 82 | return nil, err 83 | } 84 | return factory(config.Options) 85 | } 86 | 87 | // EmbeddedChunk represents a chunk of text along with its vector embeddings 88 | // and associated metadata. This is the core data structure for storing 89 | // and retrieving embedded content. 90 | type EmbeddedChunk struct { 91 | // Text is the original text content that was embedded 92 | Text string `json:"text"` 93 | // Embeddings maps embedding types to their vector representations 94 | // Multiple embeddings can exist for different models or purposes 95 | Embeddings map[string][]float64 `json:"embeddings"` 96 | // Metadata stores additional information about the chunk 97 | // This can include source document info, timestamps, etc. 98 | Metadata map[string]interface{} `json:"metadata"` 99 | } 100 | 101 | // EmbeddingService handles the process of converting text chunks into 102 | // vector embeddings. It encapsulates the embedding provider and provides 103 | // a high-level interface for embedding operations. 104 | type EmbeddingService struct { 105 | embedder providers.Embedder 106 | } 107 | 108 | // NewEmbeddingService creates a new embedding service with the specified embedder. 109 | // The embedder must be properly configured and ready to generate embeddings. 110 | func NewEmbeddingService(embedder providers.Embedder) *EmbeddingService { 111 | return &EmbeddingService{embedder: embedder} 112 | } 113 | 114 | // EmbedChunks processes a slice of text chunks and generates embeddings for each one. 115 | // It handles the embedding process in sequence, with debug output for monitoring. 116 | // The function: 117 | // 1. Allocates space for the results 118 | // 2. Processes each chunk through the embedder 119 | // 3. Creates EmbeddedChunk instances with the results 120 | // 4. Provides progress information via debug output 121 | // 122 | // Returns an error if any chunk fails to embed properly. 123 | func (s *EmbeddingService) EmbedChunks(ctx context.Context, chunks []Chunk) ([]EmbeddedChunk, error) { 124 | embeddedChunks := make([]EmbeddedChunk, 0, len(chunks)) 125 | 126 | // Debug output 127 | fmt.Printf("Processing %d chunks for embedding\n", len(chunks)) 128 | 129 | for i, chunk := range chunks { 130 | // Debug output for each chunk 131 | fmt.Printf("Processing chunk %d/%d (length: %d)\n", i+1, len(chunks), len(chunk.Text)) 132 | fmt.Printf("Chunk preview: %s\n", truncateString(chunk.Text, 100)) 133 | 134 | embedding, err := s.embedder.Embed(ctx, chunk.Text) 135 | if err != nil { 136 | return nil, fmt.Errorf("error embedding chunk %d: %w", i+1, err) 137 | } 138 | 139 | embeddedChunk := EmbeddedChunk{ 140 | Text: chunk.Text, 141 | Embeddings: map[string][]float64{ 142 | "default": embedding, 143 | }, 144 | Metadata: map[string]interface{}{ 145 | "token_size": chunk.TokenSize, 146 | "start_sentence": chunk.StartSentence, 147 | "end_sentence": chunk.EndSentence, 148 | "chunk_index": i, 149 | }, 150 | } 151 | embeddedChunks = append(embeddedChunks, embeddedChunk) 152 | 153 | // Debug output for successful embedding 154 | fmt.Printf("Successfully embedded chunk %d (embedding size: %d)\n", i+1, len(embedding)) 155 | } 156 | 157 | return embeddedChunks, nil 158 | } 159 | 160 | // truncateString shortens a string to the specified length, adding an ellipsis 161 | // if the string was truncated. This is used for debug output to keep log 162 | // messages readable. 163 | func truncateString(s string, n int) string { 164 | if len(s) <= n { 165 | return s 166 | } 167 | return s[:n] + "..." 168 | } 169 | -------------------------------------------------------------------------------- /rag/example_vectordb.go: -------------------------------------------------------------------------------- 1 | // Package rag provides retrieval-augmented generation capabilities. 2 | package rag 3 | 4 | import ( 5 | "context" 6 | "fmt" 7 | "sync" 8 | ) 9 | 10 | // ExampleDB demonstrates how to implement a new vector database provider. 11 | // This template shows the minimum required functionality and common patterns. 12 | // 13 | // Key Features to Implement: 14 | // - Thread-safe operations 15 | // - Vector similarity search 16 | // - Collection management 17 | // - Data persistence (if applicable) 18 | // - Error handling and logging 19 | type ExampleDB struct { 20 | // Configuration 21 | config *Config 22 | 23 | // Connection state 24 | isConnected bool 25 | 26 | // Collection management 27 | collections map[string]interface{} // Replace interface{} with your collection type 28 | mu sync.RWMutex // Protects concurrent access to collections 29 | 30 | // Search configuration 31 | columnNames []string // Names of columns to retrieve in search results 32 | dimension int // Vector dimension for embeddings 33 | } 34 | 35 | // newExampleDB creates a new ExampleDB instance with the given configuration. 36 | // Initialize your database connection and any required resources here. 37 | func newExampleDB(cfg *Config) (*ExampleDB, error) { 38 | // Get dimension from config parameters (example) 39 | dimension, ok := cfg.Parameters["dimension"].(int) 40 | if !ok { 41 | dimension = 1536 // Default dimension 42 | } 43 | 44 | return &ExampleDB{ 45 | config: cfg, 46 | collections: make(map[string]interface{}), 47 | dimension: dimension, 48 | }, nil 49 | } 50 | 51 | // Connect establishes a connection to the database. 52 | // Implement your connection logic here. 53 | func (db *ExampleDB) Connect(ctx context.Context) error { 54 | GlobalLogger.Debug("Connecting to example database", "address", db.config.Address) 55 | 56 | // Add your connection logic here 57 | // Example: 58 | // client, err := yourdb.Connect(db.config.Address) 59 | // if err != nil { 60 | // return fmt.Errorf("failed to connect: %w", err) 61 | // } 62 | 63 | db.isConnected = true 64 | return nil 65 | } 66 | 67 | // Close terminates the database connection. 68 | // Clean up any resources here. 69 | func (db *ExampleDB) Close() error { 70 | if !db.isConnected { 71 | return nil 72 | } 73 | 74 | // Add your cleanup logic here 75 | db.isConnected = false 76 | return nil 77 | } 78 | 79 | // HasCollection checks if a collection exists. 80 | func (db *ExampleDB) HasCollection(ctx context.Context, name string) (bool, error) { 81 | db.mu.RLock() 82 | defer db.mu.RUnlock() 83 | 84 | _, exists := db.collections[name] 85 | return exists, nil 86 | } 87 | 88 | // CreateCollection initializes a new collection. 89 | func (db *ExampleDB) CreateCollection(ctx context.Context, name string, schema Schema) error { 90 | db.mu.Lock() 91 | defer db.mu.Unlock() 92 | 93 | // Validate collection doesn't exist 94 | if _, exists := db.collections[name]; exists { 95 | return fmt.Errorf("collection %s already exists", name) 96 | } 97 | 98 | // Initialize your collection here 99 | // Example: 100 | // collection, err := db.client.CreateCollection(name, schema) 101 | // if err != nil { 102 | // return fmt.Errorf("failed to create collection: %w", err) 103 | // } 104 | // db.collections[name] = collection 105 | 106 | return nil 107 | } 108 | 109 | // DropCollection removes a collection. 110 | func (db *ExampleDB) DropCollection(ctx context.Context, name string) error { 111 | db.mu.Lock() 112 | defer db.mu.Unlock() 113 | 114 | delete(db.collections, name) 115 | return nil 116 | } 117 | 118 | // Insert adds new records to a collection. 119 | func (db *ExampleDB) Insert(ctx context.Context, collectionName string, data []Record) error { 120 | // Example vector conversion if needed: 121 | // vectors := make([][]float32, len(data)) 122 | // for i, record := range data { 123 | // vectors[i] = toFloat32Slice(record.Vector) 124 | // } 125 | 126 | // Add your insert logic here 127 | return nil 128 | } 129 | 130 | // Search performs vector similarity search. 131 | func (db *ExampleDB) Search(ctx context.Context, collectionName string, vectors map[string]Vector, topK int, metricType string, searchParams map[string]interface{}) ([]SearchResult, error) { 132 | // Example implementation steps: 133 | // 1. Convert vectors if needed 134 | // 2. Perform search 135 | // 3. Format results 136 | 137 | return nil, fmt.Errorf("not implemented") 138 | } 139 | 140 | // HybridSearch combines vector and keyword search (optional). 141 | func (db *ExampleDB) HybridSearch(ctx context.Context, collectionName string, vectors map[string]Vector, topK int, metricType string, searchParams map[string]interface{}, reranker interface{}) ([]SearchResult, error) { 142 | return nil, fmt.Errorf("hybrid search not supported") 143 | } 144 | 145 | // Additional optional methods: 146 | 147 | // Flush ensures data persistence (if applicable). 148 | func (db *ExampleDB) Flush(ctx context.Context, collectionName string) error { 149 | return nil 150 | } 151 | 152 | // CreateIndex builds search indexes (if applicable). 153 | func (db *ExampleDB) CreateIndex(ctx context.Context, collectionName, field string, index Index) error { 154 | return nil 155 | } 156 | 157 | // LoadCollection prepares a collection for searching (if needed). 158 | func (db *ExampleDB) LoadCollection(ctx context.Context, name string) error { 159 | return nil 160 | } 161 | 162 | // SetColumnNames configures which fields to return in search results. 163 | func (db *ExampleDB) SetColumnNames(names []string) { 164 | db.columnNames = names 165 | } 166 | 167 | // Helper functions 168 | 169 | // exampleToFloat32Slice converts vectors if your database needs a different format. 170 | // Note: If you need float32 conversion, consider using the existing toFloat32Slice 171 | // function from the rag package instead of implementing your own. 172 | func exampleToFloat32Slice(v Vector) []float32 { 173 | result := make([]float32, len(v)) 174 | for i, val := range v { 175 | result[i] = float32(val) 176 | } 177 | return result 178 | } 179 | -------------------------------------------------------------------------------- /rag/load.go: -------------------------------------------------------------------------------- 1 | // Package rag provides document loading functionality for the Raggo framework. 2 | // The loader component handles various input sources including local files, 3 | // directories, and URLs, with support for concurrent operations and 4 | // configurable timeouts. 5 | package rag 6 | 7 | import ( 8 | "context" 9 | "io" 10 | "net/http" 11 | "os" 12 | "path/filepath" 13 | "time" 14 | ) 15 | 16 | // Loader represents the internal loader implementation. 17 | // It provides methods for loading documents from various sources 18 | // with configurable HTTP client, timeout settings, and temporary 19 | // storage management. The loader is designed to be thread-safe 20 | // and can handle concurrent loading operations. 21 | type Loader struct { 22 | client *http.Client // HTTP client for URL downloads 23 | timeout time.Duration // Timeout for operations 24 | tempDir string // Directory for temporary files 25 | logger Logger // Logger for operation tracking 26 | } 27 | 28 | // NewLoader creates a new Loader with the given options. 29 | // It initializes a loader with default settings and applies 30 | // any provided options. Default settings include: 31 | // - Standard HTTP client 32 | // - 30-second timeout 33 | // - System temporary directory 34 | // - Global logger instance 35 | func NewLoader(opts ...LoaderOption) *Loader { 36 | l := &Loader{ 37 | client: http.DefaultClient, 38 | timeout: 30 * time.Second, 39 | tempDir: os.TempDir(), 40 | logger: GlobalLogger, 41 | } 42 | 43 | for _, opt := range opts { 44 | opt(l) 45 | } 46 | 47 | return l 48 | } 49 | 50 | // LoaderOption is a functional option for configuring a Loader. 51 | // It follows the functional options pattern to provide a clean 52 | // and extensible way to configure the loader. 53 | type LoaderOption func(*Loader) 54 | 55 | // WithHTTPClient sets a custom HTTP client for the Loader. 56 | // This allows customization of the HTTP client used for URL downloads, 57 | // enabling features like custom transport settings, proxies, or 58 | // authentication mechanisms. 59 | func WithHTTPClient(client *http.Client) LoaderOption { 60 | return func(l *Loader) { 61 | l.client = client 62 | } 63 | } 64 | 65 | // WithTimeout sets a custom timeout for the Loader. 66 | // This timeout applies to all operations including: 67 | // - URL downloads 68 | // - File operations 69 | // - Directory traversal 70 | func WithTimeout(timeout time.Duration) LoaderOption { 71 | return func(l *Loader) { 72 | l.timeout = timeout 73 | } 74 | } 75 | 76 | // WithTempDir sets the temporary directory for downloaded files. 77 | // This directory is used to store: 78 | // - Downloaded files from URLs 79 | // - Copies of local files for processing 80 | // - Temporary files during directory operations 81 | func WithTempDir(dir string) LoaderOption { 82 | return func(l *Loader) { 83 | l.tempDir = dir 84 | } 85 | } 86 | 87 | // WithLogger sets a custom logger for the Loader. 88 | // The logger is used to track operations and debug issues 89 | // across all loading operations. 90 | func WithLogger(logger Logger) LoaderOption { 91 | return func(l *Loader) { 92 | l.logger = logger 93 | } 94 | } 95 | 96 | // LoadURL downloads a file from the given URL and stores it in the temporary directory. 97 | // The function: 98 | // 1. Creates a context with the configured timeout 99 | // 2. Downloads the file using the HTTP client 100 | // 3. Stores the file in the temporary directory 101 | // 4. Returns the path to the downloaded file 102 | // 103 | // The downloaded file's name is derived from the URL's base name. 104 | func (l *Loader) LoadURL(ctx context.Context, url string) (string, error) { 105 | l.logger.Debug("Starting LoadURL", "url", url) 106 | ctx, cancel := context.WithTimeout(ctx, l.timeout) 107 | defer cancel() 108 | 109 | req, err := http.NewRequestWithContext(ctx, "GET", url, nil) 110 | if err != nil { 111 | l.logger.Error("Failed to create request", "url", url, "error", err) 112 | return "", err 113 | } 114 | 115 | resp, err := l.client.Do(req) 116 | if err != nil { 117 | l.logger.Error("Failed to execute request", "url", url, "error", err) 118 | return "", err 119 | } 120 | defer resp.Body.Close() 121 | 122 | filename := filepath.Base(url) 123 | destPath := filepath.Join(l.tempDir, filename) 124 | 125 | out, err := os.Create(destPath) 126 | if err != nil { 127 | l.logger.Error("Failed to create file", "path", destPath, "error", err) 128 | return "", err 129 | } 130 | defer out.Close() 131 | 132 | _, err = io.Copy(out, resp.Body) 133 | if err != nil { 134 | l.logger.Error("Failed to write file content", "path", destPath, "error", err) 135 | return "", err 136 | } 137 | 138 | l.logger.Debug("Successfully loaded URL", "url", url, "path", destPath) 139 | return destPath, nil 140 | } 141 | 142 | // LoadFile copies a file to the temporary directory and returns its path. 143 | // The function: 144 | // 1. Verifies the source file exists 145 | // 2. Creates a copy in the temporary directory 146 | // 3. Returns the path to the copied file 147 | // 148 | // This ensures that the original file remains unchanged during processing. 149 | func (l *Loader) LoadFile(ctx context.Context, path string) (string, error) { 150 | l.logger.Debug("Starting LoadFile", "path", path) 151 | 152 | _, err := os.Stat(path) 153 | if err != nil { 154 | l.logger.Error("File does not exist", "path", path, "error", err) 155 | return "", err 156 | } 157 | 158 | filename := filepath.Base(path) 159 | destPath := filepath.Join(l.tempDir, filename) 160 | 161 | src, err := os.Open(path) 162 | if err != nil { 163 | l.logger.Error("Failed to open source file", "path", path, "error", err) 164 | return "", err 165 | } 166 | defer src.Close() 167 | 168 | dest, err := os.Create(destPath) 169 | if err != nil { 170 | l.logger.Error("Failed to create destination file", "path", destPath, "error", err) 171 | return "", err 172 | } 173 | defer dest.Close() 174 | 175 | _, err = io.Copy(dest, src) 176 | if err != nil { 177 | l.logger.Error("Failed to copy file", "source", path, "destination", destPath, "error", err) 178 | return "", err 179 | } 180 | 181 | l.logger.Debug("Successfully loaded file", "source", path, "destination", destPath) 182 | return destPath, nil 183 | } 184 | 185 | // LoadDir recursively processes all files in a directory. 186 | // The function: 187 | // 1. Walks through the directory tree 188 | // 2. Processes each file encountered 189 | // 3. Returns paths to all processed files 190 | // 191 | // Files that fail to load are logged but don't stop the process. 192 | // The function continues with the next file on error. 193 | func (l *Loader) LoadDir(ctx context.Context, dir string) ([]string, error) { 194 | l.logger.Debug("Starting LoadDir", "dir", dir) 195 | 196 | var loadedFiles []string 197 | err := filepath.Walk(dir, func(path string, info os.FileInfo, err error) error { 198 | if err != nil { 199 | l.logger.Error("Error accessing path", "path", path, "error", err) 200 | return err 201 | } 202 | if !info.IsDir() { 203 | l.logger.Debug("Processing file", "path", path) 204 | loadedPath, err := l.LoadFile(ctx, path) 205 | if err != nil { 206 | l.logger.Warn("Failed to load file", "path", path, "error", err) 207 | return nil // Continue with next file 208 | } 209 | loadedFiles = append(loadedFiles, loadedPath) 210 | } 211 | return nil 212 | }) 213 | 214 | if err != nil { 215 | l.logger.Error("Error walking directory", "dir", dir, "error", err) 216 | return nil, err 217 | } 218 | 219 | l.logger.Debug("Successfully loaded directory", "dir", dir, "fileCount", len(loadedFiles)) 220 | return loadedFiles, nil 221 | } 222 | -------------------------------------------------------------------------------- /rag/log.go: -------------------------------------------------------------------------------- 1 | // Package rag provides a flexible logging system for the Raggo framework. 2 | // It supports multiple log levels, structured logging with key-value pairs, 3 | // and can be easily extended with custom logger implementations. 4 | package rag 5 | 6 | import ( 7 | "fmt" 8 | "log" 9 | "os" 10 | "strings" 11 | ) 12 | 13 | // LogLevel represents the severity level of a log message. 14 | // Higher values indicate more verbose logging. 15 | type LogLevel int 16 | 17 | const ( 18 | // LogLevelOff disables all logging 19 | LogLevelOff LogLevel = iota 20 | // LogLevelError enables only error messages 21 | LogLevelError 22 | // LogLevelWarn enables error and warning messages 23 | LogLevelWarn 24 | // LogLevelInfo enables error, warning, and info messages 25 | LogLevelInfo 26 | // LogLevelDebug enables all messages including debug 27 | LogLevelDebug 28 | ) 29 | 30 | // Logger defines the interface for logging operations. 31 | // Implementations must support multiple severity levels and 32 | // structured logging with key-value pairs. 33 | type Logger interface { 34 | // Debug logs a message at debug level with optional key-value pairs 35 | Debug(msg string, keysAndValues ...interface{}) 36 | // Info logs a message at info level with optional key-value pairs 37 | Info(msg string, keysAndValues ...interface{}) 38 | // Warn logs a message at warning level with optional key-value pairs 39 | Warn(msg string, keysAndValues ...interface{}) 40 | // Error logs a message at error level with optional key-value pairs 41 | Error(msg string, keysAndValues ...interface{}) 42 | // SetLevel changes the current logging level 43 | SetLevel(level LogLevel) 44 | } 45 | 46 | // DefaultLogger provides a basic implementation of the Logger interface 47 | // using the standard library's log package. It supports: 48 | // - Multiple log levels 49 | // - Structured logging with key-value pairs 50 | // - Output to os.Stderr by default 51 | // - Standard timestamp format 52 | type DefaultLogger struct { 53 | logger *log.Logger 54 | level LogLevel 55 | } 56 | 57 | // NewLogger creates a new DefaultLogger instance with the specified log level. 58 | // The logger writes to os.Stderr using the standard log package format: 59 | // timestamp + message + key-value pairs. 60 | func NewLogger(level LogLevel) Logger { 61 | return &DefaultLogger{ 62 | logger: log.New(os.Stderr, "", log.LstdFlags), 63 | level: level, 64 | } 65 | } 66 | 67 | // SetLevel updates the logging level of the DefaultLogger. 68 | // Messages below this level will not be logged. 69 | func (l *DefaultLogger) SetLevel(level LogLevel) { 70 | l.level = level 71 | } 72 | 73 | // log is an internal helper that handles the actual logging operation. 74 | // It checks the log level and formats the message with key-value pairs. 75 | func (l *DefaultLogger) log(level LogLevel, msg string, keysAndValues ...interface{}) { 76 | if level <= l.level { 77 | l.logger.Printf("%s: %s %v", level, msg, keysAndValues) 78 | } 79 | } 80 | 81 | // Debug logs a message at debug level. This level should be used for 82 | // detailed information needed for debugging purposes. 83 | func (l *DefaultLogger) Debug(msg string, keysAndValues ...interface{}) { 84 | l.log(LogLevelDebug, msg, keysAndValues...) 85 | } 86 | 87 | // Info logs a message at info level. This level should be used for 88 | // general operational information. 89 | func (l *DefaultLogger) Info(msg string, keysAndValues ...interface{}) { 90 | l.log(LogLevelInfo, msg, keysAndValues...) 91 | } 92 | 93 | // Warn logs a message at warning level. This level should be used for 94 | // potentially harmful situations that don't prevent normal operation. 95 | func (l *DefaultLogger) Warn(msg string, keysAndValues ...interface{}) { 96 | l.log(LogLevelWarn, msg, keysAndValues...) 97 | } 98 | 99 | // Error logs a message at error level. This level should be used for 100 | // error conditions that affect normal operation. 101 | func (l *DefaultLogger) Error(msg string, keysAndValues ...interface{}) { 102 | l.log(LogLevelError, msg, keysAndValues...) 103 | } 104 | 105 | // String returns the string representation of a LogLevel. 106 | // This is used for formatting log messages and configuration. 107 | func (l LogLevel) String() string { 108 | return [...]string{"OFF", "ERROR", "WARN", "INFO", "DEBUG"}[l] 109 | } 110 | 111 | // UnmarshalText implements the encoding.TextUnmarshaler interface. 112 | // It allows LogLevel to be configured from string values in configuration 113 | // files or environment variables. 114 | func (l *LogLevel) UnmarshalText(text []byte) error { 115 | switch strings.ToUpper(string(text)) { 116 | case "OFF": 117 | *l = LogLevelOff 118 | case "ERROR": 119 | *l = LogLevelError 120 | case "WARN": 121 | *l = LogLevelWarn 122 | case "INFO": 123 | *l = LogLevelInfo 124 | case "DEBUG": 125 | *l = LogLevelDebug 126 | default: 127 | return fmt.Errorf("invalid log level: %s", string(text)) 128 | } 129 | return nil 130 | } 131 | 132 | // GlobalLogger is the package-level logger instance used by default. 133 | // It can be accessed and modified by other packages using the rag framework. 134 | var GlobalLogger Logger 135 | 136 | // init initializes the global logger with a default configuration. 137 | // By default, it logs at INFO level to os.Stderr. 138 | func init() { 139 | GlobalLogger = NewLogger(LogLevelInfo) 140 | } 141 | 142 | // SetGlobalLogLevel sets the log level for the global logger instance. 143 | // This function provides a convenient way to control logging verbosity 144 | // across the entire application. 145 | func SetGlobalLogLevel(level LogLevel) { 146 | GlobalLogger.SetLevel(level) 147 | } 148 | -------------------------------------------------------------------------------- /rag/parse.go: -------------------------------------------------------------------------------- 1 | // Package rag provides document parsing capabilities for various file formats. 2 | // The parsing system is designed to be extensible, allowing users to add custom parsers 3 | // for different file types while maintaining a consistent interface. 4 | package rag 5 | 6 | import ( 7 | "fmt" 8 | "os" 9 | "path/filepath" 10 | "strings" 11 | 12 | "github.com/ledongthuc/pdf" 13 | ) 14 | 15 | // Document represents a parsed document with its content and associated metadata. 16 | // The Content field contains the extracted text, while Metadata stores additional 17 | // information about the document such as file type and path. 18 | type Document struct { 19 | Content string // The extracted text content of the document 20 | Metadata map[string]string // Additional metadata about the document 21 | } 22 | 23 | // Parser defines the interface for document parsing implementations. 24 | // Any type that implements this interface can be registered with the ParserManager 25 | // to handle specific file types. 26 | type Parser interface { 27 | // Parse processes a file at the given path and returns a Document. 28 | // It returns an error if the parsing operation fails. 29 | Parse(filePath string) (Document, error) 30 | } 31 | 32 | // ParserManager coordinates document parsing by managing different Parser implementations 33 | // and routing files to the appropriate parser based on their type. 34 | type ParserManager struct { 35 | // fileTypeDetector determines the file type based on the file path. 36 | fileTypeDetector func(string) string 37 | // parsers stores the registered parsers for different file types. 38 | parsers map[string]Parser 39 | } 40 | 41 | // NewParserManager creates a new ParserManager initialized with default settings 42 | // and parsers for common file types (PDF and text files). 43 | func NewParserManager() *ParserManager { 44 | pm := &ParserManager{ 45 | fileTypeDetector: defaultFileTypeDetector, 46 | parsers: make(map[string]Parser), 47 | } 48 | 49 | // Add default parsers 50 | pm.parsers["pdf"] = NewPDFParser() 51 | pm.parsers["text"] = NewTextParser() 52 | 53 | return pm 54 | } 55 | 56 | // Parse processes a document using the appropriate parser based on the file type. 57 | // It uses the configured fileTypeDetector to determine the file type and then 58 | // delegates to the corresponding parser. Returns an error if no suitable parser 59 | // is found or if parsing fails. 60 | func (pm *ParserManager) Parse(filePath string) (Document, error) { 61 | GlobalLogger.Debug("Starting to parse file", "path", filePath) 62 | fileType := pm.fileTypeDetector(filePath) 63 | parser, ok := pm.parsers[fileType] 64 | if !ok { 65 | GlobalLogger.Error("No parser available for file type", "type", fileType) 66 | return Document{}, fmt.Errorf("no parser available for file type: %s", fileType) 67 | } 68 | doc, err := parser.Parse(filePath) 69 | if err != nil { 70 | GlobalLogger.Error("Failed to parse document", "path", filePath, "error", err) 71 | return Document{}, err 72 | } 73 | GlobalLogger.Debug("Successfully parsed document", "path", filePath, "type", fileType) 74 | return doc, nil 75 | } 76 | 77 | // defaultFileTypeDetector determines file type based on file extension. 78 | // Currently supports .pdf and .txt files, returning "unknown" for other extensions. 79 | func defaultFileTypeDetector(filePath string) string { 80 | ext := strings.ToLower(filepath.Ext(filePath)) 81 | switch ext { 82 | case ".pdf": 83 | return "pdf" 84 | case ".txt": 85 | return "text" 86 | default: 87 | return "unknown" 88 | } 89 | } 90 | 91 | // SetFileTypeDetector allows customization of how file types are detected. 92 | // This can be used to implement more sophisticated file type detection beyond 93 | // simple extension matching. 94 | func (pm *ParserManager) SetFileTypeDetector(detector func(string) string) { 95 | pm.fileTypeDetector = detector 96 | } 97 | 98 | // AddParser registers a new parser for a specific file type. 99 | // This allows users to extend the system with custom parsers for additional 100 | // file formats. 101 | func (pm *ParserManager) AddParser(fileType string, parser Parser) { 102 | pm.parsers[fileType] = parser 103 | } 104 | 105 | // PDFParser implements the Parser interface for PDF files using the 106 | // ledongthuc/pdf library for text extraction. 107 | type PDFParser struct{} 108 | 109 | // NewPDFParser creates a new PDFParser instance. 110 | func NewPDFParser() *PDFParser { 111 | return &PDFParser{} 112 | } 113 | 114 | // Parse implements the Parser interface for PDF files. 115 | // It extracts text content from the PDF and returns it along with basic metadata. 116 | // Returns an error if the PDF cannot be processed. 117 | func (p *PDFParser) Parse(filePath string) (Document, error) { 118 | GlobalLogger.Debug("Starting to parse PDF", "path", filePath) 119 | content, err := p.extractText(filePath) 120 | if err != nil { 121 | GlobalLogger.Error("Failed to extract text from PDF", "path", filePath, "error", err) 122 | return Document{}, fmt.Errorf("failed to extract text: %w", err) 123 | } 124 | GlobalLogger.Debug("Successfully parsed PDF", "path", filePath) 125 | return Document{ 126 | Content: content, 127 | Metadata: map[string]string{ 128 | "file_type": "pdf", 129 | "file_path": filePath, 130 | }, 131 | }, nil 132 | } 133 | 134 | // extractText performs the actual text extraction from a PDF file. 135 | // It processes the PDF page by page, concatenating the extracted text. 136 | // Returns an error if any part of the extraction process fails. 137 | func (p *PDFParser) extractText(filePath string) (string, error) { 138 | file, err := os.Open(filePath) 139 | if err != nil { 140 | return "", fmt.Errorf("failed to open file: %w", err) 141 | } 142 | defer file.Close() 143 | 144 | fileInfo, err := file.Stat() 145 | if err != nil { 146 | return "", fmt.Errorf("failed to get file info: %w", err) 147 | } 148 | 149 | reader, err := pdf.NewReader(file, fileInfo.Size()) 150 | if err != nil { 151 | return "", fmt.Errorf("failed to create PDF reader: %w", err) 152 | } 153 | 154 | var textBuilder strings.Builder 155 | numPages := reader.NumPage() 156 | for i := 1; i <= numPages; i++ { 157 | page := reader.Page(i) 158 | if page.V.IsNull() { 159 | continue 160 | } 161 | content, err := page.GetPlainText(nil) 162 | if err != nil { 163 | return "", fmt.Errorf("failed to extract text from page %d: %w", i, err) 164 | } 165 | textBuilder.WriteString(content) 166 | textBuilder.WriteString("\n\n") 167 | } 168 | 169 | return textBuilder.String(), nil 170 | } 171 | 172 | // TextParser implements the Parser interface for plain text files. 173 | type TextParser struct{} 174 | 175 | // NewTextParser creates a new TextParser instance. 176 | func NewTextParser() *TextParser { 177 | return &TextParser{} 178 | } 179 | 180 | // Parse implements the Parser interface for text files. 181 | // It reads the entire file content and returns it along with basic metadata. 182 | // Returns an error if the file cannot be read. 183 | func (p *TextParser) Parse(filePath string) (Document, error) { 184 | GlobalLogger.Debug("Starting to parse text file", "path", filePath) 185 | content, err := os.ReadFile(filePath) 186 | if err != nil { 187 | GlobalLogger.Error("Failed to read text file", "path", filePath, "error", err) 188 | return Document{}, fmt.Errorf("failed to read file: %w", err) 189 | } 190 | GlobalLogger.Debug("Successfully parsed text file", "path", filePath) 191 | return Document{ 192 | Content: string(content), 193 | Metadata: map[string]string{ 194 | "file_type": "text", 195 | "file_path": filePath, 196 | }, 197 | }, nil 198 | } 199 | -------------------------------------------------------------------------------- /rag/providers/example_provider.go: -------------------------------------------------------------------------------- 1 | // Package providers includes this example to demonstrate how to implement 2 | // new embedding providers for the Raggo framework. This file shows the 3 | // recommended patterns and best practices for creating a provider that 4 | // integrates seamlessly with the system. 5 | package providers 6 | 7 | import ( 8 | "fmt" 9 | ) 10 | 11 | // ExampleProvider demonstrates how to implement a new embedding provider. 12 | // Replace this with your actual provider implementation. Your provider 13 | // should handle: 14 | // - Connection management 15 | // - Resource cleanup 16 | // - Error handling 17 | // - Rate limiting (if applicable) 18 | // - Batching (if supported) 19 | type ExampleProvider struct { 20 | // Add fields needed by your provider 21 | apiKey string 22 | model string 23 | dimension int 24 | // Add any connection or state management fields 25 | client interface{} 26 | } 27 | 28 | // NewExampleProvider shows how to create a new provider instance. 29 | // Your initialization function should: 30 | // 1. Validate the configuration 31 | // 2. Set up any connections or resources 32 | // 3. Initialize internal state 33 | // 4. Return a fully configured provider 34 | func NewExampleProvider(cfg *Config) (*ExampleProvider, error) { 35 | // Validate required configuration 36 | if cfg.APIKey == "" { 37 | return nil, fmt.Errorf("API key is required") 38 | } 39 | 40 | // Initialize your provider 41 | provider := &ExampleProvider{ 42 | apiKey: cfg.APIKey, 43 | model: cfg.Model, 44 | dimension: cfg.Dimension, 45 | } 46 | 47 | // Set up any connections or resources 48 | // Example: 49 | // client, err := yourapi.NewClient(cfg.APIKey) 50 | // if err != nil { 51 | // return nil, fmt.Errorf("failed to create client: %w", err) 52 | // } 53 | // provider.client = client 54 | 55 | return provider, nil 56 | } 57 | 58 | // Embed generates embeddings for a batch of input texts. 59 | // Your implementation should: 60 | // 1. Validate the inputs 61 | // 2. Prepare the batch request 62 | // 3. Call your embedding service 63 | // 4. Handle errors appropriately 64 | // 5. Return the vector representations as float32 arrays 65 | // 66 | // Note: The method accepts a slice of strings and returns a slice of float32 vectors 67 | // to match the Provider interface requirements. 68 | func (p *ExampleProvider) Embed(texts []string) ([][]float32, error) { 69 | // Validate input 70 | if len(texts) == 0 { 71 | return nil, fmt.Errorf("empty input texts") 72 | } 73 | 74 | // Initialize result slice 75 | result := make([][]float32, len(texts)) 76 | 77 | // Process each text in the batch 78 | for i, text := range texts { 79 | if text == "" { 80 | return nil, fmt.Errorf("empty text at position %d", i) 81 | } 82 | 83 | // Call your embedding service 84 | // Example: 85 | // response, err := p.client.CreateEmbedding(&Request{ 86 | // Text: text, 87 | // Model: p.model, 88 | // }) 89 | // if err != nil { 90 | // return nil, fmt.Errorf("embedding creation failed for text %d: %w", i, err) 91 | // } 92 | 93 | // For this example, return a mock vector 94 | mockVector := make([]float32, p.dimension) 95 | for j := range mockVector { 96 | mockVector[j] = 0.1 // Replace with actual embedding values 97 | } 98 | result[i] = mockVector 99 | } 100 | 101 | return result, nil 102 | } 103 | 104 | // GetDimension demonstrates how to implement the dimension reporting method. 105 | // Your implementation should: 106 | // 1. Return the correct dimension for your model 107 | // 2. Handle any model-specific variations 108 | // 3. Return an error if the dimension cannot be determined 109 | func (p *ExampleProvider) GetDimension() (int, error) { 110 | if p.dimension == 0 { 111 | return 0, fmt.Errorf("dimension not set") 112 | } 113 | return p.dimension, nil 114 | } 115 | 116 | // Close demonstrates how to implement resource cleanup. 117 | // Your implementation should: 118 | // 1. Close any open connections 119 | // 2. Release any held resources 120 | // 3. Handle cleanup errors appropriately 121 | func (p *ExampleProvider) Close() error { 122 | // Clean up your resources 123 | // Example: 124 | // if p.client != nil { 125 | // if err := p.client.Close(); err != nil { 126 | // return fmt.Errorf("failed to close client: %w", err) 127 | // } 128 | // } 129 | return nil 130 | } 131 | 132 | func init() { 133 | // Register your provider with a unique name 134 | Register("example", func(cfg *Config) (Provider, error) { 135 | return NewExampleProvider(cfg) 136 | }) 137 | } 138 | -------------------------------------------------------------------------------- /rag/providers/openai.go: -------------------------------------------------------------------------------- 1 | // Package providers implements embedding service providers for the Raggo framework. 2 | // The OpenAI provider offers high-quality text embeddings through OpenAI's API, 3 | // supporting models like text-embedding-3-small and text-embedding-3-large. 4 | package providers 5 | 6 | import ( 7 | "bytes" 8 | "context" 9 | "encoding/json" 10 | "fmt" 11 | "io" 12 | "net/http" 13 | "time" 14 | ) 15 | 16 | func init() { 17 | // Register the OpenAI provider when the package is initialized 18 | RegisterEmbedder("openai", NewOpenAIEmbedder) 19 | } 20 | 21 | // Default settings for the OpenAI embedder 22 | const ( 23 | // defaultEmbeddingAPI is the endpoint for OpenAI's embedding service 24 | defaultEmbeddingAPI = "https://api.openai.com/v1/embeddings" 25 | // defaultModelName is the recommended model for most use cases 26 | defaultModelName = "text-embedding-3-small" 27 | ) 28 | 29 | // OpenAIEmbedder implements the Embedder interface using OpenAI's API. 30 | // It supports various embedding models and handles API communication, 31 | // rate limiting, and error recovery. The embedder is designed to be 32 | // thread-safe and can be used concurrently. 33 | type OpenAIEmbedder struct { 34 | apiKey string // API key for authentication 35 | client *http.Client // HTTP client with timeout 36 | apiURL string // API endpoint URL 37 | modelName string // Selected embedding model 38 | } 39 | 40 | // NewOpenAIEmbedder creates a new OpenAI embedding provider with the given 41 | // configuration. The provider requires an API key and optionally accepts: 42 | // - model: The embedding model to use (defaults to text-embedding-3-small) 43 | // - api_url: Custom API endpoint URL 44 | // - timeout: Custom timeout duration 45 | // 46 | // Example config: 47 | // 48 | // config := map[string]interface{}{ 49 | // "api_key": "your-api-key", 50 | // "model": "text-embedding-3-small", 51 | // "timeout": 30 * time.Second, 52 | // } 53 | func NewOpenAIEmbedder(config map[string]interface{}) (Embedder, error) { 54 | apiKey, ok := config["api_key"].(string) 55 | if !ok || apiKey == "" { 56 | return nil, fmt.Errorf("API key is required for OpenAI embedder") 57 | } 58 | 59 | e := &OpenAIEmbedder{ 60 | apiKey: apiKey, 61 | client: &http.Client{Timeout: 30 * time.Second}, 62 | apiURL: defaultEmbeddingAPI, 63 | modelName: defaultModelName, 64 | } 65 | 66 | if model, ok := config["model"].(string); ok && model != "" { 67 | e.modelName = model 68 | } 69 | 70 | if apiURL, ok := config["api_url"].(string); ok && apiURL != "" { 71 | e.apiURL = apiURL 72 | } 73 | 74 | if timeout, ok := config["timeout"].(time.Duration); ok { 75 | e.client.Timeout = timeout 76 | } 77 | 78 | return e, nil 79 | } 80 | 81 | // embeddingRequest represents the JSON structure for API requests 82 | type embeddingRequest struct { 83 | Input string `json:"input"` // Text to embed 84 | Model string `json:"model"` // Model to use 85 | } 86 | 87 | // embeddingResponse represents the JSON structure for API responses 88 | type embeddingResponse struct { 89 | Data []struct { 90 | Embedding []float64 `json:"embedding"` // Vector representation 91 | } `json:"data"` 92 | } 93 | 94 | // Embed converts the input text into a vector representation using the 95 | // configured OpenAI model. The method handles: 96 | // - Request preparation and validation 97 | // - API communication with retry logic 98 | // - Response parsing and error handling 99 | // 100 | // The resulting vector captures the semantic meaning of the input text 101 | // and can be used for similarity search operations. 102 | func (e *OpenAIEmbedder) Embed(ctx context.Context, text string) ([]float64, error) { 103 | reqBody, err := json.Marshal(embeddingRequest{ 104 | Input: text, 105 | Model: e.modelName, 106 | }) 107 | if err != nil { 108 | return nil, fmt.Errorf("error marshaling request: %w", err) 109 | } 110 | 111 | req, err := http.NewRequestWithContext(ctx, "POST", e.apiURL, bytes.NewBuffer(reqBody)) 112 | if err != nil { 113 | return nil, fmt.Errorf("error creating request: %w", err) 114 | } 115 | 116 | req.Header.Set("Content-Type", "application/json") 117 | req.Header.Set("Authorization", "Bearer "+e.apiKey) 118 | 119 | resp, err := e.client.Do(req) 120 | if err != nil { 121 | return nil, fmt.Errorf("error sending request: %w", err) 122 | } 123 | defer resp.Body.Close() 124 | 125 | body, err := io.ReadAll(resp.Body) 126 | if err != nil { 127 | return nil, fmt.Errorf("error reading response body: %w", err) 128 | } 129 | 130 | if resp.StatusCode != http.StatusOK { 131 | return nil, fmt.Errorf("API request failed with status code %d: %s", resp.StatusCode, resp.Status) 132 | } 133 | 134 | var embeddingResp embeddingResponse 135 | err = json.Unmarshal(body, &embeddingResp) 136 | if err != nil { 137 | return nil, fmt.Errorf("error unmarshaling response: %w", err) 138 | } 139 | 140 | if len(embeddingResp.Data) == 0 { 141 | return nil, fmt.Errorf("no embedding data in response") 142 | } 143 | 144 | return embeddingResp.Data[0].Embedding, nil 145 | } 146 | 147 | // GetDimension returns the output dimension for the current embedding model. 148 | // Each model produces vectors of a fixed size: 149 | // - text-embedding-3-small: 1536 dimensions 150 | // - text-embedding-3-large: 3072 dimensions 151 | // - text-embedding-ada-002: 1536 dimensions 152 | // 153 | // This information is crucial for configuring vector databases and ensuring 154 | // compatibility across the system. 155 | func (e *OpenAIEmbedder) GetDimension() (int, error) { 156 | switch e.modelName { 157 | case "text-embedding-3-small": 158 | return 1536, nil 159 | case "text-embedding-3-large": 160 | return 3072, nil 161 | case "text-embedding-ada-002": 162 | return 1536, nil 163 | default: 164 | return 0, fmt.Errorf("unknown model: %s", e.modelName) 165 | } 166 | } 167 | -------------------------------------------------------------------------------- /rag/providers/register.go: -------------------------------------------------------------------------------- 1 | // Package providers implements a flexible system for managing different embedding 2 | // service providers in the Raggo framework. Each provider offers unique capabilities 3 | // for converting text into vector representations that capture semantic meaning. 4 | // The registration system allows new providers to be easily added and configured 5 | // while maintaining a consistent interface for the rest of the system. 6 | package providers 7 | 8 | import ( 9 | "context" 10 | "fmt" 11 | "sync" 12 | ) 13 | 14 | // EmbedderFactory is a function type that creates a new Embedder 15 | type EmbedderFactory func(config map[string]interface{}) (Embedder, error) 16 | 17 | var ( 18 | embedderFactories = make(map[string]EmbedderFactory) 19 | mu sync.RWMutex 20 | ) 21 | 22 | // RegisterEmbedder registers a new embedder factory 23 | func RegisterEmbedder(name string, factory EmbedderFactory) { 24 | mu.Lock() 25 | defer mu.Unlock() 26 | embedderFactories[name] = factory 27 | } 28 | 29 | // GetEmbedderFactory returns the factory for the given embedder name 30 | func GetEmbedderFactory(name string) (EmbedderFactory, error) { 31 | mu.RLock() 32 | defer mu.RUnlock() 33 | factory, ok := embedderFactories[name] 34 | if !ok { 35 | return nil, fmt.Errorf("embedder not found: %s", name) 36 | } 37 | return factory, nil 38 | } 39 | 40 | // Embedder interface defines the contract for embedding implementations 41 | type Embedder interface { 42 | // Embed generates embeddings for the given text 43 | Embed(ctx context.Context, text string) ([]float64, error) 44 | 45 | // GetDimension returns the dimension of the embeddings for the current model 46 | GetDimension() (int, error) 47 | } 48 | 49 | // Provider defines the interface that all embedding providers must implement. 50 | // This abstraction ensures that different providers can be used interchangeably 51 | // while providing their own specific implementation details. A provider is 52 | // responsible for converting text into vector representations that can be used 53 | // for semantic similarity search. 54 | type Provider interface { 55 | // Embed converts a slice of text inputs into their vector representations. 56 | // The method should handle batching and rate limiting internally. It returns 57 | // a slice of vectors, where each vector corresponds to the input text at the 58 | // same index. An error is returned if the embedding process fails. 59 | Embed(inputs []string) ([][]float32, error) 60 | 61 | // Close releases any resources held by the provider, such as API connections 62 | // or cached data. This method should be called when the provider is no longer 63 | // needed to prevent resource leaks. 64 | Close() error 65 | } 66 | 67 | // Config holds the configuration settings for an embedding provider. 68 | // Different providers may use different subsets of these settings, but 69 | // the configuration structure remains consistent to simplify provider 70 | // management and initialization. 71 | type Config struct { 72 | // APIKey is used for authentication with the provider's service. 73 | // For local models, this may be left empty. 74 | APIKey string 75 | 76 | // Model specifies which embedding model to use. Each provider may 77 | // offer multiple models with different characteristics. 78 | Model string 79 | 80 | // BatchSize determines how many texts can be embedded in a single API call. 81 | // This helps optimize performance and manage rate limits. 82 | BatchSize int 83 | 84 | // Dimension specifies the size of the output vectors. This must match 85 | // the chosen model's output dimension. 86 | Dimension int 87 | 88 | // Additional provider-specific settings can be added here 89 | Settings map[string]interface{} 90 | } 91 | 92 | // registry maintains a thread-safe map of provider factories. Each factory 93 | // is a function that creates a new instance of a specific provider type 94 | // using the provided configuration. 95 | type registry struct { 96 | mu sync.RWMutex 97 | factories map[string]func(cfg *Config) (Provider, error) 98 | } 99 | 100 | // The global registry instance that maintains all registered provider factories. 101 | var globalRegistry = ®istry{ 102 | factories: make(map[string]func(cfg *Config) (Provider, error)), 103 | } 104 | 105 | // Register adds a new provider factory to the global registry. The factory 106 | // function should create and configure a new instance of the provider when 107 | // called. If a provider with the same name already exists, it will be 108 | // overwritten, allowing for provider updates and replacements. 109 | func Register(name string, factory func(cfg *Config) (Provider, error)) { 110 | globalRegistry.mu.Lock() 111 | defer globalRegistry.mu.Unlock() 112 | globalRegistry.factories[name] = factory 113 | } 114 | 115 | // Get retrieves a provider factory from the registry and creates a new provider 116 | // instance using the supplied configuration. If the requested provider is not 117 | // found in the registry, an error is returned. This method is thread-safe and 118 | // can be called from multiple goroutines. 119 | func Get(name string, cfg *Config) (Provider, error) { 120 | globalRegistry.mu.RLock() 121 | factory, ok := globalRegistry.factories[name] 122 | globalRegistry.mu.RUnlock() 123 | 124 | if !ok { 125 | return nil, fmt.Errorf("provider not found: %s", name) 126 | } 127 | 128 | return factory(cfg) 129 | } 130 | 131 | // List returns the names of all registered providers. This is useful for 132 | // discovering available providers and validating provider names before 133 | // attempting to create instances. 134 | func List() []string { 135 | globalRegistry.mu.RLock() 136 | defer globalRegistry.mu.RUnlock() 137 | 138 | providers := make([]string, 0, len(globalRegistry.factories)) 139 | for name := range globalRegistry.factories { 140 | providers = append(providers, name) 141 | } 142 | return providers 143 | } 144 | -------------------------------------------------------------------------------- /rag/reranker.go: -------------------------------------------------------------------------------- 1 | // Package rag provides retrieval-augmented generation capabilities. 2 | package rag 3 | 4 | import ( 5 | "context" 6 | "sort" 7 | ) 8 | 9 | // RRFReranker implements Reciprocal Rank Fusion (RRF) for combining and reranking search results. 10 | // RRF is a robust rank fusion method that effectively combines results from different retrieval systems 11 | // without requiring score normalization. It uses the formula: RRF(d) = Σ 1/(k + r(d)) 12 | // where d is a document, k is a constant, and r(d) is the rank of document d in each result list. 13 | type RRFReranker struct { 14 | k float64 // k is a constant that prevents division by zero and controls the influence of high-ranked items 15 | } 16 | 17 | // NewRRFReranker creates a new RRF reranker with the specified k parameter. 18 | // The k parameter controls ranking influence - higher values of k decrease the 19 | // influence of high-ranked items. If k <= 0, it defaults to 60 (from the original RRF paper). 20 | // 21 | // Typical k values: 22 | // - k = 60: Standard value from RRF literature, good general-purpose setting 23 | // - k < 60: Increases influence of top-ranked items 24 | // - k > 60: More weight to lower-ranked items, smoother ranking distribution 25 | func NewRRFReranker(k float64) *RRFReranker { 26 | if k <= 0 { 27 | k = 60 // Default value from RRF paper 28 | } 29 | return &RRFReranker{k: k} 30 | } 31 | 32 | // Rerank combines and reranks results using Reciprocal Rank Fusion. 33 | // It takes results from dense (semantic) and sparse (lexical) search and combines 34 | // them using weighted RRF scores. The method handles cases where documents appear 35 | // in both result sets by combining their weighted scores. 36 | // 37 | // Parameters: 38 | // - ctx: Context for potential future extensions (e.g., timeouts, cancellation) 39 | // - query: The original search query (reserved for future extensions) 40 | // - denseResults: Results from dense/semantic search 41 | // - sparseResults: Results from sparse/lexical search 42 | // - denseWeight: Weight for dense search results (normalized internally) 43 | // - sparseWeight: Weight for sparse search results (normalized internally) 44 | // 45 | // Returns: 46 | // - []SearchResult: Reranked results sorted by combined score 47 | // - error: Currently always nil, reserved for future extensions 48 | // 49 | // The reranking process: 50 | // 1. Normalizes weights to sum to 1.0 51 | // 2. Calculates RRF scores for each result based on rank 52 | // 3. Applies weights to scores based on result source (dense/sparse) 53 | // 4. Combines scores for documents appearing in both result sets 54 | // 5. Sorts final results by combined score 55 | func (r *RRFReranker) Rerank( 56 | ctx context.Context, 57 | query string, 58 | denseResults, sparseResults []SearchResult, 59 | denseWeight, sparseWeight float64, 60 | ) ([]SearchResult, error) { 61 | // Normalize weights 62 | totalWeight := denseWeight + sparseWeight 63 | if totalWeight > 0 { 64 | denseWeight /= totalWeight 65 | sparseWeight /= totalWeight 66 | } else { 67 | denseWeight = 0.5 68 | sparseWeight = 0.5 69 | } 70 | 71 | // Create maps to store combined scores 72 | scores := make(map[int64]float64) 73 | docMap := make(map[int64]SearchResult) 74 | 75 | // Process dense results 76 | for rank, result := range denseResults { 77 | rrf := 1.0 / (float64(rank+1) + r.k) 78 | scores[result.ID] = rrf * denseWeight 79 | docMap[result.ID] = result 80 | } 81 | 82 | // Process sparse results 83 | for rank, result := range sparseResults { 84 | rrf := 1.0 / (float64(rank+1) + r.k) 85 | if score, exists := scores[result.ID]; exists { 86 | scores[result.ID] = score + rrf*sparseWeight 87 | } else { 88 | scores[result.ID] = rrf * sparseWeight 89 | docMap[result.ID] = result 90 | } 91 | } 92 | 93 | // Convert scores back to results 94 | results := make([]SearchResult, 0, len(scores)) 95 | for id, score := range scores { 96 | result := docMap[id] 97 | result.Score = score 98 | results = append(results, result) 99 | } 100 | 101 | // Sort by combined score 102 | sort.Slice(results, func(i, j int) bool { 103 | return results[i].Score > results[j].Score 104 | }) 105 | 106 | // Return all reranked results 107 | return results, nil 108 | } 109 | -------------------------------------------------------------------------------- /rag/vector_interface.go: -------------------------------------------------------------------------------- 1 | // Package rag provides a unified interface for interacting with vector databases, 2 | // offering a clean abstraction layer for vector similarity search operations. 3 | package rag 4 | 5 | import ( 6 | "context" 7 | "fmt" 8 | "time" 9 | ) 10 | 11 | // VectorDB defines the standard interface that all vector database implementations must implement. 12 | // It provides operations for managing collections, inserting data, and performing vector similarity searches. 13 | type VectorDB interface { 14 | // Connect establishes a connection to the vector database. 15 | Connect(ctx context.Context) error 16 | 17 | // Close terminates the connection to the vector database. 18 | Close() error 19 | 20 | // HasCollection checks if a collection with the given name exists. 21 | HasCollection(ctx context.Context, name string) (bool, error) 22 | 23 | // DropCollection removes a collection and all its data. 24 | DropCollection(ctx context.Context, name string) error 25 | 26 | // CreateCollection creates a new collection with the specified schema. 27 | CreateCollection(ctx context.Context, name string, schema Schema) error 28 | 29 | // Insert adds new records to the specified collection. 30 | Insert(ctx context.Context, collectionName string, data []Record) error 31 | 32 | // Flush ensures all pending writes are committed to storage. 33 | Flush(ctx context.Context, collectionName string) error 34 | 35 | // CreateIndex builds an index on the specified field to optimize search operations. 36 | CreateIndex(ctx context.Context, collectionName, field string, index Index) error 37 | 38 | // LoadCollection loads a collection into memory for faster access. 39 | LoadCollection(ctx context.Context, name string) error 40 | 41 | // Search performs a vector similarity search in the specified collection. 42 | Search(ctx context.Context, collectionName string, vectors map[string]Vector, topK int, metricType string, searchParams map[string]interface{}) ([]SearchResult, error) 43 | 44 | // HybridSearch combines vector similarity search with additional filtering or reranking. 45 | HybridSearch(ctx context.Context, collectionName string, vectors map[string]Vector, topK int, metricType string, searchParams map[string]interface{}, reranker interface{}) ([]SearchResult, error) 46 | 47 | // SetColumnNames configures the column names for the database operations. 48 | SetColumnNames(names []string) 49 | } 50 | 51 | // SearchParam defines the parameters for vector similarity search operations. 52 | type SearchParam struct { 53 | // MetricType specifies the distance metric to use (e.g., "L2", "IP", "COSINE") 54 | MetricType string 55 | // Params contains additional search parameters specific to the database implementation 56 | Params map[string]interface{} 57 | } 58 | 59 | // Schema defines the structure of a collection in the vector database. 60 | type Schema struct { 61 | // Name is the identifier for the schema 62 | Name string 63 | // Description provides additional information about the schema 64 | Description string 65 | // Fields defines the structure of the data in the collection 66 | Fields []Field 67 | } 68 | 69 | // Field represents a single field in a schema, which can be a vector or scalar value. 70 | type Field struct { 71 | // Name is the identifier for the field 72 | Name string 73 | // DataType specifies the type of data stored in the field 74 | DataType string 75 | // PrimaryKey indicates if this field is the primary key 76 | PrimaryKey bool 77 | // AutoID indicates if the field value should be automatically generated 78 | AutoID bool 79 | // Dimension specifies the size of the vector (for vector fields) 80 | Dimension int 81 | // MaxLength specifies the maximum length for variable-length fields 82 | MaxLength int 83 | } 84 | 85 | // Record represents a single data entry in the vector database. 86 | type Record struct { 87 | // Fields maps field names to their values 88 | Fields map[string]interface{} 89 | } 90 | 91 | // Vector represents a mathematical vector as a slice of float64 values. 92 | type Vector []float64 93 | 94 | // Index defines the parameters for building an index on a field. 95 | type Index struct { 96 | // Type specifies the type of index to build (e.g., "IVF", "IVFPQ") 97 | Type string 98 | // Metric specifies the distance metric to use for the index 99 | Metric string 100 | // Parameters contains additional index parameters specific to the database implementation 101 | Parameters map[string]interface{} 102 | } 103 | 104 | // SearchResult represents a single result from a vector similarity search. 105 | type SearchResult struct { 106 | // ID is the identifier for the result 107 | ID int64 108 | // Score is the similarity score for the result 109 | Score float64 110 | // Fields contains additional information about the result 111 | Fields map[string]interface{} 112 | } 113 | 114 | // Config defines the configuration for a vector database connection. 115 | type Config struct { 116 | // Type specifies the type of vector database to connect to (e.g., "milvus", "memory") 117 | Type string 118 | // Address specifies the address of the vector database 119 | Address string 120 | // MaxPoolSize specifies the maximum number of connections to the database 121 | MaxPoolSize int 122 | // Timeout specifies the timeout for database operations 123 | Timeout time.Duration 124 | // Parameters contains additional configuration parameters specific to the database implementation 125 | Parameters map[string]interface{} 126 | } 127 | 128 | // Option defines a function that can be used to configure a Config. 129 | type Option func(*Config) 130 | 131 | // SetType sets the type of vector database to connect to. 132 | func (c *Config) SetType(dbType string) *Config { 133 | c.Type = dbType 134 | return c 135 | } 136 | 137 | // SetAddress sets the address of the vector database. 138 | func (c *Config) SetAddress(address string) *Config { 139 | c.Address = address 140 | return c 141 | } 142 | 143 | // SetMaxPoolSize sets the maximum number of connections to the database. 144 | func (c *Config) SetMaxPoolSize(size int) *Config { 145 | c.MaxPoolSize = size 146 | return c 147 | } 148 | 149 | // SetTimeout sets the timeout for database operations. 150 | func (c *Config) SetTimeout(timeout time.Duration) *Config { 151 | c.Timeout = timeout 152 | return c 153 | } 154 | 155 | // NewVectorDB creates a new VectorDB instance based on the provided configuration. 156 | func NewVectorDB(cfg *Config) (VectorDB, error) { 157 | switch cfg.Type { 158 | case "milvus": 159 | return newMilvusDB(cfg) 160 | case "memory": 161 | return newMemoryDB(cfg) 162 | case "chromem": 163 | return newChromemDB(cfg) 164 | default: 165 | return nil, fmt.Errorf("unsupported database type: %s", cfg.Type) 166 | } 167 | } 168 | --------------------------------------------------------------------------------