app/llm/agents/policy.py

from typing import Dict, List
import os

from app.llm.client import llm_client
from app.services.rag_service import rag_service


class PolicyAgent:
    """Agent for construction policy and regulatory queries using official policy documents."""
    
    def __init__(self):
        """Initialize policy agent with prompt template."""
        prompt_path = os.path.join(
            os.path.dirname(__file__), 
            "..", 
            "prompts", 
            "policy.txt"
        )
        with open(prompt_path, "r") as f:
            self.system_prompt = f.read()
    
    def answer(self, query: str, context_chunks: List[Dict] = None) -> Dict[str, any]:
        """
        Generate answer for policy/regulatory queries using official policy documents.
        
        Args:
            query: User query string
            context_chunks: Retrieved policy chunks from RAG search
            
        Returns:
            Dictionary with 'answer', 'agent', and 'sources' keys
        """
        try:
            # If no context provided, search all official policies (retrieve more chunks for better coverage)
            if context_chunks is None:
                context_chunks = rag_service.collection.query(
                    query_embeddings=[rag_service.embedding_generator.generate_embedding(query)],
                    n_results=10,  # Increased from 5 to 10 for better coverage
                    where={"user_id": "official_policies"}
                )
                
                # Format results
                if context_chunks and context_chunks['documents']:
                    context_chunks = [
                        {
                            "content": context_chunks['documents'][0][i],
                            "metadata": context_chunks['metadatas'][0][i]
                        }
                        for i in range(len(context_chunks['documents'][0]))
                    ]
                else:
                    context_chunks = []
            
            # Build context from chunks
            if not context_chunks:
                return {
                    "answer": "I don't have any official policy documents to answer this question. Please ensure policies are uploaded in the Admin Panel.",
                    "agent": "policy",
                    "sources": []
                }
            
            # Format context with clear chunk numbering
            context_sections = []
            for i, chunk in enumerate(context_chunks, 1):
                doc_id = chunk['metadata'].get('document_id', 'unknown')
                filename = chunk['metadata'].get('filename', 'Official Policy')
                chunk_idx = chunk['metadata'].get('chunk_index', '?')
                
                context_sections.append(
                    f"=== EXCERPT {i} ===\n"
                    f"Document: {filename}\n"
                    f"Document ID: {doc_id}\n"
                    f"Section: Chunk {chunk_idx}\n"
                    f"---\n"
                    f"{chunk['content']}\n"
                )
            
            context_text = "\n".join(context_sections)
            
            # Create strict user message
            user_message = f"""DOCUMENT EXCERPTS FROM OFFICIAL POLICY:

{context_text}

========================================
USER QUESTION: {query}
========================================

REMEMBER: 
- Answer using ONLY the excerpts above
- Include clause/section numbers if present in the text
- Quote exact definitions or requirements
- If the answer is not in the excerpts, say "The provided document sections do not contain this information"
- Do NOT use external knowledge from other building codes

Now provide your answer:"""
            
            messages = [
                {"role": "system", "content": self.system_prompt},
                {"role": "user", "content": user_message}
            ]
            
            answer = llm_client.get_completion(
                messages=messages,
                temperature=0.1,  # Very low temperature for maximum accuracy and minimal creativity
                max_tokens=2000
            )
            
            # Extract sources and policy names
            sources = []
            policy_names = set()
            
            # Get policy titles from database
            from app.database.connection import SessionLocal
            from app.database.models import OfficialPolicy
            
            db = SessionLocal()
            try:
                # Collect unique document IDs
                doc_ids = set()
                for chunk in context_chunks:
                    doc_id = chunk["metadata"].get("document_id", "")
                    if doc_id:
                        doc_ids.add(doc_id)
                
                # Fetch policy titles from database
                policy_title_map = {}
                if doc_ids:
                    policies = db.query(OfficialPolicy).filter(
                        OfficialPolicy.id.in_(doc_ids)
                    ).all()
                    policy_title_map = {p.id: p.title for p in policies}
                
                # Build sources and collect policy names
                for chunk in context_chunks:
                    doc_id = chunk["metadata"].get("document_id", "")
                    policy_title = policy_title_map.get(doc_id, chunk["metadata"].get("filename", "Official Policy"))
                    
                    sources.append({
                        "content": chunk["content"][:300] + "...",
                        "document_id": doc_id,
                        "filename": chunk["metadata"].get("filename", "Official Policy"),
                        "title": policy_title,
                        "chunk_index": chunk["metadata"].get("chunk_index", 0)
                    })
                    
                    # Collect unique policy titles (not filenames)
                    if policy_title:
                        policy_names.add(policy_title)
            
            finally:
                db.close()
            
            print(f"[Policy Agent] Returning policy_names: {list(policy_names)}")
            
            return {
                "answer": answer,
                "agent": "policy",
                "sources": sources,
                "policy_names": list(policy_names)  # List of policy titles used
            }
            
        except Exception as e:
            print(f"Policy agent error: {e}")
            return {
                "answer": "I encountered an error while processing your policy question. Please try again.",
                "agent": "policy",
                "sources": []
            }


# Global policy agent instance
policy_agent = PolicyAgent()