Spaces:
Sleeping
Sleeping
Add RAG pipeline for context-aware agent operations with document retrieval and citation support
536d6c7 verified | """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") | |
| class Document: | |
| """Document container with metadata""" | |
| content: str | |
| source: str | |
| metadata: Dict[str, Any] | |
| embedding: Optional[np.ndarray] = None | |
| 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 | |