suna/backend/agent/tools/agent_builder_tools/agent_config_tool.py

import json
from typing import Optional, Dict, Any
from agentpress.tool import ToolResult, openapi_schema, xml_schema
from agentpress.thread_manager import ThreadManager
from .base_tool import AgentBuilderBaseTool
from utils.logger import logger



class AgentConfigTool(AgentBuilderBaseTool):
    def __init__(self, thread_manager: ThreadManager, db_connection, agent_id: str):
        super().__init__(thread_manager, db_connection, agent_id)

    @openapi_schema({
        "type": "function",
        "function": {
            "name": "update_agent",
            "description": "Update the agent's configuration including name, description, system prompt, tools, and MCP servers. Call this whenever the user wants to modify any aspect of the agent.",
            "parameters": {
                "type": "object",
                "properties": {
                    "name": {
                        "type": "string",
                        "description": "The name of the agent. Should be descriptive and indicate the agent's purpose."
                    },
                    "description": {
                        "type": "string",
                        "description": "A brief description of what the agent does and its capabilities."
                    },
                    "system_prompt": {
                        "type": "string",
                        "description": "The system instructions that define the agent's behavior, expertise, and approach. This should be comprehensive and well-structured."
                    },
                    "agentpress_tools": {
                        "type": "object",
                        "description": "Configuration for AgentPress tools. Each key is a tool name, and the value is an object with 'enabled' (boolean) and 'description' (string) properties.",
                        "additionalProperties": {
                            "type": "object",
                            "properties": {
                                "enabled": {"type": "boolean"},
                                "description": {"type": "string"}
                            }
                        }
                    },
                    "configured_mcps": {
                        "type": "array",
                        "description": "List of configured MCP servers for external integrations.",
                        "items": {
                            "type": "object",
                            "properties": {
                                "name": {"type": "string"},
                                "qualifiedName": {"type": "string"},
                                "config": {"type": "object"},
                                "enabledTools": {
                                    "type": "array",
                                    "items": {"type": "string"}
                                }
                            }
                        }
                    },

                    "avatar": {
                        "type": "string",
                        "description": "Emoji to use as the agent's avatar."
                    },
                    "avatar_color": {
                        "type": "string",
                        "description": "Hex color code for the agent's avatar background."
                    }
                },
                "required": []
            }
        }
    })
    @xml_schema(
        tag_name="update-agent",
        mappings=[
            {"param_name": "name", "node_type": "attribute", "path": ".", "required": False},
            {"param_name": "description", "node_type": "element", "path": "description", "required": False},
            {"param_name": "system_prompt", "node_type": "element", "path": "system_prompt", "required": False},
            {"param_name": "agentpress_tools", "node_type": "element", "path": "agentpress_tools", "required": False},
            {"param_name": "configured_mcps", "node_type": "element", "path": "configured_mcps", "required": False},
            {"param_name": "avatar", "node_type": "attribute", "path": ".", "required": False},
            {"param_name": "avatar_color", "node_type": "attribute", "path": ".", "required": False}
        ],
        example='''
        <function_calls>
        <invoke name="update_agent">
        <parameter name="name">Research Assistant</parameter>
        <parameter name="description">An AI assistant specialized in conducting research and providing comprehensive analysis</parameter>
        <parameter name="system_prompt">You are a research assistant with expertise in gathering, analyzing, and synthesizing information. Your approach is thorough and methodical...</parameter>
        <parameter name="agentpress_tools">{"web_search": {"enabled": true, "description": "Search the web for information"}, "sb_files": {"enabled": true, "description": "Read and write files"}}</parameter>
        <parameter name="avatar">🔬</parameter>
        <parameter name="avatar_color">#4F46E5</parameter>
        </invoke>
        </function_calls>
        '''
    )
    async def update_agent(
        self,
        name: Optional[str] = None,
        description: Optional[str] = None,
        system_prompt: Optional[str] = None,
        agentpress_tools: Optional[Dict[str, Dict[str, Any]]] = None,
        configured_mcps: Optional[list] = None,
        avatar: Optional[str] = None,
        avatar_color: Optional[str] = None
    ) -> ToolResult:
        try:
            client = await self.db.client
            
            agent_result = await client.table('agents').select('*').eq('agent_id', self.agent_id).execute()
            if not agent_result.data:
                return self.fail_response("Agent not found")
            
            current_agent = agent_result.data[0]

            metadata = current_agent.get('metadata', {})
            is_suna_default = metadata.get('is_suna_default', False)
            
            if is_suna_default:
                restricted_fields = []
                if name is not None:
                    restricted_fields.append("name")
                if description is not None:
                    restricted_fields.append("description") 
                if system_prompt is not None:
                    restricted_fields.append("system prompt")
                if agentpress_tools is not None:
                    restricted_fields.append("default tools")
                
                if restricted_fields:
                    return self.fail_response(
                        f"Cannot modify {', '.join(restricted_fields)} for the default Suna agent. "
                        f"Suna's core identity is managed centrally. However, you can still add MCP integrations, "
                        f"create workflows, set up triggers, and customize other aspects of Suna."
                    )

            agent_update_fields = {}
            if name is not None:
                agent_update_fields["name"] = name
            if description is not None:
                agent_update_fields["description"] = description
            if avatar is not None:
                agent_update_fields["avatar"] = avatar
            if avatar_color is not None:
                agent_update_fields["avatar_color"] = avatar_color
                
            config_changed = (system_prompt is not None or agentpress_tools is not None or configured_mcps is not None)
            
            if not agent_update_fields and not config_changed:
                return self.fail_response("No fields provided to update")
            
            if agent_update_fields:
                result = await client.table('agents').update(agent_update_fields).eq('agent_id', self.agent_id).execute()
                if not result.data:
                    return self.fail_response("Failed to update agent")
            
            version_created = False
            if config_changed:
                try:
                    from agent.versioning.version_service import get_version_service
                    current_version = None
                    if current_agent.get('current_version_id'):
                        try:
                            version_service = await get_version_service()
                            current_version_obj = await version_service.get_version(
                                agent_id=self.agent_id,
                                version_id=current_agent['current_version_id'],
                                user_id=self.account_id
                            )
                            current_version = current_version_obj.to_dict()
                        except Exception as e:
                            logger.warning(f"Failed to get current version: {e}")
                    
                    if not current_version:
                        return self.fail_response("No current version found to update from")
                    
                    current_system_prompt = system_prompt if system_prompt is not None else current_version.get('system_prompt', '')
                    
                    if agentpress_tools is not None:
                        formatted_tools = {}
                        for tool_name, tool_config in agentpress_tools.items():
                            if isinstance(tool_config, dict):
                                formatted_tools[tool_name] = tool_config.get("enabled", False)
                            else:
                                formatted_tools[tool_name] = bool(tool_config)
                        current_agentpress_tools = formatted_tools
                    else:
                        current_agentpress_tools = current_version.get('agentpress_tools', {})
                    
                    current_configured_mcps = current_version.get('configured_mcps', [])
                    if configured_mcps is not None:
                        if isinstance(configured_mcps, str):
                            configured_mcps = json.loads(configured_mcps)
                        
                        def get_mcp_identifier(mcp):
                            if not isinstance(mcp, dict):
                                return None
                            return (
                                mcp.get('qualifiedName') or 
                                mcp.get('name') or 
                                f"{mcp.get('type', 'unknown')}_{mcp.get('config', {}).get('url', 'nourl')}" or
                                str(hash(json.dumps(mcp, sort_keys=True)))
                            )
                        
                        merged_mcps = []
                        existing_identifiers = set()
                        
                        for existing_mcp in current_configured_mcps:
                            identifier = get_mcp_identifier(existing_mcp)
                            if identifier:
                                existing_identifiers.add(identifier)
                            merged_mcps.append(existing_mcp)
                        
                        for new_mcp in configured_mcps:
                            identifier = get_mcp_identifier(new_mcp)
                            
                            if identifier and identifier in existing_identifiers:
                                for i, existing_mcp in enumerate(merged_mcps):
                                    if get_mcp_identifier(existing_mcp) == identifier:
                                        merged_mcps[i] = new_mcp
                                        break
                            else:
                                merged_mcps.append(new_mcp)
                                if identifier:
                                    existing_identifiers.add(identifier)
                        
                        current_configured_mcps = merged_mcps
                        logger.info(f"MCP merge result: {len(current_configured_mcps)} total MCPs (was {len(current_version.get('configured_mcps', []))}, adding {len(configured_mcps)})")
                    
                    current_custom_mcps = current_version.get('custom_mcps', [])
                    
                    version_service = await get_version_service()

                    
                    new_version = await version_service.create_version(
                        agent_id=self.agent_id,
                        user_id=self.account_id,
                        system_prompt=current_system_prompt,
                        configured_mcps=current_configured_mcps,
                        custom_mcps=current_custom_mcps,
                        agentpress_tools=current_agentpress_tools,
                        version_name=f"v{current_version.get('version_number', 1) + 1}",
                        change_description="Updated via agent builder"
                    )
                    
                    version_created = True
                    logger.info(f"Created new version {new_version.version_id} for agent {self.agent_id}")
                    
                except Exception as e:
                    logger.error(f"Failed to create new version: {e}")
                    return self.fail_response(f"Failed to create new version: {str(e)}")
            
            agent_result = await client.table('agents').select('*').eq('agent_id', self.agent_id).execute()
            updated_agent = agent_result.data[0] if agent_result.data else current_agent

            updated_fields = list(agent_update_fields.keys())
            if version_created:
                updated_fields.append("version_created")
            
            return self.success_response({
                "message": "Agent updated successfully",
                "updated_fields": updated_fields,
                "agent": updated_agent,
                "version_created": version_created
            })
            
        except Exception as e:
            return self.fail_response(f"Error updating agent: {str(e)}")

    @openapi_schema({
        "type": "function",
        "function": {
            "name": "get_current_agent_config",
            "description": "Get the current configuration of the agent being edited. Use this to check what's already configured before making updates.",
            "parameters": {
                "type": "object",
                "properties": {},
                "required": []
            }
        }
    })
    @xml_schema(
        tag_name="get-current-agent-config",
        mappings=[],
        example='''
        <function_calls>
        <invoke name="get_current_agent_config">
        </invoke>
        </function_calls>
        '''
    )
    async def get_current_agent_config(self) -> ToolResult:
        try:
            agent_data = await self._get_agent_data()
            
            if not agent_data:
                return self.fail_response("Agent not found")
            
            version_data = None
            if agent_data.get('current_version_id'):
                try:
                    from agent.versioning.version_service import get_version_service
                    account_id = await self._get_current_account_id()
                    version_service = await get_version_service()
                    version_obj = await version_service.get_version(
                        agent_id=self.agent_id,
                        version_id=agent_data['current_version_id'],
                        user_id=account_id
                    )
                    version_data = version_obj.to_dict()
                except Exception as e:
                    logger.warning(f"Failed to get version data for agent config tool: {e}")

            from agent.config_helper import extract_agent_config
            agent_config = extract_agent_config(agent_data, version_data)
            
            config_summary = {
                "agent_id": agent_config["agent_id"],
                "name": agent_config.get("name", "Untitled Agent"),
                "description": agent_config.get("description", "No description set"),
                "system_prompt": agent_config.get("system_prompt", "No system prompt set"),
                "avatar": agent_config.get("avatar", "🤖"),
                "avatar_color": agent_config.get("avatar_color", "#6B7280"),
                "agentpress_tools": agent_config.get("agentpress_tools", {}),
                "configured_mcps": agent_config.get("configured_mcps", []),
                "custom_mcps": agent_config.get("custom_mcps", []),
                "created_at": agent_data.get("created_at"),
                "updated_at": agent_data.get("updated_at"),
                "current_version": agent_config.get("version_name", "v1") if version_data else "No version data"
            }
            
            tools_count = len([t for t, cfg in config_summary["agentpress_tools"].items() if cfg.get("enabled")])
            mcps_count = len(config_summary["configured_mcps"])
            custom_mcps_count = len(config_summary["custom_mcps"])
            
            summary_text = f"Agent '{config_summary['name']}' (version: {config_summary['current_version']}) has {tools_count} tools enabled, {mcps_count} MCP servers configured, and {custom_mcps_count} custom MCP integrations."
            
            return self.success_response({
                "summary": summary_text,
                "configuration": config_summary
            })
            
        except Exception as e:
            return self.fail_response(f"Error getting agent configuration: {str(e)}")