refactor: split flat core into capability packages, layer the conversations service, English defaults for every model-facing text

This commit is contained in:
hh
2026-09-02 00:13:20 +02:00
parent b96714338f
commit cae2ed4161
77 changed files with 2987 additions and 2944 deletions
View File
+619
View File
@@ -0,0 +1,619 @@
"""Jobs and deferred injects on pgqueuer (§3.6, §4.5).
A job is a name, a handler and its triggers: a cron expression, the
webhook ``/hooks/<name>``, gateway bus events. The executor is pgqueuer on
the gateway's own Postgres (a dedicated autocommit connection for
LISTEN/NOTIFY), so cron ticks, webhook deliveries and the ``schedule``
tool's one-off injects all live in one table and survive a restart. A
handler only queues work for a conversation and returns; the turn itself
runs in the conversation's worker. Non-critical jobs step aside while the
subscription window is past its threshold.
"""
from __future__ import annotations
import asyncio
import contextlib
import functools
import json
import logging
import traceback
from dataclasses import dataclass
from datetime import UTC, datetime, timedelta, tzinfo
from typing import TYPE_CHECKING, Any, cast
from zoneinfo import ZoneInfo
from croniter import croniter
from pgqueuer import PgQueuer, Queries
from pgqueuer.core.executors import ScheduleExecutor
from pgqueuer.models import JobId
from sqlalchemy import func
from sqlmodel import col, select
from starlette.applications import Starlette
from starlette.responses import JSONResponse
from starlette.routing import Route
from beaver_gateway.conversations.service import parse_at
from beaver_gateway.storage.models import JobRunRecord
if TYPE_CHECKING:
from collections.abc import Awaitable, Callable, Coroutine, Sequence
from pgqueuer.models import Job as PgJob
from pgqueuer.models import Schedule as PgSchedule
from pgqueuer.ports.driver import Driver
from starlette.requests import Request
from beaver_gateway.conversations.distill import LineCap
from beaver_gateway.conversations.injects import Priority
from beaver_gateway.conversations.rotation import Rotation
from beaver_gateway.conversations.service import Conversations, DistillResult
from beaver_gateway.storage.models import Conversation
__all__ = ["INJECT", "Budget", "Job", "JobRun", "LocalCron", "Scheduler", "next_run"]
_log = logging.getLogger("beaver_gateway.jobs.scheduler")
INJECT = "inject"
RETRY = timedelta(minutes=15)
ERROR_MAX = 2000
def next_run(expression: str, tz: tzinfo, now: datetime | None = None) -> datetime:
"""The cron's next fire time, read in ``tz``, returned in UTC."""
start = (now or datetime.now(UTC)).astimezone(tz)
return croniter(expression, start_time=start).get_next(datetime).astimezone(UTC)
@dataclass
class LocalCron(ScheduleExecutor):
tz: tzinfo = UTC
def get_next(self) -> datetime:
return next_run(self.parameters.expression, self.tz)
@dataclass(frozen=True, slots=True)
class Job:
name: str
run: Callable[[JobRun], Awaitable[None]]
cron: str | None = None
webhook: bool = False
events: tuple[str, ...] = ()
critical: bool = True
dedupe: bool = True
@property
def entrypoint(self) -> str:
return f"job:{self.name}"
@dataclass(frozen=True, slots=True)
class Budget:
threshold: float = 0.7
tokens: int | None = None
window: timedelta = timedelta(hours=5)
limit_window: str = "five_hour"
@dataclass(slots=True)
class JobRun:
job: Job
trigger: str
payload: dict[str, Any]
scheduler: Scheduler
@property
def conversations(self) -> Conversations:
return self.scheduler.conversations
async def master(self) -> Conversation | None:
masters = await self.conversations.find(kind="master", status="open", limit=1)
return masters[0] if masters else None
async def inject_master(
self, text: str, *, urgency: Priority = "normal", origin: str | None = None
) -> bool:
master = await self.master()
if master is None:
_log.error(
"job %s: no open master, inject lost: %s", self.job.name, text[:200]
)
return False
await self.conversations.inject(
master, text, urgency=urgency, origin=origin or self.job.name
)
return True
async def spawn_job(
self,
*,
agent: str,
text: str,
title: str | None = None,
line_cap: LineCap | None = None,
) -> Conversation:
"""A headless job turn; ``line_cap`` bounces a rewrite past the cap."""
return await self.conversations.spawn(
kind="job",
agent=agent,
seed="brief",
text=text,
title=title,
origin="job",
flags={"line_cap": line_cap.as_flags()} if line_cap else None,
)
async def close_idle(
self,
*,
kind: str = "deep",
days: int = 2,
limit: int = 3,
since: datetime | None = None,
) -> list[DistillResult]:
"""§4.5: close chats quiet for ``days``, at most ``limit`` per run."""
out: list[DistillResult] = []
for conv in await self.conversations.idle(
kind=kind, days=days, since=since, limit=limit
):
try:
out.append(
await self.conversations.distill(conv, reason=f"idle {days}d")
)
except Exception: # noqa: BLE001
_log.exception("closing idle %s failed", conv.external_id)
return out
async def retry_in(self, delay: timedelta) -> None:
await self.scheduler.trigger(
self.job, self.payload, delay=delay, trigger=self.trigger
)
async def rotate(self) -> list[Conversation]:
rotation = self.scheduler.rotation
return await rotation.tick() if rotation is not None else []
def background(self, coro: Coroutine[Any, Any, Any]) -> None:
self.scheduler.background(coro)
class Scheduler:
def __init__(
self,
*,
conversations: Conversations,
jobs: Sequence[Job] = (),
driver: Driver | None = None,
budget: Budget | None = None,
rotation: Rotation | None = None,
heartbeat: timedelta = timedelta(seconds=30),
tz: str = "UTC",
) -> None:
self.conversations = conversations
self.rotation = rotation
self.tz = ZoneInfo(tz)
self.budget = budget or Budget()
self._jobs = {job.name: job for job in jobs}
self._driver = driver
self._queries: Queries | None = None
self._heartbeat = heartbeat
self._tasks: set[asyncio.Task[Any]] = set()
self._runs: dict[str, dict[str, Any]] = {}
@property
def enabled(self) -> bool:
return self._queries is not None
@property
def jobs(self) -> list[Job]:
return list(self._jobs.values())
def job(self, name: str) -> Job | None:
return self._jobs.get(name)
# ---- lifecycle -----------------------------------------------------
async def start(self) -> None:
if self._driver is not None:
queries = Queries(self._driver)
if not await queries.has_table("pgqueuer"):
await queries.install()
pgq = PgQueuer(self._driver, queries=queries)
self._register(pgq)
self._queries = queries
self._spawn(pgq.run(heartbeat_timeout=self._heartbeat))
else:
_log.warning("scheduler: no postgres driver - cron and webhooks are off")
self._spawn(self._events())
async def stop(self) -> None:
for task in list(self._tasks):
task.cancel()
for task in list(self._tasks):
with contextlib.suppress(BaseException):
await task
self._tasks.clear()
self._queries = None
def _register(self, pgq: PgQueuer) -> None:
@pgq.entrypoint(INJECT)
async def deliver(job: PgJob) -> None:
await self._deliver(job)
for job in self._jobs.values():
self._register_job(pgq, job)
def _register_job(self, pgq: PgQueuer, job: Job) -> None:
@pgq.entrypoint(job.entrypoint)
async def queued(pg_job: PgJob) -> None:
data = _decode(pg_job.payload)
await self._dispatch(
job, trigger=str(data.pop("trigger", "queue")), payload=data
)
if job.cron:
@pgq.schedule(
job.name,
job.cron,
clean_old=True,
executor_factory=functools.partial(LocalCron, tz=self.tz),
)
async def cron(_schedule: PgSchedule) -> None:
await self._dispatch(job, trigger="cron", payload={})
async def _events(self) -> None:
listeners = [job for job in self._jobs.values() if job.events]
if not listeners:
return
async for event in self.conversations.bus.stream():
for job in listeners:
if event.get("type") in job.events:
self._spawn(
self._dispatch(job, trigger="event", payload=dict(event))
)
# ---- dispatch ------------------------------------------------------
async def _dispatch(
self, job: Job, *, trigger: str, payload: dict[str, Any]
) -> None:
bus = self.conversations.bus
if not job.critical and await self.throttled():
bus.publish("job.deferred", job=job.name, trigger=trigger)
if trigger != "cron":
await self.trigger(job, payload, delay=RETRY, trigger=trigger)
return
started = datetime.now(UTC)
bus.publish("job.start", job=job.name, trigger=trigger)
status, error = "done", None
try:
await job.run(JobRun(job, trigger, payload, self))
except Exception as exc: # noqa: BLE001
status, error = "failed", _error_text(exc)
_log.exception("job %s (%s) failed", job.name, trigger)
record = JobRunRecord(
job=job.name,
trigger=trigger,
started_at=started,
finished_at=datetime.now(UTC),
status=status,
error=error,
payload=payload,
)
await self._record(record)
self._runs[job.name] = self._run_public(record)
bus.publish("job.end", job=job.name, trigger=trigger, status=status)
async def _record(self, record: JobRunRecord) -> None:
try:
async with self.conversations.db.session() as session:
session.add(record)
await session.commit()
except Exception: # noqa: BLE001
_log.exception("job %s: run not recorded", record.job)
async def trigger(
self,
job: Job,
payload: dict[str, Any] | None = None,
*,
delay: timedelta | None = None,
trigger: str = "manual",
) -> int | None:
"""Queue one run of ``job``; runs inline when there is no executor."""
data = {**(payload or {}), "trigger": trigger}
if self._queries is None:
self._spawn(
self._dispatch(job, trigger=trigger, payload=payload or {}), delay=delay
)
return None
ids = await self._queries.enqueue(
job.entrypoint, _encode(data), execute_after=delay
)
return int(ids[0]) if ids and ids[0] is not None else None
async def hook(self, name: str, payload: dict[str, Any]) -> int | None:
job = self._jobs.get(name)
if job is None or not job.webhook:
msg = f"no webhook job {name!r}"
raise LookupError(msg)
self.conversations.bus.publish("hook", job=name)
if self._queries is None:
return await self.trigger(job, payload, trigger="webhook")
ids = await self._queries.enqueue(
job.entrypoint,
_encode({**payload, "trigger": "webhook"}),
dedupe_key=f"hook:{name}" if job.dedupe else None,
on_conflict="skip" if job.dedupe else "raise",
)
return int(ids[0]) if ids and ids[0] is not None else None
# ---- deferred injects ----------------------------------------------
async def schedule(
self,
conv: Conversation,
at: str,
text: str,
*,
urgency: Priority = "wake",
dedupe_key: str | None = None,
) -> tuple[int | None, datetime]:
"""A deferred inject.
``wake`` by default: it was promised for a time, so it starts a turn
then instead of waiting for the normal window.
"""
if self._queries is None:
msg = "scheduler needs postgres; `schedule` is unavailable"
raise RuntimeError(msg)
when = parse_at(at, self.tz).astimezone(self.tz)
delay = max(when - datetime.now(UTC), timedelta(0))
payload = _encode(
{
"conversation": conv.external_id,
"text": text,
"urgency": urgency,
"at": when.isoformat(timespec="minutes"),
}
)
ids = await self._queries.enqueue(
INJECT,
payload,
execute_after=delay,
dedupe_key=dedupe_key,
on_conflict="skip" if dedupe_key else "raise",
)
job_id = int(ids[0]) if ids and ids[0] is not None else None
self.conversations.bus.publish(
"schedule.created",
conversation_id=conv.external_id,
job=job_id,
execute_at=when.isoformat(timespec="minutes"),
)
return job_id, when
async def cancel(self, job_id: int) -> bool:
if self._queries is None:
return False
row = await self._queries.queue_job_by_id(JobId(job_id))
if row is None:
return False
await self._queries.mark_job_as_cancelled([JobId(job_id)])
self.conversations.bus.publish("schedule.cancelled", job=job_id)
return True
async def _deliver(self, job: PgJob) -> None:
data = _decode(job.payload)
key = str(data.get("conversation", ""))
conv = await self.conversations.resolve(key)
if conv is None:
_log.error(
"scheduled inject #%s lost: conversation %r is gone: %s",
job.id,
key,
str(data.get("text", ""))[:200],
)
return
await self.conversations.inject(
conv,
str(data.get("text", "")),
urgency=cast("Priority", data.get("urgency") or "wake"),
origin="schedule",
)
# ---- budget --------------------------------------------------------
async def utilization(self) -> float | None:
now = datetime.now(UTC)
values: list[float] = []
for row in await self.conversations.rate_limits(limit=100):
if row.window != self.budget.limit_window or row.utilization is None:
continue
fresh = now - _aware(row.ts) < self.budget.window
live = row.resets_at is None or _aware(row.resets_at) > now
if fresh and live:
values.append(float(row.utilization))
break
if self.budget.tokens:
tokens = await self.conversations.usage_tokens(now - self.budget.window)
values.append(tokens / self.budget.tokens)
return max(values) if values else None
async def throttled(self) -> bool:
utilization = await self.utilization()
return utilization is not None and utilization > self.budget.threshold
# ---- introspection -------------------------------------------------
async def snapshot(self) -> dict[str, Any]:
crons: dict[str, PgSchedule] = {}
queue: list[dict[str, Any]] = []
if self._queries is not None:
for row in await self._queries.peek_schedule():
crons[str(row.entrypoint)] = row
queue = [
_job_public(row, self.tz)
for row in await self._queries.browse_queue(limit=200)
]
utilization = await self.utilization()
last = await self._last_runs()
return {
"enabled": self.enabled,
"tz": str(self.tz),
"utilization": utilization,
"throttled": utilization is not None
and utilization > self.budget.threshold,
"threshold": self.budget.threshold,
"jobs": [
self._job_public(job, crons.get(job.name), last.get(job.name))
for job in self.jobs
],
"queue": queue,
}
async def scheduled(self, conv: Conversation | None = None) -> list[dict[str, Any]]:
if self._queries is None:
return []
rows = await self._queries.browse_queue(limit=500, entrypoints=[INJECT])
out = [_job_public(row, self.tz) for row in rows]
if conv is not None:
out = [
j for j in out if j["payload"].get("conversation") == conv.external_id
]
return out
async def runs(self, name: str, limit: int = 50) -> list[dict[str, Any]]:
stmt = (
select(JobRunRecord)
.where(JobRunRecord.job == name)
.order_by(col(JobRunRecord.id).desc())
.limit(limit)
)
async with self.conversations.db.session() as session:
rows = (await session.exec(stmt)).all()
return [self._run_public(row) for row in rows]
async def _last_runs(self) -> dict[str, dict[str, Any]]:
newest = select(func.max(JobRunRecord.id)).group_by(JobRunRecord.job)
stmt = select(JobRunRecord).where(col(JobRunRecord.id).in_(newest))
try:
async with self.conversations.db.session() as session:
rows = (await session.exec(stmt)).all()
except Exception: # noqa: BLE001
_log.exception("job runs unreadable, using memory")
return dict(self._runs)
return {row.job: self._run_public(row) for row in rows}
def _job_public(
self, job: Job, cron: PgSchedule | None, run: dict[str, Any] | None
) -> dict[str, Any]:
return {
"name": job.name,
"cron": job.cron,
"webhook": job.webhook,
"events": list(job.events),
"critical": job.critical,
"next_run": _iso(cron.next_run, self.tz) if cron is not None else None,
"last_run": _iso(cron.last_run, self.tz) if cron is not None else None,
"status": str(cron.status) if cron is not None else None,
"run": run or self._runs.get(job.name),
}
def _run_public(self, row: JobRunRecord) -> dict[str, Any]:
return {
"id": row.id,
"job": row.job,
"trigger": row.trigger,
"started_at": _iso(row.started_at, self.tz),
"finished_at": _iso(row.finished_at, self.tz),
"status": row.status,
"error": row.error,
"payload": row.payload,
}
# ---- http ----------------------------------------------------------
def app(self, authorize: Callable[[Request], Awaitable[Any]]) -> Starlette:
async def hook(request: Request) -> JSONResponse:
await authorize(request)
name = request.path_params["name"]
payload = await _payload(request)
try:
job_id = await self.hook(name, payload)
except LookupError as exc:
return JSONResponse({"error": str(exc)}, status_code=404)
return JSONResponse({"job": job_id, "name": name}, status_code=202)
return Starlette(routes=[Route("/{name}", hook, methods=["POST"])])
# ---- internals -----------------------------------------------------
def background(self, coro: Coroutine[Any, Any, Any]) -> None:
self._spawn(coro)
def _spawn(
self, coro: Coroutine[Any, Any, Any], *, delay: timedelta | None = None
) -> None:
async def later() -> None:
if delay is not None:
await asyncio.sleep(delay.total_seconds())
await coro
task = asyncio.create_task(later() if delay is not None else coro)
self._tasks.add(task)
task.add_done_callback(self._tasks.discard)
async def _payload(request: Request) -> dict[str, Any]:
body = await request.body()
if not body:
return dict(request.query_params)
try:
data = json.loads(body)
except ValueError:
return {"raw": body.decode("utf-8", errors="replace")}
return data if isinstance(data, dict) else {"raw": data}
def _encode(data: dict[str, Any]) -> bytes:
return json.dumps(data, ensure_ascii=False, default=str).encode("utf-8")
def _decode(payload: bytes | None) -> dict[str, Any]:
if not payload:
return {}
try:
data = json.loads(payload)
except ValueError:
return {"raw": payload.decode("utf-8", errors="replace")}
return data if isinstance(data, dict) else {"raw": data}
def _job_public(row: PgJob, tz: tzinfo) -> dict[str, Any]:
return {
"id": int(row.id),
"entrypoint": row.entrypoint,
"status": str(row.status),
"execute_after": _iso(row.execute_after, tz),
"created": _iso(row.created, tz),
"attempts": row.attempts,
"payload": _decode(row.payload),
}
def _error_text(exc: BaseException) -> str:
head = f"{type(exc).__name__}: {exc}"
return f"{head}\n{''.join(traceback.format_exception(exc))}"[:ERROR_MAX]
def _aware(value: datetime) -> datetime:
return value if value.tzinfo is not None else value.replace(tzinfo=UTC)
def _iso(value: datetime | None, tz: tzinfo = UTC) -> str | None:
if value is None:
return None
return _aware(value).astimezone(tz).isoformat(timespec="seconds")