agent/a11y_agent.py

"""A11y Expert - Main accessibility question-answering agent."""

from typing import Optional, Generator
from openai import OpenAI
from langdetect import detect, LangDetectException
from config import get_settings
from agent.prompts import get_system_prompt
from agent.tools import search_knowledge_base
from database.vector_store_client import VectorStoreClient
from loguru import logger


class A11yExpertAgent:
    """Accessibility expert agent using OpenAI with RAG."""
    
    def __init__(
        self,
        vector_store: VectorStoreClient,
        llm_client: OpenAI,
        language: str = "en",
        expertise: str = "general"
    ):
        """
        Initialize the A11y Expert agent.
        
        Args:
            vector_store: An instance of VectorStoreClient.
            llm_client: An instance of OpenAI client.
            language: 'pl' for Polish, 'en' for English.
            expertise: 'general', 'wcag', or 'aria'.
        """
        self.vector_store = vector_store
        self.llm_client = llm_client
        self.language = language
        self.expertise = expertise
        self.conversation_history = []
        
        settings = get_settings()
        self.model = settings.llm_model
        self.system_prompt = get_system_prompt(language, expertise)
        
        logger.info(f"A11yExpertAgent initialized (lang={language}, expertise={expertise})")
    
    def close(self):
        """Close agent resources."""
        try:
            if self.vector_store:
                self.vector_store.close()
            if hasattr(self.llm_client, 'close'):
                self.llm_client.close()
            logger.info("A11yExpertAgent resources closed")
        except Exception as e:
            logger.warning(f"Error closing A11yExpertAgent: {e}")
    
    def ask(self, question: str) -> Generator[str, None, None]:
        """
        Ask a question and get a streaming answer with RAG.
        
        Args:
            question: Question about accessibility
            
        Yields:
            Answer chunks from the agent
        """
        logger.info(f"Question: {question}")
        
        try:
            detected_lang = detect(question)
            language = "pl" if detected_lang.startswith("pl") else "en"
        except LangDetectException:
            language = self.language
        
        logger.info(f"Detected language: {language}")
        
        # Dynamically update system prompt based on detected language
        current_system_prompt = get_system_prompt(language, self.expertise)
            
        logger.info("Searching knowledge base...")
        context = search_knowledge_base(question, self.vector_store, language=language)
        
        messages = [
            {"role": "system", "content": current_system_prompt},
            *self.conversation_history[-4:],
            {"role": "user", "content": self._build_prompt_with_context(question, context, language)}
        ]
        
        full_answer = ""
        try:
            response_stream = self.llm_client.chat.completions.create(
                model=self.model,
                messages=messages,
                temperature=0.7,
                max_tokens=1500,
                top_p=0.9,
                stream=True
            )
            
            for chunk in response_stream:
                content = chunk.choices[0].delta.content
                if content:
                    full_answer += content
                    yield content
            
            self.conversation_history.append({"role": "user", "content": question})
            self.conversation_history.append({"role": "assistant", "content": full_answer})
            
            logger.info(f"Answer generated ({len(full_answer)} chars)")
            
        except Exception as e:
            logger.error(f"OpenAI API error: {e}")
            yield f"Error during response generation: {e}"

    def _build_prompt_with_context(self, question: str, context: str, language: str) -> str:
        """Build the prompt with context and language-specific instructions."""
        
        if language == "pl":
            return f"""
Na podstawie poniższego kontekstu z bazy wiedzy o dostępności, odpowiedz na pytanie.

=== KONTEKST Z BAZY WIEDZY ===
{context}

=== PYTANIE ===
{question}

=== ODPOWIEDŹ ===
KRYTYCZNE: Odpowiadaj WYŁĄCZNIE PO POLSKU. To pytanie jest po polsku, więc cała odpowiedź MUSI być po polsku.

Pamiętaj aby:
- Odpowiadać TYLKO po polsku (to jest najważniejsze!)
- Cytować konkretne kryteria i źródła
- Podawać praktyczne przykłady jeśli są istotne
- Być jasnym i zwięzłym
"""
        else:
            return f"""
Based on the following accessibility knowledge base context, answer the question.

=== KNOWLEDGE BASE CONTEXT ===
{context}

=== QUESTION ===
{question}

=== ANSWER ===
CRITICAL: Respond ONLY in ENGLISH. This question is in English, so your entire response MUST be in English.

Remember to:
- Answer ONLY in English (this is most important!)
- Cite specific criteria and sources
- Provide practical examples if relevant
- Be clear and concise
"""
    
    def clear_history(self):
        """Clear conversation history."""
        self.conversation_history = []
        logger.info("Conversation history cleared")
    
    def batch_ask(self, questions: list[str]) -> list[dict]:
        """Ask multiple questions in sequence."""
        results = []
        for question in questions:
            try:
                answer_chunks = [chunk for chunk in self.ask(question)]
                answer = "".join(answer_chunks)
                results.append({"question": question, "answer": answer, "success": True})
            except Exception as e:
                logger.error(f"Failed to answer '{question}': {e}")
                results.append({"question": question, "answer": str(e), "success": False})
        return results


def create_agent(language: Optional[str] = None) -> A11yExpertAgent:
    """Factory function to create and initialize agent."""
    language = language or "en"
    
    logger.info(f"Creating agent with language: {language}")
    settings = get_settings()
    
    # Create vector store with lazy connection (no DB access yet)
    logger.info("Initializing vector store client...")
    vector_store = VectorStoreClient(uri=settings.lancedb_uri)
    
    api_key = settings.openai_api_key
    
    logger.info("Initializing OpenAI client...")
    client_args = {"api_key": api_key}
    if settings.llm_base_url:
        client_args["base_url"] = settings.llm_base_url
        
    llm_client = OpenAI(**client_args)
    
    logger.info("Creating A11yExpertAgent instance...")
    agent = A11yExpertAgent(
        vector_store=vector_store,
        llm_client=llm_client,
        language=language
    )
    logger.info("Agent creation complete")
    return agent