beaver_project/app-instance/backend/beaver/team_workflows/mcp_tools.py

"""MCP schema tools for local team workflow graph builders."""

from __future__ import annotations

import json
from typing import Any, Callable

from beaver.coordinator.models import ExecutionGraph
from beaver.tools.base import BaseTool, ToolContext, ToolResult, ToolSpec

from . import agent_rearrange, concurrent, graph, mixture_of_agents, sequential

GraphBuilder = Callable[..., ExecutionGraph]


def create_team_workflow_tools() -> list[BaseTool]:
    return [
        TeamWorkflowSchemaTool(
            name="SequentialWorkflow",
            description=(
                "Build a sequential Beaver team workflow graph. Use this for staged work "
                "where each agent depends on the previous agent's output."
            ),
            input_schema=_sequential_schema(),
            builder=sequential.build_graph,
        ),
        TeamWorkflowSchemaTool(
            name="ConcurrentWorkflow",
            description=(
                "Build a concurrent Beaver team workflow graph. Use this only when agents "
                "can work independently on the same task."
            ),
            input_schema=_concurrent_schema(),
            builder=concurrent.build_graph,
        ),
        TeamWorkflowSchemaTool(
            name="MixtureOfAgents",
            description=(
                "Build a mixture-of-agents Beaver team workflow graph where independent "
                "expert agents feed one aggregator agent."
            ),
            input_schema=_mixture_schema(),
            builder=mixture_of_agents.build_graph,
        ),
        TeamWorkflowSchemaTool(
            name="AgentRearrange",
            description=(
                "Build a Beaver team workflow graph from strict flow syntax. Use '->' for "
                "stage order and ',' for agents in the same parallel stage."
            ),
            input_schema=_agent_rearrange_schema(),
            builder=agent_rearrange.build_graph,
        ),
        TeamWorkflowSchemaTool(
            name="GraphWorkflow",
            description=(
                "Build an explicit Beaver DAG workflow graph. Use this advanced tool only "
                "when the dependency edges must be specified directly."
            ),
            input_schema=_graph_schema(),
            builder=graph.build_graph,
        ),
    ]


class TeamWorkflowSchemaTool(BaseTool):
    def __init__(
        self,
        *,
        name: str,
        description: str,
        input_schema: dict[str, Any],
        builder: GraphBuilder,
    ) -> None:
        self._spec = ToolSpec(
            name=name,
            description=description,
            input_schema=input_schema,
            toolset="team_workflow",
            always_available=False,
            metadata={"category": "team_workflow"},
        )
        self._builder = builder

    @property
    def spec(self) -> ToolSpec:
        return self._spec

    async def invoke(self, arguments: dict[str, Any], context: ToolContext) -> ToolResult:
        del context
        try:
            graph = self._builder(**dict(arguments or {}))
            payload = {
                "success": True,
                "workflow": self.spec.name,
                "graph": _graph_to_dict(graph),
            }
            return ToolResult(
                success=True,
                content=json.dumps(payload, ensure_ascii=False),
                tool_name=self.spec.name,
                raw_output=payload,
            )
        except Exception as exc:
            payload = {"success": False, "workflow": self.spec.name, "error": str(exc)}
            return ToolResult(
                success=False,
                content=json.dumps(payload, ensure_ascii=False),
                tool_name=self.spec.name,
                error=str(exc),
                raw_output=payload,
            )


def _graph_to_dict(graph: ExecutionGraph) -> dict[str, Any]:
    return {
        "strategy": graph.strategy,
        "nodes": [
            {
                "node_id": node.node_id,
                "task": node.task,
                "depends_on": list(node.depends_on),
                "allowed_tool_names": (
                    None if node.allowed_tool_names is None else list(node.allowed_tool_names)
                ),
                "required_evidence": list(node.required_evidence),
                "evidence_contract": dict(node.evidence_contract),
                "validation_rules": list(node.validation_rules),
                "required_for_completion": node.required_for_completion,
                "block_downstream_on_partial": node.block_downstream_on_partial,
                "max_tool_iterations": node.max_tool_iterations,
                "metadata": dict(node.agent.metadata),
            }
            for node in graph.nodes
        ],
    }


def _sequential_schema() -> dict[str, Any]:
    return {
        "type": "object",
        "properties": {
            "task": _task_schema(),
            "agents": _agents_schema(),
        },
        "required": ["task", "agents"],
        "additionalProperties": False,
    }


def _concurrent_schema() -> dict[str, Any]:
    return {
        "type": "object",
        "properties": {
            "task": _task_schema(),
            "agents": _agents_schema(),
        },
        "required": ["task", "agents"],
        "additionalProperties": False,
    }


def _mixture_schema() -> dict[str, Any]:
    return {
        "type": "object",
        "properties": {
            "task": _task_schema(),
            "agents": _agents_schema(description="Expert agents that run independently before aggregation."),
            "aggregator": _agent_schema(description="Aggregator agent that synthesizes expert outputs."),
        },
        "required": ["task", "agents", "aggregator"],
        "additionalProperties": False,
    }


def _agent_rearrange_schema() -> dict[str, Any]:
    return {
        "type": "object",
        "properties": {
            "task": _task_schema(),
            "agents": _agents_schema(),
            "flow": {
                "type": "string",
                "description": "Strict flow syntax, e.g. 'collector -> tactics, players -> synthesizer'.",
            },
        },
        "required": ["task", "agents", "flow"],
        "additionalProperties": False,
    }


def _graph_schema() -> dict[str, Any]:
    return {
        "type": "object",
        "properties": {
            "task": _task_schema(),
            "agents": _agents_schema(),
            "edges": {
                "type": "array",
                "description": "Directed dependency edges as [source_agent, target_agent] pairs.",
                "items": {
                    "type": "array",
                    "minItems": 2,
                    "maxItems": 2,
                    "items": {"type": "string"},
                },
            },
            "output_agent": {
                "type": "string",
                "description": "Final output/synthesis agent. Must be reachable from upstream agents.",
            },
            "allow_disconnected": {
                "type": "boolean",
                "description": "Allow agents that are not connected to output_agent. Defaults to false.",
            },
        },
        "required": ["task", "agents", "edges", "output_agent"],
        "additionalProperties": False,
    }


def _task_schema() -> dict[str, Any]:
    return {
        "type": "string",
        "description": "Overall user task this workflow supports.",
    }


def _agents_schema(*, description: str = "Workflow agents in the order or set used by this workflow.") -> dict[str, Any]:
    return {
        "type": "array",
        "description": description,
        "items": _agent_schema(),
        "minItems": 1,
    }


def _agent_schema(*, description: str = "One workflow agent slot.") -> dict[str, Any]:
    return {
        "type": "object",
        "description": description,
        "properties": {
            "name": {"type": "string"},
            "instruction": {"type": "string"},
            "use_skill": {"type": "string"},
            "skill_query": {"type": "string"},
            "allowed_tool_names": {"type": "array", "items": {"type": "string"}},
            "required_evidence": {"type": "array", "items": {"type": "string"}},
            "evidence_contract": {"type": "object"},
            "validation_rules": {"type": "array", "items": {"type": "string"}},
            "required_for_completion": {"type": "boolean"},
            "block_downstream_on_partial": {"type": "boolean"},
            "max_tool_iterations": {"type": "integer"},
            "constraints": {"type": "array", "items": {"type": "string"}},
            "expected_output": {"type": "string"},
            "input_contract": {"type": "object"},
            "output_contract": {"type": "object"},
        },
        "required": ["name", "instruction"],
        "additionalProperties": False,
    }