""" 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"]