Răsfoiți Sursa

feat: add tool registry management

zhenyu.hu 3 săptămâni în urmă
părinte
comite
43b0cc8cc3

+ 10 - 18
src/agent_lab/application/event_agent.py

@@ -1,13 +1,20 @@
 import json
 from collections.abc import Iterable
+from typing import Any
 
+from agent_lab.application.tools import ToolRegistry, build_default_tool_registry
 from agent_lab.domain.events import ToolCallEvent
 from agent_lab.domain.messages import ChatMessage
 
 
 class EventAgent:
-    def __init__(self, enabled_tools: Iterable[str]) -> None:
+    def __init__(
+        self,
+        enabled_tools: Iterable[str],
+        registry: ToolRegistry | None = None,
+    ) -> None:
         self.enabled_tools = set(enabled_tools)
+        self.registry = registry or build_default_tool_registry()
 
     async def handle(self, event: ToolCallEvent) -> ChatMessage:
         if event.name not in self.enabled_tools:
@@ -19,24 +26,9 @@ class EventAgent:
                 },
             )
 
-        if event.name == "handoff_note":
-            return self._tool_reply(
-                event,
-                {
-                    "tool": "handoff_note",
-                    "message": str(event.arguments.get("message", "")),
-                },
-            )
-
-        return self._tool_reply(
-            event,
-            {
-                "tool": event.name,
-                "error": "unknown tool",
-            },
-        )
+        return self._tool_reply(event, self.registry.handle(event))
 
-    def _tool_reply(self, event: ToolCallEvent, payload: dict[str, str]) -> ChatMessage:
+    def _tool_reply(self, event: ToolCallEvent, payload: dict[str, Any]) -> ChatMessage:
         return ChatMessage(
             role="tool",
             content=json.dumps(payload, ensure_ascii=False),

+ 8 - 23
src/agent_lab/application/runtime.py

@@ -4,6 +4,7 @@ from typing import Any, Protocol
 from agent_lab.application.contracts import AgentParams, DebugRunRequest
 from agent_lab.application.event_agent import EventAgent
 from agent_lab.application.queues import RuntimeQueues
+from agent_lab.application.tools import ToolRegistry, build_default_tool_registry
 from agent_lab.domain.messages import ChatMessage, StreamItem
 
 
@@ -22,16 +23,21 @@ class DebugRuntime:
         self,
         chat_client: ChatClient,
         queues: RuntimeQueues | None = None,
+        registry: ToolRegistry | None = None,
     ) -> None:
         self.chat_client = chat_client
         self.queues = queues
+        self.registry = registry or build_default_tool_registry()
 
     async def run(self, request: DebugRunRequest) -> AsyncIterator[dict[str, Any]]:
         queues = self.queues or RuntimeQueues()
         messages = self._build_initial_messages(request)
         await queues.input.put(ChatMessage(role="user", content=request.user_message))
-        tools = self._build_tools(request.event_agent.enabled_tools)
-        event_agent = EventAgent(request.event_agent.enabled_tools)
+        tools = self.registry.chat_tools(request.event_agent.enabled_tools)
+        event_agent = EventAgent(
+            request.event_agent.enabled_tools,
+            registry=self.registry,
+        )
 
         yield await self._emit(queues, {"type": "session_started"})
 
@@ -128,24 +134,3 @@ class DebugRuntime:
     ) -> dict[str, Any]:
         await queues.output.put(payload)
         return await queues.output.get()
-
-    def _build_tools(self, enabled_tools: list[str]) -> list[dict[str, Any]]:
-        tools: list[dict[str, Any]] = []
-        if "handoff_note" in enabled_tools:
-            tools.append(
-                {
-                    "type": "function",
-                    "function": {
-                        "name": "handoff_note",
-                        "description": "Send a note to the event agent.",
-                        "parameters": {
-                            "type": "object",
-                            "properties": {
-                                "message": {"type": "string"},
-                            },
-                            "required": ["message"],
-                        },
-                    },
-                }
-            )
-        return tools

+ 82 - 0
src/agent_lab/application/tools.py

@@ -0,0 +1,82 @@
+from collections.abc import Callable, Iterable
+from copy import deepcopy
+from dataclasses import dataclass
+from typing import Any
+
+from agent_lab.domain.events import ToolCallEvent
+
+
+ToolHandler = Callable[[ToolCallEvent], dict[str, Any]]
+
+
+@dataclass(frozen=True)
+class ToolDefinition:
+    name: str
+    description: str
+    parameters: dict[str, Any]
+    handler: ToolHandler
+
+
+class ToolRegistry:
+    def __init__(self, definitions: Iterable[ToolDefinition]) -> None:
+        self._definitions = {definition.name: definition for definition in definitions}
+
+    def available_tools(self) -> list[dict[str, Any]]:
+        return [
+            {
+                "name": definition.name,
+                "description": definition.description,
+                "parameters": deepcopy(definition.parameters),
+            }
+            for definition in self._definitions.values()
+        ]
+
+    def chat_tools(self, enabled_names: Iterable[str]) -> list[dict[str, Any]]:
+        enabled = set(enabled_names)
+        return [
+            {
+                "type": "function",
+                "function": {
+                    "name": definition.name,
+                    "description": definition.description,
+                    "parameters": deepcopy(definition.parameters),
+                },
+            }
+            for definition in self._definitions.values()
+            if definition.name in enabled
+        ]
+
+    def handle(self, event: ToolCallEvent) -> dict[str, Any]:
+        definition = self._definitions.get(event.name)
+        if definition is None:
+            return {
+                "tool": event.name,
+                "error": "unknown tool",
+            }
+        return definition.handler(event)
+
+
+def build_default_tool_registry() -> ToolRegistry:
+    return ToolRegistry(
+        [
+            ToolDefinition(
+                name="handoff_note",
+                description="Send a note to the event agent.",
+                parameters={
+                    "type": "object",
+                    "properties": {
+                        "message": {"type": "string"},
+                    },
+                    "required": ["message"],
+                },
+                handler=_handle_handoff_note,
+            )
+        ]
+    )
+
+
+def _handle_handoff_note(event: ToolCallEvent) -> dict[str, Any]:
+    return {
+        "tool": "handoff_note",
+        "message": str(event.arguments.get("message", "")),
+    }

+ 45 - 3
src/agent_lab/presentation/static/app.js

@@ -5,6 +5,7 @@ const userMessage = document.querySelector("#user-message");
 const systemPrompts = document.querySelector("#system-prompts");
 const preMessages = document.querySelector("#pre-messages");
 const preMessageTemplate = document.querySelector("#pre-message-template");
+const toolList = document.querySelector("#tool-list");
 
 let socket = null;
 let activeAssistant = null;
@@ -27,6 +28,43 @@ chatForm.addEventListener("submit", (event) => {
   runDebugSession();
 });
 
+loadTools();
+
+async function loadTools() {
+  try {
+    const response = await fetch("/api/tools");
+    if (!response.ok) {
+      throw new Error(`Tool API failed: ${response.status}`);
+    }
+    renderTools(await response.json());
+  } catch (error) {
+    toolList.textContent = "Failed to load tools";
+  }
+}
+
+function renderTools(tools) {
+  toolList.textContent = "";
+  if (!tools.length) {
+    toolList.textContent = "No tools available";
+    return;
+  }
+
+  tools.forEach((tool) => {
+    const label = document.createElement("label");
+    label.className = "checkbox";
+    label.title = tool.description || tool.name;
+
+    const checkbox = document.createElement("input");
+    checkbox.type = "checkbox";
+    checkbox.className = "tool-toggle";
+    checkbox.value = tool.name;
+    checkbox.checked = true;
+
+    label.append(checkbox, ` ${tool.name}`);
+    toolList.append(label);
+  });
+}
+
 function runDebugSession() {
   closeSocket();
   resetRun();
@@ -67,14 +105,18 @@ function buildRequest() {
       max_tokens: Number(document.querySelector("#max-tokens").value),
     },
     event_agent: {
-      enabled_tools: document.querySelector("#tool-handoff-note").checked
-        ? ["handoff_note"]
-        : [],
+      enabled_tools: selectedTools(),
       max_event_loops: Number(document.querySelector("#max-event-loops").value),
     },
   };
 }
 
+function selectedTools() {
+  return [...toolList.querySelectorAll(".tool-toggle:checked")].map(
+    (input) => input.value,
+  );
+}
+
 function handleServerMessage(message) {
   if (message.type === "session_started") {
     appendLog("session", "Session started");

+ 1 - 4
src/agent_lab/presentation/static/index.html

@@ -42,10 +42,7 @@
         <section>
           <h2>Event Agent</h2>
           <label>Max Loops <input id="max-event-loops" type="number" min="1" step="1" value="3" /></label>
-          <label class="checkbox">
-            <input id="tool-handoff-note" type="checkbox" checked />
-            handoff_note
-          </label>
+          <div id="tool-list" class="stack">Loading tools...</div>
         </section>
       </aside>
 

+ 15 - 3
src/agent_lab/presentation/web.py

@@ -9,6 +9,7 @@ from pydantic import ValidationError
 
 from agent_lab.application.contracts import DebugRunRequest
 from agent_lab.application.runtime import DebugRuntime
+from agent_lab.application.tools import ToolRegistry, build_default_tool_registry
 from agent_lab.infrastructure.chat_client import OpenAICompatibleChatClient
 from agent_lab.settings import Settings
 
@@ -23,7 +24,11 @@ def create_app(
     runtime_factory: RuntimeFactory | None = None,
 ) -> FastAPI:
     resolved_settings = settings or Settings()
-    resolved_runtime_factory = runtime_factory or _runtime_factory(resolved_settings)
+    tool_registry = build_default_tool_registry()
+    resolved_runtime_factory = runtime_factory or _runtime_factory(
+        resolved_settings,
+        tool_registry,
+    )
 
     app = FastAPI(title="Agent Lab")
     app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")
@@ -36,6 +41,10 @@ def create_app(
     async def index() -> FileResponse:
         return FileResponse(STATIC_DIR / "index.html")
 
+    @app.get("/api/tools")
+    async def tools() -> list[dict[str, Any]]:
+        return tool_registry.available_tools()
+
     @app.websocket("/ws/debug")
     async def debug(websocket: WebSocket) -> None:
         await websocket.accept()
@@ -61,7 +70,10 @@ def create_app(
     return app
 
 
-def _runtime_factory(settings: Settings) -> Callable[[], DebugRuntime]:
+def _runtime_factory(
+    settings: Settings,
+    tool_registry: ToolRegistry,
+) -> Callable[[], DebugRuntime]:
     def build_runtime() -> DebugRuntime:
         chat_client = OpenAICompatibleChatClient(
             api_key=settings.openai_api_key,
@@ -70,7 +82,7 @@ def _runtime_factory(settings: Settings) -> Callable[[], DebugRuntime]:
             request_timeout_seconds=settings.request_timeout_seconds,
             include_usage=settings.openai_include_usage,
         )
-        return DebugRuntime(chat_client)
+        return DebugRuntime(chat_client, registry=tool_registry)
 
     return build_runtime
 

+ 66 - 0
tests/test_debug_runtime.py

@@ -7,6 +7,7 @@ import pytest
 
 from agent_lab.application.contracts import AgentParams, DebugRunRequest, EventAgentParams
 from agent_lab.application.runtime import DebugRuntime
+from agent_lab.application.tools import ToolDefinition, ToolRegistry
 from agent_lab.domain.events import ToolCallEvent
 from agent_lab.domain.messages import ChatMessage, StreamItem
 
@@ -96,6 +97,20 @@ class StrictHistoryChatClient:
         yield StreamItem.message_delta("final answer")
 
 
+class ToolCapturingChatClient:
+    def __init__(self) -> None:
+        self.tools: list[dict[str, Any]] = []
+
+    async def stream_chat(
+        self,
+        messages: list[ChatMessage],
+        tools: list[dict],
+        params: AgentParams,
+    ) -> AsyncIterator[StreamItem]:
+        self.tools = list(tools)
+        yield StreamItem.message_delta("final answer")
+
+
 def test_runtime_queues_exposes_input_output_and_events_queues():
     RuntimeQueues = _runtime_queues_class()
 
@@ -257,3 +272,54 @@ async def test_runtime_preserves_assistant_tool_calls_before_tool_reply():
     ]
     assert tool_message.tool_call_id == "call_1"
     assert outputs[-1] == {"type": "done"}
+
+
+@pytest.mark.asyncio
+async def test_runtime_passes_selected_tool_schema_from_registry_to_chat_agent():
+    registry = ToolRegistry(
+        [
+            ToolDefinition(
+                name="handoff_note",
+                description="Registry-owned handoff tool.",
+                parameters={
+                    "type": "object",
+                    "properties": {
+                        "message": {"type": "string"},
+                        "priority": {"type": "number"},
+                    },
+                    "required": ["message"],
+                },
+                handler=lambda event: {"tool": event.name, "message": "handled"},
+            )
+        ]
+    )
+    request = DebugRunRequest(
+        user_message="debug this",
+        system_prompts=[],
+        pre_messages=[],
+        chat_agent=AgentParams(model="fake-model", temperature=0.1, max_tokens=200),
+        event_agent=EventAgentParams(enabled_tools=["handoff_note"], max_event_loops=2),
+    )
+    client = ToolCapturingChatClient()
+    runtime = DebugRuntime(client, registry=registry)
+
+    outputs = [message async for message in runtime.run(request)]
+
+    assert outputs[-1] == {"type": "done"}
+    assert client.tools == [
+        {
+            "type": "function",
+            "function": {
+                "name": "handoff_note",
+                "description": "Registry-owned handoff tool.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "message": {"type": "string"},
+                        "priority": {"type": "number"},
+                    },
+                    "required": ["message"],
+                },
+            },
+        }
+    ]

+ 46 - 0
tests/test_event_agent.py

@@ -3,6 +3,7 @@ import json
 import pytest
 
 from agent_lab.application.event_agent import EventAgent
+from agent_lab.application.tools import ToolDefinition, ToolRegistry
 from agent_lab.domain.events import ToolCallEvent
 
 
@@ -26,3 +27,48 @@ async def test_event_agent_executes_enabled_tool_as_tool_reply_message():
         "message": "inspect this event",
     }
 
+
+@pytest.mark.asyncio
+async def test_event_agent_returns_registry_errors_for_disabled_and_unknown_tools():
+    registry = ToolRegistry(
+        [
+            ToolDefinition(
+                name="handoff_note",
+                description="Send a note to the event agent.",
+                parameters={"type": "object"},
+                handler=lambda event: {"tool": event.name, "message": "handled"},
+            )
+        ]
+    )
+
+    disabled_reply = await EventAgent(
+        enabled_tools=[],
+        registry=registry,
+    ).handle(
+        ToolCallEvent(
+            id="call_1",
+            name="handoff_note",
+            arguments={"message": "inspect this event"},
+            raw_arguments='{"message":"inspect this event"}',
+        )
+    )
+    unknown_reply = await EventAgent(
+        enabled_tools=["missing_tool"],
+        registry=registry,
+    ).handle(
+        ToolCallEvent(
+            id="call_2",
+            name="missing_tool",
+            arguments={},
+            raw_arguments="{}",
+        )
+    )
+
+    assert json.loads(disabled_reply.content) == {
+        "tool": "handoff_note",
+        "error": "tool disabled",
+    }
+    assert json.loads(unknown_reply.content) == {
+        "tool": "missing_tool",
+        "error": "unknown tool",
+    }

+ 33 - 0
tests/test_websocket_api.py

@@ -52,6 +52,28 @@ def test_health_returns_ok():
     assert response.json() == {"status": "ok"}
 
 
+def test_api_tools_returns_handoff_note_metadata():
+    app = create_app(runtime_factory=FakeRuntime)
+    client = TestClient(app)
+
+    response = client.get("/api/tools")
+
+    assert response.status_code == 200
+    assert response.json() == [
+        {
+            "name": "handoff_note",
+            "description": "Send a note to the event agent.",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "message": {"type": "string"},
+                },
+                "required": ["message"],
+            },
+        }
+    ]
+
+
 def test_websocket_debug_streams_runtime_messages():
     runtime = FakeRuntime()
     app = create_app(runtime_factory=lambda: runtime)
@@ -329,6 +351,17 @@ def test_static_pre_message_role_selector_does_not_offer_tool():
     assert '<option value="tool">tool</option>' not in html
 
 
+def test_static_tools_ui_is_dynamic_and_not_hardcoded_to_handoff_note_checkbox():
+    html = Path("src/agent_lab/presentation/static/index.html").read_text()
+    js = Path("src/agent_lab/presentation/static/app.js").read_text()
+
+    assert 'id="tool-handoff-note"' not in html
+    assert "#tool-handoff-note" not in js
+    assert 'id="tool-list"' in html
+    assert '"/api/tools"' in js
+    assert "fetch" in js
+
+
 @pytest.mark.asyncio
 async def test_openai_chat_client_omits_stream_options_when_usage_disabled():
     captured_payloads: list[dict] = []