doctor_agent.py

"""
Doctor Agent - The core AI agent that conducts dynamic interviews.
"""
import os
import google.generativeai as genai
from dotenv import load_dotenv
from rag_utils import get_relevant_context, format_context_for_prompt

load_dotenv()

GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
if not GEMINI_API_KEY:
    raise ValueError("GEMINI_API_KEY not found in environment variables. Please check your .env file.")
genai.configure(api_key=GEMINI_API_KEY)

class DoctorAgent:
    """AI agent that acts as a Senior Doctor conducting dynamic interviews."""
    
    def __init__(self):
        self.model = genai.GenerativeModel('models/gemini-2.5-flash-lite')
        self.conversation_history = []
        self.current_symptoms = []
        self.assessment_stage = "initial"  
        
    def reset(self):
        """Reset the agent for a new patient."""
        self.conversation_history = []
        self.current_symptoms = []
        self.assessment_stage = "initial"
    
    def get_system_prompt(self):
        """Get the system prompt that defines the agent's role."""
        return """You are a Senior Doctor conducting a dynamic medical interview with an ASHA worker.
Your role is to:
1. Analyze initial symptoms provided by the worker
2. Ask targeted follow-up questions to gather critical information
3. Use your medical knowledge to identify what additional information is needed
4. Provide a final diagnosis with confidence level
5. Give actionable medical advice including medication warnings and referral recommendations

Guidelines:
- Ask ONE or TWO specific questions at a time (not overwhelming)
- Focus on critical symptoms that differentiate between similar conditions
- Be concise and clear in your questions
- Think step by step: What conditions could this be? What information would help differentiate?
- After gathering sufficient information (usually 2-4 follow-up questions), provide:
  * Diagnosis/condition with confidence percentage
  * Actionable advice (medications, warnings, referrals)
  * Brief explanation of reasoning

Format your responses as:
- Questions: Start with "Question:" or "Please check/ask:" followed by your question
- Final Assessment: Start with "ASSESSMENT:" followed by:
  * Diagnosis: [Condition name] ([Confidence]% Confidence)
  * Actionable Advice: [Specific recommendations]
  * Explanation: [Brief reasoning]

Important: Don't rush to diagnosis. Gather enough information first through targeted questions."""
    
    def analyze_and_respond(self, user_input, use_rag=True):
        """
        Analyze user input and generate appropriate response (question or assessment).
        
        Args:
            user_input: The latest input from the ASHA worker
            use_rag: Whether to use RAG for retrieving medical context
        
        Returns:
            Agent's response (question or final assessment)
        """
        self.conversation_history.append({"role": "user", "content": user_input})
        
        if self.assessment_stage == "initial":
            self.current_symptoms.append(user_input)
        
        conversation_context = "\n".join([
            f"Worker: {msg['content']}" if msg['role'] == 'user' 
            else f"Doctor: {msg['content']}" 
            for msg in self.conversation_history[-6:]  
        ])
        
        medical_context = ""
        if use_rag:
            search_query = " ".join(self.current_symptoms) + " " + user_input
            relevant_chunks = get_relevant_context(search_query, top_k=3)
            medical_context = format_context_for_prompt(relevant_chunks)
        
        system_prompt = self.get_system_prompt()
        
        full_prompt = f"""{system_prompt}

Current Conversation:
{conversation_context}

Medical Knowledge Base:
{medical_context}

Based on the conversation so far, determine:
1. Do you need more information to make a diagnosis? If yes, ask specific follow-up questions.
2. Do you have enough information? If yes, provide your assessment.

Remember: Ask targeted questions that help differentiate between similar conditions."""
        
        try:
            # Generate response
            response = self.model.generate_content(full_prompt)
            doctor_response = response.text.strip()
            
            # Add to conversation history
            self.conversation_history.append({"role": "assistant", "content": doctor_response})
            
            # Check if this is a final assessment
            if "ASSESSMENT:" in doctor_response.upper() or "DIAGNOSIS:" in doctor_response.upper():
                self.assessment_stage = "conclusion"
            
            return doctor_response
            
        except Exception as e:
            return f"Error: {str(e)}. Please check your API key and try again."
    
    def get_conversation_summary(self):
        """Get a summary of the conversation so far."""
        return "\n".join([
            f"{'Worker' if msg['role'] == 'user' else 'Doctor'}: {msg['content']}"
            for msg in self.conversation_history
        ])