Files
beaver-gateway/src/beaver_gateway/core/rotation.py
T

132 lines
4.8 KiB
Python

"""Master rotation (§4.5, §8.1, §8.3).
One logical master thread, many physical sessions: when the policy says
so, a new master is spawned and takes over the window atomically, the old
one writes its handout as its last turn, closes, its finished branches get
marked in their windows, its queued normal injects move over, and the new
one receives "new day". Silence is measured by the user's messages only -
injects never extend a day.
"""
from __future__ import annotations
import logging
from dataclasses import dataclass
from datetime import UTC, date, datetime, timedelta
from typing import TYPE_CHECKING
from zoneinfo import ZoneInfo
if TYPE_CHECKING:
from beaver_gateway.core.conversations import Conversations
from beaver_gateway.storage.models import Conversation
__all__ = ["HandoutContext", "Rotation", "RotationPolicy"]
_log = logging.getLogger("beaver_gateway.core.rotation")
@dataclass(frozen=True, slots=True)
class RotationPolicy:
tz: str = "UTC"
hour: int = 4
night_silence: timedelta = timedelta(hours=3)
max_age: timedelta = timedelta(hours=36)
short_silence: timedelta = timedelta(minutes=30)
max_context_tokens: int = 80_000
def reason(
self, master: Conversation, *, now: datetime, context_tokens: int
) -> str | None:
zone = ZoneInfo(self.tz)
created = _aware(master.created_at)
silence = now - _aware(master.last_user_activity_at or master.created_at)
boundary = now.astimezone(zone).replace(
hour=self.hour, minute=0, second=0, microsecond=0
)
if now.astimezone(zone) < boundary:
boundary -= timedelta(days=1)
if created < boundary and silence > self.night_silence:
return "ночь"
if now - created > self.max_age and silence > self.short_silence:
return "возраст"
if context_tokens > self.max_context_tokens and silence > self.short_silence:
return "транскрипт"
return None
def day_of(self, master: Conversation) -> date:
return _aware(master.created_at).astimezone(ZoneInfo(self.tz)).date()
@dataclass(frozen=True, slots=True)
class HandoutContext:
day: date
master: Conversation
reason: str
class Rotation:
def __init__(
self, conversations: Conversations, policy: RotationPolicy | None = None
) -> None:
self._conversations = conversations
self.policy = policy or RotationPolicy()
async def due(self, now: datetime | None = None) -> list[tuple[Conversation, str]]:
now = now or datetime.now(UTC)
out: list[tuple[Conversation, str]] = []
for master in await self._conversations.find(kind="master", status="open"):
tokens = await self._conversations.context_tokens(master)
reason = self.policy.reason(master, now=now, context_tokens=tokens)
if reason is not None:
out.append((master, reason))
return out
async def tick(self, now: datetime | None = None) -> list[Conversation]:
rotated: list[Conversation] = []
for master, reason in await self.due(now):
new = await self.rotate(master, reason)
if new is not None:
rotated.append(new)
return rotated
async def rotate(self, old: Conversation, reason: str) -> Conversation | None:
conversations = self._conversations
if await conversations.busy(old):
_log.info("rotation of %s skipped: busy", old.external_id)
return None
new = await conversations.spawn(
kind="master", agent=old.agent_name, seed="morning", origin="rotation"
)
day = self.policy.day_of(old)
_log.info(
"rotation (%s): %s -> %s, handout for %s",
reason,
old.external_id,
new.external_id,
day,
)
await conversations.handout(
old, HandoutContext(day=day, master=old, reason=reason)
)
await conversations.close(old)
for branch in await conversations.find(parent=old, limit=1000):
if branch.status == "open":
await conversations.reparent(branch, new)
elif not branch.running_turn:
await conversations.mark_closed(branch)
moved = await conversations.queue.move(old, new, priority="normal")
await conversations.new_day(new, moved=moved)
conversations.bus.publish(
"conversation.rotated",
conversation_id=new.external_id,
closed=old.external_id,
reason=reason,
handout_day=day.isoformat(),
moved=moved,
)
return new
def _aware(value: datetime) -> datetime:
return value if value.tzinfo is not None else value.replace(tzinfo=UTC)