src/tools/default_tools/agent_creater_tool.py

from src.manager.agent_manager import AgentManager
from src.manager.config.model_selector import choose_best_model
from src.manager.utils.runtime_selector import detect_runtime_environment
__all__ = ['AgentCreator']

class AgentCreator():
    dependencies = ["ollama==0.4.7",
                    "pydantic==2.11.1",
                    "pydantic_core==2.33.0"]

    inputSchema = {
        "name": "AgentCreator",
        "description": "Creates an AI agent for you. Please make sure to invoke the created agent using the AskAgent tool.",
        "parameters": {
            "type": "object",
            "properties":{
                "agent_name": {
                    "type": "string",
                    "description": "Name of the AI agent that is to be created. This name cannot have spaces or special characters. It should be a single word.",
                },
                "base_model": {
                    "type": "string",
                    "description": "A base model from which the new agent mode is to be created. Available models are: llama3.2, mistral, gemini-2.5-flash-preview-04-17, gemini-2.5-pro-preview-03-25, gemini-2.0-flash, gemini-2.0-flash-lite, gemini-1.5-flash, gemini-1.5-flash-8b, gemini-1.5-pro, and gemini-2.0-flash-live-001"
                },
                "system_prompt": {
                    "type": "string",
                    "description": "This is the system prompt that will be used to create the agent. It should be a string that describes the role of the agent and its capabilities."
                },
                "description": {
                    "type": "string",
                    "description": "Description of the agent. This is a string that describes the agent and its capabilities. It should be a single line description.",
                },
            },
            "required": ["agent_name", "base_model", "system_prompt", "description"],
        },
        "creates": {
            "selector": "base_model",
            "types": {
                "llama3.2":{
                    "description": "3 Billion parameter model",
                    "create_cost": 10,
                    "invoke_cost": 20,
                },
                "mistral":{
                    "description": "7 Billion parameter model",
                    "create_cost": 20,
                    "invoke_cost": 50,
                },
                "gemini-2.5-flash-preview-04-17": {
                    "description": "Adaptive thinking, cost efficiency",
                    "create_cost": 20,
                    "invoke_cost": 50
                },
                "gemini-2.5-pro-preview-03-25": {
                    "description": "Enhanced thinking and reasoning, multimodal understanding, advanced coding, and more",
                    "create_cost": 20,
                    "invoke_cost": 50
                },
                "gemini-2.0-flash": {
                    "description": "Next generation features, speed, thinking, realtime streaming, and multimodal generation",
                    "create_cost": 20,
                    "invoke_cost": 50
                },
                "gemini-2.0-flash-lite": {
                    "description": "Cost efficiency and low latency",
                    "create_cost": 20,
                    "invoke_cost": 50
                },
                "gemini-1.5-flash": {
                    "description": "Fast and versatile performance across a diverse variety of tasks",
                    "create_cost": 20,
                    "invoke_cost": 50
                },
                "gemini-1.5-flash-8b": {
                    "description": "High volume and lower intelligence tasks",
                    "create_cost": 20,
                    "invoke_cost": 50
                },
                "gemini-1.5-pro": {
                    "description": "Complex reasoning tasks requiring more intelligence",
                    "create_cost": 20,
                    "invoke_cost": 50
                },
                # "gemini-embedding-exp": {
                #     "description": "Measuring the relatedness of text strings",
                #     "create_cost": 20,
                #     "invoke_cost": 50
                # },
                # "imagen-3.0-generate-002": {
                #     "description": "Our most advanced image generation model",
                #     "create_cost": 20,
                #     "invoke_cost": 50
                # },
                # "veo-2.0-generate-001": {
                #     "description": "High quality video generation",
                #     "create_cost": 20,
                #     "invoke_cost": 50
                # },
                "gemini-2.0-flash-live-001": {
                    "description": "Low-latency bidirectional voice and video interactions",
                    "create_cost": 20,
                    "invoke_cost": 50
                }
            }
        }
    }


    def run(self, **kwargs):
        print("Running Agent Creator")
        agent_name = kwargs.get("agent_name")
        base_model = kwargs.get("base_model")

        print(f"[DEBUG] Selected Model: {base_model}")

        system_prompt = kwargs.get("system_prompt")
        description = kwargs.get("description")
        create_cost = self.inputSchema["creates"]["types"][base_model]["create_cost"]
        invoke_cost = self.inputSchema["creates"]["types"][base_model]["invoke_cost"]

        agent_manager = AgentManager()
        try:
            _, remaining_budget = agent_manager.create_agent(
                agent_name=agent_name,
                base_model=base_model,
                system_prompt=system_prompt,
                description=description,
                create_cost=create_cost,
                invoke_cost=invoke_cost
            )
        except ValueError as e:
            return {
                "status": "error",
                "message": f"Error occurred: {str(e)}",
                "output": None
            }
        
        return {
            "status": "success",
            "message": "Agent successfully created",
            "remaining_budget": remaining_budget,
        }