92 lines
2.7 KiB
Python
92 lines
2.7 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
from ont_platform.core.extraction.lightweight_extractor import ExtractionResult
|
||
|
|
from ont_platform.storage.candidate_repository import CandidateRepository
|
||
|
|
from ont_platform.storage.models import Base, CandidateSource, ReviewStatus
|
||
|
|
from sqlalchemy import create_engine
|
||
|
|
from sqlalchemy.orm import Session, sessionmaker
|
||
|
|
|
||
|
|
|
||
|
|
def _session() -> Session:
|
||
|
|
engine = create_engine("sqlite:///:memory:")
|
||
|
|
Base.metadata.create_all(engine)
|
||
|
|
return sessionmaker(bind=engine)()
|
||
|
|
|
||
|
|
|
||
|
|
def test_repository_saves_lightweight_candidates_with_evidence() -> None:
|
||
|
|
db = _session()
|
||
|
|
repository = CandidateRepository(db)
|
||
|
|
result = ExtractionResult(
|
||
|
|
entities=[
|
||
|
|
{
|
||
|
|
"id": "E_alice",
|
||
|
|
"label": "Alice",
|
||
|
|
"type": "person",
|
||
|
|
"confidence": 0.91,
|
||
|
|
"evidence_ids": ["EV_1"],
|
||
|
|
}
|
||
|
|
],
|
||
|
|
relations=[],
|
||
|
|
evidence_spans=[
|
||
|
|
{
|
||
|
|
"id": "EV_1",
|
||
|
|
"text": "Alice works at Acme.",
|
||
|
|
"start_offset": 0,
|
||
|
|
"end_offset": 20,
|
||
|
|
}
|
||
|
|
],
|
||
|
|
warnings=[],
|
||
|
|
)
|
||
|
|
|
||
|
|
batch = repository.save_lightweight_result(
|
||
|
|
project_id="proj_1",
|
||
|
|
document_id="doc_1",
|
||
|
|
result=result,
|
||
|
|
source_trust=0.8,
|
||
|
|
validation_passed=True,
|
||
|
|
)
|
||
|
|
|
||
|
|
assert len(batch.entities) == 1
|
||
|
|
entity = batch.entities[0]
|
||
|
|
assert entity.source_type == CandidateSource.LIGHTWEIGHT
|
||
|
|
assert entity.review_status == ReviewStatus.PENDING
|
||
|
|
assert entity.evidence_ids == ["EV_1"]
|
||
|
|
assert repository.candidate_has_valid_evidence(entity)
|
||
|
|
|
||
|
|
|
||
|
|
def test_repository_saves_ontocast_candidates_on_separate_source_path() -> None:
|
||
|
|
db = _session()
|
||
|
|
repository = CandidateRepository(db)
|
||
|
|
batch = repository.save_ontocast_result(
|
||
|
|
project_id="proj_1",
|
||
|
|
document_id="doc_1",
|
||
|
|
result={
|
||
|
|
"entities": [
|
||
|
|
{
|
||
|
|
"id": "E_graph",
|
||
|
|
"label": "GraphUpdate",
|
||
|
|
"entity_type": "concept",
|
||
|
|
"confidence": 0.72,
|
||
|
|
"evidence_ids": ["EV_graph"],
|
||
|
|
}
|
||
|
|
],
|
||
|
|
"relations": [],
|
||
|
|
"evidence_spans": [
|
||
|
|
{
|
||
|
|
"id": "EV_graph",
|
||
|
|
"text": "OntoCast proposed a graph update.",
|
||
|
|
"start_offset": 0,
|
||
|
|
"end_offset": 34,
|
||
|
|
}
|
||
|
|
],
|
||
|
|
},
|
||
|
|
source_trust=0.7,
|
||
|
|
validation_passed=False,
|
||
|
|
)
|
||
|
|
|
||
|
|
entity = batch.entities[0]
|
||
|
|
assert entity.source_type == CandidateSource.ONTOCAST
|
||
|
|
assert entity.created_by == "ontocast"
|
||
|
|
assert entity.validation_passed is False
|
||
|
|
assert entity.metadata_["source_type"] == "ontocast"
|