agent-lab/src/agent_lab/application/event_agent.py

import asyncio
import json
from collections.abc import AsyncIterator, Iterable, Sequence
from dataclasses import dataclass
from typing import Any, Protocol

from agent_lab.application.contracts import AgentParams
from agent_lab.application.tools import (
    ToolExecutionContext,
    ToolRegistry,
    build_default_tool_registry,
)
from agent_lab.domain.events import ToolCallEvent
from agent_lab.domain.messages import ChatMessage, StreamItem


class EventAgentChatClient(Protocol):
    async def stream_chat(
        self,
        messages: list[ChatMessage],
        tools: list[dict[str, Any]],
        params: AgentParams,
    ) -> AsyncIterator[StreamItem]:
        ...


@dataclass(frozen=True)
class EventAgentRequest:
    events: list[ToolCallEvent]
    history: list[ChatMessage]
    system_prompt: str = ""
    extra_body: dict[str, Any] | None = None
    session_id: str | None = None
    turn_index: int | None = None
    round_index: int | None = None
    turn_started_at: float | None = None


class EventAgent:
    def __init__(
        self,
        enabled_tools: Iterable[str],
        registry: ToolRegistry | None = None,
        chat_client: EventAgentChatClient | None = None,
        params: AgentParams | None = None,
    ) -> None:
        self.enabled_tools = set(enabled_tools)
        self.registry = registry or build_default_tool_registry()
        self.chat_client = chat_client
        self.params = params or AgentParams()
        self._raw_chunks_by_event: dict[str, list[dict[str, Any]]] = {}

    async def handle(
        self,
        event: ToolCallEvent,
        history: Sequence[ChatMessage] = (),
        system_prompt: str = "",
        extra_body: dict[str, Any] | None = None,
    ) -> ChatMessage:
        self._raw_chunks_by_event[event.id] = []
        if event.name not in self.enabled_tools:
            return self._tool_reply(
                event,
                {
                    "tool": event.name,
                    "error": "tool disabled",
                },
            )

        try:
            if self.chat_client is None:
                payload = self.registry.handle(
                    event,
                    context=ToolExecutionContext(
                        history=history,
                        system_prompt=system_prompt,
                        extra_body=extra_body or {},
                    ),
                )
            else:
                resolved_event = await self._resolve_event_with_llm(
                    event,
                    history,
                    system_prompt,
                    extra_body,
                )
                if resolved_event is None:
                    payload = self.registry.handle(
                        event,
                        context=ToolExecutionContext(
                            history=history,
                            system_prompt=system_prompt,
                            extra_body=extra_body or {},
                        ),
                    )
                else:
                    payload = self.registry.execute(resolved_event)
        except Exception as exc:
            payload = {
                "tool": event.name,
                "error": f"tool handler failed: {exc}",
            }
        return self._tool_reply(event, payload)

    async def handle_many(
        self,
        events: Sequence[ToolCallEvent],
        history: Sequence[ChatMessage],
        system_prompt: str = "",
        extra_body: dict[str, Any] | None = None,
    ) -> list[ChatMessage]:
        for event in events:
            self._raw_chunks_by_event[event.id] = []
        return await asyncio.gather(
            *[
                self.handle(
                    event,
                    history=history,
                    system_prompt=system_prompt,
                    extra_body=extra_body,
                )
                for event in events
            ]
        )

    def raw_model_chunks(
        self,
        events: Sequence[ToolCallEvent],
    ) -> list[dict[str, Any]]:
        return [
            {
                "event_id": event.id,
                "event_name": event.name,
                "chunks": self._raw_chunks_by_event.get(event.id, []),
            }
            for event in events
        ]

    async def _resolve_event_with_llm(
        self,
        event: ToolCallEvent,
        history: Sequence[ChatMessage],
        system_prompt: str,
        extra_body: dict[str, Any] | None,
    ) -> ToolCallEvent | None:
        assert self.chat_client is not None
        tool = self.registry.tool_schema(event.name)
        if tool is None:
            return event

        messages = self._build_argument_messages(event, history, system_prompt)
        params = self.params.model_copy(
            update={
                "extra_body": (
                    extra_body if extra_body is not None else self.params.extra_body
                ),
            }
        )
        raw_chunks: list[dict[str, Any]] = []
        async for item in self.chat_client.stream_chat(
            messages=messages,
            tools=[tool],
            params=params,
        ):
            if item.kind == "raw_chunk" and item.raw_chunk is not None:
                raw_chunks.append(item.raw_chunk)
                continue

            if (
                item.kind in {"provider_tool_call", "event"}
                and item.event is not None
            ):
                self._raw_chunks_by_event[event.id] = raw_chunks
                return event.model_copy(
                    update={
                        "arguments": item.event.arguments,
                        "raw_arguments": item.event.raw_arguments,
                    }
                )

        self._raw_chunks_by_event[event.id] = raw_chunks
        return None

    def _build_argument_messages(
        self,
        event: ToolCallEvent,
        history: Sequence[ChatMessage],
        system_prompt: str,
    ) -> list[ChatMessage]:
        messages = [
            ChatMessage(
                role="system",
                content=(
                    "Generate parameters for the requested EventAgent tool. "
                    "Call exactly the provided tool with valid JSON arguments."
                ),
            )
        ]
        if system_prompt.strip():
            messages.append(ChatMessage(role="system", content=system_prompt))
        messages.extend(
            message
            for message in history
            if message.role in {"system", "user", "assistant"}
        )
        messages.append(
            ChatMessage(
                role="user",
                content=f"Resolve arguments for event `{event.name}`.",
            )
        )
        return messages

    def summarize_replies(self, replies: Sequence[ChatMessage]) -> ChatMessage | None:
        if not replies:
            return None
        return ChatMessage(
            role="user",
            content=(
                "EventAgent results:\n"
                + "\n".join(reply.content for reply in replies)
            ),
            name="event_agent",
        )

    def _tool_reply(self, event: ToolCallEvent, payload: dict[str, Any]) -> ChatMessage:
        return ChatMessage(
            role="tool",
            content=json.dumps(payload, ensure_ascii=False),
            name=event.name,
            tool_call_id=event.id,
        )