insucompass/core/agents/plan_agent.py

import logging
import json
from typing import List, Dict, Any

from langchain_core.documents import Document

from insucompass.services import llm_provider
from insucompass.config import settings
from insucompass.prompts.prompt_loader import load_prompt

# Configure logging
logging.basicConfig(level=settings.LOG_LEVEL, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

llm = llm_provider.get_gemini_llm()

class PlanAgent:
    """
    An agent that analyzes a user profile and retrieved documents to
    recommend the top 3 most suitable health insurance plans.
    """

    def __init__(self):
        """Initializes the PlanAgent."""
        try:
            self.agent_prompt = load_prompt("plan_agent")
            logger.info("PlanAgent initialized successfully.")
        except FileNotFoundError:
            logger.critical("PlanAgent prompt file not found. The agent cannot function.")
            raise

    def generate_recommendations(
        self,
        user_profile: Dict[str, Any],
        documents: List[Document]
    ) -> Dict[str, Any]:
        """
        Generates plan recommendations in a structured JSON format.
        """
        if not documents:
            logger.warning("PlanAgent received no documents. Cannot generate recommendations.")
            return {"recommendations": []}

        profile_str = json.dumps(user_profile, indent=2)
        context_str = "\n\n---\n\n".join([d.page_content for d in documents])
        
        full_prompt = (
            f"{self.agent_prompt}\n\n"
            f"### CONTEXT FOR YOUR RECOMMENDATION\n"
            f"user_profile: {profile_str}\n\n"
            f"retrieved_context:\n{context_str}"
        )
        
        logger.info("Generating plan recommendations with PlanAgent...")
        try:
            response = llm.invoke(full_prompt)
            # Clean the response to ensure it's valid JSON
            response_content = response.content.strip().replace("```json", "").replace("```", "")
            recommendations = json.loads(response_content)
            logger.info(f"Successfully generated structured plan recommendations \n\n {recommendations}.")
            return recommendations
        except json.JSONDecodeError as e:
            logger.error(f"Failed to decode JSON from PlanAgent: {e}\nResponse was: {response.content}")
            return {"recommendations": [], "error": "Failed to generate plan recommendations."}
        except Exception as e:
            logger.error(f"Error during plan recommendation generation: {e}")
            return {"recommendations": [], "error": "An error occurred while generating plans."}

# Singleton instance
planner = PlanAgent()