70 lines
2.8 KiB
Python
70 lines
2.8 KiB
Python
"""Cost plan for the residual AI work items, with budget guard (section 18).
|
|
|
|
Runs BEFORE any real execution: UNIQUE_AI_WORK_ITEMS, token estimates per task type,
|
|
estimated cost, and the authorize/refuse verdict against AI_MAX_RUN_COST_USD.
|
|
"""
|
|
import argparse
|
|
import json
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
|
|
|
from sqlalchemy import select
|
|
from central.app.ai.provider import BudgetGuard
|
|
from central.app.config import settings
|
|
from central.app.db import SessionLocal
|
|
from central.app.models import AiWorkItem
|
|
|
|
TOKENS_PER_ITEM = { # conservative prompt+schema overhead per single work item
|
|
"INVALID_NCM": (700, 120),
|
|
"CEST": (500, 80),
|
|
"EX_TIPI": (500, 80),
|
|
"PRODUCT_ROLE": (400, 60),
|
|
}
|
|
|
|
|
|
def main(output="docs/AI_COST_PLAN_V1.json") -> dict:
|
|
session = SessionLocal()
|
|
items = session.scalars(select(AiWorkItem)).all()
|
|
guard = BudgetGuard(settings.ai_max_run_cost_usd, settings.ai_price_in_per_m, settings.ai_price_out_per_m)
|
|
per_task = {}
|
|
in_total = out_total = 0
|
|
resolved = {"KB_HIT", "KB_EXACT_HIT", "KB_FINGERPRINT_HIT", "HIGH_CONFIDENCE_RETRIEVAL", "SINGLE_OFFICIAL_CANDIDATE", "DONE", "SKIPPED", "CONFLICT"}
|
|
for item in items:
|
|
if item.status in resolved:
|
|
continue # cache/official candidate resolved — no AI call needed
|
|
tokens_in, tokens_out = TOKENS_PER_ITEM.get(item.task_type, (500, 80))
|
|
bucket = per_task.setdefault(item.task_type, {"items": 0, "input_tokens": 0, "output_tokens": 0, "cells_covered": 0})
|
|
bucket["items"] += 1; bucket["input_tokens"] += tokens_in; bucket["output_tokens"] += tokens_out
|
|
bucket["cells_covered"] += item.cells_covered
|
|
in_total += tokens_in; out_total += tokens_out
|
|
plan = guard.estimate(in_total, out_total)
|
|
report = {
|
|
"artifact": "AI_COST_PLAN_V1",
|
|
"unique_ai_work_items": sum(bucket["items"] for bucket in per_task.values()),
|
|
"per_task_type": per_task,
|
|
"estimated_input_tokens": plan["input_tokens"],
|
|
"estimated_output_tokens": plan["output_tokens"],
|
|
"estimated_cost_usd": plan["estimated_cost"],
|
|
"ai_max_run_cost_usd": guard.max_run_cost_usd,
|
|
"budget_authorized": guard.authorize(plan),
|
|
"provider_configured": settings.ai_provider,
|
|
"model_configured": settings.ai_model,
|
|
"ai_enabled": settings.ai_enabled,
|
|
"max_batch_size": settings.ai_max_batch_size,
|
|
"min_confidence": settings.ai_min_confidence,
|
|
}
|
|
Path(output).write_text(json.dumps(report, indent=2, ensure_ascii=False), encoding="utf-8")
|
|
session.close()
|
|
return report
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import argparse
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--output", default="/tmp/ai_cost_plan.json")
|
|
args = parser.parse_args()
|
|
print(json.dumps(main(args.output), indent=2))
|
|
|