a2d46a7c6e
Ruff autofix removed pre-existing F401 unused imports across adapters, context, compiler, extractor, scheduler, web, and main. Also reordered models.py so its SQLAlchemy imports sit at the top of the file (E402 was triggered by the inline `utcnow` helper definition). Added a [dependency-groups] dev block (pytest, pytest-cov, ruff) so the zero-check test gauntlet — which runs `pytest --cov --cov-report=term-missing` — can resolve its plugins without a manual `uv pip install pytest-cov`.
195 lines
7.4 KiB
Python
195 lines
7.4 KiB
Python
"""
|
||
Overnight scheduler — rate-limit smoother.
|
||
|
||
Spreads extraction calls across a configurable overnight window, enforcing
|
||
a minimum gap between calls and a cap on calls-per-hour.
|
||
|
||
Two throttle modes, picked from `config.extractor.backend`:
|
||
- CLI backend (default): throttle by request count. Max OAuth doesn't bill
|
||
per-token; the relevant limit is the 5h-window message cap and the weekly
|
||
cap, both of which are call-shaped. `scheduler.max_calls_per_hour` is the
|
||
knob.
|
||
- API backend: keep the legacy token budget — `scheduler.max_tokens_per_hour`.
|
||
|
||
Source-status fix vs the previous version: this loop processes TRANSCRIBED
|
||
sources and transitions them to ANALYZED (the prior code looked for ANALYZED
|
||
sources and re-ran extraction on them, which was a no-op + drift bug).
|
||
Idempotent: safe to restart mid-run.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import time
|
||
from datetime import datetime, time as dtime
|
||
|
||
from second_brain.config import Config
|
||
from second_brain.database import get_database
|
||
from second_brain.models import Source, SourceStatus, utcnow
|
||
|
||
|
||
class Scheduler:
|
||
"""Process pending sources inside a time window, respecting rate caps."""
|
||
|
||
def __init__(self, config: Config) -> None:
|
||
self.config = config
|
||
sched = config.scheduler
|
||
self.window_start: dtime = _parse_time(sched.get("window_start", "22:00"))
|
||
self.window_end: dtime = _parse_time(sched.get("window_end", "06:00"))
|
||
self.min_gap_seconds: int = int(sched.get("min_gap_seconds", 120))
|
||
|
||
ex_cfg = config.extractor
|
||
self.backend: str = ex_cfg["backend"]
|
||
|
||
# Throttle knobs (only the one matching the backend gets enforced)
|
||
self.max_calls_per_hour: int = int(sched.get("max_calls_per_hour", 30))
|
||
self.max_tokens_per_hour: int = int(sched.get("max_tokens_per_hour", 100_000))
|
||
|
||
def run(self, dry_run: bool = False) -> None:
|
||
"""Process TRANSCRIBED sources until window closes or queue empties."""
|
||
from second_brain.context.assembler import ContextAssembler
|
||
from second_brain.extractor.engine import ExtractionEngine
|
||
from second_brain.models import Extraction
|
||
|
||
db = get_database()
|
||
assembler = ContextAssembler(self.config)
|
||
ex_cfg = self.config.extractor
|
||
engine = ExtractionEngine(
|
||
backend=ex_cfg["backend"],
|
||
model=ex_cfg["model"],
|
||
api_key=self.config.anthropic_api_key,
|
||
cli_binary=ex_cfg["cli_binary"],
|
||
timeout_sec=ex_cfg["timeout_sec"],
|
||
)
|
||
|
||
calls_this_hour: int = 0
|
||
tokens_this_hour: int = 0
|
||
hour_start: float = time.time()
|
||
processed: int = 0
|
||
|
||
cap_desc = (
|
||
f"{self.max_calls_per_hour} calls/hr"
|
||
if self.backend == "cli"
|
||
else f"{self.max_tokens_per_hour:,} tokens/hr"
|
||
)
|
||
print(
|
||
f"[Scheduler] backend={self.backend} window "
|
||
f"{_fmt(self.window_start)}–{_fmt(self.window_end)} cap {cap_desc} "
|
||
f"gap {self.min_gap_seconds}s"
|
||
)
|
||
|
||
while self._in_window():
|
||
# Reset hourly buckets
|
||
elapsed = time.time() - hour_start
|
||
if elapsed >= 3600:
|
||
calls_this_hour = 0
|
||
tokens_this_hour = 0
|
||
hour_start = time.time()
|
||
elapsed = 0
|
||
|
||
# Pre-call budget check
|
||
if self._over_budget(calls_this_hour, tokens_this_hour):
|
||
wait = int(3600 - elapsed) + 5
|
||
print(f"[Scheduler] Hourly cap reached. Sleeping {wait}s…")
|
||
if not dry_run:
|
||
time.sleep(wait)
|
||
calls_this_hour = 0
|
||
tokens_this_hour = 0
|
||
hour_start = time.time()
|
||
continue
|
||
|
||
with db.session() as sess:
|
||
source = (
|
||
sess.query(Source)
|
||
.filter(Source.status == SourceStatus.TRANSCRIBED)
|
||
.order_by(Source.ingested_at.asc())
|
||
.first()
|
||
)
|
||
if source is None:
|
||
print("[Scheduler] No pending sources. Done.")
|
||
break
|
||
|
||
if dry_run:
|
||
print(f"[Scheduler] DRY RUN — would process source {source.id}")
|
||
break
|
||
|
||
try:
|
||
prompt = assembler.assemble(source)
|
||
result = engine.extract(prompt, source)
|
||
|
||
existing = sess.query(Extraction).filter_by(
|
||
source_id=source.id
|
||
).first()
|
||
fields = result.to_db_dict()
|
||
if existing:
|
||
for k, v in fields.items():
|
||
setattr(existing, k, v)
|
||
existing.raw_response = getattr(result, "_raw_response", None)
|
||
existing.input_tokens = getattr(result, "_input_tokens", None)
|
||
existing.output_tokens = getattr(result, "_output_tokens", None)
|
||
ext_row = existing
|
||
else:
|
||
ext_row = Extraction(
|
||
source_id=source.id,
|
||
raw_response=getattr(result, "_raw_response", None),
|
||
input_tokens=getattr(result, "_input_tokens", None),
|
||
output_tokens=getattr(result, "_output_tokens", None),
|
||
**fields,
|
||
)
|
||
sess.add(ext_row)
|
||
|
||
source.status = SourceStatus.ANALYZED
|
||
source.updated_at = utcnow()
|
||
|
||
# Flush so the embedding step can key on the new PK.
|
||
sess.flush()
|
||
if self.config.embeddings.get("enabled", True):
|
||
from second_brain.embeddings import embed_extraction
|
||
|
||
embed_extraction(ext_row.id, ext_row.summary)
|
||
|
||
calls_this_hour += 1
|
||
used = (getattr(result, "_input_tokens", 0) or 0) + (
|
||
getattr(result, "_output_tokens", 0) or 0
|
||
)
|
||
tokens_this_hour += used
|
||
processed += 1
|
||
|
||
except Exception as exc:
|
||
print(f"[Scheduler] Error on source {source.id}: {exc}")
|
||
source.error_message = str(exc)
|
||
source.status = SourceStatus.FAILED
|
||
|
||
time.sleep(self.min_gap_seconds)
|
||
|
||
print(f"[Scheduler] Finished. Processed {processed} source(s).")
|
||
|
||
# ------------------------------------------------------------------
|
||
# Helpers
|
||
# ------------------------------------------------------------------
|
||
|
||
def _over_budget(self, calls: int, tokens: int) -> bool:
|
||
if self.backend == "cli":
|
||
return calls + 1 > self.max_calls_per_hour
|
||
# API backend — assume ~2k tokens per call when projecting
|
||
return tokens + 2000 > self.max_tokens_per_hour
|
||
|
||
def _in_window(self) -> bool:
|
||
now = datetime.now().time().replace(second=0, microsecond=0)
|
||
start, end = self.window_start, self.window_end
|
||
if start <= end:
|
||
return start <= now <= end
|
||
# crosses midnight
|
||
return now >= start or now <= end
|
||
|
||
|
||
def _parse_time(s: str) -> dtime:
|
||
h, m = s.split(":")
|
||
return dtime(int(h), int(m))
|
||
|
||
|
||
def _fmt(t: dtime) -> str:
|
||
return t.strftime("%H:%M")
|
||
|
||
|
||
__all__ = ["Scheduler"]
|