Pinscope open-source core
Agentic schematic validation: datasheet extraction via Claude Console Skills, netlist/BOM design graph, per-IC direct datasheet review with page citations, capacitor derating, Next.js report UI. Extracted from the Pinscope cloud codebase. Auth and billing live in the private gateway repo behind stable seams (billing_hook.py, adapter files listed in CLAUDE.md).
This commit is contained in:
Siddharth Kothari
@@ -0,0 +1,48 @@
|
||||
"""Global admin settings, persisted via StorageBackend.
|
||||
|
||||
Settings are stored at ``admin/settings.json`` in storage (GCS or local
|
||||
``data/``). The module mirrors the pattern in ``limits.py``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from packaging.version import Version
|
||||
|
||||
from backend.services.storage import StorageBackend
|
||||
|
||||
_SETTINGS_KEY = "admin/settings.json"
|
||||
|
||||
_DEFAULTS: dict[str, str] = {
|
||||
"min_model_version": "0.0.0", # no threshold by default
|
||||
}
|
||||
|
||||
|
||||
def get_admin_settings(storage: StorageBackend) -> dict:
|
||||
"""Return the full admin settings dict, with defaults."""
|
||||
if storage.exists(_SETTINGS_KEY):
|
||||
data = storage.read_json(_SETTINGS_KEY)
|
||||
return {**_DEFAULTS, **data}
|
||||
return dict(_DEFAULTS)
|
||||
|
||||
|
||||
def get_min_model_version(storage: StorageBackend) -> str:
|
||||
"""Return the min_model_version threshold."""
|
||||
return get_admin_settings(storage).get("min_model_version", "0.0.0")
|
||||
|
||||
|
||||
def set_min_model_version(storage: StorageBackend, version: str) -> None:
|
||||
"""Set the min_model_version threshold. Validates semver format."""
|
||||
Version(version) # raises InvalidVersion if bad
|
||||
data = get_admin_settings(storage)
|
||||
data["min_model_version"] = version
|
||||
storage.write_json(_SETTINGS_KEY, data)
|
||||
|
||||
|
||||
def version_is_stale(component_version: str, min_version: str) -> bool:
|
||||
"""Return True if *component_version* < *min_version* (semver)."""
|
||||
if min_version == "0.0.0":
|
||||
return False
|
||||
try:
|
||||
return Version(component_version) < Version(min_version)
|
||||
except Exception:
|
||||
return True # unparseable → treat as stale
|
||||
@@ -0,0 +1,106 @@
|
||||
"""Per-project API call logging.
|
||||
|
||||
Captures metadata for every LLM API call made during a pipeline run and
|
||||
serialises to JSONL for storage alongside other project artefacts. Pricing
|
||||
lives in ``backend.services.llm.pricing`` and is provider-aware.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from datetime import datetime, timezone
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
# Re-exported for callers (pipeline.total_cost) — provider-aware now
|
||||
from backend.services.llm.pricing import cost_for_entry, total_cost # noqa: F401
|
||||
|
||||
|
||||
class ApiLogEntry(BaseModel):
|
||||
timestamp: str
|
||||
stage: str # pintable | rules | pattern | validation | ...
|
||||
identifier: str # MPN or component designator
|
||||
model: str
|
||||
provider: str = "anthropic" # anthropic | gemini
|
||||
input_tokens: int
|
||||
output_tokens: int
|
||||
cache_creation_input_tokens: int = 0
|
||||
cache_read_input_tokens: int = 0
|
||||
duration_ms: int
|
||||
stop_reason: str
|
||||
skill_id: str | None = None
|
||||
turns: int | None = None
|
||||
error: str | None = None
|
||||
cost_usd: float | None = None
|
||||
credits_charged: float | None = None
|
||||
# True when the call ran in an admin-initiated free context (e.g. regen)
|
||||
# — the raw USD cost is still recorded for accounting, but no credits
|
||||
# are charged to the user.
|
||||
free: bool = False
|
||||
|
||||
|
||||
@dataclass
|
||||
class CallMeta:
|
||||
"""Metadata returned alongside every Claude API call result."""
|
||||
input_tokens: int
|
||||
output_tokens: int
|
||||
cache_creation_input_tokens: int
|
||||
cache_read_input_tokens: int
|
||||
duration_ms: int
|
||||
stop_reason: str
|
||||
turns: int = 1
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Logger
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@dataclass
|
||||
class ApiLogger:
|
||||
"""Collects API call log entries during a pipeline run.
|
||||
|
||||
``free=True`` marks every entry as admin-initiated and zeros the
|
||||
``credits_charged`` field so downstream charging / reporting treats the
|
||||
run as free to the user. The underlying USD cost is still recorded.
|
||||
"""
|
||||
entries: list[dict] = field(default_factory=list)
|
||||
free: bool = False
|
||||
|
||||
def log(self, **kwargs: object) -> None:
|
||||
kwargs.setdefault("timestamp", datetime.now(timezone.utc).isoformat())
|
||||
entry = ApiLogEntry(**kwargs) # type: ignore[arg-type]
|
||||
d = entry.model_dump()
|
||||
d["cost_usd"] = round(cost_for_entry(d), 6)
|
||||
if self.free:
|
||||
d["credits_charged"] = 0.0
|
||||
d["free"] = True
|
||||
else:
|
||||
# Attribute credits to this call using the same margin used by
|
||||
# the credit service. Local import to avoid a module-load cycle.
|
||||
from backend.services.billing_hook import get_billing
|
||||
|
||||
d["credits_charged"] = get_billing().credits_for_api_cost(d["cost_usd"])
|
||||
self.entries.append(d)
|
||||
|
||||
def to_jsonl(self) -> str:
|
||||
if not self.entries:
|
||||
return ""
|
||||
return "\n".join(json.dumps(e) for e in self.entries) + "\n"
|
||||
|
||||
def flush(self, storage, user_id: str, project_id: str) -> None:
|
||||
"""Write the current entries to ``api_logs.jsonl`` in storage.
|
||||
|
||||
Called periodically during a pipeline run so a preempted worker
|
||||
doesn't lose billing data. Idempotent — safe to call repeatedly;
|
||||
each flush overwrites the prior copy with the latest entries.
|
||||
"""
|
||||
text = self.to_jsonl()
|
||||
if not text:
|
||||
return
|
||||
# Local import avoids a cycle with services.projects (which imports
|
||||
# from services.storage which imports from here transitively).
|
||||
from backend.services.projects import project_prefix
|
||||
|
||||
key = f"{project_prefix(user_id, project_id)}/api_logs.jsonl"
|
||||
storage.write_text(key, text)
|
||||
@@ -0,0 +1,192 @@
|
||||
"""Open-core billing seam.
|
||||
|
||||
Everything outside the billing modules (``credits``, ``credit_grants``,
|
||||
``stripe_billing``, ``stripe_customer_map``, ``auto_topup`` and the
|
||||
``billing``/``credits`` routers) talks to billing exclusively through
|
||||
:func:`get_billing`. With ``BILLING_ENABLED=false`` the returned
|
||||
:class:`NullBilling` makes every pipeline run free — the same shape as the
|
||||
existing admin ``free=True`` path — so the core can run with no credits
|
||||
ledger, no Stripe, and no billing routes mounted.
|
||||
|
||||
This module must stay a leaf: no billing module is imported at module
|
||||
level (``CreditsBilling`` lazy-imports inside each method), so the core
|
||||
never touches the Stripe SDK when billing is disabled.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, Protocol
|
||||
|
||||
from backend.config import settings
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from backend.services.storage import StorageBackend
|
||||
|
||||
|
||||
class InsufficientCredits(RuntimeError):
|
||||
"""Raised when a charge would drop the balance below zero."""
|
||||
|
||||
def __init__(self, required: float, available: float) -> None:
|
||||
super().__init__(
|
||||
f"Insufficient credits: required {required}, available {available}"
|
||||
)
|
||||
self.required = required
|
||||
self.available = available
|
||||
|
||||
|
||||
class BillingHook(Protocol):
|
||||
"""The full billing surface the core is allowed to depend on."""
|
||||
|
||||
def credits_for_api_cost(self, cost_usd: float) -> float: ...
|
||||
|
||||
def get_balance(self, storage: "StorageBackend", user_id: str) -> float: ...
|
||||
|
||||
def charge(
|
||||
self,
|
||||
storage: "StorageBackend",
|
||||
user_id: str,
|
||||
amount: float,
|
||||
*,
|
||||
reason: str = "pipeline_charge",
|
||||
run_id: str | None = None,
|
||||
unit_id: str | None = None,
|
||||
allow_overdraft: bool = False,
|
||||
) -> None: ...
|
||||
|
||||
def ensure_trial_grant(self, storage: "StorageBackend", user_id: str) -> bool: ...
|
||||
|
||||
def list_user_ids(self, storage: "StorageBackend") -> list[str]: ...
|
||||
|
||||
async def maybe_auto_topup(
|
||||
self, storage: "StorageBackend", user_id: str
|
||||
) -> dict | None: ...
|
||||
|
||||
|
||||
class NullBilling:
|
||||
"""Billing disabled: everything is free and nothing is written.
|
||||
|
||||
``credits_for_api_cost`` returning 0.0 is the linchpin — every
|
||||
``ApiLogger`` entry gets ``credits_charged=0``, so the pipeline's
|
||||
charge path early-returns and the credit gate always allows.
|
||||
"""
|
||||
|
||||
def credits_for_api_cost(self, cost_usd: float) -> float:
|
||||
return 0.0
|
||||
|
||||
def get_balance(self, storage: "StorageBackend", user_id: str) -> float:
|
||||
return 0.0
|
||||
|
||||
def charge(
|
||||
self,
|
||||
storage: "StorageBackend",
|
||||
user_id: str,
|
||||
amount: float,
|
||||
*,
|
||||
reason: str = "pipeline_charge",
|
||||
run_id: str | None = None,
|
||||
unit_id: str | None = None,
|
||||
allow_overdraft: bool = False,
|
||||
) -> None:
|
||||
return None
|
||||
|
||||
def ensure_trial_grant(self, storage: "StorageBackend", user_id: str) -> bool:
|
||||
return False
|
||||
|
||||
def list_user_ids(self, storage: "StorageBackend") -> list[str]:
|
||||
return []
|
||||
|
||||
async def maybe_auto_topup(
|
||||
self, storage: "StorageBackend", user_id: str
|
||||
) -> dict | None:
|
||||
return None
|
||||
|
||||
|
||||
class CreditsBilling:
|
||||
"""Production billing: delegates to the credits ledger + auto top-up."""
|
||||
|
||||
def credits_for_api_cost(self, cost_usd: float) -> float:
|
||||
from backend.services import credits as credits_svc
|
||||
|
||||
return credits_svc.credits_for_api_cost(cost_usd)
|
||||
|
||||
def get_balance(self, storage: "StorageBackend", user_id: str) -> float:
|
||||
from backend.services import credits as credits_svc
|
||||
|
||||
return credits_svc.get_balance(storage, user_id)
|
||||
|
||||
def charge(
|
||||
self,
|
||||
storage: "StorageBackend",
|
||||
user_id: str,
|
||||
amount: float,
|
||||
*,
|
||||
reason: str = "pipeline_charge",
|
||||
run_id: str | None = None,
|
||||
unit_id: str | None = None,
|
||||
allow_overdraft: bool = False,
|
||||
) -> None:
|
||||
from backend.services import credits as credits_svc
|
||||
|
||||
credits_svc.charge(
|
||||
storage, user_id, amount,
|
||||
reason=reason,
|
||||
run_id=run_id,
|
||||
unit_id=unit_id,
|
||||
allow_overdraft=allow_overdraft,
|
||||
)
|
||||
|
||||
def ensure_trial_grant(self, storage: "StorageBackend", user_id: str) -> bool:
|
||||
from backend.services import credits as credits_svc
|
||||
|
||||
return credits_svc.ensure_trial_grant(storage, user_id)
|
||||
|
||||
def list_user_ids(self, storage: "StorageBackend") -> list[str]:
|
||||
from backend.services import credits as credits_svc
|
||||
|
||||
return credits_svc.list_user_ids(storage)
|
||||
|
||||
async def maybe_auto_topup(
|
||||
self, storage: "StorageBackend", user_id: str
|
||||
) -> dict | None:
|
||||
"""Run an auto top-up attempt if configured.
|
||||
|
||||
Returns ``{"reason", "amount_usd"}`` when this call produced a NEW
|
||||
failed attempt (so the caller can notify the user), else None.
|
||||
"""
|
||||
from backend.services.auto_topup import get_config, maybe_trigger
|
||||
|
||||
before = get_config(storage, user_id).last_attempt_ts
|
||||
try:
|
||||
await maybe_trigger(storage, user_id)
|
||||
except Exception:
|
||||
return None
|
||||
after = get_config(storage, user_id)
|
||||
if (
|
||||
after.last_attempt_status == "failed"
|
||||
and after.last_attempt_ts
|
||||
and after.last_attempt_ts != before
|
||||
):
|
||||
return {
|
||||
"reason": after.last_failure_reason or "unknown",
|
||||
"amount_usd": after.amount_usd,
|
||||
}
|
||||
return None
|
||||
|
||||
|
||||
_NULL = NullBilling()
|
||||
_credits_billing: CreditsBilling | None = None
|
||||
|
||||
|
||||
def get_billing() -> BillingHook:
|
||||
"""Return the active billing implementation.
|
||||
|
||||
Selected per call (not at import) so the ``billing_enabled`` setting
|
||||
can be monkeypatched in tests and so importing this module never pulls
|
||||
in billing code.
|
||||
"""
|
||||
if not settings.billing_enabled:
|
||||
return _NULL
|
||||
global _credits_billing
|
||||
if _credits_billing is None:
|
||||
_credits_billing = CreditsBilling()
|
||||
return _credits_billing
|
||||
@@ -0,0 +1,398 @@
|
||||
"""Pre-flight cost estimator for pipeline runs.
|
||||
|
||||
Walks the uploaded BOM + library cache and returns a low/high credit
|
||||
range the user will see *before* they start a run. Pure read-only:
|
||||
no storage writes, no API calls.
|
||||
|
||||
The estimator is intentionally conservative. Low/high bounds are
|
||||
bracketed around a central estimate (0.7× / 1.4×) so the user always
|
||||
sees a plausible range rather than a false-precision single number.
|
||||
|
||||
Per-call USD is computed from per-stage **token baselines**
|
||||
(``STAGE_TOKEN_BASELINES``) multiplied by the runtime-resolved
|
||||
provider+model rate from ``backend.services.llm.pricing.PRICING``.
|
||||
This means a change to ``PROVIDER_VALIDATION`` / ``MODEL_VALIDATION``
|
||||
(or any other per-stage routing env var) automatically updates the
|
||||
estimate — no constant-bumping required. The baselines themselves
|
||||
are hand-tuned from historical ``api_logs.jsonl`` aggregates and
|
||||
should be recalibrated periodically.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any, Literal
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
from backend.config import settings
|
||||
from backend.pinscopex.parsers import parse_bom
|
||||
from backend.pinscopex.resolve_passives import resolve_mpn
|
||||
from backend.pinscopex.taxonomy import SIMPLE_TYPES, type_for_ref
|
||||
from backend.pinscopex.utils import safe_mpn
|
||||
from backend.services import projects as proj_svc
|
||||
from backend.services.billing_hook import get_billing
|
||||
from backend.services.llm.pricing import CACHE_RATES, PRICING
|
||||
from backend.services.storage import StorageBackend
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Per-stage token baselines (model-aware estimator)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
# Average tokens per call for a single sub-unit of each stage. Values
|
||||
# come from aggregating recent ``api_logs.jsonl`` runs across staging +
|
||||
# prod (see ``scripts/recalibrate_estimator_baselines.py`` follow-up;
|
||||
# until that lands, eyeball + paste from gcloud-mined stats).
|
||||
#
|
||||
# ``settings_stage`` is the key passed to ``settings.model_for_stage`` /
|
||||
# ``settings.provider_for_stage``. The estimator-stage names mirror the
|
||||
# ``CostItem.kind`` Literal so the breakdown stays self-consistent.
|
||||
STAGE_TOKEN_BASELINES: dict[str, dict[str, int | str]] = {
|
||||
"ic_extraction": {
|
||||
"settings_stage": "pintable",
|
||||
"input": 100, "output": 2000,
|
||||
"cache_create": 80_000, "cache_read": 170_000,
|
||||
},
|
||||
"simple_extraction": {
|
||||
"settings_stage": "specs",
|
||||
"input": 100, "output": 1000,
|
||||
"cache_create": 20_000, "cache_read": 20_000,
|
||||
},
|
||||
"passive_pattern": {
|
||||
"settings_stage": "pattern",
|
||||
"input": 100, "output": 7000,
|
||||
"cache_create": 60_000, "cache_read": 330_000,
|
||||
},
|
||||
"digikey_resolve": {
|
||||
"settings_stage": "auto_resolve",
|
||||
"input": 2000, "output": 200,
|
||||
"cache_create": 0, "cache_read": 0,
|
||||
},
|
||||
# Validation review per IC. The multi-turn validator pulls the
|
||||
# cached system + graph + datasheet on each turn (~4-5 turns/IC),
|
||||
# so cache_read dominates. Page-aware scaling was abandoned — token
|
||||
# counts already encompass the PDF via cache reuse, and IC
|
||||
# complexity correlates more weakly with raw page count than the
|
||||
# old per-page heuristic assumed.
|
||||
"review": {
|
||||
"settings_stage": "validation",
|
||||
"input": 13_500, "output": 2000,
|
||||
"cache_create": 110_000, "cache_read": 300_000,
|
||||
},
|
||||
# Per-IC normalize pass — dedup + severity re-grade. Runs on the
|
||||
# already-structured findings (no PDF, no graph tools), one turn,
|
||||
# validation-class model (Sonnet). A few hundred input tokens for the
|
||||
# rubric, a few hundred output tokens for the normalized list.
|
||||
"normalize": {
|
||||
"settings_stage": "normalize",
|
||||
"input": 1500, "output": 600,
|
||||
"cache_create": 0, "cache_read": 0,
|
||||
},
|
||||
# Cross-IC dedup — one call per run over all findings (no PDF/graph),
|
||||
# validation-class model (Sonnet). Slightly larger input than per-IC
|
||||
# normalize since it sees every IC's findings at once.
|
||||
"cross_ic_dedup": {
|
||||
"settings_stage": "normalize",
|
||||
"input": 2500, "output": 700,
|
||||
"cache_create": 0, "cache_read": 0,
|
||||
},
|
||||
}
|
||||
|
||||
LOW_MULT: float = 0.7
|
||||
HIGH_MULT: float = 1.4
|
||||
|
||||
|
||||
def estimate_stage_cost_usd(stage: str) -> float:
|
||||
"""Per-call USD for one sub-unit of ``stage``, model-aware.
|
||||
|
||||
Resolves provider+model from ``settings`` and computes
|
||||
``(input_tokens × rate) + ...`` using the same ``PRICING`` /
|
||||
``CACHE_RATES`` tables that real billing in
|
||||
``services.llm.pricing.cost_for_entry`` reads. Changing a
|
||||
``MODEL_*`` / ``PROVIDER_*`` env var therefore updates the estimate
|
||||
automatically.
|
||||
|
||||
Falls through to ``PRICING[provider]["default"]`` when the resolved
|
||||
model is missing from the table — same fallback semantics as the
|
||||
real billing code, so a missing pricing entry surfaces uniformly
|
||||
everywhere instead of crashing the estimator.
|
||||
|
||||
Raises ``KeyError`` only when ``stage`` itself is unknown.
|
||||
"""
|
||||
base = STAGE_TOKEN_BASELINES[stage]
|
||||
settings_stage = str(base["settings_stage"])
|
||||
provider = settings.provider_for_stage(settings_stage)
|
||||
model = settings.model_for_stage(settings_stage)
|
||||
table = PRICING.get(provider) or PRICING["anthropic"]
|
||||
rates = table.get(model, table["default"])
|
||||
cache = CACHE_RATES.get(provider, CACHE_RATES["anthropic"])
|
||||
return (
|
||||
int(base["input"]) * rates["input"]
|
||||
+ int(base["output"]) * rates["output"]
|
||||
+ int(base["cache_create"]) * rates["input"] * cache["create"]
|
||||
+ int(base["cache_read"]) * rates["input"] * cache["read"]
|
||||
) / 1_000_000
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Data classes
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
UnitKind = Literal[
|
||||
"ic_extraction",
|
||||
"simple_extraction",
|
||||
"passive_pattern",
|
||||
"digikey_resolve",
|
||||
"review",
|
||||
]
|
||||
|
||||
|
||||
class CostItem(BaseModel):
|
||||
identifier: str # MPN, ref, or a fixed token like a stage name
|
||||
kind: UnitKind
|
||||
api_cost_usd: float
|
||||
source: Literal["cache_hit", "api_call", "api_call_estimated"]
|
||||
note: str | None = None
|
||||
|
||||
|
||||
class CostEstimate(BaseModel):
|
||||
"""What the pipeline will likely cost this run."""
|
||||
api_cost_low: float
|
||||
api_cost_high: float
|
||||
api_cost_mid: float
|
||||
credits_low: float
|
||||
credits_high: float
|
||||
credits_mid: float
|
||||
breakdown: list[CostItem]
|
||||
ic_count: int
|
||||
simple_count: int
|
||||
passive_count: int
|
||||
cached_ic_count: int
|
||||
cached_simple_count: int
|
||||
cached_passive_count: int
|
||||
review_ic_count: int
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Internals
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _load_library_patterns(storage: StorageBackend):
|
||||
"""Load passive patterns from the library to test cache hits.
|
||||
|
||||
This mirrors the pipeline's own seeding behaviour but keeps the
|
||||
estimator synchronous and side-effect free (downloads to a local
|
||||
tempdir only if the backend is remote).
|
||||
"""
|
||||
try:
|
||||
return proj_svc.load_library_patterns(storage)
|
||||
except Exception:
|
||||
return []
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Public entry point
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def estimate_pipeline_cost(
|
||||
storage: StorageBackend,
|
||||
user_id: str,
|
||||
project_id: str,
|
||||
) -> CostEstimate:
|
||||
"""Produce a CostEstimate for the given project without running anything.
|
||||
|
||||
The BOM must already be uploaded; the netlist may or may not be. If
|
||||
the BOM is missing, raises FileNotFoundError.
|
||||
"""
|
||||
bom_key = proj_svc.get_bom_key(storage, user_id, project_id)
|
||||
if not bom_key:
|
||||
raise FileNotFoundError("BOM not uploaded for this project")
|
||||
|
||||
meta = proj_svc.get_project(storage, user_id, project_id)
|
||||
col_map = (meta.bom_columns if meta else None) or {}
|
||||
ref_col = col_map.get("reference", "Reference")
|
||||
mpn_col = col_map.get("mpn", "Manufacturer Part Number")
|
||||
|
||||
# Download BOM to a local temp path so parse_bom can read it.
|
||||
import tempfile
|
||||
|
||||
with tempfile.NamedTemporaryFile(suffix=".csv", delete=False) as tmp:
|
||||
tmp.write(storage.read_bytes(bom_key))
|
||||
bom_local_path = Path(tmp.name)
|
||||
|
||||
try:
|
||||
bom = parse_bom(str(bom_local_path), reference_col=ref_col, mpn_col=mpn_col)
|
||||
finally:
|
||||
bom_local_path.unlink(missing_ok=True)
|
||||
|
||||
# Classify unique MPNs by type.
|
||||
ic_mpns: set[str] = set()
|
||||
simple_mpns: set[str] = set()
|
||||
passive_mpns: set[str] = set()
|
||||
for ref, info in bom.items():
|
||||
mpn = info.get("mpn")
|
||||
if not mpn:
|
||||
continue
|
||||
typ = type_for_ref(ref)
|
||||
if typ == "ic":
|
||||
ic_mpns.add(mpn)
|
||||
elif typ == "passive":
|
||||
passive_mpns.add(mpn)
|
||||
elif typ in SIMPLE_TYPES:
|
||||
simple_mpns.add(mpn)
|
||||
|
||||
# Load library patterns once so we can resolve passives against cache.
|
||||
patterns = _load_library_patterns(storage)
|
||||
|
||||
breakdown: list[CostItem] = []
|
||||
cached_ic = 0
|
||||
cached_simple = 0
|
||||
cached_passive = 0
|
||||
|
||||
# IC extraction
|
||||
for mpn in sorted(ic_mpns):
|
||||
if proj_svc.library_has_extraction(storage, mpn):
|
||||
breakdown.append(CostItem(
|
||||
identifier=mpn, kind="ic_extraction",
|
||||
api_cost_usd=0.0, source="cache_hit",
|
||||
note="library hit",
|
||||
))
|
||||
cached_ic += 1
|
||||
else:
|
||||
breakdown.append(CostItem(
|
||||
identifier=mpn, kind="ic_extraction",
|
||||
api_cost_usd=round(estimate_stage_cost_usd("ic_extraction"), 4),
|
||||
source="api_call_estimated",
|
||||
))
|
||||
|
||||
# Simple component specs
|
||||
for mpn in sorted(simple_mpns):
|
||||
if proj_svc.library_has_model(storage, mpn):
|
||||
breakdown.append(CostItem(
|
||||
identifier=mpn, kind="simple_extraction",
|
||||
api_cost_usd=0.0, source="cache_hit",
|
||||
))
|
||||
cached_simple += 1
|
||||
else:
|
||||
breakdown.append(CostItem(
|
||||
identifier=mpn, kind="simple_extraction",
|
||||
api_cost_usd=round(estimate_stage_cost_usd("simple_extraction"), 4),
|
||||
source="api_call_estimated",
|
||||
))
|
||||
|
||||
# Passives — pattern resolution covers many MPNs with one pattern
|
||||
unresolved_passives: list[str] = []
|
||||
for mpn in sorted(passive_mpns):
|
||||
if patterns and resolve_mpn(mpn, patterns) is not None:
|
||||
breakdown.append(CostItem(
|
||||
identifier=mpn, kind="passive_pattern",
|
||||
api_cost_usd=0.0, source="cache_hit",
|
||||
note="pattern match",
|
||||
))
|
||||
cached_passive += 1
|
||||
continue
|
||||
if proj_svc.library_has_passive_model(storage, mpn):
|
||||
breakdown.append(CostItem(
|
||||
identifier=mpn, kind="passive_pattern",
|
||||
api_cost_usd=0.0, source="cache_hit",
|
||||
note="cached passive model",
|
||||
))
|
||||
cached_passive += 1
|
||||
continue
|
||||
unresolved_passives.append(mpn)
|
||||
|
||||
# Each unresolved passive MPN may contribute one pattern extraction.
|
||||
# Heuristic: N unique first-7-char prefixes = N new patterns.
|
||||
prefixes = {m[:7] for m in unresolved_passives}
|
||||
for prefix in sorted(prefixes):
|
||||
sample_mpn = next(m for m in unresolved_passives if m.startswith(prefix))
|
||||
breakdown.append(CostItem(
|
||||
identifier=sample_mpn, kind="passive_pattern",
|
||||
api_cost_usd=round(estimate_stage_cost_usd("passive_pattern"), 4),
|
||||
source="api_call_estimated",
|
||||
note=f"may cover {sum(1 for m in unresolved_passives if m.startswith(prefix))} MPNs",
|
||||
))
|
||||
|
||||
# Direct datasheet review — per IC that has a datasheet available.
|
||||
# Cost is the flat per-IC observed average; multi-turn cache reuse
|
||||
# makes this less page-sensitive than the old per-page heuristic
|
||||
# implied.
|
||||
review_per_ic = estimate_stage_cost_usd("review")
|
||||
if settings.normalize_findings_enabled:
|
||||
review_per_ic += estimate_stage_cost_usd("normalize")
|
||||
review_ic_count = 0
|
||||
for mpn in sorted(ic_mpns):
|
||||
pdf_path = _locate_datasheet_local(storage, user_id, project_id, mpn)
|
||||
has_pdf = pdf_path is not None
|
||||
has_library_pdf = (
|
||||
proj_svc.library_has_datasheet(storage, mpn) is not None
|
||||
if not has_pdf else False
|
||||
)
|
||||
if not (has_pdf or has_library_pdf):
|
||||
continue # Skipped in pipeline — no datasheet, no review
|
||||
breakdown.append(CostItem(
|
||||
identifier=mpn, kind="review",
|
||||
api_cost_usd=round(review_per_ic, 4),
|
||||
source="api_call_estimated",
|
||||
))
|
||||
review_ic_count += 1
|
||||
|
||||
# One cross-IC dedup call per run, only when ≥2 ICs get reviewed (a
|
||||
# single-IC run has no cross-IC pair to merge — see _maybe_dedupe_cross_ic).
|
||||
if settings.cross_ic_dedup_enabled and review_ic_count > 1:
|
||||
breakdown.append(CostItem(
|
||||
identifier="cross-IC dedup", kind="review",
|
||||
api_cost_usd=round(estimate_stage_cost_usd("cross_ic_dedup"), 4),
|
||||
source="api_call_estimated",
|
||||
note="collapses one interface defect reported from both ICs",
|
||||
))
|
||||
|
||||
api_total = sum(item.api_cost_usd for item in breakdown)
|
||||
api_low = round(api_total * LOW_MULT, 4)
|
||||
api_high = round(api_total * HIGH_MULT, 4)
|
||||
|
||||
billing = get_billing()
|
||||
return CostEstimate(
|
||||
api_cost_low=api_low,
|
||||
api_cost_high=api_high,
|
||||
api_cost_mid=round(api_total, 4),
|
||||
credits_low=billing.credits_for_api_cost(api_low),
|
||||
credits_high=billing.credits_for_api_cost(api_high),
|
||||
credits_mid=billing.credits_for_api_cost(api_total),
|
||||
breakdown=breakdown,
|
||||
ic_count=len(ic_mpns),
|
||||
simple_count=len(simple_mpns),
|
||||
passive_count=len(passive_mpns),
|
||||
cached_ic_count=cached_ic,
|
||||
cached_simple_count=cached_simple,
|
||||
cached_passive_count=cached_passive,
|
||||
review_ic_count=review_ic_count,
|
||||
)
|
||||
|
||||
|
||||
def _locate_datasheet_local(
|
||||
storage: StorageBackend, user_id: str, project_id: str, mpn: str,
|
||||
) -> Path | None:
|
||||
"""Return a local Path to the datasheet PDF if it can be read quickly.
|
||||
|
||||
For LocalStorageBackend, reads directly from disk. For remote backends
|
||||
we skip the page-count read (returns None) — estimator will fall back
|
||||
to the mid-cap heuristic rather than downloading the PDF during pre-flight.
|
||||
"""
|
||||
from backend.services.storage import LocalStorageBackend
|
||||
|
||||
if not isinstance(storage, LocalStorageBackend):
|
||||
return None
|
||||
safe = safe_mpn(mpn)
|
||||
key = f"users/{user_id}/projects/{project_id}/uploads/datasheets/{safe}.pdf"
|
||||
if storage.exists(key):
|
||||
return storage._path(key) # type: ignore[attr-defined]
|
||||
legacy = f"library/datasheets/{safe}.pdf"
|
||||
if storage.exists(legacy):
|
||||
return storage._path(legacy) # type: ignore[attr-defined]
|
||||
return None
|
||||
@@ -0,0 +1,168 @@
|
||||
"""Content-addressed datasheet storage for the shared library.
|
||||
|
||||
Stores PDF blobs by their MD5 hash and creates lightweight JSON ref files
|
||||
that map MPN names to blob keys. This deduplicates identical PDFs that
|
||||
were previously stored under different human-readable names.
|
||||
|
||||
Layout::
|
||||
|
||||
library/datasheets/
|
||||
blobs/{md5hash}.pdf -- unique PDF content, stored once
|
||||
refs/{safe_mpn}.json -- per-MPN pointer: {"hash": "...", "blob_key": "..."}
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
from pathlib import Path
|
||||
|
||||
from backend.pinscopex.utils import safe_mpn
|
||||
from backend.services.storage import StorageBackend
|
||||
|
||||
BLOB_PREFIX = "library/datasheets/blobs/"
|
||||
REF_PREFIX = "library/datasheets/refs/"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Hashing helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def compute_md5_from_path(local_path: Path) -> str:
|
||||
"""Return the hex MD5 digest of a local file (chunked read)."""
|
||||
h = hashlib.md5()
|
||||
with open(local_path, "rb") as f:
|
||||
for chunk in iter(lambda: f.read(8192), b""):
|
||||
h.update(chunk)
|
||||
return h.hexdigest()
|
||||
|
||||
|
||||
def compute_md5_from_bytes(data: bytes) -> str:
|
||||
"""Return the hex MD5 digest of in-memory bytes."""
|
||||
return hashlib.md5(data).hexdigest()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Key construction
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def blob_key(md5: str) -> str:
|
||||
"""Storage key for a content-addressed PDF blob."""
|
||||
return f"{BLOB_PREFIX}{md5}.pdf"
|
||||
|
||||
|
||||
def ref_key(mpn: str) -> str:
|
||||
"""Storage key for an MPN → blob ref file."""
|
||||
return f"{REF_PREFIX}{safe_mpn(mpn)}.json"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Store / resolve / delete
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def store_datasheet(
|
||||
storage: StorageBackend,
|
||||
local_path: Path,
|
||||
mpn: str,
|
||||
) -> str:
|
||||
"""Store a datasheet PDF by content hash and create an MPN ref.
|
||||
|
||||
Idempotent: skips blob upload if it already exists, always writes the ref.
|
||||
Returns the blob storage key.
|
||||
"""
|
||||
md5 = compute_md5_from_path(local_path)
|
||||
bk = blob_key(md5)
|
||||
if not storage.exists(bk):
|
||||
storage.upload_from_local(local_path, bk)
|
||||
storage.write_json(ref_key(mpn), {"hash": md5, "blob_key": bk})
|
||||
return bk
|
||||
|
||||
|
||||
def store_datasheet_bytes(
|
||||
storage: StorageBackend,
|
||||
data: bytes,
|
||||
mpn: str,
|
||||
) -> str:
|
||||
"""Same as :func:`store_datasheet` but from in-memory bytes."""
|
||||
md5 = compute_md5_from_bytes(data)
|
||||
bk = blob_key(md5)
|
||||
if not storage.exists(bk):
|
||||
storage.write_bytes(bk, data)
|
||||
storage.write_json(ref_key(mpn), {"hash": md5, "blob_key": bk})
|
||||
return bk
|
||||
|
||||
|
||||
def resolve_datasheet(storage: StorageBackend, mpn: str) -> str | None:
|
||||
"""Look up the blob key for an MPN via its ref file.
|
||||
|
||||
Returns the blob key if the ref exists *and* the blob exists, else None.
|
||||
"""
|
||||
rk = ref_key(mpn)
|
||||
if not storage.exists(rk):
|
||||
return None
|
||||
ref = storage.read_json(rk)
|
||||
bk = ref.get("blob_key")
|
||||
if bk and storage.exists(bk):
|
||||
return bk
|
||||
return None
|
||||
|
||||
|
||||
def delete_datasheet_ref(storage: StorageBackend, mpn: str) -> str | None:
|
||||
"""Delete the ref for an MPN. Returns the blob key if a ref existed.
|
||||
|
||||
Does **not** delete the blob — other refs may point to it. Use
|
||||
:func:`gc_orphan_blobs` to clean up unreferenced blobs.
|
||||
"""
|
||||
rk = ref_key(mpn)
|
||||
if not storage.exists(rk):
|
||||
return None
|
||||
ref = storage.read_json(rk)
|
||||
bk = ref.get("blob_key")
|
||||
storage.delete_key(rk)
|
||||
return bk
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Maintenance
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def gc_orphan_blobs(
|
||||
storage: StorageBackend, *, dry_run: bool = True,
|
||||
) -> list[str]:
|
||||
"""Find blobs not referenced by any ref file. Optionally delete them.
|
||||
|
||||
Also checks pattern ``datasheet_key`` values so blobs referenced only
|
||||
by patterns (not MPN refs) are kept.
|
||||
|
||||
Intended for maintenance scripts, not hot paths.
|
||||
"""
|
||||
# Collect all hashes referenced by ref files
|
||||
referenced_hashes: set[str] = set()
|
||||
for rk in storage.list_recursive(REF_PREFIX):
|
||||
if rk.endswith(".json"):
|
||||
ref = storage.read_json(rk)
|
||||
h = ref.get("hash")
|
||||
if h:
|
||||
referenced_hashes.add(h)
|
||||
|
||||
# Also collect hashes from pattern datasheet_key values
|
||||
for pk in storage.list_recursive("library/patterns/"):
|
||||
if pk.endswith(".json"):
|
||||
pat = storage.read_json(pk)
|
||||
ds_key = pat.get("datasheet_key", "")
|
||||
if ds_key.startswith(BLOB_PREFIX) and ds_key.endswith(".pdf"):
|
||||
h = ds_key.removeprefix(BLOB_PREFIX).removesuffix(".pdf")
|
||||
referenced_hashes.add(h)
|
||||
|
||||
# Find orphan blobs
|
||||
orphans: list[str] = []
|
||||
for bk in storage.list_recursive(BLOB_PREFIX):
|
||||
if not bk.endswith(".pdf"):
|
||||
continue
|
||||
filename = bk.rsplit("/", 1)[-1]
|
||||
h = filename.removesuffix(".pdf")
|
||||
if h not in referenced_hashes:
|
||||
orphans.append(bk)
|
||||
if not dry_run:
|
||||
storage.delete_key(bk)
|
||||
|
||||
return orphans
|
||||
@@ -0,0 +1,398 @@
|
||||
"""Cross-IC dedup pass — collapse one physical defect reported from both ends.
|
||||
|
||||
Direct datasheet review runs once per IC, in isolation. An interface defect
|
||||
(e.g. a 5V driver into a non-5V-tolerant input on the U2↔U3 UART) is therefore
|
||||
discovered twice — once when reviewing U2, once when reviewing U3 — and the
|
||||
per-IC normalize pass cannot collapse them because it only sees one IC's
|
||||
findings at a time. The two copies land in the report as separate findings,
|
||||
double-counting the same problem and (worse) sometimes disagreeing with each
|
||||
other.
|
||||
|
||||
This module runs a single small LLM call over the *concatenated* findings from
|
||||
all ICs (no PDF, no graph tools) to merge findings that describe the same
|
||||
physical defect on the same net/interface/component. It is the cross-IC analog
|
||||
of ``normalize_findings`` and follows the same fail-soft contract: on any LLM
|
||||
error, schema violation, or coverage gap, the original findings are returned
|
||||
unchanged. It never drops findings (that is normalize's job) and never raises a
|
||||
merged finding's severity above the highest of its members.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
from datetime import datetime, timezone
|
||||
from typing import Awaitable, Callable
|
||||
|
||||
from backend.pinscopex.models import Finding
|
||||
from backend.services.api_logs import ApiLogger
|
||||
from backend.services.llm import Message, TextBlock
|
||||
from backend.services.llm.factory import call_with_fallback
|
||||
from backend.services.llm.types import ToolSchema
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
# Severity ordering — a merged group's severity is capped at the highest
|
||||
# severity among its members (downgrade-only, same principle as normalize).
|
||||
_INFO, _WARN, _ERR = 0, 1, 2
|
||||
_SEVERITY_RANK = {"INFO": _INFO, "WARNING": _WARN, "ERROR": _ERR}
|
||||
_RANK_TO_SEV = {_INFO: "INFO", _WARN: "WARNING", _ERR: "ERROR"}
|
||||
|
||||
|
||||
def _is_unverified(why: str | None) -> bool:
|
||||
return (why or "").lstrip().lower().startswith("unverified:")
|
||||
|
||||
|
||||
SYSTEM_PROMPT = """\
|
||||
You deduplicate hardware-review findings across multiple ICs.
|
||||
|
||||
Each finding was produced by reviewing one IC in isolation, so a defect on
|
||||
the interface *between* two ICs is reported twice — once from each side. Your
|
||||
job is to group findings that describe the SAME physical defect and merge each
|
||||
group into one finding. You may NOT invent findings, drop findings, or change
|
||||
the engineering substance.
|
||||
|
||||
### When two findings are the same defect (merge)
|
||||
|
||||
Merge when they describe the same physical problem at the same place:
|
||||
- the same net or signal (e.g. both flag over-voltage on `/UART0.NCTS`),
|
||||
- the same component pair / interface (e.g. "U2 RTS# drives U3 PA14" and
|
||||
"U3 PA14 is driven by U2's 5V output" are one interface defect seen from
|
||||
each end),
|
||||
- the same shared part with the same fix.
|
||||
|
||||
A merged group is resolved by ONE change to the design. Name that interface or
|
||||
root cause once in the merged `finding`; restate each side's consequence in
|
||||
`why`.
|
||||
|
||||
### When findings are NOT the same defect (keep separate)
|
||||
|
||||
Do NOT merge findings that need different fixes, even if they touch the same
|
||||
component or net:
|
||||
- different pins / different signals on the same IC,
|
||||
- a decoupling issue and a voltage issue on the same supply,
|
||||
- two unrelated problems that happen to involve the same part.
|
||||
|
||||
When in doubt, keep them separate. Over-merging hides distinct problems and is
|
||||
worse than a visible duplicate.
|
||||
|
||||
### Severity
|
||||
|
||||
Use the HIGHEST severity among a group's members. Never grade a merged finding
|
||||
above its strongest member. If any member's `why` begins with `Unverified:`,
|
||||
keep that prefix and do not grade the merged finding above WARNING.
|
||||
|
||||
### Output
|
||||
|
||||
Call `submit_deduped` exactly once. Provide a `groups` array. Every original
|
||||
finding (numbered 1..N) must appear in exactly one group's `member_indices`,
|
||||
and no index may appear twice.
|
||||
- A group of ONE index is a passthrough — it is kept unchanged (you do not
|
||||
need to restate its text).
|
||||
- A group of MORE THAN ONE index is a merge — supply the merged `finding`,
|
||||
`why`, `status`, `recommendation`, and a `primary_index` (one of the group's
|
||||
members) whose datasheet citation/page is the strongest evidence; that
|
||||
member supplies the finding's component attribution and source reference.
|
||||
"""
|
||||
|
||||
|
||||
SUBMIT_DEDUPED_SCHEMA = ToolSchema(
|
||||
name="submit_deduped",
|
||||
description=(
|
||||
"Submit the cross-IC deduplicated findings. Every original finding "
|
||||
"(1..N) must appear in exactly one group's `member_indices`."
|
||||
),
|
||||
input_schema={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"groups": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"member_indices": {
|
||||
"type": "array",
|
||||
"items": {"type": "integer"},
|
||||
"minItems": 1,
|
||||
"description": (
|
||||
"1-indexed positions in the original findings "
|
||||
"list this group represents. Length 1 = "
|
||||
"passthrough; length > 1 = merge."
|
||||
),
|
||||
},
|
||||
"primary_index": {
|
||||
"type": ["integer", "null"],
|
||||
"description": (
|
||||
"REQUIRED when member_indices has length > 1: "
|
||||
"the member whose datasheet citation/source is "
|
||||
"the strongest. Supplies the merged finding's "
|
||||
"component attribution and source reference. "
|
||||
"Must be one of member_indices."
|
||||
),
|
||||
},
|
||||
"finding": {"type": "string"},
|
||||
"why": {"type": "string"},
|
||||
"status": {
|
||||
"type": "string",
|
||||
"enum": ["ERROR", "WARNING", "INFO"],
|
||||
},
|
||||
"recommendation": {"type": "string"},
|
||||
"change_rationale": {
|
||||
"type": "string",
|
||||
"description": (
|
||||
"≤1 line: 'passthrough', or 'merged N+M: "
|
||||
"<shared interface/root cause>'."
|
||||
),
|
||||
},
|
||||
},
|
||||
"required": ["member_indices", "change_rationale"],
|
||||
},
|
||||
},
|
||||
},
|
||||
"required": ["groups"],
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _serialize_findings_for_prompt(findings: list[Finding]) -> str:
|
||||
"""Number findings 1..N with their IC, severity, and text.
|
||||
|
||||
Unlike the per-IC normalize pass, the designator IS included — it is the
|
||||
primary signal for spotting that two findings sit on opposite ends of one
|
||||
interface.
|
||||
"""
|
||||
rows: list[dict] = []
|
||||
for i, f in enumerate(findings, start=1):
|
||||
rows.append({
|
||||
"index": i,
|
||||
"ic": f.designator,
|
||||
"mpn": f.mpn,
|
||||
"reviewer_severity": f.status,
|
||||
"finding": f.finding,
|
||||
"why": f.why,
|
||||
"recommendation": f.recommendation,
|
||||
"source_page": f.source_page,
|
||||
"reference": f.reference,
|
||||
})
|
||||
return json.dumps(rows, indent=2)
|
||||
|
||||
|
||||
def _build_deduped(
|
||||
raw_groups: list[dict],
|
||||
originals: list[Finding],
|
||||
) -> list[Finding] | None:
|
||||
"""Validate the tool output and reconstruct the deduped finding list.
|
||||
|
||||
Returns the kept/merged findings, or ``None`` if coverage/schema
|
||||
validation fails (caller falls back to originals). A merge that omits a
|
||||
valid ``primary_index`` is not a hard failure — that group falls back to
|
||||
its per-index originals (un-merged), preserving coverage and severities.
|
||||
"""
|
||||
n = len(originals)
|
||||
seen: set[int] = set()
|
||||
result: list[Finding] = []
|
||||
|
||||
for group in raw_groups:
|
||||
if not isinstance(group, dict):
|
||||
return None
|
||||
member_indices = group.get("member_indices") or []
|
||||
if not isinstance(member_indices, list) or not member_indices:
|
||||
return None
|
||||
try:
|
||||
indices = [int(x) for x in member_indices]
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
for idx in indices:
|
||||
if idx < 1 or idx > n or idx in seen:
|
||||
return None
|
||||
seen.add(idx)
|
||||
|
||||
# Passthrough — keep the original verbatim. No laundering of text or
|
||||
# severity for a finding the model chose not to merge.
|
||||
if len(indices) == 1:
|
||||
result.append(originals[indices[0] - 1])
|
||||
continue
|
||||
|
||||
# Merge — needs a valid primary_index naming the canonical member.
|
||||
# Missing/invalid → un-merge to per-index originals (coverage kept).
|
||||
primary_raw = group.get("primary_index")
|
||||
try:
|
||||
primary = int(primary_raw)
|
||||
except (TypeError, ValueError):
|
||||
primary = None
|
||||
if primary not in indices:
|
||||
log.warning(
|
||||
"dedupe: merge of %s has invalid primary_index %r — "
|
||||
"falling back to per-index originals (un-merging)",
|
||||
indices, primary_raw,
|
||||
)
|
||||
for idx in indices:
|
||||
result.append(originals[idx - 1])
|
||||
continue
|
||||
|
||||
canon = originals[primary - 1]
|
||||
members = [originals[i - 1] for i in indices]
|
||||
ceiling = max(_SEVERITY_RANK.get(m.status, _ERR) for m in members)
|
||||
unverified = any(_is_unverified(m.why) for m in members)
|
||||
if unverified:
|
||||
ceiling = min(ceiling, _WARN)
|
||||
proposed = str(group.get("status") or canon.status)
|
||||
final_status = _RANK_TO_SEV[
|
||||
min(_SEVERITY_RANK.get(proposed, ceiling), ceiling)
|
||||
]
|
||||
|
||||
new_why = str(group.get("why") or canon.why)
|
||||
if unverified and not _is_unverified(new_why):
|
||||
new_why = "Unverified: " + new_why
|
||||
|
||||
try:
|
||||
result.append(Finding(
|
||||
finding_id=canon.finding_id,
|
||||
designator=canon.designator,
|
||||
mpn=canon.mpn,
|
||||
aspect=canon.aspect,
|
||||
finding=str(group.get("finding") or canon.finding),
|
||||
why=new_why,
|
||||
source_page=group.get("source_page", canon.source_page),
|
||||
source_quote=canon.source_quote,
|
||||
source_designator=canon.source_designator,
|
||||
status=final_status,
|
||||
recommendation=str(
|
||||
group.get("recommendation") or canon.recommendation
|
||||
),
|
||||
reference=str(group.get("reference") or canon.reference),
|
||||
source=canon.source,
|
||||
))
|
||||
except Exception:
|
||||
log.exception("dedupe: failed to build merged Finding")
|
||||
return None
|
||||
|
||||
if seen != set(range(1, n + 1)):
|
||||
return None
|
||||
return result
|
||||
|
||||
|
||||
async def dedupe_cross_ic_findings_async(
|
||||
findings: list[Finding],
|
||||
*,
|
||||
api_logger: ApiLogger | None = None,
|
||||
on_progress: Callable[[str, int, str, str], Awaitable[None]] | None = None,
|
||||
) -> tuple[list[Finding], dict]:
|
||||
"""Run the cross-IC dedup pass over findings from every IC.
|
||||
|
||||
Returns ``(deduped_findings, trace)``. On any failure (LLM error, schema
|
||||
violation, coverage gap) returns the original findings unchanged with an
|
||||
``error`` field set in the trace.
|
||||
"""
|
||||
trace: dict = {
|
||||
"timestamp": datetime.now(timezone.utc).isoformat(),
|
||||
"input_findings": [f.model_dump(mode="json") for f in findings],
|
||||
"output_findings": None,
|
||||
"submission": None,
|
||||
"model": None,
|
||||
"provider": None,
|
||||
"duration_ms": None,
|
||||
"error": None,
|
||||
}
|
||||
|
||||
# Nothing to merge across fewer than two findings.
|
||||
if len(findings) < 2:
|
||||
trace["output_findings"] = trace["input_findings"]
|
||||
trace["error"] = "skipped: <2 findings"
|
||||
return findings, trace
|
||||
|
||||
user_text = (
|
||||
f"There are {len(findings)} findings across all reviewed ICs. "
|
||||
f"Indices are 1-based. Group findings that describe the same physical "
|
||||
f"defect (especially the same interface seen from both ICs) and call "
|
||||
f"submit_deduped.\n\n{_serialize_findings_for_prompt(findings)}"
|
||||
)
|
||||
|
||||
t0 = time.monotonic()
|
||||
|
||||
async def _run(provider, model):
|
||||
trace["model"] = model
|
||||
trace["provider"] = provider.name
|
||||
session = await provider.create_session(
|
||||
model=model,
|
||||
system=SYSTEM_PROMPT,
|
||||
max_tokens=4096,
|
||||
temperature=0.0,
|
||||
)
|
||||
try:
|
||||
completion = await session.complete(
|
||||
messages=[Message(
|
||||
role="user",
|
||||
content=[TextBlock(text=user_text, cacheable=False)],
|
||||
)],
|
||||
tools=[SUBMIT_DEDUPED_SCHEMA],
|
||||
tool_choice={"name": "submit_deduped"},
|
||||
)
|
||||
if api_logger:
|
||||
api_logger.log(
|
||||
stage="cross_ic_dedupe",
|
||||
identifier="all",
|
||||
model=model,
|
||||
provider=provider.name,
|
||||
input_tokens=completion.usage.input_tokens,
|
||||
output_tokens=completion.usage.output_tokens,
|
||||
cache_creation_input_tokens=completion.usage.cache_creation_tokens,
|
||||
cache_read_input_tokens=completion.usage.cache_read_tokens,
|
||||
duration_ms=int((time.monotonic() - t0) * 1000),
|
||||
stop_reason="submit_deduped",
|
||||
turns=1,
|
||||
)
|
||||
for tc in completion.tool_calls:
|
||||
if tc.name == "submit_deduped":
|
||||
return tc.input
|
||||
return None
|
||||
finally:
|
||||
await session.close()
|
||||
|
||||
try:
|
||||
# Reuse the "normalize" stage config (validation-class Sonnet model +
|
||||
# any configured fallback); the log entry above is stamped
|
||||
# "cross_ic_dedupe" so cost accounting still distinguishes it.
|
||||
submission = await call_with_fallback("normalize", _run)
|
||||
except Exception as exc:
|
||||
log.exception("dedupe: call failed")
|
||||
trace["error"] = f"{type(exc).__name__}: {exc}"
|
||||
trace["duration_ms"] = int((time.monotonic() - t0) * 1000)
|
||||
trace["output_findings"] = trace["input_findings"]
|
||||
return findings, trace
|
||||
|
||||
trace["duration_ms"] = int((time.monotonic() - t0) * 1000)
|
||||
trace["submission"] = submission
|
||||
|
||||
if not submission or not isinstance(submission, dict):
|
||||
trace["error"] = "no submission"
|
||||
trace["output_findings"] = trace["input_findings"]
|
||||
return findings, trace
|
||||
|
||||
raw_groups = submission.get("groups") or []
|
||||
if not isinstance(raw_groups, list):
|
||||
trace["error"] = "submission.groups not a list"
|
||||
trace["output_findings"] = trace["input_findings"]
|
||||
return findings, trace
|
||||
|
||||
built = _build_deduped(raw_groups, findings)
|
||||
if built is None:
|
||||
trace["error"] = "invalid index coverage or schema"
|
||||
trace["output_findings"] = trace["input_findings"]
|
||||
log.warning(
|
||||
"dedupe: invalid output (%d originals, %d groups) — "
|
||||
"falling back to originals", len(findings), len(raw_groups),
|
||||
)
|
||||
return findings, trace
|
||||
|
||||
trace["output_findings"] = [f.model_dump(mode="json") for f in built]
|
||||
if on_progress:
|
||||
try:
|
||||
await on_progress(
|
||||
"cross_ic_dedupe", 0, "deduped",
|
||||
f"{len(findings)} → {len(built)} findings",
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
return built, trace
|
||||
@@ -0,0 +1,321 @@
|
||||
"""DigiKey API integration — fetch datasheet PDFs and product parameters by MPN.
|
||||
|
||||
Uses DigiKey Product Information API v4 with OAuth2 client credentials.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
import httpx
|
||||
|
||||
from backend.config import settings
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# OAuth2 token cache
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_token_cache: dict[str, str | float] = {"access_token": "", "expires_at": 0.0}
|
||||
|
||||
_BASE_URLS = {
|
||||
"production": "https://api.digikey.com",
|
||||
"sandbox": "https://sandbox-api.digikey.com",
|
||||
}
|
||||
|
||||
|
||||
async def _get_access_token() -> str:
|
||||
"""Get a DigiKey OAuth2 access token, refreshing if expired."""
|
||||
now = time.time()
|
||||
if _token_cache["access_token"] and float(_token_cache["expires_at"]) > now + 60:
|
||||
return str(_token_cache["access_token"])
|
||||
|
||||
base = _BASE_URLS.get(settings.digikey_environment, _BASE_URLS["production"])
|
||||
async with httpx.AsyncClient(timeout=15) as client:
|
||||
resp = await client.post(
|
||||
f"{base}/v1/oauth2/token",
|
||||
data={
|
||||
"client_id": settings.digikey_client_id,
|
||||
"client_secret": settings.digikey_client_secret,
|
||||
"grant_type": "client_credentials",
|
||||
},
|
||||
)
|
||||
resp.raise_for_status()
|
||||
data = resp.json()
|
||||
|
||||
_token_cache["access_token"] = data["access_token"]
|
||||
_token_cache["expires_at"] = now + data.get("expires_in", 3600)
|
||||
logger.info("DigiKey OAuth token refreshed (expires in %ds)", data.get("expires_in", 3600))
|
||||
return str(_token_cache["access_token"])
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Product search
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _get_mpn(product: dict) -> str:
|
||||
return product.get("ManufacturerProductNumber") or product.get("ManufacturerPartNumber") or ""
|
||||
|
||||
|
||||
def _get_ds_url(product: dict) -> str:
|
||||
url = product.get("DatasheetUrl") or product.get("PrimaryDatasheet") or ""
|
||||
# DigiKey sometimes returns protocol-relative URLs
|
||||
if url.startswith("//"):
|
||||
url = "https:" + url
|
||||
return url
|
||||
|
||||
|
||||
async def _keyword_search(mpn: str) -> list[dict]:
|
||||
"""Run a DigiKey keyword search and return the raw products list."""
|
||||
base = _BASE_URLS.get(settings.digikey_environment, _BASE_URLS["production"])
|
||||
token = await _get_access_token()
|
||||
|
||||
headers = {
|
||||
"Authorization": f"Bearer {token}",
|
||||
"X-DIGIKEY-Client-Id": settings.digikey_client_id,
|
||||
"X-DIGIKEY-Locale-Site": settings.digikey_locale_site,
|
||||
"X-DIGIKEY-Locale-Language": settings.digikey_locale_language,
|
||||
"X-DIGIKEY-Locale-Currency": settings.digikey_locale_currency,
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
body = {
|
||||
"Keywords": mpn,
|
||||
"Limit": 5,
|
||||
"Offset": 0,
|
||||
"ExcludeMarketPlaceProducts": True,
|
||||
}
|
||||
|
||||
async with httpx.AsyncClient(timeout=20) as client:
|
||||
resp = await client.post(
|
||||
f"{base}/products/v4/search/keyword",
|
||||
headers=headers,
|
||||
json=body,
|
||||
)
|
||||
resp.raise_for_status()
|
||||
data = resp.json()
|
||||
|
||||
return data.get("Products") or data.get("products") or []
|
||||
|
||||
|
||||
def _find_product(mpn: str, products: list[dict]) -> dict | None:
|
||||
"""Find the product whose MPN exactly matches ``mpn`` (case/space-insensitive).
|
||||
|
||||
Returns None when no result has a matching MPN. We intentionally do NOT
|
||||
fall back to ``products[0]`` — keyword-search hits without an MPN match
|
||||
are usually for a different part, and silently returning them has
|
||||
polluted the library with wrong specs for non-MPN tokens like ``10uF``.
|
||||
"""
|
||||
if not products:
|
||||
return None
|
||||
|
||||
mpn_upper = mpn.upper().replace(" ", "")
|
||||
for product in products:
|
||||
if _get_mpn(product).upper().replace(" ", "") == mpn_upper:
|
||||
return product
|
||||
return None
|
||||
|
||||
|
||||
async def _search_mpn(mpn: str) -> str | None:
|
||||
"""Search DigiKey for an MPN and return the primary datasheet URL, or None."""
|
||||
products = await _keyword_search(mpn)
|
||||
product = _find_product(mpn, products)
|
||||
if not product:
|
||||
return None
|
||||
url = _get_ds_url(product)
|
||||
return url or None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# PDF download + validation
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_PDF_MAGIC = b"%PDF-"
|
||||
_MIN_PDF_SIZE = 5_000 # 5 KB — anything smaller is probably an error page
|
||||
|
||||
|
||||
async def _download_pdf(url: str) -> bytes:
|
||||
"""Download a PDF from a URL and validate it.
|
||||
|
||||
Raises ValueError if the file isn't a valid PDF or is too small.
|
||||
Raises httpx.HTTPStatusError on 4xx/5xx responses.
|
||||
"""
|
||||
async with httpx.AsyncClient(timeout=15, follow_redirects=True) as client:
|
||||
resp = await client.get(url)
|
||||
resp.raise_for_status()
|
||||
data = resp.content
|
||||
|
||||
if not data.startswith(_PDF_MAGIC):
|
||||
raise ValueError("Downloaded file is not a valid PDF (bad magic bytes)")
|
||||
|
||||
if len(data) < _MIN_PDF_SIZE:
|
||||
raise ValueError(f"PDF too small ({len(data)} bytes) — likely an error page")
|
||||
|
||||
return data
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Public API
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class DatasheetFetchResult:
|
||||
"""Result of a datasheet fetch attempt."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
mpn: str,
|
||||
pdf_bytes: bytes | None = None,
|
||||
error: str | None = None,
|
||||
url: str | None = None,
|
||||
):
|
||||
self.mpn = mpn
|
||||
self.pdf_bytes = pdf_bytes
|
||||
self.error = error
|
||||
self.url = url # DigiKey datasheet URL (present even when PDF download fails)
|
||||
|
||||
@property
|
||||
def ok(self) -> bool:
|
||||
return self.pdf_bytes is not None
|
||||
|
||||
|
||||
async def fetch_datasheet(mpn: str) -> DatasheetFetchResult:
|
||||
"""Fetch a datasheet PDF for the given MPN from DigiKey.
|
||||
|
||||
Returns a DatasheetFetchResult with either pdf_bytes or an error message.
|
||||
The `url` field is set whenever DigiKey returns a datasheet link, even if
|
||||
the PDF download itself fails.
|
||||
Never raises — all errors are captured in the result.
|
||||
"""
|
||||
if not settings.use_digikey:
|
||||
return DatasheetFetchResult(mpn, error="DigiKey API not configured")
|
||||
|
||||
try:
|
||||
url = await _search_mpn(mpn)
|
||||
except httpx.HTTPStatusError as e:
|
||||
logger.warning("DigiKey search failed for %s: %s", mpn, e)
|
||||
return DatasheetFetchResult(mpn, error=f"DigiKey search failed ({e.response.status_code})")
|
||||
except Exception as e:
|
||||
msg = str(e) or type(e).__name__
|
||||
logger.warning("DigiKey search error for %s: %s", mpn, msg)
|
||||
return DatasheetFetchResult(mpn, error=f"DigiKey search error: {msg}")
|
||||
|
||||
if not url:
|
||||
return DatasheetFetchResult(mpn, error="No datasheet found on DigiKey")
|
||||
|
||||
try:
|
||||
pdf_bytes = await _download_pdf(url)
|
||||
except httpx.HTTPStatusError as e:
|
||||
logger.warning("Datasheet download blocked for %s (%s): %s", mpn, url, e)
|
||||
return DatasheetFetchResult(mpn, error=f"Download blocked ({e.response.status_code})", url=url)
|
||||
except ValueError as e:
|
||||
logger.warning("Invalid PDF for %s (%s): %s", mpn, url, e)
|
||||
return DatasheetFetchResult(mpn, error=str(e), url=url)
|
||||
except httpx.TimeoutException:
|
||||
logger.warning("Datasheet download timed out for %s (%s)", mpn, url)
|
||||
return DatasheetFetchResult(mpn, error="Download timed out", url=url)
|
||||
except Exception as e:
|
||||
msg = str(e) or type(e).__name__
|
||||
logger.warning("Datasheet download failed for %s (%s): %s", mpn, url, msg)
|
||||
return DatasheetFetchResult(mpn, error=f"Download failed: {msg}", url=url)
|
||||
|
||||
logger.info("Fetched datasheet for %s (%d KB)", mpn, len(pdf_bytes) // 1024)
|
||||
return DatasheetFetchResult(mpn, pdf_bytes=pdf_bytes, url=url)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Product parameters
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@dataclass
|
||||
class ProductParams:
|
||||
"""Structured product parameters from a DigiKey search result."""
|
||||
|
||||
mpn: str
|
||||
parameters: list[dict[str, str]] = field(default_factory=list) # [{"name": ..., "value": ...}]
|
||||
category: str = ""
|
||||
description: str = ""
|
||||
|
||||
|
||||
class ParamsFetchResult:
|
||||
"""Result of a product parameters fetch attempt."""
|
||||
|
||||
def __init__(self, mpn: str, params: ProductParams | None = None, error: str | None = None):
|
||||
self.mpn = mpn
|
||||
self.params = params
|
||||
self.error = error
|
||||
|
||||
@property
|
||||
def ok(self) -> bool:
|
||||
return self.params is not None
|
||||
|
||||
|
||||
def _parse_product_params(mpn: str, product: dict) -> ProductParams:
|
||||
"""Extract structured parameters from a DigiKey product dict."""
|
||||
raw_params = product.get("Parameters") or product.get("parameters") or []
|
||||
parameters = []
|
||||
for p in raw_params:
|
||||
name = p.get("ParameterText") or p.get("parameterText") or ""
|
||||
value = p.get("ValueText") or p.get("valueText") or ""
|
||||
if name and value and value != "-":
|
||||
parameters.append({"name": name, "value": value})
|
||||
|
||||
# Category
|
||||
cat = product.get("Category") or product.get("category") or {}
|
||||
category = cat.get("Name") or cat.get("name") or ""
|
||||
|
||||
# Description
|
||||
desc_obj = product.get("Description") or product.get("description") or {}
|
||||
if isinstance(desc_obj, str):
|
||||
description = desc_obj
|
||||
else:
|
||||
description = (
|
||||
desc_obj.get("ProductDescription")
|
||||
or desc_obj.get("productDescription")
|
||||
or desc_obj.get("DetailedDescription")
|
||||
or desc_obj.get("detailedDescription")
|
||||
or ""
|
||||
)
|
||||
|
||||
return ProductParams(
|
||||
mpn=mpn,
|
||||
parameters=parameters,
|
||||
category=category,
|
||||
description=description,
|
||||
)
|
||||
|
||||
|
||||
async def fetch_params(mpn: str) -> ParamsFetchResult:
|
||||
"""Fetch DigiKey product parameters for the given MPN.
|
||||
|
||||
Returns structured parameter data (no PDF download needed).
|
||||
Never raises — all errors are captured in the result.
|
||||
"""
|
||||
if not settings.use_digikey:
|
||||
return ParamsFetchResult(mpn, error="DigiKey API not configured")
|
||||
|
||||
try:
|
||||
products = await _keyword_search(mpn)
|
||||
except httpx.HTTPStatusError as e:
|
||||
logger.warning("DigiKey search failed for %s: %s", mpn, e)
|
||||
return ParamsFetchResult(mpn, error=f"DigiKey search failed ({e.response.status_code})")
|
||||
except Exception as e:
|
||||
msg = str(e) or type(e).__name__
|
||||
logger.warning("DigiKey search error for %s: %s", mpn, msg)
|
||||
return ParamsFetchResult(mpn, error=f"DigiKey search error: {msg}")
|
||||
|
||||
product = _find_product(mpn, products)
|
||||
if not product:
|
||||
return ParamsFetchResult(mpn, error="No results found on DigiKey")
|
||||
|
||||
params = _parse_product_params(mpn, product)
|
||||
if not params.parameters:
|
||||
return ParamsFetchResult(mpn, error="No parameters available on DigiKey")
|
||||
|
||||
logger.info("Fetched %d params for %s (category: %s)", len(params.parameters), mpn, params.category)
|
||||
return ParamsFetchResult(mpn, params=params)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,177 @@
|
||||
"""Cross-process event bridge for pipeline progress.
|
||||
|
||||
Today the FastAPI API process and the pipeline worker (Cloud Run Job
|
||||
execution, or a local subprocess in dev) live in different processes, so
|
||||
the in-memory ``EventBroker`` in ``services.pipeline`` can't span them.
|
||||
|
||||
The bridge:
|
||||
|
||||
* Worker writes one object per event to
|
||||
``users/{user_id}/projects/{project_id}/events/{seq:010d}.json``.
|
||||
The object holds ``{seq, ts, event, data}``. The worker is the only
|
||||
writer for a given run, so its local monotonic ``seq`` counter
|
||||
needs no coordination.
|
||||
|
||||
* API SSE handler tails the same prefix via ``StorageBackend.list_prefix_after``,
|
||||
yielding events in order until a terminal one arrives or the caller
|
||||
cancels.
|
||||
|
||||
This avoids appending to a single JSONL on GCS (no append API; full
|
||||
rewrite or compose-per-event has worse semantics) and naturally survives
|
||||
SSE reconnects (consumer just resumes from its last seen ``seq``).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
from datetime import datetime, timezone
|
||||
from typing import AsyncIterator
|
||||
|
||||
from backend.services.projects import project_prefix
|
||||
from backend.services.storage import StorageBackend
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# Filename pattern: 10-digit zero-padded seq + .json. Lexicographic order
|
||||
# matches numeric order so list_prefix_after pages cleanly.
|
||||
_SEQ_WIDTH = 10
|
||||
_FILENAME_FMT = f"{{seq:0{_SEQ_WIDTH}d}}.json"
|
||||
|
||||
# Terminal event names — the SSE loop stops on these.
|
||||
TERMINAL_EVENTS = frozenset({
|
||||
"pipeline_complete",
|
||||
"pipeline_error",
|
||||
"pipeline_cancelled",
|
||||
"pipeline_paused",
|
||||
})
|
||||
|
||||
|
||||
def _events_prefix(user_id: str, project_id: str) -> str:
|
||||
return f"{project_prefix(user_id, project_id)}/events/"
|
||||
|
||||
|
||||
def _seq_from_key(key: str) -> int | None:
|
||||
"""Extract the integer seq from an event key; ``None`` on parse failure."""
|
||||
name = key.rsplit("/", 1)[-1]
|
||||
if not name.endswith(".json"):
|
||||
return None
|
||||
stem = name[:-5]
|
||||
try:
|
||||
return int(stem)
|
||||
except ValueError:
|
||||
return None
|
||||
|
||||
|
||||
class GCSEventBroker:
|
||||
"""Drop-in for the in-memory ``EventBroker`` that persists to storage.
|
||||
|
||||
Same interface (``publish``, ``subscribe``, ``unsubscribe``,
|
||||
``clear_history``) so the worker can swap it in for the module-level
|
||||
``broker`` singleton without touching call sites. Subscription is a
|
||||
no-op — the API consumes events via :func:`tail_events` instead.
|
||||
"""
|
||||
|
||||
def __init__(self, storage: StorageBackend, user_id: str) -> None:
|
||||
self.storage = storage
|
||||
self.user_id = user_id
|
||||
# Per-project local counter. Workers handle one project per
|
||||
# execution, but the dict shape keeps parity with ``EventBroker``.
|
||||
self._seq: dict[str, int] = {}
|
||||
|
||||
def subscribe(self, project_id: str) -> asyncio.Queue:
|
||||
# Workers never subscribe — only the API tails the GCS event log.
|
||||
# Returning an unfed queue is acceptable but raising is more
|
||||
# honest about the contract.
|
||||
raise NotImplementedError(
|
||||
"GCSEventBroker is publish-only; subscribers should call "
|
||||
"event_bridge.tail_events(...) instead."
|
||||
)
|
||||
|
||||
def unsubscribe(self, project_id: str, q: asyncio.Queue) -> None:
|
||||
# No-op for symmetry with the in-memory broker.
|
||||
return
|
||||
|
||||
def clear_history(self, project_id: str) -> None:
|
||||
"""Wipe all prior event objects for this project.
|
||||
|
||||
Called at the start of a fresh run so resumed/restarted runs
|
||||
don't intermix with stale events from earlier attempts.
|
||||
"""
|
||||
prefix = _events_prefix(self.user_id, project_id)
|
||||
try:
|
||||
self.storage.delete_prefix(prefix)
|
||||
except Exception:
|
||||
logger.exception("failed to clear event history at %s", prefix)
|
||||
self._seq[project_id] = 0
|
||||
|
||||
def publish(self, project_id: str, event: str, data: dict) -> None:
|
||||
seq = self._seq.get(project_id, 0)
|
||||
self._seq[project_id] = seq + 1
|
||||
key = _events_prefix(self.user_id, project_id) + _FILENAME_FMT.format(seq=seq)
|
||||
msg = {
|
||||
"seq": seq,
|
||||
"ts": datetime.now(timezone.utc).isoformat(),
|
||||
"event": event,
|
||||
"data": data,
|
||||
}
|
||||
try:
|
||||
self.storage.write_json(key, msg)
|
||||
except Exception:
|
||||
# An event-write failure should never crash the pipeline.
|
||||
logger.exception("failed to write event %s to %s", event, key)
|
||||
|
||||
|
||||
async def tail_events(
|
||||
storage: StorageBackend,
|
||||
user_id: str,
|
||||
project_id: str,
|
||||
*,
|
||||
poll_interval: float = 0.5,
|
||||
heartbeat_interval: float = 15.0,
|
||||
) -> AsyncIterator[dict]:
|
||||
"""Yield events from the GCS-backed event log in order.
|
||||
|
||||
Stops yielding after a terminal event (``pipeline_complete``,
|
||||
``pipeline_error``, ``pipeline_cancelled``). Emits a
|
||||
``{"event": "heartbeat", "data": {}}`` synthetic event roughly every
|
||||
``heartbeat_interval`` seconds when no real events arrive, matching
|
||||
the behaviour of the in-memory broker's SSE loop.
|
||||
|
||||
The caller is expected to handle disconnects/cancellations and
|
||||
secondary terminal-detection (``meta.status``, Cloud Run execution
|
||||
state) on top of this iterator.
|
||||
"""
|
||||
prefix = _events_prefix(user_id, project_id)
|
||||
last_seen_key: str | None = None
|
||||
last_emit_ts = 0.0
|
||||
|
||||
while True:
|
||||
try:
|
||||
keys = storage.list_prefix_after(prefix, after_key=last_seen_key)
|
||||
except Exception:
|
||||
logger.exception("event tail: list_prefix_after failed for %s", prefix)
|
||||
keys = []
|
||||
|
||||
emitted_any = False
|
||||
for key in keys:
|
||||
try:
|
||||
msg = storage.read_json(key)
|
||||
except Exception:
|
||||
logger.exception("event tail: read_json failed for %s", key)
|
||||
continue
|
||||
yield msg
|
||||
emitted_any = True
|
||||
last_seen_key = key
|
||||
last_emit_ts = asyncio.get_event_loop().time()
|
||||
if msg.get("event") in TERMINAL_EVENTS:
|
||||
return
|
||||
|
||||
now = asyncio.get_event_loop().time()
|
||||
if not emitted_any and now - last_emit_ts >= heartbeat_interval:
|
||||
yield {"event": "heartbeat", "data": {}}
|
||||
last_emit_ts = now
|
||||
|
||||
await asyncio.sleep(poll_interval)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,327 @@
|
||||
"""Pipeline-worker dispatcher.
|
||||
|
||||
In production: enqueues a Cloud Run Job execution that runs the
|
||||
``backend.pipeline_worker`` entrypoint with project_id/user_id/resume/free
|
||||
passed as env-var overrides.
|
||||
|
||||
In local dev (no ``GCS_BUCKET``): launches the worker as a child process
|
||||
so the same code path runs end-to-end. Removes the in-process
|
||||
``BackgroundTask`` divergence between dev and prod.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import threading
|
||||
from typing import Literal
|
||||
|
||||
from backend.config import settings
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
ExecutionState = Literal[
|
||||
"pending", "running", "succeeded", "failed", "cancelled", "unknown"
|
||||
]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Local subprocess fallback (dev mode)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
# Track child processes so the API can query "is it still running?" in
|
||||
# dev. In prod the Cloud Run Jobs admin API answers the same question.
|
||||
_local_procs: dict[str, subprocess.Popen] = {}
|
||||
_local_procs_lock = threading.Lock()
|
||||
|
||||
|
||||
def _local_execution_name(project_id: str) -> str:
|
||||
"""Stable synthetic execution name for the dev subprocess path.
|
||||
|
||||
Lets the rest of the codebase treat dev runs uniformly with prod
|
||||
runs (we always have an ``execution_name`` to store on ProjectMeta
|
||||
and pass to status / cancel calls).
|
||||
"""
|
||||
return f"local/projects/{project_id}"
|
||||
|
||||
|
||||
def _spawn_local_subprocess(
|
||||
project_id: str,
|
||||
user_id: str,
|
||||
*,
|
||||
resume: bool,
|
||||
free: bool,
|
||||
mode: str = "run",
|
||||
regen_stages: list[str] | None = None,
|
||||
) -> str:
|
||||
name = _local_execution_name(project_id)
|
||||
env = os.environ.copy()
|
||||
env["PROJECT_ID"] = project_id
|
||||
env["USER_ID"] = user_id
|
||||
env["RESUME"] = "1" if resume else "0"
|
||||
env["FREE"] = "1" if free else "0"
|
||||
env["MODE"] = mode
|
||||
if regen_stages:
|
||||
env["REGEN_STAGES"] = ",".join(regen_stages)
|
||||
env["EXECUTION_NAME"] = name
|
||||
proc = subprocess.Popen(
|
||||
[sys.executable, "-m", "backend.pipeline_worker"],
|
||||
env=env,
|
||||
# Inherit stdout/stderr so logs appear in the dev terminal
|
||||
stdin=subprocess.DEVNULL,
|
||||
)
|
||||
with _local_procs_lock:
|
||||
# Reap any old proc for the same project before tracking the new one.
|
||||
prior = _local_procs.pop(project_id, None)
|
||||
if prior is not None:
|
||||
try:
|
||||
prior.terminate()
|
||||
except Exception:
|
||||
pass
|
||||
_local_procs[project_id] = proc
|
||||
logger.info("dev: spawned worker subprocess pid=%s for %s", proc.pid, project_id)
|
||||
return name
|
||||
|
||||
|
||||
def _local_state(project_id: str) -> ExecutionState:
|
||||
with _local_procs_lock:
|
||||
proc = _local_procs.get(project_id)
|
||||
if proc is None:
|
||||
return "unknown"
|
||||
rc = proc.poll()
|
||||
if rc is None:
|
||||
return "running"
|
||||
if rc == 0:
|
||||
return "succeeded"
|
||||
if rc < 0:
|
||||
# Negative return = terminated by signal
|
||||
return "cancelled"
|
||||
return "failed"
|
||||
|
||||
|
||||
def _local_cancel(project_id: str) -> None:
|
||||
with _local_procs_lock:
|
||||
proc = _local_procs.get(project_id)
|
||||
if proc is None or proc.poll() is not None:
|
||||
return
|
||||
try:
|
||||
proc.terminate()
|
||||
except Exception:
|
||||
logger.exception("dev: failed to terminate worker subprocess for %s", project_id)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Cloud Run Jobs (prod path)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _gcp_project() -> str:
|
||||
"""Resolve the GCP project id for the Cloud Run Jobs admin API."""
|
||||
if settings.pipeline_worker_project:
|
||||
return settings.pipeline_worker_project
|
||||
proj = os.environ.get("GOOGLE_CLOUD_PROJECT") or os.environ.get("GCLOUD_PROJECT")
|
||||
if proj:
|
||||
return proj
|
||||
# Fall back to the metadata server (works on Cloud Run).
|
||||
try:
|
||||
import requests # type: ignore[import-not-found]
|
||||
|
||||
resp = requests.get(
|
||||
"http://metadata.google.internal/computeMetadata/v1/project/project-id",
|
||||
headers={"Metadata-Flavor": "Google"},
|
||||
timeout=2.0,
|
||||
)
|
||||
if resp.ok:
|
||||
return resp.text.strip()
|
||||
except Exception:
|
||||
pass
|
||||
raise RuntimeError(
|
||||
"Could not resolve GCP project for Cloud Run Jobs. Set "
|
||||
"PIPELINE_WORKER_PROJECT or GOOGLE_CLOUD_PROJECT."
|
||||
)
|
||||
|
||||
|
||||
def _job_resource_name() -> str:
|
||||
return (
|
||||
f"projects/{_gcp_project()}/locations/{settings.pipeline_worker_region}"
|
||||
f"/jobs/{settings.pipeline_worker_job_name}"
|
||||
)
|
||||
|
||||
|
||||
def _jobs_client():
|
||||
# Lazy import: keeps the API process startup fast in local dev where
|
||||
# google-cloud-run isn't even installed (it's an optional dep there).
|
||||
from google.cloud import run_v2 # type: ignore[import-not-found]
|
||||
|
||||
return run_v2.JobsClient()
|
||||
|
||||
|
||||
def _executions_client():
|
||||
from google.cloud import run_v2 # type: ignore[import-not-found]
|
||||
|
||||
return run_v2.ExecutionsClient()
|
||||
|
||||
|
||||
def _enqueue_cloud_run_job(
|
||||
project_id: str,
|
||||
user_id: str,
|
||||
*,
|
||||
resume: bool,
|
||||
free: bool,
|
||||
mode: str = "run",
|
||||
regen_stages: list[str] | None = None,
|
||||
) -> str:
|
||||
"""Issue ``RunJob`` with env-var overrides; return the execution name."""
|
||||
from google.cloud import run_v2 # type: ignore[import-not-found]
|
||||
|
||||
env_overrides = [
|
||||
run_v2.EnvVar(name="PROJECT_ID", value=project_id),
|
||||
run_v2.EnvVar(name="USER_ID", value=user_id),
|
||||
run_v2.EnvVar(name="RESUME", value="1" if resume else "0"),
|
||||
run_v2.EnvVar(name="FREE", value="1" if free else "0"),
|
||||
run_v2.EnvVar(name="MODE", value=mode),
|
||||
]
|
||||
if regen_stages:
|
||||
env_overrides.append(
|
||||
run_v2.EnvVar(name="REGEN_STAGES", value=",".join(regen_stages)),
|
||||
)
|
||||
overrides = run_v2.RunJobRequest.Overrides(
|
||||
container_overrides=[
|
||||
run_v2.RunJobRequest.Overrides.ContainerOverride(env=env_overrides),
|
||||
],
|
||||
)
|
||||
request = run_v2.RunJobRequest(name=_job_resource_name(), overrides=overrides)
|
||||
operation = _jobs_client().run_job(request=request)
|
||||
# Don't wait for completion — fire and forget. The metadata is enough
|
||||
# to extract the execution resource name.
|
||||
metadata = operation.metadata
|
||||
name = getattr(metadata, "name", None) if metadata is not None else None
|
||||
if not name:
|
||||
# As a fallback, peek at the operation; on Cloud Run RunJob this
|
||||
# is a long-running op whose initial metadata holds the execution.
|
||||
name = operation.operation.name # type: ignore[union-attr]
|
||||
if not name:
|
||||
raise RuntimeError("Cloud Run RunJob returned no execution name")
|
||||
logger.info("enqueued Cloud Run Job execution %s for project %s", name, project_id)
|
||||
return name
|
||||
|
||||
|
||||
def _cloud_run_state(execution_name: str) -> ExecutionState:
|
||||
"""Map Cloud Run Execution state to our enum."""
|
||||
try:
|
||||
from google.cloud import run_v2 # type: ignore[import-not-found]
|
||||
|
||||
client = _executions_client()
|
||||
ex = client.get_execution(name=execution_name)
|
||||
except Exception:
|
||||
logger.exception("get_execution failed for %s", execution_name)
|
||||
return "unknown"
|
||||
|
||||
# An Execution has reconciliation_started, completion_time, conditions.
|
||||
# Map to our enum based on completion + conditions.
|
||||
if ex.completion_time is None or ex.completion_time.seconds == 0:
|
||||
if ex.start_time and ex.start_time.seconds:
|
||||
return "running"
|
||||
return "pending"
|
||||
# Completed — figure out success vs failure.
|
||||
failed = int(getattr(ex, "failed_count", 0) or 0)
|
||||
cancelled = int(getattr(ex, "cancelled_count", 0) or 0)
|
||||
succeeded = int(getattr(ex, "succeeded_count", 0) or 0)
|
||||
if cancelled > 0 and succeeded == 0:
|
||||
return "cancelled"
|
||||
if failed > 0:
|
||||
return "failed"
|
||||
if succeeded > 0:
|
||||
return "succeeded"
|
||||
return "unknown"
|
||||
|
||||
|
||||
def _cloud_run_cancel(execution_name: str) -> None:
|
||||
try:
|
||||
from google.cloud import run_v2 # type: ignore[import-not-found]
|
||||
|
||||
request = run_v2.CancelExecutionRequest(name=execution_name)
|
||||
_executions_client().cancel_execution(request=request)
|
||||
except Exception:
|
||||
logger.exception("cancel_execution failed for %s", execution_name)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Public API
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def use_cloud_run_jobs() -> bool:
|
||||
"""True iff we should dispatch via Cloud Run Jobs.
|
||||
|
||||
Tied to whether GCS storage is configured — Jobs and GCS go together
|
||||
in prod, and local dev uses neither.
|
||||
"""
|
||||
return bool(settings.gcs_bucket)
|
||||
|
||||
|
||||
def enqueue_pipeline(
|
||||
project_id: str, user_id: str, *, resume: bool = False, free: bool = False,
|
||||
) -> str:
|
||||
"""Dispatch a pipeline run.
|
||||
|
||||
In prod, returns the Cloud Run Execution resource name. In dev,
|
||||
returns a synthetic ``local/projects/{id}`` name. Either way, callers
|
||||
should persist the returned name on ``ProjectMeta.execution_name``.
|
||||
"""
|
||||
if use_cloud_run_jobs():
|
||||
return _enqueue_cloud_run_job(project_id, user_id, resume=resume, free=free)
|
||||
return _spawn_local_subprocess(project_id, user_id, resume=resume, free=free)
|
||||
|
||||
|
||||
def enqueue_pipeline_regen(
|
||||
project_id: str, user_id: str, *, stages: list[str],
|
||||
) -> str:
|
||||
"""Dispatch a regen run (graph + selected stages, free).
|
||||
|
||||
Same image, same worker; differs only in the env-var-driven mode.
|
||||
"""
|
||||
if not stages:
|
||||
raise ValueError("regen requires at least one stage")
|
||||
if use_cloud_run_jobs():
|
||||
return _enqueue_cloud_run_job(
|
||||
project_id, user_id, resume=False, free=True,
|
||||
mode="regen", regen_stages=stages,
|
||||
)
|
||||
return _spawn_local_subprocess(
|
||||
project_id, user_id, resume=False, free=True,
|
||||
mode="regen", regen_stages=stages,
|
||||
)
|
||||
|
||||
|
||||
def get_execution_state(execution_name: str | None) -> ExecutionState:
|
||||
"""Return current state of a previously-enqueued execution.
|
||||
|
||||
Used by the SSE handler's hard-crash escape hatch and by the
|
||||
stale-running sweeper. ``None`` -> ``"unknown"``.
|
||||
"""
|
||||
if not execution_name:
|
||||
return "unknown"
|
||||
if execution_name.startswith("local/projects/"):
|
||||
project_id = execution_name.split("/", 2)[-1]
|
||||
return _local_state(project_id)
|
||||
return _cloud_run_state(execution_name)
|
||||
|
||||
|
||||
def cancel_execution(execution_name: str | None) -> None:
|
||||
"""Hard-cancel an execution (Cloud Run cancel or local SIGTERM).
|
||||
|
||||
Best-effort. Soft cancel via ``meta.cancel_requested`` is preferred —
|
||||
only fall back to this when the worker has already gone unresponsive.
|
||||
"""
|
||||
if not execution_name:
|
||||
return
|
||||
if execution_name.startswith("local/projects/"):
|
||||
project_id = execution_name.split("/", 2)[-1]
|
||||
_local_cancel(project_id)
|
||||
return
|
||||
_cloud_run_cancel(execution_name)
|
||||
@@ -0,0 +1,36 @@
|
||||
"""Provider-agnostic LLM client layer.
|
||||
|
||||
All Claude API calls in the backend route through this package via the
|
||||
``LLMProvider`` interface. The default provider is Anthropic; per-stage
|
||||
overrides via ``Settings.provider_*`` env vars route specific stages to
|
||||
other providers (currently Anthropic + Gemini).
|
||||
"""
|
||||
|
||||
from backend.services.llm.factory import call_with_fallback, get_provider
|
||||
from backend.services.llm.types import (
|
||||
Completion,
|
||||
ContentBlock,
|
||||
Message,
|
||||
PdfBlock,
|
||||
TextBlock,
|
||||
ToolCall,
|
||||
ToolChoice,
|
||||
ToolResultBlock,
|
||||
ToolSchema,
|
||||
Usage,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"Completion",
|
||||
"ContentBlock",
|
||||
"Message",
|
||||
"PdfBlock",
|
||||
"TextBlock",
|
||||
"ToolCall",
|
||||
"ToolChoice",
|
||||
"ToolResultBlock",
|
||||
"ToolSchema",
|
||||
"Usage",
|
||||
"call_with_fallback",
|
||||
"get_provider",
|
||||
]
|
||||
@@ -0,0 +1,363 @@
|
||||
"""Anthropic provider — wraps AsyncAnthropic + Console Skills.
|
||||
|
||||
Translates the unified ``Message`` / ``Completion`` shapes into Anthropic's
|
||||
native message-block format and back. Caching is per-block via
|
||||
``cache_control: ephemeral``."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import re
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import anthropic
|
||||
|
||||
from backend.config import settings
|
||||
from backend.services.llm.base import LLMProvider, LLMSession
|
||||
from backend.services.llm.types import (
|
||||
Completion,
|
||||
ContentBlock,
|
||||
Message,
|
||||
PdfBlock,
|
||||
TextBlock,
|
||||
ToolCall,
|
||||
ToolChoice,
|
||||
ToolResultBlock,
|
||||
ToolSchema,
|
||||
Usage,
|
||||
)
|
||||
|
||||
|
||||
_SKILL_MAX_TURNS = 10
|
||||
|
||||
# Sampling params were removed on newer Claude models (Sonnet 5, Opus 4.7+,
|
||||
# Fable/Mythos 5) — sending `temperature` returns 400 "`temperature` is
|
||||
# deprecated for this model". Allowlist the families that still accept it so
|
||||
# unknown/future models fail safe (omit → default sampling) instead of
|
||||
# 400-ing every call in the session.
|
||||
_TEMPERATURE_OK = re.compile(r"^claude-(3-|opus-4-[0-6]|sonnet-4-|haiku-)")
|
||||
|
||||
|
||||
def _model_accepts_temperature(model: str) -> bool:
|
||||
return bool(_TEMPERATURE_OK.match(model))
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Translation helpers — unified types ↔ Anthropic dicts
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _encode_pdf_block(path: Path | str, *, cache: bool) -> dict:
|
||||
data = base64.standard_b64encode(Path(path).read_bytes()).decode()
|
||||
block: dict = {
|
||||
"type": "document",
|
||||
"source": {"type": "base64", "media_type": "application/pdf", "data": data},
|
||||
}
|
||||
if cache:
|
||||
block["cache_control"] = {"type": "ephemeral"}
|
||||
return block
|
||||
|
||||
|
||||
def _to_anthropic_block(b: ContentBlock) -> dict:
|
||||
if isinstance(b, TextBlock):
|
||||
d: dict = {"type": "text", "text": b.text}
|
||||
if b.cacheable:
|
||||
d["cache_control"] = {"type": "ephemeral"}
|
||||
return d
|
||||
if isinstance(b, PdfBlock):
|
||||
return _encode_pdf_block(b.path, cache=b.cacheable)
|
||||
if isinstance(b, ToolCall):
|
||||
return {"type": "tool_use", "id": b.id, "name": b.name, "input": b.input}
|
||||
if isinstance(b, ToolResultBlock):
|
||||
return {
|
||||
"type": "tool_result",
|
||||
"tool_use_id": b.tool_use_id,
|
||||
"content": b.content,
|
||||
}
|
||||
raise TypeError(f"Unknown ContentBlock: {type(b).__name__}")
|
||||
|
||||
|
||||
def _to_anthropic_message(m: Message) -> dict:
|
||||
return {"role": m.role, "content": [_to_anthropic_block(b) for b in m.content]}
|
||||
|
||||
|
||||
# Anthropic allows at most 4 cache_control breakpoints per request. The system
|
||||
# prompt always consumes one (see AnthropicSession.complete), leaving 3 for
|
||||
# message content. A multi-turn review attaches a cacheable PDF for each
|
||||
# get_datasheet_excerpt fetch (validation_tools.py), so a hub IC that verifies
|
||||
# two interface excerpts produced 5 breakpoints — system + initial PDF + initial
|
||||
# context + 2 excerpts — and the API rejected the request with
|
||||
# "A maximum of 4 blocks with cache_control may be provided. Found 5."
|
||||
#
|
||||
# Cap the message-block breakpoints in the translated request, keeping the most
|
||||
# valuable ones: the first cacheable block (the full-datasheet anchor — a stable,
|
||||
# guaranteed cache hit every turn) plus the two most recent (incremental caching
|
||||
# of the growing tail). Any caller-set cache_control beyond that is dropped.
|
||||
_MAX_MESSAGE_CACHE_BREAKPOINTS = 3
|
||||
|
||||
|
||||
def _enforce_cache_breakpoint_limit(messages: list[dict]) -> None:
|
||||
"""Strip excess cache_control markers from message blocks in place so that
|
||||
system(1) + message breakpoints never exceed Anthropic's per-request limit."""
|
||||
marked: list[dict] = []
|
||||
for m in messages:
|
||||
content = m.get("content")
|
||||
if not isinstance(content, list):
|
||||
continue
|
||||
for block in content:
|
||||
if isinstance(block, dict) and "cache_control" in block:
|
||||
marked.append(block)
|
||||
if len(marked) <= _MAX_MESSAGE_CACHE_BREAKPOINTS:
|
||||
return
|
||||
keep = {id(marked[0]), id(marked[-1]), id(marked[-2])}
|
||||
for block in marked:
|
||||
if id(block) not in keep:
|
||||
block.pop("cache_control", None)
|
||||
|
||||
|
||||
def _to_anthropic_tool(t: ToolSchema) -> dict:
|
||||
return {"name": t.name, "description": t.description, "input_schema": t.input_schema}
|
||||
|
||||
|
||||
def _to_anthropic_tool_choice(c: ToolChoice) -> dict:
|
||||
if c == "auto":
|
||||
return {"type": "auto"}
|
||||
if c == "none":
|
||||
return {"type": "none"}
|
||||
if isinstance(c, dict) and "name" in c:
|
||||
return {"type": "tool", "name": c["name"]}
|
||||
raise ValueError(f"Invalid tool_choice: {c!r}")
|
||||
|
||||
|
||||
def _from_anthropic_response(resp) -> Completion:
|
||||
"""Parse an Anthropic message response into a unified Completion."""
|
||||
text_parts: list[str] = []
|
||||
tool_calls: list[ToolCall] = []
|
||||
raw_blocks: list[ContentBlock] = []
|
||||
|
||||
for block in resp.content:
|
||||
btype = getattr(block, "type", None)
|
||||
if btype == "text":
|
||||
text_parts.append(block.text)
|
||||
raw_blocks.append(TextBlock(text=block.text))
|
||||
elif btype == "tool_use":
|
||||
tc = ToolCall(id=block.id, name=block.name, input=dict(block.input))
|
||||
tool_calls.append(tc)
|
||||
raw_blocks.append(tc)
|
||||
# Other block types (server tool calls etc.) are pass-through ignored
|
||||
|
||||
usage = Usage(
|
||||
input_tokens=resp.usage.input_tokens,
|
||||
output_tokens=resp.usage.output_tokens,
|
||||
cache_creation_tokens=getattr(resp.usage, "cache_creation_input_tokens", 0) or 0,
|
||||
cache_read_tokens=getattr(resp.usage, "cache_read_input_tokens", 0) or 0,
|
||||
)
|
||||
|
||||
return Completion(
|
||||
text="".join(text_parts),
|
||||
tool_calls=tool_calls,
|
||||
usage=usage,
|
||||
stop_reason=resp.stop_reason or "unknown",
|
||||
raw_assistant_blocks=raw_blocks,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Session
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class AnthropicSession(LLMSession):
|
||||
provider_name = "anthropic"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
client: anthropic.AsyncAnthropic,
|
||||
model: str,
|
||||
system: str,
|
||||
max_tokens: int,
|
||||
temperature: float | None = None,
|
||||
) -> None:
|
||||
self._client = client
|
||||
self.model = model
|
||||
self._system = system
|
||||
self._max_tokens = max_tokens
|
||||
self._temperature = temperature
|
||||
|
||||
async def complete(
|
||||
self,
|
||||
*,
|
||||
messages: list[Message],
|
||||
tools: list[ToolSchema] | None = None,
|
||||
tool_choice: ToolChoice = "auto",
|
||||
) -> Completion:
|
||||
kwargs: dict = {
|
||||
"model": self.model,
|
||||
"max_tokens": self._max_tokens,
|
||||
"system": [{
|
||||
"type": "text",
|
||||
"text": self._system,
|
||||
"cache_control": {"type": "ephemeral"},
|
||||
}],
|
||||
"messages": [_to_anthropic_message(m) for m in messages],
|
||||
}
|
||||
_enforce_cache_breakpoint_limit(kwargs["messages"])
|
||||
if self._temperature is not None and _model_accepts_temperature(self.model):
|
||||
kwargs["temperature"] = self._temperature
|
||||
if tools:
|
||||
kwargs["tools"] = [_to_anthropic_tool(t) for t in tools]
|
||||
kwargs["tool_choice"] = _to_anthropic_tool_choice(tool_choice)
|
||||
|
||||
# Streaming, not create(): SDK 0.83+ raises ValueError pre-flight on
|
||||
# `messages.create` whenever max_tokens crosses ~21k for Sonnet
|
||||
# (the "may take longer than 10 minutes" guard). Review uses 32k
|
||||
# max_tokens for Gemini thinking headroom; streaming bypasses that
|
||||
# client-side timeout cap. get_final_message() returns the same
|
||||
# shape as create(), so _from_anthropic_response is reused as-is.
|
||||
async with self._client.messages.stream(**kwargs) as stream:
|
||||
resp = await stream.get_final_message()
|
||||
return _from_anthropic_response(resp)
|
||||
|
||||
async def close(self) -> None:
|
||||
# Anthropic ephemeral cache cleans up on its own (5-min TTL).
|
||||
pass
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Provider
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class AnthropicProvider(LLMProvider):
|
||||
name = "anthropic"
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._client = anthropic.AsyncAnthropic(api_key=settings.anthropic_api_key)
|
||||
|
||||
async def create_session(
|
||||
self,
|
||||
*,
|
||||
model: str,
|
||||
system: str,
|
||||
max_tokens: int = 4096,
|
||||
temperature: float | None = None,
|
||||
) -> LLMSession:
|
||||
return AnthropicSession(
|
||||
client=self._client,
|
||||
model=model,
|
||||
system=system,
|
||||
max_tokens=max_tokens,
|
||||
temperature=temperature,
|
||||
)
|
||||
|
||||
async def run_skill(
|
||||
self,
|
||||
*,
|
||||
skill_name: str,
|
||||
model: str,
|
||||
system: str,
|
||||
user_text: str,
|
||||
pdf_path: str | None,
|
||||
output_tool: ToolSchema,
|
||||
) -> tuple[dict, Completion]:
|
||||
"""Anthropic Console Skills — multi-turn skill execution with the
|
||||
``skills-2025-10-02`` + ``code-execution-2025-08-25`` betas.
|
||||
|
||||
Skill mounts in a per-call container; the model reads ``SKILL.md``,
|
||||
runs ``validate.py`` server-side via code_execution, and voluntarily
|
||||
calls ``output_tool`` once it has well-formed data.
|
||||
"""
|
||||
skill_id, version = settings.get_skill(skill_name)
|
||||
|
||||
# Build initial user content
|
||||
user_content: list[dict] = []
|
||||
if pdf_path:
|
||||
user_content.append(_encode_pdf_block(pdf_path, cache=True))
|
||||
user_content.append({"type": "text", "text": user_text})
|
||||
|
||||
messages: list[dict] = [{"role": "user", "content": user_content}]
|
||||
container: dict = {
|
||||
"skills": [{
|
||||
"type": "custom",
|
||||
"skill_id": skill_id,
|
||||
"version": version,
|
||||
}],
|
||||
}
|
||||
|
||||
total_input = 0
|
||||
total_output = 0
|
||||
total_cache_creation = 0
|
||||
total_cache_read = 0
|
||||
t0 = time.monotonic()
|
||||
last_resp = None
|
||||
|
||||
for turn in range(_SKILL_MAX_TURNS):
|
||||
resp = await self._client.beta.messages.create(
|
||||
model=model,
|
||||
max_tokens=16384,
|
||||
system=[{
|
||||
"type": "text",
|
||||
"text": system,
|
||||
"cache_control": {"type": "ephemeral"},
|
||||
}],
|
||||
tools=[
|
||||
{"type": "code_execution_20250825", "name": "code_execution"},
|
||||
_to_anthropic_tool(output_tool),
|
||||
],
|
||||
container=container,
|
||||
messages=messages,
|
||||
betas=["skills-2025-10-02", "code-execution-2025-08-25"],
|
||||
)
|
||||
last_resp = resp
|
||||
|
||||
total_input += resp.usage.input_tokens
|
||||
total_output += resp.usage.output_tokens
|
||||
total_cache_creation += getattr(resp.usage, "cache_creation_input_tokens", 0) or 0
|
||||
total_cache_read += getattr(resp.usage, "cache_read_input_tokens", 0) or 0
|
||||
|
||||
# Reuse container for subsequent turns
|
||||
if hasattr(resp, "container") and resp.container:
|
||||
container = {"id": resp.container.id}
|
||||
|
||||
for block in resp.content:
|
||||
if (
|
||||
getattr(block, "type", None) == "tool_use"
|
||||
and block.name == output_tool.name
|
||||
):
|
||||
completion = Completion(
|
||||
text="",
|
||||
tool_calls=[ToolCall(id=block.id, name=block.name, input=dict(block.input))],
|
||||
usage=Usage(
|
||||
input_tokens=total_input,
|
||||
output_tokens=total_output,
|
||||
cache_creation_tokens=total_cache_creation,
|
||||
cache_read_tokens=total_cache_read,
|
||||
),
|
||||
stop_reason=resp.stop_reason or "unknown",
|
||||
)
|
||||
# Stash turns count via attribute for callers that need it
|
||||
completion.turns = turn + 1 # type: ignore[attr-defined]
|
||||
completion.duration_ms = int((time.monotonic() - t0) * 1000) # type: ignore[attr-defined]
|
||||
return dict(block.input), completion
|
||||
|
||||
messages.append({"role": "assistant", "content": resp.content})
|
||||
|
||||
if resp.stop_reason == "pause_turn":
|
||||
continue
|
||||
|
||||
if resp.stop_reason == "end_turn":
|
||||
messages.append({
|
||||
"role": "user",
|
||||
"content": f"Please call {output_tool.name} with the extracted data.",
|
||||
})
|
||||
continue
|
||||
|
||||
# tool_use from code_execution — let the loop continue
|
||||
continue
|
||||
|
||||
raise RuntimeError(
|
||||
f"Skill {skill_name!r} did not produce {output_tool.name} "
|
||||
f"in {_SKILL_MAX_TURNS} turns"
|
||||
)
|
||||
@@ -0,0 +1,100 @@
|
||||
"""Abstract LLMProvider + LLMSession interfaces."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Protocol
|
||||
|
||||
from backend.services.llm.types import (
|
||||
Completion,
|
||||
Message,
|
||||
ToolChoice,
|
||||
ToolSchema,
|
||||
)
|
||||
|
||||
|
||||
class LLMSession(ABC):
|
||||
"""A multi-turn conversation with provider-specific cache lifecycle.
|
||||
|
||||
Lifecycle::
|
||||
|
||||
session = await provider.create_session(model=..., system=...)
|
||||
try:
|
||||
messages = [Message("user", [
|
||||
PdfBlock(path, cacheable=True),
|
||||
TextBlock(context, cacheable=True),
|
||||
])]
|
||||
for turn in range(N):
|
||||
completion = await session.complete(
|
||||
messages=messages, tools=..., tool_choice=...,
|
||||
)
|
||||
# process tool_calls, append to messages, repeat
|
||||
finally:
|
||||
await session.close()
|
||||
|
||||
Caching: blocks with ``cacheable=True`` participate in provider caching.
|
||||
Anthropic stamps ``cache_control: ephemeral`` on each cacheable block on
|
||||
every call. Gemini collects all cacheable blocks (plus the system prompt)
|
||||
on the first ``complete()`` call into a ``CachedContent`` object and
|
||||
references it on subsequent calls. The system prompt is always cached.
|
||||
"""
|
||||
|
||||
provider_name: str
|
||||
"""Provider identifier ("anthropic", "gemini") — used for api_logs."""
|
||||
model: str
|
||||
|
||||
@abstractmethod
|
||||
async def complete(
|
||||
self,
|
||||
*,
|
||||
messages: list[Message],
|
||||
tools: list[ToolSchema] | None = None,
|
||||
tool_choice: ToolChoice = "auto",
|
||||
) -> Completion:
|
||||
"""Run one inference turn."""
|
||||
|
||||
@abstractmethod
|
||||
async def close(self) -> None:
|
||||
"""Release any provider-side resources (e.g. delete a cache object).
|
||||
Safe to call multiple times."""
|
||||
|
||||
|
||||
class LLMProvider(Protocol):
|
||||
"""Top-level provider interface."""
|
||||
|
||||
name: str
|
||||
|
||||
async def create_session(
|
||||
self,
|
||||
*,
|
||||
model: str,
|
||||
system: str,
|
||||
max_tokens: int = 4096,
|
||||
temperature: float | None = None,
|
||||
) -> LLMSession:
|
||||
"""Construct a session. ``system`` is always cached by the session.
|
||||
|
||||
``temperature`` — if not None, applied to every ``complete()`` call on
|
||||
this session. ``None`` means use the provider's default. Set to 0.0
|
||||
for deterministic-as-possible behavior in agentic loops where the same
|
||||
inputs should produce the same outputs."""
|
||||
...
|
||||
|
||||
async def run_skill(
|
||||
self,
|
||||
*,
|
||||
skill_name: str,
|
||||
model: str,
|
||||
system: str,
|
||||
user_text: str,
|
||||
pdf_path: str | None,
|
||||
output_tool: ToolSchema,
|
||||
) -> tuple[dict, "Completion"]:
|
||||
"""Execute a managed Skill and return (forced-tool input, Completion).
|
||||
|
||||
Anthropic uses Console Skills (skill_id + container + code_execution
|
||||
beta). Gemini raises ``NotImplementedError`` — there is no
|
||||
Gemini-managed-Skill equivalent today; if you want a Gemini path for
|
||||
skill-style extraction, inline the SKILL.md content as ``system`` and
|
||||
run validation locally."""
|
||||
...
|
||||
@@ -0,0 +1,73 @@
|
||||
"""Provider factory + per-stage routing."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from functools import lru_cache
|
||||
from typing import Awaitable, Callable, TypeVar
|
||||
|
||||
from backend.config import settings
|
||||
from backend.services.llm.base import LLMProvider
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
@lru_cache(maxsize=4)
|
||||
def get_provider_by_name(name: str) -> LLMProvider:
|
||||
"""Return a singleton provider instance for ``name`` ("anthropic" |
|
||||
"gemini"). Used by :func:`get_provider` and :func:`call_with_fallback`."""
|
||||
if name == "anthropic":
|
||||
from backend.services.llm.anthropic_provider import AnthropicProvider
|
||||
return AnthropicProvider()
|
||||
if name == "gemini":
|
||||
from backend.services.llm.gemini_provider import GeminiProvider
|
||||
return GeminiProvider()
|
||||
raise ValueError(f"Unknown LLM provider: {name!r}")
|
||||
|
||||
|
||||
# Backwards-compatible alias
|
||||
_get_provider_by_name = get_provider_by_name
|
||||
|
||||
|
||||
def get_provider(stage: str) -> LLMProvider:
|
||||
"""Return the provider configured for ``stage``.
|
||||
|
||||
Falls back to ``settings.provider_default`` if no per-stage override.
|
||||
Providers are cached per-name, so repeated calls return the same
|
||||
instance (and share the underlying SDK client)."""
|
||||
name = settings.provider_for_stage(stage)
|
||||
return get_provider_by_name(name)
|
||||
|
||||
|
||||
async def call_with_fallback(
|
||||
stage: str,
|
||||
body: Callable[[LLMProvider, str], Awaitable[T]],
|
||||
) -> T:
|
||||
"""Run ``body(provider, model)`` for ``stage``; on any exception,
|
||||
retry once with the fallback provider/model if one is configured via
|
||||
``FALLBACK_PROVIDER_<STAGE>`` / ``FALLBACK_MODEL_<STAGE>``.
|
||||
|
||||
The fallback runs ``body`` from scratch — any tokens spent in the
|
||||
primary attempt are lost (and not logged). ``asyncio.CancelledError``
|
||||
is always re-raised so cancellation still works.
|
||||
"""
|
||||
primary_provider = get_provider(stage)
|
||||
primary_model = settings.model_for_stage(stage)
|
||||
try:
|
||||
return await body(primary_provider, primary_model)
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception as exc:
|
||||
fb = settings.fallback_for_stage(stage)
|
||||
if fb is None:
|
||||
raise
|
||||
log.warning(
|
||||
"[%s] primary %s/%s failed (%s) — falling back to %s/%s",
|
||||
stage, primary_provider.name, primary_model,
|
||||
exc, fb[0], fb[1],
|
||||
)
|
||||
fallback_provider = get_provider_by_name(fb[0])
|
||||
return await body(fallback_provider, fb[1])
|
||||
@@ -0,0 +1,379 @@
|
||||
"""Gemini provider — wraps google-genai async client.
|
||||
|
||||
Translates the unified ``Message`` / ``Completion`` shapes into Gemini's
|
||||
native ``Content`` / ``Part`` format. Caching uses ``CachedContent``: on the
|
||||
first ``complete()`` call, cacheable blocks (system + any block flagged
|
||||
``cacheable=True`` in the first user message) are uploaded as a
|
||||
``CachedContent`` with TTL=30min; subsequent calls reference the cache by
|
||||
name. On ``close()`` the cache is deleted. If creation fails (e.g.
|
||||
sub-threshold token count), the session falls back to inline content with no
|
||||
caching for the remainder of the conversation.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from google import genai
|
||||
from google.genai import types as gtypes
|
||||
|
||||
from backend.config import settings
|
||||
from backend.services.llm.base import LLMProvider, LLMSession
|
||||
from backend.services.llm.types import (
|
||||
Completion,
|
||||
ContentBlock,
|
||||
Message,
|
||||
PdfBlock,
|
||||
TextBlock,
|
||||
ToolCall,
|
||||
ToolChoice,
|
||||
ToolResultBlock,
|
||||
ToolSchema,
|
||||
Usage,
|
||||
)
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
_CACHE_TTL = "1800s" # 30 min — covers our longest agent loop with margin
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Translation helpers — unified types ↔ Gemini Parts/Contents
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _block_to_part(b: ContentBlock) -> gtypes.Part:
|
||||
if isinstance(b, TextBlock):
|
||||
return gtypes.Part(
|
||||
text=b.text,
|
||||
thought_signature=b.thought_signature,
|
||||
)
|
||||
if isinstance(b, PdfBlock):
|
||||
return gtypes.Part(
|
||||
inline_data=gtypes.Blob(
|
||||
mime_type="application/pdf",
|
||||
data=Path(b.path).read_bytes(),
|
||||
),
|
||||
)
|
||||
if isinstance(b, ToolCall):
|
||||
return gtypes.Part(
|
||||
function_call=gtypes.FunctionCall(
|
||||
id=b.id or None,
|
||||
name=b.name,
|
||||
args=b.input,
|
||||
),
|
||||
thought_signature=b.thought_signature,
|
||||
)
|
||||
if isinstance(b, ToolResultBlock):
|
||||
return gtypes.Part(
|
||||
function_response=gtypes.FunctionResponse(
|
||||
id=b.tool_use_id or None,
|
||||
name=b.name,
|
||||
# FunctionResponse.response is a dict — wrap string content
|
||||
response={"result": b.content},
|
||||
),
|
||||
)
|
||||
raise TypeError(f"Unknown ContentBlock: {type(b).__name__}")
|
||||
|
||||
|
||||
def _message_to_content(m: Message) -> gtypes.Content:
|
||||
# Gemini uses "user" and "model" (not "assistant")
|
||||
role = "model" if m.role == "assistant" else "user"
|
||||
return gtypes.Content(
|
||||
role=role,
|
||||
parts=[_block_to_part(b) for b in m.content],
|
||||
)
|
||||
|
||||
|
||||
def _tool_to_function_declaration(t: ToolSchema) -> gtypes.FunctionDeclaration:
|
||||
return gtypes.FunctionDeclaration(
|
||||
name=t.name,
|
||||
description=t.description,
|
||||
parameters_json_schema=t.input_schema,
|
||||
)
|
||||
|
||||
|
||||
def _tools_to_gemini(tools: list[ToolSchema]) -> list[gtypes.Tool]:
|
||||
return [
|
||||
gtypes.Tool(
|
||||
function_declarations=[_tool_to_function_declaration(t) for t in tools],
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def _tool_choice_to_config(c: ToolChoice) -> gtypes.ToolConfig:
|
||||
if c == "auto":
|
||||
return gtypes.ToolConfig(
|
||||
function_calling_config=gtypes.FunctionCallingConfig(mode="AUTO"),
|
||||
)
|
||||
if c == "none":
|
||||
return gtypes.ToolConfig(
|
||||
function_calling_config=gtypes.FunctionCallingConfig(mode="NONE"),
|
||||
)
|
||||
if isinstance(c, dict) and "name" in c:
|
||||
return gtypes.ToolConfig(
|
||||
function_calling_config=gtypes.FunctionCallingConfig(
|
||||
mode="ANY",
|
||||
allowed_function_names=[c["name"]],
|
||||
),
|
||||
)
|
||||
raise ValueError(f"Invalid tool_choice: {c!r}")
|
||||
|
||||
|
||||
def _from_gemini_response(resp: Any) -> Completion:
|
||||
"""Parse a Gemini GenerateContentResponse into a unified Completion."""
|
||||
text_parts: list[str] = []
|
||||
tool_calls: list[ToolCall] = []
|
||||
raw_blocks: list[ContentBlock] = []
|
||||
stop_reason = "unknown"
|
||||
|
||||
candidates = getattr(resp, "candidates", None) or []
|
||||
if candidates:
|
||||
cand = candidates[0]
|
||||
finish = getattr(cand, "finish_reason", None)
|
||||
if finish:
|
||||
stop_reason = str(finish).lower().split(".")[-1]
|
||||
content = getattr(cand, "content", None)
|
||||
if content and content.parts:
|
||||
for part in content.parts:
|
||||
# Preserve thought_signature (Gemini 3 thinking-mode) for
|
||||
# exact replay on subsequent turns; missing signatures cause
|
||||
# 400 INVALID_ARGUMENT on the next call.
|
||||
sig = getattr(part, "thought_signature", None)
|
||||
if getattr(part, "text", None):
|
||||
text_parts.append(part.text)
|
||||
raw_blocks.append(TextBlock(
|
||||
text=part.text, thought_signature=sig,
|
||||
))
|
||||
elif getattr(part, "function_call", None):
|
||||
fc = part.function_call
|
||||
tc = ToolCall(
|
||||
id=fc.id or f"{fc.name}_{len(tool_calls)}",
|
||||
name=fc.name,
|
||||
input=dict(fc.args or {}),
|
||||
thought_signature=sig,
|
||||
)
|
||||
tool_calls.append(tc)
|
||||
raw_blocks.append(tc)
|
||||
|
||||
usage_md = getattr(resp, "usage_metadata", None)
|
||||
if usage_md is not None:
|
||||
prompt_tokens = usage_md.prompt_token_count or 0
|
||||
cached_tokens = usage_md.cached_content_token_count or 0
|
||||
# Gemini reports prompt_token_count as the TOTAL prompt tokens —
|
||||
# cached tokens are billed at the cache-read rate, the rest at the
|
||||
# input rate. Subtract so they don't double-count.
|
||||
non_cached = max(0, prompt_tokens - cached_tokens)
|
||||
# Thinking-mode models (2.5 Pro, 3 series) report reasoning tokens
|
||||
# in thoughts_token_count, billed at the output rate. Fold into
|
||||
# output_tokens so cost accounting matches Gemini's actual bill.
|
||||
thoughts_tokens = getattr(usage_md, "thoughts_token_count", 0) or 0
|
||||
usage = Usage(
|
||||
input_tokens=non_cached,
|
||||
output_tokens=(usage_md.candidates_token_count or 0) + thoughts_tokens,
|
||||
cache_creation_tokens=0, # Gemini doesn't expose this separately
|
||||
cache_read_tokens=cached_tokens,
|
||||
)
|
||||
else:
|
||||
usage = Usage()
|
||||
|
||||
return Completion(
|
||||
text="".join(text_parts),
|
||||
tool_calls=tool_calls,
|
||||
usage=usage,
|
||||
stop_reason=stop_reason,
|
||||
raw_assistant_blocks=raw_blocks,
|
||||
)
|
||||
|
||||
|
||||
def _is_first_user_message_fully_cacheable(messages: list[Message]) -> bool:
|
||||
"""We cache only when EVERY block in the very first user message is
|
||||
flagged cacheable. This matches our actual usage (validation + power
|
||||
tree both pass entirely cacheable initial messages) and avoids brittle
|
||||
partial-cache scenarios."""
|
||||
if not messages:
|
||||
return False
|
||||
first = messages[0]
|
||||
if first.role != "user" or not first.content:
|
||||
return False
|
||||
return all(
|
||||
isinstance(b, (TextBlock, PdfBlock)) and b.cacheable
|
||||
for b in first.content
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Session
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class GeminiSession(LLMSession):
|
||||
provider_name = "gemini"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
client: genai.Client,
|
||||
model: str,
|
||||
system: str,
|
||||
max_tokens: int,
|
||||
temperature: float | None = None,
|
||||
) -> None:
|
||||
self._client = client
|
||||
self.model = model
|
||||
self._system = system
|
||||
self._max_tokens = max_tokens
|
||||
self._temperature = temperature
|
||||
self._cache_name: str | None = None
|
||||
self._cache_attempted = False
|
||||
|
||||
async def _try_create_cache(self, first_msg: Message) -> str | None:
|
||||
"""Attempt to create a CachedContent from system + first user message.
|
||||
Returns the cache name on success, None on failure."""
|
||||
try:
|
||||
parts = [_block_to_part(b) for b in first_msg.content]
|
||||
cache = await self._client.aio.caches.create(
|
||||
model=self.model,
|
||||
config=gtypes.CreateCachedContentConfig(
|
||||
system_instruction=self._system,
|
||||
contents=[gtypes.Content(role="user", parts=parts)],
|
||||
ttl=_CACHE_TTL,
|
||||
),
|
||||
)
|
||||
log.info(
|
||||
"Gemini cache created (%s, model=%s, ttl=%s)",
|
||||
cache.name, self.model, _CACHE_TTL,
|
||||
)
|
||||
return cache.name
|
||||
except Exception as exc:
|
||||
log.info(
|
||||
"Gemini cache creation skipped (%s) — falling back to inline",
|
||||
exc,
|
||||
)
|
||||
return None
|
||||
|
||||
async def complete(
|
||||
self,
|
||||
*,
|
||||
messages: list[Message],
|
||||
tools: list[ToolSchema] | None = None,
|
||||
tool_choice: ToolChoice = "auto",
|
||||
) -> Completion:
|
||||
if not messages:
|
||||
raise ValueError("Gemini complete() requires at least one message")
|
||||
|
||||
# First call: decide whether to cache
|
||||
if not self._cache_attempted:
|
||||
self._cache_attempted = True
|
||||
if _is_first_user_message_fully_cacheable(messages):
|
||||
self._cache_name = await self._try_create_cache(messages[0])
|
||||
|
||||
# Build per-call contents
|
||||
if self._cache_name:
|
||||
# Skip the cached first message — its contents are in the cache
|
||||
contents = [_message_to_content(m) for m in messages[1:]]
|
||||
else:
|
||||
contents = [_message_to_content(m) for m in messages]
|
||||
|
||||
# Build config
|
||||
config_kwargs: dict[str, Any] = {
|
||||
"max_output_tokens": self._max_tokens,
|
||||
}
|
||||
if self._temperature is not None:
|
||||
config_kwargs["temperature"] = self._temperature
|
||||
if self._cache_name:
|
||||
config_kwargs["cached_content"] = self._cache_name
|
||||
else:
|
||||
config_kwargs["system_instruction"] = self._system
|
||||
if tools:
|
||||
config_kwargs["tools"] = _tools_to_gemini(tools)
|
||||
config_kwargs["tool_config"] = _tool_choice_to_config(tool_choice)
|
||||
|
||||
config = gtypes.GenerateContentConfig(**config_kwargs)
|
||||
|
||||
# When using cached_content, Gemini still requires non-empty contents.
|
||||
# If the cached path leaves us with no per-call contents (only happens
|
||||
# on the very first turn with a cached initial message), seed with a
|
||||
# minimal continuation prompt.
|
||||
if self._cache_name and not contents:
|
||||
contents = [gtypes.Content(role="user", parts=[gtypes.Part(text="Continue.")])]
|
||||
|
||||
try:
|
||||
resp = await self._client.aio.models.generate_content(
|
||||
model=self.model,
|
||||
contents=contents,
|
||||
config=config,
|
||||
)
|
||||
except Exception as exc:
|
||||
# Cache may have expired mid-loop — drop it and retry inline once
|
||||
if self._cache_name and "cache" in str(exc).lower():
|
||||
log.warning("Gemini cache failed (%s) — retrying inline", exc)
|
||||
self._cache_name = None
|
||||
return await self.complete(
|
||||
messages=messages, tools=tools, tool_choice=tool_choice,
|
||||
)
|
||||
raise
|
||||
|
||||
return _from_gemini_response(resp)
|
||||
|
||||
async def close(self) -> None:
|
||||
if self._cache_name:
|
||||
try:
|
||||
await self._client.aio.caches.delete(name=self._cache_name)
|
||||
except Exception as exc:
|
||||
log.warning("Gemini cache delete failed (%s): %s", self._cache_name, exc)
|
||||
finally:
|
||||
self._cache_name = None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Provider
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class GeminiProvider(LLMProvider):
|
||||
name = "gemini"
|
||||
|
||||
def __init__(self) -> None:
|
||||
api_key = settings.gemini_api_key
|
||||
if not api_key:
|
||||
raise RuntimeError(
|
||||
"GEMINI_API_KEY is not set. Either set it in .env or route "
|
||||
"this stage to Anthropic via PROVIDER_<STAGE>=anthropic."
|
||||
)
|
||||
self._client = genai.Client(api_key=api_key)
|
||||
|
||||
async def create_session(
|
||||
self,
|
||||
*,
|
||||
model: str,
|
||||
system: str,
|
||||
max_tokens: int = 4096,
|
||||
temperature: float | None = None,
|
||||
) -> LLMSession:
|
||||
return GeminiSession(
|
||||
client=self._client,
|
||||
model=model,
|
||||
system=system,
|
||||
max_tokens=max_tokens,
|
||||
temperature=temperature,
|
||||
)
|
||||
|
||||
async def run_skill(
|
||||
self,
|
||||
*,
|
||||
skill_name: str,
|
||||
model: str,
|
||||
system: str,
|
||||
user_text: str,
|
||||
pdf_path: str | None,
|
||||
output_tool: ToolSchema,
|
||||
) -> tuple[dict, Completion]:
|
||||
raise NotImplementedError(
|
||||
f"GeminiProvider.run_skill() not implemented (skill={skill_name!r}). "
|
||||
f"Anthropic Console Skills have no Gemini equivalent. To migrate "
|
||||
f"this skill to Gemini, inline its SKILL.md as the system prompt "
|
||||
f"and run validate.py locally."
|
||||
)
|
||||
@@ -0,0 +1,81 @@
|
||||
"""Per-provider pricing tables and cost computation.
|
||||
|
||||
Replaces the flat ``PRICING`` dict that used to live in
|
||||
``backend/services/api_logs.py``. Indexed by (provider, model).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
|
||||
# Per-million-token USD rates. Source-of-truth links:
|
||||
# Anthropic: https://docs.anthropic.com/en/docs/about-claude/pricing
|
||||
# Google: https://ai.google.dev/pricing
|
||||
# Last updated: 2026-07-01
|
||||
PRICING: dict[str, dict[str, dict[str, float]]] = {
|
||||
"anthropic": {
|
||||
"claude-opus-4-6": {"input": 5.00, "output": 25.00},
|
||||
"claude-opus-4-5": {"input": 5.00, "output": 25.00},
|
||||
"claude-opus-4-1": {"input": 15.00, "output": 75.00},
|
||||
"claude-opus-4": {"input": 15.00, "output": 75.00},
|
||||
"claude-sonnet-4-6": {"input": 3.00, "output": 15.00},
|
||||
# Sonnet 5 standard rate (== Sonnet 4.6). Introductory pricing of
|
||||
# $2/$10 runs through 2026-08-31; intentionally NOT tracked here —
|
||||
# chosen set-and-forget so no dated bump is needed on 2026-09-01.
|
||||
# (New tokenizer emits ~30% more tokens, so per-run cost still rises.)
|
||||
"claude-sonnet-5": {"input": 3.00, "output": 15.00},
|
||||
"claude-sonnet-4-5": {"input": 3.00, "output": 15.00},
|
||||
"claude-sonnet-4": {"input": 3.00, "output": 15.00},
|
||||
"claude-haiku-4-5-20251001": {"input": 1.00, "output": 5.00},
|
||||
"claude-haiku-4-5": {"input": 1.00, "output": 5.00},
|
||||
"claude-haiku-3-5": {"input": 0.80, "output": 4.00},
|
||||
"default": {"input": 3.00, "output": 15.00},
|
||||
},
|
||||
"gemini": {
|
||||
# Gemini 3 Flash pricing (per 1M tokens). Preview alias mirrors GA.
|
||||
"gemini-3-flash-preview": {"input": 0.30, "output": 2.50},
|
||||
"gemini-3-flash": {"input": 0.30, "output": 2.50},
|
||||
"gemini-flash-latest": {"input": 0.30, "output": 2.50},
|
||||
"gemini-2.5-flash": {"input": 0.30, "output": 2.50},
|
||||
"gemini-2.5-pro": {"input": 1.25, "output": 10.00},
|
||||
# Gemini 3.1 Pro Preview — standard tier, prompts ≤200k tokens.
|
||||
# Above 200k Google charges $4.00/$18.00; we don't yet split by
|
||||
# prompt size, so we use the smaller-tier rate. Almost every
|
||||
# pipeline call here is well under 200k.
|
||||
"gemini-3.1-pro-preview": {"input": 2.00, "output": 12.00},
|
||||
"gemini-3-pro-preview": {"input": 2.00, "output": 12.00},
|
||||
"default": {"input": 0.30, "output": 2.50},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
# Per-provider cache token multipliers, applied on top of the input rate.
|
||||
# create: cost when a cache is *written* (Anthropic charges 1.25× input;
|
||||
# Gemini charges 1.0× input — caching writes are billed as a
|
||||
# normal input pass)
|
||||
# read: cost when a cached prefix is *reused* (much cheaper)
|
||||
CACHE_RATES: dict[str, dict[str, float]] = {
|
||||
"anthropic": {"create": 1.25, "read": 0.10},
|
||||
"gemini": {"create": 1.00, "read": 0.25},
|
||||
}
|
||||
|
||||
|
||||
def cost_for_entry(entry: dict) -> float:
|
||||
"""USD cost for an api_logs entry. Reads ``provider`` (default
|
||||
``anthropic`` for legacy entries) and ``model`` to pick rates."""
|
||||
provider = entry.get("provider") or "anthropic"
|
||||
table = PRICING.get(provider) or PRICING["anthropic"]
|
||||
rates = table.get(entry.get("model", ""), table["default"])
|
||||
cache_rates = CACHE_RATES.get(provider, CACHE_RATES["anthropic"])
|
||||
input_rate = rates["input"]
|
||||
output_rate = rates["output"]
|
||||
return (
|
||||
entry.get("input_tokens", 0) * input_rate
|
||||
+ entry.get("cache_creation_input_tokens", 0) * input_rate * cache_rates["create"]
|
||||
+ entry.get("cache_read_input_tokens", 0) * input_rate * cache_rates["read"]
|
||||
+ entry.get("output_tokens", 0) * output_rate
|
||||
) / 1_000_000
|
||||
|
||||
|
||||
def total_cost(entries: list[dict]) -> float:
|
||||
"""Sum USD across entries."""
|
||||
return round(sum(cost_for_entry(e) for e in entries), 6)
|
||||
@@ -0,0 +1,116 @@
|
||||
"""Provider-agnostic message and completion types.
|
||||
|
||||
These dataclasses are the lingua franca between calling code and providers.
|
||||
Each provider implementation translates these into its native shape on the
|
||||
way out and back into these on the way in.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any, Literal
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Content blocks — what goes inside a Message
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@dataclass
|
||||
class TextBlock:
|
||||
text: str
|
||||
cacheable: bool = False
|
||||
# Gemini 3 / thinking-mode: opaque bytes the model returns alongside text
|
||||
# parts that came from internal reasoning. Must be replayed verbatim when
|
||||
# this turn is fed back into the conversation, or the next call 400s.
|
||||
# Anthropic: always None.
|
||||
thought_signature: bytes | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class PdfBlock:
|
||||
"""Inline PDF document. Provider encodes as base64 (Anthropic) or
|
||||
inline_data (Gemini) and applies caching policy if cacheable=True."""
|
||||
path: Path
|
||||
cacheable: bool = False
|
||||
|
||||
|
||||
@dataclass
|
||||
class ToolCall:
|
||||
"""Assistant turn: model called a tool."""
|
||||
id: str
|
||||
name: str
|
||||
input: dict[str, Any]
|
||||
# Same purpose as TextBlock.thought_signature — Gemini 3 attaches one
|
||||
# to every function_call part when thinking is on. Round-trip required.
|
||||
thought_signature: bytes | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class ToolResultBlock:
|
||||
"""User turn: result fed back from a tool the model invoked previously."""
|
||||
tool_use_id: str
|
||||
name: str
|
||||
content: str
|
||||
|
||||
|
||||
ContentBlock = TextBlock | PdfBlock | ToolCall | ToolResultBlock
|
||||
|
||||
|
||||
@dataclass
|
||||
class Message:
|
||||
role: Literal["user", "assistant"]
|
||||
content: list[ContentBlock]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tools
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@dataclass
|
||||
class ToolSchema:
|
||||
"""JSON-schema tool definition. Both providers accept the same shape."""
|
||||
name: str
|
||||
description: str
|
||||
input_schema: dict[str, Any]
|
||||
|
||||
|
||||
# Tool choice: "auto" (model picks), "none" (no tools), or a forced name
|
||||
ToolChoice = Literal["auto", "none"] | dict # {"name": "save_xyz"}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Completion / usage
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@dataclass
|
||||
class Usage:
|
||||
"""Token usage normalised across providers.
|
||||
|
||||
Anthropic exposes cache_creation_input_tokens (write) and
|
||||
cache_read_input_tokens (hit). Gemini only exposes a cache hit count
|
||||
(cached_content_token_count) — its cache writes don't bill as input.
|
||||
|
||||
For Gemini, ``cache_creation_tokens`` is always 0; ``cache_read_tokens``
|
||||
holds the cached hit count when a cache was used.
|
||||
"""
|
||||
input_tokens: int = 0
|
||||
output_tokens: int = 0
|
||||
cache_creation_tokens: int = 0
|
||||
cache_read_tokens: int = 0
|
||||
|
||||
|
||||
@dataclass
|
||||
class Completion:
|
||||
"""Result of a single provider.complete() / session.complete() call."""
|
||||
text: str # any text block(s) concatenated
|
||||
tool_calls: list[ToolCall]
|
||||
usage: Usage
|
||||
stop_reason: str
|
||||
raw_assistant_blocks: list[ContentBlock] = field(default_factory=list)
|
||||
"""The full assistant message, in our normalised content-block form, so
|
||||
callers can append it back to the conversation history when continuing
|
||||
the loop."""
|
||||
@@ -0,0 +1,565 @@
|
||||
"""Per-IC normalize pass — dedup findings with a shared root cause and
|
||||
re-grade severity against a fixed rubric.
|
||||
|
||||
Two runs of the reviewer on identical inputs can produce different *judgments*
|
||||
(severity choices, finding-splitting) even when they reach the same underlying
|
||||
observations. This module runs a single small LLM call against the structured
|
||||
findings (no PDF, no graph tools) to:
|
||||
|
||||
- merge findings that describe the same defect from different angles, and
|
||||
- re-grade each remaining finding's severity using an anchored rubric.
|
||||
|
||||
It is intentionally conservative: if the call fails, the schema is malformed,
|
||||
or the index coverage is invalid, the original findings are returned
|
||||
unchanged. A normalize failure must never break the per-IC review.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
from datetime import datetime, timezone
|
||||
from typing import Awaitable, Callable
|
||||
|
||||
from backend.config import settings
|
||||
from backend.pinscopex.models import Finding
|
||||
from backend.services.api_logs import ApiLogger
|
||||
from backend.services.llm import Message, TextBlock
|
||||
from backend.services.llm.factory import call_with_fallback
|
||||
from backend.services.llm.types import ToolSchema
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
# Severity ordering for the downgrade-only clamp. Normalize may lower a
|
||||
# finding's severity but never raise it above what the reviewer chose — the
|
||||
# reviewer had the datasheet + graph; this pass sees only text.
|
||||
_INFO, _WARN, _ERR = 0, 1, 2
|
||||
_SEVERITY_RANK = {"INFO": _INFO, "WARNING": _WARN, "ERROR": _ERR}
|
||||
_RANK_TO_SEV = {_INFO: "INFO", _WARN: "WARNING", _ERR: "ERROR"}
|
||||
|
||||
|
||||
def _is_unverified(why: str | None) -> bool:
|
||||
"""True when a finding's ``why`` is flagged ``Unverified:`` — the reviewer
|
||||
could not confirm the spec from the datasheet and deliberately hedged."""
|
||||
return (why or "").lstrip().lower().startswith("unverified:")
|
||||
|
||||
|
||||
SYSTEM_PROMPT = """\
|
||||
You normalize a single IC's review findings for a hardware design review tool.
|
||||
|
||||
Three operations:
|
||||
1. **Drop** self-cancelling findings whose own analysis confirms the \
|
||||
design is correct.
|
||||
2. **Merge** findings that share a single-fix root cause (atomic-fix test).
|
||||
3. **Re-grade severity** independently against the rubric below.
|
||||
|
||||
You CANNOT invent new findings or new facts. Every original finding \
|
||||
(numbered 1..N) must end up in exactly one of:
|
||||
- a kept/merged entry in `findings` (referenced by `merged_from`), or
|
||||
- a dropped entry in `dropped` (referenced by `index`).
|
||||
|
||||
You ARE shown the reviewer's original severity. The reviewer graded each \
|
||||
finding with the datasheet PDF and the design graph in front of it; you \
|
||||
see only the finding text. You may **lower** a severity when the rubric \
|
||||
clearly supports a milder grade — over-stated, conditional, or the \
|
||||
`why` itself flags incomplete evidence — but you must **never raise** a \
|
||||
finding above the reviewer's grade. Upgrading is where you have the \
|
||||
least evidence and do the most damage: a normalize pass that promotes a \
|
||||
hedged WARNING into a confident ERROR is the exact failure this rule \
|
||||
exists to prevent.
|
||||
|
||||
### Drop rule (self-cancelling findings)
|
||||
|
||||
A finding is self-cancelling when its own `why` confirms the requirement \
|
||||
is met or no issue actually exists. The surface reading suggested a \
|
||||
problem; the analysis itself proved otherwise. Examples:
|
||||
|
||||
- "Output cap C1 (100 nF) is below the 1 µF minimum, but C24 (1 µF) in \
|
||||
parallel satisfies the spec." → drop. Total Cout meets spec; no issue.
|
||||
- "No dedicated input decoupling cap directly at VIN — but C3 (1 µF) is \
|
||||
on the VIN net and satisfies the requirement." → drop. C3 IS the input \
|
||||
cap, in the correct place.
|
||||
- "Pin X appears unconnected, however net Y shows it is grounded." → drop.
|
||||
|
||||
Drop these via the `dropped` array with a short `reason`. Do NOT keep \
|
||||
them as INFO — they dilute the signal of real issues. If the `why` \
|
||||
contains "satisfies", "meets the requirement", "is in the correct \
|
||||
place", "no issue", or equivalent language confirming the design is \
|
||||
correct, the finding is almost certainly self-cancelling.
|
||||
|
||||
A finding that flags a real concern but acknowledges *partial* \
|
||||
mitigation or *conditional* validity ("works at low load only", "meets \
|
||||
spec only at room temperature") is NOT self-cancelling — keep it.
|
||||
|
||||
### Root-cause merge rule (atomic-fix test)
|
||||
|
||||
Two findings share a root cause if a SINGLE atomic change resolves both. \
|
||||
The atomic-fix test: can you describe the fix in `single_fix` as ONE \
|
||||
action — remove X, replace X with Y, rewire X to Z, or add X — without \
|
||||
using "and", "also", or describing multiple steps?
|
||||
|
||||
If yes: merge. Write the combined `finding` title naming the root cause \
|
||||
once. Restate downstream consequences inside `why`. Keep `source_page`, \
|
||||
`source_quote`, and `reference` from the original with the strongest \
|
||||
evidence.
|
||||
|
||||
If no: do NOT merge. Two defects involving the same component, the same \
|
||||
net, or the same fix-area are still separate root causes when they \
|
||||
require separate changes.
|
||||
|
||||
**Invalid merge example**: combining "R1 (17.8Ω) in series with VIN \
|
||||
causes dropout" with "EN tied to VIN — no independent enable" into a \
|
||||
single ERROR with `single_fix` = "Remove R1 AND route EN from a \
|
||||
separate GPIO." That is TWO changes (remove R1; rewire EN). Keep these \
|
||||
as two separate findings — the dropout finding alone may be ERROR or \
|
||||
WARNING; the EN finding is INFO.
|
||||
|
||||
When you merge, you MUST populate `single_fix` with the one atomic \
|
||||
action. If you cannot, do not merge.
|
||||
|
||||
### Severity rubric (grade independently)
|
||||
|
||||
- **ERROR**: The circuit, as wired, will not function correctly. The \
|
||||
output won't reach spec, the regulator won't regulate, the signal \
|
||||
won't reach the destination, abs-max is exceeded with a strict \
|
||||
inequality (actual > limit), or a required pin is left undriven. A \
|
||||
concrete failure mode is reachable from the design as drawn.
|
||||
|
||||
- **WARNING**: The circuit functions but has reduced margin, degraded \
|
||||
performance, or conditional malfunction (depends on load, \
|
||||
temperature, or firmware state). A recommended-but-not-required \
|
||||
component is missing. The finding is "unverified" because evidence \
|
||||
was incomplete.
|
||||
|
||||
- **INFO**: A valid topology choice that disables an optional \
|
||||
feature, or a documentation/layout observation that cannot be \
|
||||
verified from a netlist. Examples: EN tied to VIN to use the LDO's \
|
||||
always-on mode (firmware shutdown unavailable but the chip works), \
|
||||
an optional bypass cap omitted on a non-critical pin.
|
||||
|
||||
Grade each kept finding against this rubric, but only ever *downward* \
|
||||
from the reviewer's original severity (shown to you). A merged \
|
||||
finding's severity may not exceed the highest original severity among \
|
||||
its members. If a finding's `why` begins with `Unverified:`, the \
|
||||
reviewer could not confirm the spec from the datasheet — keep the \
|
||||
`Unverified:` prefix and never grade it above WARNING.
|
||||
|
||||
### Output
|
||||
|
||||
Call the `submit_normalized` tool exactly once with:
|
||||
- `findings`: kept and merged entries (each with `merged_from` indices, \
|
||||
`single_fix` if merged, and a re-graded `status`).
|
||||
- `dropped`: self-cancelling entries (each with `index` and `reason`).
|
||||
|
||||
Every original index 1..N must appear in exactly one location across \
|
||||
both arrays. No index may appear twice.
|
||||
"""
|
||||
|
||||
|
||||
SUBMIT_NORMALIZED_SCHEMA = ToolSchema(
|
||||
name="submit_normalized",
|
||||
description=(
|
||||
"Submit the normalized findings. Every original finding (1..N) "
|
||||
"must appear in exactly one location across `findings.merged_from` "
|
||||
"or `dropped.index`."
|
||||
),
|
||||
input_schema={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"findings": {
|
||||
"type": "array",
|
||||
"description": (
|
||||
"Kept and merged findings. Re-graded severity; merged "
|
||||
"entries must include `single_fix`."
|
||||
),
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"merged_from": {
|
||||
"type": "array",
|
||||
"items": {"type": "integer"},
|
||||
"minItems": 1,
|
||||
"description": (
|
||||
"1-indexed positions in the original "
|
||||
"findings list this output entry "
|
||||
"represents. Length 1 = passed through; "
|
||||
"length > 1 = merged."
|
||||
),
|
||||
},
|
||||
"finding": {"type": "string"},
|
||||
"why": {"type": "string"},
|
||||
"status": {
|
||||
"type": "string",
|
||||
"enum": ["ERROR", "WARNING", "INFO"],
|
||||
},
|
||||
"recommendation": {"type": "string"},
|
||||
"source_page": {"type": ["integer", "null"]},
|
||||
"source_quote": {"type": "string"},
|
||||
"reference": {"type": "string"},
|
||||
"single_fix": {
|
||||
"type": "string",
|
||||
"description": (
|
||||
"REQUIRED when merged_from has length > 1. "
|
||||
"The single atomic component or net change "
|
||||
"that resolves ALL members of the merge "
|
||||
"(remove X, replace X with Y, rewire X to "
|
||||
"Z, or add X). If you cannot write the fix "
|
||||
"in one sentence without 'and' / 'also' / "
|
||||
"multiple steps, do NOT merge."
|
||||
),
|
||||
},
|
||||
"change_rationale": {
|
||||
"type": "string",
|
||||
"description": (
|
||||
"≤1 line: 'unchanged', or what changed "
|
||||
"and why (merged X+Y, graded <sev> per "
|
||||
"rubric because <reason>, ...)."
|
||||
),
|
||||
},
|
||||
},
|
||||
"required": [
|
||||
"merged_from",
|
||||
"finding",
|
||||
"why",
|
||||
"status",
|
||||
"recommendation",
|
||||
"change_rationale",
|
||||
],
|
||||
},
|
||||
},
|
||||
"dropped": {
|
||||
"type": "array",
|
||||
"description": (
|
||||
"Self-cancelling findings whose own `why` confirms "
|
||||
"the design is correct. These should NOT appear in "
|
||||
"`findings` — they are removed entirely from the "
|
||||
"report. Use this rather than demoting to INFO."
|
||||
),
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"index": {
|
||||
"type": "integer",
|
||||
"description": (
|
||||
"1-indexed position of the original "
|
||||
"finding being dropped."
|
||||
),
|
||||
},
|
||||
"reason": {
|
||||
"type": "string",
|
||||
"description": (
|
||||
"Short explanation of why the finding is "
|
||||
"self-cancelling (e.g., 'C1<1µF but C24 "
|
||||
"in parallel meets spec', 'C3 is the "
|
||||
"input cap, already in the correct "
|
||||
"place')."
|
||||
),
|
||||
},
|
||||
},
|
||||
"required": ["index", "reason"],
|
||||
},
|
||||
},
|
||||
},
|
||||
"required": ["findings"],
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _serialize_findings_for_prompt(findings: list[Finding]) -> str:
|
||||
"""Number the original findings 1..N and emit a compact JSON block.
|
||||
|
||||
The reviewer's `status` IS included: normalize re-grades only *downward*
|
||||
from it (the reviewer had the datasheet + graph; this pass sees only
|
||||
text). A deterministic clamp in ``_build_normalized`` enforces the
|
||||
downgrade-only invariant regardless of what the model returns.
|
||||
"""
|
||||
rows: list[dict] = []
|
||||
for i, f in enumerate(findings, start=1):
|
||||
rows.append({
|
||||
"index": i,
|
||||
"reviewer_severity": f.status,
|
||||
"finding": f.finding,
|
||||
"why": f.why,
|
||||
"recommendation": f.recommendation,
|
||||
"source_page": f.source_page,
|
||||
"source_quote": f.source_quote,
|
||||
"reference": f.reference,
|
||||
})
|
||||
return json.dumps(rows, indent=2)
|
||||
|
||||
|
||||
def _build_normalized(
|
||||
raw_findings: list[dict],
|
||||
raw_dropped: list[dict],
|
||||
originals: list[Finding],
|
||||
) -> tuple[list[Finding], list[dict]] | None:
|
||||
"""Validate the tool output and reconstruct Finding objects.
|
||||
|
||||
Returns ``(kept_findings, dropped_records)`` or ``None`` if coverage /
|
||||
schema validation fails (caller falls back to originals).
|
||||
|
||||
A merge with ``len(merged_from) > 1`` that omits ``single_fix`` is not
|
||||
a hard failure — the merge is rejected and its members fall back to
|
||||
their per-index originals. Self-cancelling drops require a non-empty
|
||||
`reason`; missing reason = treat as ungrouped and fail coverage.
|
||||
|
||||
The `change_rationale` and `single_fix` fields are informational and
|
||||
are not carried onto Finding objects; the full normalize trace keeps
|
||||
them for forensics.
|
||||
"""
|
||||
n = len(originals)
|
||||
seen: set[int] = set()
|
||||
result: list[Finding] = []
|
||||
dropped_records: list[dict] = []
|
||||
|
||||
# Process explicit drops first so indices are reserved before any
|
||||
# accidental double-coverage from a merge.
|
||||
for d in raw_dropped or []:
|
||||
if not isinstance(d, dict):
|
||||
return None
|
||||
try:
|
||||
idx = int(d.get("index"))
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
if idx < 1 or idx > n or idx in seen:
|
||||
return None
|
||||
reason = str(d.get("reason") or "").strip()
|
||||
if not reason:
|
||||
return None
|
||||
seen.add(idx)
|
||||
dropped_records.append({
|
||||
"index": idx,
|
||||
"reason": reason,
|
||||
"original_finding": originals[idx - 1].model_dump(mode="json"),
|
||||
})
|
||||
|
||||
for entry in raw_findings:
|
||||
if not isinstance(entry, dict):
|
||||
return None
|
||||
merged_from = entry.get("merged_from") or []
|
||||
if not isinstance(merged_from, list) or not merged_from:
|
||||
return None
|
||||
try:
|
||||
indices = [int(x) for x in merged_from]
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
for idx in indices:
|
||||
if idx < 1 or idx > n or idx in seen:
|
||||
return None
|
||||
seen.add(idx)
|
||||
|
||||
# Atomic-fix test: a merge (len > 1) must populate `single_fix`.
|
||||
# If missing, reject the merge and fall back to the per-index
|
||||
# originals — preserves coverage but un-merges. The reviewer's
|
||||
# original severity is preserved on the fallback path because we
|
||||
# construct each Finding directly from `originals[i-1]`.
|
||||
if len(indices) > 1:
|
||||
single_fix = str(entry.get("single_fix") or "").strip()
|
||||
if not single_fix:
|
||||
log.warning(
|
||||
"normalize: merge of %s lacks single_fix — falling "
|
||||
"back to per-index originals (un-merging)",
|
||||
indices,
|
||||
)
|
||||
for idx in indices:
|
||||
result.append(originals[idx - 1])
|
||||
continue
|
||||
|
||||
# Use the first original in the group as the canonical source for
|
||||
# fields the normalize layer doesn't own (designator, mpn, aspect,
|
||||
# finding_id). These are identical across an IC's findings anyway
|
||||
# since normalize is per-IC.
|
||||
canon = originals[indices[0] - 1]
|
||||
|
||||
# Severity safety net — downgrade-only. Normalize may lower a
|
||||
# finding's severity but never raise it above the reviewer's
|
||||
# calibrated grade (the reviewer had the datasheet + graph; this
|
||||
# pass sees only text). Cap at the highest original severity among
|
||||
# merged members; findings the reviewer marked "Unverified:" are
|
||||
# capped at WARNING and keep that prefix. This deterministic clamp
|
||||
# holds even when the model ignores the prompt instruction.
|
||||
members = [originals[i - 1] for i in indices]
|
||||
ceiling = max(_SEVERITY_RANK.get(m.status, _ERR) for m in members)
|
||||
unverified = any(_is_unverified(m.why) for m in members)
|
||||
if unverified:
|
||||
ceiling = min(ceiling, _WARN)
|
||||
proposed = str(entry.get("status") or canon.status)
|
||||
final_status = _RANK_TO_SEV[
|
||||
min(_SEVERITY_RANK.get(proposed, ceiling), ceiling)
|
||||
]
|
||||
|
||||
new_why = str(entry.get("why") or canon.why)
|
||||
if unverified and not _is_unverified(new_why):
|
||||
new_why = "Unverified: " + new_why
|
||||
|
||||
try:
|
||||
result.append(Finding(
|
||||
finding_id=canon.finding_id,
|
||||
designator=canon.designator,
|
||||
mpn=canon.mpn,
|
||||
aspect=canon.aspect,
|
||||
finding=str(entry.get("finding") or canon.finding),
|
||||
why=new_why,
|
||||
source_page=entry.get("source_page", canon.source_page),
|
||||
source_quote=str(entry.get("source_quote") or canon.source_quote),
|
||||
source_designator=canon.source_designator,
|
||||
status=final_status,
|
||||
recommendation=str(entry.get("recommendation") or canon.recommendation),
|
||||
reference=str(entry.get("reference") or canon.reference),
|
||||
))
|
||||
except Exception:
|
||||
log.exception("normalize: failed to build merged Finding")
|
||||
return None
|
||||
if seen != set(range(1, n + 1)):
|
||||
return None
|
||||
return result, dropped_records
|
||||
|
||||
|
||||
async def normalize_findings_async(
|
||||
ic_ref: str,
|
||||
mpn: str,
|
||||
findings: list[Finding],
|
||||
*,
|
||||
api_logger: ApiLogger | None = None,
|
||||
on_progress: Callable[[str, int, str, str], Awaitable[None]] | None = None,
|
||||
) -> tuple[list[Finding], dict]:
|
||||
"""Run the per-IC normalize pass.
|
||||
|
||||
Returns ``(normalized_findings, trace)``. On any failure (LLM error,
|
||||
schema violation, index coverage gap), returns the original findings
|
||||
unchanged with an ``error`` field set in the trace.
|
||||
"""
|
||||
trace: dict = {
|
||||
"ic_ref": ic_ref,
|
||||
"mpn": mpn,
|
||||
"timestamp": datetime.now(timezone.utc).isoformat(),
|
||||
"input_findings": [f.model_dump(mode="json") for f in findings],
|
||||
"output_findings": None,
|
||||
"dropped_findings": None,
|
||||
"submission": None,
|
||||
"model": None,
|
||||
"provider": None,
|
||||
"duration_ms": None,
|
||||
"error": None,
|
||||
}
|
||||
|
||||
# Nothing to do for 0 findings. With 1 finding there is no merge to
|
||||
# consider but the drop and re-grade rules still apply — let it
|
||||
# through to the LLM call.
|
||||
if not findings:
|
||||
trace["output_findings"] = []
|
||||
trace["dropped_findings"] = []
|
||||
trace["error"] = "skipped: 0 findings"
|
||||
return findings, trace
|
||||
|
||||
user_text = (
|
||||
f"Original findings for IC {ic_ref} ({mpn}). "
|
||||
f"There are {len(findings)} findings. "
|
||||
f"Indices are 1-based.\n\n"
|
||||
f"{_serialize_findings_for_prompt(findings)}\n\n"
|
||||
f"Normalize them per the rubric and call submit_normalized."
|
||||
)
|
||||
|
||||
t0 = time.monotonic()
|
||||
|
||||
async def _run(provider, model):
|
||||
trace["model"] = model
|
||||
trace["provider"] = provider.name
|
||||
session = await provider.create_session(
|
||||
model=model,
|
||||
system=SYSTEM_PROMPT,
|
||||
max_tokens=4096,
|
||||
temperature=0.0,
|
||||
)
|
||||
try:
|
||||
completion = await session.complete(
|
||||
messages=[Message(
|
||||
role="user",
|
||||
content=[TextBlock(text=user_text, cacheable=False)],
|
||||
)],
|
||||
tools=[SUBMIT_NORMALIZED_SCHEMA],
|
||||
tool_choice={"name": "submit_normalized"},
|
||||
)
|
||||
if api_logger:
|
||||
api_logger.log(
|
||||
stage="normalize",
|
||||
identifier=ic_ref,
|
||||
model=model,
|
||||
provider=provider.name,
|
||||
input_tokens=completion.usage.input_tokens,
|
||||
output_tokens=completion.usage.output_tokens,
|
||||
cache_creation_input_tokens=completion.usage.cache_creation_tokens,
|
||||
cache_read_input_tokens=completion.usage.cache_read_tokens,
|
||||
duration_ms=int((time.monotonic() - t0) * 1000),
|
||||
stop_reason="submit_normalized",
|
||||
turns=1,
|
||||
)
|
||||
for tc in completion.tool_calls:
|
||||
if tc.name == "submit_normalized":
|
||||
return tc.input
|
||||
return None
|
||||
finally:
|
||||
await session.close()
|
||||
|
||||
try:
|
||||
submission = await call_with_fallback("normalize", _run)
|
||||
except Exception as exc:
|
||||
log.exception("normalize: call failed for %s", ic_ref)
|
||||
trace["error"] = f"{type(exc).__name__}: {exc}"
|
||||
trace["duration_ms"] = int((time.monotonic() - t0) * 1000)
|
||||
trace["output_findings"] = trace["input_findings"]
|
||||
return findings, trace
|
||||
|
||||
trace["duration_ms"] = int((time.monotonic() - t0) * 1000)
|
||||
trace["submission"] = submission
|
||||
|
||||
if not submission or not isinstance(submission, dict):
|
||||
trace["error"] = "no submission"
|
||||
trace["output_findings"] = trace["input_findings"]
|
||||
return findings, trace
|
||||
|
||||
raw_findings = submission.get("findings") or []
|
||||
if not isinstance(raw_findings, list):
|
||||
trace["error"] = "submission.findings not a list"
|
||||
trace["output_findings"] = trace["input_findings"]
|
||||
return findings, trace
|
||||
|
||||
raw_dropped = submission.get("dropped") or []
|
||||
if not isinstance(raw_dropped, list):
|
||||
trace["error"] = "submission.dropped not a list"
|
||||
trace["output_findings"] = trace["input_findings"]
|
||||
return findings, trace
|
||||
|
||||
built = _build_normalized(raw_findings, raw_dropped, findings)
|
||||
if built is None:
|
||||
trace["error"] = "invalid index coverage or schema"
|
||||
trace["output_findings"] = trace["input_findings"]
|
||||
log.warning(
|
||||
"normalize: invalid output for %s (%d originals, %d kept, "
|
||||
"%d dropped) — falling back to originals",
|
||||
ic_ref, len(findings), len(raw_findings), len(raw_dropped),
|
||||
)
|
||||
if on_progress:
|
||||
try:
|
||||
await on_progress(
|
||||
ic_ref, 0, "normalize_skipped",
|
||||
f"invalid output, kept {len(findings)} originals",
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
return findings, trace
|
||||
|
||||
normalized, dropped_records = built
|
||||
trace["output_findings"] = [f.model_dump(mode="json") for f in normalized]
|
||||
trace["dropped_findings"] = dropped_records
|
||||
if on_progress:
|
||||
try:
|
||||
await on_progress(
|
||||
ic_ref, 0, "normalized",
|
||||
f"{len(findings)} → {len(normalized)} kept, "
|
||||
f"{len(dropped_records)} dropped",
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
return normalized, trace
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,802 @@
|
||||
"""Project storage via StorageBackend.
|
||||
|
||||
Each project lives at users/{user_id}/projects/{id}/ with:
|
||||
project.json — metadata
|
||||
uploads/bom.csv — uploaded BOM
|
||||
uploads/netlist.asc — uploaded netlist
|
||||
uploads/datasheets/*.pdf — uploaded datasheets
|
||||
extracted/ — IC extraction output
|
||||
patterns/ — passive patterns
|
||||
models/ — cached component specs
|
||||
design_graph.json — graph output
|
||||
report.json — validation report
|
||||
|
||||
Library (global, shared across users):
|
||||
library/extracted/{mpn}.json
|
||||
library/patterns/{mfr}_{type}.json
|
||||
library/datasheets/{mpn}.pdf
|
||||
library/models/{mpn}.json — discrete/connector/crystal specs
|
||||
library/passives/{mpn}.json — DigiKey-resolved passive specs
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import uuid
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
from backend.pinscopex.utils import safe_mpn
|
||||
from backend.services.storage import StaleGeneration, StorageBackend
|
||||
|
||||
|
||||
class ProjectNotFound(Exception):
|
||||
"""An operation targeted a project whose metadata is gone.
|
||||
|
||||
Raised when ``project.json`` is missing — e.g. the project was deleted
|
||||
while a slow request (a large BOM upload) was still in flight. Callers /
|
||||
the global handler map this to a clean 404 instead of letting the raw
|
||||
storage NotFound bubble up as a 500 (which tears down the HTTP/2 stream
|
||||
mid-upload and surfaces in the browser as ERR_HTTP2_PROTOCOL_ERROR).
|
||||
"""
|
||||
|
||||
def __init__(self, project_id: str):
|
||||
self.project_id = project_id
|
||||
super().__init__(f"Project {project_id} not found")
|
||||
|
||||
|
||||
# Statuses
|
||||
STATUS_DRAFT = "draft"
|
||||
STATUS_QUEUED = "queued"
|
||||
STATUS_RUNNING = "running"
|
||||
STATUS_COMPLETE = "complete"
|
||||
STATUS_ERROR = "error"
|
||||
STATUS_CANCELLED = "cancelled"
|
||||
STATUS_PAUSED = "paused_insufficient_credits"
|
||||
|
||||
TERMINAL_STATUSES = frozenset({
|
||||
STATUS_COMPLETE, STATUS_ERROR, STATUS_CANCELLED, STATUS_PAUSED,
|
||||
})
|
||||
|
||||
|
||||
class StatusConflict(Exception):
|
||||
"""Raised when a status transition's preconditions don't hold.
|
||||
|
||||
Either the current status is not in ``from_status`` or another writer
|
||||
won the optimistic-concurrency race.
|
||||
"""
|
||||
|
||||
|
||||
class ProjectMeta(BaseModel):
|
||||
id: str
|
||||
name: str
|
||||
user_id: str = ""
|
||||
# draft | running | complete | error | cancelled
|
||||
# | paused_insufficient_credits | paused_by_user
|
||||
status: str = "draft"
|
||||
created: str = ""
|
||||
updated: str = ""
|
||||
has_bom: bool = False
|
||||
has_netlist: bool = False
|
||||
# "pads" | "edif" | None — None for legacy projects (pre-EDIF-support).
|
||||
# Legacy reads fall back to looking for netlist.asc on disk.
|
||||
netlist_format: str | None = None
|
||||
# When the EDIF file contains 2+ sub-designs, this is the list of
|
||||
# sub-design IDs (e.g. ["&0441"]) the user picked. None means "include
|
||||
# everything found in the file" — also the value when the netlist has a
|
||||
# single sub-design and no choice was offered.
|
||||
netlist_subdesigns: list[str] | None = None
|
||||
datasheet_count: int = 0
|
||||
summary: dict[str, int] | None = None
|
||||
component_mpns: dict[str, list[str]] | None = None # {ic: [...], passive: [...]}
|
||||
bom_columns: dict[str, str] | None = None # {reference: "...", mpn: "..."}
|
||||
# LCSC id → resolved manufacturer part number, populated by upload_bom when
|
||||
# the MPN column is detected as entirely LCSC ids (^C\d+$). The wizard UI
|
||||
# uses this to show "C12044 → STM32F103C8T6" alongside each row.
|
||||
lcsc_to_mpn: dict[str, str] | None = None
|
||||
# LCSC id → full purple-parts payload (mpn, manufacturer, package, description,
|
||||
# category, subcategory). Cached at upload time so the wizard's
|
||||
# /lcsc/resolve-passive endpoint can synthesize an auto-resolve call without
|
||||
# a second purple-parts round trip.
|
||||
lcsc_payloads: dict[str, dict] | None = None
|
||||
skipped_components: list[dict[str, str]] | None = None # [{identifier, stage, error}]
|
||||
pipeline_state: dict[str, Any] | None = None
|
||||
total_cost_usd: float | None = None
|
||||
collaborators: list[str] = [] # Clerk user_ids with access to this project
|
||||
tier: str = "demo" # user tier at project creation; "demo" default for pre-existing projects
|
||||
|
||||
# Credit-system fields
|
||||
credits_spent: float = 0.0
|
||||
estimate: dict[str, Any] | None = None # CostEstimate snapshot
|
||||
pause_checkpoint: dict[str, Any] | None = None # PauseCheckpoint on paused runs
|
||||
pause_reason: str | None = None
|
||||
completed_review_refs: list[str] = [] # IC refs already reviewed (persists across pauses)
|
||||
|
||||
# Pinscope app version that generated the project's report.
|
||||
# Stamped on the first /start transition and preserved thereafter.
|
||||
pinscope_version: str | None = None
|
||||
|
||||
# Worker bookkeeping (set by the API on enqueue, read by /events SSE
|
||||
# and by the stale-running sweeper).
|
||||
execution_name: str | None = None
|
||||
queued_at: str | None = None
|
||||
# User-initiated cancel signal — the worker reads this in its cancel
|
||||
# gate (inside _charge_for_logs) and exits cleanly.
|
||||
cancel_requested: bool = False
|
||||
|
||||
|
||||
def _project_prefix(user_id: str, project_id: str) -> str:
|
||||
return f"users/{user_id}/projects/{project_id}"
|
||||
|
||||
|
||||
def _meta_key(user_id: str, project_id: str) -> str:
|
||||
return f"{_project_prefix(user_id, project_id)}/project.json"
|
||||
|
||||
|
||||
def _read_meta(storage: StorageBackend, user_id: str, project_id: str) -> ProjectMeta:
|
||||
data = storage.read_json(_meta_key(user_id, project_id))
|
||||
return ProjectMeta.model_validate(data)
|
||||
|
||||
|
||||
def _read_meta_with_generation(
|
||||
storage: StorageBackend, user_id: str, project_id: str
|
||||
) -> tuple[ProjectMeta, int]:
|
||||
data, gen = storage.read_json_with_generation(_meta_key(user_id, project_id))
|
||||
return ProjectMeta.model_validate(data), gen
|
||||
|
||||
|
||||
def _write_meta(storage: StorageBackend, meta: ProjectMeta) -> None:
|
||||
meta.updated = datetime.now(timezone.utc).isoformat()
|
||||
storage.write_json(
|
||||
_meta_key(meta.user_id, meta.id),
|
||||
meta.model_dump(),
|
||||
)
|
||||
|
||||
|
||||
def transition_status(
|
||||
storage: StorageBackend,
|
||||
user_id: str,
|
||||
project_id: str,
|
||||
*,
|
||||
from_status: str | set[str] | frozenset[str],
|
||||
to_status: str,
|
||||
**fields: Any,
|
||||
) -> ProjectMeta:
|
||||
"""Move a project from ``from_status`` → ``to_status`` atomically.
|
||||
|
||||
Reads the meta with its GCS generation, refuses the write if the
|
||||
current status isn't in ``from_status``, then issues a conditional
|
||||
write that fails if another writer raced in. Retries up to a few
|
||||
times on generation mismatch caused by unrelated field updates.
|
||||
|
||||
Raises :class:`StatusConflict` when the current status doesn't match.
|
||||
"""
|
||||
allowed: frozenset[str]
|
||||
if isinstance(from_status, str):
|
||||
allowed = frozenset({from_status})
|
||||
else:
|
||||
allowed = frozenset(from_status)
|
||||
|
||||
last_exc: Exception | None = None
|
||||
for _ in range(5):
|
||||
meta, gen = _read_meta_with_generation(storage, user_id, project_id)
|
||||
if meta.status not in allowed:
|
||||
raise StatusConflict(
|
||||
f"project {project_id} is in status {meta.status!r}; "
|
||||
f"expected one of {sorted(allowed)} for transition to {to_status!r}"
|
||||
)
|
||||
meta.status = to_status
|
||||
for k, v in fields.items():
|
||||
setattr(meta, k, v)
|
||||
meta.updated = datetime.now(timezone.utc).isoformat()
|
||||
try:
|
||||
storage.write_json_if_match(
|
||||
_meta_key(user_id, project_id), meta.model_dump(), gen,
|
||||
)
|
||||
return meta
|
||||
except StaleGeneration as exc:
|
||||
last_exc = exc
|
||||
continue
|
||||
raise StatusConflict(
|
||||
f"project {project_id}: lost optimistic-concurrency race after retries"
|
||||
) from last_exc
|
||||
|
||||
|
||||
def request_cancel(
|
||||
storage: StorageBackend, user_id: str, project_id: str
|
||||
) -> ProjectMeta:
|
||||
"""Set ``cancel_requested = True`` so the worker's cancel gate trips.
|
||||
|
||||
Does not touch ``status`` — the worker is responsible for moving the
|
||||
project to ``cancelled`` when it observes the flag.
|
||||
"""
|
||||
return update_project(
|
||||
storage, user_id, project_id, cancel_requested=True,
|
||||
)
|
||||
|
||||
|
||||
def mark_stale_running(
|
||||
storage: StorageBackend, user_id: str, project_id: str, error: str,
|
||||
) -> ProjectMeta | None:
|
||||
"""Flip a stale ``running`` project to ``error``. No-op otherwise.
|
||||
|
||||
Returns the updated meta on success; ``None`` if the project's status
|
||||
was already terminal or the project no longer exists.
|
||||
"""
|
||||
try:
|
||||
return transition_status(
|
||||
storage, user_id, project_id,
|
||||
from_status={STATUS_RUNNING, STATUS_QUEUED},
|
||||
to_status=STATUS_ERROR,
|
||||
pipeline_state={"error": error},
|
||||
cancel_requested=False,
|
||||
)
|
||||
except StatusConflict:
|
||||
return None
|
||||
|
||||
|
||||
# --- CRUD ---
|
||||
|
||||
|
||||
def create_project(storage: StorageBackend, user_id: str, name: str) -> ProjectMeta:
|
||||
project_id = uuid.uuid4().hex[:12]
|
||||
meta = ProjectMeta(
|
||||
id=project_id,
|
||||
name=name,
|
||||
user_id=user_id,
|
||||
created=datetime.now(timezone.utc).isoformat(),
|
||||
)
|
||||
_write_meta(storage, meta)
|
||||
return meta
|
||||
|
||||
|
||||
def list_projects(storage: StorageBackend, user_id: str) -> list[ProjectMeta]:
|
||||
prefix = f"users/{user_id}/projects/"
|
||||
projects: list[ProjectMeta] = []
|
||||
for entry in storage.list_prefix(prefix):
|
||||
# entry is like users/{uid}/projects/{pid} (a directory)
|
||||
# or users/{uid}/projects/{pid}/project.json (a file)
|
||||
meta_key = f"{entry}/project.json" if not entry.endswith("/project.json") else entry
|
||||
if storage.exists(meta_key):
|
||||
data = storage.read_json(meta_key)
|
||||
projects.append(ProjectMeta.model_validate(data))
|
||||
return projects
|
||||
|
||||
|
||||
def get_project(
|
||||
storage: StorageBackend, user_id: str, project_id: str
|
||||
) -> ProjectMeta | None:
|
||||
key = _meta_key(user_id, project_id)
|
||||
if not storage.exists(key):
|
||||
return None
|
||||
return _read_meta(storage, user_id, project_id)
|
||||
|
||||
|
||||
def update_project(
|
||||
storage: StorageBackend, user_id: str, project_id: str, **fields: Any
|
||||
) -> ProjectMeta:
|
||||
if not storage.exists(_meta_key(user_id, project_id)):
|
||||
raise ProjectNotFound(project_id)
|
||||
meta = _read_meta(storage, user_id, project_id)
|
||||
for k, v in fields.items():
|
||||
setattr(meta, k, v)
|
||||
_write_meta(storage, meta)
|
||||
return meta
|
||||
|
||||
|
||||
def delete_project(
|
||||
storage: StorageBackend, user_id: str, project_id: str
|
||||
) -> bool:
|
||||
key = _meta_key(user_id, project_id)
|
||||
if not storage.exists(key):
|
||||
return False
|
||||
# Clean up shared references for all collaborators before deleting
|
||||
meta = _read_meta(storage, user_id, project_id)
|
||||
for collab_id in meta.collaborators:
|
||||
ref_key = _shared_ref_key(collab_id, project_id)
|
||||
if storage.exists(ref_key):
|
||||
storage.delete_key(ref_key)
|
||||
storage.delete_prefix(_project_prefix(user_id, project_id))
|
||||
return True
|
||||
|
||||
|
||||
def clear_project_extractions(
|
||||
storage: StorageBackend, user_id: str, project_id: str
|
||||
) -> None:
|
||||
"""Delete per-project extraction JSONs and derived artifacts.
|
||||
|
||||
Clears extracted/, patterns/, and models/ plus derived files (graph,
|
||||
power tree, BOM summary, derating, report, API logs) so the next
|
||||
pipeline run starts from fresh per-project data. The global library
|
||||
(library/*) is untouched — shared entries remain reusable.
|
||||
|
||||
Meta fields tied to the prior run (summary, skipped list, review
|
||||
checkpoint, error state) are reset; historical spend fields
|
||||
(total_cost_usd, credits_spent) are preserved.
|
||||
"""
|
||||
prefix = _project_prefix(user_id, project_id)
|
||||
for subdir in ("extracted", "patterns", "models"):
|
||||
storage.delete_prefix(f"{prefix}/{subdir}")
|
||||
for name in (
|
||||
"design_graph.json",
|
||||
"bom_summary.json",
|
||||
"derating.json",
|
||||
"report.json",
|
||||
"api_logs.jsonl",
|
||||
"graph_voltage_updates.json",
|
||||
):
|
||||
key = f"{prefix}/{name}"
|
||||
if storage.exists(key):
|
||||
storage.delete_key(key)
|
||||
update_project(
|
||||
storage, user_id, project_id,
|
||||
summary=None,
|
||||
skipped_components=None,
|
||||
pipeline_state=None,
|
||||
pause_checkpoint=None,
|
||||
pause_reason=None,
|
||||
completed_review_refs=[],
|
||||
)
|
||||
|
||||
|
||||
def reopen_project(
|
||||
storage: StorageBackend, user_id: str, project_id: str
|
||||
) -> ProjectMeta:
|
||||
"""Reset a finished/cancelled/errored project back to a draft-like state.
|
||||
|
||||
Clears derived artifacts (graph, report, etc.) and the pause/review
|
||||
bookkeeping so the next pipeline run starts fresh, but preserves uploads,
|
||||
column mappings, and the extraction cache so the rerun reuses prior work
|
||||
cheaply.
|
||||
"""
|
||||
prefix = _project_prefix(user_id, project_id)
|
||||
for name in (
|
||||
"design_graph.json",
|
||||
"bom_summary.json",
|
||||
"derating.json",
|
||||
"report.json",
|
||||
"api_logs.jsonl",
|
||||
"graph_voltage_updates.json",
|
||||
):
|
||||
key = f"{prefix}/{name}"
|
||||
if storage.exists(key):
|
||||
storage.delete_key(key)
|
||||
return update_project(
|
||||
storage, user_id, project_id,
|
||||
status="draft",
|
||||
summary=None,
|
||||
skipped_components=None,
|
||||
pipeline_state=None,
|
||||
pause_checkpoint=None,
|
||||
pause_reason=None,
|
||||
completed_review_refs=[],
|
||||
)
|
||||
|
||||
|
||||
def list_project_datasheets(
|
||||
storage: StorageBackend, user_id: str, project_id: str
|
||||
) -> list[str]:
|
||||
"""Return the safe-MPN stems of datasheet PDFs stored for a project."""
|
||||
ds_prefix = f"{_project_prefix(user_id, project_id)}/uploads/datasheets/"
|
||||
stems: list[str] = []
|
||||
for key in storage.list_prefix(ds_prefix):
|
||||
if key.endswith(".pdf"):
|
||||
stems.append(key.rsplit("/", 1)[-1][:-4])
|
||||
return stems
|
||||
|
||||
|
||||
# --- Collaborator access resolution ---
|
||||
|
||||
|
||||
def _shared_ref_key(user_id: str, project_id: str) -> str:
|
||||
return f"users/{user_id}/shared/{project_id}.json"
|
||||
|
||||
|
||||
def resolve_project_access(
|
||||
storage: StorageBackend, caller_user_id: str, project_id: str
|
||||
) -> tuple[str, ProjectMeta] | None:
|
||||
"""Resolve project access for a user — checks ownership then collaborator refs.
|
||||
|
||||
Returns (owner_user_id, ProjectMeta) or None if no access.
|
||||
"""
|
||||
# 1. Direct ownership
|
||||
meta = get_project(storage, caller_user_id, project_id)
|
||||
if meta is not None:
|
||||
return (caller_user_id, meta)
|
||||
|
||||
# 2. Shared reference
|
||||
ref_key = _shared_ref_key(caller_user_id, project_id)
|
||||
if not storage.exists(ref_key):
|
||||
return None
|
||||
ref = storage.read_json(ref_key)
|
||||
owner_id = ref.get("owner_user_id")
|
||||
if not owner_id:
|
||||
return None
|
||||
meta = get_project(storage, owner_id, project_id)
|
||||
if meta is None:
|
||||
return None
|
||||
# Verify caller is still in collaborators list
|
||||
if caller_user_id not in meta.collaborators:
|
||||
# Stale reference — clean up
|
||||
storage.delete_key(ref_key)
|
||||
return None
|
||||
return (owner_id, meta)
|
||||
|
||||
|
||||
def find_project_any_user(
|
||||
storage: StorageBackend, project_id: str
|
||||
) -> tuple[str, ProjectMeta] | None:
|
||||
"""Scan all users to find a project by ID (for admin access).
|
||||
|
||||
Returns (owner_user_id, ProjectMeta) or None.
|
||||
"""
|
||||
seen_uids: set[str] = set()
|
||||
for entry in storage.list_prefix("users/"):
|
||||
parts = entry.split("/")
|
||||
if len(parts) >= 2:
|
||||
uid = parts[1]
|
||||
if uid in seen_uids:
|
||||
continue
|
||||
seen_uids.add(uid)
|
||||
meta = get_project(storage, uid, project_id)
|
||||
if meta is not None:
|
||||
return (uid, meta)
|
||||
return None
|
||||
|
||||
|
||||
def add_collaborator(
|
||||
storage: StorageBackend, owner_user_id: str, project_id: str, collaborator_user_id: str
|
||||
) -> ProjectMeta:
|
||||
"""Add a collaborator to a project and write a shared reference."""
|
||||
meta = _read_meta(storage, owner_user_id, project_id)
|
||||
if collaborator_user_id not in meta.collaborators:
|
||||
meta.collaborators.append(collaborator_user_id)
|
||||
_write_meta(storage, meta)
|
||||
# Write reverse reference for the collaborator
|
||||
ref_key = _shared_ref_key(collaborator_user_id, project_id)
|
||||
storage.write_json(ref_key, {"owner_user_id": owner_user_id})
|
||||
return meta
|
||||
|
||||
|
||||
def remove_collaborator(
|
||||
storage: StorageBackend, owner_user_id: str, project_id: str, collaborator_user_id: str
|
||||
) -> ProjectMeta:
|
||||
"""Remove a collaborator from a project and delete the shared reference."""
|
||||
meta = _read_meta(storage, owner_user_id, project_id)
|
||||
meta.collaborators = [c for c in meta.collaborators if c != collaborator_user_id]
|
||||
_write_meta(storage, meta)
|
||||
# Delete reverse reference
|
||||
ref_key = _shared_ref_key(collaborator_user_id, project_id)
|
||||
if storage.exists(ref_key):
|
||||
storage.delete_key(ref_key)
|
||||
return meta
|
||||
|
||||
|
||||
def transfer_ownership(
|
||||
storage: StorageBackend,
|
||||
current_owner_user_id: str,
|
||||
project_id: str,
|
||||
new_owner_user_id: str,
|
||||
) -> ProjectMeta:
|
||||
"""Make an existing collaborator the new owner of a project.
|
||||
|
||||
Swaps roles: ``new_owner_user_id`` becomes the owner, the previous owner
|
||||
is appended to ``collaborators``. All project files are physically moved
|
||||
from ``users/{old}/projects/{id}/`` to ``users/{new}/projects/{id}/`` so
|
||||
that the storage layout (which keys off the owner) stays consistent.
|
||||
Shared references are rewritten — the new owner's ref is deleted, the
|
||||
old owner gets one, and every remaining collaborator's ref is repointed
|
||||
at the new owner.
|
||||
|
||||
Raises ``ValueError`` if the target is already the owner or is not a
|
||||
current collaborator.
|
||||
"""
|
||||
meta = _read_meta(storage, current_owner_user_id, project_id)
|
||||
|
||||
if new_owner_user_id == current_owner_user_id:
|
||||
raise ValueError("target user is already the owner")
|
||||
if new_owner_user_id not in meta.collaborators:
|
||||
raise ValueError("target user must currently be a collaborator")
|
||||
|
||||
new_collaborators = [c for c in meta.collaborators if c != new_owner_user_id]
|
||||
if current_owner_user_id not in new_collaborators:
|
||||
new_collaborators.append(current_owner_user_id)
|
||||
|
||||
meta.user_id = new_owner_user_id
|
||||
meta.collaborators = new_collaborators
|
||||
meta.updated = datetime.now(timezone.utc).isoformat()
|
||||
|
||||
old_prefix = _project_prefix(current_owner_user_id, project_id)
|
||||
new_prefix = _project_prefix(new_owner_user_id, project_id)
|
||||
new_meta_key = _meta_key(new_owner_user_id, project_id)
|
||||
|
||||
# Copy every file under the old prefix to the corresponding new key.
|
||||
# The old project.json is copied too — we overwrite it below with the
|
||||
# refreshed meta so the new location is authoritative even if a partial
|
||||
# failure leaves the old prefix in place.
|
||||
for old_key in storage.list_recursive(old_prefix):
|
||||
rel = old_key[len(old_prefix):].lstrip("/")
|
||||
storage.copy_object(old_key, f"{new_prefix}/{rel}")
|
||||
|
||||
storage.write_json(new_meta_key, meta.model_dump())
|
||||
storage.delete_prefix(old_prefix)
|
||||
|
||||
# Reverse references: new owner no longer needs one; old owner now does;
|
||||
# every other collaborator's existing ref must point at the new owner.
|
||||
new_owner_ref = _shared_ref_key(new_owner_user_id, project_id)
|
||||
if storage.exists(new_owner_ref):
|
||||
storage.delete_key(new_owner_ref)
|
||||
storage.write_json(
|
||||
_shared_ref_key(current_owner_user_id, project_id),
|
||||
{"owner_user_id": new_owner_user_id},
|
||||
)
|
||||
for collab_id in new_collaborators:
|
||||
if collab_id == current_owner_user_id:
|
||||
continue
|
||||
storage.write_json(
|
||||
_shared_ref_key(collab_id, project_id),
|
||||
{"owner_user_id": new_owner_user_id},
|
||||
)
|
||||
|
||||
return meta
|
||||
|
||||
|
||||
def list_shared_projects(storage: StorageBackend, user_id: str) -> list[ProjectMeta]:
|
||||
"""List projects shared with a user (where they are a collaborator)."""
|
||||
prefix = f"users/{user_id}/shared/"
|
||||
shared: list[ProjectMeta] = []
|
||||
for entry in storage.list_prefix(prefix):
|
||||
if not entry.endswith(".json"):
|
||||
continue
|
||||
try:
|
||||
ref = storage.read_json(entry)
|
||||
owner_id = ref.get("owner_user_id")
|
||||
if not owner_id:
|
||||
continue
|
||||
# Extract project_id from the key: users/{uid}/shared/{project_id}.json
|
||||
filename = entry.rsplit("/", 1)[-1]
|
||||
project_id = filename.replace(".json", "")
|
||||
meta = get_project(storage, owner_id, project_id)
|
||||
if meta and user_id in meta.collaborators:
|
||||
shared.append(meta)
|
||||
except Exception:
|
||||
continue
|
||||
return shared
|
||||
|
||||
|
||||
# --- File operations ---
|
||||
|
||||
|
||||
def save_bom(
|
||||
storage: StorageBackend, user_id: str, project_id: str, data: bytes
|
||||
) -> str:
|
||||
key = f"{_project_prefix(user_id, project_id)}/uploads/bom.csv"
|
||||
storage.write_bytes(key, data)
|
||||
update_project(storage, user_id, project_id, has_bom=True)
|
||||
return key
|
||||
|
||||
|
||||
_NETLIST_EXT = {"pads": "asc", "edif": "edn"}
|
||||
|
||||
|
||||
def _netlist_key(user_id: str, project_id: str, fmt: str) -> str:
|
||||
ext = _NETLIST_EXT.get(fmt, "asc")
|
||||
return f"{_project_prefix(user_id, project_id)}/uploads/netlist.{ext}"
|
||||
|
||||
|
||||
def save_netlist(
|
||||
storage: StorageBackend,
|
||||
user_id: str,
|
||||
project_id: str,
|
||||
data: bytes,
|
||||
*,
|
||||
fmt: str = "pads",
|
||||
) -> str:
|
||||
"""Persist the uploaded netlist with the extension matching ``fmt``.
|
||||
|
||||
Also clears any previously-saved netlist in the *other* format so we
|
||||
never have stale ``.asc`` and ``.edn`` files side-by-side (e.g. user
|
||||
re-uploads with a different format).
|
||||
"""
|
||||
key = _netlist_key(user_id, project_id, fmt)
|
||||
storage.write_bytes(key, data)
|
||||
other_fmt = "edif" if fmt == "pads" else "pads"
|
||||
other_key = _netlist_key(user_id, project_id, other_fmt)
|
||||
if storage.exists(other_key):
|
||||
storage.delete_key(other_key)
|
||||
# Reset sub-design selection on every upload — the prior selection may
|
||||
# reference IDs that no longer exist in the new file. Frontend resets
|
||||
# the picker after upload too; this keeps backend in sync.
|
||||
update_project(
|
||||
storage, user_id, project_id,
|
||||
has_netlist=True, netlist_format=fmt, netlist_subdesigns=None,
|
||||
)
|
||||
return key
|
||||
|
||||
|
||||
def save_datasheet(
|
||||
storage: StorageBackend, user_id: str, project_id: str, mpn: str, data: bytes
|
||||
) -> str:
|
||||
"""Save a datasheet PDF to the project uploads directory.
|
||||
|
||||
Library writes happen during pattern extraction (one PDF per pattern series).
|
||||
"""
|
||||
safe = safe_mpn(mpn)
|
||||
key = f"{_project_prefix(user_id, project_id)}/uploads/datasheets/{safe}.pdf"
|
||||
storage.write_bytes(key, data)
|
||||
# Count datasheets
|
||||
ds_prefix = f"{_project_prefix(user_id, project_id)}/uploads/datasheets/"
|
||||
count = sum(1 for k in storage.list_prefix(ds_prefix) if k.endswith(".pdf"))
|
||||
update_project(storage, user_id, project_id, datasheet_count=count)
|
||||
return key
|
||||
|
||||
|
||||
def get_bom_key(
|
||||
storage: StorageBackend, user_id: str, project_id: str
|
||||
) -> str | None:
|
||||
key = f"{_project_prefix(user_id, project_id)}/uploads/bom.csv"
|
||||
return key if storage.exists(key) else None
|
||||
|
||||
|
||||
def get_netlist_key(
|
||||
storage: StorageBackend, user_id: str, project_id: str
|
||||
) -> str | None:
|
||||
"""Return the storage key of whichever netlist file exists (.asc or .edn)."""
|
||||
for fmt in ("pads", "edif"):
|
||||
key = _netlist_key(user_id, project_id, fmt)
|
||||
if storage.exists(key):
|
||||
return key
|
||||
return None
|
||||
|
||||
|
||||
def get_datasheet_key(
|
||||
storage: StorageBackend, user_id: str, project_id: str, mpn: str
|
||||
) -> str | None:
|
||||
safe = safe_mpn(mpn)
|
||||
key = f"{_project_prefix(user_id, project_id)}/uploads/datasheets/{safe}.pdf"
|
||||
return key if storage.exists(key) else None
|
||||
|
||||
|
||||
def project_prefix(user_id: str, project_id: str) -> str:
|
||||
"""Return the storage prefix for a project (for use by pipeline/routers)."""
|
||||
return _project_prefix(user_id, project_id)
|
||||
|
||||
|
||||
# --- Library operations ---
|
||||
|
||||
|
||||
def library_has_extraction(
|
||||
storage: StorageBackend, mpn: str, min_version: str | None = None,
|
||||
) -> str | None:
|
||||
"""Check if library has a complete extraction (with pintable) for this MPN.
|
||||
|
||||
If *min_version* is set, also checks that the extraction's
|
||||
``model_version`` meets the minimum threshold.
|
||||
Returns the key if found and valid, None otherwise.
|
||||
"""
|
||||
safe = safe_mpn(mpn)
|
||||
key = f"library/extracted/{safe}.json"
|
||||
if not storage.exists(key):
|
||||
return None
|
||||
data = storage.read_json(key)
|
||||
if not data.get("pintable"):
|
||||
return None
|
||||
if min_version:
|
||||
from backend.services.admin_settings import version_is_stale
|
||||
|
||||
component_version = data.get("model_version", "0.0.0")
|
||||
if version_is_stale(component_version, min_version):
|
||||
return None
|
||||
return key
|
||||
|
||||
|
||||
def library_has_datasheet(
|
||||
storage: StorageBackend, mpn: str, patterns: list | None = None,
|
||||
) -> str | None:
|
||||
"""Check if library has a datasheet PDF for this MPN.
|
||||
|
||||
Checks content-addressed refs first, then falls back to legacy flat
|
||||
files (for pre-migration data), then pattern-based lookup.
|
||||
|
||||
Returns the storage key if found, None otherwise.
|
||||
"""
|
||||
from backend.services.datasheet_store import resolve_datasheet
|
||||
|
||||
# 1. Content-addressed ref lookup
|
||||
resolved = resolve_datasheet(storage, mpn)
|
||||
if resolved:
|
||||
return resolved
|
||||
# 2. Legacy flat file fallback (remove after migration confirmed)
|
||||
safe = safe_mpn(mpn)
|
||||
key = f"library/datasheets/{safe}.pdf"
|
||||
if storage.exists(key):
|
||||
return key
|
||||
# 3. Pattern-based fallback for passives
|
||||
if patterns:
|
||||
from backend.pinscopex.resolve_passives import resolve_mpn
|
||||
|
||||
match = resolve_mpn(mpn, patterns)
|
||||
if match is not None:
|
||||
pat = match[0]
|
||||
ds_key = pat.datasheet_key
|
||||
if ds_key and storage.exists(ds_key):
|
||||
return ds_key
|
||||
return None
|
||||
|
||||
|
||||
def library_has_model(storage: StorageBackend, mpn: str) -> str | None:
|
||||
"""Check if library has a ComponentModel (specs) for this MPN.
|
||||
|
||||
Returns the key if found, None otherwise.
|
||||
"""
|
||||
safe = safe_mpn(mpn)
|
||||
key = f"library/models/{safe}.json"
|
||||
return key if storage.exists(key) else None
|
||||
|
||||
|
||||
def library_has_passive_model(storage: StorageBackend, mpn: str) -> str | None:
|
||||
"""Check if library has a DigiKey-resolved passive model for this MPN.
|
||||
|
||||
Checks library/passives/ first, then falls back to library/models/
|
||||
for pre-migration data. Returns the key if found, None otherwise.
|
||||
"""
|
||||
safe = safe_mpn(mpn)
|
||||
key = f"library/passives/{safe}.json"
|
||||
if storage.exists(key):
|
||||
return key
|
||||
# Fallback: pre-migration passive specs may still be in library/models/
|
||||
legacy_key = f"library/models/{safe}.json"
|
||||
return legacy_key if storage.exists(legacy_key) else None
|
||||
|
||||
|
||||
def save_to_library(
|
||||
storage: StorageBackend, src_key: str, category: str, filename: str
|
||||
) -> str:
|
||||
"""Copy a file to the shared library."""
|
||||
dst_key = f"library/{category}/{filename}"
|
||||
storage.copy_object(src_key, dst_key)
|
||||
return dst_key
|
||||
|
||||
|
||||
def list_library_patterns(storage: StorageBackend) -> list[str]:
|
||||
"""List all pattern keys in the library."""
|
||||
prefix = "library/patterns/"
|
||||
return [k for k in storage.list_prefix(prefix) if k.endswith(".json")]
|
||||
|
||||
|
||||
def load_library_patterns(storage: StorageBackend):
|
||||
"""Load and parse all passive patterns from the library.
|
||||
|
||||
For local backend, delegates to pinscopex. For GCS, downloads to temp first.
|
||||
This function is only used by the library/check endpoint — during pipeline
|
||||
execution, patterns are loaded from the workspace temp directory.
|
||||
"""
|
||||
from backend.pinscopex.resolve_passives import load_patterns
|
||||
|
||||
from backend.services.storage import LocalStorageBackend
|
||||
|
||||
if isinstance(storage, LocalStorageBackend):
|
||||
d = storage._path("library/patterns")
|
||||
if not d.is_dir():
|
||||
return []
|
||||
return load_patterns(str(d))
|
||||
|
||||
# GCS: download patterns to a temp directory
|
||||
import tempfile
|
||||
|
||||
pattern_keys = list_library_patterns(storage)
|
||||
if not pattern_keys:
|
||||
return []
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
tmp_path = Path(tmpdir) / "patterns"
|
||||
tmp_path.mkdir()
|
||||
for key in pattern_keys:
|
||||
filename = key.rsplit("/", 1)[-1]
|
||||
storage.download_to_local(key, tmp_path / filename)
|
||||
return load_patterns(str(tmp_path))
|
||||
|
||||
|
||||
# Re-export for convenience
|
||||
from pathlib import Path # noqa: E402
|
||||
@@ -0,0 +1,334 @@
|
||||
"""Purple Parts API client — LCSC code → MPN resolution.
|
||||
|
||||
Wraps the external `purple-parts` HTTP service (a read-only API over the
|
||||
jlcparts/LCSC catalogue, deployed at the URL in `settings.purple_parts_url`).
|
||||
Used by the BOM-parse stage to convert LCSC codes (e.g. "C12345") into
|
||||
manufacturer part numbers before the DigiKey resolver runs.
|
||||
|
||||
The remote service is Cloud Run with IAM auth, so calls send a Google
|
||||
identity token (audience = purple_parts_url) plus an X-API-Key header. In
|
||||
Cloud Run the identity token is minted automatically via ADC + the
|
||||
metadata server; locally `fetch_id_token` only works if
|
||||
GOOGLE_APPLICATION_CREDENTIALS points at a service-account key file. On
|
||||
local dev with user creds the helper logs a debug line and the call is
|
||||
skipped (returns an empty result), which the caller treats as a no-op.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import re
|
||||
import time
|
||||
from typing import Optional
|
||||
|
||||
import httpx
|
||||
|
||||
from backend.config import settings
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_LCSC_RE = re.compile(r"^C\d+$", re.IGNORECASE)
|
||||
|
||||
# Identity tokens are valid for ~1h; refresh ~10 min early.
|
||||
_TOKEN_TTL_SECONDS = 50 * 60
|
||||
_token_cache: dict[str, float | str] = {"token": "", "expires_at": 0.0}
|
||||
_token_lock = asyncio.Lock()
|
||||
|
||||
# Conservative batch size — purple-parts accepts up to 500 per request.
|
||||
_BATCH_SIZE = 400
|
||||
|
||||
|
||||
def is_lcsc_code(value: str | None) -> bool:
|
||||
"""Return True if `value` looks like an LCSC part number (e.g. C12345)."""
|
||||
if not value:
|
||||
return False
|
||||
return bool(_LCSC_RE.match(value.strip()))
|
||||
|
||||
|
||||
async def _get_identity_token() -> str | None:
|
||||
"""Mint a Google ID token for the purple-parts audience, cached.
|
||||
|
||||
Returns None when credentials don't support identity-token minting
|
||||
(typical for local dev with `gcloud auth application-default login` user
|
||||
creds). Caller should treat None as "skip the purple-parts call."
|
||||
"""
|
||||
now = time.time()
|
||||
cached = _token_cache.get("token", "")
|
||||
if cached and float(_token_cache.get("expires_at", 0.0)) > now:
|
||||
return str(cached)
|
||||
|
||||
async with _token_lock:
|
||||
cached = _token_cache.get("token", "")
|
||||
if cached and float(_token_cache.get("expires_at", 0.0)) > now:
|
||||
return str(cached)
|
||||
|
||||
try:
|
||||
from google.auth.transport.requests import Request
|
||||
from google.oauth2 import id_token as gid_token
|
||||
except ImportError:
|
||||
logger.warning("google-auth not installed; purple-parts disabled")
|
||||
return None
|
||||
|
||||
loop = asyncio.get_running_loop()
|
||||
try:
|
||||
token = await loop.run_in_executor(
|
||||
None,
|
||||
lambda: gid_token.fetch_id_token(Request(), settings.purple_parts_url),
|
||||
)
|
||||
except Exception as e:
|
||||
logger.debug(
|
||||
"purple-parts: identity-token mint failed (%s: %s) — "
|
||||
"expected for local user creds, skipping",
|
||||
type(e).__name__, e,
|
||||
)
|
||||
return None
|
||||
|
||||
_token_cache["token"] = token
|
||||
_token_cache["expires_at"] = now + _TOKEN_TTL_SECONDS
|
||||
return token
|
||||
|
||||
|
||||
def detect_lcsc_column(csv_bytes: bytes, mpn_col: str) -> bool:
|
||||
"""Return True when every non-empty value in `mpn_col` matches `^C\\d+$`.
|
||||
|
||||
Used by the upload endpoint to auto-detect when the user's chosen MPN
|
||||
column is actually an LCSC column (i.e. the user pasted LCSC ids into
|
||||
the MPN slot, or labeled their LCSC column as "Manufacturer Part Number").
|
||||
Column-level — a single non-LCSC entry disqualifies the column so that
|
||||
BOMs mixing real MPNs with LCSC ids aren't silently mangled.
|
||||
"""
|
||||
import csv as csv_mod
|
||||
import io
|
||||
|
||||
text = csv_bytes.decode("utf-8", errors="replace")
|
||||
reader = csv_mod.DictReader(io.StringIO(text))
|
||||
if not reader.fieldnames or mpn_col not in reader.fieldnames:
|
||||
return False
|
||||
|
||||
seen_any = False
|
||||
for row in reader:
|
||||
val = (row.get(mpn_col) or "").strip()
|
||||
if not val:
|
||||
continue
|
||||
if not is_lcsc_code(val):
|
||||
return False
|
||||
seen_any = True
|
||||
return seen_any
|
||||
|
||||
|
||||
async def resolve_lcsc_column_bytes(
|
||||
csv_bytes: bytes,
|
||||
*,
|
||||
mpn_col: str = "Manufacturer Part Number",
|
||||
) -> tuple[bytes, int, dict[str, str], dict[str, dict]]:
|
||||
"""Replace every value in `mpn_col` with the manufacturer part number
|
||||
resolved via purple-parts.
|
||||
|
||||
Returns `(new_csv_bytes, rows_updated, lcsc_to_mpn_map, lcsc_payloads_map)`.
|
||||
The first map is keyed by LCSC id (e.g. "C12044") → resolved MPN string,
|
||||
so the caller can surface "C12044 → STM32F103C8T6" in the UI. The second
|
||||
map is keyed by the same LCSC id → the full purple-parts payload (mpn,
|
||||
manufacturer, package, description, category, subcategory, ...) so the
|
||||
caller can cache it on the project for the wizard's per-row resolve
|
||||
endpoint. Preserves column order, headers, and untouched cells. No-op
|
||||
when purple-parts isn't configured.
|
||||
"""
|
||||
import csv as csv_mod
|
||||
import io
|
||||
|
||||
if not settings.use_purple_parts:
|
||||
return csv_bytes, 0, {}, {}
|
||||
|
||||
text = csv_bytes.decode("utf-8", errors="replace")
|
||||
reader = csv_mod.DictReader(io.StringIO(text))
|
||||
fieldnames = reader.fieldnames or []
|
||||
rows = list(reader)
|
||||
|
||||
if not rows or mpn_col not in fieldnames:
|
||||
return csv_bytes, 0, {}, {}
|
||||
|
||||
todo: list[tuple[int, str]] = []
|
||||
for i, row in enumerate(rows):
|
||||
code = (row.get(mpn_col) or "").strip()
|
||||
if is_lcsc_code(code):
|
||||
todo.append((i, code))
|
||||
|
||||
if not todo:
|
||||
return csv_bytes, 0, {}, {}
|
||||
|
||||
unique_codes = sorted({c for _, c in todo})
|
||||
resolved = await lookup_lcsc_batch(unique_codes)
|
||||
|
||||
updated = 0
|
||||
lcsc_to_mpn: dict[str, str] = {}
|
||||
lcsc_payloads: dict[str, dict] = {}
|
||||
for i, code in todo:
|
||||
part = resolved.get(code)
|
||||
if part and part.get("mpn"):
|
||||
rows[i][mpn_col] = part["mpn"]
|
||||
lcsc_to_mpn[code] = part["mpn"]
|
||||
lcsc_payloads[code] = dict(part)
|
||||
updated += 1
|
||||
|
||||
if updated == 0:
|
||||
return csv_bytes, 0, {}, {}
|
||||
|
||||
out = io.StringIO()
|
||||
writer = csv_mod.DictWriter(out, fieldnames=fieldnames, extrasaction="ignore")
|
||||
writer.writeheader()
|
||||
writer.writerows(rows)
|
||||
return out.getvalue().encode("utf-8"), updated, lcsc_to_mpn, lcsc_payloads
|
||||
|
||||
|
||||
async def lookup_lcsc_batch(lcsc_codes: list[str]) -> dict[str, Optional[dict]]:
|
||||
"""Batch LCSC → MPN lookup.
|
||||
|
||||
Returns `{lcsc_code: part_dict_or_None}` for every code in input. Misses,
|
||||
invalid codes, and (after warning) total failures all return None values
|
||||
so the caller can treat the result as a uniform per-code map. The pipeline
|
||||
never aborts on a purple-parts miss; the row simply stays unresolved and
|
||||
the existing DigiKey/Haiku paths handle it.
|
||||
|
||||
Part dict shape: {lcsc, mpn, manufacturer, package, description, stock,
|
||||
basic, preferred}.
|
||||
"""
|
||||
if not settings.use_purple_parts:
|
||||
return {c: None for c in lcsc_codes}
|
||||
|
||||
codes = [c for c in (raw.strip() for raw in lcsc_codes) if c]
|
||||
if not codes:
|
||||
return {}
|
||||
|
||||
token = await _get_identity_token()
|
||||
if token is None:
|
||||
logger.info("purple-parts: no identity token, skipping batch of %d", len(codes))
|
||||
return {c: None for c in codes}
|
||||
|
||||
base_url = settings.purple_parts_url.rstrip("/")
|
||||
headers = {
|
||||
"Authorization": f"Bearer {token}",
|
||||
"X-API-Key": settings.purple_parts_api_key,
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
results: dict[str, Optional[dict]] = {c: None for c in codes}
|
||||
|
||||
async with httpx.AsyncClient(timeout=15) as client:
|
||||
for i in range(0, len(codes), _BATCH_SIZE):
|
||||
chunk = codes[i:i + _BATCH_SIZE]
|
||||
try:
|
||||
resp = await client.post(
|
||||
f"{base_url}/v1/parts/by-lcsc/batch",
|
||||
headers=headers,
|
||||
json={"ids": chunk},
|
||||
)
|
||||
resp.raise_for_status()
|
||||
body = resp.json()
|
||||
except httpx.HTTPStatusError as e:
|
||||
logger.warning(
|
||||
"purple-parts: batch call failed %s for chunk of %d",
|
||||
e.response.status_code, len(chunk),
|
||||
)
|
||||
continue
|
||||
except Exception as e:
|
||||
logger.warning("purple-parts: batch call error: %s", e)
|
||||
continue
|
||||
|
||||
for code, part in (body.get("results") or {}).items():
|
||||
results[code] = part
|
||||
|
||||
return results
|
||||
|
||||
|
||||
def _norm_mpn(value: str | None) -> str:
|
||||
"""Normalize an MPN for comparison: drop whitespace, uppercase."""
|
||||
return "".join((value or "").split()).upper()
|
||||
|
||||
|
||||
def _pick_exact(query: str, candidates: list[dict]) -> Optional[dict]:
|
||||
"""Return the candidate whose ``mpn`` exactly matches ``query``.
|
||||
|
||||
Match is case- and whitespace-insensitive. purple-parts' ``by-mpn``
|
||||
endpoint returns exact matches first and then prefix matches, but we
|
||||
re-check rather than trust ordering — a prefix-only hit (e.g. a series
|
||||
family for a more specific MPN) must be treated as a miss so it can't
|
||||
pollute the shared passive library. Mirrors the exact-MPN discipline of
|
||||
``services.digikey._find_product``.
|
||||
"""
|
||||
q = _norm_mpn(query)
|
||||
for part in candidates:
|
||||
if part and _norm_mpn(part.get("mpn")) == q:
|
||||
return part
|
||||
return None
|
||||
|
||||
|
||||
async def lookup_mpn_batch(mpns: list[str]) -> dict[str, Optional[dict]]:
|
||||
"""Reverse lookup: manufacturer part number → LCSC catalogue record.
|
||||
|
||||
Fans the unique MPNs out to purple-parts' batch endpoint
|
||||
(``POST /v1/parts/by-mpn/batch``) in chunks of ``_BATCH_SIZE`` — one indexed
|
||||
query per chunk instead of a GET per MPN, which is what stalled huge-BOM
|
||||
uploads when the by-mpn query was seq-scanning. Returns
|
||||
``{mpn: part_dict_or_None}`` keyed by the *input* MPN string.
|
||||
|
||||
The endpoint is exact-match only, and we additionally run :func:`_pick_exact`
|
||||
over each MPN's candidate list (case/whitespace-insensitive) to keep the
|
||||
exact-MPN discipline — a prefix / family hit can carry the wrong
|
||||
voltage / dielectric / package and must never reach the shared
|
||||
``library/passives``. Misses, missing creds (no identity token), and per-chunk
|
||||
failures all come back as ``None`` so the caller can treat the map uniformly.
|
||||
No-op (all ``None``) when purple-parts isn't configured.
|
||||
|
||||
Part dict shape matches :func:`lookup_lcsc_batch`: {lcsc, mpn, manufacturer,
|
||||
package, description, category, subcategory, stock, basic, preferred}.
|
||||
"""
|
||||
if not settings.use_purple_parts:
|
||||
return {m: None for m in mpns}
|
||||
|
||||
# Preserve input keys but query each unique, non-empty MPN once.
|
||||
names = list(dict.fromkeys(m.strip() for m in mpns if m and m.strip()))
|
||||
if not names:
|
||||
return {}
|
||||
|
||||
token = await _get_identity_token()
|
||||
if token is None:
|
||||
logger.info("purple-parts: no identity token, skipping by-mpn batch of %d", len(names))
|
||||
return {m: None for m in names}
|
||||
|
||||
base_url = settings.purple_parts_url.rstrip("/")
|
||||
headers = {
|
||||
"Authorization": f"Bearer {token}",
|
||||
"X-API-Key": settings.purple_parts_api_key,
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
results: dict[str, Optional[dict]] = {m: None for m in names}
|
||||
|
||||
async with httpx.AsyncClient(timeout=15) as client:
|
||||
for i in range(0, len(names), _BATCH_SIZE):
|
||||
chunk = names[i:i + _BATCH_SIZE]
|
||||
try:
|
||||
resp = await client.post(
|
||||
f"{base_url}/v1/parts/by-mpn/batch",
|
||||
headers=headers,
|
||||
json={"mpns": chunk},
|
||||
)
|
||||
resp.raise_for_status()
|
||||
body = resp.json()
|
||||
except httpx.HTTPStatusError as e:
|
||||
logger.warning(
|
||||
"purple-parts: by-mpn batch call failed %s for chunk of %d",
|
||||
e.response.status_code, len(chunk),
|
||||
)
|
||||
continue
|
||||
except Exception as e:
|
||||
msg = str(e) or type(e).__name__
|
||||
logger.warning("purple-parts: by-mpn batch call error: %s", msg)
|
||||
continue
|
||||
|
||||
# Each MPN maps to a candidate list; keep only the exact match.
|
||||
for mpn, candidates in (body.get("results") or {}).items():
|
||||
results[mpn] = _pick_exact(mpn, candidates or [])
|
||||
|
||||
return results
|
||||
@@ -0,0 +1,264 @@
|
||||
"""Storage abstraction layer.
|
||||
|
||||
Provides a StorageBackend protocol with two implementations:
|
||||
- LocalStorageBackend: maps GCS-style keys to local filesystem paths (dev/test)
|
||||
- GCSStorageBackend: uses Google Cloud Storage (production)
|
||||
|
||||
Keys use forward-slash-separated paths like GCS object names:
|
||||
users/{user_id}/projects/{project_id}/project.json
|
||||
library/extracted/{safe_mpn}.json
|
||||
taxonomy/ic.json
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import shutil
|
||||
import threading
|
||||
from pathlib import Path
|
||||
from typing import Protocol, runtime_checkable
|
||||
|
||||
|
||||
# Sentinel used by conditional writes to require that the object does not yet
|
||||
# exist (matches GCS ``if_generation_match=0`` semantics).
|
||||
GENERATION_NEW = 0
|
||||
|
||||
|
||||
class StaleGeneration(Exception):
|
||||
"""Raised when a conditional write loses an optimistic-concurrency race."""
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class StorageBackend(Protocol):
|
||||
"""Abstract storage interface used by all backend services."""
|
||||
|
||||
def read_json(self, key: str) -> dict:
|
||||
"""Read and parse a JSON object."""
|
||||
...
|
||||
|
||||
def write_json(self, key: str, data: dict) -> None:
|
||||
"""Serialize and write a JSON object."""
|
||||
...
|
||||
|
||||
def read_bytes(self, key: str) -> bytes:
|
||||
"""Read raw bytes."""
|
||||
...
|
||||
|
||||
def write_bytes(self, key: str, data: bytes) -> None:
|
||||
"""Write raw bytes."""
|
||||
...
|
||||
|
||||
def read_text(self, key: str) -> str:
|
||||
"""Read as UTF-8 text."""
|
||||
...
|
||||
|
||||
def write_text(self, key: str, text: str) -> None:
|
||||
"""Write UTF-8 text."""
|
||||
...
|
||||
|
||||
def exists(self, key: str) -> bool:
|
||||
"""Check if an object exists."""
|
||||
...
|
||||
|
||||
def list_prefix(self, prefix: str) -> list[str]:
|
||||
"""List all keys under a prefix (non-recursive by default).
|
||||
|
||||
Returns keys that are direct children of the prefix — i.e. one level
|
||||
deep. For example, listing ``users/abc/projects/`` returns keys like
|
||||
``users/abc/projects/p1/project.json`` but NOT keys nested further.
|
||||
|
||||
To list all keys recursively, use list_recursive().
|
||||
"""
|
||||
...
|
||||
|
||||
def list_recursive(self, prefix: str) -> list[str]:
|
||||
"""List all keys under a prefix, recursively."""
|
||||
...
|
||||
|
||||
def list_prefix_after(self, prefix: str, after_key: str | None = None) -> list[str]:
|
||||
"""List keys under ``prefix`` whose name lexicographically follows
|
||||
``after_key``. Used by the GCS-backed event tail (worker writes one
|
||||
object per event with a zero-padded sequence number; the SSE
|
||||
consumer pages through new files only).
|
||||
"""
|
||||
...
|
||||
|
||||
def read_json_with_generation(self, key: str) -> tuple[dict, int]:
|
||||
"""Read JSON and return ``(data, generation)``.
|
||||
|
||||
``generation`` is an opaque token that callers pass back to
|
||||
``write_json_if_match`` to detect lost-update races.
|
||||
"""
|
||||
...
|
||||
|
||||
def write_json_if_match(self, key: str, data: dict, generation: int) -> int:
|
||||
"""Write JSON only if the current generation equals ``generation``.
|
||||
|
||||
Pass ``GENERATION_NEW`` (0) to require that the key does not exist.
|
||||
Returns the new generation. Raises :class:`StaleGeneration` when the
|
||||
precondition fails (loser of a race).
|
||||
"""
|
||||
...
|
||||
|
||||
def delete_key(self, key: str) -> None:
|
||||
"""Delete a single object."""
|
||||
...
|
||||
|
||||
def delete_prefix(self, prefix: str) -> None:
|
||||
"""Delete all objects under a prefix (recursive)."""
|
||||
...
|
||||
|
||||
def copy_object(self, src_key: str, dst_key: str) -> None:
|
||||
"""Copy an object from src to dst."""
|
||||
...
|
||||
|
||||
def download_to_local(self, key: str, local_path: Path) -> Path:
|
||||
"""Download an object to a local file path. Returns the local path."""
|
||||
...
|
||||
|
||||
def upload_from_local(self, local_path: Path, key: str) -> None:
|
||||
"""Upload a local file to storage."""
|
||||
...
|
||||
|
||||
def signed_url(self, key: str, expiration_minutes: int = 15) -> str:
|
||||
"""Generate a time-limited URL for direct access to an object.
|
||||
|
||||
For LocalStorageBackend, returns a backend-proxied URL.
|
||||
For GCS, returns a signed GCS URL.
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
class LocalStorageBackend:
|
||||
"""Maps GCS-style keys to local filesystem paths under a base directory.
|
||||
|
||||
Key ``users/abc/projects/p1/project.json`` becomes
|
||||
``{base_dir}/users/abc/projects/p1/project.json``.
|
||||
"""
|
||||
|
||||
def __init__(self, base_dir: Path) -> None:
|
||||
self._base = base_dir
|
||||
# In-memory generation counter for optimistic-concurrency parity with
|
||||
# GCS. Single-process only; subprocess-based local workers run in a
|
||||
# different process and will collide on the meta key. The local
|
||||
# subprocess path is dev-only and rarely concurrent, so we accept it.
|
||||
self._generations: dict[str, int] = {}
|
||||
self._gen_lock = threading.Lock()
|
||||
|
||||
def _path(self, key: str) -> Path:
|
||||
return self._base / key
|
||||
|
||||
def read_json(self, key: str) -> dict:
|
||||
return json.loads(self._path(key).read_text())
|
||||
|
||||
def write_json(self, key: str, data: dict) -> None:
|
||||
p = self._path(key)
|
||||
p.parent.mkdir(parents=True, exist_ok=True)
|
||||
p.write_text(json.dumps(data, indent=2) + "\n")
|
||||
|
||||
def read_bytes(self, key: str) -> bytes:
|
||||
return self._path(key).read_bytes()
|
||||
|
||||
def write_bytes(self, key: str, data: bytes) -> None:
|
||||
p = self._path(key)
|
||||
p.parent.mkdir(parents=True, exist_ok=True)
|
||||
p.write_bytes(data)
|
||||
|
||||
def read_text(self, key: str) -> str:
|
||||
return self._path(key).read_text()
|
||||
|
||||
def write_text(self, key: str, text: str) -> None:
|
||||
p = self._path(key)
|
||||
p.parent.mkdir(parents=True, exist_ok=True)
|
||||
p.write_text(text)
|
||||
|
||||
def exists(self, key: str) -> bool:
|
||||
return self._path(key).is_file()
|
||||
|
||||
def list_prefix(self, prefix: str) -> list[str]:
|
||||
d = self._path(prefix)
|
||||
if not d.is_dir():
|
||||
return []
|
||||
keys: list[str] = []
|
||||
for child in sorted(d.iterdir()):
|
||||
rel = child.relative_to(self._base)
|
||||
keys.append(str(rel))
|
||||
return keys
|
||||
|
||||
def list_recursive(self, prefix: str) -> list[str]:
|
||||
d = self._path(prefix)
|
||||
if not d.is_dir():
|
||||
return []
|
||||
keys: list[str] = []
|
||||
for child in sorted(d.rglob("*")):
|
||||
if child.is_file():
|
||||
rel = child.relative_to(self._base)
|
||||
keys.append(str(rel))
|
||||
return keys
|
||||
|
||||
def list_prefix_after(self, prefix: str, after_key: str | None = None) -> list[str]:
|
||||
d = self._path(prefix)
|
||||
if not d.is_dir():
|
||||
return []
|
||||
keys: list[str] = []
|
||||
for child in sorted(d.iterdir()):
|
||||
if not child.is_file():
|
||||
continue
|
||||
rel = str(child.relative_to(self._base))
|
||||
if after_key is not None and rel <= after_key:
|
||||
continue
|
||||
keys.append(rel)
|
||||
return keys
|
||||
|
||||
def read_json_with_generation(self, key: str) -> tuple[dict, int]:
|
||||
data = json.loads(self._path(key).read_text())
|
||||
with self._gen_lock:
|
||||
gen = self._generations.get(key, 1)
|
||||
return data, gen
|
||||
|
||||
def write_json_if_match(self, key: str, data: dict, generation: int) -> int:
|
||||
p = self._path(key)
|
||||
with self._gen_lock:
|
||||
current = self._generations.get(key, 0 if not p.is_file() else 1)
|
||||
if generation != current:
|
||||
raise StaleGeneration(
|
||||
f"generation mismatch on {key}: expected {generation}, current {current}"
|
||||
)
|
||||
new_gen = current + 1
|
||||
p.parent.mkdir(parents=True, exist_ok=True)
|
||||
p.write_text(json.dumps(data, indent=2) + "\n")
|
||||
self._generations[key] = new_gen
|
||||
return new_gen
|
||||
|
||||
def delete_key(self, key: str) -> None:
|
||||
p = self._path(key)
|
||||
if p.is_file():
|
||||
p.unlink()
|
||||
with self._gen_lock:
|
||||
self._generations.pop(key, None)
|
||||
|
||||
def delete_prefix(self, prefix: str) -> None:
|
||||
d = self._path(prefix)
|
||||
if d.is_dir():
|
||||
shutil.rmtree(d)
|
||||
|
||||
def copy_object(self, src_key: str, dst_key: str) -> None:
|
||||
src = self._path(src_key)
|
||||
dst = self._path(dst_key)
|
||||
dst.parent.mkdir(parents=True, exist_ok=True)
|
||||
shutil.copy2(src, dst)
|
||||
|
||||
def download_to_local(self, key: str, local_path: Path) -> Path:
|
||||
src = self._path(key)
|
||||
local_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
shutil.copy2(src, local_path)
|
||||
return local_path
|
||||
|
||||
def upload_from_local(self, local_path: Path, key: str) -> None:
|
||||
dst = self._path(key)
|
||||
dst.parent.mkdir(parents=True, exist_ok=True)
|
||||
shutil.copy2(local_path, dst)
|
||||
|
||||
def signed_url(self, key: str, expiration_minutes: int = 15) -> str:
|
||||
# Local dev: return a path that the backend can serve directly
|
||||
return f"/api/datasheets/_local/{key}"
|
||||
@@ -0,0 +1,126 @@
|
||||
"""Google Cloud Storage backend for StorageBackend protocol."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
from google.api_core.exceptions import PreconditionFailed
|
||||
from google.cloud import storage as gcs
|
||||
|
||||
from backend.services.storage import StaleGeneration
|
||||
|
||||
|
||||
class GCSStorageBackend:
|
||||
"""StorageBackend implementation using Google Cloud Storage."""
|
||||
|
||||
def __init__(self, bucket_name: str) -> None:
|
||||
self._client = gcs.Client()
|
||||
self._bucket = self._client.bucket(bucket_name)
|
||||
|
||||
def _blob(self, key: str) -> gcs.Blob:
|
||||
return self._bucket.blob(key)
|
||||
|
||||
def read_json(self, key: str) -> dict:
|
||||
text = self._blob(key).download_as_text()
|
||||
return json.loads(text)
|
||||
|
||||
def write_json(self, key: str, data: dict) -> None:
|
||||
text = json.dumps(data, indent=2) + "\n"
|
||||
self._blob(key).upload_from_string(text, content_type="application/json")
|
||||
|
||||
def read_bytes(self, key: str) -> bytes:
|
||||
return self._blob(key).download_as_bytes()
|
||||
|
||||
def write_bytes(self, key: str, data: bytes) -> None:
|
||||
self._blob(key).upload_from_string(data)
|
||||
|
||||
def read_text(self, key: str) -> str:
|
||||
return self._blob(key).download_as_text()
|
||||
|
||||
def write_text(self, key: str, text: str) -> None:
|
||||
self._blob(key).upload_from_string(text, content_type="text/plain")
|
||||
|
||||
def exists(self, key: str) -> bool:
|
||||
return self._blob(key).exists()
|
||||
|
||||
def list_prefix(self, prefix: str) -> list[str]:
|
||||
# List immediate children (one level) using delimiter
|
||||
blobs = self._client.list_blobs(
|
||||
self._bucket, prefix=prefix, delimiter="/",
|
||||
)
|
||||
keys: list[str] = []
|
||||
# Files directly under prefix
|
||||
for blob in blobs:
|
||||
keys.append(blob.name)
|
||||
# "Subdirectories" — strip trailing slash for consistency
|
||||
for pfx in blobs.prefixes:
|
||||
keys.append(pfx.rstrip("/"))
|
||||
return sorted(keys)
|
||||
|
||||
def list_recursive(self, prefix: str) -> list[str]:
|
||||
blobs = self._client.list_blobs(self._bucket, prefix=prefix)
|
||||
return sorted(blob.name for blob in blobs)
|
||||
|
||||
def list_prefix_after(self, prefix: str, after_key: str | None = None) -> list[str]:
|
||||
# Use GCS ``start_offset`` to skip already-seen keys server-side. We
|
||||
# ask for the next-after value; since after_key may be the last seen
|
||||
# key, advance one byte so the listing excludes it.
|
||||
kwargs: dict = {"prefix": prefix, "delimiter": "/"}
|
||||
if after_key is not None:
|
||||
# Request keys strictly greater than after_key. Append a NUL byte
|
||||
# so GCS treats start_offset as "after" rather than "starting at".
|
||||
kwargs["start_offset"] = after_key + "\x00"
|
||||
blobs = self._client.list_blobs(self._bucket, **kwargs)
|
||||
return sorted(blob.name for blob in blobs)
|
||||
|
||||
def read_json_with_generation(self, key: str) -> tuple[dict, int]:
|
||||
blob = self._blob(key)
|
||||
text = blob.download_as_text()
|
||||
# download_as_text populates blob.generation as a side effect.
|
||||
gen = int(blob.generation) if blob.generation is not None else 0
|
||||
return json.loads(text), gen
|
||||
|
||||
def write_json_if_match(self, key: str, data: dict, generation: int) -> int:
|
||||
text = json.dumps(data, indent=2) + "\n"
|
||||
blob = self._blob(key)
|
||||
try:
|
||||
blob.upload_from_string(
|
||||
text,
|
||||
content_type="application/json",
|
||||
if_generation_match=generation,
|
||||
)
|
||||
except PreconditionFailed as exc:
|
||||
raise StaleGeneration(
|
||||
f"generation mismatch on {key}: expected {generation}"
|
||||
) from exc
|
||||
# blob.generation is set by upload_from_string on success.
|
||||
return int(blob.generation) if blob.generation is not None else 0
|
||||
|
||||
def delete_key(self, key: str) -> None:
|
||||
blob = self._blob(key)
|
||||
if blob.exists():
|
||||
blob.delete()
|
||||
|
||||
def delete_prefix(self, prefix: str) -> None:
|
||||
blobs = list(self._client.list_blobs(self._bucket, prefix=prefix))
|
||||
if blobs:
|
||||
self._bucket.delete_blobs(blobs)
|
||||
|
||||
def copy_object(self, src_key: str, dst_key: str) -> None:
|
||||
src_blob = self._blob(src_key)
|
||||
self._bucket.copy_blob(src_blob, self._bucket, dst_key)
|
||||
|
||||
def download_to_local(self, key: str, local_path: Path) -> Path:
|
||||
local_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
self._blob(key).download_to_filename(str(local_path))
|
||||
return local_path
|
||||
|
||||
def upload_from_local(self, local_path: Path, key: str) -> None:
|
||||
self._blob(key).upload_from_filename(str(local_path))
|
||||
|
||||
def signed_url(self, key: str, expiration_minutes: int = 15) -> str:
|
||||
# Not used for GCS on Cloud Run — the backend proxies PDFs directly
|
||||
# via the /datasheet-proxy/ endpoint instead. Kept for interface
|
||||
# compatibility.
|
||||
raise NotImplementedError("Use read_bytes() and proxy instead")
|
||||
@@ -0,0 +1,98 @@
|
||||
"""Onboarding survey — appends responses to a Google Sheet and tracks completion."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from backend.config import settings
|
||||
from backend.services.storage import StorageBackend
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_SURVEY_PREFIX = "admin/survey/"
|
||||
|
||||
|
||||
def _status_key(user_id: str) -> str:
|
||||
return f"{_SURVEY_PREFIX}{user_id}.json"
|
||||
|
||||
|
||||
def is_completed(storage: StorageBackend, user_id: str) -> bool:
|
||||
return storage.exists(_status_key(user_id))
|
||||
|
||||
|
||||
def _mark_completed(storage: StorageBackend, user_id: str) -> None:
|
||||
payload = {"completed": True, "timestamp": datetime.now(timezone.utc).isoformat()}
|
||||
storage.write_json(_status_key(user_id), payload)
|
||||
|
||||
|
||||
def _build_sheets_service():
|
||||
"""Build an authenticated Google Sheets API service.
|
||||
|
||||
On Cloud Run, google.auth.default() returns Compute Engine credentials
|
||||
which are auto-scoped. We just need the Sheets API enabled in the GCP
|
||||
project and the service account shared on the sheet.
|
||||
"""
|
||||
try:
|
||||
import google.auth
|
||||
from googleapiclient.discovery import build
|
||||
except ImportError:
|
||||
logger.warning("google-api-python-client not installed; survey sheet disabled")
|
||||
return None
|
||||
|
||||
try:
|
||||
credentials, project = google.auth.default()
|
||||
logger.debug("Sheets: credentials type=%s project=%s", type(credentials).__name__, project)
|
||||
# Compute Engine credentials don't need explicit scopes — they use
|
||||
# the access scopes set on the instance (which default to cloud-platform).
|
||||
# For user/SA key credentials, we need to scope them.
|
||||
if hasattr(credentials, "with_scopes"):
|
||||
credentials = credentials.with_scopes(
|
||||
["https://www.googleapis.com/auth/spreadsheets"]
|
||||
)
|
||||
return build("sheets", "v4", credentials=credentials, cache_discovery=False)
|
||||
except Exception:
|
||||
logger.exception("Could not build Sheets service")
|
||||
return None
|
||||
|
||||
|
||||
async def append_to_sheet(
|
||||
user_id: str,
|
||||
email: str,
|
||||
name: str,
|
||||
referral_source: str,
|
||||
user_profile: str,
|
||||
) -> bool:
|
||||
"""Append a survey row to the configured Google Sheet. Returns True on success."""
|
||||
sheet_id = settings.survey_sheet_id
|
||||
if not sheet_id:
|
||||
logger.warning("SURVEY_SHEET_ID not set; skipping sheet append for user %s", user_id)
|
||||
return False
|
||||
|
||||
service = _build_sheets_service()
|
||||
if not service:
|
||||
return False
|
||||
|
||||
timestamp = datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S UTC")
|
||||
row = [timestamp, user_id, email, name, referral_source, user_profile]
|
||||
|
||||
try:
|
||||
await asyncio.to_thread(
|
||||
service.spreadsheets()
|
||||
.values()
|
||||
.append(
|
||||
spreadsheetId=sheet_id,
|
||||
range="Sheet1!A:F",
|
||||
valueInputOption="RAW",
|
||||
insertDataOption="INSERT_ROWS",
|
||||
body={"values": [row]},
|
||||
)
|
||||
.execute
|
||||
)
|
||||
logger.info("Survey response appended for user %s", user_id)
|
||||
return True
|
||||
except Exception:
|
||||
logger.exception("Failed to append survey response to Google Sheet for user %s", user_id)
|
||||
return False
|
||||
@@ -0,0 +1,977 @@
|
||||
"""Async direct datasheet review — per-IC with graph tools.
|
||||
|
||||
Each IC gets a review call with its datasheet PDF and circuit neighborhood.
|
||||
ICs run concurrently with a semaphore. Provider-agnostic — routes through
|
||||
the LLM provider abstraction so a stage env var (PROVIDER_VALIDATION) can
|
||||
flip between Anthropic and Gemini without code changes.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
import tempfile
|
||||
import time
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Awaitable, Callable
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
from backend.pinscopex.models import (
|
||||
ComponentConstraints,
|
||||
ComponentType,
|
||||
DesignGraph,
|
||||
Finding,
|
||||
NetType,
|
||||
ValidationReport,
|
||||
)
|
||||
from backend.pinscopex.validate import (
|
||||
SYSTEM_PROMPT,
|
||||
_MAX_REVIEW_TURNS,
|
||||
ReviewResult,
|
||||
_load_datasheets,
|
||||
_match_constraints,
|
||||
_build_constraints_map,
|
||||
assign_finding_ids,
|
||||
build_component_context,
|
||||
_parse_review,
|
||||
)
|
||||
from backend.pinscopex.utils import safe_mpn
|
||||
from backend.pinscopex.pin_mux_check import check_pin_mux_feasibility
|
||||
from backend.pinscopex.led_current_check import check_led_current
|
||||
|
||||
TRACE_VERSION = 1
|
||||
|
||||
|
||||
def _is_deterministic(f: Finding) -> bool:
|
||||
"""True for a finding produced by a deterministic check (not the LLM review)."""
|
||||
return bool(getattr(f, "source", None)) and f.source != "review"
|
||||
|
||||
|
||||
def _run_deterministic_checks(
|
||||
graph: DesignGraph, constraints_map: dict
|
||||
) -> list[Finding]:
|
||||
"""Run the deterministic graph checks, fail-soft per check — a check bug
|
||||
can never break the review or the report."""
|
||||
out: list[Finding] = []
|
||||
for name, fn in (
|
||||
("pin_mux_check", lambda: check_pin_mux_feasibility(graph, constraints_map)),
|
||||
("led_current_check", lambda: check_led_current(graph)),
|
||||
):
|
||||
try:
|
||||
out.extend(fn())
|
||||
except Exception:
|
||||
log.exception("deterministic check %s failed — skipping", name)
|
||||
return out
|
||||
|
||||
|
||||
def _assistant_text(blocks) -> str:
|
||||
"""Best-effort extraction of text content from a completion's raw
|
||||
assistant blocks. Provider-agnostic and never raises."""
|
||||
parts: list[str] = []
|
||||
try:
|
||||
for b in blocks or []:
|
||||
txt = getattr(b, "text", None)
|
||||
if txt is None and isinstance(b, dict):
|
||||
txt = b.get("text") if b.get("type") == "text" else None
|
||||
elif getattr(b, "type", None) not in (None, "text"):
|
||||
txt = None
|
||||
if isinstance(txt, str) and txt:
|
||||
parts.append(txt)
|
||||
except Exception:
|
||||
log.exception("trace: assistant_text extraction failed")
|
||||
return "\n".join(parts)
|
||||
from backend.pinscopex.validation_tools import (
|
||||
ALL_TOOLS,
|
||||
SUBMIT_REVIEW_SCHEMA,
|
||||
ConstraintsMap,
|
||||
ExcerptState,
|
||||
execute_tool,
|
||||
)
|
||||
from backend.pinscopex.utils import safe_mpn
|
||||
|
||||
from backend.config import settings
|
||||
from backend.services.api_logs import ApiLogger
|
||||
from backend.services.normalize_findings import normalize_findings_async
|
||||
from backend.services.dedupe_findings import dedupe_cross_ic_findings_async
|
||||
from backend.services.llm import (
|
||||
Message,
|
||||
PdfBlock,
|
||||
TextBlock,
|
||||
ToolCall,
|
||||
ToolResultBlock,
|
||||
ToolSchema,
|
||||
call_with_fallback,
|
||||
)
|
||||
|
||||
# Type for progress callback: (ref, turn, tool_name_or_status, detail)
|
||||
ProgressCallback = Callable[[str, int, str, str], Awaitable[None]]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tool schemas — defined as dicts in validation_tools.py, converted here
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _to_tool_schema(d: dict) -> ToolSchema:
|
||||
return ToolSchema(
|
||||
name=d["name"],
|
||||
description=d["description"],
|
||||
input_schema=d["input_schema"],
|
||||
)
|
||||
|
||||
|
||||
_ALL_TOOL_SCHEMAS = [_to_tool_schema(t) for t in ALL_TOOLS]
|
||||
_SUBMIT_TOOL_SCHEMA = _to_tool_schema(SUBMIT_REVIEW_SCHEMA)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Review keywords for PDF page trimming
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_REVIEW_KEYWORDS = re.compile(
|
||||
r"pin\s+(out|diagram|configuration|description|assignment|function|name|table|map)"
|
||||
r"|ball\s+map|package\s+(pin|drawing|outline)|signal\s+description"
|
||||
r"|absolute\s+maximum|recommended\s+operating|electrical\s+characteristics"
|
||||
r"|power\s+supply|thermal\s+(resistance|shutdown|pad)|ESD\s+(rating|tolerance)"
|
||||
r"|decoupling|bypass\s+capacitor|layout\s+(guideline|recommendation)"
|
||||
r"|application\s+(circuit|schematic|information|note)"
|
||||
r"|typical\s+application|reference\s+design",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
|
||||
_MAX_PDF_PAGES = 90
|
||||
|
||||
# Per-review excerpt budget — keeps fan-out cost bounded on hub ICs (e.g. an
|
||||
# MCU connected to many neighbors). On exhaustion, the tool returns a budget
|
||||
# message and the model is steered to submit WARNING with Unverified:
|
||||
# assumption rather than fetching more.
|
||||
#
|
||||
# The global page budget got raised from 25→60 and gained a per-neighbor
|
||||
# sub-budget after the U2-001 / U3-001 false positives: a single 25-page
|
||||
# global cap was exhausted by one neighbor's pin_voltage_levels excerpt
|
||||
# before the abs-max table could be read, so the reviewer was forced to
|
||||
# guess at the very moment it was trying to verify a damage claim. 30 pages
|
||||
# per neighbor fits the ~3 topic fetches (pin levels + abs-max + electrical)
|
||||
# one interface check needs; 60 global allows ~2 such neighbors before the
|
||||
# fan-out ceiling kicks in.
|
||||
_PER_REVIEW_FETCH_BUDGET = 8
|
||||
_PER_REVIEW_PAGE_BUDGET = 60
|
||||
_PER_NEIGHBOR_PAGE_BUDGET = 30
|
||||
|
||||
# A signal net with more components than this is treated as a hub/bus and
|
||||
# excluded from the neighbor set even if classified as "signal". Bounds
|
||||
# fan-out on designs that use an oversized common signal (rare but possible).
|
||||
_SIGNAL_NET_MAX_COMPONENTS = 8
|
||||
|
||||
|
||||
def _signal_neighbors(graph: DesignGraph, ic_ref: str) -> set[str]:
|
||||
"""Return the set of designators that share at least one *signal* net
|
||||
with ``ic_ref``. Excludes power/ground rails (which connect every IC and
|
||||
would otherwise fan the neighbor set out across the whole design) and
|
||||
excludes the IC under review itself.
|
||||
"""
|
||||
comp = graph.components.get(ic_ref)
|
||||
if not comp:
|
||||
return set()
|
||||
neighbors: set[str] = set()
|
||||
for net_name in set(comp.pins.values()):
|
||||
net = graph.nets.get(net_name)
|
||||
if not net:
|
||||
continue
|
||||
if net.net_type in (NetType.POWER, NetType.GROUND):
|
||||
continue
|
||||
refs_on_net = {pc.component_ref for pc in net.pins}
|
||||
if len(refs_on_net) > _SIGNAL_NET_MAX_COMPONENTS:
|
||||
continue
|
||||
for ref in refs_on_net:
|
||||
if ref != ic_ref:
|
||||
neighbors.add(ref)
|
||||
return neighbors
|
||||
|
||||
|
||||
def _select_review_pages(pdf_path: str) -> str:
|
||||
"""Trim a datasheet PDF to pages relevant for design review.
|
||||
|
||||
Returns path to trimmed PDF (or original if already small enough).
|
||||
|
||||
Note: the reviewer cites the datasheet's *printed* page number (read from
|
||||
the page content/footer), not the page's physical position in the trimmed
|
||||
file — so `source_page` already matches the full original PDF the frontend
|
||||
serves. No trimmed→original remap is applied (an earlier remap attempt
|
||||
corrupted correct citations on large datasheets).
|
||||
"""
|
||||
from pypdf import PdfReader, PdfWriter
|
||||
|
||||
reader = PdfReader(pdf_path)
|
||||
total = len(reader.pages)
|
||||
if total <= _MAX_PDF_PAGES:
|
||||
return pdf_path
|
||||
|
||||
# Always keep first 5 pages (title, TOC, overview)
|
||||
keep: set[int] = set(range(min(5, total)))
|
||||
|
||||
# Keyword-matched pages + neighbors
|
||||
for i, page in enumerate(reader.pages):
|
||||
text = page.extract_text() or ""
|
||||
if _REVIEW_KEYWORDS.search(text):
|
||||
for neighbor in (i - 1, i, i + 1):
|
||||
if 0 <= neighbor < total:
|
||||
keep.add(neighbor)
|
||||
|
||||
# Pad from front if under budget
|
||||
if len(keep) < _MAX_PDF_PAGES:
|
||||
for i in range(total):
|
||||
if len(keep) >= _MAX_PDF_PAGES:
|
||||
break
|
||||
keep.add(i)
|
||||
|
||||
selected = sorted(keep)[:_MAX_PDF_PAGES]
|
||||
|
||||
writer = PdfWriter()
|
||||
for i in selected:
|
||||
writer.add_page(reader.pages[i])
|
||||
|
||||
tmp = tempfile.NamedTemporaryFile(suffix=".pdf", delete=False)
|
||||
writer.write(tmp)
|
||||
tmp.close()
|
||||
return tmp.name
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Per-IC async review
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def review_ic_async(
|
||||
graph: DesignGraph,
|
||||
constraints_map: ConstraintsMap,
|
||||
ic_ref: str,
|
||||
pdf_path: str,
|
||||
on_progress: ProgressCallback | None = None,
|
||||
api_logger: ApiLogger | None = None,
|
||||
trace_git_commit: str = "unknown",
|
||||
pdf_dir: Path | None = None,
|
||||
storage=None,
|
||||
excerpt_cache: dict | None = None,
|
||||
) -> tuple[ReviewResult, dict]:
|
||||
"""Review one IC against its datasheet. Async, multi-turn.
|
||||
|
||||
Returns ``(ReviewResult, trace)`` — ``trace`` is a transcript dict of the
|
||||
full agentic loop (turns, tool calls + outputs, final submission) for
|
||||
offline inspection. Trace assembly is best-effort and never affects the
|
||||
review result.
|
||||
"""
|
||||
comp = graph.components[ic_ref]
|
||||
mpn = comp.mpn or comp.value
|
||||
|
||||
# Datasheet identity for the trace — hash the original PDF, not the
|
||||
# trimmed copy, so the reference is stable across trim-heuristic changes.
|
||||
try:
|
||||
ds_md5 = hashlib.md5(Path(pdf_path).read_bytes()).hexdigest()
|
||||
except Exception:
|
||||
log.exception("trace: datasheet md5 failed for %s", ic_ref)
|
||||
ds_md5 = None
|
||||
|
||||
# Pre-compute which designators the excerpt tool will accept for this
|
||||
# review (neighbors via signal nets only — power/GND fan-out filtered).
|
||||
connected_designators = _signal_neighbors(graph, ic_ref)
|
||||
|
||||
# Designator -> MPN, so a finding citing a neighbor's datasheet excerpt
|
||||
# (source_designator) is referenced against — and viewed from — that
|
||||
# neighbor's datasheet rather than this IC's.
|
||||
mpn_by_designator = {
|
||||
ref: comp.mpn
|
||||
for ref, comp in graph.components.items()
|
||||
if comp.mpn
|
||||
}
|
||||
|
||||
# Build the per-review state for the excerpt tool. ``cache`` is shared
|
||||
# across ICs in the same validate_design_async run so symmetric checks
|
||||
# (U2 fetches U3@abs_max, then U3 fetches U2@abs_max) don't redo pypdf
|
||||
# work.
|
||||
excerpt_state = ExcerptState(
|
||||
current_ic=ic_ref,
|
||||
connected_designators=connected_designators,
|
||||
graph=graph,
|
||||
pdf_dir=pdf_dir or Path(pdf_path).parent,
|
||||
storage=storage,
|
||||
cache=excerpt_cache if excerpt_cache is not None else {},
|
||||
fetch_budget=_PER_REVIEW_FETCH_BUDGET,
|
||||
page_budget=_PER_REVIEW_PAGE_BUDGET,
|
||||
per_neighbor_page_budget=_PER_NEIGHBOR_PAGE_BUDGET,
|
||||
)
|
||||
|
||||
# Trim PDF up-front — both primary and fallback attempts share it.
|
||||
trimmed_pdf = _select_review_pages(pdf_path)
|
||||
try:
|
||||
async def _run(provider, model) -> tuple[ReviewResult, dict]:
|
||||
t0 = time.monotonic()
|
||||
total_input = 0
|
||||
total_output = 0
|
||||
total_cache_creation = 0
|
||||
total_cache_read = 0
|
||||
turns = 0
|
||||
|
||||
session = await provider.create_session(
|
||||
model=model,
|
||||
system=SYSTEM_PROMPT,
|
||||
# Gemini 2.5/3 thinking models count thoughts against this cap.
|
||||
# 4096 was too tight: U3 (largest IC) burned the entire budget
|
||||
# on thinking and emitted zero visible output, dropping its
|
||||
# review silently.
|
||||
max_tokens=32768,
|
||||
# Deterministic sampling: same inputs → same findings across
|
||||
# reruns. The default temperature of 1.0 caused identical
|
||||
# netlists to produce very different reports (different
|
||||
# findings + severities) run-to-run.
|
||||
temperature=0.0,
|
||||
)
|
||||
try:
|
||||
context = build_component_context(graph, constraints_map, ic_ref)
|
||||
|
||||
initial_msg = Message(
|
||||
role="user",
|
||||
content=[
|
||||
PdfBlock(path=Path(trimmed_pdf), cacheable=True),
|
||||
TextBlock(
|
||||
text=f"Review this component's usage:\n\n{context}",
|
||||
cacheable=True,
|
||||
),
|
||||
],
|
||||
)
|
||||
messages: list[Message] = [initial_msg]
|
||||
|
||||
trace: dict = {
|
||||
"trace_version": TRACE_VERSION,
|
||||
"ic_ref": ic_ref,
|
||||
"mpn": mpn,
|
||||
"model": model,
|
||||
"provider": provider.name,
|
||||
"git_commit": trace_git_commit,
|
||||
"datasheet": {"md5": ds_md5, "safe_mpn": safe_mpn(mpn)},
|
||||
"timestamp": datetime.now(timezone.utc).isoformat(),
|
||||
"max_turns": _MAX_REVIEW_TURNS,
|
||||
"turns": [],
|
||||
"final_submission": None,
|
||||
"result": None,
|
||||
"stop_reason": None,
|
||||
"error": None,
|
||||
"duration_ms": None,
|
||||
}
|
||||
|
||||
# Set after a turn produces zero tool calls (model wrote
|
||||
# text only). Next turn is forced to submit_review so any
|
||||
# findings drafted as prose still make it to the report.
|
||||
force_submit_next_turn = False
|
||||
|
||||
for turn in range(_MAX_REVIEW_TURNS):
|
||||
is_last_turn = turn == _MAX_REVIEW_TURNS - 1
|
||||
|
||||
if is_last_turn or force_submit_next_turn:
|
||||
tools = [_SUBMIT_TOOL_SCHEMA]
|
||||
tool_choice: dict | str = {"name": "submit_review"}
|
||||
else:
|
||||
tools = _ALL_TOOL_SCHEMAS
|
||||
tool_choice = "auto"
|
||||
|
||||
completion = await session.complete(
|
||||
messages=messages,
|
||||
tools=tools,
|
||||
tool_choice=tool_choice,
|
||||
)
|
||||
turns += 1
|
||||
total_input += completion.usage.input_tokens
|
||||
total_output += completion.usage.output_tokens
|
||||
total_cache_creation += completion.usage.cache_creation_tokens
|
||||
total_cache_read += completion.usage.cache_read_tokens
|
||||
|
||||
turn_record: dict = {
|
||||
"index": turn,
|
||||
"assistant_text": _assistant_text(
|
||||
completion.raw_assistant_blocks
|
||||
),
|
||||
"tool_calls": [],
|
||||
"usage": {
|
||||
"input_tokens": completion.usage.input_tokens,
|
||||
"output_tokens": completion.usage.output_tokens,
|
||||
"cache_creation_tokens": completion.usage.cache_creation_tokens,
|
||||
"cache_read_tokens": completion.usage.cache_read_tokens,
|
||||
},
|
||||
}
|
||||
try:
|
||||
trace["turns"].append(turn_record)
|
||||
except Exception:
|
||||
log.exception("trace: turn append failed for %s", ic_ref)
|
||||
|
||||
# Check for submit_review
|
||||
for tc in completion.tool_calls:
|
||||
if tc.name == "submit_review":
|
||||
result = _parse_review(
|
||||
tc.input, ic_ref, mpn,
|
||||
mpn_by_designator=mpn_by_designator,
|
||||
connected=connected_designators,
|
||||
)
|
||||
turn_record["tool_calls"].append({
|
||||
"name": "submit_review",
|
||||
"input": tc.input,
|
||||
"output": None,
|
||||
"duration_ms": None,
|
||||
})
|
||||
trace["final_submission"] = tc.input
|
||||
trace["stop_reason"] = "submit_review"
|
||||
trace["result"] = {
|
||||
"findings_count": len(result.findings),
|
||||
"checked_areas": result.checked_areas,
|
||||
}
|
||||
trace["duration_ms"] = int((time.monotonic() - t0) * 1000)
|
||||
if on_progress:
|
||||
await on_progress(
|
||||
ic_ref, turn, "submit_review",
|
||||
f"{len(result.findings)} findings",
|
||||
)
|
||||
if api_logger:
|
||||
api_logger.log(
|
||||
stage="review", identifier=ic_ref,
|
||||
model=model, provider=provider.name,
|
||||
input_tokens=total_input, output_tokens=total_output,
|
||||
cache_creation_input_tokens=total_cache_creation,
|
||||
cache_read_input_tokens=total_cache_read,
|
||||
duration_ms=int((time.monotonic() - t0) * 1000),
|
||||
stop_reason="submit_review", turns=turns,
|
||||
)
|
||||
if settings.normalize_findings_enabled:
|
||||
try:
|
||||
normalized, norm_trace = await normalize_findings_async(
|
||||
ic_ref, mpn, result.findings,
|
||||
api_logger=api_logger,
|
||||
on_progress=on_progress,
|
||||
)
|
||||
trace["normalize"] = norm_trace
|
||||
result.findings = normalized
|
||||
trace["result"]["findings_count"] = len(normalized)
|
||||
except Exception:
|
||||
log.exception(
|
||||
"normalize: unexpected failure for %s "
|
||||
"— keeping reviewer findings",
|
||||
ic_ref,
|
||||
)
|
||||
return result, trace
|
||||
|
||||
# Process graph tool calls
|
||||
tool_results: list[ToolResultBlock] = []
|
||||
attached_pdfs: list[PdfBlock] = []
|
||||
for tc in completion.tool_calls:
|
||||
_tc_t0 = time.monotonic()
|
||||
result_text, attachment = execute_tool(
|
||||
graph, constraints_map, tc.name, tc.input,
|
||||
state=excerpt_state,
|
||||
)
|
||||
turn_record["tool_calls"].append({
|
||||
"name": tc.name,
|
||||
"input": tc.input,
|
||||
"output": result_text,
|
||||
"duration_ms": int((time.monotonic() - _tc_t0) * 1000),
|
||||
})
|
||||
if on_progress:
|
||||
await on_progress(
|
||||
ic_ref, turn, tc.name, json.dumps(tc.input),
|
||||
)
|
||||
tool_results.append(ToolResultBlock(
|
||||
tool_use_id=tc.id,
|
||||
name=tc.name,
|
||||
content=result_text,
|
||||
))
|
||||
if attachment is not None:
|
||||
attached_pdfs.append(attachment)
|
||||
|
||||
if not tool_results:
|
||||
# Model emitted text but called no tools. This is a
|
||||
# known failure mode (esp. with reasoning models)
|
||||
# where the model writes findings as a JSON code
|
||||
# block in prose instead of calling submit_review.
|
||||
# Don't drop the work — append a nudge and force
|
||||
# submit_review on the next iteration.
|
||||
if not is_last_turn and not force_submit_next_turn:
|
||||
messages.append(Message(
|
||||
role="assistant",
|
||||
content=completion.raw_assistant_blocks,
|
||||
))
|
||||
messages.append(Message(
|
||||
role="user",
|
||||
content=[TextBlock(
|
||||
text=(
|
||||
"You produced text but did not call "
|
||||
"any tool. Findings only reach the "
|
||||
"report when submitted via the "
|
||||
"submit_review tool — text JSON is "
|
||||
"ignored. Call submit_review now with "
|
||||
"the findings you identified (or an "
|
||||
"empty findings array if none) and "
|
||||
"your checked_areas list."
|
||||
),
|
||||
)],
|
||||
))
|
||||
force_submit_next_turn = True
|
||||
continue
|
||||
break
|
||||
|
||||
# Reset recovery flag once the model is calling tools again.
|
||||
force_submit_next_turn = False
|
||||
|
||||
messages.append(Message(role="assistant", content=completion.raw_assistant_blocks))
|
||||
# tool_result blocks first, then any PdfBlocks the tools
|
||||
# attached (excerpt fetches). The Anthropic provider
|
||||
# encodes each block independently — mixed-block user
|
||||
# messages are supported and the cached initial PDF is
|
||||
# not invalidated by appending uncached/cached content.
|
||||
messages.append(Message(
|
||||
role="user",
|
||||
content=[*tool_results, *attached_pdfs],
|
||||
))
|
||||
|
||||
# Fell through without submitting
|
||||
trace["stop_reason"] = "no_submission"
|
||||
trace["result"] = {"findings_count": 0, "checked_areas": []}
|
||||
trace["duration_ms"] = int((time.monotonic() - t0) * 1000)
|
||||
if api_logger:
|
||||
api_logger.log(
|
||||
stage="review", identifier=ic_ref,
|
||||
model=model, provider=provider.name,
|
||||
input_tokens=total_input, output_tokens=total_output,
|
||||
cache_creation_input_tokens=total_cache_creation,
|
||||
cache_read_input_tokens=total_cache_read,
|
||||
duration_ms=int((time.monotonic() - t0) * 1000),
|
||||
stop_reason="no_submission", turns=turns,
|
||||
)
|
||||
return ReviewResult([], []), trace
|
||||
finally:
|
||||
await session.close()
|
||||
|
||||
return await call_with_fallback("validation", _run)
|
||||
finally:
|
||||
if trimmed_pdf != pdf_path:
|
||||
Path(trimmed_pdf).unlink(missing_ok=True)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# PDF resolution
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _find_pdf(
|
||||
mpn: str,
|
||||
pdf_dir: Path,
|
||||
storage=None,
|
||||
) -> Path | None:
|
||||
"""Find the datasheet PDF for an MPN. Checks local dir first,
|
||||
then tries to download from the library.
|
||||
"""
|
||||
safe = safe_mpn(mpn)
|
||||
local = pdf_dir / f"{safe}.pdf"
|
||||
if local.is_file():
|
||||
return local
|
||||
|
||||
if storage:
|
||||
from backend.services import projects as proj_svc
|
||||
lib_key = proj_svc.library_has_datasheet(storage, mpn)
|
||||
if lib_key:
|
||||
storage.download_to_local(lib_key, local)
|
||||
if local.is_file():
|
||||
return local
|
||||
|
||||
return None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Main entrypoint
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
BeforeIcCallback = Callable[[str], Awaitable[bool]]
|
||||
"""Gate callback — called with the IC ref before review. Return False to pause."""
|
||||
|
||||
OnIcDoneCallback = Callable[[str, "ReviewResult", "ApiLogger | None"], Awaitable[None]]
|
||||
"""Callback after each IC finishes successfully — used to charge credits.
|
||||
|
||||
Receives the IC's private ``ApiLogger`` (the calls made during this review)
|
||||
so the charge can be attributed to exactly this IC under concurrency."""
|
||||
|
||||
OnIcErrorCallback = Callable[[str, BaseException], Awaitable[None]]
|
||||
"""Callback after an IC review raises — used to record a SkippedItem so the
|
||||
failure surfaces in the project's skipped_components list."""
|
||||
|
||||
OnDedupeDoneCallback = Callable[["ApiLogger | None"], Awaitable[None]]
|
||||
"""Callback after the cross-IC dedup pass finishes — used to charge for that
|
||||
single LLM call (it runs once at end-of-run, outside any per-IC logger)."""
|
||||
|
||||
|
||||
async def validate_design_async(
|
||||
graph_path: str,
|
||||
output_path: str,
|
||||
datasheets_dir: str = "datasheets/extracted",
|
||||
pdf_dir: str = "uploads/datasheets",
|
||||
on_progress: ProgressCallback | None = None,
|
||||
api_logger: ApiLogger | None = None,
|
||||
storage=None,
|
||||
skip_refs: set[str] | None = None,
|
||||
before_ic: BeforeIcCallback | None = None,
|
||||
on_ic_done: OnIcDoneCallback | None = None,
|
||||
on_ic_error: OnIcErrorCallback | None = None,
|
||||
on_dedupe_done: OnDedupeDoneCallback | None = None,
|
||||
project_prefix: str | None = None,
|
||||
run_meta: dict | None = None,
|
||||
) -> ValidationReport:
|
||||
"""Review every IC against its datasheet.
|
||||
|
||||
By default runs concurrently via an asyncio.Semaphore. When ``before_ic``
|
||||
is supplied, reviews are executed sequentially so the callback can
|
||||
decide whether to pause the run between ICs. In that mode the report
|
||||
is written incrementally after each IC so a pause preserves all
|
||||
completed findings.
|
||||
|
||||
``skip_refs`` is consumed on the first pass — any IC in the set is
|
||||
skipped without starting a review (used to resume a paused run).
|
||||
"""
|
||||
skip_refs = skip_refs or set()
|
||||
|
||||
raw = json.loads(Path(graph_path).read_text())
|
||||
graph = DesignGraph.model_validate(raw)
|
||||
datasheets = _load_datasheets(datasheets_dir)
|
||||
constraints_map = _build_constraints_map(datasheets)
|
||||
|
||||
# Deterministic graph checks (pin-mux feasibility, LED current). Pure
|
||||
# functions of the graph; fail-soft. Seeded into all_findings below.
|
||||
deterministic_findings = _run_deterministic_checks(graph, constraints_map)
|
||||
|
||||
pdf_dir_path = Path(pdf_dir)
|
||||
|
||||
# Collect ICs that have a datasheet PDF available
|
||||
ic_tasks: list[tuple[str, str]] = [] # (ref, pdf_path)
|
||||
not_reviewed: list[dict] = [] # ICs skipped for lack of a datasheet PDF
|
||||
for ref, comp in sorted(graph.components.items()):
|
||||
if comp.component_type != ComponentType.IC:
|
||||
continue
|
||||
mpn = comp.mpn or comp.value
|
||||
pdf = _find_pdf(mpn, pdf_dir_path, storage=storage)
|
||||
if pdf:
|
||||
ic_tasks.append((ref, str(pdf)))
|
||||
else:
|
||||
not_reviewed.append({"designator": ref, "reason": "no datasheet PDF"})
|
||||
if on_progress:
|
||||
await on_progress(ref, 0, "skipped", "no datasheet PDF")
|
||||
|
||||
# Load any previously-written report so we can accumulate findings
|
||||
# across a pause/resume cycle without losing prior results.
|
||||
existing_path = Path(output_path)
|
||||
preserved_findings: list[Finding] = []
|
||||
preserved_coverage: dict[str, list[str]] = {}
|
||||
preserved_comments = None
|
||||
if existing_path.is_file():
|
||||
try:
|
||||
existing = json.loads(existing_path.read_text())
|
||||
preserved_comments = existing.get("comments")
|
||||
if before_ic is not None:
|
||||
# Resume mode — keep findings for refs we're about to skip
|
||||
for f in existing.get("findings", []):
|
||||
ref = f.get("component_ref") or f.get("designator") or ""
|
||||
if ref in skip_refs:
|
||||
preserved_findings.append(Finding.model_validate(f))
|
||||
for ref, areas in (existing.get("coverage") or {}).items():
|
||||
if ref in skip_refs:
|
||||
preserved_coverage[ref] = list(areas)
|
||||
except (json.JSONDecodeError, OSError):
|
||||
pass
|
||||
|
||||
# Seed deterministic findings exactly once. On resume, preserved_findings may
|
||||
# already contain them (they were written to the prior report), so strip any
|
||||
# deterministic findings before re-seeding to avoid double-counting.
|
||||
preserved_review = [f for f in preserved_findings if not _is_deterministic(f)]
|
||||
all_findings: list[Finding] = list(preserved_review) + list(deterministic_findings)
|
||||
all_coverage: dict[str, list[str]] = dict(preserved_coverage)
|
||||
review_errors: dict[str, str] = {}
|
||||
|
||||
def _sanitize_coverage(src: dict[str, list[str]]) -> dict[str, list[str]]:
|
||||
"""Drop any entries that aren't a list of strings so one IC's bad
|
||||
payload can't fail the whole ValidationReport validation."""
|
||||
clean: dict[str, list[str]] = {}
|
||||
for ref, areas in src.items():
|
||||
if isinstance(areas, list) and all(isinstance(a, str) for a in areas):
|
||||
clean[ref] = areas
|
||||
else:
|
||||
print(f"[validation] dropping coverage for {ref}: {areas!r}")
|
||||
return clean
|
||||
|
||||
def _write_report(paused: bool = False) -> ValidationReport:
|
||||
assign_finding_ids(all_findings)
|
||||
summary = {"total": len(all_findings), "ERROR": 0, "WARNING": 0, "INFO": 0}
|
||||
for f in all_findings:
|
||||
summary[f.status] = summary.get(f.status, 0) + 1
|
||||
try:
|
||||
report = ValidationReport(
|
||||
project=Path(graph_path).stem,
|
||||
timestamp=datetime.now(timezone.utc).isoformat(),
|
||||
findings=all_findings,
|
||||
summary=summary,
|
||||
coverage=_sanitize_coverage(all_coverage),
|
||||
review_errors=dict(review_errors),
|
||||
not_reviewed=not_reviewed,
|
||||
)
|
||||
except Exception as exc:
|
||||
print(f"[validation] report build failed, retrying without coverage: {exc}")
|
||||
report = ValidationReport(
|
||||
project=Path(graph_path).stem,
|
||||
timestamp=datetime.now(timezone.utc).isoformat(),
|
||||
findings=all_findings,
|
||||
summary=summary,
|
||||
coverage={},
|
||||
review_errors=dict(review_errors),
|
||||
not_reviewed=not_reviewed,
|
||||
)
|
||||
report_dict = json.loads(report.model_dump_json(indent=2))
|
||||
if preserved_comments is not None:
|
||||
report_dict["comments"] = preserved_comments
|
||||
if paused:
|
||||
report_dict["partial"] = True
|
||||
existing_path.write_text(json.dumps(report_dict, indent=2))
|
||||
return report
|
||||
|
||||
git_commit = (run_meta or {}).get("git_commit", "unknown")
|
||||
|
||||
def _write_trace(trace: dict, ref: str) -> None:
|
||||
"""Persist a per-IC review trace. Best-effort: a trace failure must
|
||||
never break the review, the report, or the pipeline."""
|
||||
if not storage or not project_prefix or not trace:
|
||||
return
|
||||
try:
|
||||
key = f"{project_prefix}/review_traces/{safe_mpn(ref)}.json"
|
||||
storage.write_json(key, trace)
|
||||
except Exception:
|
||||
log.exception("trace: write failed for %s", ref)
|
||||
|
||||
async def _maybe_dedupe_cross_ic() -> None:
|
||||
"""Collapse one interface defect reported from both ICs into a single
|
||||
finding. Runs once, after all per-IC reviews, when findings span ≥2
|
||||
ICs. Mutates ``all_findings`` in place. Best-effort: any failure keeps
|
||||
the per-IC findings (the dedup function is itself fail-soft)."""
|
||||
if not settings.cross_ic_dedup_enabled:
|
||||
return
|
||||
# Deterministic findings never enter the LLM dedupe — it has no datasheet
|
||||
# basis to judge a pin-mux/LED finding, and merging could mangle them.
|
||||
review = [f for f in all_findings if not _is_deterministic(f)]
|
||||
deterministic = [f for f in all_findings if _is_deterministic(f)]
|
||||
if len({f.designator for f in review}) < 2:
|
||||
return # nothing cross-IC to merge
|
||||
# Gated path: charge via a private logger merged by on_dedupe_done.
|
||||
# Legacy path (no callback): log straight to the shared logger so the
|
||||
# call still shows up in api_logs even though nothing is charged.
|
||||
private = (
|
||||
ApiLogger(free=api_logger.free)
|
||||
if (api_logger is not None and on_dedupe_done is not None)
|
||||
else None
|
||||
)
|
||||
try:
|
||||
deduped, dedupe_trace = await dedupe_cross_ic_findings_async(
|
||||
review,
|
||||
api_logger=private if private is not None else api_logger,
|
||||
on_progress=on_progress,
|
||||
)
|
||||
except Exception:
|
||||
log.exception("cross-IC dedupe failed — keeping per-IC findings")
|
||||
return
|
||||
all_findings[:] = deduped + deterministic
|
||||
if storage and project_prefix and dedupe_trace:
|
||||
try:
|
||||
storage.write_json(
|
||||
f"{project_prefix}/review_traces/_cross_ic_dedupe.json",
|
||||
dedupe_trace,
|
||||
)
|
||||
except Exception:
|
||||
log.exception("trace: cross-IC dedupe write failed")
|
||||
# Charge for the single dedup call (gated path only — the private
|
||||
# logger merges into the shared log and bills exactly this call).
|
||||
if private is not None and on_dedupe_done is not None:
|
||||
try:
|
||||
await on_dedupe_done(private)
|
||||
except Exception:
|
||||
log.exception("on_dedupe_done callback failed")
|
||||
|
||||
def _stub_trace(ref: str, error: str) -> dict:
|
||||
"""Minimal trace for an IC whose review raised before producing one,
|
||||
so an eval harness still sees a record for every attempted IC."""
|
||||
try:
|
||||
comp = graph.components.get(ref)
|
||||
mpn = (comp.mpn or comp.value) if comp else ref
|
||||
except Exception:
|
||||
mpn = ref
|
||||
return {
|
||||
"trace_version": TRACE_VERSION,
|
||||
"ic_ref": ref,
|
||||
"mpn": mpn,
|
||||
"git_commit": git_commit,
|
||||
"datasheet": {"md5": None, "safe_mpn": safe_mpn(mpn)},
|
||||
"timestamp": datetime.now(timezone.utc).isoformat(),
|
||||
"turns": [],
|
||||
"final_submission": None,
|
||||
"result": None,
|
||||
"stop_reason": "error",
|
||||
"error": error,
|
||||
"duration_ms": None,
|
||||
}
|
||||
|
||||
# Cross-IC excerpt cache — symmetric interface checks (U2 fetches U3@X,
|
||||
# U3 fetches U2@X) reuse the trimmed PDF instead of redoing pypdf work.
|
||||
# LLM-side ephemeral cache can't span ICs (different conversation prefix),
|
||||
# so the win here is purely pypdf I/O.
|
||||
excerpt_cache: dict = {}
|
||||
|
||||
def _cleanup_excerpt_cache() -> None:
|
||||
for entry in excerpt_cache.values():
|
||||
try:
|
||||
if isinstance(entry, tuple) and len(entry) == 2:
|
||||
Path(entry[0]).unlink(missing_ok=True)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if before_ic is None:
|
||||
# Legacy concurrent path (no credit gate)
|
||||
sem = asyncio.Semaphore(settings.ic_concurrency)
|
||||
|
||||
async def _review_one(ref: str, pdf_path: str) -> tuple[ReviewResult, dict]:
|
||||
async with sem:
|
||||
return await review_ic_async(
|
||||
graph, constraints_map, ref, pdf_path,
|
||||
on_progress=on_progress, api_logger=api_logger,
|
||||
trace_git_commit=git_commit,
|
||||
pdf_dir=pdf_dir_path, storage=storage,
|
||||
excerpt_cache=excerpt_cache,
|
||||
)
|
||||
|
||||
results = await asyncio.gather(
|
||||
*(_review_one(ref, pdf) for ref, pdf in ic_tasks if ref not in skip_refs),
|
||||
return_exceptions=True,
|
||||
)
|
||||
remaining_tasks = [t for t in ic_tasks if t[0] not in skip_refs]
|
||||
for i, result in enumerate(results):
|
||||
ref = remaining_tasks[i][0]
|
||||
if isinstance(result, BaseException):
|
||||
msg = f"{type(result).__name__}: {result}"
|
||||
log.exception("Review failed for %s", ref, exc_info=result)
|
||||
review_errors[ref] = msg
|
||||
_write_trace(_stub_trace(ref, msg), ref)
|
||||
if on_progress:
|
||||
await on_progress(ref, 0, "error", msg)
|
||||
if on_ic_error is not None:
|
||||
try:
|
||||
await on_ic_error(ref, result)
|
||||
except Exception:
|
||||
log.exception("on_ic_error callback failed for %s", ref)
|
||||
elif isinstance(result, tuple):
|
||||
rr, trace = result
|
||||
_write_trace(trace, ref)
|
||||
all_findings.extend(rr.findings)
|
||||
if rr.checked_areas:
|
||||
all_coverage[ref] = rr.checked_areas
|
||||
await _maybe_dedupe_cross_ic()
|
||||
try:
|
||||
return _write_report(paused=False)
|
||||
finally:
|
||||
_cleanup_excerpt_cache()
|
||||
|
||||
# Gated concurrent path — used by the pipeline with credit enforcement.
|
||||
# Runs up to ``ic_concurrency`` reviews in parallel while keeping the
|
||||
# per-IC credit gate, incremental report/trace writes, and the charging
|
||||
# callback. Each IC reviews against a private ApiLogger so concurrent
|
||||
# reviews don't interleave their API entries — on_ic_done charges exactly
|
||||
# that IC's calls.
|
||||
sem = asyncio.Semaphore(settings.ic_concurrency)
|
||||
stop = False # set once a gate trips — stops *starting* new reviews
|
||||
|
||||
async def _gated_review_one(ref: str, pdf_path: str) -> None:
|
||||
nonlocal stop
|
||||
async with sem:
|
||||
if stop:
|
||||
return
|
||||
try:
|
||||
ok = await before_ic(ref)
|
||||
except Exception:
|
||||
ok = True
|
||||
if not ok:
|
||||
# Out of credits — don't start this or any further IC.
|
||||
stop = True
|
||||
return
|
||||
private = ApiLogger(free=api_logger.free) if api_logger is not None else None
|
||||
try:
|
||||
result, trace = await review_ic_async(
|
||||
graph, constraints_map, ref, pdf_path,
|
||||
on_progress=on_progress, api_logger=private,
|
||||
trace_git_commit=git_commit,
|
||||
pdf_dir=pdf_dir_path, storage=storage,
|
||||
excerpt_cache=excerpt_cache,
|
||||
)
|
||||
except Exception as exc:
|
||||
msg = f"{type(exc).__name__}: {exc}"
|
||||
log.exception("Review failed for %s", ref)
|
||||
review_errors[ref] = msg
|
||||
_write_trace(_stub_trace(ref, msg), ref)
|
||||
if on_progress:
|
||||
await on_progress(ref, 0, "error", msg)
|
||||
if on_ic_error is not None:
|
||||
try:
|
||||
await on_ic_error(ref, exc)
|
||||
except Exception:
|
||||
log.exception("on_ic_error callback failed for %s", ref)
|
||||
# Persist the error into the report so the run finishes with a
|
||||
# complete picture even if every IC fails.
|
||||
try:
|
||||
_write_report(paused=False)
|
||||
except Exception:
|
||||
log.exception("incremental report write failed after error on %s", ref)
|
||||
return
|
||||
# Merge results — synchronous block, atomic under asyncio (no await
|
||||
# until the trailing callbacks), so concurrent completions can't
|
||||
# corrupt all_findings / all_coverage.
|
||||
all_findings.extend(result.findings)
|
||||
if result.checked_areas:
|
||||
all_coverage[ref] = result.checked_areas
|
||||
# Incremental write — preserves state if the process dies.
|
||||
# Never let a single IC's bad payload kill the whole pipeline.
|
||||
try:
|
||||
_write_report(paused=False)
|
||||
except Exception as exc:
|
||||
print(f"[validation] incremental write failed after {ref}: {exc}")
|
||||
all_coverage.pop(ref, None)
|
||||
if on_progress:
|
||||
await on_progress(ref, 0, "warning", f"report write failed: {exc}")
|
||||
# Per-IC trace flush — written as each IC completes so a cancel/pause
|
||||
# preserves every completed trace.
|
||||
_write_trace(trace, ref)
|
||||
if on_ic_done is not None:
|
||||
try:
|
||||
await on_ic_done(ref, result, private)
|
||||
except Exception:
|
||||
log.exception("on_ic_done callback failed for %s", ref)
|
||||
|
||||
results = await asyncio.gather(
|
||||
*(_gated_review_one(ref, pdf) for ref, pdf in ic_tasks if ref not in skip_refs),
|
||||
return_exceptions=True,
|
||||
)
|
||||
# Surface a hard cancellation so the pipeline's run handler cleans up.
|
||||
# Per-IC review failures stay isolated (captured into review_errors above).
|
||||
for r in results:
|
||||
if isinstance(r, asyncio.CancelledError):
|
||||
raise r
|
||||
|
||||
# Dedup only a *complete* run — a paused/partial run may gain more
|
||||
# findings on resume, and merging now could collapse a pair before its
|
||||
# counterpart exists.
|
||||
if not stop:
|
||||
await _maybe_dedupe_cross_ic()
|
||||
try:
|
||||
return _write_report(paused=bool(stop))
|
||||
finally:
|
||||
_cleanup_excerpt_cache()
|
||||
Reference in New Issue
Block a user