from __future__ import annotations import json import re from dataclasses import dataclass from pathlib import Path from bom_assistant.ai.client import complete from bom_assistant.session.models import BomRow @dataclass class ResolvedCategory: name: str id: int def resolve_category( row: BomRow, subcats: list[dict], supplier: str, cache_path: Path, ) -> list[ResolvedCategory]: """Map a BomRow to up to 3 supplier subcategories ranked by confidence. Cache hit → instant.""" cache = _load(cache_path) key = _cache_key(row) if key in cache: return _load_entry(cache[key]) if not subcats: return [] results = _ai_resolve(row, subcats, supplier) if results: cache[key] = [{"name": r.name, "id": r.id} for r in results] _save(cache, cache_path) return results def _load_entry(e) -> list[ResolvedCategory]: if isinstance(e, dict): return [ResolvedCategory(e["name"], e["id"])] return [ResolvedCategory(x["name"], x["id"]) for x in e] def _cache_key(row: BomRow) -> str: p = row.normalized_params return f"{row.category.value}|{p.package or ''}|{p.value_str or row.raw_value}" def _fmt_subcat(s: dict) -> str: line = f" {s['name']}" desc = s.get("description", "") examples = s.get("examples", []) if desc or examples: line += f" — {desc}" if examples: line += f" e.g. {', '.join(examples)}" return line def _ai_resolve(row: BomRow, subcats: list[dict], supplier: str) -> list[ResolvedCategory]: p = row.normalized_params lines: list[str] = [f" type: {row.category.value}"] if p.value_str: lines.append(f" value: {p.value_str}") if p.package: lines.append(f" package: {p.package}") if p.voltage_rating is not None: lines.append(f" voltage_rating: {p.voltage_rating}V") if p.current_rating is not None: lines.append(f" current_rating: {p.current_rating}A") if p.part_number: lines.append(f" part_number: {p.part_number}") if row.raw_value and row.raw_value != p.value_str: lines.append(f" raw_value: {row.raw_value!r}") subcat_lines = "\n".join(_fmt_subcat(s) for s in subcats) prompt = ( f"You are an electronics sourcing expert. Rank the top 3 {supplier.upper()} subcategories for this component, best match first.\n\n" "Component:\n" + "\n".join(lines) + "\n\nAvailable subcategories:\n" + subcat_lines + '\n\nReply ONLY with JSON: {"subcategories": ["", "", ""]} — use exact names from the list above, ordered by confidence' ) try: text = complete(prompt, max_tokens=150) m = re.search(r'"subcategories"\s*:\s*\[([^\]]+)\]', text, re.DOTALL) if not m: return [] names = re.findall(r'"([^"]+)"', m.group(1)) results = [] for name in names[:3]: r = _lookup(name.strip(), subcats) if r: results.append(r) return results except Exception as e: print(f" [resolver error] {e}") return [] def _lookup(chosen: str, subcats: list[dict]) -> ResolvedCategory | None: cl = chosen.lower() for s in subcats: if s["name"].lower() == cl: return ResolvedCategory(s["name"], s["id"]) for s in subcats: if cl in s["name"].lower() or s["name"].lower() in cl: return ResolvedCategory(s["name"], s["id"]) return None def _load(path: Path) -> dict: return json.loads(path.read_text(encoding="utf-8")) if path.exists() else {} def _save(cache: dict, path: Path) -> None: path.write_text(json.dumps(cache, indent=2, ensure_ascii=False), encoding="utf-8")