From 64505e1057817938e2322ccfda227a8a4cf5e3e7 Mon Sep 17 00:00:00 2001 From: Prajna1999 Date: Thu, 10 Sep 2026 09:47:48 +0530 Subject: [PATCH 1/4] feat: SQL query for batchwise llm_call input/content redaction --- .env.example | 2 ++ backend/app/api/routes/cron.py | 42 ++++++++++++++++++++++ backend/app/core/config.py | 10 ++++++ backend/app/crud/llm.py | 47 ++++++++++++++++++++++++ backend/app/services/llm/retention.py | 51 +++++++++++++++++++++++++++ 5 files changed, 152 insertions(+) create mode 100644 backend/app/services/llm/retention.py diff --git a/.env.example b/.env.example index 5a041fe26..e5b74445d 100644 --- a/.env.example +++ b/.env.example @@ -110,3 +110,5 @@ SMTP_PASSWORD= EMAILS_FROM_EMAIL= EMAILS_FROM_NAME=Kaapi FRONTEND_HOST= + +DELETE_ROLLING_WINDOW_HOURS= diff --git a/backend/app/api/routes/cron.py b/backend/app/api/routes/cron.py index 9e11d033a..d2ccd09a7 100644 --- a/backend/app/api/routes/cron.py +++ b/backend/app/api/routes/cron.py @@ -9,6 +9,7 @@ from app.core.config import settings from app.crud.evaluations import process_all_pending_evaluations from app.services.job_monitoring import monitor_pending_jobs +from app.services.llm.retention import redact_aged_llm_calls from app.crud.stats import StatRow, get_daily_stats from app.services.stats import format_sections, post_to_discord @@ -59,6 +60,17 @@ } +LLM_CALL_RETENTION_CRON_MONITOR_CONFIG: MonitorConfig = { + "schedule": {"type": "crontab", "value": "0 3 * * *"}, + "timezone": "UTC", + "checkin_margin": 5, + # Generous ceiling: with no index on updated_at the batch scan is full-table. + "max_runtime": 30, + "failure_issue_threshold": 1, + "recovery_threshold": 1, +} + + @router.get( "/cron/evaluations", include_in_schema=False, @@ -158,6 +170,36 @@ def daily_stats_cron_job(session: SessionDep) -> dict[str, list[StatRow]]: raise +@router.get( + "/cron/llm-call-retention", + include_in_schema=False, + dependencies=[Depends(require_permission(Permission.SUPERUSER))], +) +@sentry_sdk.monitor( + monitor_slug="llm-call-retention-cron-job", + monitor_config=LLM_CALL_RETENTION_CRON_MONITOR_CONFIG, +) +def llm_call_retention_cron_job(session: SessionDep) -> dict: + logger.info("[llm_call_retention_cron_job] Cron job invoked") + + try: + result = redact_aged_llm_calls(session=session) + logger.info( + f"[llm_call_retention_cron_job] Completed: " + f"rows_redacted={result['rows_redacted']}, " + f"batches_run={result['batches_run']}, " + f"cutoff={result['cutoff']}" + ) + return result + except Exception as e: + logger.error( + f"[llm_call_retention_cron_job] Error executing cron job: {e}", + exc_info=True, + ) + sentry_sdk.capture_exception(e) + raise + + @router.get( "/cron/pending-jobs", include_in_schema=False, diff --git a/backend/app/core/config.py b/backend/app/core/config.py index cc105fbec..18b6ecd0d 100644 --- a/backend/app/core/config.py +++ b/backend/app/core/config.py @@ -2,6 +2,7 @@ import os import secrets import warnings +from datetime import timedelta from typing import Any, Literal, Self from pydantic import ( @@ -196,6 +197,15 @@ def AWS_S3_BUCKET(self) -> str: EVAL_FAST_STALL_THRESHOLD_MINUTES: int = 15 PENDING_JOB_QUERY_TIMEOUT_MS: int = 1000 + # Rolling retention window for llm_call payload redaction: rows older than this + # keep their metadata but lose `input` and the inner `content` value. + DELETE_ROLLING_WINDOW_HOURS: int = 168 + + @computed_field # type: ignore[prop-decorator] + @property + def DELETE_ROLLING_WINDOW_TIMEDELTA(self) -> timedelta: + return timedelta(hours=self.DELETE_ROLLING_WINDOW_HOURS) + # AI-assisted prompt improvement settings. # See docs/srd-ai-prompt-improvement.md for the full design rationale. # Platform-owned Anthropic key shared by every org/project for this feature, diff --git a/backend/app/crud/llm.py b/backend/app/crud/llm.py index 44e605fb7..6724c046c 100644 --- a/backend/app/crud/llm.py +++ b/backend/app/crud/llm.py @@ -1,9 +1,11 @@ import logging import base64 import json +from datetime import datetime from uuid import UUID from typing import Any, Literal +from sqlalchemy import text from sqlmodel import Session, select from app.core.util import now @@ -335,6 +337,51 @@ def get_llm_call_by_id( return session.exec(statement).first() +# `llm_call.input` is NOT NULL at the DB level, so redaction writes this +# sentinel rather than NULL (which would raise NotNullViolation on every row). +REDACTED_INPUT_SENTINEL = "[redacted]" + +REDACT_LLM_CALL_BATCH_SQL = text( + """ + WITH batch AS ( + SELECT id FROM llm_call + WHERE updated_at <= :cutoff + -- Skips already-redacted rows, so repeated runs are idempotent. + AND (input <> :redacted_sentinel OR content -> 'content' ->> 'value' IS NOT NULL) + ORDER BY updated_at + LIMIT :batch_size + FOR UPDATE SKIP LOCKED + ) + UPDATE llm_call + SET input = :redacted_sentinel, + -- create_missing=false so legacy rows lacking that path aren't given a fabricated one. + content = jsonb_set(content, '{content,value}', 'null'::jsonb, false) + WHERE id IN (SELECT id FROM batch) + """ +) + + +def redact_llm_call_batch( + *, session: Session, cutoff: datetime, batch_size: int +) -> int: + result = session.connection().execute( + REDACT_LLM_CALL_BATCH_SQL, + { + "cutoff": cutoff, + "batch_size": batch_size, + "redacted_sentinel": REDACTED_INPUT_SENTINEL, + }, + ) + session.commit() + + logger.info( + f"[redact_llm_call_batch] Redacted batch | rows: {result.rowcount} | " + f"cutoff: {cutoff.isoformat()} | batch_size: {batch_size}" + ) + + return result.rowcount + + def get_llm_calls_by_job_id( session: Session, job_id: UUID, project_id: int ) -> list[LlmCall]: diff --git a/backend/app/services/llm/retention.py b/backend/app/services/llm/retention.py new file mode 100644 index 000000000..bd2b1ca82 --- /dev/null +++ b/backend/app/services/llm/retention.py @@ -0,0 +1,51 @@ +import logging + +from sqlmodel import Session + +from app.core.config import settings +from app.core.util import now +from app.crud.llm import redact_llm_call_batch + +logger = logging.getLogger(__name__) + +# Rows updated per statement; small enough to keep each row-lock window short. +LLM_CALL_REDACTION_BATCH_SIZE = 2000 + + +def redact_aged_llm_calls(*, session: Session) -> dict[str, int | str]: + cutoff = now() - settings.DELETE_ROLLING_WINDOW_TIMEDELTA + + logger.info( + f"[redact_aged_llm_calls] Starting redaction | cutoff: {cutoff.isoformat()} | " + f"batch_size: {LLM_CALL_REDACTION_BATCH_SIZE}" + ) + + total_redacted = 0 + batches_run = 0 + + while True: + redacted = redact_llm_call_batch( + session=session, + cutoff=cutoff, + batch_size=LLM_CALL_REDACTION_BATCH_SIZE, + ) + if redacted == 0: + break + + total_redacted += redacted + batches_run += 1 + logger.info( + f"[redact_aged_llm_calls] Batch complete | batch: {batches_run} | " + f"rows: {redacted} | total_rows: {total_redacted}" + ) + + logger.info( + f"[redact_aged_llm_calls] Completed | rows_redacted: {total_redacted} | " + f"batches_run: {batches_run} | cutoff: {cutoff.isoformat()}" + ) + + return { + "rows_redacted": total_redacted, + "batches_run": batches_run, + "cutoff": cutoff.isoformat(), + } From afa509c9f6ff0a802c30423ae5d2076771398bc0 Mon Sep 17 00:00:00 2001 From: Prajna1999 Date: Thu, 10 Sep 2026 09:59:07 +0530 Subject: [PATCH 2/4] test cases --- backend/app/api/routes/cron.py | 10 +- backend/app/core/config.py | 2 - backend/app/crud/llm.py | 10 +- backend/app/models/llm/response.py | 8 + backend/app/services/llm/retention.py | 14 +- backend/app/tests/api/routes/test_cron.py | 60 +++++++ backend/app/tests/crud/test_llm_retention.py | 152 ++++++++++++++++++ .../app/tests/services/llm/test_retention.py | 91 +++++++++++ backend/app/tests/utils/llm.py | 42 ++++- 9 files changed, 369 insertions(+), 20 deletions(-) create mode 100644 backend/app/tests/crud/test_llm_retention.py create mode 100644 backend/app/tests/services/llm/test_retention.py diff --git a/backend/app/api/routes/cron.py b/backend/app/api/routes/cron.py index d2ccd09a7..defbff177 100644 --- a/backend/app/api/routes/cron.py +++ b/backend/app/api/routes/cron.py @@ -8,6 +8,7 @@ from app.api.permissions import Permission, require_permission from app.core.config import settings from app.crud.evaluations import process_all_pending_evaluations +from app.models.llm.response import LlmCallRedactionResult from app.services.job_monitoring import monitor_pending_jobs from app.services.llm.retention import redact_aged_llm_calls from app.crud.stats import StatRow, get_daily_stats @@ -64,7 +65,6 @@ "schedule": {"type": "crontab", "value": "0 3 * * *"}, "timezone": "UTC", "checkin_margin": 5, - # Generous ceiling: with no index on updated_at the batch scan is full-table. "max_runtime": 30, "failure_issue_threshold": 1, "recovery_threshold": 1, @@ -179,16 +179,16 @@ def daily_stats_cron_job(session: SessionDep) -> dict[str, list[StatRow]]: monitor_slug="llm-call-retention-cron-job", monitor_config=LLM_CALL_RETENTION_CRON_MONITOR_CONFIG, ) -def llm_call_retention_cron_job(session: SessionDep) -> dict: +def llm_call_retention_cron_job(session: SessionDep) -> LlmCallRedactionResult: logger.info("[llm_call_retention_cron_job] Cron job invoked") try: result = redact_aged_llm_calls(session=session) logger.info( f"[llm_call_retention_cron_job] Completed: " - f"rows_redacted={result['rows_redacted']}, " - f"batches_run={result['batches_run']}, " - f"cutoff={result['cutoff']}" + f"rows_redacted={result.rows_redacted}, " + f"batches_run={result.batches_run}, " + f"cutoff={result.cutoff}" ) return result except Exception as e: diff --git a/backend/app/core/config.py b/backend/app/core/config.py index 18b6ecd0d..104e2c7b7 100644 --- a/backend/app/core/config.py +++ b/backend/app/core/config.py @@ -197,8 +197,6 @@ def AWS_S3_BUCKET(self) -> str: EVAL_FAST_STALL_THRESHOLD_MINUTES: int = 15 PENDING_JOB_QUERY_TIMEOUT_MS: int = 1000 - # Rolling retention window for llm_call payload redaction: rows older than this - # keep their metadata but lose `input` and the inner `content` value. DELETE_ROLLING_WINDOW_HOURS: int = 168 @computed_field # type: ignore[prop-decorator] diff --git a/backend/app/crud/llm.py b/backend/app/crud/llm.py index 6724c046c..c017874da 100644 --- a/backend/app/crud/llm.py +++ b/backend/app/crud/llm.py @@ -337,8 +337,6 @@ def get_llm_call_by_id( return session.exec(statement).first() -# `llm_call.input` is NOT NULL at the DB level, so redaction writes this -# sentinel rather than NULL (which would raise NotNullViolation on every row). REDACTED_INPUT_SENTINEL = "[redacted]" REDACT_LLM_CALL_BATCH_SQL = text( @@ -346,7 +344,6 @@ def get_llm_call_by_id( WITH batch AS ( SELECT id FROM llm_call WHERE updated_at <= :cutoff - -- Skips already-redacted rows, so repeated runs are idempotent. AND (input <> :redacted_sentinel OR content -> 'content' ->> 'value' IS NOT NULL) ORDER BY updated_at LIMIT :batch_size @@ -354,8 +351,11 @@ def get_llm_call_by_id( ) UPDATE llm_call SET input = :redacted_sentinel, - -- create_missing=false so legacy rows lacking that path aren't given a fabricated one. - content = jsonb_set(content, '{content,value}', 'null'::jsonb, false) + content = CASE + WHEN jsonb_typeof(content) = 'object' + THEN jsonb_set(content, '{content,value}', 'null'::jsonb, false) + ELSE content + END WHERE id IN (SELECT id FROM batch) """ ) diff --git a/backend/app/models/llm/response.py b/backend/app/models/llm/response.py index 439a6adfc..70bd76c83 100644 --- a/backend/app/models/llm/response.py +++ b/backend/app/models/llm/response.py @@ -19,6 +19,14 @@ class Usage(SQLModel): reasoning_tokens: int | None = None +class LlmCallRedactionResult(SQLModel): + """Outcome of a single llm_call retention redaction run.""" + + rows_redacted: int + batches_run: int + cutoff: datetime + + class TextOutput(SQLModel): type: Literal["text"] = "text" content: TextContent diff --git a/backend/app/services/llm/retention.py b/backend/app/services/llm/retention.py index bd2b1ca82..f3bd8046a 100644 --- a/backend/app/services/llm/retention.py +++ b/backend/app/services/llm/retention.py @@ -5,14 +5,14 @@ from app.core.config import settings from app.core.util import now from app.crud.llm import redact_llm_call_batch +from app.models.llm.response import LlmCallRedactionResult logger = logging.getLogger(__name__) -# Rows updated per statement; small enough to keep each row-lock window short. LLM_CALL_REDACTION_BATCH_SIZE = 2000 -def redact_aged_llm_calls(*, session: Session) -> dict[str, int | str]: +def redact_aged_llm_calls(*, session: Session) -> LlmCallRedactionResult: cutoff = now() - settings.DELETE_ROLLING_WINDOW_TIMEDELTA logger.info( @@ -44,8 +44,8 @@ def redact_aged_llm_calls(*, session: Session) -> dict[str, int | str]: f"batches_run: {batches_run} | cutoff: {cutoff.isoformat()}" ) - return { - "rows_redacted": total_redacted, - "batches_run": batches_run, - "cutoff": cutoff.isoformat(), - } + return LlmCallRedactionResult( + rows_redacted=total_redacted, + batches_run=batches_run, + cutoff=cutoff, + ) diff --git a/backend/app/tests/api/routes/test_cron.py b/backend/app/tests/api/routes/test_cron.py index 399ed6eee..32fd4f471 100644 --- a/backend/app/tests/api/routes/test_cron.py +++ b/backend/app/tests/api/routes/test_cron.py @@ -1,3 +1,4 @@ +from datetime import datetime from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -5,6 +6,7 @@ from app.api.routes import cron from app.core.config import settings +from app.models.llm.response import LlmCallRedactionResult from app.tests.utils.auth import TestAuthContext @@ -326,6 +328,63 @@ def test_daily_stats_cron_job_captures_and_reraises_on_error() -> None: sentry.capture_exception.assert_called_once() +def test_llm_call_retention_cron_job_success( + client: TestClient, + superuser_api_key: TestAuthContext, +) -> None: + """Returns the redaction summary produced by the retention service.""" + result = LlmCallRedactionResult( + rows_redacted=4137, + batches_run=3, + cutoff=datetime(2026, 2, 22, 12, 0, 0), + ) + with patch( + "app.api.routes.cron.redact_aged_llm_calls", + return_value=result, + ) as redact: + response = client.get( + f"{settings.API_V1_STR}/cron/llm-call-retention", + headers={"X-API-KEY": superuser_api_key.key}, + ) + + assert response.status_code == 200 + assert response.json() == { + "rows_redacted": 4137, + "batches_run": 3, + "cutoff": "2026-02-22T12:00:00", + } + redact.assert_called_once() + + +def test_llm_call_retention_cron_job_requires_superuser( + client: TestClient, + user_api_key: TestAuthContext, +) -> None: + """Non-superuser cannot access the llm_call retention cron endpoint.""" + response = client.get( + f"{settings.API_V1_STR}/cron/llm-call-retention", + headers={"X-API-KEY": user_api_key.key}, + ) + + assert response.status_code == 403 + assert "Insufficient permissions" in response.json()["error"] + + +def test_llm_call_retention_cron_job_captures_and_reraises_on_error() -> None: + """On failure the job reports to Sentry and re-raises.""" + with ( + patch( + "app.api.routes.cron.redact_aged_llm_calls", + side_effect=RuntimeError("boom"), + ), + patch("app.api.routes.cron.sentry_sdk") as sentry, + ): + with pytest.raises(RuntimeError): + cron.llm_call_retention_cron_job(session=MagicMock()) + + sentry.capture_exception.assert_called_once() + + def test_evaluation_cron_job_not_in_schema( client: TestClient, ) -> None: @@ -340,6 +399,7 @@ def test_evaluation_cron_job_not_in_schema( assert f"{settings.API_V1_STR}/cron/evaluations" not in paths assert f"{settings.API_V1_STR}/cron/pending-jobs" not in paths assert f"{settings.API_V1_STR}/cron/daily-stats" not in paths + assert f"{settings.API_V1_STR}/cron/llm-call-retention" not in paths def test_cron_intervals_match_to_prevent_sentry_monitor_drift() -> None: diff --git a/backend/app/tests/crud/test_llm_retention.py b/backend/app/tests/crud/test_llm_retention.py new file mode 100644 index 000000000..787d9b895 --- /dev/null +++ b/backend/app/tests/crud/test_llm_retention.py @@ -0,0 +1,152 @@ +from datetime import datetime, timedelta + +import pytest +from sqlalchemy import text +from sqlmodel import Session + +from app.models import Job +from app.crud.llm import REDACTED_INPUT_SENTINEL, redact_llm_call_batch +from app.models.llm import LlmCall +from app.tests.utils.llm import create_aged_llm_call, create_llm_job + +CUTOFF = datetime(2026, 1, 8, 0, 0, 0) +AGED = CUTOFF - timedelta(days=1) +RECENT = CUTOFF + timedelta(minutes=1) + +TEXT_CONTENT = {"type": "text", "content": {"format": "text", "value": "hello"}} + + +@pytest.fixture +def job(db: Session) -> Job: + # Other llm_call rows (seed data, sibling fixtures) would also match the + # cutoff and skew the rowcount assertions, so start from an empty table. + db.exec(text("DELETE FROM llm_call")) + db.commit() + return create_llm_job(db) + + +def _reload(db: Session, llm_call: LlmCall) -> LlmCall: + db.expire_all() + return db.get(LlmCall, llm_call.id) + + +def test_redacts_input_and_content_value_only(db: Session, job: Job) -> None: + llm_call = create_aged_llm_call( + db, + updated_at=AGED, + job_id=job.id, + input="my social security number is 123-45-6789", + content=TEXT_CONTENT, + ) + + redacted = redact_llm_call_batch(session=db, cutoff=CUTOFF, batch_size=10) + + assert redacted == 1 + row = _reload(db, llm_call) + assert row.input == "[redacted]" + assert row.content["content"]["value"] is None + assert row.content["content"]["format"] == "text" + assert row.content["type"] == "text" + + +def test_leaves_rows_newer_than_cutoff_untouched(db: Session, job: Job) -> None: + llm_call = create_aged_llm_call( + db, + updated_at=RECENT, + job_id=job.id, + input="still within the retention window", + content=TEXT_CONTENT, + ) + + redacted = redact_llm_call_batch(session=db, cutoff=CUTOFF, batch_size=10) + + assert redacted == 0 + row = _reload(db, llm_call) + assert row.input == "still within the retention window" + assert row.content["content"]["value"] == "hello" + + +def test_row_exactly_at_cutoff_is_redacted(db: Session, job: Job) -> None: + llm_call = create_aged_llm_call( + db, + updated_at=CUTOFF, + job_id=job.id, + content=TEXT_CONTENT, + ) + + assert redact_llm_call_batch(session=db, cutoff=CUTOFF, batch_size=10) == 1 + assert _reload(db, llm_call).input == "[redacted]" + + +def test_second_run_skips_already_redacted_rows(db: Session, job: Job) -> None: + create_aged_llm_call(db, updated_at=AGED, job_id=job.id, content=TEXT_CONTENT) + + assert redact_llm_call_batch(session=db, cutoff=CUTOFF, batch_size=10) == 1 + assert redact_llm_call_batch(session=db, cutoff=CUTOFF, batch_size=10) == 0 + + +def test_row_with_null_content_is_redacted_without_error(db: Session, job: Job) -> None: + llm_call = create_aged_llm_call( + db, + updated_at=AGED, + job_id=job.id, + input="no response was ever recorded", + content=None, + ) + # A call that never got a response stores the JSONB scalar 'null', not SQL + # NULL; jsonb_set errors on scalars, so this must take a guarded path. + assert ( + db.exec( + text( + "SELECT jsonb_typeof(content) FROM llm_call WHERE id = :id" + ).bindparams(id=llm_call.id) + ).scalar_one() + == "null" + ) + + assert redact_llm_call_batch(session=db, cutoff=CUTOFF, batch_size=10) == 1 + row = _reload(db, llm_call) + assert row.input == "[redacted]" + assert row.content is None + + +def test_batch_size_caps_rows_touched_per_call(db: Session, job: Job) -> None: + for _ in range(5): + create_aged_llm_call(db, updated_at=AGED, job_id=job.id, content=TEXT_CONTENT) + + assert redact_llm_call_batch(session=db, cutoff=CUTOFF, batch_size=2) == 2 + assert redact_llm_call_batch(session=db, cutoff=CUTOFF, batch_size=2) == 2 + assert redact_llm_call_batch(session=db, cutoff=CUTOFF, batch_size=2) == 1 + assert redact_llm_call_batch(session=db, cutoff=CUTOFF, batch_size=2) == 0 + + +def test_returns_zero_when_no_rows_match(db: Session, job: Job) -> None: + assert redact_llm_call_batch(session=db, cutoff=CUTOFF, batch_size=10) == 0 + + +def test_redacts_only_aged_rows_in_a_mixed_table(db: Session, job: Job) -> None: + aged = create_aged_llm_call( + db, updated_at=AGED, job_id=job.id, input="aged", content=TEXT_CONTENT + ) + recent = create_aged_llm_call( + db, updated_at=RECENT, job_id=job.id, input="recent", content=TEXT_CONTENT + ) + + assert redact_llm_call_batch(session=db, cutoff=CUTOFF, batch_size=10) == 1 + assert _reload(db, aged).input == REDACTED_INPUT_SENTINEL + assert _reload(db, recent).input == "recent" + + +def test_redacts_content_value_of_row_whose_input_is_already_sentinel( + db: Session, job: Job +) -> None: + llm_call = create_aged_llm_call( + db, + updated_at=AGED, + job_id=job.id, + input=REDACTED_INPUT_SENTINEL, + content={"type": "audio", "content": {"format": "uri", "value": "s3://a.wav"}}, + ) + + assert redact_llm_call_batch(session=db, cutoff=CUTOFF, batch_size=10) == 1 + assert _reload(db, llm_call).content["content"]["value"] is None diff --git a/backend/app/tests/services/llm/test_retention.py b/backend/app/tests/services/llm/test_retention.py new file mode 100644 index 000000000..7a9963218 --- /dev/null +++ b/backend/app/tests/services/llm/test_retention.py @@ -0,0 +1,91 @@ +from datetime import datetime, timedelta + +import pytest +from sqlmodel import Session + +from app.core.config import settings +from app.services.llm import retention +from app.services.llm.retention import ( + LLM_CALL_REDACTION_BATCH_SIZE, + redact_aged_llm_calls, +) + + +@pytest.fixture +def batch_calls(monkeypatch: pytest.MonkeyPatch) -> list[dict]: + """Records the kwargs of each redact_llm_call_batch call; returns 0 by default.""" + calls: list[dict] = [] + + def fake_batch(**kwargs) -> int: + calls.append(kwargs) + return 0 + + monkeypatch.setattr(retention, "redact_llm_call_batch", fake_batch) + return calls + + +def _stub_returns( + monkeypatch: pytest.MonkeyPatch, returns: list[int], calls: list[dict] +) -> None: + remaining = list(returns) + + def fake_batch(**kwargs) -> int: + calls.append(kwargs) + return remaining.pop(0) + + monkeypatch.setattr(retention, "redact_llm_call_batch", fake_batch) + + +def test_sums_rows_across_batches_until_a_batch_is_empty( + db: Session, monkeypatch: pytest.MonkeyPatch +) -> None: + calls: list[dict] = [] + _stub_returns(monkeypatch, [2000, 2000, 137, 0], calls) + + result = redact_aged_llm_calls(session=db) + + assert result.rows_redacted == 4137 + assert result.batches_run == 3 + assert len(calls) == 4 + + +def test_no_matching_rows(db: Session, batch_calls: list[dict]) -> None: + result = redact_aged_llm_calls(session=db) + + assert result.rows_redacted == 0 + assert result.batches_run == 0 + assert len(batch_calls) == 1 + + +def test_cutoff_is_naive_utc(db: Session, batch_calls: list[dict]) -> None: + # llm_call.updated_at is a naive TIMESTAMP; a tz-aware cutoff would raise on compare. + result = redact_aged_llm_calls(session=db) + + cutoff = batch_calls[0]["cutoff"] + assert cutoff.tzinfo is None + assert result.cutoff == cutoff + + +def test_cutoff_trails_now_by_the_configured_rolling_window( + db: Session, batch_calls: list[dict], monkeypatch: pytest.MonkeyPatch +) -> None: + frozen = datetime(2026, 3, 1, 12, 0, 0) + monkeypatch.setattr(retention, "now", lambda: frozen) + monkeypatch.setattr(settings, "DELETE_ROLLING_WINDOW_HOURS", 168) + + redact_aged_llm_calls(session=db) + + assert batch_calls[0]["cutoff"] == datetime(2026, 2, 22, 12, 0, 0) + + +def test_passes_session_and_module_batch_size_to_crud( + db: Session, batch_calls: list[dict] +) -> None: + redact_aged_llm_calls(session=db) + + assert batch_calls[0]["session"] is db + assert batch_calls[0]["batch_size"] == LLM_CALL_REDACTION_BATCH_SIZE == 2000 + + +def test_default_rolling_window_is_one_week() -> None: + assert settings.DELETE_ROLLING_WINDOW_TIMEDELTA == timedelta(days=7) diff --git a/backend/app/tests/utils/llm.py b/backend/app/tests/utils/llm.py index a60059b87..da84528c2 100644 --- a/backend/app/tests/utils/llm.py +++ b/backend/app/tests/utils/llm.py @@ -1,8 +1,13 @@ +from datetime import datetime +from typing import Any +from uuid import UUID + from sqlmodel import Session from app.crud import JobCrud from app.crud.llm import create_llm_call, update_llm_call_response from app.models import JobType, Job +from app.models.llm import LlmCall from app.models.llm.response import LLMCallResponse from app.models.llm.request import ( ConfigBlob, @@ -10,10 +15,45 @@ QueryParams, build_kaapi_completion_config, ) -from app.tests.utils.utils import get_project +from app.tests.utils.utils import get_project, get_organization from app.models.llm import LLMCallRequest +def create_aged_llm_call( + db: Session, + *, + updated_at: datetime, + job_id: UUID, + input: str = "what is the capital of France?", + content: dict[str, Any] | None = None, +) -> LlmCall: + """Persist an LlmCall with an explicit `updated_at`, for retention tests.""" + project = get_project(db, "Dalgo") + organization = get_organization(db) + + llm_call = LlmCall( + job_id=job_id, + project_id=project.id, + organization_id=organization.id, + input=input, + input_type="text", + output_type="text", + provider="openai", + model="gpt-4o", + content=content, + ) + db.add(llm_call) + db.commit() + + # updated_at has a server/model default, so it is overwritten after insert. + llm_call.updated_at = updated_at + db.add(llm_call) + db.commit() + db.refresh(llm_call) + + return llm_call + + def create_llm_job(db: Session) -> Job: """Create a persisted LLM_API job for use in tests.""" project = get_project(db, "Dalgo") From da21bd8e9104e7243d45ed9452071ec78bd05663 Mon Sep 17 00:00:00 2001 From: Prajna1999 Date: Fri, 11 Sep 2026 10:05:38 +0530 Subject: [PATCH 3/4] fix: tz from UTC to IST, remove redudant comments --- backend/app/api/routes/cron.py | 9 ++------- backend/app/core/config.py | 4 ---- 2 files changed, 2 insertions(+), 11 deletions(-) diff --git a/backend/app/api/routes/cron.py b/backend/app/api/routes/cron.py index defbff177..2c9c1c3bb 100644 --- a/backend/app/api/routes/cron.py +++ b/backend/app/api/routes/cron.py @@ -25,15 +25,10 @@ "value": settings.CRON_INTERVAL_MINUTES, "unit": "minute", }, - # Timezone for the schedule (only affects crontab-style schedules). "timezone": "UTC", - # Grace period (minutes) before a late check-in is marked as missed. "checkin_margin": 2, - # Max runtime (minutes) before an in-progress run is marked as timed out. "max_runtime": 2 * settings.CRON_INTERVAL_MINUTES, - # Consecutive failures/missed/timeouts required to open a Sentry issue. "failure_issue_threshold": 2, - # Consecutive successful check-ins required to auto-resolve the issue. "recovery_threshold": 1, } @@ -62,8 +57,8 @@ LLM_CALL_RETENTION_CRON_MONITOR_CONFIG: MonitorConfig = { - "schedule": {"type": "crontab", "value": "0 3 * * *"}, - "timezone": "UTC", + "schedule": {"type": "crontab", "value": "0 9 * * *"}, + "timezone": "Asia/Kolkata", "checkin_margin": 5, "max_runtime": 30, "failure_issue_threshold": 1, diff --git a/backend/app/core/config.py b/backend/app/core/config.py index 104e2c7b7..10b652bbe 100644 --- a/backend/app/core/config.py +++ b/backend/app/core/config.py @@ -204,10 +204,6 @@ def AWS_S3_BUCKET(self) -> str: def DELETE_ROLLING_WINDOW_TIMEDELTA(self) -> timedelta: return timedelta(hours=self.DELETE_ROLLING_WINDOW_HOURS) - # AI-assisted prompt improvement settings. - # See docs/srd-ai-prompt-improvement.md for the full design rationale. - # Platform-owned Anthropic key shared by every org/project for this feature, - # so prompt improvement works without per-project credentials. ANTHROPIC_API_KEY: str = "" PROMPT_IMPROVEMENT_MODEL: str = "claude-opus-4-8" From 49afb1e8b34922d40528ca1b13878409682df425 Mon Sep 17 00:00:00 2001 From: Prajna1999 Date: Fri, 11 Sep 2026 11:50:53 +0530 Subject: [PATCH 4/4] feat: remove batched updates in favour of direct updates --- backend/app/api/routes/cron.py | 1 - backend/app/crud/llm.py | 33 +++------- backend/app/models/llm/response.py | 1 - backend/app/services/llm/retention.py | 36 ++--------- backend/app/tests/api/routes/test_cron.py | 2 - backend/app/tests/crud/test_llm_retention.py | 53 ++++++++-------- .../app/tests/services/llm/test_retention.py | 62 +++++++------------ scripts/python/invoke-cron.py | 1 + 8 files changed, 65 insertions(+), 124 deletions(-) diff --git a/backend/app/api/routes/cron.py b/backend/app/api/routes/cron.py index 2c9c1c3bb..5100ac930 100644 --- a/backend/app/api/routes/cron.py +++ b/backend/app/api/routes/cron.py @@ -182,7 +182,6 @@ def llm_call_retention_cron_job(session: SessionDep) -> LlmCallRedactionResult: logger.info( f"[llm_call_retention_cron_job] Completed: " f"rows_redacted={result.rows_redacted}, " - f"batches_run={result.batches_run}, " f"cutoff={result.cutoff}" ) return result diff --git a/backend/app/crud/llm.py b/backend/app/crud/llm.py index c017874da..c807abc6a 100644 --- a/backend/app/crud/llm.py +++ b/backend/app/crud/llm.py @@ -337,46 +337,33 @@ def get_llm_call_by_id( return session.exec(statement).first() -REDACTED_INPUT_SENTINEL = "[redacted]" +REDACTED_SENTINEL = "[redacted]" -REDACT_LLM_CALL_BATCH_SQL = text( +REDACT_LLM_CALLS_SQL = text( """ - WITH batch AS ( - SELECT id FROM llm_call - WHERE updated_at <= :cutoff - AND (input <> :redacted_sentinel OR content -> 'content' ->> 'value' IS NOT NULL) - ORDER BY updated_at - LIMIT :batch_size - FOR UPDATE SKIP LOCKED - ) UPDATE llm_call SET input = :redacted_sentinel, content = CASE WHEN jsonb_typeof(content) = 'object' - THEN jsonb_set(content, '{content,value}', 'null'::jsonb, false) + THEN jsonb_set(content, '{content,value}', to_jsonb(CAST(:redacted_sentinel AS text)), false) ELSE content END - WHERE id IN (SELECT id FROM batch) + WHERE updated_at <= :cutoff + AND (input <> :redacted_sentinel OR content -> 'content' ->> 'value' IS DISTINCT FROM :redacted_sentinel) """ ) -def redact_llm_call_batch( - *, session: Session, cutoff: datetime, batch_size: int -) -> int: +def redact_llm_calls(*, session: Session, cutoff: datetime) -> int: result = session.connection().execute( - REDACT_LLM_CALL_BATCH_SQL, - { - "cutoff": cutoff, - "batch_size": batch_size, - "redacted_sentinel": REDACTED_INPUT_SENTINEL, - }, + REDACT_LLM_CALLS_SQL, + {"cutoff": cutoff, "redacted_sentinel": REDACTED_SENTINEL}, ) session.commit() logger.info( - f"[redact_llm_call_batch] Redacted batch | rows: {result.rowcount} | " - f"cutoff: {cutoff.isoformat()} | batch_size: {batch_size}" + f"[redact_llm_calls] Redacted rows | rows: {result.rowcount} | " + f"cutoff: {cutoff.isoformat()}" ) return result.rowcount diff --git a/backend/app/models/llm/response.py b/backend/app/models/llm/response.py index 70bd76c83..afc6905ed 100644 --- a/backend/app/models/llm/response.py +++ b/backend/app/models/llm/response.py @@ -23,7 +23,6 @@ class LlmCallRedactionResult(SQLModel): """Outcome of a single llm_call retention redaction run.""" rows_redacted: int - batches_run: int cutoff: datetime diff --git a/backend/app/services/llm/retention.py b/backend/app/services/llm/retention.py index f3bd8046a..c7592cb71 100644 --- a/backend/app/services/llm/retention.py +++ b/backend/app/services/llm/retention.py @@ -4,48 +4,24 @@ from app.core.config import settings from app.core.util import now -from app.crud.llm import redact_llm_call_batch +from app.crud.llm import redact_llm_calls from app.models.llm.response import LlmCallRedactionResult logger = logging.getLogger(__name__) -LLM_CALL_REDACTION_BATCH_SIZE = 2000 - def redact_aged_llm_calls(*, session: Session) -> LlmCallRedactionResult: cutoff = now() - settings.DELETE_ROLLING_WINDOW_TIMEDELTA logger.info( - f"[redact_aged_llm_calls] Starting redaction | cutoff: {cutoff.isoformat()} | " - f"batch_size: {LLM_CALL_REDACTION_BATCH_SIZE}" + f"[redact_aged_llm_calls] Starting redaction | cutoff: {cutoff.isoformat()}" ) - total_redacted = 0 - batches_run = 0 - - while True: - redacted = redact_llm_call_batch( - session=session, - cutoff=cutoff, - batch_size=LLM_CALL_REDACTION_BATCH_SIZE, - ) - if redacted == 0: - break - - total_redacted += redacted - batches_run += 1 - logger.info( - f"[redact_aged_llm_calls] Batch complete | batch: {batches_run} | " - f"rows: {redacted} | total_rows: {total_redacted}" - ) + rows_redacted = redact_llm_calls(session=session, cutoff=cutoff) logger.info( - f"[redact_aged_llm_calls] Completed | rows_redacted: {total_redacted} | " - f"batches_run: {batches_run} | cutoff: {cutoff.isoformat()}" + f"[redact_aged_llm_calls] Completed | rows_redacted: {rows_redacted} | " + f"cutoff: {cutoff.isoformat()}" ) - return LlmCallRedactionResult( - rows_redacted=total_redacted, - batches_run=batches_run, - cutoff=cutoff, - ) + return LlmCallRedactionResult(rows_redacted=rows_redacted, cutoff=cutoff) diff --git a/backend/app/tests/api/routes/test_cron.py b/backend/app/tests/api/routes/test_cron.py index 32fd4f471..9868eec77 100644 --- a/backend/app/tests/api/routes/test_cron.py +++ b/backend/app/tests/api/routes/test_cron.py @@ -335,7 +335,6 @@ def test_llm_call_retention_cron_job_success( """Returns the redaction summary produced by the retention service.""" result = LlmCallRedactionResult( rows_redacted=4137, - batches_run=3, cutoff=datetime(2026, 2, 22, 12, 0, 0), ) with patch( @@ -350,7 +349,6 @@ def test_llm_call_retention_cron_job_success( assert response.status_code == 200 assert response.json() == { "rows_redacted": 4137, - "batches_run": 3, "cutoff": "2026-02-22T12:00:00", } redact.assert_called_once() diff --git a/backend/app/tests/crud/test_llm_retention.py b/backend/app/tests/crud/test_llm_retention.py index 787d9b895..f8888c9ef 100644 --- a/backend/app/tests/crud/test_llm_retention.py +++ b/backend/app/tests/crud/test_llm_retention.py @@ -5,7 +5,7 @@ from sqlmodel import Session from app.models import Job -from app.crud.llm import REDACTED_INPUT_SENTINEL, redact_llm_call_batch +from app.crud.llm import REDACTED_SENTINEL, redact_llm_calls from app.models.llm import LlmCall from app.tests.utils.llm import create_aged_llm_call, create_llm_job @@ -18,11 +18,18 @@ @pytest.fixture def job(db: Session) -> Job: + # The statement targets llm_call_2, a copy table that only exists in the + # deployed databases; an auto-updatable view over llm_call gives the tests + # that name while keeping the real column types and the ORM factories. + db.exec(text("DROP VIEW IF EXISTS llm_call_2")) # Other llm_call rows (seed data, sibling fixtures) would also match the # cutoff and skew the rowcount assertions, so start from an empty table. db.exec(text("DELETE FROM llm_call")) + db.exec(text("CREATE VIEW llm_call_2 AS SELECT * FROM llm_call")) + db.commit() + yield create_llm_job(db) + db.exec(text("DROP VIEW IF EXISTS llm_call_2")) db.commit() - return create_llm_job(db) def _reload(db: Session, llm_call: LlmCall) -> LlmCall: @@ -39,12 +46,12 @@ def test_redacts_input_and_content_value_only(db: Session, job: Job) -> None: content=TEXT_CONTENT, ) - redacted = redact_llm_call_batch(session=db, cutoff=CUTOFF, batch_size=10) + redacted = redact_llm_calls(session=db, cutoff=CUTOFF) assert redacted == 1 row = _reload(db, llm_call) - assert row.input == "[redacted]" - assert row.content["content"]["value"] is None + assert row.input == REDACTED_SENTINEL + assert row.content["content"]["value"] == REDACTED_SENTINEL assert row.content["content"]["format"] == "text" assert row.content["type"] == "text" @@ -58,7 +65,7 @@ def test_leaves_rows_newer_than_cutoff_untouched(db: Session, job: Job) -> None: content=TEXT_CONTENT, ) - redacted = redact_llm_call_batch(session=db, cutoff=CUTOFF, batch_size=10) + redacted = redact_llm_calls(session=db, cutoff=CUTOFF) assert redacted == 0 row = _reload(db, llm_call) @@ -74,15 +81,15 @@ def test_row_exactly_at_cutoff_is_redacted(db: Session, job: Job) -> None: content=TEXT_CONTENT, ) - assert redact_llm_call_batch(session=db, cutoff=CUTOFF, batch_size=10) == 1 - assert _reload(db, llm_call).input == "[redacted]" + assert redact_llm_calls(session=db, cutoff=CUTOFF) == 1 + assert _reload(db, llm_call).input == REDACTED_SENTINEL def test_second_run_skips_already_redacted_rows(db: Session, job: Job) -> None: create_aged_llm_call(db, updated_at=AGED, job_id=job.id, content=TEXT_CONTENT) - assert redact_llm_call_batch(session=db, cutoff=CUTOFF, batch_size=10) == 1 - assert redact_llm_call_batch(session=db, cutoff=CUTOFF, batch_size=10) == 0 + assert redact_llm_calls(session=db, cutoff=CUTOFF) == 1 + assert redact_llm_calls(session=db, cutoff=CUTOFF) == 0 def test_row_with_null_content_is_redacted_without_error(db: Session, job: Job) -> None: @@ -104,24 +111,14 @@ def test_row_with_null_content_is_redacted_without_error(db: Session, job: Job) == "null" ) - assert redact_llm_call_batch(session=db, cutoff=CUTOFF, batch_size=10) == 1 + assert redact_llm_calls(session=db, cutoff=CUTOFF) == 1 row = _reload(db, llm_call) - assert row.input == "[redacted]" + assert row.input == REDACTED_SENTINEL assert row.content is None -def test_batch_size_caps_rows_touched_per_call(db: Session, job: Job) -> None: - for _ in range(5): - create_aged_llm_call(db, updated_at=AGED, job_id=job.id, content=TEXT_CONTENT) - - assert redact_llm_call_batch(session=db, cutoff=CUTOFF, batch_size=2) == 2 - assert redact_llm_call_batch(session=db, cutoff=CUTOFF, batch_size=2) == 2 - assert redact_llm_call_batch(session=db, cutoff=CUTOFF, batch_size=2) == 1 - assert redact_llm_call_batch(session=db, cutoff=CUTOFF, batch_size=2) == 0 - - def test_returns_zero_when_no_rows_match(db: Session, job: Job) -> None: - assert redact_llm_call_batch(session=db, cutoff=CUTOFF, batch_size=10) == 0 + assert redact_llm_calls(session=db, cutoff=CUTOFF) == 0 def test_redacts_only_aged_rows_in_a_mixed_table(db: Session, job: Job) -> None: @@ -132,8 +129,8 @@ def test_redacts_only_aged_rows_in_a_mixed_table(db: Session, job: Job) -> None: db, updated_at=RECENT, job_id=job.id, input="recent", content=TEXT_CONTENT ) - assert redact_llm_call_batch(session=db, cutoff=CUTOFF, batch_size=10) == 1 - assert _reload(db, aged).input == REDACTED_INPUT_SENTINEL + assert redact_llm_calls(session=db, cutoff=CUTOFF) == 1 + assert _reload(db, aged).input == REDACTED_SENTINEL assert _reload(db, recent).input == "recent" @@ -144,9 +141,9 @@ def test_redacts_content_value_of_row_whose_input_is_already_sentinel( db, updated_at=AGED, job_id=job.id, - input=REDACTED_INPUT_SENTINEL, + input=REDACTED_SENTINEL, content={"type": "audio", "content": {"format": "uri", "value": "s3://a.wav"}}, ) - assert redact_llm_call_batch(session=db, cutoff=CUTOFF, batch_size=10) == 1 - assert _reload(db, llm_call).content["content"]["value"] is None + assert redact_llm_calls(session=db, cutoff=CUTOFF) == 1 + assert _reload(db, llm_call).content["content"]["value"] == REDACTED_SENTINEL diff --git a/backend/app/tests/services/llm/test_retention.py b/backend/app/tests/services/llm/test_retention.py index 7a9963218..8d5881174 100644 --- a/backend/app/tests/services/llm/test_retention.py +++ b/backend/app/tests/services/llm/test_retention.py @@ -5,69 +5,56 @@ from app.core.config import settings from app.services.llm import retention -from app.services.llm.retention import ( - LLM_CALL_REDACTION_BATCH_SIZE, - redact_aged_llm_calls, -) +from app.services.llm.retention import redact_aged_llm_calls -@pytest.fixture -def batch_calls(monkeypatch: pytest.MonkeyPatch) -> list[dict]: - """Records the kwargs of each redact_llm_call_batch call; returns 0 by default.""" - calls: list[dict] = [] - - def fake_batch(**kwargs) -> int: +def _stub_redact( + monkeypatch: pytest.MonkeyPatch, calls: list[dict], returns: int = 0 +) -> None: + def fake_redact(**kwargs) -> int: calls.append(kwargs) - return 0 - - monkeypatch.setattr(retention, "redact_llm_call_batch", fake_batch) - return calls + return returns + monkeypatch.setattr(retention, "redact_llm_calls", fake_redact) -def _stub_returns( - monkeypatch: pytest.MonkeyPatch, returns: list[int], calls: list[dict] -) -> None: - remaining = list(returns) - def fake_batch(**kwargs) -> int: - calls.append(kwargs) - return remaining.pop(0) - - monkeypatch.setattr(retention, "redact_llm_call_batch", fake_batch) +@pytest.fixture +def redact_calls(monkeypatch: pytest.MonkeyPatch) -> list[dict]: + calls: list[dict] = [] + _stub_redact(monkeypatch, calls) + return calls -def test_sums_rows_across_batches_until_a_batch_is_empty( +def test_reports_rows_redacted_from_a_single_crud_call( db: Session, monkeypatch: pytest.MonkeyPatch ) -> None: calls: list[dict] = [] - _stub_returns(monkeypatch, [2000, 2000, 137, 0], calls) + _stub_redact(monkeypatch, calls, returns=4137) result = redact_aged_llm_calls(session=db) assert result.rows_redacted == 4137 - assert result.batches_run == 3 - assert len(calls) == 4 + assert len(calls) == 1 -def test_no_matching_rows(db: Session, batch_calls: list[dict]) -> None: +def test_no_matching_rows(db: Session, redact_calls: list[dict]) -> None: result = redact_aged_llm_calls(session=db) assert result.rows_redacted == 0 - assert result.batches_run == 0 - assert len(batch_calls) == 1 + assert len(redact_calls) == 1 -def test_cutoff_is_naive_utc(db: Session, batch_calls: list[dict]) -> None: +def test_cutoff_is_naive_utc(db: Session, redact_calls: list[dict]) -> None: # llm_call.updated_at is a naive TIMESTAMP; a tz-aware cutoff would raise on compare. result = redact_aged_llm_calls(session=db) - cutoff = batch_calls[0]["cutoff"] + cutoff = redact_calls[0]["cutoff"] assert cutoff.tzinfo is None assert result.cutoff == cutoff def test_cutoff_trails_now_by_the_configured_rolling_window( - db: Session, batch_calls: list[dict], monkeypatch: pytest.MonkeyPatch + db: Session, redact_calls: list[dict], monkeypatch: pytest.MonkeyPatch ) -> None: frozen = datetime(2026, 3, 1, 12, 0, 0) monkeypatch.setattr(retention, "now", lambda: frozen) @@ -75,16 +62,13 @@ def test_cutoff_trails_now_by_the_configured_rolling_window( redact_aged_llm_calls(session=db) - assert batch_calls[0]["cutoff"] == datetime(2026, 2, 22, 12, 0, 0) + assert redact_calls[0]["cutoff"] == datetime(2026, 2, 22, 12, 0, 0) -def test_passes_session_and_module_batch_size_to_crud( - db: Session, batch_calls: list[dict] -) -> None: +def test_passes_session_to_crud(db: Session, redact_calls: list[dict]) -> None: redact_aged_llm_calls(session=db) - assert batch_calls[0]["session"] is db - assert batch_calls[0]["batch_size"] == LLM_CALL_REDACTION_BATCH_SIZE == 2000 + assert redact_calls[0]["session"] is db def test_default_rolling_window_is_one_week() -> None: diff --git a/scripts/python/invoke-cron.py b/scripts/python/invoke-cron.py index e589c65aa..2156b00ad 100644 --- a/scripts/python/invoke-cron.py +++ b/scripts/python/invoke-cron.py @@ -19,6 +19,7 @@ "/api/v1/cron/evaluations", "/api/v1/cron/pending-jobs", "/api/v1/cron/daily-stats", + "/api/v1/cron/llm-call-retention", ] REQUEST_TIMEOUT = 30 # Timeout for requests in seconds