instant-bom/bom_assistant/suppliers/category_resolver.py

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,
) -> ResolvedCategory | None:
    """Map a BomRow to a supplier subcategory using LLM. Cache hit → instant; miss → LLM call → cached."""
    cache = _load(cache_path)
    key = _cache_key(row)
    if key in cache:
        e = cache[key]
        return ResolvedCategory(e["name"], e["id"])
    if not subcats:
        return None
    result = _ai_resolve(row, subcats, supplier)
    if result:
        cache[key] = {"name": result.name, "id": result.id}
        _save(cache, cache_path)
    return result


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 _ai_resolve(row: BomRow, subcats: list[dict], supplier: str) -> ResolvedCategory | None:
    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(f"  {s['name']}" for s in subcats)
    prompt = (
        f"You are an electronics sourcing expert. Pick the best {supplier.upper()} subcategory for this component.\n\n"
        "Component:\n"
        + "\n".join(lines)
        + "\n\nAvailable subcategories:\n"
        + subcat_lines
        + '\n\nReply ONLY with JSON: {"subcategory": "<exact name from list above>"}'
    )
    try:
        text = complete(prompt, max_tokens=100)
        m = re.search(r'"subcategory"\s*:\s*"([^"]+)"', text)
        if not m:
            return None
        return _lookup(m.group(1).strip(), subcats)
    except Exception:
        return None


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")