117 lines
5.5 KiB
Python
117 lines
5.5 KiB
Python
"""Semantic candidate retrieval V2: hierarchical official corpus + KB-neighbor bridge + hybrid ranking."""
|
|
from sqlalchemy import create_engine
|
|
from sqlalchemy.orm import Session
|
|
|
|
from central.app.models import Base, Dataset, KbKnowledge, KbObservation, KbProduct, NcmEntry, Tenant
|
|
|
|
|
|
def _engine():
|
|
engine = create_engine("sqlite://")
|
|
Base.metadata.create_all(engine)
|
|
return engine
|
|
|
|
|
|
def _seed(session):
|
|
ncm = Dataset(kind="NCM", provider="Siscomex", sha256="n" * 64, version_label="v", act="a", source_url="u", status="ACTIVE", record_count=7)
|
|
session.add(ncm); session.flush()
|
|
rows = [
|
|
("22", "BEBIDAS, LIQUIDOS ALCOOLICOS E VINAGRES"),
|
|
("2202", "AGUAS, INCLUIDAS AS AGUAS GASEIFICADAS E MINERAIS"),
|
|
("2202.10", "AGUAS MINERAIS E GASEIFICADAS"),
|
|
("22021000", "AGUAS MINERAIS NAO GASEIFICADAS"),
|
|
("61", "VESTUARIO E SEUS ACESSORIOS, DE MALHA"),
|
|
("6105", "CAMISAS DE MALHA, DE USO MASCULINO"),
|
|
("61051000", "DE ALGODAO"),
|
|
]
|
|
for code, description in rows:
|
|
session.add(NcmEntry(dataset_id=ncm.id, code=code.replace(".", ""), description=description, valid_from="", valid_to=""))
|
|
session.flush()
|
|
# two tenants already selling mineral water with a validated NCM
|
|
tenant_a, tenant_b = Tenant(name="A"), Tenant(name="B")
|
|
session.add_all([tenant_a, tenant_b]); session.flush()
|
|
p1 = KbProduct(fingerprint="f1", normalized_description="AGUA MINERAL SEM GAS 500ML")
|
|
p2 = KbProduct(fingerprint="f2", normalized_description="AGUA MINERAL GARRAFA 1L")
|
|
p3 = KbProduct(fingerprint="f3", normalized_description="CAMISA MASCULINA MALHA ALGODAO")
|
|
session.add_all([p1, p2, p3]); session.flush()
|
|
for product, ncm_value, tenant in ((p1, "22021000", tenant_a), (p2, "22021000", tenant_b), (p3, "61051000", tenant_a)):
|
|
session.add(KbKnowledge(kb_product_id=product.id, knowledge_key="validated_ncm", value=ncm_value, source="OFFICIAL_DATASET", confidence=1.0))
|
|
session.add(KbObservation(kb_product_id=p1.id, tenant_id=tenant_a.id, company_key="A", erp_code="1", database_label="A", observed_fields={}))
|
|
session.add(KbObservation(kb_product_id=p2.id, tenant_id=tenant_b.id, company_key="B", erp_code="2", database_label="B", observed_fields={}))
|
|
session.add(KbObservation(kb_product_id=p3.id, tenant_id=tenant_a.id, company_key="A", erp_code="3", database_label="A", observed_fields={}))
|
|
session.commit()
|
|
return ncm.sha256
|
|
|
|
|
|
def test_enriched_corpus_carries_ancestor_context():
|
|
from central.app.ai.retrieval import EnrichedNcmCorpus
|
|
engine = _engine()
|
|
with Session(engine) as session:
|
|
_seed(session)
|
|
corpus = EnrichedNcmCorpus.build(session)
|
|
doc = corpus.document("22021000")
|
|
# chapter + position + subposition context enriches the very short legal text (stemmed)
|
|
assert "BEBIDA" in doc and "AGUA" in doc
|
|
assert "VESTUARIO" not in doc
|
|
|
|
|
|
def test_hierarchical_retrieval_prefers_right_positions():
|
|
from central.app.ai.retrieval import EnrichedNcmCorpus, hierarchical_candidates
|
|
engine = _engine()
|
|
with Session(engine) as session:
|
|
_seed(session)
|
|
corpus = EnrichedNcmCorpus.build(session)
|
|
codes = [code for code, _, _ in hierarchical_candidates(corpus, "AGUA MINERAL", k=2)]
|
|
assert codes[0] == "22021000"
|
|
|
|
|
|
def test_kb_neighbor_bridge_maps_commercial_to_ncm():
|
|
from central.app.ai.retrieval import KbNeighborIndex, kb_neighbor_candidates
|
|
engine = _engine()
|
|
with Session(engine) as session:
|
|
_seed(session)
|
|
index = KbNeighborIndex.build(session)
|
|
neighbors = kb_neighbor_candidates(index, "AGUA MINERAL PET 500ML", k=5)
|
|
assert neighbors and neighbors[0][0] == "22021000"
|
|
# cross-tenant agreement is reported, never auto-majority
|
|
assert neighbors[0][3] >= 2
|
|
|
|
|
|
def test_hybrid_merges_official_and_kb_rankings():
|
|
from central.app.ai.retrieval import EnrichedNcmCorpus, KbNeighborIndex, hybrid_candidates
|
|
engine = _engine()
|
|
with Session(engine) as session:
|
|
_seed(session)
|
|
corpus = EnrichedNcmCorpus.build(session)
|
|
index = KbNeighborIndex.build(session)
|
|
hybrid = hybrid_candidates(corpus, index, "CAMISA MASCULINA ALGODAO", k=2)
|
|
assert hybrid[0][0] == "61051000"
|
|
scores = hybrid_candidates(corpus, index, "AGUA MINERAL SEM GAS", k=2)
|
|
assert scores[0][0] == "22021000"
|
|
|
|
|
|
def test_lexicon_expands_abbreviations_supported_by_kb():
|
|
from central.app.ai.retrieval import derive_abbreviation_lexicon, expand_stems
|
|
corpus = {
|
|
"REFRIGERANTE COCA 2L": "22021000",
|
|
"REFRI GUARANA 1L": "22021000",
|
|
"REFRI LARANJA 350ML": "22021000",
|
|
"CAMISA POLO ALGODAO": "61051000",
|
|
"CAM POLO MASC": "61051000",
|
|
"CAM POLO AZUL": "61051000",
|
|
"BOTA COURO MASC": "64039900",
|
|
}
|
|
lexicon = derive_abbreviation_lexicon(corpus, min_support=2)
|
|
assert lexicon.get("REFRI") == "REFRIGERANTE"
|
|
assert lexicon.get("CAM") == "CAMISA"
|
|
expanded = expand_stems({"REFRI", "COLA"}, lexicon)
|
|
assert "REFRIGERANTE" in expanded
|
|
|
|
|
|
def test_commercial_tokens_strip_units_and_sizes():
|
|
from central.app.ai.retrieval import commercial_tokens
|
|
tokens = commercial_tokens("CREME TRATAMENTO NOVEX 1KG CX 12 COR AZUL TAM M REF 453")
|
|
assert "NOVEX" in tokens and "CREME" in tokens and "TRATAMENTO" in tokens
|
|
assert "1KG" not in tokens and "CX" not in tokens and "COR" not in tokens and "TAM" not in tokens
|
|
assert "AZUL" not in tokens and "REF" not in tokens and "453" not in tokens
|
|
assert "12" not in tokens
|