GTI_GERENCIADOR_FISCAL/fiscal/central/tests/test_ai_semantic.py

"""AI semantic layer V1: fingerprint, KB, dedup, provider contract, budget guard, candidates."""
import json

from sqlalchemy import create_engine, select
from sqlalchemy.orm import Session

from central.app.models import Base, Dataset, NcmEntry, CestEntry, CestNcmEntry, TipiEntry


def _engine():
    engine = create_engine("sqlite://")
    Base.metadata.create_all(engine)
    return engine


def _ncm_dataset(session, entries):
    dataset = Dataset(kind="NCM", provider="Siscomex", sha256="a" * 64, version_label="v", act="a", source_url="u", status="ACTIVE", record_count=len(entries))
    session.add(dataset); session.flush()
    for code, description in entries:
        session.add(NcmEntry(dataset_id=dataset.id, code=code, description=description, valid_from="", valid_to=""))
    session.flush()
    return dataset


# ---------------- fingerprint / identity ----------------

def test_fingerprint_is_stable_and_ignores_unstable_context():
    from central.app.ai.fingerprint import normalize_description, product_fingerprint, validate_gtin
    a = product_fingerprint("Coca-Cola 2LT", "COCA-COLA", "UN", "2")
    b = product_fingerprint("coca cola 2lt  ", "coca-cola", "un", "2.0")
    assert a == b
    # fiscal/company context never participates in global identity
    c = product_fingerprint("Coca-Cola 2LT", "COCA-COLA", "UN", "2")
    assert a == c
    assert normalize_description("Água Mineral 500ml") == "AGUA MINERAL 500ML"
    assert validate_gtin("7891234567895") == "7891234567895"   # valid GTIN-13
    assert validate_gtin("7891234567890") is None              # bad checksum
    assert validate_gtin("123") is None


# ---------------- knowledge base ----------------

def test_kb_exact_fingerprint_hit_and_conflict_without_majority():
    from central.app.ai.kb import KnowledgeBase, EXACT, FINGERPRINT
    engine = _engine()
    with Session(engine) as session:
        kb = KnowledgeBase(session)
        p1 = kb.upsert_product(fingerprint="f1", gtin="7891234567895", normalized_description="AGUA MINERAL 500ML")
        # same GTIN arrives with a different fingerprint -> same product, exact hit
        p2, hit = kb.resolve(gtin="7891234567895", fingerprint="other")
        assert p2.id == p1.id and hit == EXACT
        # same fingerprint, no GTIN -> fingerprint hit
        p3, hit = kb.resolve(gtin=None, fingerprint="f2")
        assert p3 is None and hit is None
        p4 = kb.upsert_product(fingerprint="f2", normalized_description="OUTRO")
        _, hit = kb.resolve(gtin=None, fingerprint="f2")
        assert hit == FINGERPRINT
        kb.add_knowledge(p1.id, "candidate_ncm", "22011000", source="OFFICIAL_DATASET", confidence=1.0)
        kb.add_knowledge(p1.id, "candidate_ncm", "22019000", source="OFFICIAL_DATASET", confidence=1.0)
        current = kb.current_knowledge(p1.id)
        assert current["candidate_ncm"]["conflict"] is True          # no majority auto-pick
        assert current["candidate_ncm"]["values"] == {"22011000", "22019000"}
        stats = kb.stats()
        assert stats["products"] == 2 and stats["gtin_products"] == 1 and stats["conflicts"] == 1


def test_kb_aggregates_observations_across_tenants():
    from central.app.ai.kb import KnowledgeBase
    engine = _engine()
    with Session(engine) as session:
        kb = KnowledgeBase(session)
        p = kb.upsert_product(fingerprint="f1", normalized_description="REFRIGERANTE 2L")
        kb.record_observation(p.id, tenant_id=1, company_key="MMMODAS", erp_code="000001", database_label="MMMODAS", observed_fields={"ncm": "22021000"})
        kb.record_observation(p.id, tenant_id=2, company_key="UNIAO", erp_code="000009", database_label="UNIAO", observed_fields={"ncm": "22021000"})
        agg = kb.observation_summary(p.id)
        assert agg["observation_count"] == 2 and agg["tenant_count"] == 2
        assert agg["agreement"]["ncm"] == {"22021000": 2} and agg["conflicts"] == {}


# ---------------- work item dedup ----------------

def test_work_items_dedup_by_knowledge_needed():
    from central.app.ai.workitems import build_work_items
    questions = [
        {"QUESTION_ID": "Q1", "DATABASE": "MMMODAS", "COMPANY": "MMMODAS", "SCOPE": "NCM", "GROUP": "", "NCM": "33059000", "KEY": "ex_tipi", "CELLS": 10},
        {"QUESTION_ID": "Q2", "DATABASE": "UNIAO", "COMPANY": "UNIAO", "SCOPE": "NCM", "GROUP": "", "NCM": "33059000", "KEY": "ex_tipi", "CELLS": 5},
        {"QUESTION_ID": "Q3", "DATABASE": "RINCAO", "COMPANY": "RINCAO", "SCOPE": "GROUP", "GROUP": "000015", "KEY": "product_role", "CELLS": 40},
        {"QUESTION_ID": "Q4", "DATABASE": "CONSTRUCAO", "COMPANY": "CONSTRUCAO", "SCOPE": "PRODUCT", "GROUP": "", "NCM": "", "KEY": "ncm", "CELLS": 8, "QUESTION": "produto PRODUCT 000054", "DATABASE_CODE": "000054"},
    ]
    items = build_work_items(questions, product_fingerprints={"CONSTRUCAO:000054": "fp-a"})
    # 4 raw questions -> 3 deduped work items (two ex_tipi questions for NCM 33059000 collapse)
    by_key = {i.dedup_key: i for i in items}
    assert len(items) == 3
    assert by_key["ex_tipi:33059000"].cells_covered == 15
    assert by_key["ex_tipi:33059000"].question_refs == ["Q1", "Q2"]
    assert by_key["product_role:RINCAO:000015"].cells_covered == 40
    assert by_key["invalid_ncm:fp-a"].cells_covered == 8
    # AI-ineligible keys never become work items
    assert all(i.task_type in ("INVALID_NCM", "CEST", "EX_TIPI", "PRODUCT_ROLE") for i in items)


def test_work_items_accept_canonical_identity_map():
    from central.app.ai.workitems import build_work_items
    questions = [
        {"QUESTION_ID": "Q1", "DATABASE": "A", "SCOPE": "PRODUCT", "KEY": "ncm", "CELLS": 1, "DATABASE_CODE": "1"},
        {"QUESTION_ID": "Q2", "DATABASE": "B", "SCOPE": "PRODUCT", "KEY": "ncm", "CELLS": 1, "DATABASE_CODE": "2"},
    ]
    items = build_work_items(questions, canonical_products={"A:1": "GTIN:7891234567895", "B:2": "GTIN:7891234567895"})
    assert len(items) == 1 and items[0].cells_covered == 2


# ---------------- provider contract / structured output / budget ----------------

def test_structured_output_schema_and_baseline_provider():
    from central.app.ai.provider import BaselineProvider, result_from_dict, SemanticOutputError
    provider = BaselineProvider()
    item = type("W", (), {"task_type": "INVALID_NCM", "dedup_key": "invalid_ncm:fp-a", "payload": {"description": "AGUA MINERAL NATURAL 500ML"}})()
    candidates = [("22011000", "AGUAS MINERAIS NATURAIS", 0.75), ("22019000", "AGUAS ... OUTRAS", 0.25)]
    result = provider.select(item, candidates, dataset_versions={"ncm": "sha-ncm"})
    assert result.task_type == "INVALID_NCM" and result.selected_candidate == "22011000"
    assert result.candidate_ids_seen == ["22011000", "22019000"]
    assert result.model == "baseline-deterministic-v1" and result.prompt_version == "ai-semantic-v1"
    assert result.dataset_versions and result.engine_version
    as_dict = result.to_dict()
    assert result_from_dict(as_dict).selected_candidate == "22011000"
    # schema fail-closed: a provider answer outside the candidate list is invalid
    try:
        result_from_dict({**as_dict, "selected_candidate": "99999999"})
        raise AssertionError("must reject")
    except SemanticOutputError:
        pass
    # no candidates at all -> NONE, never invented
    none_result = provider.select(item, [])
    assert none_result.selected_candidate == "NONE"


def test_budget_guard_blocks_run_over_limit():
    from central.app.ai.provider import BudgetGuard
    guard = BudgetGuard(max_run_cost_usd=1.0, price_in_per_m=1.0, price_out_per_m=2.0)
    plan = guard.estimate(items_in_tokens=600_000, items_out_tokens=300_000)
    assert plan["estimated_cost"] > 0
    assert guard.authorize(plan) is False           # (600k*1 + 300k*2)/1M = 1.2 > 1.0
    assert guard.authorize(guard.estimate(100_000, 50_000)) is True


# ---------------- deterministic candidates ----------------

def test_ncm_top_k_and_ex_tipi_candidates():
    from central.app.ai.candidates import ncm_top_k, ex_tipi_candidates_for_ncm, cest_candidates_for_ncm
    engine = _engine()
    with Session(engine) as session:
        _ncm_dataset(session, [("22011000", "AGUAS MINERAIS NATURAIS"), ("22019000", "AGUAS MINERAIS GASEIFICADAS OUTRAS"), ("08043000", "ABACAXIS FRESCOS")])
        top = ncm_top_k(session, "AGUA MINERAL NATURAL", k=2)
        assert [c[0] for c in top] == ["22011000", "22019000"]
        dataset = Dataset(kind="TIPI_IPI", provider="RFB", sha256="b" * 64, version_label="v", act="a", source_url="u", status="ACTIVE", record_count=2)
        session.add(dataset); session.flush()
        session.add(TipiEntry(dataset_id=dataset.id, ncm="33059000", ex=None, description="outros", ipi_rate=None, ipi_rate_text="NT", source_row=1))
        session.add(TipiEntry(dataset_id=dataset.id, ncm="33059000", ex="01", description="preparacoes para alisar", ipi_rate=None, ipi_rate_text="NT", source_row=2))
        cest_dataset = Dataset(kind="CEST", provider="CONFAZ", sha256="c" * 64, version_label="v", act="a", source_url="u", status="ACTIVE", record_count=1)
        session.add(cest_dataset); session.flush()
        cest = CestEntry(dataset_id=cest_dataset.id, code="03.999.99", description="x", annex="II")
        session.add(cest); session.flush()
        session.add(CestNcmEntry(cest_entry_id=cest.id, ncm_code="22011000"))
        session.commit()
        assert [e.ex for e in ex_tipi_candidates_for_ncm(session, "33059000")] == ["01"]
        assert [c.code for c in cest_candidates_for_ncm(session, "22011000")] == ["03.999.99"]


def test_kb_persist_roundtrip_models():
    from central.app.models import KbProduct, KbKnowledge, KbObservation, TenantProductContext, AiWorkItem, AiRun
    engine = _engine()
    with Session(engine) as session:
        p = KbProduct(fingerprint="f", normalized_description="D")
        session.add(p); session.flush()
        session.add(KbKnowledge(kb_product_id=p.id, knowledge_key="candidate_ncm", value="22011000", source="OFFICIAL_DATASET", confidence=1.0))
        session.add(KbObservation(kb_product_id=p.id, tenant_id=1, company_key="C", erp_code="1", database_label="C", observed_fields={}))
        session.add(TenantProductContext(tenant_id=1, company_key="C", erp_code="1", kb_product_id=p.id))
        session.add(AiWorkItem(task_type="CEST", dedup_key="cest:1", payload={}, question_refs=[], cells_covered=1))
        session.add(AiRun(provider="baseline", model="m", prompt_version="v1", input_tokens=1, output_tokens=1, cost_usd=0.0, status="DONE", results=[]))
        session.commit()
        assert session.query(KbProduct).count() == 1 and session.query(AiWorkItem).count() == 1