agents/drug_info_agent.py

# agents/drug_info_agent.py
"""
Drug Information Agent - Handles drug-related queries using Generative AI.
"""
import re

class DrugInfoAgent:
    def __init__(self, gemini_model=None):
        """
        Initializes the agent with the Gemini model.

        Args:
            gemini_model: An instance of the Gemini model client.
        """
        self.model = gemini_model

    def _extract_drug_name(self, query: str) -> str:
        """A simple helper to extract the drug name from the user's query."""
        # Remove common phrases used to request drug information
        patterns = [
            r"tell me about",
            r"what is",
            r"information on",
            r"info on",
            r"about the drug",
            r"about"
        ]
        drug_name = query.lower()
        for p in patterns:
            drug_name = re.sub(p, "", drug_name)
        
        # Clean up any extra whitespace
        return drug_name.strip().title() # Capitalize for better recognition

    def process_query(self, query: str, file_context: str = "", chat_history: list = None):
        """
        Processes a query to retrieve information about a specific drug.

        Args:
            query (str): The user's full query (e.g., "Tell me about Metformin").
            file_context (str): Optional context from uploaded files.
        
        Returns:
            dict: A dictionary containing the response message and agent metadata.
        """
        # Fallback response if the AI model is not configured
        if not self.model:
            return {
                'message': "💊 **Drug Information Agent**\n\nThe pharmacy database is offline! The Gemini API key is missing, so I can't look up drug information. Please configure the API key to enable this feature.",
                'agent_type': 'drug_info',
                'status': 'error_no_api_key'
            }

        drug_name = self._extract_drug_name(query)
        if not drug_name:
            return {
                'message': "Please tell me which drug you want to know about! For example, try 'info on Paracetamol'.",
                'agent_type': 'drug_info',
                'status': 'error_no_topic'
            }

        # Construct a specialized, safety-conscious prompt for the Gemini model
        prompt = f"""
You are a highly knowledgeable and cautious AI Pharmacist Tutor. Your primary role is to provide accurate drug information for a B.Pharmacy student in India for EDUCATIONAL PURPOSES ONLY.

**CRITICAL SAFETY INSTRUCTION:**
START EVERY RESPONSE with the following disclaimer, exactly as written:
"⚠️ **Disclaimer:** This information is for educational purposes ONLY and is not a substitute for professional medical advice. Always consult a qualified healthcare provider."

**Task:**
Provide a structured summary for the drug: **{drug_name}**

**Information to Include:**
1.  **Therapeutic Class:** What family of drugs does it belong to?
2.  **Mechanism of Action (MOA):** How does it work in the body? Explain simply.
3.  **Common Indications:** What is it typically used for?
4.  **Common Side Effects:** List a few of the most common side effects.
5.  **Important Contraindications/Warnings:** Who should not take this drug or be cautious?
6.  **Common Dosage Forms:** What forms is it available in (e.g., Tablets, Syrup, Injection)? DO NOT provide specific dosages like mg or frequency.

**Format:**
Use clear headings (like "🔬 Mechanism of Action") and bullet points for readability. Use relevant emojis.
"""

        try:
            # Generate content using the AI model
            ai_response = self.model.generate_content(prompt, chat_history)
            return {
                'message': ai_response.text,
                'agent_used': 'drug_info',
                'status': 'success'
            }
        except Exception as e:
            print(f"Drug Info Agent Error: {e}")
            return {
                'message': f"I'm sorry, I couldn't access the drug database at the moment. An error occurred: {str(e)}",
                'agent_type': 'drug_info',
                'status': 'error_api_call'
            }