diff --git a/apps/worker/scripts/debug_text_track.py b/apps/worker/scripts/debug_text_track.py index e14176a1f..d85fc9a99 100644 --- a/apps/worker/scripts/debug_text_track.py +++ b/apps/worker/scripts/debug_text_track.py @@ -120,6 +120,10 @@ def _apply_token_usage_to_outputs( trace: dict[str, Any], usage: dict[str, Any], ) -> None: + from shared.services.storage.zip_manifest_schema import ( + enrich_manifest_with_token_cost_estimate, + ) + trace["token_usage"] = usage _write_json(out_dir / "trace.json", trace) manifest_path = out_dir / "manifest.json" @@ -131,7 +135,10 @@ def _apply_token_usage_to_outputs( processing = manifest.setdefault("processing", {}) if isinstance(processing, dict): processing["token_usage"] = usage - _write_json(manifest_path, manifest) + _write_json( + manifest_path, + enrich_manifest_with_token_cost_estimate(manifest), + ) # ── Stage 1: Profile + Shard Plan (PDF only) ─────────────────────────────── diff --git a/apps/worker/scripts/page_memory/debug_pm_stage5_tagging_finalize.py b/apps/worker/scripts/page_memory/debug_pm_stage5_tagging_finalize.py index 079a3349e..09113e4aa 100644 --- a/apps/worker/scripts/page_memory/debug_pm_stage5_tagging_finalize.py +++ b/apps/worker/scripts/page_memory/debug_pm_stage5_tagging_finalize.py @@ -611,7 +611,10 @@ def _load_selected_scope(scope_id: str) -> ScopeResult: report_path.write_text(report, encoding="utf-8") if args.finalize: - from shared.services.storage.zip_manifest_schema import ZipManifestBuilder + from shared.services.storage.zip_manifest_schema import ( + ZipManifestBuilder, + enrich_manifest_with_token_cost_estimate, + ) manifest = ZipManifestBuilder().generate_manifest( job_id=filename, @@ -625,7 +628,11 @@ def _load_selected_scope(scope_id: str) -> ScopeResult: hierarchy=hierarchy_dict, ) (out_dir / "manifest.json").write_text( - json.dumps(manifest, ensure_ascii=False, indent=2), + json.dumps( + enrich_manifest_with_token_cost_estimate(manifest), + ensure_ascii=False, + indent=2, + ), encoding="utf-8", ) diff --git a/packages/shared-python/shared/services/storage/zip_manifest_schema.py b/packages/shared-python/shared/services/storage/zip_manifest_schema.py index 4190fe4e2..adc76513a 100644 --- a/packages/shared-python/shared/services/storage/zip_manifest_schema.py +++ b/packages/shared-python/shared/services/storage/zip_manifest_schema.py @@ -2,12 +2,51 @@ from __future__ import annotations +from copy import deepcopy from typing import Any from shared.services.ai.token_costing import build_token_cost_estimate from shared.utils.utc_now import utc_now_naive +def extract_manifest_token_usage(manifest: dict[str, Any]) -> dict[str, Any]: + """Return token usage embedded in a manifest, if present.""" + processing = manifest.get("processing") + if not isinstance(processing, dict): + return {} + + stages = processing.get("stages") + if isinstance(stages, dict): + usage = stages.get("token_usage") + if isinstance(usage, dict): + return usage + + usage = processing.get("token_usage") + if isinstance(usage, dict): + return usage + return {} + + +def enrich_manifest_with_token_cost_estimate(manifest: dict[str, Any]) -> dict[str, Any]: + """Add LLM cost_estimate to a local debug manifest copy.""" + enriched = deepcopy(manifest) + processing = enriched.setdefault("processing", {}) + if not isinstance(processing, dict): + return enriched + token_usage = extract_manifest_token_usage(enriched) + processing["cost_estimate"] = build_token_cost_estimate(token_usage) + return enriched + + +def strip_manifest_cost_fields(manifest: dict[str, Any]) -> dict[str, Any]: + """Remove internal LLM cost fields before writing manifest into a ZIP.""" + stripped = deepcopy(manifest) + processing = stripped.get("processing") + if isinstance(processing, dict): + processing.pop("cost_estimate", None) + return stripped + + class ZipManifestBuilder: def generate_manifest( self, @@ -20,7 +59,6 @@ def generate_manifest( hierarchy: dict[str, Any] | None = None, ) -> dict[str, Any]: stages = job_metadata.get("stages", {}) - token_usage = stages.get("token_usage") if isinstance(stages, dict) else {} return { "version": "2.0", "job_id": job_id, @@ -34,7 +72,6 @@ def generate_manifest( "micro_dollars": job_metadata.get("billing_amount_micro_dollars"), "credits": job_metadata.get("billing_credits"), }, - "cost_estimate": build_token_cost_estimate(token_usage), "timing": { "started_at": job_metadata.get("processing_started_at"), "completed_at": job_metadata.get("processing_completed_at"), diff --git a/packages/shared-python/shared/services/storage/zip_result_service.py b/packages/shared-python/shared/services/storage/zip_result_service.py index c1808002f..160223a9e 100644 --- a/packages/shared-python/shared/services/storage/zip_result_service.py +++ b/packages/shared-python/shared/services/storage/zip_result_service.py @@ -24,6 +24,7 @@ ) from shared.services.storage.zip_result_resources import ZipResourceCollector from shared.services.storage.zip_result_schema import ZipResultSchemaBuilder +from shared.services.storage.zip_manifest_schema import strip_manifest_cost_fields class ZipResultService: @@ -83,6 +84,7 @@ def generate_zip_package( job_metadata=job_metadata, hierarchy=hierarchy, ) + manifest = strip_manifest_cost_fields(manifest) parse_track = str((job_metadata or {}).get("parse_track") or "") artifact = self._writer.write( ZipPackageWriteRequest( diff --git a/packages/shared-python/shared/tests/test_zip_manifest_schema.py b/packages/shared-python/shared/tests/test_zip_manifest_schema.py new file mode 100644 index 000000000..ad64bd7d2 --- /dev/null +++ b/packages/shared-python/shared/tests/test_zip_manifest_schema.py @@ -0,0 +1,73 @@ +from __future__ import annotations + +import os + +os.environ.setdefault("DATABASE_URL", "postgresql+asyncpg://test:test@localhost/test") +os.environ.setdefault("TMP_PATH", "/tmp/knowhere-test") +os.environ.setdefault("S3_BUCKET_NAME", "test-uploads") +os.environ.setdefault("S3_TEMP_PATH", "/tmp") + +from shared.services.storage.zip_manifest_schema import ( + ZipManifestBuilder, + enrich_manifest_with_token_cost_estimate, + strip_manifest_cost_fields, +) + + +def test_generate_manifest_omits_llm_cost_estimate() -> None: + manifest = ZipManifestBuilder().generate_manifest( + job_id="job_test", + data_id="data_test", + source_file_name="sample.pdf", + statistics={"total_chunks": 1}, + job_metadata={ + "page_count": 10, + "billing_status": "charged", + "billing_amount_micro_dollars": 150_000, + "billing_credits": 0.15, + "stages": { + "token_usage": { + "prompt_tokens": 100, + "completion_tokens": 20, + "total_tokens": 120, + "calls": 2, + "by_model": {"deepseek-chat": {"total_tokens": 120, "calls": 2}}, + } + }, + }, + hierarchy={"Root": {}}, + ) + + processing = manifest["processing"] + assert "cost_estimate" not in processing + assert processing["cost"]["credits"] == 0.15 + assert processing["stages"]["token_usage"]["by_model"] + + +def test_enrich_and_strip_manifest_cost_fields() -> None: + manifest = ZipManifestBuilder().generate_manifest( + job_id="job_test", + data_id=None, + source_file_name="sample.pdf", + statistics={}, + job_metadata={ + "stages": { + "token_usage": { + "prompt_tokens": 10, + "completion_tokens": 4, + "total_tokens": 14, + "calls": 1, + "by_model": {}, + "by_task": {}, + } + } + }, + ) + + enriched = enrich_manifest_with_token_cost_estimate(manifest) + assert "cost_estimate" in enriched["processing"] + assert enriched["processing"]["cost_estimate"]["currency"] == "USD" + + stripped = strip_manifest_cost_fields(enriched) + assert "cost_estimate" not in stripped["processing"] + assert stripped["processing"]["stages"]["token_usage"]["total_tokens"] == 14