"""Tests para extract_relations_llm.""" import sys import os # Rutas para importar desde el registry REGISTRY_ROOT = os.path.join(os.path.dirname(__file__), "..", "..", "..", "") sys.path.insert(0, REGISTRY_ROOT) sys.path.insert(0, os.path.dirname(__file__)) from python.types.datascience.entity_candidate import EntityCandidate from python.types.datascience.relation_candidate import RelationCandidate from extract_relations_llm import extract_relations_llm def _make_entity(name: str, type_label: str = "Entity") -> EntityCandidate: return EntityCandidate(name=name, type_label=type_label, confidence=0.9) def _make_llm(response: dict): """Crea un stub de llm_chat_json que retorna la respuesta fija.""" def llm_chat_json(messages: list[dict]) -> dict: return response return llm_chat_json def test_texto_con_dos_entidades_relacionadas(): entities = [_make_entity("Acme Corp", "Organization"), _make_entity("John Smith", "Person")] relation_types = ["employs", "funds", "related_to"] llm_response = { "relations": [ { "from_name": "Acme Corp", "to_name": "John Smith", "relation_type": "employs", "description": "Acme Corp employs John Smith as CEO", "confidence": 0.9, } ] } result = extract_relations_llm( text="Acme Corp employs John Smith as CEO.", entities=entities, relation_types=relation_types, llm_chat_json=_make_llm(llm_response), ) assert len(result) == 1 rel = result[0] assert rel.from_name == "Acme Corp" assert rel.to_name == "John Smith" assert rel.relation_type == "employs" assert rel.confidence == 0.9 assert "CEO" in rel.description def test_texto_con_entidades_pero_sin_relacion(): entities = [_make_entity("Alice", "Person"), _make_entity("Bob", "Person")] relation_types = ["funds", "employs"] llm_response = {"relations": []} result = extract_relations_llm( text="Alice and Bob both attended the conference.", entities=entities, relation_types=relation_types, llm_chat_json=_make_llm(llm_response), ) assert result == [] def test_menos_de_dos_entidades_retorna_lista_vacia(): entities = [_make_entity("Solo Corp", "Organization")] relation_types = ["employs", "funds"] # El LLM nunca deberia ser llamado, pero si lo fuera retornaria relaciones llm_response = { "relations": [ {"from_name": "Solo Corp", "to_name": "Nobody", "relation_type": "employs", "confidence": 0.9} ] } result = extract_relations_llm( text="Solo Corp is a company.", entities=entities, relation_types=relation_types, llm_chat_json=_make_llm(llm_response), ) assert result == [] def test_llm_inventa_entidad_que_no_existe_se_descarta(): entities = [_make_entity("Alice", "Person"), _make_entity("Bob", "Person")] relation_types = ["funds", "employs", "related_to"] llm_response = { "relations": [ # Valida — Alice y Bob existen { "from_name": "Alice", "to_name": "Bob", "relation_type": "funds", "description": "Alice funds Bob", "confidence": 0.8, }, # Invalida — "Charlie" no esta en entities { "from_name": "Alice", "to_name": "Charlie", "relation_type": "employs", "description": "Alice employs Charlie", "confidence": 0.7, }, # Invalida — "Unknown Corp" no esta en entities { "from_name": "Unknown Corp", "to_name": "Bob", "relation_type": "related_to", "description": "...", "confidence": 0.6, }, ] } result = extract_relations_llm( text="Alice funds Bob. Alice also employs Charlie from Unknown Corp.", entities=entities, relation_types=relation_types, llm_chat_json=_make_llm(llm_response), ) # Solo la primera relacion es valida assert len(result) == 1 assert result[0].from_name == "Alice" assert result[0].to_name == "Bob" assert result[0].relation_type == "funds"