553 lines
20 KiB
Python
553 lines
20 KiB
Python
"""Running turns: one worker per conversation over its queue, the backend call."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import base64
|
|
import contextlib
|
|
import inspect
|
|
import logging
|
|
from datetime import UTC, datetime
|
|
from pathlib import Path
|
|
from typing import TYPE_CHECKING, Any, cast
|
|
from uuid import uuid4
|
|
|
|
from claude_agent_sdk import (
|
|
AssistantMessage,
|
|
RateLimitEvent,
|
|
ResultMessage,
|
|
StreamEvent,
|
|
ToolResultBlock,
|
|
ToolUseBlock,
|
|
UserMessage,
|
|
)
|
|
|
|
from beaver_gateway.backends.capture import TurnCapture
|
|
from beaver_gateway.backends.transcript import reply_text, text_of
|
|
from beaver_gateway.conversations import injects
|
|
from beaver_gateway.conversations.seeds import Seeds
|
|
from beaver_gateway.conversations.state import Runner, aware, iso
|
|
from beaver_gateway.conversations.texts import UserSaid
|
|
from beaver_gateway.frontends.accumulate import StreamAccumulator
|
|
from beaver_gateway.storage.models import Conversation, InjectQueueItem, RateLimit
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import AsyncIterator, Callable, Sequence
|
|
|
|
from beaver_gateway.events.stream import MessageStreamEvent
|
|
|
|
__all__ = ["Turns", "content_of", "image_block"]
|
|
|
|
_log = logging.getLogger(__name__)
|
|
|
|
|
|
class Turns(Seeds):
|
|
async def turn(
|
|
self,
|
|
conv: Conversation,
|
|
*,
|
|
messages: Sequence[Any],
|
|
origin: str,
|
|
capture: TurnCapture | None = None,
|
|
session_id: str | None = None,
|
|
use_session: bool = True,
|
|
tools: bool = True,
|
|
turn_id: str | None = None,
|
|
item_origin: str | None = None,
|
|
shown: str | None = None,
|
|
attach: bool = False,
|
|
) -> AsyncIterator[MessageStreamEvent]:
|
|
"""Run one turn under the conversation's lock; the only path to the backend.
|
|
|
|
``shown`` is what frontends display as the turn's prompt (the user's
|
|
words, not the envelope). ``attach`` sends nothing: the turn is one
|
|
the live session started by itself, a subagent reporting back.
|
|
"""
|
|
row_id = cast("int", conv.id)
|
|
runner = self._runner(row_id)
|
|
backend = self._backend(conv.agent_name)
|
|
turn_id = turn_id or f"turn_{uuid4().hex[:12]}"
|
|
capture = capture or TurnCapture()
|
|
async with runner.lock:
|
|
runner.turn_id = turn_id
|
|
runner.origin = origin
|
|
runner.text = shown if shown is not None else _prompt_preview(messages)
|
|
runner.started_at = datetime.now(UTC)
|
|
runner.tools = {}
|
|
row = await self._mark_running(conv, turn_id)
|
|
resume = session_id if session_id is not None else row.session_id
|
|
before = await self.before_turn(conv)
|
|
self._bus.publish(
|
|
"turn.start",
|
|
conversation_id=conv.external_id,
|
|
turn_id=turn_id,
|
|
origin=origin,
|
|
item_origin=item_origin,
|
|
text=runner.text,
|
|
)
|
|
stop = "error"
|
|
cut = False
|
|
try:
|
|
observer = self._observer(conv, runner, turn_id, origin)
|
|
events = (
|
|
backend.attach(
|
|
key=conv.external_id,
|
|
capture=capture,
|
|
observer=observer,
|
|
turn_id=turn_id,
|
|
)
|
|
if attach
|
|
else backend.complete(
|
|
agent=self._claude_agent(conv.agent_name),
|
|
messages=messages,
|
|
conversation_id=conv.external_id,
|
|
session_id=resume if use_session else None,
|
|
reseed=not use_session,
|
|
capture=capture,
|
|
kind=conv.kind,
|
|
pinned=conv.kind == "master",
|
|
tools=tools,
|
|
observer=observer,
|
|
turn_id=turn_id,
|
|
)
|
|
)
|
|
async for event in events:
|
|
yield event
|
|
stop = "interrupted" if capture.interrupted else "end_turn"
|
|
except asyncio.CancelledError:
|
|
cut = True
|
|
raise
|
|
finally:
|
|
runner.turn_id = None
|
|
await self._mark_done(conv, capture, cut=cut)
|
|
if stop != "error":
|
|
try:
|
|
await self.after_turn(conv, before)
|
|
except Exception: # noqa: BLE001
|
|
_log.exception("after-turn hook on %s failed", conv.external_id)
|
|
self._bus.publish(
|
|
"turn.end",
|
|
conversation_id=conv.external_id,
|
|
turn_id=turn_id,
|
|
origin=origin,
|
|
item_origin=item_origin,
|
|
stop=stop,
|
|
usage=_usage_dict(capture),
|
|
)
|
|
|
|
async def run_text_turn(
|
|
self,
|
|
conv: Conversation,
|
|
text: str,
|
|
*,
|
|
origin: str,
|
|
tools: bool = True,
|
|
turn_id: str | None = None,
|
|
item_origin: str | None = None,
|
|
attachments: Sequence[dict[str, Any]] | None = None,
|
|
shown: str | None = None,
|
|
attach: bool = False,
|
|
) -> tuple[str, TurnCapture]:
|
|
capture = TurnCapture()
|
|
acc = StreamAccumulator()
|
|
agent = self._claude_agent(conv.agent_name)
|
|
async for event in self.turn(
|
|
conv,
|
|
messages=[{"role": "user", "content": content_of(text, attachments)}],
|
|
origin=origin,
|
|
capture=capture,
|
|
tools=tools,
|
|
turn_id=turn_id,
|
|
item_origin=item_origin,
|
|
shown=shown,
|
|
attach=attach,
|
|
):
|
|
acc.feed(event)
|
|
message = acc.finalize(model=agent.model)
|
|
return reply_text(
|
|
message.content, thinking_is_progress=agent.thinking_is_progress
|
|
), capture
|
|
|
|
async def before_turn(self, conv: Conversation) -> str | None: # noqa: ARG002
|
|
return None
|
|
|
|
async def after_turn(self, conv: Conversation, before: str | None) -> None: # noqa: ARG002
|
|
return
|
|
|
|
def turn_origin(self, conv: Conversation) -> str | None:
|
|
runner = self._runners.get(cast("int", conv.id))
|
|
return runner.origin if runner is not None and runner.turn_id else None
|
|
|
|
async def _woken(self, key: str, note: str) -> None:
|
|
"""The live session of ``key`` started a turn by itself; run it next."""
|
|
conv = await self.get(key)
|
|
if conv is None:
|
|
_log.warning("self-started turn on unknown conversation %s: %s", key, note)
|
|
return
|
|
_log.info("conversation %s woke by itself: %s", key, note)
|
|
self._ensure_worker(cast("int", conv.id))
|
|
|
|
async def _run_pending(self, conv: Conversation) -> None:
|
|
turn_id = f"turn_{uuid4().hex[:12]}"
|
|
note = self._backend(conv.agent_name).pending_note(conv.external_id) or ""
|
|
try:
|
|
text, _ = await self.run_text_turn(
|
|
conv,
|
|
note,
|
|
origin="task",
|
|
turn_id=turn_id,
|
|
item_origin="task",
|
|
shown=note,
|
|
attach=True,
|
|
)
|
|
except Exception: # noqa: BLE001
|
|
_log.exception(
|
|
"self-started turn %s on %s failed", turn_id, conv.external_id
|
|
)
|
|
return
|
|
if not text.strip():
|
|
return
|
|
self._bus.publish(
|
|
"reply",
|
|
conversation_id=conv.external_id,
|
|
turn_id=turn_id,
|
|
item=None,
|
|
item_origin="task",
|
|
source="task",
|
|
prompt=self._texts.task_prompt.format(note=note),
|
|
user_text=None,
|
|
text=text,
|
|
)
|
|
|
|
def _ensure_worker(self, row_id: int) -> None:
|
|
runner = self._runner(row_id)
|
|
runner.wake.set()
|
|
if runner.task is None or runner.task.done():
|
|
runner.task = asyncio.create_task(self._worker(row_id))
|
|
self._track(runner.task)
|
|
|
|
async def _worker(self, row_id: int) -> None:
|
|
runner = self._runner(row_id)
|
|
while True:
|
|
conv = await self.get_row(row_id)
|
|
if conv is not None and self._has_pending(conv):
|
|
await self._run_pending(conv)
|
|
continue
|
|
items = await self._queue.pending(row_id)
|
|
batch, wait = self._pick(items)
|
|
if batch is None:
|
|
runner.wake.clear()
|
|
if wait is None:
|
|
return
|
|
with contextlib.suppress(TimeoutError):
|
|
await asyncio.wait_for(runner.wake.wait(), timeout=wait)
|
|
continue
|
|
conv = await self.get_row(row_id)
|
|
if conv is None:
|
|
await self._queue.finish(batch, status="failed")
|
|
return
|
|
await self._run_batch(conv, batch)
|
|
|
|
def _pick(
|
|
self, items: list[InjectQueueItem]
|
|
) -> tuple[list[InjectQueueItem] | None, float | None]:
|
|
if not items:
|
|
return None, None
|
|
head = items[0]
|
|
if head.priority == "urgent":
|
|
return [head], None
|
|
tail = [i for i in items if i is not head and i.priority in ("wake", "normal")]
|
|
if head.priority in ("user", "wake"):
|
|
return [head, *tail], None
|
|
age = (datetime.now(UTC) - aware(head.created_at)).total_seconds()
|
|
if age >= self._normal_window:
|
|
return [head, *tail], None
|
|
return None, max(self._normal_window - age, 1.0)
|
|
|
|
async def _run_batch(
|
|
self, conv: Conversation, batch: list[InjectQueueItem]
|
|
) -> None:
|
|
head = batch[0]
|
|
turn_id = f"turn_{uuid4().hex[:12]}"
|
|
await self._queue.start(batch, turn_id)
|
|
if head.priority == "user":
|
|
origin = "user"
|
|
prompt = head.text
|
|
await self._note_user(conv, head.text)
|
|
envelope = self._envelope_for(conv, head.text)
|
|
if envelope:
|
|
prompt += "\n\n" + envelope
|
|
if len(batch) > 1:
|
|
prompt += "\n\n" + self._bundle(batch[1:])
|
|
else:
|
|
origin = "inject"
|
|
prompt = "\n\n".join(
|
|
f"{self._texts.inject_header(injects.context_of(i))}\n{i.text}"
|
|
for i in batch
|
|
)
|
|
seed = await self.pending_seed(conv)
|
|
if seed:
|
|
prompt = f"{seed}\n\n{prompt}"
|
|
try:
|
|
text, capture = await self.run_text_turn(
|
|
conv,
|
|
prompt,
|
|
origin=origin,
|
|
turn_id=turn_id,
|
|
item_origin=head.origin,
|
|
attachments=head.attachments if origin == "user" else None,
|
|
shown=head.text if origin == "user" else None,
|
|
)
|
|
except Exception: # noqa: BLE001
|
|
_log.exception("turn %s on %s failed", turn_id, conv.external_id)
|
|
await self._queue.finish(batch, status="failed")
|
|
return
|
|
await self._queue.finish(
|
|
batch, status="interrupted" if capture.interrupted else "done"
|
|
)
|
|
if origin == "user":
|
|
self._bus.publish(
|
|
"reply",
|
|
conversation_id=conv.external_id,
|
|
turn_id=turn_id,
|
|
item=head.id,
|
|
item_origin=head.origin,
|
|
source="queue",
|
|
prompt=prompt,
|
|
user_text=head.text,
|
|
text=text,
|
|
)
|
|
|
|
def _has_pending(self, conv: Conversation) -> bool:
|
|
try:
|
|
return self._backend(conv.agent_name).has_pending(conv.external_id)
|
|
except LookupError:
|
|
return False
|
|
|
|
def _bundle(self, items: Sequence[InjectQueueItem]) -> str:
|
|
lines = [self._texts.bundle_header.format(since=iso(items[0].created_at))]
|
|
lines.extend(f"- [{i.origin}] {i.text}" for i in items)
|
|
return "\n".join(lines)
|
|
|
|
def _envelope_for(self, conv: Conversation, text: str = "") -> str | None:
|
|
if self._envelope is None:
|
|
return None
|
|
if conv.kind == "master":
|
|
return self._envelope.build(text=text, kind="master", agent=conv.agent_name)
|
|
if conv.kind == "branch":
|
|
return self._envelope.recall_only(
|
|
text=text, kind="branch", agent=conv.agent_name
|
|
)
|
|
return None
|
|
|
|
async def _note_user(self, conv: Conversation, text: str) -> None:
|
|
if self._user_sink is None:
|
|
return
|
|
message = UserSaid(
|
|
conversation_id=conv.external_id,
|
|
kind=conv.kind,
|
|
title=conv.title,
|
|
text=text,
|
|
at=datetime.now(UTC),
|
|
agent=conv.agent_name,
|
|
)
|
|
try:
|
|
result = self._user_sink(message)
|
|
if inspect.isawaitable(result):
|
|
await result
|
|
except Exception: # noqa: BLE001
|
|
_log.exception("user sink failed for %s", conv.external_id)
|
|
|
|
def _observer(
|
|
self, conv: Conversation, runner: Runner, turn_id: str, origin: str
|
|
) -> Callable[[Any], None]:
|
|
conversation_id = conv.external_id
|
|
|
|
def observe(message: Any) -> None:
|
|
parent = getattr(message, "parent_tool_use_id", None)
|
|
if isinstance(message, RateLimitEvent):
|
|
self._observe_rate_limit(conv, message)
|
|
elif isinstance(message, StreamEvent):
|
|
self._bus.publish(
|
|
"stream",
|
|
conversation_id=conversation_id,
|
|
turn_id=turn_id,
|
|
origin=origin,
|
|
parent_tool_use_id=parent,
|
|
event=message.event,
|
|
)
|
|
elif isinstance(message, AssistantMessage):
|
|
for block in message.content:
|
|
if isinstance(block, ToolUseBlock):
|
|
event = self._bus.publish(
|
|
"tool",
|
|
conversation_id=conversation_id,
|
|
turn_id=turn_id,
|
|
origin=origin,
|
|
parent_tool_use_id=parent,
|
|
tool_use_id=block.id,
|
|
name=block.name,
|
|
input=block.input,
|
|
)
|
|
runner.tools[block.id] = {
|
|
"tool_use_id": block.id,
|
|
"name": block.name,
|
|
"input": block.input,
|
|
"parent_tool_use_id": parent,
|
|
"started_at": event["ts"],
|
|
"ended_at": None,
|
|
"is_error": None,
|
|
"content": None,
|
|
}
|
|
elif isinstance(message, UserMessage):
|
|
blocks = message.content if isinstance(message.content, list) else ()
|
|
for block in blocks:
|
|
if isinstance(block, ToolResultBlock):
|
|
event = self._bus.publish(
|
|
"tool.result",
|
|
conversation_id=conversation_id,
|
|
turn_id=turn_id,
|
|
origin=origin,
|
|
parent_tool_use_id=parent,
|
|
tool_use_id=block.tool_use_id,
|
|
is_error=bool(block.is_error),
|
|
content=_result_preview(block.content),
|
|
)
|
|
node = runner.tools.get(block.tool_use_id)
|
|
if node is not None:
|
|
node["ended_at"] = event["ts"]
|
|
node["is_error"] = event["is_error"]
|
|
node["content"] = event["content"]
|
|
elif isinstance(message, ResultMessage) and parent is None:
|
|
self._bus.publish(
|
|
"result",
|
|
conversation_id=conversation_id,
|
|
turn_id=turn_id,
|
|
origin=origin,
|
|
subtype=message.subtype,
|
|
is_error=message.is_error,
|
|
num_turns=message.num_turns,
|
|
)
|
|
|
|
return observe
|
|
|
|
def _observe_rate_limit(self, conv: Conversation, message: RateLimitEvent) -> None:
|
|
info = message.rate_limit_info
|
|
row = RateLimit(
|
|
window=info.rate_limit_type or "unknown",
|
|
status=info.status,
|
|
utilization=info.utilization,
|
|
resets_at=_from_unix(info.resets_at),
|
|
overage_status=info.overage_status,
|
|
overage_resets_at=_from_unix(info.overage_resets_at),
|
|
agent_name=conv.agent_name,
|
|
session_id=message.session_id,
|
|
raw=dict(info.raw),
|
|
)
|
|
self._bus.publish(
|
|
"rate_limit",
|
|
conversation_id=conv.external_id,
|
|
window=row.window,
|
|
status=row.status,
|
|
utilization=row.utilization,
|
|
resets_at=iso(row.resets_at),
|
|
overage_status=row.overage_status,
|
|
)
|
|
self._track(asyncio.create_task(self._record_rate_limit(row)))
|
|
|
|
async def _record_rate_limit(self, row: RateLimit) -> None:
|
|
try:
|
|
async with self._db.session() as session:
|
|
session.add(row)
|
|
await session.commit()
|
|
except Exception: # noqa: BLE001
|
|
_log.exception("rate limit write failed")
|
|
|
|
async def _mark_running(self, conv: Conversation, turn_id: str) -> Conversation:
|
|
async def apply(row: Conversation) -> None:
|
|
row.running_turn = turn_id
|
|
row.last_activity_at = datetime.now(UTC)
|
|
|
|
return await self._update(conv, apply)
|
|
|
|
async def _mark_done(
|
|
self, conv: Conversation, capture: TurnCapture, *, cut: bool = False
|
|
) -> None:
|
|
async def apply(row: Conversation) -> None:
|
|
if not cut:
|
|
row.running_turn = None
|
|
row.last_activity_at = datetime.now(UTC)
|
|
if capture.session_id is not None:
|
|
row.session_id = capture.session_id
|
|
|
|
await self._update(conv, apply)
|
|
|
|
|
|
IMAGE_MAX_BYTES = 5 * 1024 * 1024
|
|
|
|
|
|
def content_of(
|
|
text: str, attachments: Sequence[dict[str, Any]] | None
|
|
) -> str | list[dict[str, Any]]:
|
|
"""Plain text, or text plus the images that ride with it."""
|
|
blocks = [b for a in attachments or () if (b := image_block(a)) is not None]
|
|
if not blocks:
|
|
return text
|
|
return [{"type": "text", "text": text}, *blocks]
|
|
|
|
|
|
def image_block(attachment: dict[str, Any]) -> dict[str, Any] | None:
|
|
media_type = str(attachment.get("media_type") or "")
|
|
path = Path(str(attachment.get("path") or ""))
|
|
if not media_type.startswith("image/") or not path.is_file():
|
|
return None
|
|
if path.stat().st_size > IMAGE_MAX_BYTES:
|
|
return None
|
|
data = base64.b64encode(path.read_bytes()).decode("ascii")
|
|
return {
|
|
"type": "image",
|
|
"source": {"type": "base64", "media_type": media_type, "data": data},
|
|
}
|
|
|
|
|
|
def _prompt_preview(messages: Sequence[Any], limit: int = 400) -> str | None:
|
|
if not messages:
|
|
return None
|
|
text = text_of(messages[-1].get("content"))
|
|
return text[:limit] if text else None
|
|
|
|
|
|
def _from_unix(value: int | None) -> datetime | None:
|
|
return datetime.fromtimestamp(value, tz=UTC) if value is not None else None
|
|
|
|
|
|
def _result_preview(
|
|
content: str | list[dict[str, Any]] | None, limit: int = 400
|
|
) -> str:
|
|
if content is None:
|
|
return ""
|
|
text = (
|
|
content
|
|
if isinstance(content, str)
|
|
else "\n".join(
|
|
str(part.get("text", ""))
|
|
for part in content
|
|
if isinstance(part, dict) and part.get("type") == "text"
|
|
)
|
|
)
|
|
return text if len(text) <= limit else text[:limit] + "…"
|
|
|
|
|
|
def _usage_dict(capture: TurnCapture) -> dict[str, Any] | None:
|
|
usage = capture.usage
|
|
if usage is None:
|
|
return None
|
|
return {
|
|
"input": usage.input_tokens,
|
|
"output": usage.output_tokens,
|
|
"cache_read": usage.cache_read_tokens,
|
|
"cache_creation": usage.cache_creation_tokens,
|
|
"cost_usd": usage.cost_usd,
|
|
"duration_ms": usage.duration_ms,
|
|
}
|