fn_registry/python/functions/datascience/extract_relations_llm_test.py

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