tariffwise / ai_layer.py
Mussez's picture
Upload 6 files
f4dcb85 verified
Raw
History Blame Contribute Delete
8.6 kB
"""TariffWise AI layer. The model does the language work.
The model has three jobs. It extracts facts from the user's description.
It phrases the engine's named gap as a plain question. It reads each
answer back into schedule terms. The engine makes every decision. No
model output assigns a code.
Provenance is recorded per fact. A fact is user_stated when extracted
from the user's own words, and answered when it came from an interview
reply. The audit record keeps both.
"""
import json
import os
from engine import walk, rationale, FACT_KEYS
MODEL_ID = "llama-3.3-70b-versatile"
BOOL_KEYS = {k for k, v in FACT_KEYS.items() if v.startswith("true")}
ENUM_KEYS = {
"sole_material": ["rubber_plastics", "leather", "composition_leather", "other"],
"upper_material": ["rubber_plastics", "leather", "composition_leather", "textile", "other"],
"upper_textile_kind": ["vegetable_fibers", "wool_felt", "other_textile"],
"gender_age": ["men", "women", "youths_boys", "misses", "children", "infants", "unisex"],
}
EXTRACT_SYSTEM = (
"You extract footwear facts for US tariff classification under HTS Chapter 64. "
"Apply the chapter notes. The upper material is the constituent material with the "
"greatest external surface area, disregarding accessories such as laces, eyelet stays, "
"and ornamentation. Textile with a visible external layer of rubber or plastics counts "
"as rubber_plastics. Return only a JSON object. Include a key only when the text "
"clearly states or strongly implies it. Omit anything not determinable.\n"
"Keys and values:\n"
+ "\n".join(
f" {k}: {'true or false, ' + v[5:] if k in BOOL_KEYS else ('one of ' + ', '.join(ENUM_KEYS[k]) if k in ENUM_KEYS else 'number, ' + v)}"
for k, v in FACT_KEYS.items()
)
)
def get_client():
from groq import Groq
key = os.environ.get("GROQ_API_KEY", "").strip()
if not key:
raise RuntimeError("GROQ_API_KEY is not set")
return Groq(api_key=key)
def _chat(client, system, user, json_mode=True, max_tokens=400):
kwargs = dict(
model=MODEL_ID,
messages=[{"role": "system", "content": system},
{"role": "user", "content": user}],
temperature=0,
max_tokens=max_tokens,
)
if json_mode:
kwargs["response_format"] = {"type": "json_object"}
out = client.chat.completions.create(**kwargs)
return out.choices[0].message.content
def normalize(raw):
"""Keep only known keys with valid values."""
clean = {}
for k, v in raw.items():
if k in BOOL_KEYS:
if isinstance(v, bool):
clean[k] = v
elif str(v).lower() in ("true", "yes"):
clean[k] = True
elif str(v).lower() in ("false", "no"):
clean[k] = False
elif k in ENUM_KEYS:
if str(v).lower() in ENUM_KEYS[k]:
clean[k] = str(v).lower()
elif k == "value_per_pair":
try:
clean[k] = float(v)
except (TypeError, ValueError):
pass
return clean
def extract_facts(client, text):
raw = _chat(client, EXTRACT_SYSTEM, text)
try:
return normalize(json.loads(raw))
except json.JSONDecodeError:
return {}
def ask_gap(client, fact_key):
"""Phrase the engine's named gap as one plain question."""
system = ("You interview a small importer about one footwear product. "
"Ask one short plain question that determines the stated fact. "
"One sentence. No preamble.")
user = f"The fact to determine: {FACT_KEYS[fact_key]}."
return _chat(client, system, user, json_mode=False, max_tokens=60).strip()
def ask_branch(client, branch_text):
"""Phrase an unparsed schedule branch as a yes or no question."""
system = ("You interview a small importer about one footwear product. "
"Restate the tariff condition below as one plain yes or no question "
"about their product. One sentence. No preamble.")
return _chat(client, system, f"Condition: {branch_text}", json_mode=False, max_tokens=80).strip()
def read_answer(client, fact_key, answer):
"""Read an interview answer back into the one fact it addresses."""
user = f"The question was about: {FACT_KEYS[fact_key]}.\nThe importer answered: {answer}"
raw = _chat(client, EXTRACT_SYSTEM, user)
try:
got = normalize(json.loads(raw))
except json.JSONDecodeError:
return {}
return {fact_key: got[fact_key]} if fact_key in got else {}
def read_branch_answer(client, branch_text, answer):
"""Read a yes or no reply to a branch condition. Returns True, False, or None."""
system = ("Decide whether the importer's answer means the condition holds. "
'Return only JSON like {"holds": true} or {"holds": false}. '
'Return {"holds": null} when the answer does not settle it.')
raw = _chat(client, system, f"Condition: {branch_text}\nAnswer: {answer}")
try:
v = json.loads(raw).get("holds")
return v if isinstance(v, bool) else None
except json.JSONDecodeError:
return None
# ---------------------------------------------------------------------
# The interview loop. ask_user is any callable that shows a question and
# returns the reply, which keeps the loop usable in a notebook or a web app.
# ---------------------------------------------------------------------
MAX_TURNS = 12
def classify_interactive(client, description, ask_user):
facts = extract_facts(client, description)
provenance = {k: "user_stated" for k in facts}
transcript = []
for _ in range(MAX_TURNS):
r = walk(facts)
if r["status"] == "code":
return {"status": "code", "code": r["code"], "rate": r["general"],
"facts": facts, "provenance": provenance,
"transcript": transcript, "path": r["path"],
"rationale": rationale(r, facts)}
if r["status"] == "need_fact":
q = ask_gap(client, r["fact"])
a = ask_user(q)
transcript.append({"q": q, "a": a, "about": r["fact"]})
got = read_answer(client, r["fact"], a)
if got:
facts.update(got)
provenance[r["fact"]] = "answered"
continue
q2 = f"{q} Please answer directly."
a2 = ask_user(q2)
transcript.append({"q": q2, "a": a2, "about": r["fact"]})
got = read_answer(client, r["fact"], a2)
if got:
facts.update(got)
provenance[r["fact"]] = "answered"
continue
return {"status": "review", "reason": f"the fact {r['fact']} could not be established",
"facts": facts, "provenance": provenance, "transcript": transcript}
if r["status"] == "review" and "at" in r:
q = ask_branch(client, r["at"])
a = ask_user(q)
transcript.append({"q": q, "a": a, "about": r["at"][:60]})
holds = read_branch_answer(client, r["at"], a)
if holds is None:
return {"status": "review",
"reason": "a schedule condition needs expert judgment",
"at": r["at"], "facts": facts,
"provenance": provenance, "transcript": transcript}
facts.setdefault("_overrides", {})[r["at"]] = holds
provenance[f"branch: {r['at'][:50]}"] = "answered"
continue
return {"status": "review", "reason": r.get("reason", "unresolved"),
"facts": facts, "provenance": provenance, "transcript": transcript}
return {"status": "review", "reason": "the interview exceeded its turn limit",
"facts": facts, "provenance": provenance, "transcript": transcript}
def audit_record(description, result):
rec = {
"input": description,
"status": result["status"],
"code": result.get("code"),
"rate": result.get("rate"),
"facts": {k: v for k, v in result["facts"].items() if k != "_overrides"},
"provenance": result["provenance"],
"interview": result["transcript"],
"path": [f"{s['htsno'] or 'grouping'} {s['desc']}" for s in result.get("path", [])],
"basis": "HTS Chapter 64, 2026 revision, walked per the chapter notes",
"disclaimer": "Decision support only. The importer or a licensed broker makes the final decision.",
}
return rec