Oliver-ai-bot_2.0/backend/app/tools/agent_builder.py

"""
Agent Builder tools for Personal Assistant conversational agent creation.
Registered for the 'assistant' mode so PA can create/list/edit agents on behalf of the user.
"""
import re
import json
import logging
from uuid import uuid4
from typing import Any, Dict
from sqlalchemy import select

from app.tools.base import BaseTool, ToolContext, ToolResult
from app.models.agent import Agent

logger = logging.getLogger(__name__)

_DEFAULT_PROVIDER = "anthropic"
_DEFAULT_MODEL = "claude-sonnet-4-6"
_DEFAULT_TEMPERATURE = 0.7


class CreateAgentTool(BaseTool):
    """Tool for PA to create a new private agent draft."""

    @property
    def name(self) -> str:
        return "create_agent"

    @property
    def display_name(self) -> str:
        return "Create Agent"

    @property
    def description(self) -> str:
        return (
            "Create a new AI agent with a custom name, purpose, and instructions. "
            "The agent is saved as a private draft and appears in the user's sidebar. "
            "Use this when the user asks to create, build, or set up a new AI agent or bot."
        )

    @property
    def parameters_schema(self) -> Dict[str, Any]:
        return {
            "type": "object",
            "properties": {
                "name": {
                    "type": "string",
                    "description": "The agent's display name (e.g. 'HR Onboarding Bot')",
                },
                "description": {
                    "type": "string",
                    "description": "Short description of what the agent does",
                },
                "system_prompt": {
                    "type": "string",
                    "description": "Detailed instructions that define the agent's behaviour",
                },
                "category": {
                    "type": "string",
                    "enum": ["general", "hr", "finance", "it", "operations", "custom"],
                    "description": "Agent category",
                },
                "enable_rag": {
                    "type": "boolean",
                    "description": "Whether the agent should search the corporate knowledge base",
                    "default": False,
                },
                "welcome_message": {
                    "type": "string",
                    "description": "Optional greeting shown when the agent is first opened",
                },
                "suggested_prompts": {
                    "type": "array",
                    "items": {"type": "string"},
                    "description": "Optional list of example prompts to show the user",
                },
            },
            "required": ["name", "system_prompt"],
        }

    @property
    def category(self) -> str:
        return "agent_management"

    async def execute(self, arguments: Dict[str, Any], context: ToolContext) -> ToolResult:
        name: str = arguments.get("name", "")
        description: str = arguments.get("description", "")
        system_prompt: str = arguments.get("system_prompt", "")
        category: str = arguments.get("category", "general")
        enable_rag: bool = arguments.get("enable_rag", False)
        welcome_message: str = arguments.get("welcome_message", "")
        suggested_prompts = arguments.get("suggested_prompts", [])

        if not name or not system_prompt:
            return ToolResult(success=False, error="Name and system_prompt are required")

        # Generate unique slug
        base_slug = re.sub(r'[^a-z0-9]+', '-', name.lower()).strip('-')[:45]
        slug = base_slug
        counter = 1
        db = context.db_session
        while True:
            existing = await db.execute(select(Agent).where(Agent.slug == slug))
            if not existing.scalar_one_or_none():
                break
            slug = f"{base_slug}-{counter}"
            counter += 1

        agent = Agent(
            id=uuid4(),
            slug=slug,
            name=name,
            description=description or None,
            icon="bot",
            color="primary",
            category=category,
            llm_provider=_DEFAULT_PROVIDER,
            llm_model=_DEFAULT_MODEL,
            temperature=_DEFAULT_TEMPERATURE,
            system_prompt=system_prompt,
            welcome_message=welcome_message or None,
            suggested_prompts=suggested_prompts,
            tool_ids=[],
            enable_rag=enable_rag,
            knowledge_scope={"all": True} if enable_rag else {},
            enable_file_upload=True,
            capabilities=[],
            visibility="private",
            allowed_department_ids=[],
            is_template=False,
            created_by=context.user_id,
            status="draft",
            is_system=False,
        )
        db.add(agent)
        await db.commit()
        await db.refresh(agent)

        return ToolResult(
            success=True,
            data={"slug": agent.slug, "id": str(agent.id), "name": agent.name},
            display=(
                f"✓ Created agent **{agent.name}** (slug: `{agent.slug}`). "
                f"It's saved as a private draft — click it in your sidebar to start chatting, "
                f"or go to Admin → Agents to publish it for your team."
            ),
        )


class ListAgentsTool(BaseTool):
    """Tool for PA to list agents visible to the user."""

    @property
    def name(self) -> str:
        return "list_agents"

    @property
    def display_name(self) -> str:
        return "List Agents"

    @property
    def description(self) -> str:
        return "List all AI agents available to the current user, including their names, slugs, and status."

    @property
    def parameters_schema(self) -> Dict[str, Any]:
        return {
            "type": "object",
            "properties": {
                "status_filter": {
                    "type": "string",
                    "enum": ["all", "active", "draft"],
                    "description": "Filter by agent status. Default: all",
                    "default": "all",
                }
            },
            "required": [],
        }

    @property
    def category(self) -> str:
        return "agent_management"

    async def execute(self, arguments: Dict[str, Any], context: ToolContext) -> ToolResult:
        status_filter = arguments.get("status_filter", "all")
        db = context.db_session

        query = select(Agent).where(
            (Agent.visibility == "public") |
            (Agent.created_by == context.user_id)
        ).order_by(Agent.is_system.desc(), Agent.name)

        if status_filter != "all":
            query = query.where(Agent.status == status_filter)

        result = await db.execute(query)
        agents = result.scalars().all()

        items = [
            {"name": a.name, "slug": a.slug, "status": a.status, "category": a.category,
             "description": a.description or ""}
            for a in agents
        ]

        if not items:
            return ToolResult(success=True, data=[], display="No agents found.")

        lines = [f"- **{a['name']}** (`{a['slug']}`) [{a['status']}] — {a['description']}" for a in items]
        return ToolResult(
            success=True,
            data=items,
            display="**Available agents:**\n" + "\n".join(lines),
        )


# Register tools when this module is imported
def register_agent_builder_tools():
    from app.tools.registry import ToolRegistry
    ToolRegistry.register(CreateAgentTool())
    ToolRegistry.register(ListAgentsTool())