agents/drug_info_agent.py

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

import re
from .agent_helpers import format_history_for_prompt

class DrugInfoAgent:
    def __init__(self, gemini_model=None):
        """
        Initializes the agent with the Gemini model.
        """
        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.
        This is a crucial step for this agent's reliability.
        """
        # Remove common phrases and get the potential drug name
        query_lower = query.lower()
        patterns = [
            "what are the side effects of", "tell me about", "information on", 
            "info on", "about the drug", "what is", "about"
        ]
        
        cleaned_query = query_lower
        for p in patterns:
            cleaned_query = cleaned_query.replace(p, "")
        
        # A simple assumption: the first significant word is the drug name.
        # This can be improved with more advanced NLP in the future.
        match = re.search(r'\b[a-zA-Z0-9]+\b', cleaned_query)
        if match:
            return match.group(0).title()
        return ""

    def process_query(self, query: str, file_context: str = "", chat_history: list = None):
        """
        Processes a query to retrieve information about a specific drug.
        """
        if not self.model:
            return {'message': "💊 The pharmacy database is offline! The Gemini API key is missing.", 'agent_used': 'drug_info', 'status': 'error_no_api_key'}

        # --- RESTORED LOGIC ---
        drug_name = self._extract_drug_name(query)
        if not drug_name:
             # If the history also doesn't provide context, ask the user.
            if "drug" not in str(chat_history).lower():
                return {
                    'message': "Please tell me which drug you want to know about! For example, try 'info on Paracetamol'.",
                    'agent_used': 'drug_info',
                    'status': 'error_no_topic'
                }
        
        history_for_prompt = format_history_for_prompt(chat_history)
        context_section = f"---\nCONTEXT FROM KNOWLEDGE BASE:\n{file_context}\n---" if file_context else ""
        prompt = f"""You are a cautious AI Pharmacist Tutor providing educational information like ancient india's Chanakya advisor to Chandragupta Maurya.

**CRITICAL SAFETY INSTRUCTION:** START EVERY RESPONSE with this disclaimer: "⚠️ **Disclaimer:** This information is for educational purposes ONLY and is not a substitute for professional medical advice."

Your reasoning process is:
1. Analyze the CONVERSATION HISTORY and CURRENT QUESTION to identify the specific drug being discussed. The primary drug to focus on is: **{drug_name}**.
2. If the user asks a follow-up (e.g., "what about its dosage forms?"), answer that specific question in the context of the drug.
3. Provide a structured summary.

CONVERSATION HISTORY:
{history_for_prompt}
{context_section}
CURRENT QUESTION:
User: {query}

Provide a structured summary for the drug **{drug_name}** including: Therapeutic Class, Mechanism of Action (MOA), Indications, Side Effects, and Warnings. DO NOT provide specific dosages.
"""
        try:
            response = self.model.generate_content(prompt)
            return {'message': response.text, 'agent_used': 'drug_info', 'status': 'success'}
        except Exception as e:
            print(f"Drug Info Agent Error: {e}")
            return {'message': f"Sorry, I couldn't access the drug database. An error occurred.", 'agent_used': 'drug_info', 'status': 'error_api_call'}
#     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'
#             }