Spaces:
Sleeping
Sleeping
File size: 4,773 Bytes
116524e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 | from __future__ import annotations
import json
from importlib import import_module
from typing import Any
from ace.integrations.mcp.handlers import MCPHandlers
from ace.integrations.mcp.models import (
AskRequest,
LearnFeedbackRequest,
LearnSampleRequest,
SkillbookGetRequest,
SkillbookLoadRequest,
SkillbookSaveRequest,
)
from ace.integrations.mcp.errors import map_error_to_mcp
_MCP_INSTALL_HINT = (
"ACE MCP support is optional. Install it with "
'`pip install "ace-framework[mcp]"` or `uv add "ace-framework[mcp]"`.'
)
def _load_mcp_types():
try:
return import_module("mcp.types")
except ModuleNotFoundError as exc:
if (exc.name or "").split(".")[0] == "mcp":
raise RuntimeError(_MCP_INSTALL_HINT) from exc
raise
def _mcp_schema(model: Any) -> dict[str, Any]:
"""Return an MCP-friendly JSON schema for a Pydantic model.
Some MCP clients (e.g. the Inspector) don't resolve ``$defs``/``$ref``
correctly and reject valid input. This helper inlines all ``$ref``
pointers so the schema is self-contained.
We keep ``additionalProperties: false`` in the published schema so
clients know that extra fields will be rejected by validation.
"""
schema = model.model_json_schema()
defs = schema.pop("$defs", {})
def _resolve(obj: Any) -> Any:
if isinstance(obj, dict):
if "$ref" in obj:
ref_name = obj["$ref"].rsplit("/", 1)[-1]
return _resolve(defs.get(ref_name, {}))
return {k: _resolve(v) for k, v in obj.items()}
if isinstance(obj, list):
return [_resolve(item) for item in obj]
return obj
return _resolve(schema)
# ── Tool dispatch table ────────────────────────────────────────────
_TOOL_DISPATCH: dict[str, tuple[type, str]] = {
"ace.ask": (AskRequest, "handle_ask"),
"ace.learn.sample": (LearnSampleRequest, "handle_learn_sample"),
"ace.learn.feedback": (LearnFeedbackRequest, "handle_learn_feedback"),
"ace.skillbook.get": (SkillbookGetRequest, "handle_skillbook_get"),
"ace.skillbook.save": (SkillbookSaveRequest, "handle_skillbook_save"),
"ace.skillbook.load": (SkillbookLoadRequest, "handle_skillbook_load"),
}
def register_tools(server: Any, handlers: MCPHandlers) -> None:
types = _load_mcp_types()
@server.list_tools()
async def handle_list_tools():
return [
types.Tool(
name="ace.ask",
description="Ask a question and get a response from ACE.",
inputSchema=_mcp_schema(AskRequest),
),
types.Tool(
name="ace.learn.sample",
description="Provide sample questions/answers for ACE to learn from.",
inputSchema=_mcp_schema(LearnSampleRequest),
),
types.Tool(
name="ace.learn.feedback",
description="Provide feedback on an ACE answer.",
inputSchema=_mcp_schema(LearnFeedbackRequest),
),
types.Tool(
name="ace.skillbook.get",
description="Get statistics and skills from the active skillbook.",
inputSchema=_mcp_schema(SkillbookGetRequest),
),
types.Tool(
name="ace.skillbook.save",
description="Save the active skillbook to disk.",
inputSchema=_mcp_schema(SkillbookSaveRequest),
),
types.Tool(
name="ace.skillbook.load",
description="Load a skillbook from disk into the session.",
inputSchema=_mcp_schema(SkillbookLoadRequest),
),
]
@server.call_tool()
async def handle_call_tool(name: str, arguments: dict | None):
args = arguments or {}
try:
entry = _TOOL_DISPATCH.get(name)
if entry is None:
raise ValueError(f"Unknown tool: {name}")
request_cls, handler_method = entry
req = request_cls(**args)
resp = await getattr(handlers, handler_method)(req)
return types.CallToolResult(
content=[types.TextContent(type="text", text=resp.model_dump_json())],
)
except Exception as e:
mcp_err = map_error_to_mcp(e)
return types.CallToolResult(
isError=True,
content=[types.TextContent(type="text", text=json.dumps(mcp_err))],
)
|