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