agents/router_agent.py

# """
# Router Agent - The Coordinator
# Classifies user queries and routes them to appropriate specialist agents
# Now with Gemini API integration!
# """

# import re
# from .academic_agent import AcademicAgent
# from .drug_info_agent import DrugInfoAgent  
# from .quiz_agent import QuizAgent
# from .mnemonic_agent import MnemonicAgent
# from .viva_agent import VivaAgent

# class RouterAgent:
#     def __init__(self, gemini_model=None):
#         # Store Gemini model
#         self.model = gemini_model
        
#         # Initialize specialist agents with Gemini model
#         self.academic_agent = AcademicAgent(gemini_model)
#         self.drug_info_agent = DrugInfoAgent(gemini_model)
#         self.quiz_agent = QuizAgent(gemini_model)
#         self.mnemonic_agent = MnemonicAgent(gemini_model)
#         self.viva_agent = VivaAgent(gemini_model)
        
#         # Define keywords for each agent type (Free-tier friendly classification)
#         self.agent_keywords = {
#             'drug_info': [
#                 'drug', 'medicine', 'medication', 'side effects', 'dosage', 
#                 'contraindication', 'interaction', 'pharmacology', 'therapeutic',
#                 'adverse', 'mechanism', 'action', 'indication', 'prescription'
#             ],
#             'quiz_generation': [
#                 'quiz', 'test', 'questions', 'mcq', 'multiple choice', 
#                 'flashcard', 'practice', 'exam', 'assessment', 'evaluate'
#             ],
#             'mnemonic_creation': [
#                 'mnemonic', 'remember', 'memory', 'trick', 'acronym',
#                 'rhyme', 'shortcut', 'memorize', 'recall', 'aide'
#             ],
#             'viva_practice': [
#                 'viva', 'oral', 'interview', 'practice session', 'mock',
#                 'question answer', 'preparation', 'rehearse'
#             ]
#         }
    
#     def classify_query_with_ai(self, query):
#         """Use Gemini AI to classify queries more accurately"""
#         if not self.model:
#             return self.classify_query(query)  # Fallback to keyword matching
        
#         try:
#             classification_prompt = f"""
#             You are a query classifier for a pharmacy education AI assistant. 
#             Classify this user query into ONE of these categories:
            
#             1. academic_query - General academic questions about pharmacy, chemistry, biology, mechanisms
#             2. drug_info - Specific questions about drugs, medicines, side effects, dosages, interactions
#             3. quiz_generation - Requests to create quizzes, tests, MCQs, practice questions
#             4. mnemonic_creation - Requests for memory aids, mnemonics, acronyms, memory tricks
#             5. viva_practice - Requests for mock interviews, viva practice, oral exam preparation
            
#             User Query: "{query}"
            
#             Respond with ONLY the category name (e.g., "academic_query")
#             """
            
#             response = self.model.generate_content(classification_prompt)
#             classification = response.text.strip().lower()
            
#             # Validate the classification
#             valid_types = ['academic_query', 'drug_info', 'quiz_generation', 'mnemonic_creation', 'viva_practice']
#             if classification in valid_types:
#                 return classification
#             else:
#                 return 'academic_query'  # Default fallback
                
#         except Exception as e:
#             print(f"AI classification failed: {e}")
#             return self.classify_query(query)  # Fallback to keyword matching
#         """
#         Classify the user query into one of the agent categories
#         Uses keyword matching for free-tier efficiency
#         """
#         query_lower = query.lower()
        
#         # Count keyword matches for each agent type
#         scores = {}
        
#         for agent_type, keywords in self.agent_keywords.items():
#             score = sum(1 for keyword in keywords if keyword in query_lower)
#             scores[agent_type] = score
        
#         # Special pattern matching for better accuracy
#         if re.search(r'\b(what|explain|definition|mechanism|process|how does)\b', query_lower):
#             scores['drug_info'] += 1 if any(drug_word in query_lower for drug_word in ['drug', 'medicine', 'pharmacology']) else 0
#             scores.setdefault('academic_query', 0)
#             scores['academic_query'] += 1
            
#         if re.search(r'\b(create|make|generate|give me)\s+(quiz|questions|mcq)\b', query_lower):
#             scores['quiz_generation'] += 2
            
#         if re.search(r'\b(help.*remember|memory.*trick|mnemonic.*for)\b', query_lower):
#             scores['mnemonic_creation'] += 2
            
#         if re.search(r'\b(practice.*viva|mock.*interview|oral.*exam)\b', query_lower):
#             scores['viva_practice'] += 2
        
#         # Find the highest scoring agent type
#         if max(scores.values()) == 0:
#             return 'academic_query'  # Default to academic for general questions
        
#         return max(scores, key=scores.get)
    
#     def route_query(self, query):
#         """
#         Route the query to the appropriate specialist agent
#         """
#         agent_type = self.classify_query(query)
        
#         try:
#             if agent_type == 'drug_info':
#                 response = self.drug_info_agent.process_query(query)
#             elif agent_type == 'quiz_generation':
#                 response = self.quiz_agent.process_query(query)
#             elif agent_type == 'mnemonic_creation':
#                 response = self.mnemonic_agent.process_query(query)
#             elif agent_type == 'viva_practice':
#                 response = self.viva_agent.process_query(query)
#             else:  # academic_query or default
#                 response = self.academic_agent.process_query(query)
#                 agent_type = 'academic_query'
            
#             # Add metadata to response
#             if isinstance(response, dict):
#                 response['agent_type'] = agent_type
#                 return response
#             else:
#                 return {
#                     'message': response,
#                     'agent_type': agent_type,
#                     'success': True
#                 }
                
#         except Exception as e:
#             return {
#                 'message': f"Router Error: {str(e)}",
#                 'agent_type': 'error',
#                 'success': False
#             }






# # """
# # Router Agent - The Coordinator
# # Classifies user queries and routes them to appropriate specialist agents
# # """

# # import re
# # from .academic_agent import AcademicAgent
# # from .drug_info_agent import DrugInfoAgent  
# # from .quiz_agent import QuizAgent
# # from .mnemonic_agent import MnemonicAgent
# # from .viva_agent import VivaAgent

# # class RouterAgent:
# #     def __init__(self):
# #         # Initialize specialist agents
# #         self.academic_agent = AcademicAgent()
# #         self.drug_info_agent = DrugInfoAgent()
# #         self.quiz_agent = QuizAgent()
# #         self.mnemonic_agent = MnemonicAgent()
# #         self.viva_agent = VivaAgent()
        
# #         # Define keywords for each agent type (Free-tier friendly classification)
# #         self.agent_keywords = {
# #             'drug_info': [
# #                 'drug', 'medicine', 'medication', 'side effects', 'dosage', 
# #                 'contraindication', 'interaction', 'pharmacology', 'therapeutic',
# #                 'adverse', 'mechanism', 'action', 'indication', 'prescription'
# #             ],
# #             'quiz_generation': [
# #                 'quiz', 'test', 'questions', 'mcq', 'multiple choice', 
# #                 'flashcard', 'practice', 'exam', 'assessment', 'evaluate'
# #             ],
# #             'mnemonic_creation': [
# #                 'mnemonic', 'remember', 'memory', 'trick', 'acronym',
# #                 'rhyme', 'shortcut', 'memorize', 'recall', 'aide'
# #             ],
# #             'viva_practice': [
# #                 'viva', 'oral', 'interview', 'practice session', 'mock',
# #                 'question answer', 'preparation', 'rehearse'
# #             ]
# #         }
    
# #     def classify_query(self, query):
# #         """
# #         Classify the user query into one of the agent categories
# #         Uses keyword matching for free-tier efficiency
# #         """
# #         query_lower = query.lower()
        
# #         # Count keyword matches for each agent type
# #         scores = {}
        
# #         for agent_type, keywords in self.agent_keywords.items():
# #             score = sum(1 for keyword in keywords if keyword in query_lower)
# #             scores[agent_type] = score
        
# #         # Special pattern matching for better accuracy
# #         if re.search(r'\b(what|explain|definition|mechanism|process|how does)\b', query_lower):
# #             scores['drug_info'] += 1 if any(drug_word in query_lower for drug_word in ['drug', 'medicine', 'pharmacology']) else 0
# #             scores.setdefault('academic_query', 0)
# #             scores['academic_query'] += 1
            
# #         if re.search(r'\b(create|make|generate|give me)\s+(quiz|questions|mcq)\b', query_lower):
# #             scores['quiz_generation'] += 2
            
# #         if re.search(r'\b(help.*remember|memory.*trick|mnemonic.*for)\b', query_lower):
# #             scores['mnemonic_creation'] += 2
            
# #         if re.search(r'\b(practice.*viva|mock.*interview|oral.*exam)\b', query_lower):
# #             scores['viva_practice'] += 2
        
# #         # Find the highest scoring agent type
# #         if max(scores.values()) == 0:
# #             return 'academic_query'  # Default to academic for general questions
        
# #         return max(scores, key=scores.get)
    
# #     def route_query(self, query):
# #         """
# #         Route the query to the appropriate specialist agent
# #         """
# #         agent_type = self.classify_query(query)
        
# #         try:
# #             if agent_type == 'drug_info':
# #                 response = self.drug_info_agent.process_query(query)
# #             elif agent_type == 'quiz_generation':
# #                 response = self.quiz_agent.process_query(query)
# #             elif agent_type == 'mnemonic_creation':
# #                 response = self.mnemonic_agent.process_query(query)
# #             elif agent_type == 'viva_practice':
# #                 response = self.viva_agent.process_query(query)
# #             else:  # academic_query or default
# #                 response = self.academic_agent.process_query(query)
# #                 agent_type = 'academic_query'
            
# #             # Add metadata to response
# #             if isinstance(response, dict):
# #                 response['agent_type'] = agent_type
# #                 return response
# #             else:
# #                 return {
# #                     'message': response,
# #                     'agent_type': agent_type,
# #                     'success': True
# #                 }
                
# #         except Exception as e:
# #             return {
# #                 'message': f"Router Error: {str(e)}",
# #                 'agent_type': 'error',
# #                 'success': False
# #             }

# agents/router_agent.py
"""
Router Agent - Directs queries to the appropriate specialist agent.
"""
import re

# Import all specialist agents
from .academic_agent import AcademicAgent
from .drug_info_agent import DrugInfoAgent
from .mnemonic_agent import MnemonicAgent
from .quiz_agent import QuizAgent
from .viva_agent import VivaAgent

class RouterAgent:
    def __init__(self, gemini_model=None):
        """
        Initializes the router and all specialist agents, passing the model to them.
        """
        self.model = gemini_model
        # Instantiate all agents
        self.academic_agent = AcademicAgent(gemini_model)
        self.drug_info_agent = DrugInfoAgent(gemini_model)
        self.mnemonic_agent = MnemonicAgent(gemini_model)
        self.quiz_agent = QuizAgent(gemini_model)
        self.viva_agent = VivaAgent(gemini_model)
        self.default_agent = self.academic_agent

    def route_query(self, query: str, file_context: str = "", viva_state: dict = None):
        """
        Determines the user's intent and routes the query to the correct agent.
        
        Args:
            query (str): The user's input query.
            file_context (str): Text content from any uploaded files.
            viva_state (dict): The current state of the viva session.

        Returns:
            dict: The response dictionary from the selected agent.
        """
        query_lower = query.lower()

        # --- Intent Detection Logic ---

        # 1. Viva Agent: High priority to catch session-based commands
        # If a viva session is active, or user wants to start/end one.
        if viva_state and viva_state.get('active'):
            # The VivaAgent itself handles all logic when a session is active
            return self.viva_agent.process_query(query, file_context, viva_state)
        if any(cmd in query_lower for cmd in ["viva", "interview"]):
            return self.viva_agent.process_query(query, file_context, viva_state)

        # 2. Mnemonic Agent
        if any(cmd in query_lower for cmd in ["mnemonic", "memory aid", "remember"]):
            return self.mnemonic_agent.process_query(query, file_context)

        # 3. Quiz Agent
        if any(cmd in query_lower for cmd in ["quiz", "test me", "flashcard"]):
            return self.quiz_agent.process_query(query, file_context)

        # 4. Drug Info Agent
        # Uses keywords and also checks for common drug endings like 'ol', 'in', 'am'
        if any(cmd in query_lower for cmd in ["drug", "medicine", "medication", "side effect", "dosage", "interaction"]):
            return self.drug_info_agent.process_query(query, file_context)
        if re.search(r'\b(paracetamol|ibuprofen|metformin|aspirin|amoxicillin)\b', query_lower):
            return self.drug_info_agent.process_query(query, file_context)

        # 5. Default to Academic Agent
        # If no other intent is detected, it's likely a general academic question.
        return self.academic_agent.process_query(query, file_context)