"""B1-P small provider-limit diagnostic. Not a 120-prompt run. Does not change retry.""" from __future__ import annotations import argparse import asyncio import hashlib import inspect import json import sys import time from collections import Counter from datetime import UTC, datetime from pathlib import Path from types import SimpleNamespace from uuid import UUID ROOT = Path(__file__).resolve().parents[1] sys.path.insert(0, str(ROOT)) from sqlalchemy import text # noqa: E402 from src.config import get_settings # noqa: E402 from src.core.provider_capacity import watchdog_seconds # noqa: E402 from src.domain.projects.discovery_b1 import capacity_trace_fields # noqa: E402 from src.domain.projects.discovery_catalog import build_discovery_payload # noqa: E402 from src.integrations.project_intake_client import ( # noqa: E402 DISCOVERY_PROMPT, discover_project, provider_error_snapshot, ) from src.workers.tasks.project_discovery import reserve_discovery_call # noqa: E402 EVALUATION_PROJECT_ID = "1f6556ff-4a16-4143-9b12-e9d9249d1d3d" EVALUATION_WORKSPACE_ID = "5e890ff3-5c7b-43f5-a4ad-75c47f90bdba" FIXED_PROMPT = ( "Our sales team needs a private customer directory with names and phone " "numbers. Staff can add and edit customers; everyone works in one company " "workspace." ) PROTECTED_OUTPUTS = { ROOT / "docs/audits/fixtures/13a_run_01", ROOT / "docs/audits/fixtures/13a_run_02", ROOT / "docs/audits/fixtures/13b_b1_admission_gated_01", ROOT / "docs/audits/fixtures/13b_b1_admission_gated_prod_host_01", ROOT / "docs/audits/fixtures/13b_b1_admission_gated_prod_host_02", ROOT / "docs/audits/fixtures/13b_b1_extraction_prod_host_01", ROOT / "docs/audits/fixtures/13b_b1_provider_limit_prod_host_01", ROOT / "docs/audits/fixtures/13b_b1_provider_limit_prod_host_02", ROOT / "docs/audits/fixtures/13b_b1_provider_limit_prod_host_03", } def write_json(path: Path, value) -> None: path.parent.mkdir(parents=True, exist_ok=True) path.write_text(json.dumps(value, ensure_ascii=False, indent=2) + "\n", encoding="utf-8") def _refuse_protected(output: Path) -> None: resolved = output.resolve() for path in PROTECTED_OUTPUTS: target = path.resolve() if resolved == target or target in resolved.parents: raise ValueError("Refusing to write into a protected historical fixture") def preflight(): settings = get_settings() if "EXTRACTION DISCIPLINE" not in DISCOVERY_PROMPT: raise ValueError("Production measurement requires the authorized extraction contract") if "INTENT FIDELITY" not in DISCOVERY_PROMPT: raise ValueError("Production measurement requires the intent-fidelity contract") if settings.project_intake_timeout_seconds != 8: raise ValueError("Diagnostic requires the frozen 8s timeout") if settings.project_intake_model != "gpt-4o-mini": raise ValueError("Model differs from the frozen baseline") if not settings.openai_api_key: raise ValueError("Diagnostic requires the configured provider credential") if not settings.redis_url: raise ValueError("Diagnostic requires configured Redis") if settings.project_intake_calls_per_minute != 60: raise ValueError("Expected unchanged 60/min admission policy") source = inspect.getsource(reserve_discovery_call) if "project-intake:budget:" not in source: raise ValueError("Admission implementation is not the production reserve_discovery_call") return { "preflight": "passed", "measurement_kind": "provider_capacity_localization", "max_calls_default": 64, "stop_after_rate_limit_snapshots": 3, "concurrency_default": 4, "start_interval_seconds": 1.1, "timeout_seconds": 8, "retry_policy_changed": False, "admission_policy_changed": False, "planner_scored": False, "provider_calls": 0, "required_environment_names": ["REDIS_URL", "DATABASE_URL", "OPENAI_API_KEY"], "script_sha256": hashlib.sha256(Path(__file__).read_bytes()).hexdigest(), } def verify(output: Path): scorecard = json.loads((output / "scorecard.json").read_text(encoding="utf-8")) metadata = json.loads((output / "metadata.json").read_text(encoding="utf-8")) assert scorecard["measurement_kind"] == "provider_capacity_localization" assert metadata["environment"] == "production" assert metadata["retry_policy_changed"] is False assert metadata["admission_policy_changed"] is False assert metadata["planner_scored"] is False return { "verification": "passed", "concurrency": metadata["concurrency"], "rate_limit_snapshots": scorecard["rate_limit_snapshots"], "calls": scorecard["calls"], "started": scorecard.get("started_calls"), "skipped_after_stop": scorecard.get("skipped_after_stop"), } async def live_context(): from src.core.database import async_session_factory, engine from src.domain.projects.discovery_service import available_templates engine.echo = False try: async with async_session_factory() as db: await db.execute(text("SET TRANSACTION READ ONLY")) project = ( ( await db.execute( text( "SELECT type, source_kind FROM projects " "WHERE id=:pid AND workspace_id=:wid" ), { "pid": UUID(EVALUATION_PROJECT_ID), "wid": UUID(EVALUATION_WORKSPACE_ID), }, ) ) .mappings() .first() ) if not project: raise ValueError("Project does not belong to the evaluation workspace") templates = [ SimpleNamespace( **{ key: getattr(item, key) for key in ("ref", "name", "description", "framework", "tags") } ) for item in await available_templates(db, UUID(EVALUATION_WORKSPACE_ID)) ] await db.rollback() return dict(project), templates finally: await engine.dispose() def _persistable_row(row: dict) -> dict: return { key: value for key, value in row.items() if key not in {"trace_request", "raw_turn", "result"} } async def measure( output: Path, environment: str, max_calls: int, stop_after: int, concurrency: int, start_interval: float, ): _refuse_protected(output) if output.exists(): raise ValueError("Select a new experiment directory") if concurrency < 1: raise ValueError("concurrency must be at least 1") if start_interval < 1.1: raise ValueError("Must not exceed the baseline request start rate") if max_calls > 64: raise ValueError("This diagnostic ceiling is 64 calls; do not expand toward 120") settings = get_settings() preflight() project, templates = await live_context() payload = build_discovery_payload( project["type"], project["source_kind"], {"messages": [{"role": "user", "content": FIXED_PROMPT}]}, templates, ) metadata = { "phase": "13B", "gate": "B1-P", "measurement_kind": "provider_capacity_localization", "environment": environment, "max_calls": max_calls, "stop_after_rate_limit_snapshots": stop_after, "concurrency": concurrency, "start_interval_seconds": start_interval, "timeout_seconds": 8.0, "retry_policy_changed": False, "admission_policy_changed": False, "timeout_policy_changed": False, "planner_scored": False, "model": settings.project_intake_model, "started_at": datetime.now(UTC).isoformat(), "status": "running", "do_not": [ "overwrite historical provider-limit or 120-prompt fixtures", "continue after 3 rate-limit snapshots", "expand into a 120-prompt run", "change timeout, admission, or retry", ], } write_json(output / "metadata.json", metadata) semaphore, lock = asyncio.Semaphore(concurrency), asyncio.Lock() last_start = 0.0 in_flight = 0 stop = False snapshots: list[dict] = [] calls: list[dict] = [] skipped = 0 async def one(index: int): nonlocal last_start, in_flight, stop, skipped async with semaphore: async with lock: if stop or len(snapshots) >= stop_after: skipped += 1 return await asyncio.sleep(max(0, last_start + start_interval - time.monotonic())) if stop or len(snapshots) >= stop_after: skipped += 1 return last_start = time.monotonic() in_flight += 1 inflight_now = in_flight started_at = datetime.now(UTC).isoformat() trace: dict = {} admission_started = time.monotonic() try: await reserve_discovery_call(settings) admission = { "attempted": True, "granted": True, "check_seconds": round(time.monotonic() - admission_started, 4), } except ValueError as exc: admission = { "attempted": True, "granted": False, "denial_kind": "policy", "reason": str(exc), "check_seconds": round(time.monotonic() - admission_started, 4), } except Exception as exc: admission = { "attempted": True, "granted": False, "denial_kind": "infrastructure", "reason": type(exc).__name__, "check_seconds": round(time.monotonic() - admission_started, 4), } row = { "index": index, "started_at": started_at, "in_flight": inflight_now, "attempt_number": 1, "retry_occurred": False, "retry_delay_seconds": None, "retry_outcome": None, "model": settings.project_intake_model, "admission": admission, } if not admission["granted"]: row["status"] = "admission_denied" else: started = time.monotonic() try: await asyncio.wait_for( discover_project(payload, settings, trace=trace), timeout=watchdog_seconds(settings, 8), ) row.update( status="success", elapsed_seconds=round(time.monotonic() - started, 3), usage=trace.get("usage"), **capacity_trace_fields(trace), ) except Exception as exc: snapshot = trace.get("provider_error") or provider_error_snapshot(exc) snapshot["attempt_number"] = 1 snapshot["retry_occurred"] = False snapshot["retry_outcome"] = None snapshot["model"] = settings.project_intake_model row.update( status="error", error_type=type(exc).__name__, elapsed_seconds=round(time.monotonic() - started, 3), provider_error=snapshot, **capacity_trace_fields(trace), ) if type(exc).__name__ == "RateLimitError": async with lock: snapshots.append(snapshot) if len(snapshots) >= stop_after: stop = True row["finished_at"] = datetime.now(UTC).isoformat() persistable = _persistable_row(row) write_json(output / "calls" / f"{index:02}.json", persistable) async with lock: calls.append(persistable) in_flight = max(0, in_flight - 1) print( json.dumps( { "index": index, "status": row["status"], "in_flight": inflight_now, "snapshots": len(snapshots), } ), flush=True, ) await asyncio.gather(*(one(index) for index in range(1, max_calls + 1))) started_calls = [row for row in calls] scorecard = { "measurement_kind": "provider_capacity_localization", "concurrency": concurrency, "start_interval_seconds": start_interval, "max_calls": max_calls, "started_calls": len(started_calls), "skipped_after_stop": skipped, "calls": len(started_calls), "successes": sum(1 for row in started_calls if row["status"] == "success"), "admission_denied": sum( 1 for row in started_calls if row["status"] == "admission_denied" ), "rate_limit_snapshots": len(snapshots), "error_types": dict( Counter(row.get("error_type") for row in started_calls if row.get("error_type")) ), "usage_on_success": [ {"index": row["index"], "usage": row.get("usage"), "in_flight": row.get("in_flight")} for row in started_calls if row["status"] == "success" ], "snapshots": snapshots, "retry_policy_changed": False, "planner_scored": False, } metadata.update(status="complete", finished_at=datetime.now(UTC).isoformat()) write_json(output / "metadata.json", metadata) write_json(output / "scorecard.json", scorecard) print( json.dumps( { "rate_limit_snapshots": len(snapshots), "calls": len(started_calls), "skipped_after_stop": skipped, "concurrency": concurrency, }, indent=2, ) ) if __name__ == "__main__": parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--out", type=Path) parser.add_argument("--environment", choices=("local", "production"), default="local") parser.add_argument("--max-calls", type=int, default=64) parser.add_argument("--stop-after", type=int, default=3) parser.add_argument("--concurrency", type=int, default=4) parser.add_argument("--start-interval", type=float, default=1.1) parser.add_argument("--preflight", action="store_true") parser.add_argument("--verify", type=Path) args = parser.parse_args() try: if args.preflight: print(json.dumps(preflight(), indent=2)) elif args.verify: print(json.dumps(verify(args.verify), indent=2)) elif args.out: asyncio.run( measure( args.out, args.environment, args.max_calls, args.stop_after, args.concurrency, args.start_interval, ) ) else: parser.error("Supply --preflight, --verify RECEIPT_DIR, or --out NEW_DIRECTORY") except Exception as ext: print( json.dumps({"diagnostic_failed": type(ext).__name__, "detail": str(ext)}), file=sys.stderr, ) raise SystemExit(1) from None