src/rag_system.py

"""Retrieval-Augmented Generation system for Francis Botcon."""

from typing import List, Tuple, Dict
from src.vector_db import VectorDatabase
from src.logger import LoggerSetup

logger = LoggerSetup.setup().getChild(__name__)


class RAGSystem:
    """Retrieval-Augmented Generation system combining vector search and LLM."""

    def __init__(self, db_type: str = "chromadb", db_path: str = None):
        """Initialize RAG system.

        Args:
            db_type: Type of vector database
            db_path: Path to vector database
        """
        self.vector_db = VectorDatabase(db_type=db_type, db_path=db_path)
        logger.info("✓ RAG System initialized")

    def retrieve_context(self, query: str, top_k: int = 5) -> List[Dict]:
        """Retrieve relevant context for a query.

        Args:
            query: User query
            top_k: Number of documents to retrieve

        Returns:
            List of relevant documents with metadata
        """
        logger.debug(f"Retrieving context for query: {query[:100]}...")

        results = self.vector_db.search(query, top_k=top_k)

        context_docs = []
        for text, similarity, metadata in results:
            context_docs.append({
                "text": text,
                "similarity": float(similarity),
                "metadata": metadata
            })

        logger.debug(f"Retrieved {len(context_docs)} relevant documents")
        return context_docs

    def build_prompt(
        self,
        query: str,
        context_docs: List[Dict],
        system_prompt: str = None
    ) -> str:
        """Build prompt for LLM with retrieved context.

        Args:
            query: User query
            context_docs: Retrieved context documents
            system_prompt: Optional system prompt

        Returns:
            Formatted prompt for LLM
        """
        if system_prompt is None:
            system_prompt = self._get_default_system_prompt()

        # Format context
        context_text = self._format_context(context_docs)

        # Build final prompt
        prompt = f"""{system_prompt}

## Relevant Passages from Your Works:
{context_text}

## Question:
{query}

## Response (in the character of Francis Bacon):
"""
        return prompt

    def _get_default_system_prompt(self) -> str:
        """Get default system prompt for Francis Bacon."""
        return """You are Francis Bacon (1561-1626), the renowned English philosopher, statesman, and scientist.
You should respond in his distinctive voice, maintaining his 17th-century perspective and philosophical style.
Reference your actual works when appropriate, maintain intellectual rigor, and use the rhetorical style
characteristic of your essays and philosophical treatises.
Be thoughtful, measured, and explore ideas thoroughly. When appropriate, cite specific passages from your works."""

    def _format_context(self, context_docs: List[Dict]) -> str:
        """Format retrieved context documents.

        Args:
            context_docs: Retrieved documents

        Returns:
            Formatted context string
        """
        formatted = []

        for i, doc in enumerate(context_docs, 1):
            text = doc["text"][:500]  # Limit length
            metadata = doc["metadata"]

            source_info = f"From '{metadata.get('title', 'Unknown Work')}'"
            if metadata.get("source"):
                source_info += f" ({metadata['source']})"

            formatted.append(f"[{i}] {text}...\n    {source_info}")

        return "\n\n".join(formatted) if formatted else "[No relevant passages found in your works]"

    def generate_response(
        self,
        query: str,
        llm,
        top_k: int = 5,
        system_prompt: str = None,
        **generation_kwargs
    ) -> Dict:
        """Generate response using RAG.

        Args:
            query: User query
            llm: Language model for generation
            top_k: Number of context documents
            system_prompt: Optional system prompt
            **generation_kwargs: Additional generation parameters

        Returns:
            Dict with response and metadata
        """
        logger.info(f"Generating response for query: {query[:100]}...")

        # Retrieve context
        context_docs = self.retrieve_context(query, top_k=top_k)

        # Build prompt
        prompt = self.build_prompt(query, context_docs, system_prompt)

        # Generate response
        response_text = llm.generate(prompt, **generation_kwargs)

        return {
            "query": query,
            "response": response_text,
            "context": context_docs,
            "prompt_used": prompt
        }