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.events import (
    EventDefinition,
    EventExecutionContext,
    EventKernel,
    EventRequest,
    ResolvedEventArguments,
)
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,
        tool_choice: dict[str, Any] | None = None,
    ) -> 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.kernel = EventKernel(
            self.registry.event_registry,
            argument_fallback=(
                self._resolve_arguments_with_llm if chat_client is not None else None
            ),
        )
        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] = []
        result = await self.kernel.execute(
            self.registry.event_request(event),
            enabled_names=self.enabled_tools,
            context=ToolExecutionContext(
                history=history,
                system_prompt=system_prompt,
                extra_body=(
                    self.params.extra_body if extra_body is None else extra_body
                ),
            ),
        )
        payload = self.registry.tool_payload(result)
        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_arguments_with_llm(
        self,
        definition: EventDefinition,
        request: EventRequest,
        context: EventExecutionContext,
    ) -> ResolvedEventArguments | None:
        assert self.chat_client is not None
        tool = self.registry.tool_schema(request.name)
        if tool is None:
            return None

        event = self.registry._tool_event(request)
        messages = self._build_argument_messages(
            event,
            context.history,
            context.system_prompt,
        )
        params = self.params.model_copy(
            update={
                "extra_body": (
                    context.extra_body
                    if context.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,
            tool_choice={
                "type": "function",
                "function": {"name": event.name},
            },
        ):
            if item.kind == "raw_chunk" and item.raw_chunk is not None:
                raw_chunks.append(item.raw_chunk)
                continue

            if item.kind == "provider_tool_call" and item.event is not None:
                self._raw_chunks_by_event[event.id] = raw_chunks
                return ResolvedEventArguments(
                    event_name=item.event.name,
                    arguments=dict(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(self._project_visible_history(history))
        messages.append(
            ChatMessage(
                role="user",
                content=f"Resolve arguments for event `{event.name}`.",
            )
        )
        return messages

    def _project_visible_history(
        self,
        history: Sequence[ChatMessage],
    ) -> list[ChatMessage]:
        projected: list[ChatMessage] = []
        for message in history:
            if message.role == "tool":
                continue
            if message.role == "assistant":
                if not message.content:
                    continue
                projected.append(
                    ChatMessage(
                        role="assistant",
                        content=message.content,
                        name=message.name,
                    )
                )
                continue
            projected.append(message)
        return projected

    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,
        )