Add RAG pipeline for context-aware agent operations with document retrieval and citation support

smolagents/rag_pipeline.py

@@ -0,0 +1,361 @@
+"""RAG Pipeline for Context-Aware Agent Operations
+
+Retrieval-Augmented Generation pipeline that enriches agent responses
+with relevant context from docs/assets/test data corpus.
+"""
+
+import os
+import json
+from pathlib import Path
+from typing import List, Dict, Any, Optional, Tuple
+import numpy as np
+from dataclasses import dataclass
+
+try:
+    from sentence_transformers import SentenceTransformer
+    EMBEDDINGS_AVAILABLE = True
+except ImportError:
+    EMBEDDINGS_AVAILABLE = False
+    print("Warning: sentence-transformers not available. Install with: pip install sentence-transformers")
+
+
+@dataclass
+class Document:
+    """Document container with metadata"""
+    content: str
+    source: str
+    metadata: Dict[str, Any]
+    embedding: Optional[np.ndarray] = None
+
+
+@dataclass
+class RetrievalResult:
+    """Result from RAG retrieval with context and citations"""
+    documents: List[Document]
+    context: str
+    citations: List[str]
+    scores: List[float]
+
+
+class RAGPipeline:
+    """Retrieval-Augmented Generation Pipeline
+    
+    Provides context-aware enrichment for agent operations by:
+    - Indexing documents from docs/assets/test data
+    - Retrieving relevant context for queries
+    - Generating citations and evidence for responses
+    """
+    
+    def __init__(
+        self,
+        corpus_dir: str = "./docs",
+        model_name: str = "all-MiniLM-L6-v2",
+        top_k: int = 5,
+        cache_dir: str = "./rag_cache"
+    ):
+        """Initialize RAG pipeline
+        
+        Args:
+            corpus_dir: Directory containing documents to index
+            model_name: Sentence transformer model for embeddings
+            top_k: Number of top documents to retrieve
+            cache_dir: Directory for caching embeddings
+        """
+        self.corpus_dir = Path(corpus_dir)
+        self.top_k = top_k
+        self.cache_dir = Path(cache_dir)
+        self.cache_dir.mkdir(exist_ok=True)
+        
+        # Initialize embedding model
+        if EMBEDDINGS_AVAILABLE:
+            try:
+                self.model = SentenceTransformer(model_name)
+                self.embedding_dim = self.model.get_sentence_embedding_dimension()
+            except Exception as e:
+                print(f"Error loading embedding model: {e}")
+                self.model = None
+        else:
+            self.model = None
+        
+        # Document store
+        self.documents: List[Document] = []
+        self.embeddings: Optional[np.ndarray] = None
+        
+        # Load and index documents
+        self._load_corpus()
+    
+    def _load_corpus(self):
+        """Load documents from corpus directory"""
+        if not self.corpus_dir.exists():
+            print(f"Warning: Corpus directory {self.corpus_dir} does not exist")
+            return
+        
+        # Load all text-based files
+        extensions = [".txt", ".md", ".json", ".csv", ".py"]
+        
+        for ext in extensions:
+            for file_path in self.corpus_dir.rglob(f"*{ext}"):
+                try:
+                    self._load_document(file_path)
+                except Exception as e:
+                    print(f"Error loading {file_path}: {e}")
+        
+        print(f"Loaded {len(self.documents)} documents from corpus")
+        
+        # Generate embeddings
+        if self.model and self.documents:
+            self._generate_embeddings()
+    
+    def _load_document(self, file_path: Path):
+        """Load a single document"""
+        with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
+            content = f.read()
+        
+        # Skip empty files
+        if not content.strip():
+            return
+        
+        # Create document
+        doc = Document(
+            content=content,
+            source=str(file_path.relative_to(self.corpus_dir)),
+            metadata={
+                "file_type": file_path.suffix,
+                "file_name": file_path.name,
+                "path": str(file_path)
+            }
+        )
+        
+        self.documents.append(doc)
+    
+    def _generate_embeddings(self):
+        """Generate embeddings for all documents"""
+        if not self.model:
+            return
+        
+        # Check cache
+        cache_file = self.cache_dir / "embeddings.npy"
+        cache_meta = self.cache_dir / "embeddings_meta.json"
+        
+        if cache_file.exists() and cache_meta.exists():
+            try:
+                self.embeddings = np.load(cache_file)
+                with open(cache_meta, 'r') as f:
+                    meta = json.load(f)
+                if meta.get("num_docs") == len(self.documents):
+                    print("Loaded embeddings from cache")
+                    return
+            except Exception as e:
+                print(f"Error loading cache: {e}")
+        
+        # Generate new embeddings
+        print("Generating document embeddings...")
+        texts = [doc.content for doc in self.documents]
+        
+        try:
+            self.embeddings = self.model.encode(
+                texts,
+                show_progress_bar=True,
+                convert_to_numpy=True
+            )
+            
+            # Cache embeddings
+            np.save(cache_file, self.embeddings)
+            with open(cache_meta, 'w') as f:
+                json.dump({"num_docs": len(self.documents)}, f)
+            
+            print("Embeddings generated and cached")
+        except Exception as e:
+            print(f"Error generating embeddings: {e}")
+            self.embeddings = None
+    
+    def retrieve(
+        self,
+        query: str,
+        top_k: Optional[int] = None,
+        filter_metadata: Optional[Dict[str, Any]] = None
+    ) -> RetrievalResult:
+        """Retrieve relevant documents for a query
+        
+        Args:
+            query: Search query
+            top_k: Number of documents to retrieve (default: self.top_k)
+            filter_metadata: Optional metadata filters
+            
+        Returns:
+            RetrievalResult with documents, context, and citations
+        """
+        if top_k is None:
+            top_k = self.top_k
+        
+        if not self.documents:
+            return RetrievalResult(
+                documents=[],
+                context="No documents available in corpus.",
+                citations=[],
+                scores=[]
+            )
+        
+        # Use semantic search if embeddings available
+        if self.model and self.embeddings is not None:
+            return self._semantic_retrieve(query, top_k, filter_metadata)
+        else:
+            # Fallback to keyword search
+            return self._keyword_retrieve(query, top_k, filter_metadata)
+    
+    def _semantic_retrieve(
+        self,
+        query: str,
+        top_k: int,
+        filter_metadata: Optional[Dict[str, Any]]
+    ) -> RetrievalResult:
+        """Semantic retrieval using embeddings"""
+        # Encode query
+        query_embedding = self.model.encode([query], convert_to_numpy=True)[0]
+        
+        # Compute similarities
+        similarities = np.dot(self.embeddings, query_embedding)
+        
+        # Apply metadata filters
+        valid_indices = range(len(self.documents))
+        if filter_metadata:
+            valid_indices = [
+                i for i in valid_indices
+                if all(
+                    self.documents[i].metadata.get(k) == v
+                    for k, v in filter_metadata.items()
+                )
+            ]
+        
+        # Get top-k indices
+        if valid_indices:
+            filtered_sims = [(i, similarities[i]) for i in valid_indices]
+            top_indices = sorted(filtered_sims, key=lambda x: x[1], reverse=True)[:top_k]
+            indices, scores = zip(*top_indices) if top_indices else ([], [])
+        else:
+            indices, scores = [], []
+        
+        # Build result
+        retrieved_docs = [self.documents[i] for i in indices]
+        
+        return self._build_result(retrieved_docs, list(scores))
+    
+    def _keyword_retrieve(
+        self,
+        query: str,
+        top_k: int,
+        filter_metadata: Optional[Dict[str, Any]]
+    ) -> RetrievalResult:
+        """Simple keyword-based retrieval fallback"""
+        query_terms = set(query.lower().split())
+        
+        # Score documents by keyword overlap
+        scores = []
+        for doc in self.documents:
+            doc_terms = set(doc.content.lower().split())
+            overlap = len(query_terms & doc_terms)
+            scores.append(overlap)
+        
+        # Apply metadata filters
+        valid_indices = range(len(self.documents))
+        if filter_metadata:
+            valid_indices = [
+                i for i in valid_indices
+                if all(
+                    self.documents[i].metadata.get(k) == v
+                    for k, v in filter_metadata.items()
+                )
+            ]
+        
+        # Get top-k
+        if valid_indices:
+            filtered_scores = [(i, scores[i]) for i in valid_indices]
+            top_indices = sorted(filtered_scores, key=lambda x: x[1], reverse=True)[:top_k]
+            indices, top_scores = zip(*top_indices) if top_indices else ([], [])
+        else:
+            indices, top_scores = [], []
+        
+        retrieved_docs = [self.documents[i] for i in indices]
+        
+        return self._build_result(retrieved_docs, list(top_scores))
+    
+    def _build_result(self, documents: List[Document], scores: List[float]) -> RetrievalResult:
+        """Build retrieval result with context and citations"""
+        # Build context string
+        context_parts = []
+        citations = []
+        
+        for i, (doc, score) in enumerate(zip(documents, scores), 1):
+            # Truncate long documents
+            content_preview = doc.content[:500]
+            if len(doc.content) > 500:
+                content_preview += "..."
+            
+            context_parts.append(
+                f"[{i}] {doc.source}\n{content_preview}\n"
+            )
+            citations.append(f"[{i}] {doc.source} (relevance: {score:.2f})")
+        
+        context = "\n".join(context_parts) if context_parts else "No relevant context found."
+        
+        return RetrievalResult(
+            documents=documents,
+            context=context,
+            citations=citations,
+            scores=scores
+        )
+    
+    def enrich_query(self, query: str, max_context_length: int = 2000) -> Dict[str, Any]:
+        """Enrich a query with retrieved context
+        
+        Args:
+            query: User query
+            max_context_length: Maximum context length to include
+            
+        Returns:
+            Dict with enriched_query, context, and citations
+        """
+        result = self.retrieve(query)
+        
+        # Truncate context if needed
+        context = result.context
+        if len(context) > max_context_length:
+            context = context[:max_context_length] + "..."
+        
+        enriched_query = f"""{query}
+
+Relevant Context:
+{context}
+
+Please use the above context to provide an informed response and cite sources where applicable."""
+        
+        return {
+            "enriched_query": enriched_query,
+            "context": context,
+            "citations": result.citations,
+            "num_sources": len(result.documents)
+        }
+
+
+# Singleton instance
+_rag_pipeline = None
+
+def get_rag_pipeline(
+    corpus_dir: str = "./docs",
+    force_reload: bool = False
+) -> RAGPipeline:
+    """Get or create RAG pipeline singleton
+    
+    Args:
+        corpus_dir: Directory containing documents
+        force_reload: Force reload of the pipeline
+        
+    Returns:
+        RAGPipeline instance
+    """
+    global _rag_pipeline
+    
+    if _rag_pipeline is None or force_reload:
+        _rag_pipeline = RAGPipeline(corpus_dir=corpus_dir)
+    
+    return _rag_pipeline