Untitled

class AgentModels:
    def __init__(self):
        # Dictionary of instantiated agents with default configurations
        self.agents = {
            "Navi": NaviBot(llm=model_categories["m_model"], tools=[rag_tool, trim_tokens_web_search, generate_image, get_todays_date, scrape_webpages]),
            "Support": SupportBot(model=model_categories["m_model"][0], tools=[create_support_request]),
            "Help": HelpBot(llm=model_categories["m_model"], tools=[rag_tool]),
            "InfoGathering": InfoGatheringAgent(team="Planning", system_prompt="Plan the project", tools=[search_tool], model=model_categories["l_model"][0]),
            "Checkpointer": CheckpointerAgent(),
            "Medscribe": MedscribeAgent(model=model_categories["l_model"][0], tools=[rag_tool, get_todays_date]),
            "Analyst": AnalystBot(llm=model_categories["l_model"][0]),
            # You can add more agents here...
        }

        # Dictionary mapping the same keys to class references without instantiation
        self.agent_classes = {
            "Navi": NaviBot,
            "Support": SupportBot,
            "Help": HelpBot,
            "InfoGathering": InfoGatheringAgent,
            "Checkpointer": CheckpointerAgent,
            "Medscribe": MedscribeAgent,
            "Analyst": AnalystBot,
            # Add more agent class references as needed...
        }

    def create_agent(self, agent_type, **kwargs):
        # Dynamically instantiate the agent using the class reference from self.agent_classes
        if agent_type in self.agent_classes:
            return self.agent_classes[agent_type](**kwargs)
        else:
            return "Agent type not implemented"

    def set_agent(self, account_id, tier_choice, agent_type, custom_tools=None):
        if agent_type not in self.agent_classes:
            return "Agent type not implemented"

        model_categories = get_round_robin_llm_deployments_based_on_tier(account_id)
        llm_choice = model_categories[tier_choice]

        # Build the kwargs for the new agent instantiation
        agent_kwargs = {
            "llm": llm_choice,
            "tools": custom_tools  # Assume tools can be customized
        }

        # Instantiate the agent dynamically with custom parameters
        new_agent = self.create_agent(agent_type, **agent_kwargs)
        return new_agent