app/agent.py

import os
import json
import requests
from dotenv import load_dotenv

from app.models import Message, ChatResponse
from app.catalog import Catalog, get_catalog
from app.retrieval import HybridRetriever
from app.analyzer import ConversationAnalyzer
from app.safety import SafetyGuard
from app.prompts import build_system_prompt, build_catalog_context
from app.validator import ResponseValidator

load_dotenv()


class SHLAgent:
    """
    The core SHL Assessment Recommender agent powered by Gemma via NVIDIA NIM API.
    """
    
    def __init__(self):
        # Initialize components
        self.catalog: Catalog = get_catalog()
        self.retriever = HybridRetriever(self.catalog)
        self.analyzer = ConversationAnalyzer()
        self.safety = SafetyGuard()
        self.validator = ResponseValidator(self.catalog)
        
        # Configure NVIDIA API
        self.api_key = os.getenv("NVIDIA_API_KEY", "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")
        self.invoke_url = "https://integrate.api.nvidia.com/v1/chat/completions"
        self.model_name = "google/diffusiongemma-26b-a4b-it"
        
        # Build retrieval index at initialization
        self.retriever.build_index()
        
        print(f"[Agent] Initialized successfully with NVIDIA NIM API using {self.model_name}")
    
    def process(self, messages: list[Message]) -> ChatResponse:
        """
        Process a chat request and return a response.
        """
        try:
            return self._process_internal(messages)
        except Exception as e:
            import traceback
            print(f"[Agent] Top-level error: {e}")
            traceback.print_exc()
            return self.validator.create_safe_response(
                reply=f"Sorry, an internal error occurred: {str(e)}"
            )
    
    def _process_internal(self, messages: list[Message]) -> ChatResponse:
        """Internal processing pipeline."""
        
        # --- Step 1: Safety Check ---
        print("[Agent] Step 1: Safety check...")
        latest_user_msg = self.analyzer.get_latest_user_message(messages)
        print(f"[Agent] Latest user msg: {latest_user_msg[:100] if latest_user_msg else 'EMPTY'}")
        is_safe, refusal_msg = self.safety.check(latest_user_msg)
        if not is_safe:
            print(f"[Agent] Safety refusal: {refusal_msg}")
            return self.validator.create_refusal_response(refusal_msg)
        print("[Agent] Step 1: PASSED")
        
        # --- Step 2: Extract Slots ---
        print("[Agent] Step 2: Extract slots...")
        slots = self.analyzer.extract_slots(messages)
        print(f"[Agent] Step 2: DONE - confidence={slots.confidence}")
        
        # --- Step 3: Classify Intent ---
        print("[Agent] Step 3: Classify intent...")
        intent = self.analyzer.classify_intent(messages)
        print(f"[Agent] Step 3: DONE - intent={intent}")
        
        # --- Step 4: Retrieve Relevant Assessments ---
        print("[Agent] Step 4: Retrieve assessments...")
        search_queries = slots.build_search_queries()
        if latest_user_msg:
            search_queries.append(latest_user_msg)
        print(f"[Agent] Search queries: {search_queries}")
        
        retrieved_items = self.retriever.search_multi_query(
            search_queries, top_k=8
        )
        print(f"[Agent] Step 4: DONE - retrieved {len(retrieved_items)} items")
        
        # --- Step 5: Build Prompt ---
        print("[Agent] Step 5: Build prompt...")
        slots_summary = slots.to_summary()
        catalog_context = build_catalog_context(retrieved_items)
        
        system_prompt = build_system_prompt(
            slots_summary=slots_summary,
            catalog_context=catalog_context,
            intent=intent,
        )
        print(f"[Agent] Step 5: DONE - prompt length={len(system_prompt)} chars")
        
        # --- Step 6: LLM Call ---
        print("[Agent] Step 6: LLM call...")
        raw_response = self._call_llm(messages, system_prompt)
        print(f"[Agent] Step 6: DONE - response length={len(raw_response)} chars")
        print(f"[Agent] Raw LLM response (first 300 chars): {raw_response[:300].encode('ascii', 'replace').decode('ascii')}")
        
        # --- Step 7: Validate and Return ---
        print("[Agent] Step 7: Parse and validate...")
        try:
            parsed = self.validator.parse_llm_output(raw_response)
            response = self.validator.validate_and_fix(parsed)
            print(f"[Agent] Step 7: DONE - reply length={len(response.reply)}, recs={len(response.recommendations)}")
        except Exception as e:
            print(f"[Agent] LLM output parsing failed: {e}")
            # Fallback: use the raw LLM text directly instead of generic error
            response = self.validator.create_safe_response(
                reply=raw_response if raw_response.strip() else None
            )
        
        return response
    
    def _call_llm(self, messages: list[Message], system_prompt: str) -> str:
        """
        Make a single call to diffusiongemma-26b-a4b-it via NVIDIA API.
        """
        try:
            headers = {
                "Authorization": f"Bearer {self.api_key}",
                "Accept": "application/json",
                "Content-Type": "application/json"
            }
            
            # Construct model chat payload
            api_messages = [
                {"role": "system", "content": system_prompt}
            ]
            for msg in messages:
                role = "assistant" if msg.role == "assistant" else "user"
                api_messages.append({"role": role, "content": msg.content})
            
            payload = {
                "model": self.model_name,
                "messages": api_messages,
                "max_tokens": 500,
                "temperature": 0.3,
                "top_p": 0.95,
                "stream": False,
            }
            
            response = requests.post(self.invoke_url, headers=headers, json=payload)
            response.raise_for_status()
            
            result = response.json()
            return result["choices"][0]["message"]["content"]
            
        except Exception as e:
            print(f"[Agent] NVIDIA LLM call failed: {e}")
            raise


# Singleton agent instance
_agent: SHLAgent | None = None


def get_agent() -> SHLAgent:
    """Get or initialize the agent singleton."""
    global _agent
    if _agent is None:
        _agent = SHLAgent()
    return _agent