|
|
@@ -1,12 +1,22 @@
|
|
|
import asyncio
|
|
|
-import json
|
|
|
import logging
|
|
|
import time
|
|
|
from collections.abc import AsyncIterator, Callable
|
|
|
from typing import Any, Protocol
|
|
|
|
|
|
-from agent_lab.application.contracts import AgentParams, DebugRunRequest
|
|
|
+from agent_lab.application.contracts import (
|
|
|
+ AgentParams,
|
|
|
+ DebugRunRequest,
|
|
|
+ EventAgentParams,
|
|
|
+)
|
|
|
from agent_lab.application.event_agent import EventAgent, EventAgentRequest
|
|
|
+from agent_lab.application.events import (
|
|
|
+ EventBatchDecision,
|
|
|
+ EventBatchExecutor,
|
|
|
+ EventBatchResult,
|
|
|
+ EventKernel,
|
|
|
+ EventSource,
|
|
|
+)
|
|
|
from agent_lab.application.queues import RuntimeQueues
|
|
|
from agent_lab.application.session_store import SessionStore
|
|
|
from agent_lab.application.tools import ToolRegistry, build_default_tool_registry
|
|
|
@@ -43,6 +53,7 @@ class DebugRuntime:
|
|
|
self.session_store = session_store
|
|
|
self.clock = clock
|
|
|
self._tasks: list[asyncio.Task[None]] = []
|
|
|
+ self._batch_executors: dict[tuple[int, float], EventBatchExecutor] = {}
|
|
|
|
|
|
def start(self, request: DebugRunRequest) -> RuntimeQueues:
|
|
|
queues = self.queues or RuntimeQueues()
|
|
|
@@ -161,6 +172,8 @@ class DebugRuntime:
|
|
|
assistant_content: list[str] = []
|
|
|
events: list[ToolCallEvent] = []
|
|
|
tool_replies: list[ChatMessage] = []
|
|
|
+ batch_result: EventBatchResult | None = None
|
|
|
+ decision: EventBatchDecision | None = None
|
|
|
message_stream_started = False
|
|
|
message_delta_count = 0
|
|
|
configured_events = request.event_agent.enabled_tools
|
|
|
@@ -309,13 +322,27 @@ class DebugRuntime:
|
|
|
messages.append(assistant_message)
|
|
|
|
|
|
if events:
|
|
|
+ batch_scope = self._event_scope(request, queues, turn_index=1)
|
|
|
+ await self._audit(
|
|
|
+ queues,
|
|
|
+ "event_batch_started",
|
|
|
+ turn_started_at=turn_started_at,
|
|
|
+ round_index=round_index,
|
|
|
+ scope=batch_scope,
|
|
|
+ event_names=[event.name for event in events],
|
|
|
+ max_parallel_events=request.event_agent.max_parallel_events,
|
|
|
+ batch_timeout_seconds=request.event_agent.batch_timeout_seconds,
|
|
|
+ )
|
|
|
if request.tool_invocation_mode == "chat_agent_tools":
|
|
|
- tool_replies = await self._execute_provider_tools(
|
|
|
+ batch_result, tool_replies = await self._execute_provider_tools(
|
|
|
events,
|
|
|
enabled_names=request.event_agent.enabled_tools,
|
|
|
+ params=request.event_agent,
|
|
|
+ scope=batch_scope,
|
|
|
)
|
|
|
messages.extend(tool_replies)
|
|
|
else:
|
|
|
+ response = asyncio.get_running_loop().create_future()
|
|
|
await queues.events.put(
|
|
|
EventAgentRequest(
|
|
|
events=events,
|
|
|
@@ -323,8 +350,11 @@ class DebugRuntime:
|
|
|
system_prompt=request.event_agent.system_prompt,
|
|
|
extra_body=request.event_agent.extra_body,
|
|
|
turn_started_at=turn_started_at,
|
|
|
+ scope=batch_scope,
|
|
|
+ response=response,
|
|
|
)
|
|
|
)
|
|
|
+ batch_result = await response
|
|
|
tool_replies = await self._wait_for_tool_replies(
|
|
|
queues,
|
|
|
[event.id for event in events],
|
|
|
@@ -352,10 +382,28 @@ class DebugRuntime:
|
|
|
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
|
|
|
- ),
|
|
|
)
|
|
|
+ assert batch_result is not None
|
|
|
+ await self._audit_batch_result(
|
|
|
+ queues,
|
|
|
+ batch_result,
|
|
|
+ turn_started_at=turn_started_at,
|
|
|
+ round_index=round_index,
|
|
|
+ )
|
|
|
+ decision = await self._apply_batch_policy(
|
|
|
+ queues,
|
|
|
+ messages,
|
|
|
+ batch_result,
|
|
|
+ turn_started_at=turn_started_at,
|
|
|
+ round_index=round_index,
|
|
|
+ )
|
|
|
+ if (
|
|
|
+ decision.llm_follow_up
|
|
|
+ and request.tool_invocation_mode == "dual_agent"
|
|
|
+ ):
|
|
|
+ summary = self.registry.compact_results_message(batch_result)
|
|
|
+ if summary is not None:
|
|
|
+ messages.append(summary)
|
|
|
|
|
|
elapsed_ms = self._elapsed_ms(started_at)
|
|
|
await queues.output.put(
|
|
|
@@ -385,6 +433,9 @@ class DebugRuntime:
|
|
|
break
|
|
|
|
|
|
event_loops += 1
|
|
|
+ assert decision is not None
|
|
|
+ if decision.terminate or not decision.llm_follow_up:
|
|
|
+ break
|
|
|
|
|
|
await self._audit(
|
|
|
queues,
|
|
|
@@ -425,7 +476,7 @@ class DebugRuntime:
|
|
|
turn_index += 1
|
|
|
turn_started_at = next_turn_started_at or self.clock()
|
|
|
next_turn_started_at = None
|
|
|
- await self._run_chat_turn(
|
|
|
+ terminated = await self._run_chat_turn(
|
|
|
request,
|
|
|
queues,
|
|
|
messages,
|
|
|
@@ -434,6 +485,17 @@ class DebugRuntime:
|
|
|
session_id,
|
|
|
turn_started_at,
|
|
|
)
|
|
|
+ if terminated:
|
|
|
+ await self._audit(
|
|
|
+ queues,
|
|
|
+ "session_finished",
|
|
|
+ turn_started_at=turn_started_at,
|
|
|
+ session_id=session_id,
|
|
|
+ turn_index=turn_index,
|
|
|
+ terminated=True,
|
|
|
+ )
|
|
|
+ await queues.output.put({"type": "done"})
|
|
|
+ return
|
|
|
|
|
|
async def _run_chat_turn(
|
|
|
self,
|
|
|
@@ -444,7 +506,7 @@ class DebugRuntime:
|
|
|
turn_index: int,
|
|
|
session_id: str | None = None,
|
|
|
turn_started_at: float | None = None,
|
|
|
- ) -> None:
|
|
|
+ ) -> bool:
|
|
|
turn_started_at = turn_started_at or self.clock()
|
|
|
messages.append(user_message)
|
|
|
self._start_persisted_turn(session_id, turn_index, user_message)
|
|
|
@@ -460,6 +522,7 @@ class DebugRuntime:
|
|
|
event_loops = 0
|
|
|
round_index = 0
|
|
|
deferred_user_messages: list[ChatMessage] = []
|
|
|
+ terminate_session = False
|
|
|
while True:
|
|
|
if round_index:
|
|
|
await self._drain_input(
|
|
|
@@ -477,6 +540,8 @@ class DebugRuntime:
|
|
|
saw_event = False
|
|
|
assistant_content: list[str] = []
|
|
|
events: list[ToolCallEvent] = []
|
|
|
+ batch_result: EventBatchResult | None = None
|
|
|
+ decision: EventBatchDecision | None = None
|
|
|
message_stream_started = False
|
|
|
message_delta_count = 0
|
|
|
configured_events = request.event_agent.enabled_tools
|
|
|
@@ -642,13 +707,33 @@ class DebugRuntime:
|
|
|
)
|
|
|
|
|
|
if events:
|
|
|
+ batch_scope = self._event_scope(
|
|
|
+ request,
|
|
|
+ queues,
|
|
|
+ turn_index=turn_index,
|
|
|
+ )
|
|
|
+ await self._audit(
|
|
|
+ queues,
|
|
|
+ "event_batch_started",
|
|
|
+ turn_started_at=turn_started_at,
|
|
|
+ session_id=session_id,
|
|
|
+ turn_index=turn_index,
|
|
|
+ round_index=round_index,
|
|
|
+ scope=batch_scope,
|
|
|
+ event_names=[event.name for event in events],
|
|
|
+ max_parallel_events=request.event_agent.max_parallel_events,
|
|
|
+ batch_timeout_seconds=request.event_agent.batch_timeout_seconds,
|
|
|
+ )
|
|
|
if request.tool_invocation_mode == "chat_agent_tools":
|
|
|
- tool_replies = await self._execute_provider_tools(
|
|
|
+ batch_result, tool_replies = await self._execute_provider_tools(
|
|
|
events,
|
|
|
enabled_names=request.event_agent.enabled_tools,
|
|
|
+ params=request.event_agent,
|
|
|
+ scope=batch_scope,
|
|
|
)
|
|
|
messages.extend(tool_replies)
|
|
|
else:
|
|
|
+ response = asyncio.get_running_loop().create_future()
|
|
|
await queues.events.put(
|
|
|
EventAgentRequest(
|
|
|
events=events,
|
|
|
@@ -659,8 +744,11 @@ class DebugRuntime:
|
|
|
turn_index=turn_index,
|
|
|
round_index=round_index,
|
|
|
turn_started_at=turn_started_at,
|
|
|
+ scope=batch_scope,
|
|
|
+ response=response,
|
|
|
)
|
|
|
)
|
|
|
+ batch_result = await response
|
|
|
tool_replies = await self._wait_for_tool_replies(
|
|
|
queues,
|
|
|
[event.id for event in events],
|
|
|
@@ -692,10 +780,32 @@ class DebugRuntime:
|
|
|
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
|
|
|
- ),
|
|
|
)
|
|
|
+ assert batch_result is not None
|
|
|
+ await self._audit_batch_result(
|
|
|
+ queues,
|
|
|
+ batch_result,
|
|
|
+ turn_started_at=turn_started_at,
|
|
|
+ session_id=session_id,
|
|
|
+ turn_index=turn_index,
|
|
|
+ round_index=round_index,
|
|
|
+ )
|
|
|
+ decision = await self._apply_batch_policy(
|
|
|
+ queues,
|
|
|
+ messages,
|
|
|
+ batch_result,
|
|
|
+ turn_started_at=turn_started_at,
|
|
|
+ session_id=session_id,
|
|
|
+ turn_index=turn_index,
|
|
|
+ round_index=round_index,
|
|
|
+ )
|
|
|
+ if (
|
|
|
+ decision.llm_follow_up
|
|
|
+ and request.tool_invocation_mode == "dual_agent"
|
|
|
+ ):
|
|
|
+ summary = self.registry.compact_results_message(batch_result)
|
|
|
+ if summary is not None:
|
|
|
+ messages.append(summary)
|
|
|
|
|
|
elapsed_ms = self._elapsed_ms(started_at)
|
|
|
self._append_persisted_usage(
|
|
|
@@ -740,6 +850,12 @@ class DebugRuntime:
|
|
|
break
|
|
|
|
|
|
event_loops += 1
|
|
|
+ assert decision is not None
|
|
|
+ if decision.terminate:
|
|
|
+ terminate_session = True
|
|
|
+ break
|
|
|
+ if not decision.llm_follow_up:
|
|
|
+ break
|
|
|
|
|
|
for deferred_message in deferred_user_messages:
|
|
|
await queues.input.put(deferred_message)
|
|
|
@@ -759,6 +875,7 @@ class DebugRuntime:
|
|
|
"round_count": round_index,
|
|
|
}
|
|
|
)
|
|
|
+ return terminate_session
|
|
|
|
|
|
async def _consume_events(
|
|
|
self,
|
|
|
@@ -791,28 +908,35 @@ class DebugRuntime:
|
|
|
if tool is not None
|
|
|
],
|
|
|
)
|
|
|
- replies = await event_agent.handle_many(
|
|
|
- request.events,
|
|
|
- history=request.history,
|
|
|
- system_prompt=request.system_prompt,
|
|
|
- extra_body=request.extra_body,
|
|
|
- )
|
|
|
- await self._audit(
|
|
|
- queues,
|
|
|
- "event_agent_response",
|
|
|
- turn_started_at=request.turn_started_at,
|
|
|
- session_id=request.session_id,
|
|
|
- agent="event_agent",
|
|
|
- turn_index=request.turn_index,
|
|
|
- round_index=request.round_index,
|
|
|
- replies=[reply.model_dump() for reply in replies],
|
|
|
- raw_model_chunks=event_agent.raw_model_chunks(request.events),
|
|
|
- )
|
|
|
- for reply in replies:
|
|
|
- await queues.input.put(reply)
|
|
|
- summary = event_agent.summarize_replies(replies)
|
|
|
- if summary is not None:
|
|
|
- await queues.input.put(summary)
|
|
|
+ try:
|
|
|
+ batch = await event_agent.handle_many(
|
|
|
+ request.events,
|
|
|
+ history=request.history,
|
|
|
+ system_prompt=request.system_prompt,
|
|
|
+ extra_body=request.extra_body,
|
|
|
+ scope=request.scope,
|
|
|
+ )
|
|
|
+ replies = event_agent.registry.tool_replies(request.events, batch)
|
|
|
+ await self._audit(
|
|
|
+ queues,
|
|
|
+ "event_agent_response",
|
|
|
+ turn_started_at=request.turn_started_at,
|
|
|
+ session_id=request.session_id,
|
|
|
+ agent="event_agent",
|
|
|
+ turn_index=request.turn_index,
|
|
|
+ round_index=request.round_index,
|
|
|
+ replies=[reply.model_dump() for reply in replies],
|
|
|
+ raw_model_chunks=event_agent.raw_model_chunks(request.events),
|
|
|
+ )
|
|
|
+ for reply in replies:
|
|
|
+ await queues.input.put(reply)
|
|
|
+ except BaseException as exc:
|
|
|
+ if request.response is not None and not request.response.done():
|
|
|
+ request.response.set_exception(exc)
|
|
|
+ raise
|
|
|
+ else:
|
|
|
+ if request.response is not None and not request.response.done():
|
|
|
+ request.response.set_result(batch)
|
|
|
|
|
|
async def _wait_for_tool_replies(
|
|
|
self,
|
|
|
@@ -908,25 +1032,124 @@ class DebugRuntime:
|
|
|
events: list[ToolCallEvent],
|
|
|
*,
|
|
|
enabled_names: list[str],
|
|
|
- ) -> list[ChatMessage]:
|
|
|
- payloads = await asyncio.gather(
|
|
|
- *[
|
|
|
- self.registry.execute_async(
|
|
|
+ params: EventAgentParams,
|
|
|
+ scope: str,
|
|
|
+ ) -> tuple[EventBatchResult, list[ChatMessage]]:
|
|
|
+ executor = self._batch_executor(params)
|
|
|
+ batch = await executor.execute(
|
|
|
+ [
|
|
|
+ self.registry.event_request(
|
|
|
event,
|
|
|
- enabled_names=enabled_names,
|
|
|
+ source=EventSource.PROVIDER_RESOLVED,
|
|
|
)
|
|
|
for event in events
|
|
|
- ]
|
|
|
+ ],
|
|
|
+ enabled_names=enabled_names,
|
|
|
+ scope=scope,
|
|
|
)
|
|
|
- return [
|
|
|
- ChatMessage(
|
|
|
- role="tool",
|
|
|
- content=json.dumps(payload, ensure_ascii=False),
|
|
|
- name=event.name,
|
|
|
- tool_call_id=event.id,
|
|
|
+ return batch, self.registry.tool_replies(events, batch)
|
|
|
+
|
|
|
+ def _batch_executor(self, params: EventAgentParams) -> EventBatchExecutor:
|
|
|
+ key = (params.max_parallel_events, params.batch_timeout_seconds)
|
|
|
+ executor = self._batch_executors.get(key)
|
|
|
+ if executor is None:
|
|
|
+ executor = EventBatchExecutor(
|
|
|
+ EventKernel(self.registry.event_registry),
|
|
|
+ max_parallel_events=params.max_parallel_events,
|
|
|
+ batch_timeout_seconds=params.batch_timeout_seconds,
|
|
|
)
|
|
|
- for event, payload in zip(events, payloads, strict=True)
|
|
|
- ]
|
|
|
+ self._batch_executors[key] = executor
|
|
|
+ return executor
|
|
|
+
|
|
|
+ def _event_scope(
|
|
|
+ self,
|
|
|
+ request: DebugRunRequest,
|
|
|
+ queues: RuntimeQueues,
|
|
|
+ *,
|
|
|
+ turn_index: int,
|
|
|
+ ) -> str:
|
|
|
+ session_scope = request.session_id or f"runtime-{id(queues)}"
|
|
|
+ return f"{session_scope}:turn-{turn_index}"
|
|
|
+
|
|
|
+ async def _audit_batch_result(
|
|
|
+ self,
|
|
|
+ queues: RuntimeQueues,
|
|
|
+ batch: EventBatchResult,
|
|
|
+ *,
|
|
|
+ turn_started_at: float | None,
|
|
|
+ session_id: str | None = None,
|
|
|
+ turn_index: int | None = None,
|
|
|
+ round_index: int | None = None,
|
|
|
+ ) -> None:
|
|
|
+ await self._audit(
|
|
|
+ queues,
|
|
|
+ "event_batch_results",
|
|
|
+ turn_started_at=turn_started_at,
|
|
|
+ session_id=session_id,
|
|
|
+ turn_index=turn_index,
|
|
|
+ round_index=round_index,
|
|
|
+ results=[
|
|
|
+ {
|
|
|
+ "event_id": result.event_id,
|
|
|
+ "event_name": result.event_name,
|
|
|
+ "status": result.status.value,
|
|
|
+ "error": result.error,
|
|
|
+ "result_policy": result.result_policy.value,
|
|
|
+ "terminal": result.terminal,
|
|
|
+ }
|
|
|
+ for result in batch.results
|
|
|
+ ],
|
|
|
+ timeout_count=sum(
|
|
|
+ result.status.value == "timeout" for result in batch.results
|
|
|
+ ),
|
|
|
+ coalesced_count=batch.coalesced_count,
|
|
|
+ replayed_count=batch.replayed_count,
|
|
|
+ deadline_exceeded=batch.deadline_exceeded,
|
|
|
+ )
|
|
|
+
|
|
|
+ async def _apply_batch_policy(
|
|
|
+ self,
|
|
|
+ queues: RuntimeQueues,
|
|
|
+ messages: list[ChatMessage],
|
|
|
+ batch: EventBatchResult,
|
|
|
+ *,
|
|
|
+ turn_started_at: float | None,
|
|
|
+ session_id: str | None = None,
|
|
|
+ turn_index: int | None = None,
|
|
|
+ round_index: int | None = None,
|
|
|
+ ) -> EventBatchDecision:
|
|
|
+ decision = self.registry.batch_decision(batch)
|
|
|
+ for content in decision.template_messages:
|
|
|
+ message = ChatMessage(role="assistant", content=content)
|
|
|
+ messages.append(message)
|
|
|
+ self._append_persisted_message(session_id, turn_index or 0, message)
|
|
|
+ await queues.output.put({"type": "message_delta", "content": content})
|
|
|
+ await self._audit(
|
|
|
+ queues,
|
|
|
+ "event_policy_decision",
|
|
|
+ turn_started_at=turn_started_at,
|
|
|
+ session_id=session_id,
|
|
|
+ turn_index=turn_index,
|
|
|
+ round_index=round_index,
|
|
|
+ template_count=len(decision.template_messages),
|
|
|
+ llm_follow_up=decision.llm_follow_up,
|
|
|
+ terminate=decision.terminate,
|
|
|
+ )
|
|
|
+ if decision.terminate:
|
|
|
+ await self._audit(
|
|
|
+ queues,
|
|
|
+ "terminal_completed",
|
|
|
+ turn_started_at=turn_started_at,
|
|
|
+ session_id=session_id,
|
|
|
+ turn_index=turn_index,
|
|
|
+ round_index=round_index,
|
|
|
+ event_ids=[
|
|
|
+ result.event_id
|
|
|
+ for result in batch.results
|
|
|
+ if result.terminal and result.status.value == "success"
|
|
|
+ ],
|
|
|
+ )
|
|
|
+ return decision
|
|
|
|
|
|
def _chat_messages_for_round(
|
|
|
self,
|