323 lines
12 KiB
Python
323 lines
12 KiB
Python
"""Ending conversations: the distiller, the line cap, the master handover."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import contextlib
|
|
import inspect
|
|
import logging
|
|
from dataclasses import dataclass
|
|
from datetime import UTC, datetime, timedelta
|
|
from typing import TYPE_CHECKING, Any, cast
|
|
|
|
from beaver_gateway.conversations.distill import (
|
|
Digest,
|
|
DistillContext,
|
|
LineCap,
|
|
append_index,
|
|
check_digest,
|
|
find_digest,
|
|
index_line,
|
|
trim_summary,
|
|
written_paths,
|
|
)
|
|
from beaver_gateway.conversations.questions import Questions
|
|
from beaver_gateway.conversations.state import aware
|
|
from beaver_gateway.conversations.texts import NewDayContext
|
|
|
|
if TYPE_CHECKING:
|
|
from beaver_gateway.conversations.rotation import HandoutContext
|
|
from beaver_gateway.storage.models import Conversation, InjectQueueItem
|
|
|
|
__all__ = ["Closing", "DistillResult"]
|
|
|
|
_log = logging.getLogger(__name__)
|
|
|
|
CLOSE_WAIT = 0.25
|
|
CLOSE_TRIES = 40
|
|
CAP_TRIES = 3
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class DistillResult:
|
|
conversation: Conversation
|
|
fork: Conversation
|
|
text: str
|
|
digest: Digest | None
|
|
error: str | None
|
|
trimmed: bool
|
|
|
|
|
|
class Closing(Questions):
|
|
async def close(self, conv: Conversation) -> Conversation:
|
|
row = await self.set_status(conv, "closed")
|
|
with contextlib.suppress(LookupError):
|
|
await self._backend(conv.agent_name).close(conv.external_id)
|
|
return row
|
|
|
|
async def request_close(self, conv: Conversation) -> Conversation:
|
|
"""``close_chat`` from inside a turn: the chat closes once the turn ends."""
|
|
if conv.kind != "deep":
|
|
msg = f"only deep chats close this way, {conv.external_id} is {conv.kind}"
|
|
raise ValueError(msg)
|
|
return await self.set_flags(conv, {"close_requested": True})
|
|
|
|
async def idle(
|
|
self,
|
|
*,
|
|
kind: str,
|
|
days: int,
|
|
since: datetime | None = None,
|
|
limit: int | None = None,
|
|
) -> list[Conversation]:
|
|
"""Open conversations of ``kind`` with a session, quiet for ``days``."""
|
|
now = datetime.now(UTC)
|
|
cutoff = now - timedelta(days=days)
|
|
out: list[tuple[datetime, Conversation]] = []
|
|
for conv in await self.find(status="open", kind=kind, limit=10_000):
|
|
if conv.session_id is None:
|
|
continue
|
|
last = aware(conv.last_activity_at or conv.created_at)
|
|
if last > cutoff or (since is not None and last < since):
|
|
continue
|
|
out.append((last, conv))
|
|
out.sort(key=lambda pair: pair[0])
|
|
rows = [conv for _, conv in out]
|
|
return rows[:limit] if limit is not None else rows
|
|
|
|
async def distill(
|
|
self, conv: Conversation, *, reason: str = "api"
|
|
) -> DistillResult:
|
|
"""Fork under the distiller: digest checked and indexed, merge to the master."""
|
|
if self._distiller is None:
|
|
msg = "no distiller configured (Gateway(distiller=...))"
|
|
raise RuntimeError(msg)
|
|
if conv.kind != "deep":
|
|
msg = f"only deep chats are distilled, {conv.external_id} is {conv.kind}"
|
|
raise ValueError(msg)
|
|
row = await self.get_row(cast("int", conv.id)) or conv
|
|
if row.status != "open":
|
|
msg = f"conversation {row.external_id} is {row.status}"
|
|
raise ValueError(msg)
|
|
if await self.busy(row):
|
|
msg = f"conversation {row.external_id} is busy"
|
|
raise RuntimeError(msg)
|
|
memory = bool(row.flags.get("memory", True))
|
|
chat_name = await self.chat_name(row)
|
|
ctx = DistillContext(
|
|
conversation=row,
|
|
title=await self.implied_title(row),
|
|
source=await self.window_of(row),
|
|
chat_name=chat_name,
|
|
memory=memory,
|
|
reason=reason,
|
|
day=datetime.now(UTC).astimezone().date(),
|
|
)
|
|
prompt = await self._distill_prompt(ctx)
|
|
started = datetime.now(UTC)
|
|
self._bus.publish(
|
|
"distill.start",
|
|
conversation_id=row.external_id,
|
|
reason=reason,
|
|
memory=memory,
|
|
)
|
|
result = await self.fork(
|
|
row,
|
|
prompt,
|
|
strip_tools=True,
|
|
agent=self._distiller.agent,
|
|
title=f"digest: {chat_name}",
|
|
)
|
|
text, trimmed = trim_summary(result.text)
|
|
digest: Digest | None = None
|
|
error: str | None = None
|
|
if memory:
|
|
written = written_paths(result.capture.synthesized_messages)
|
|
path = find_digest(self._distiller, since=started, written=written)
|
|
if path is None:
|
|
error = self._texts.digest_missing
|
|
else:
|
|
checked = check_digest(path, self._distiller)
|
|
if isinstance(checked, str):
|
|
error = f"{path.name}: {checked}"
|
|
else:
|
|
digest = checked
|
|
append_index(self._distiller, index_line(digest, chat_name))
|
|
if error is not None:
|
|
_log.warning("distill of %s: %s", row.external_id, error)
|
|
master = await self.open_master()
|
|
if master is not None and text:
|
|
note = self._texts.closed.format(
|
|
chat=chat_name,
|
|
digest=(
|
|
self._texts.closed_digest.format(digest=digest.path.stem)
|
|
if digest
|
|
else ""
|
|
),
|
|
text=text,
|
|
)
|
|
await self.inject(master, note, urgency="normal", origin="digest")
|
|
await self.close(row)
|
|
row = await self.set_flags(
|
|
row,
|
|
{
|
|
"close_requested": None,
|
|
"closed_reason": reason,
|
|
"digest": str(digest.path) if digest else None,
|
|
"digest_error": error,
|
|
},
|
|
)
|
|
self._bus.publish(
|
|
"conversation.distilled",
|
|
conversation_id=row.external_id,
|
|
fork=result.conversation.external_id,
|
|
reason=reason,
|
|
memory=memory,
|
|
digest=str(digest.path) if digest else None,
|
|
error=error,
|
|
text=text,
|
|
trimmed=trimmed,
|
|
master=master.external_id if master is not None else None,
|
|
)
|
|
return DistillResult(
|
|
conversation=row,
|
|
fork=result.conversation,
|
|
text=text,
|
|
digest=digest,
|
|
error=error,
|
|
trimmed=trimmed,
|
|
)
|
|
|
|
async def _distill_prompt(self, ctx: DistillContext) -> str:
|
|
source = self._texts.distill
|
|
if source is None:
|
|
template = (
|
|
self._texts.distill_prompt
|
|
if ctx.memory
|
|
else self._texts.distill_prompt_no_memory
|
|
)
|
|
return template.format(
|
|
chat=ctx.chat_name, reason=ctx.reason, day=ctx.day.isoformat()
|
|
)
|
|
produced: Any = source(ctx)
|
|
return await produced if inspect.isawaitable(produced) else produced
|
|
|
|
async def _close_after_turn(self, conv: Conversation) -> None:
|
|
for _ in range(CLOSE_TRIES):
|
|
if await self.busy(conv):
|
|
await asyncio.sleep(CLOSE_WAIT)
|
|
continue
|
|
try:
|
|
await self.distill(conv, reason="close_chat")
|
|
except RuntimeError as exc:
|
|
_log.info("closing %s: %s, retrying", conv.external_id, exc)
|
|
await asyncio.sleep(CLOSE_WAIT)
|
|
continue
|
|
except Exception: # noqa: BLE001
|
|
_log.exception("closing %s after its turn failed", conv.external_id)
|
|
return
|
|
_log.warning("closing %s: still busy, giving up", conv.external_id)
|
|
|
|
async def before_turn(self, conv: Conversation) -> str | None:
|
|
cap = LineCap.from_flags(conv.flags.get("line_cap"))
|
|
if cap is None:
|
|
return None
|
|
try:
|
|
return cap.path.read_text(encoding="utf-8") if cap.path.exists() else ""
|
|
except OSError:
|
|
_log.exception("line cap: cannot read %s", cap.path)
|
|
return None
|
|
|
|
async def after_turn(self, conv: Conversation, before: str | None) -> None:
|
|
row = await self.get_row(cast("int", conv.id))
|
|
if row is None:
|
|
return
|
|
if row.kind == "deep" and row.flags.get("close_requested"):
|
|
self._track(asyncio.create_task(self._close_after_turn(row)))
|
|
cap = LineCap.from_flags(row.flags.get("line_cap"))
|
|
if cap is not None and before is not None:
|
|
await self._enforce_cap(row, cap, before)
|
|
|
|
async def _enforce_cap(self, conv: Conversation, cap: LineCap, before: str) -> None:
|
|
if not cap.path.exists():
|
|
return
|
|
after = cap.path.read_text(encoding="utf-8")
|
|
lines = sum(1 for line in after.splitlines() if line.strip())
|
|
if lines <= cap.max_lines:
|
|
return
|
|
if before:
|
|
cap.path.write_text(before, encoding="utf-8")
|
|
else:
|
|
cap.path.unlink()
|
|
attempts = int(conv.flags.get("line_cap_attempts", 0) or 0) + 1
|
|
await self.set_flags(conv, {"line_cap_attempts": attempts})
|
|
self._bus.publish(
|
|
"line_cap.bounced",
|
|
conversation_id=conv.external_id,
|
|
path=str(cap.path),
|
|
lines=lines,
|
|
max_lines=cap.max_lines,
|
|
attempt=attempts,
|
|
)
|
|
_log.warning(
|
|
"line cap: %s came back with %d lines (cap %d), restored; attempt %d",
|
|
cap.path,
|
|
lines,
|
|
cap.max_lines,
|
|
attempts,
|
|
)
|
|
if attempts > CAP_TRIES:
|
|
return
|
|
await self.inject(
|
|
conv,
|
|
self._texts.too_long.format(
|
|
name=cap.path.name, lines=lines, max_lines=cap.max_lines
|
|
),
|
|
urgency="urgent",
|
|
origin="cap",
|
|
interrupt=False,
|
|
)
|
|
|
|
async def handout(self, conv: Conversation, ctx: HandoutContext) -> str:
|
|
"""The closing master's last turn."""
|
|
source = self._texts.handout
|
|
if isinstance(source, str):
|
|
prompt = source.format(day=ctx.day.isoformat(), reason=ctx.reason)
|
|
else:
|
|
produced: Any = source(ctx)
|
|
prompt = await produced if inspect.isawaitable(produced) else produced
|
|
self._bus.publish(
|
|
"handout.start", conversation_id=conv.external_id, day=ctx.day.isoformat()
|
|
)
|
|
try:
|
|
text, _ = await self.run_text_turn(conv, prompt, origin="handout")
|
|
except Exception: # noqa: BLE001
|
|
_log.exception("handout turn on %s failed", conv.external_id)
|
|
text = ""
|
|
self._bus.publish(
|
|
"handout.end",
|
|
conversation_id=conv.external_id,
|
|
day=ctx.day.isoformat(),
|
|
text=text[:2000],
|
|
)
|
|
return text
|
|
|
|
async def new_day(
|
|
self, conv: Conversation, *, reason: str = "night", moved: int = 0
|
|
) -> InjectQueueItem:
|
|
"""The new master's first inject."""
|
|
ctx = NewDayContext(
|
|
day=datetime.now(UTC).astimezone().date(), reason=reason, moved=moved
|
|
)
|
|
source = self._texts.new_day
|
|
if isinstance(source, str):
|
|
text = source.format(day=ctx.day.isoformat(), reason=ctx.reason)
|
|
else:
|
|
produced: Any = source(ctx)
|
|
text = await produced if inspect.isawaitable(produced) else produced
|
|
if moved:
|
|
text += self._texts.moved_injects.format(moved=moved)
|
|
return await self.inject(
|
|
conv, text, urgency="urgent", origin="rotation", interrupt=False
|
|
)
|