127 lines
3.8 KiB
Python
127 lines
3.8 KiB
Python
"""What a job is: its triggers, the budget it respects, and what a run may do."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from dataclasses import dataclass
|
|
from datetime import datetime, timedelta
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Awaitable, Callable, Coroutine
|
|
|
|
from beaver_gateway.conversations.distill import LineCap
|
|
from beaver_gateway.conversations.injects import Priority
|
|
from beaver_gateway.conversations.service import Conversations, DistillResult
|
|
from beaver_gateway.jobs.scheduler import Scheduler
|
|
from beaver_gateway.storage.models import Conversation
|
|
|
|
__all__ = ["Budget", "Job", "JobRun"]
|
|
|
|
_log = logging.getLogger(__name__)
|
|
|
|
|
|
@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]:
|
|
"""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, *, force: str | None = None) -> list[Conversation]:
|
|
rotation = self.scheduler.rotation
|
|
return await rotation.tick(force=force) if rotation is not None else []
|
|
|
|
def background(self, coro: Coroutine[Any, Any, Any]) -> None:
|
|
self.scheduler.background(coro)
|