rag-db/retriever.py

"""
RAG Retrieval Utilities for gprMax Documentation
Provides search and retrieval functions for the vector database
"""

import logging
from pathlib import Path
from typing import List, Dict, Any, Optional, Tuple
import json

import chromadb
from dataclasses import dataclass

logger = logging.getLogger(__name__)


@dataclass
class SearchResult:
    """Container for search results"""
    text: str
    score: float
    metadata: Dict[str, Any]
    
    def __str__(self) -> str:
        return f"[Score: {self.score:.3f}] {self.metadata.get('source', 'Unknown')}: {self.text[:100]}..."


# Removed QwenEmbeddingModel class - using ChromaDB's default embedding


class GprMaxRAGRetriever:
    """Retriever for gprMax documentation RAG database"""
    
    def __init__(self, db_path: Path = None):
        if db_path is None:
            db_path = Path(__file__).parent / "chroma_db"
            
        if not db_path.exists():
            raise ValueError(f"Database path {db_path} does not exist. Run generate_db.py first.")
            
        self.db_path = db_path
        
        # Load metadata
        metadata_path = db_path / "metadata.json"
        if metadata_path.exists():
            with open(metadata_path, 'r') as f:
                self.metadata = json.load(f)
        else:
            self.metadata = {}
            
        # Initialize ChromaDB client
        self.client = chromadb.PersistentClient(path=str(db_path))
        
        # Get collection
        self.collection_name = self.metadata.get("collection_name", "gprmax_docs_v1")
        try:
            print(f"[RAG] Loading collection: {self.collection_name}")
            self.collection = self.client.get_collection(self.collection_name)
            doc_count = self.collection.count()
            print(f"[RAG] Loaded collection: {self.collection_name} with {doc_count} documents")
            logger.info(f"Loaded collection: {self.collection_name} with {doc_count} documents")
        except Exception as e:
            print(f"[RAG] ERROR loading collection: {e}")
            raise ValueError(f"Failed to load collection {self.collection_name}: {e}")
        
    def search(
        self,
        query: str,
        k: int = 10,
        threshold: float = 0.0,
        filter_metadata: Optional[Dict[str, Any]] = None
    ) -> List[SearchResult]:
        """
        Search for relevant documents
        
        Args:
            query: Search query text
            k: Number of results to return
            threshold: Minimum similarity score threshold
            filter_metadata: Optional metadata filters
            
        Returns:
            List of SearchResult objects
        """
        # Search in ChromaDB (it will generate embeddings automatically)
        try:
            results = self.collection.query(
                query_texts=[query],  # Use query_texts instead of query_embeddings
                n_results=k,
                where=filter_metadata if filter_metadata else None,
                include=["documents", "metadatas", "distances"]
            )
            logger.info(f"ChromaDB query returned: {len(results.get('documents', [[]])[0]) if results.get('documents') else 0} results")
        except Exception as e:
            logger.error(f"ChromaDB query failed: {e}")
            raise
        
        # Convert to SearchResult objects
        search_results = []
        if results["documents"] and results["documents"][0]:
            for doc, meta, dist in zip(
                results["documents"][0],
                results["metadatas"][0],
                results["distances"][0]
            ):
                # Convert distance to similarity score (1 - normalized_distance)
                score = 1.0 - (dist / 2.0)  # Assuming cosine distance in [-1, 1]
                
                if score >= threshold:
                    search_results.append(SearchResult(
                        text=doc,
                        score=score,
                        metadata=meta
                    ))
                    
        return search_results
    
    def get_context(
        self,
        query: str,
        k: int = 3,
        max_context_length: int = 2000,
        format_as_markdown: bool = True
    ) -> str:
        """
        Get formatted context for a query
        
        Args:
            query: Search query
            k: Number of documents to retrieve
            max_context_length: Maximum total context length
            format_as_markdown: Format output as markdown
            
        Returns:
            Formatted context string
        """
        results = self.search(query, k=k)
        
        if not results:
            return "No relevant documentation found."
            
        context_parts = []
        total_length = 0
        
        for i, result in enumerate(results, 1):
            if total_length >= max_context_length:
                break
                
            # Truncate if needed
            text = result.text
            if total_length + len(text) > max_context_length:
                text = text[:max_context_length - total_length]
                
            if format_as_markdown:
                source = result.metadata.get("source", "Unknown")
                context_parts.append(
                    f"### Document {i} (Source: {source}, Score: {result.score:.3f})\n"
                    f"```\n{text}\n```\n"
                )
            else:
                context_parts.append(text)
                
            total_length += len(text)
            
        return "\n".join(context_parts)
    
    def get_relevant_files(self, query: str, k: int = 5) -> List[str]:
        """Get list of relevant source files for a query"""
        results = self.search(query, k=k)
        
        # Extract unique source files
        sources = set()
        for result in results:
            source = result.metadata.get("source")
            if source:
                sources.add(source)
                
        return sorted(list(sources))
    
    def search_by_file(self, file_pattern: str, k: int = 10) -> List[SearchResult]:
        """Search for documents from specific files"""
        # This would need ChromaDB's where clause with pattern matching
        # For now, we do a broad search and filter
        results = self.collection.get(
            limit=1000,  # Get many results
            include=["documents", "metadatas"]
        )
        
        filtered_results = []
        if results["documents"]:
            for doc, meta in zip(results["documents"], results["metadatas"]):
                source = meta.get("source", "")
                if file_pattern.lower() in source.lower():
                    filtered_results.append(SearchResult(
                        text=doc,
                        score=1.0,  # No score for direct retrieval
                        metadata=meta
                    ))
                    
                if len(filtered_results) >= k:
                    break
                    
        return filtered_results
    
    def get_stats(self) -> Dict[str, Any]:
        """Get database statistics"""
        stats = {
            "total_documents": self.collection.count(),
            "database_path": str(self.db_path),
            "collection_name": self.collection_name,
            "embedding_model": self.metadata.get("embedding_model", "Unknown"),
            "created_at": self.metadata.get("created_at", "Unknown"),
            "chunk_size": self.metadata.get("chunk_size", "Unknown"),
            "chunk_overlap": self.metadata.get("chunk_overlap", "Unknown")
        }
        return stats


def create_retriever(db_path: Optional[Path] = None) -> GprMaxRAGRetriever:
    """Factory function to create a retriever instance"""
    return GprMaxRAGRetriever(db_path=db_path)


if __name__ == "__main__":
    # Example usage
    import sys
    
    if len(sys.argv) > 1:
        query = " ".join(sys.argv[1:])
    else:
        query = "How to create a source in gprMax?"
        
    print(f"Testing retriever with query: '{query}'")
    print("-" * 80)
    
    try:
        retriever = create_retriever()
        
        # Get stats
        stats = retriever.get_stats()
        print(f"Database stats: {stats}")
        print("-" * 80)
        
        # Search
        results = retriever.search(query, k=3)
        print(f"Found {len(results)} results:")
        for i, result in enumerate(results, 1):
            print(f"\n{i}. {result}")
            
        # Get formatted context
        print("\n" + "=" * 80)
        print("Formatted context:")
        print(retriever.get_context(query, k=3))
        
    except Exception as e:
        print(f"Error: {e}")
        sys.exit(1)