138 lines
4.9 KiB
Python
138 lines
4.9 KiB
Python
"""Master rotation: one logical master thread, many physical sessions.
|
|
|
|
When the policy says so, a new master takes over the window; the old one
|
|
hands out, closes, and its branches and queued normal injects move over.
|
|
"""
|
|
|
|
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.conversations.service import Conversations
|
|
from beaver_gateway.storage.models import Conversation
|
|
|
|
__all__ = ["HandoutContext", "Rotation", "RotationPolicy"]
|
|
|
|
_log = logging.getLogger("beaver_gateway.conversations.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 "night"
|
|
if now - created > self.max_age and silence > self.short_silence:
|
|
return "age"
|
|
if context_tokens > self.max_context_tokens and silence > self.short_silence:
|
|
return "context"
|
|
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, *, force: str | None = None
|
|
) -> list[Conversation]:
|
|
"""Rotate the masters that are due; ``force`` rotates every open master."""
|
|
rotated: list[Conversation] = []
|
|
if force:
|
|
due = [
|
|
(m, force)
|
|
for m in await self._conversations.find(kind="master", status="open")
|
|
]
|
|
else:
|
|
due = await self.due(now)
|
|
for master, reason in due:
|
|
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, reason=reason, 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)
|