Эх сурвалжийг харах

feat: add agent config audit and mock tools

zhenyu.hu 3 долоо хоног өмнө
parent
commit
afbf19f335

+ 41 - 0
docs/plans/todo-17-agent-config-audit-tools.md

@@ -0,0 +1,41 @@
+# Todo 17 Agent Config, Audit, and Mock Tools Plan
+
+**Status:** done
+
+## Goal
+
+Align ChatAgent/EventAgent config, add extra provider body defaults, default EventAgent to one round, add mock tools, audit events, and backend logger calls.
+
+## Design
+
+- `AgentParams` becomes the shared model for model, temperature, max tokens, system prompt, and `extra_body`.
+- `EventAgentParams` extends shared agent params with enabled tools and `max_event_loops`, defaulting to one round.
+- `OpenAICompatibleChatClient` merges `extra_body` into the request payload.
+- `ToolRegistry` exposes several mock tools for the UI while still giving ChatAgent name/description-only event tools.
+- `DebugRuntime` emits `audit` output events for key pipeline transitions and logs them through `logging`.
+
+## Files
+
+- Modify: `src/agent_lab/application/contracts.py`
+- Modify: `src/agent_lab/application/runtime.py`
+- Modify: `src/agent_lab/application/tools.py`
+- Modify: `src/agent_lab/infrastructure/chat_client.py`
+- Modify tests in `tests/test_debug_runtime.py`, `tests/test_websocket_api.py`, and `tests/test_ws_smoke.py`.
+
+## Verification
+
+- Focused backend tests prove:
+  - ChatAgent and EventAgent share agent config fields.
+  - `extra_body` is merged into provider payloads.
+  - EventAgent defaults to one round.
+  - `/api/tools` exposes multiple mock tools.
+  - Runtime emits audit events and logger records key events.
+- Full suite passes.
+
+## Result
+
+- Added shared ChatAgent/EventAgent config fields: `system_prompt` and default `extra_body`.
+- Set EventAgent to one event round by default.
+- Added `mock_search` and `mock_ticket` to the default tool registry.
+- Added runtime audit output events and logger entries for session, round, event, and EventAgent transitions.
+- Verified with `uv run pytest` (`42 passed`, one existing Starlette deprecation warning).

+ 3 - 0
docs/plans/todos.md

@@ -43,3 +43,6 @@
 | 14 | done | `docs/plans/todo-14-message-role-validation.md` | Reject provider-invalid API message roles and malformed tool replies before network calls. | `uv run pytest tests/test_websocket_api.py tests/test_event_agent.py`; `uv run pytest`. |
 | 15 | done | `docs/plans/todo-15-static-package-data.md` | Include static UI assets in installed package builds. | `uv run pytest tests/test_packaging.py`; `uv run pytest`; `uv build`; wheel inspection. |
 | 16 | done | `docs/plans/todo-16-tool-error-isolation.md` | Convert EventAgent tool handler exceptions into structured tool-result errors instead of aborting the session. | `uv run pytest tests/test_event_agent.py tests/test_debug_runtime.py`; `uv run pytest`. |
+| 17 | done | `docs/plans/todo-17-agent-config-audit-tools.md` | Align ChatAgent/EventAgent config, add extra_body defaults, default EventAgent to one event round, add mock tools, audit events, and backend logging. | `uv run pytest` passes (`42 passed`, one existing Starlette deprecation warning). |
+| 18 | pending | `docs/plans/todo-18-prompt-modal-tool-ui.md` | Move prompt workspace into a modal and improve visible tool-management UI. | Static UI tests prove modal controls, tool metadata, and audit log handling exist. |
+| 19 | pending | `docs/plans/todo-19-latency-oriented-polish.md` | Do one optimization pass focused on reducing perceived reply wait time and cleaning frontend/backend rough edges. | Full test suite plus focused UI/static checks pass. |

+ 11 - 1
scripts/ws_smoke.py

@@ -1,8 +1,11 @@
 import argparse
 import asyncio
 import json
+from copy import deepcopy
 from typing import Any
 
+from agent_lab.application.contracts import DEFAULT_EXTRA_BODY
+
 
 DEFAULT_URL = "ws://127.0.0.1:8000/ws/debug"
 DEFAULT_MESSAGE = "Debug this agent handoff path."
@@ -25,10 +28,17 @@ def build_payload(
             "model": model,
             "temperature": 0.2,
             "max_tokens": 1024,
+            "system_prompt": "",
+            "extra_body": deepcopy(DEFAULT_EXTRA_BODY),
         },
         "event_agent": {
+            "model": model,
+            "temperature": 0.2,
+            "max_tokens": 1024,
+            "system_prompt": "",
+            "extra_body": deepcopy(DEFAULT_EXTRA_BODY),
             "enabled_tools": ["handoff_note"] if handoff_note_enabled else [],
-            "max_event_loops": 3,
+            "max_event_loops": 1,
         },
     }
 

+ 18 - 4
src/agent_lab/application/contracts.py

@@ -1,21 +1,35 @@
+from copy import deepcopy
+from typing import Any
+
 from pydantic import BaseModel, ConfigDict, Field, model_validator
 
 from agent_lab.domain.messages import ChatMessage
 
 
+DEFAULT_EXTRA_BODY: dict[str, Any] = {
+    "thinking": {"type": "disabled"},
+    "enable_search": False,
+    "search_options": {"forced_search": False},
+}
+
+
+def default_extra_body() -> dict[str, Any]:
+    return deepcopy(DEFAULT_EXTRA_BODY)
+
+
 class AgentParams(BaseModel):
     model_config = ConfigDict(extra="forbid")
 
     model: str | None = None
     temperature: float = 0.2
     max_tokens: int = 1024
+    system_prompt: str = ""
+    extra_body: dict[str, Any] = Field(default_factory=default_extra_body)
 
 
-class EventAgentParams(BaseModel):
-    model_config = ConfigDict(extra="forbid")
-
+class EventAgentParams(AgentParams):
     enabled_tools: list[str] = Field(default_factory=lambda: ["handoff_note"])
-    max_event_loops: int = 3
+    max_event_loops: int = 1
 
 
 class DebugRunRequest(BaseModel):

+ 56 - 1
src/agent_lab/application/runtime.py

@@ -1,4 +1,5 @@
 import asyncio
+import logging
 import time
 from collections.abc import AsyncIterator, Callable
 from typing import Any, Protocol
@@ -11,6 +12,9 @@ from agent_lab.domain.events import ToolCallEvent
 from agent_lab.domain.messages import ChatMessage, StreamItem
 
 
+logger = logging.getLogger(__name__)
+
+
 class ChatClient(Protocol):
     async def stream_chat(
         self,
@@ -78,6 +82,7 @@ class DebugRuntime:
         try:
             await self._run_chat_agent(request, queues)
         except Exception as exc:
+            logger.exception("runtime session failed")
             await queues.output.put({"type": "error", "message": str(exc)})
         finally:
             await queues.events.put(None)
@@ -92,6 +97,7 @@ class DebugRuntime:
         await queues.input.put(ChatMessage(role="user", content=request.user_message))
 
         await queues.output.put({"type": "session_started"})
+        await self._audit(queues, "session_started")
 
         event_loops = 0
         round_index = 0
@@ -114,6 +120,12 @@ class DebugRuntime:
                 if event_loops < request.event_agent.max_event_loops
                 else []
             )
+            await self._audit(
+                queues,
+                "chat_round_started",
+                round_index=round_index,
+                tools_enabled=[tool["function"]["name"] for tool in tools],
+            )
 
             async for item in self.chat_client.stream_chat(
                 messages=messages,
@@ -143,6 +155,13 @@ class DebugRuntime:
                     saw_event = True
                     event = self._event_name_only(item.event)
                     events.append(event)
+                    await self._audit(
+                        queues,
+                        "chat_event_detected",
+                        round_index=round_index,
+                        event_id=event.id,
+                        event_name=event.name,
+                    )
                     assistant_tool_calls.append(
                         {
                             "id": event.id,
@@ -181,14 +200,23 @@ class DebugRuntime:
                     await queues.output.put(
                         {"type": "tool_result", "message": reply.model_dump()}
                     )
+                await self._audit(
+                    queues,
+                    "event_agent_completed",
+                    round_index=round_index,
+                    event_names=[event.name for event in events],
+                    result_count=len(tool_replies),
+                    result_summary="\n".join(reply.content for reply in tool_replies),
+                )
             messages.extend(tool_replies)
 
+            elapsed_ms = self._elapsed_ms(started_at)
             await queues.output.put(
                 {
                     "type": "round_stats",
                     "round_index": round_index,
                     "ttft_ms": ttft_ms,
-                    "elapsed_ms": self._elapsed_ms(started_at),
+                    "elapsed_ms": elapsed_ms,
                     "prompt_tokens": prompt_tokens,
                     "completion_tokens": completion_tokens,
                     "total_tokens": total_tokens,
@@ -196,12 +224,20 @@ class DebugRuntime:
                     "had_event": saw_event,
                 },
             )
+            await self._audit(
+                queues,
+                "chat_round_finished",
+                round_index=round_index,
+                had_event=saw_event,
+                elapsed_ms=elapsed_ms,
+            )
 
             if not saw_event:
                 break
 
             event_loops += 1
 
+        await self._audit(queues, "session_finished", round_count=round_index)
         await queues.output.put({"type": "done"})
 
     async def _consume_events(
@@ -246,6 +282,10 @@ class DebugRuntime:
             ChatMessage(role="system", content=prompt)
             for prompt in request.system_prompts
         ]
+        if request.chat_agent.system_prompt.strip():
+            messages.append(
+                ChatMessage(role="system", content=request.chat_agent.system_prompt)
+            )
         messages.extend(request.pre_messages)
         return messages
 
@@ -259,3 +299,18 @@ class DebugRuntime:
 
     def _elapsed_ms(self, started_at: float) -> int:
         return round((self.clock() - started_at) * 1000)
+
+    async def _audit(
+        self,
+        queues: RuntimeQueues,
+        event: str,
+        **details: Any,
+    ) -> None:
+        logger.info("audit event=%s details=%s", event, details)
+        await queues.output.put(
+            {
+                "type": "audit",
+                "event": event,
+                "details": details,
+            }
+        )

+ 71 - 1
src/agent_lab/application/tools.py

@@ -104,7 +104,33 @@ def build_default_tool_registry() -> ToolRegistry:
                 },
                 handler=_handle_handoff_note,
                 argument_resolver=_resolve_handoff_note_arguments,
-            )
+            ),
+            ToolDefinition(
+                name="mock_search",
+                description="Search mock external knowledge for the current turn.",
+                parameters={
+                    "type": "object",
+                    "properties": {
+                        "query": {"type": "string"},
+                    },
+                    "required": ["query"],
+                },
+                handler=_handle_mock_search,
+                argument_resolver=_resolve_mock_search_arguments,
+            ),
+            ToolDefinition(
+                name="mock_ticket",
+                description="Create a mock support ticket for follow-up work.",
+                parameters={
+                    "type": "object",
+                    "properties": {
+                        "title": {"type": "string"},
+                    },
+                    "required": ["title"],
+                },
+                handler=_handle_mock_ticket,
+                argument_resolver=_resolve_mock_ticket_arguments,
+            ),
         ]
     )
 
@@ -119,6 +145,26 @@ def _resolve_handoff_note_arguments(
     return {"message": content}
 
 
+def _resolve_mock_search_arguments(
+    event: ToolCallEvent,
+    history: Sequence[ChatMessage],
+) -> dict[str, Any]:
+    query = _latest_content(history, preferred_roles=("assistant", "user"))
+    if not query and not history:
+        query = str(event.arguments.get("query", ""))
+    return {"query": query}
+
+
+def _resolve_mock_ticket_arguments(
+    event: ToolCallEvent,
+    history: Sequence[ChatMessage],
+) -> dict[str, Any]:
+    title = _latest_content(history, preferred_roles=("assistant", "user"))
+    if not title and not history:
+        title = str(event.arguments.get("title", ""))
+    return {"title": title}
+
+
 def _latest_content(
     history: Sequence[ChatMessage],
     preferred_roles: tuple[str, ...],
@@ -135,3 +181,27 @@ def _handle_handoff_note(event: ToolCallEvent) -> dict[str, Any]:
         "tool": "handoff_note",
         "message": str(event.arguments.get("message", "")),
     }
+
+
+def _handle_mock_search(event: ToolCallEvent) -> dict[str, Any]:
+    query = str(event.arguments.get("query", ""))
+    return {
+        "tool": "mock_search",
+        "query": query,
+        "results": [
+            {
+                "title": "Mock knowledge base result",
+                "snippet": f"Simulated external context for: {query}",
+            }
+        ],
+    }
+
+
+def _handle_mock_ticket(event: ToolCallEvent) -> dict[str, Any]:
+    title = str(event.arguments.get("title", ""))
+    return {
+        "tool": "mock_ticket",
+        "ticket_id": "MOCK-001",
+        "title": title,
+        "status": "created",
+    }

+ 2 - 0
src/agent_lab/infrastructure/chat_client.py

@@ -45,6 +45,8 @@ class OpenAICompatibleChatClient:
             payload["stream_options"] = {"include_usage": True}
         if tools:
             payload["tools"] = tools
+        if params.extra_body:
+            payload.update(params.extra_body)
 
         headers = {}
         if self.api_key:

+ 71 - 16
tests/test_debug_runtime.py

@@ -1,6 +1,7 @@
 import asyncio
 import importlib
 import json
+import logging
 from collections.abc import AsyncIterator
 from typing import Any
 
@@ -23,6 +24,23 @@ async def _collect_outputs(stream: AsyncIterator[dict[str, Any]]) -> list[dict[s
     return [message async for message in stream]
 
 
+def _without_audit(outputs: list[dict[str, Any]]) -> list[dict[str, Any]]:
+    return [message for message in outputs if message["type"] != "audit"]
+
+
+def _message_types(outputs: list[dict[str, Any]]) -> list[str]:
+    return [message["type"] for message in _without_audit(outputs)]
+
+
+async def _next_non_audit(
+    stream: AsyncIterator[dict[str, Any]],
+) -> dict[str, Any]:
+    while True:
+        message = await anext(stream)
+        if message["type"] != "audit":
+            return message
+
+
 class RecordingQueue(asyncio.Queue):
     def __init__(self, name: str, log: list[tuple[str, str, str]]) -> None:
         super().__init__()
@@ -258,9 +276,10 @@ async def test_runtime_routes_chat_events_through_event_agent_then_continues_cha
     runtime = DebugRuntime(client)
 
     outputs = [message async for message in runtime.run(request)]
+    business_outputs = _without_audit(outputs)
 
     assert client.calls == 2
-    assert [message["type"] for message in outputs] == [
+    assert _message_types(outputs) == [
         "session_started",
         "event",
         "tool_result",
@@ -269,8 +288,41 @@ async def test_runtime_routes_chat_events_through_event_agent_then_continues_cha
         "round_stats",
         "done",
     ]
-    assert outputs[1]["event"]["name"] == "handoff_note"
-    assert outputs[4]["content"] == "final answer"
+    assert business_outputs[1]["event"]["name"] == "handoff_note"
+    assert business_outputs[4]["content"] == "final answer"
+
+
+@pytest.mark.asyncio
+async def test_runtime_emits_audit_events_and_backend_logs(caplog):
+    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=1),
+    )
+    runtime = DebugRuntime(FakeChatClient())
+
+    with caplog.at_level(logging.INFO, logger="agent_lab.application.runtime"):
+        outputs = [message async for message in runtime.run(request)]
+
+    audit_events = [
+        message["event"]
+        for message in outputs
+        if message["type"] == "audit"
+    ]
+    assert audit_events == [
+        "session_started",
+        "chat_round_started",
+        "chat_event_detected",
+        "event_agent_completed",
+        "chat_round_finished",
+        "chat_round_started",
+        "chat_round_finished",
+        "session_finished",
+    ]
+    assert "chat_event_detected" in caplog.text
+    assert "event_agent_completed" in caplog.text
 
 
 @pytest.mark.asyncio
@@ -328,7 +380,7 @@ async def test_runtime_batches_round_events_before_continuing_chat_agent():
 
     outputs = [message async for message in runtime.run(request)]
 
-    assert [message["type"] for message in outputs] == [
+    assert _message_types(outputs) == [
         "session_started",
         "message_delta",
         "event",
@@ -390,7 +442,7 @@ async def test_runtime_start_returns_queues_for_downstream_output_consumer():
         outputs.append(message)
         if message["type"] == "done":
             break
-    assert [message["type"] for message in outputs] == [
+    assert _message_types(outputs) == [
         "session_started",
         "message_delta",
         "usage",
@@ -412,11 +464,12 @@ async def test_runtime_finalizes_chat_after_reaching_event_loop_limit():
     runtime = DebugRuntime(client)
 
     outputs = [message async for message in runtime.run(request)]
+    business_outputs = _without_audit(outputs)
 
     assert client.calls == 2
     assert client.tools_by_call[0][0]["function"]["name"] == "handoff_note"
     assert client.tools_by_call[1] == []
-    assert [message["type"] for message in outputs] == [
+    assert _message_types(outputs) == [
         "session_started",
         "event",
         "tool_result",
@@ -425,7 +478,7 @@ async def test_runtime_finalizes_chat_after_reaching_event_loop_limit():
         "round_stats",
         "done",
     ]
-    assert outputs[4]["content"] == "final after event limit"
+    assert business_outputs[4]["content"] == "final after event limit"
 
 
 @pytest.mark.asyncio
@@ -443,8 +496,8 @@ async def test_runtime_buffers_upstream_user_input_until_after_matching_tool_rep
     runtime = DebugRuntime(client, queues=queues)
     stream = runtime.run(request)
 
-    assert await anext(stream) == {"type": "session_started"}
-    event_message = await anext(stream)
+    assert await _next_non_audit(stream) == {"type": "session_started"}
+    event_message = await _next_non_audit(stream)
     assert event_message["type"] == "event"
     await queues.input.put(ChatMessage(role="user", content="follow-up while tool runs"))
     remaining = [message async for message in stream]
@@ -488,8 +541,9 @@ async def test_runtime_continues_when_event_agent_tool_handler_raises():
         _collect_outputs(runtime.run(request)),
         timeout=1,
     )
+    business_outputs = _without_audit(outputs)
 
-    assert [message["type"] for message in outputs] == [
+    assert _message_types(outputs) == [
         "session_started",
         "event",
         "tool_result",
@@ -498,7 +552,7 @@ async def test_runtime_continues_when_event_agent_tool_handler_raises():
         "round_stats",
         "done",
     ]
-    assert json.loads(outputs[2]["message"]["content"]) == {
+    assert json.loads(business_outputs[2]["message"]["content"]) == {
         "tool": "handoff_note",
         "error": "tool handler failed: boom",
     }
@@ -524,7 +578,7 @@ async def test_runtime_uses_event_and_input_queues_for_event_agent_handoff():
 
     outputs = [message async for message in runtime.run(request)]
 
-    assert [message["type"] for message in outputs] == [
+    assert _message_types(outputs) == [
         "session_started",
         "event",
         "tool_result",
@@ -567,7 +621,7 @@ async def test_runtime_run_consumes_output_queue_in_stream_order():
 
     outputs = [message async for message in runtime.run(request)]
 
-    assert [message["type"] for message in outputs] == [
+    assert _message_types(outputs) == [
         "session_started",
         "event",
         "tool_result",
@@ -582,7 +636,7 @@ async def test_runtime_run_consumes_output_queue_in_stream_order():
     output_gets = [
         entry[2] for entry in queue_log if entry[0] == "output" and entry[1] == "get"
     ]
-    assert output_puts == [
+    assert [message for message in output_puts if message != "output:audit"] == [
         "output:session_started",
         "output:event",
         "output:tool_result",
@@ -693,15 +747,16 @@ async def test_runtime_emits_round_stats_with_clock_and_usage_after_model_turn()
     runtime = DebugRuntime(RoundStatsChatClient(), clock=lambda: next(ticks))
 
     outputs = [message async for message in runtime.run(request)]
+    business_outputs = _without_audit(outputs)
 
-    assert [message["type"] for message in outputs] == [
+    assert _message_types(outputs) == [
         "session_started",
         "message_delta",
         "usage",
         "round_stats",
         "done",
     ]
-    assert outputs[3] == {
+    assert business_outputs[3] == {
         "type": "round_stats",
         "round_index": 1,
         "ttft_ms": 123,

+ 65 - 12
tests/test_websocket_api.py

@@ -72,6 +72,20 @@ def _request_payload() -> dict:
     }
 
 
+def test_agent_params_defaults_include_extra_body_and_event_agent_defaults_to_one_round():
+    chat_params = AgentParams()
+    event_params = EventAgentParams()
+
+    assert chat_params.extra_body == {
+        "thinking": {"type": "disabled"},
+        "enable_search": False,
+        "search_options": {"forced_search": False},
+    }
+    assert event_params.extra_body == chat_params.extra_body
+    assert event_params.max_event_loops == 1
+    assert event_params.system_prompt == ""
+
+
 def test_health_returns_ok():
     app = create_app(runtime_factory=FakeRuntime)
     client = TestClient(app)
@@ -89,19 +103,13 @@ def test_api_tools_returns_handoff_note_metadata():
     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"],
-            },
-        }
+    tools = response.json()
+    assert [tool["name"] for tool in tools] == [
+        "handoff_note",
+        "mock_search",
+        "mock_ticket",
     ]
+    assert tools[0]["parameters"]["required"] == ["message"]
 
 
 def test_websocket_debug_streams_runtime_messages():
@@ -508,3 +516,48 @@ async def test_openai_chat_client_omits_stream_options_when_usage_disabled():
         ]
 
     assert "stream_options" not in captured_payloads[0]
+
+
+@pytest.mark.asyncio
+async def test_openai_chat_client_merges_agent_extra_body_into_payload():
+    captured_payloads: list[dict] = []
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        captured_payloads.append(json.loads(request.content))
+        return httpx.Response(
+            200,
+            content=b'data: {"choices":[{"delta":{"content":"hi"},"finish_reason":null}]}\n\n',
+        )
+
+    async with httpx.AsyncClient(
+        transport=httpx.MockTransport(handler),
+        base_url="https://llm.test/v1",
+    ) as http_client:
+        client = OpenAICompatibleChatClient(
+            api_key="",
+            base_url="https://llm.test/v1",
+            default_model="provider-default",
+            request_timeout_seconds=5,
+            http_client=http_client,
+        )
+        [
+            item
+            async for item in client.stream_chat(
+                messages=[ChatMessage(role="user", content="hi")],
+                tools=[],
+                params=AgentParams(
+                    model="provider-default",
+                    temperature=0.3,
+                    max_tokens=50,
+                    extra_body={
+                        "thinking": {"type": "disabled"},
+                        "enable_search": False,
+                        "search_options": {"forced_search": False},
+                    },
+                ),
+            )
+        ]
+
+    assert captured_payloads[0]["thinking"] == {"type": "disabled"}
+    assert captured_payloads[0]["enable_search"] is False
+    assert captured_payloads[0]["search_options"] == {"forced_search": False}

+ 2 - 0
tests/test_ws_smoke.py

@@ -23,6 +23,8 @@ def test_build_payload_can_disable_handoff_note_tool():
     assert payload["user_message"] == "debug the event flow"
     assert payload["chat_agent"]["model"] == "compatible-model"
     assert payload["event_agent"]["enabled_tools"] == []
+    assert payload["event_agent"]["max_event_loops"] == 1
+    assert payload["chat_agent"]["extra_body"] == payload["event_agent"]["extra_body"]
 
 
 def test_parser_defaults_to_blank_model_for_backend_fallback():