Files
beaver-gateway/src/beaver_gateway/conversations/service.py
T

159 lines
5.5 KiB
Python

"""``Conversations`` - the service every frontend, job and gateway tool talks to.
Built as layers, one file each: rows → seeds → turns → messaging →
spawning → questions → closing; this file adds start, stop and restart
recovery. A turn started by a user message streams back to whoever asked;
a turn started by an inject streams nowhere and can only speak via ``say``.
"""
from __future__ import annotations
import asyncio
import contextlib
import logging
import re
from datetime import UTC, datetime, timedelta, tzinfo
from sqlmodel import col, select
from beaver_gateway.conversations.closing import Closing, DistillResult
from beaver_gateway.conversations.kinds import KINDS
from beaver_gateway.conversations.rows import context_of, implied_title
from beaver_gateway.conversations.seeds import SEEDS
from beaver_gateway.conversations.spawning import ForkResult
from beaver_gateway.conversations.state import aware
from beaver_gateway.conversations.texts import (
ConversationTexts,
NewDayContext,
SeedContext,
UserSaid,
)
from beaver_gateway.storage.models import Conversation
__all__ = [
"KINDS",
"SEEDS",
"ConversationTexts",
"Conversations",
"DistillResult",
"ForkResult",
"NewDayContext",
"SeedContext",
"UserSaid",
"context_of",
"implied_title",
"parse_at",
]
_log = logging.getLogger(__name__)
_RELATIVE = re.compile(r"^\+(\d+)\s*([smhd])$")
_UNITS = {"s": 1, "m": 60, "h": 3600, "d": 86400}
class Conversations(Closing):
async def start(self) -> None:
await self.recover()
for row_id in await self._queue.conversations_with_pending():
self._ensure_worker(row_id)
self._idle_task = asyncio.create_task(self._idle_loop())
async def stop(self) -> None:
tasks = list(self._tasks)
if self._idle_task is not None:
tasks.append(self._idle_task)
for task in tasks:
task.cancel()
for task in tasks:
with contextlib.suppress(BaseException):
await task
self._tasks.clear()
self._idle_task = None
async def recover(self) -> list[Conversation]:
"""Repair the transcripts of turns a restart cut and tell each conversation."""
async with self._db.session() as session:
result = await session.exec(
select(Conversation).where(col(Conversation.running_turn).is_not(None))
)
cut = list(result.all())
for conv in cut:
fixed = 0
if conv.session_id is not None:
backend = self._backend(conv.agent_name)
try:
fixed = await backend.repair_session(
conv.session_id, text=self._texts.interrupted
)
except Exception: # noqa: BLE001
_log.exception("repair of %s failed", conv.session_id)
turn_id = conv.running_turn
async def clear(row: Conversation) -> None:
row.running_turn = None
row.pending_question = False
await self._update(conv, clear)
note = self._texts.cut_by_restart.format(turn_id=turn_id)
if fixed:
note += self._texts.repaired_tools.format(
fixed=fixed, interrupted=self._texts.interrupted
)
await self.inject(conv, note, urgency="normal", origin="system")
_log.warning("conversation %s: %s", conv.external_id, note)
for item in await self._queue.interrupted():
_log.warning(
"queue item #%s (%s) was running at restart; marked interrupted",
item.id,
item.priority,
)
return cut
async def _idle_loop(self) -> None:
while True:
try:
await self._emit_idle()
except Exception: # noqa: BLE001
_log.exception("idle watcher failed")
await asyncio.sleep(self._idle_interval)
async def _emit_idle(self) -> None:
if not self._idle_days:
return
now = datetime.now(UTC)
for conv in await self.find(status="open", limit=10_000):
last = aware(conv.last_activity_at or conv.created_at)
days = int((now - last).total_seconds() // 86400)
due = [d for d in self._idle_days if days >= d]
if not due:
continue
notified = int(conv.flags.get("idle_notified", 0) or 0)
if due[-1] <= notified:
continue
await self.set_flags(conv, {"idle_notified": due[-1]})
bindings = await self.bindings(conv)
self._bus.publish(
"conversation.idle",
conversation_id=conv.external_id,
kind=conv.kind,
agent=conv.agent_name,
days=due[-1],
bindings=[
{"frontend": b.frontend, "external_id": b.external_id}
for b in bindings
if b.visible
],
)
def parse_at(at: str, tz: tzinfo = UTC) -> datetime:
raw = at.strip()
match = _RELATIVE.match(raw.replace(" ", ""))
if match:
amount, unit = match.groups()
return datetime.now(UTC) + timedelta(seconds=int(amount) * _UNITS[unit])
parsed = datetime.fromisoformat(raw)
if parsed.tzinfo is None:
parsed = parsed.replace(tzinfo=tz)
return parsed.astimezone(UTC)