agents/router_agent.py

import re
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):
        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)

    def route_query(self, query: str, file_context: str, viva_state: dict, chat_history: list):
        """
        Determines user intent and correctly routes the query with all
        necessary context (file_context, chat_history, etc.) to the
        correct specialist agent.
        """
        query_lower = query.lower()

        # 1. Viva Agent (High priority)
        if viva_state and viva_state.get('active'):
            return self.viva_agent.process_query(query, file_context, viva_state)
        if any(cmd in query_lower for cmd in ["viva", "interview", "start viva"]):
            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, chat_history)

        # 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, chat_history)

        # 4. Drug Info Agent
        if any(cmd in query_lower for cmd in ["drug", "medicine", "medication", "side effect", "dosage"]):
            return self.drug_info_agent.process_query(query, file_context, chat_history)

        # 5. Default to Academic Agent
        return self.academic_agent.process_query(query, file_context, chat_history)