229 lines
7.6 KiB
Python
229 lines
7.6 KiB
Python
#!/usr/bin/env python3
|
|
"""Run one session digest compression job.
|
|
|
|
Dry-run is the default: it loads one ended session, releases the DB connection,
|
|
calls the configured engine, applies the local digest quality gate, and prints a
|
|
metadata-only result. Use --apply to replace the fallback digest after acceptance.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import asyncio
|
|
import json
|
|
import os
|
|
import sys
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
|
|
REPO_ROOT = Path(__file__).resolve().parents[1]
|
|
API_ROOT = REPO_ROOT / "apps" / "api"
|
|
|
|
if hasattr(sys.stdout, "reconfigure"):
|
|
sys.stdout.reconfigure(encoding="utf-8")
|
|
if hasattr(sys.stderr, "reconfigure"):
|
|
sys.stderr.reconfigure(encoding="utf-8")
|
|
|
|
|
|
def _prepare_imports() -> None:
|
|
os.chdir(API_ROOT)
|
|
api_root = str(API_ROOT)
|
|
if api_root not in sys.path:
|
|
sys.path.insert(0, api_root)
|
|
|
|
|
|
def build_parser() -> argparse.ArgumentParser:
|
|
parser = argparse.ArgumentParser(
|
|
description=(
|
|
"Run the M2 session digest worker for one session. Dry-run is the "
|
|
"default; --apply writes accepted LLM digest output back to DB."
|
|
)
|
|
)
|
|
parser.add_argument("--session-id", required=True, help="app.sessions.id to compress")
|
|
parser.add_argument("--model", help="optional engine model override")
|
|
parser.add_argument(
|
|
"--forbidden-substring",
|
|
action="append",
|
|
default=[],
|
|
help="raw substring that must reject the candidate digest; repeatable",
|
|
)
|
|
parser.add_argument("--apply", action="store_true", help="write accepted digest output to DB")
|
|
parser.add_argument("--json", action="store_true", help="emit machine-readable JSON")
|
|
parser.add_argument(
|
|
"--show-digest",
|
|
action="store_true",
|
|
help="include the accepted digest text in output; default is metadata only",
|
|
)
|
|
parser.add_argument(
|
|
"--allow-rejected",
|
|
action="store_true",
|
|
help="exit 0 when the engine returned a rejected candidate",
|
|
)
|
|
return parser
|
|
|
|
|
|
async def _run(args: argparse.Namespace) -> dict[str, Any]:
|
|
_prepare_imports()
|
|
from app.db import acquire, close_pool, init_pool
|
|
from app.engine_client import engine_client
|
|
from app.routes.admin import apply_engine_config_from_store
|
|
from app.session_persistence import record_llm_call_audit
|
|
from app.services import session_digest_worker
|
|
|
|
started = time.perf_counter()
|
|
await init_pool()
|
|
try:
|
|
async with acquire(role="admin") as conn:
|
|
loaded = await session_digest_worker.load_session_digest_job(conn, args.session_id)
|
|
if loaded is None:
|
|
return {
|
|
"ok": False,
|
|
"found": False,
|
|
"applied": False,
|
|
"session_id": args.session_id,
|
|
"error": "session digest job not found",
|
|
}
|
|
|
|
await apply_engine_config_from_store()
|
|
await engine_client.startup()
|
|
worker = await session_digest_worker.run_session_digest_worker(
|
|
loaded.job,
|
|
engine_client,
|
|
existing_case_digest=loaded.existing_case_digest,
|
|
forbidden_substrings=tuple(args.forbidden_substring or ()),
|
|
model=args.model,
|
|
audit_hook=record_llm_call_audit,
|
|
)
|
|
|
|
applied = False
|
|
if args.apply and worker.apply_plan is not None:
|
|
async with acquire(role="admin") as conn:
|
|
applied = await session_digest_worker.apply_session_digest_plan(
|
|
conn,
|
|
worker.apply_plan,
|
|
learner_id=loaded.learner_id,
|
|
)
|
|
return _result_payload(
|
|
loaded=loaded,
|
|
worker=worker,
|
|
applied=applied,
|
|
apply_requested=bool(args.apply),
|
|
show_digest=bool(args.show_digest),
|
|
elapsed_ms=int((time.perf_counter() - started) * 1000),
|
|
)
|
|
finally:
|
|
await engine_client.shutdown()
|
|
await close_pool()
|
|
|
|
|
|
def _result_payload(
|
|
*,
|
|
loaded: Any,
|
|
worker: Any,
|
|
applied: bool,
|
|
apply_requested: bool,
|
|
show_digest: bool,
|
|
elapsed_ms: int,
|
|
) -> dict[str, Any]:
|
|
quality = worker.outcome.quality
|
|
result = worker.outcome.result
|
|
apply_plan = worker.apply_plan
|
|
payload: dict[str, Any] = {
|
|
"ok": bool(result is not None),
|
|
"found": True,
|
|
"accepted": bool(result is not None),
|
|
"applied": bool(applied),
|
|
"apply_requested": apply_requested,
|
|
"fallback_required": bool(worker.fallback_required),
|
|
"session_id": loaded.job.session_id,
|
|
"case_id": loaded.job.case_id,
|
|
"session_no": loaded.job.session_no,
|
|
"quality": {
|
|
"accepted": bool(quality.accepted),
|
|
"reason": quality.reason,
|
|
"retryable": bool(quality.retryable),
|
|
"details": list(quality.details),
|
|
},
|
|
"engine": {
|
|
"provider": worker.response.provider,
|
|
"model": worker.response.model,
|
|
"tokens_in": worker.response.tokens_in,
|
|
"tokens_out": worker.response.tokens_out,
|
|
"cost_usd": worker.response.cost_usd,
|
|
"inference_geo": worker.response.inference_geo,
|
|
},
|
|
"elapsed_ms": elapsed_ms,
|
|
}
|
|
if apply_plan is not None:
|
|
payload["apply_plan"] = {
|
|
"compressed_by": apply_plan.compressed_by,
|
|
"token_count": apply_plan.token_count,
|
|
"open_thread_count": len(apply_plan.result.open_threads),
|
|
"digest_chars": len(apply_plan.result.digest),
|
|
"case_digest_chars": len(apply_plan.case_digest or ""),
|
|
}
|
|
if show_digest:
|
|
payload["digest"] = apply_plan.result.digest
|
|
return payload
|
|
|
|
|
|
def _print_text(payload: dict[str, Any]) -> None:
|
|
if not payload.get("found"):
|
|
print(payload.get("error") or "session digest job not found", file=sys.stderr)
|
|
return
|
|
status = "accepted" if payload.get("accepted") else "rejected"
|
|
mode = "applied" if payload.get("applied") else "dry-run"
|
|
print(
|
|
"session digest worker: {status} ({mode}) session={session} reason={reason}".format(
|
|
status=status,
|
|
mode=mode,
|
|
session=payload.get("session_id"),
|
|
reason=(payload.get("quality") or {}).get("reason"),
|
|
)
|
|
)
|
|
plan = payload.get("apply_plan") or {}
|
|
if plan:
|
|
print(
|
|
"candidate: compressed_by={compressed_by} token_count={token_count} "
|
|
"digest_chars={digest_chars} case_digest_chars={case_digest_chars}".format(**plan)
|
|
)
|
|
engine = payload.get("engine") or {}
|
|
if engine:
|
|
print(
|
|
"engine: provider={provider} model={model} tokens={tokens_in}/{tokens_out} cost={cost_usd}".format(
|
|
**engine
|
|
)
|
|
)
|
|
if payload.get("digest"):
|
|
print("digest:")
|
|
print(payload["digest"])
|
|
|
|
|
|
def main(argv: list[str] | None = None) -> int:
|
|
parser = build_parser()
|
|
args = parser.parse_args(list(argv if argv is not None else sys.argv[1:]))
|
|
try:
|
|
payload = asyncio.run(_run(args))
|
|
except Exception as exc: # noqa: BLE001 - CLI should report compactly.
|
|
payload = {
|
|
"ok": False,
|
|
"found": False,
|
|
"applied": False,
|
|
"error_type": exc.__class__.__name__,
|
|
"error": str(exc),
|
|
}
|
|
if args.json:
|
|
print(json.dumps(payload, ensure_ascii=False, indent=2, sort_keys=True))
|
|
else:
|
|
_print_text(payload)
|
|
if payload.get("ok") or (args.allow_rejected and payload.get("found")):
|
|
return 0
|
|
if payload.get("found"):
|
|
return 2
|
|
return 3
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|