Untitled

class AgentModels:
    def __init__(self):
        self.agent_classes = {
                AgentTeamType.Navi.value: NaviBot,
               "checkpointer_agent": CheckpointerAgent,
                'SUPPORT': SupportBot,
                AgentTeamType.Help.value: HelpBot,
                AgentTeamType.Product.value: InfoGatheringAgent,
                AgentTeamType.Planning.value: InfoGatheringAgent,
                AgentTeamType.Design.value: InfoGatheringAgent,
                AgentTeamType.Frontend.value: InfoGatheringAgent,
                AgentTeamType.Backend.value: InfoGatheringAgent,
                AgentTeamType.DevOps.value: InfoGatheringAgent,
                AgentTeamType.QualityAssurance.value: InfoGatheringAgent,
                AgentTeamType.Brainstorm.value: InfoGatheringAgent,
                AgentTeamType.MarketResearch.value: InfoGatheringAgent,
                AgentTeamType.Analyst.value:AnalystBot,
                AgentTeamType.Medscribe.value: MedscribeAgent,
                "RADIOLOGY": NaviBot
        }

        self.required_params = {
                AgentTeamType.Navi.value: ["llm"],
               "checkpointer_agent": [],
                'SUPPORT': ["llm"],
                AgentTeamType.Help.value: ["llm"],
                AgentTeamType.Product.value: ["team_type", "initial_system_prompt", "final_solution_prompt", "llm"],
                AgentTeamType.Planning.value: ["team_type", "initial_system_prompt", "final_solution_prompt", "llm"],
                AgentTeamType.Design.value: ["team_type", "initial_system_prompt", "final_solution_prompt", "llm"],
                AgentTeamType.Frontend.value: ["team_type", "initial_system_prompt", "final_solution_prompt", "llm"],
                AgentTeamType.Backend.value: ["team_type", "initial_system_prompt", "final_solution_prompt", "llm"],
                AgentTeamType.DevOps.value: ["team_type", "initial_system_prompt", "final_solution_prompt", "llm"],
                AgentTeamType.QualityAssurance.value: ["team_type", "initial_system_prompt", "final_solution_prompt", "llm"],
                AgentTeamType.Brainstorm.value: ["team_type", "initial_system_prompt", "final_solution_prompt", "llm"],
                AgentTeamType.MarketResearch.value: ["team_type", "initial_system_prompt", "final_solution_prompt", "llm"],
                AgentTeamType.Analyst.value:["llm"],
                AgentTeamType.Medscribe.value:["llm"],
                "RADIOLOGY": ["llm"]
        }

        # 4) Optional parameters for each agent key.
        self.optional_params = {
                AgentTeamType.Navi.value: ["system_message", "tools"],
               "checkpointer_agent": [],
                'SUPPORT': [ "tools"],
                AgentTeamType.Help.value: ["system_message", "tools"],
                AgentTeamType.Product.value: ["tools"],
                AgentTeamType.Planning.value: ["tools"],
                AgentTeamType.Design.value: ["tools"],
                AgentTeamType.Frontend.value: ["tools"],
                AgentTeamType.Backend.value: ["tools"],
                AgentTeamType.DevOps.value: ["tools"],
                AgentTeamType.QualityAssurance.value: ["tools"],
                AgentTeamType.Brainstorm.value: ["tools"],
                AgentTeamType.MarketResearch.value: ["tools"],
                AgentTeamType.Analyst.value:["system_message", "tools"],
                AgentTeamType.Medscribe.value:["system_message", "tools"],
                "RADIOLOGY": ["system_message", "tools"]
        }

    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.agent_classes:
            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]
        required=self.required_params.get(agent_type,[])
        optional=self.optional_params.get(agent_type,[])
        valid_params=required+optional
        all_inputs={
            
        }
        # NaviBot(llm=model_categories["m_model"],
        #         tools=[rag_tool, trim_tokens_web_search, generate_image, get_todays_date])