app/llm/agents/rag.py

from typing import Dict, List
import os

from app.llm.client import llm_client


class RAGAgent:
    """Agent for document-based question answering using RAG."""
    
    def __init__(self):
        """Initialize RAG agent with prompt template."""
        prompt_path = os.path.join(
            os.path.dirname(__file__), 
            "..", 
            "prompts", 
            "rag.txt"
        )
        with open(prompt_path, "r") as f:
            self.system_prompt = f.read()
    
    def answer(self, query: str, context_chunks: List[Dict]) -> Dict[str, any]:
        """
        Generate answer based on retrieved document chunks.
        
        Args:
            query: User query string
            context_chunks: List of retrieved document chunks with metadata
            
        Returns:
            Dictionary with 'answer', 'sources', and 'agent' keys
        """
        try:
            # Format context for LLM
            context = self._format_context(context_chunks)
            
            messages = [
                {"role": "system", "content": self.system_prompt},
                {"role": "user", "content": f"Context:\n{context}\n\nQuery: {query}"}
            ]
            
            answer = llm_client.get_completion(
                messages=messages,
                temperature=0.3,  # Lower temperature for factual accuracy
                max_tokens=1024
            )
            
            # Extract unique sources
            sources = list({
                chunk.get("metadata", {}).get("filename", "Unknown")
                for chunk in context_chunks
            })
            
            return {
                "answer": answer,
                "sources": sources,
                "agent": "rag"
            }
            
        except Exception as e:
            print(f"RAG agent error: {e}")
            return {
                "answer": "I encountered an error while processing your question. Please try again.",
                "sources": [],
                "agent": "rag"
            }
    
    def _format_context(self, chunks: List[Dict]) -> str:
        """Format document chunks for LLM context."""
        if not chunks:
            return "No relevant documents found."
        
        formatted = []
        for i, chunk in enumerate(chunks, 1):
            metadata = chunk.get("metadata", {})
            content = chunk.get("content", "")
            filename = metadata.get("filename", "Unknown")
            
            formatted.append(
                f"[Document {i}: {filename}]\n{content}\n"
            )
        
        return "\n---\n".join(formatted)


# Global RAG agent instance
rag_agent = RAGAgent()