Untitled

 avatar
unknown
plain_text
a month ago
9.7 kB
5
Indexable
class AgentModels:
    def __init__(self):
        # 1) Dictionary with instantiated agents (default configuration).
        self.agents = {
            "global": {
                "Navi": NaviBot(llm=model_categories["m_model"], tools=[rag_tool, trim_tokens_web_search, generate_image, get_todays_date, scrape_webpages]),
                "Checkpointer": CheckpointerAgent(),
                "Support": SupportBot(model=model_categories["m_model"][0], tools=[create_support_request]),
                "Help": HelpBot(llm=model_categories["m_model"], tools=[rag_tool]),
                "Product": InfoGatheringAgent(
                    team_type="Product",
                    initial_system_prompt="BRD_Info_Node_Prompt",
                    final_system_prompt="BRD_Spec_Node_Prompt",
                    tools=[search_tool],
                    model=model_categories["l_model"][0]
                ),
                "Planning": InfoGatheringAgent(
                    team_type="Planning",
                    initial_system_prompt="Info_Node_Prompt",
                    final_system_prompt="Solution_Node_Prompt",
                    tools=[search_tool],
                    model=model_categories["l_model"][0]
                ),
                "Design": InfoGatheringAgent(
                    team_type="Design",
                    initial_system_prompt="Info_Node_Prompt",
                    final_system_prompt="Solution_Node_Prompt",
                    tools=[search_tool],
                    model=model_categories["l_model"][0]
                ),
                "Frontend": InfoGatheringAgent(
                    team_type="Frontend",
                    initial_system_prompt="Info_Node_Prompt",
                    final_system_prompt="Solution_Node_Prompt",
                    tools=[search_tool],
                    model=model_categories["l_model"][0]
                ),
                "Backend": InfoGatheringAgent(
                    team_type="Backend",
                    initial_system_prompt="Info_Node_Prompt",
                    final_system_prompt="Solution_Node_Prompt",
                    tools=[search_tool],
                    model=model_categories["l_model"][0]
                ),
                "DevOps": InfoGatheringAgent(
                    team_type="DevOps",
                    initial_system_prompt="Info_Node_Prompt",
                    final_system_prompt="Solution_Node_Prompt",
                    tools=[search_tool],
                    model=model_categories["l_model"][0]
                ),
                "QualityAssurance": InfoGatheringAgent(
                    team_type="QualityAssurance",
                    initial_system_prompt="Info_Node_Prompt",
                    final_system_prompt="Solution_Node_Prompt",
                    tools=[search_tool],
                    model=model_categories["l_model"][0]
                ),
                "Brainstorm": InfoGatheringAgent(
                    team_type="Brainstorm",
                    initial_system_prompt="brainstorm_agent_system_prompt",
                    final_system_prompt="Solution_Node_Prompt",
                    tools=[search_tool],
                    model=model_categories["l_model"][0]
                ),
                "MarketResearch": InfoGatheringAgent(
                    team_type="MarketResearch",
                    initial_system_prompt="market_analyst_prompt",
                    final_system_prompt="Solution_Node_Prompt",
                    tools=[search_tool],
                    model=model_categories["l_model"][0]
                ),
                "Analyst": AnalystBot(llm=model_categories["l_model"][0]),
                "Medscribe": 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])
            }
        }

        # 2) Dictionary mapping the same keys to class references (no parameters passed).
        self.agent_classes = {
            "global": {
                "Navi": NaviBot,
                "Checkpointer": CheckpointerAgent,
                "Support": SupportBot,
                "Help": HelpBot,
                "Product": InfoGatheringAgent,
                "Planning": InfoGatheringAgent,
                "Design": InfoGatheringAgent,
                "Frontend": InfoGatheringAgent,
                "Backend": InfoGatheringAgent,
                "DevOps": InfoGatheringAgent,
                "QualityAssurance": InfoGatheringAgent,
                "Brainstorm": InfoGatheringAgent,
                "MarketResearch": InfoGatheringAgent,
                "Analyst": AnalystBot,
                "Medscribe": MedscribeAgent,
                "Radiology": NaviBot
            }
        }

        # 3) Required parameters for each agent key.
        self.required_params = {
            "global": {
                "Navi": ["llm", "tools"],
                "Checkpointer": [],
                "Support": ["model", "tools"],
                "Help": ["llm", "tools"],
                "Product": ["team_type", "initial_system_prompt", "final_system_prompt", "model", "tools"],
                "Planning": ["team_type", "initial_system_prompt", "final_system_prompt", "model", "tools"],
                "Design": ["team_type", "initial_system_prompt", "final_system_prompt", "model", "tools"],
                "Frontend": ["team_type", "initial_system_prompt", "final_system_prompt", "model", "tools"],
                "Backend": ["team_type", "initial_system_prompt", "final_system_prompt", "model", "tools"],
                "DevOps": ["team_type", "initial_system_prompt", "final_system_prompt", "model", "tools"],
                "QualityAssurance": ["team_type", "initial_system_prompt", "final_system_prompt", "model", "tools"],
                "Brainstorm": ["team_type", "initial_system_prompt", "final_system_prompt", "model", "tools"],
                "MarketResearch": ["team_type", "initial_system_prompt", "final_system_prompt", "model", "tools"],
                "Analyst": ["llm", "tools"],
                "Medscribe": ["model", "tools"],
                "Radiology": ["llm", "tools"]
            }
        }

        # 4) Optional parameters for each agent key.
        self.optional_params = {
            "global": {
                "Navi": ["additional_info"],
                "Checkpointer": ["extra_feature"],
                "Support": ["optional_param"],
                "Help": [],
                "Product": ["extra_data"],
                "Planning": ["extra_data"],
                "Design": ["extra_data"],
                "Frontend": ["extra_data"],
                "Backend": ["extra_data"],
                "DevOps": ["extra_data"],
                "QualityAssurance": ["extra_data"],
                "Brainstorm": ["extra_data"],
                "MarketResearch": ["extra_data"],
                "Analyst": ["analysis_notes"],
                "Medscribe": ["optional_tools"],
                "Radiology": ["scan_type"]
            }
        }

    # 5) set_agent with explicit parameters for all possible agent usage.
    #    In real usage, you might have additional or fewer parameters; adapt as needed.
    def set_agent(self,
                  agent_type,
                  team_type=None,
                  initial_system_prompt=None,
                  final_system_prompt=None,
                  model=None,
                  tools=None,
                  llm=None,
                  additional_info=None,
                  extra_feature=None,
                  optional_param=None,
                  extra_data=None,
                  analysis_notes=None,
                  optional_tools=None,
                  scan_type=None):
        # First check if the agent_type is in agent_classes
        if agent_type not in self.agent_classes["global"]:
            return "Agent type not implemented"

        # Retrieve references
        agent_class = self.agent_classes["global"][agent_type]
        required = self.required_params["global"].get(agent_type, [])
        optional = self.optional_params["global"].get(agent_type, [])
        valid_params = required + optional

        # Gather all possible inputs
        all_inputs = {
            "team_type": team_type,
            "initial_system_prompt": initial_system_prompt,
            "final_system_prompt": final_system_prompt,
            "model": model,
            "tools": tools,
            "llm": llm,
            "additional_info": additional_info,
            "extra_feature": extra_feature,
            "optional_param": optional_param,
            "extra_data": extra_data,
            "analysis_notes": analysis_notes,
            "optional_tools": optional_tools,
            "scan_type": scan_type
        }

        # Filter out None values and only keep parameters valid for the chosen agent type
        filtered_params = {
            key: value for key, value in all_inputs.items()
            if value is not None and key in valid_params
        }

        # Ensure all required parameters are present
        missing_req = [param for param in required if param not in filtered_params]
        if missing_req:
            return f"Missing required parameters for {agent_type}: {', '.join(missing_req)}"

        # Instantiate the agent with the filtered parameters
        agent_instance = agent_class(**filtered_params)
        return agent_instance
Leave a Comment