spb/company-atlas

1"""Retention (spec §2.1, §95–96): history is never deleted. The only prunable rows are bookkeeping:2- `observations` older than `retention_observations_days` that were `not_modified` / unchanged, are not referenced by a snapshot and are not3  the sensor's most recent observation (aggregate counts stay in `sensors.observation_count`; pruned totals are kept in settings_kv);4- `crawl_runs` older than `retention_crawl_runs_days`;5- finished `llm_jobs` older than `retention_llm_jobs_days`, summarised into `cost_ledger` (dimension `llm`, key `archived:<model>`) first.6Snapshots, changes, events, entities, metrics and series are never touched.7"""8from __future__ import annotations910import logging11from datetime import UTC, datetime, timedelta12from typing import Any1314from companyatlas.config import settings15from companyatlas.db import execute, fetch_all, fetch_val, jsonb, transaction16from companyatlas.services.periodic import periodic1718log = logging.getLogger(__name__)1920PRUNED_KEY = "retention:pruned"21BATCH = 5000222324async def prune_observations(*, dry_run: bool = False, batch: int = BATCH) -> int:25    cutoff = datetime.now(UTC) - timedelta(days=settings.retention_observations_days)26    total = 027    async with transaction() as conn:28        while True:29            ids = [r["id"] for r in await fetch_all(conn, """30                select o.id from observations o31                where o.fetched_at < :cutoff and (o.not_modified or (o.changed = false and o.failure_class is null))32                  and not exists (select 1 from snapshots s where s.observation_id = o.id)33                  and o.id <> (select o2.id from observations o2 where o2.sensor_id = o.sensor_id order by o2.fetched_at desc limit 1)34                limit :batch""", cutoff=cutoff, batch=batch)]35            if not ids or dry_run:36                total += len(ids)37                break38            await execute(conn, "delete from observations where id = any(cast(:ids as text[]))", ids=ids)39            total += len(ids)40            if len(ids) < batch:41                break42        if total and not dry_run:43            prev = await fetch_val(conn, "select value from settings_kv where key = :k", k=PRUNED_KEY)44            state = dict(prev) if isinstance(prev, dict) else {}45            state["observations"] = int(state.get("observations", 0)) + total46            state["last_run_at"] = datetime.now(UTC).isoformat()47            await execute(conn, "insert into settings_kv (key, value, updated_at) values (:k, cast(:v as jsonb), now()) on conflict (key) do update set value = excluded.value, updated_at = now()",48                          k=PRUNED_KEY, v=jsonb(state))49    return total505152async def prune_crawl_runs(*, dry_run: bool = False) -> int:53    cutoff = datetime.now(UTC) - timedelta(days=settings.retention_crawl_runs_days)54    async with transaction() as conn:55        n = int(await fetch_val(conn, "select count(*) from crawl_runs where started_at < :c", c=cutoff) or 0)56        if n and not dry_run:57            await execute(conn, "delete from crawl_runs where started_at < :c", c=cutoff)58    return n596061async def archive_llm_jobs(*, dry_run: bool = False) -> int:62    cutoff = datetime.now(UTC) - timedelta(days=settings.retention_llm_jobs_days)63    async with transaction() as conn:64        rows = await fetch_all(conn, """select coalesce(model, 'unknown') as model, count(*) as n, coalesce(sum(request_tokens), 0) + coalesce(sum(response_tokens), 0) as tokens65                                        from llm_jobs where status in ('done', 'failed') and finished_at < :c group by 1""", c=cutoff)66        n = sum(int(r["n"]) for r in rows)67        if n and not dry_run:68            for r in rows:69                await execute(conn, """insert into cost_ledger (day, dimension, key, units, cost_estimate) values (current_date, 'llm', :key, :units, 0)70                                       on conflict (day, dimension, key) do update set units = cost_ledger.units + excluded.units""",71                              key=f"archived:{r['model']}", units=float(r["tokens"]))72                await execute(conn, """insert into cost_ledger (day, dimension, key, units, cost_estimate) values (current_date, 'llm', :key, :units, 0)73                                       on conflict (day, dimension, key) do update set units = cost_ledger.units + excluded.units""",74                              key=f"archived-jobs:{r['model']}", units=float(r["n"]))75            await execute(conn, "delete from llm_jobs where status in ('done', 'failed') and finished_at < :c", c=cutoff)76    return n777879async def run_retention(*, dry_run: bool = False) -> dict[str, Any]:80    stats = {"observations": await prune_observations(dry_run=dry_run), "crawl_runs": await prune_crawl_runs(dry_run=dry_run),81             "llm_jobs": await archive_llm_jobs(dry_run=dry_run), "dry_run": dry_run}82    log.info("retention", extra=stats)83    return stats848586@periodic("retention", cron="50 3 * * *")87async def retention_task() -> None:88    await run_retention()899091__all__ = ["archive_llm_jobs", "prune_crawl_runs", "prune_observations", "run_retention"]92