Untitled

class AgentModels:
    def __init__(self):
        self.agents = {"global": {
            AgentTeamType.Navi.value: NaviBot(llm=model_categories["m_model"],
                                              tools=[rag_tool, trim_tokens_web_search, generate_image, get_todays_date,
                                                     scrape_webpages]),
            # 'SUMMARY': UnifiedSummary(llm=model_categories["m_model"],
            #                                   tools=[rag_tool]),
            "checkpointer_agent": CheckpointerAgent(),
            'SUPPORT': SupportBot(model=model_categories["m_model"][0], tools=[create_support_request]),
            AgentTeamType.Help.value: HelpBot(llm=model_categories["m_model"], tools=[rag_tool]),
            AgentTeamType.Product.value: InfoGatheringAgent(AgentTeamType.Product.value,
                                                            agentic_flow_prompts['BRD_Info_Node_Prompt'],
                                                            agentic_flow_prompts['BRD_Spec_Node_Prompt'],
                                                            tools=[search_tool],
                                                            model=model_categories["l_model"][0]),
            AgentTeamType.Planning.value: InfoGatheringAgent(AgentTeamType.Planning.value,
                                                             agentic_flow_prompts['Info_Node_Prompt'],
                                                             agentic_flow_prompts['Solution_Node_Prompt'],
                                                             tools=[search_tool],
                                                             model=model_categories["l_model"][0]),
            AgentTeamType.Design.value: InfoGatheringAgent(AgentTeamType.Design.value,
                                                           agentic_flow_prompts['Info_Node_Prompt'],
                                                           agentic_flow_prompts['Solution_Node_Prompt'],
                                                           tools=[search_tool],
                                                           model=model_categories["l_model"][0]),
            AgentTeamType.Frontend.value: InfoGatheringAgent(AgentTeamType.Frontend.value,
                                                             agentic_flow_prompts['Info_Node_Prompt'],
                                                             agentic_flow_prompts['Solution_Node_Prompt'],
                                                             tools=[search_tool],
                                                             model=model_categories["l_model"][0]),
            AgentTeamType.Backend.value: InfoGatheringAgent(AgentTeamType.Backend.value,
                                                            agentic_flow_prompts['Info_Node_Prompt'],
                                                            agentic_flow_prompts['Solution_Node_Prompt'],
                                                            tools=[search_tool],
                                                            model=model_categories["l_model"][0]),
            AgentTeamType.DevOps.value: InfoGatheringAgent(AgentTeamType.DevOps.value,
                                                           agentic_flow_prompts['Info_Node_Prompt'],
                                                           agentic_flow_prompts['Solution_Node_Prompt'],
                                                           tools=[search_tool],
                                                           model=model_categories["l_model"][0]),
            AgentTeamType.QualityAssurance.value: InfoGatheringAgent(AgentTeamType.QualityAssurance.value,
                                                                     agentic_flow_prompts['Info_Node_Prompt'],
                                                                     agentic_flow_prompts['Solution_Node_Prompt'],
                                                                     tools=[search_tool],
                                                                     model=model_categories["l_model"][0]),
            AgentTeamType.Brainstorm.value: InfoGatheringAgent(AgentTeamType.Brainstorm.value,
                                                               agentic_flow_prompts['brainstorm_agent_system_prompt'],
                                                               agentic_flow_prompts['Solution_Node_Prompt'],
                                                               tools=[search_tool],
                                                               model=model_categories["l_model"][0]),
            AgentTeamType.MarketResearch.value: InfoGatheringAgent(AgentTeamType.MarketResearch.value,
                                                                   agentic_flow_prompts['market_analyst_prompt'],
                                                                   agentic_flow_prompts['Solution_Node_Prompt'],
                                                                   tools=[search_tool],
                                                                   model=model_categories["l_model"][0]),
            # AgentTeamType.Analyst.value: NaviBot(agentic_flow_prompts['analyst_prompt']),
            AgentTeamType.Analyst.value: AnalystBot(llm=model_categories["l_model"][0]),

            AgentTeamType.Medscribe.value: MedscribeAgent(model=model_categories["l_model"][0],
                                                          tools=[rag_tool, get_todays_date]),
            "RADIOLOGY": NaviBot(llm=model_categories["m_model"],
                                 tools=[rag_tool, trim_tokens_web_search, generate_image, get_todays_date])}}

    def set_agent(self, account_id, tier_choice, agent_type, system_prompt, tools):
        # get_model_creds_and_initialize_agent
        if agent_type not in self.agents["global"]:
            return "Agent not implemented"
        # Get the llm model_categories attached to account_id
        model_categories=get_round_robin_llm_deployments_based_on_tier(account_id)
        llm_choice_for_agent_construction=model_categories[tier_choice]
        new_agent_constructed= self.agents["global"][agent_type]