Untitled
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