Spaces:
Running
Running
File size: 7,163 Bytes
ac68ae4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 | from src.core.saving.architect_batch import (
BatchEndpoint,
BatchExtractResponse,
BatchRelationship,
event_leg_incomplete,
validate_batch_extract,
)
def _entity(uuid: str, name: str, type_: str):
return type("E", (), {"uuid": uuid, "name": name, "type": type_})()
def test_validate_accepts_grounded_event_hub():
source = "Alice hired Bob as a contractor in May 2024 for $5000."
entities = [
_entity("u1", "Alice", "PERSON"),
_entity("u2", "Hiring", "EVENT"),
_entity("u3", "Bob", "PERSON"),
]
payload = BatchExtractResponse(
relationships=[
BatchRelationship(
tail=BatchEndpoint(uuid="u1", name="Alice", type="PERSON"),
tip=BatchEndpoint(uuid="u2", name="Hiring", type="EVENT"),
name="MADE",
description="Alice hired Bob",
amount=5000,
source_span="Alice hired Bob as a contractor",
happened_at="2024-05",
),
BatchRelationship(
tail=BatchEndpoint(uuid="u2", name="Hiring", type="EVENT"),
tip=BatchEndpoint(uuid="u3", name="Bob", type="PERSON"),
name="TARGETED",
description="Alice hired Bob",
amount=5000,
source_span="Alice hired Bob as a contractor in May 2024",
),
]
)
result = validate_batch_extract(
payload, source_text=source, scout_entities=entities
)
assert len(result.accepted) == 2
assert result.rejected == []
assert not event_leg_incomplete(result.accepted, entities)
def test_validate_rejects_type_named_and_ungrounded():
source = "Alice met Bob at the park."
entities = [
_entity("u1", "Alice", "PERSON"),
_entity("u2", "Bob", "PERSON"),
]
payload = BatchExtractResponse(
relationships=[
BatchRelationship(
tail=BatchEndpoint(uuid="u1", name="PERSON", type="PERSON"),
tip=BatchEndpoint(uuid="u2", name="Bob", type="PERSON"),
name="MET",
source_span="Alice met Bob",
),
BatchRelationship(
tail=BatchEndpoint(uuid="u1", name="Alice", type="PERSON"),
tip=BatchEndpoint(uuid="u2", name="Bob", type="PERSON"),
name="MET",
source_span="totally fabricated zebra claim",
),
BatchRelationship(
# well-formed unknown endpoint is soft-admitted (no escalate)
tail=BatchEndpoint(uuid="missing", name="Carol", type="PERSON"),
tip=BatchEndpoint(uuid="u2", name="Bob", type="PERSON"),
name="MET",
source_span="Alice met Bob at the park",
),
]
)
result = validate_batch_extract(
payload, source_text=source, scout_entities=entities
)
reasons = {issue.reason for issue in result.rejected}
assert "type_named_placeholder_endpoint" in reasons
assert "source_span_not_grounded" in reasons
assert "unknown_tail_endpoint" not in reasons
assert len(result.accepted) == 1
assert result.accepted[0].tail.name == "Carol"
assert any(n.name == "Carol" for n in result.new_nodes)
def test_validate_admits_unknown_well_formed_endpoints():
source = "Alice introduced Carol to Bob at the park."
entities = [
_entity("u1", "Alice", "PERSON"),
_entity("u2", "Bob", "PERSON"),
]
payload = BatchExtractResponse(
relationships=[
BatchRelationship(
tail=BatchEndpoint(name="Carol", type="PERSON"),
tip=BatchEndpoint(uuid="u2", name="Bob", type="PERSON"),
name="MET",
source_span="Alice introduced Carol to Bob at the park",
),
]
)
result = validate_batch_extract(
payload, source_text=source, scout_entities=entities
)
assert result.usable
assert len(result.accepted) == 1
assert result.rejected == []
def test_validate_span_offsets_and_scout_uuid_reuse():
source = "On Tuesday Alice paid Bob 12 dollars."
entities = [_entity("a", "Alice", "PERSON"), _entity("b", "Bob", "PERSON")]
span = "Alice paid Bob 12 dollars"
start = source.index(span)
payload = {
"new_nodes": [],
"relationships": [
{
"tail": {"uuid": "a", "name": "Alice", "type": "PERSON"},
"tip": {"uuid": "b", "name": "Bob", "type": "PERSON"},
"name": "PAID",
"amount": 12,
"source_span": span,
"span_start": start,
"span_end": start + len(span),
}
],
}
result = validate_batch_extract(
payload, source_text=source, scout_entities=entities
)
assert len(result.accepted) == 1
assert result.accepted[0].amount == 12
def test_event_leg_incomplete_detects_untouched_event():
entities = [
_entity("e1", "Trip", "EVENT"),
_entity("p1", "Alice", "PERSON"),
]
rels = [
BatchRelationship(
tail=BatchEndpoint(uuid="p1", name="Alice", type="PERSON"),
tip=BatchEndpoint(uuid="p1", name="Alice", type="PERSON"),
name="SELF",
source_span="Alice",
)
]
assert event_leg_incomplete(rels, entities)
def test_validate_partial_keeps_usable_and_drops_bad_new_nodes():
source = "Alice hired Bob as a contractor in May 2024."
entities = [
_entity("u1", "Alice", "PERSON"),
_entity("u2", "Hiring", "EVENT"),
_entity("u3", "Bob", "PERSON"),
]
payload = BatchExtractResponse(
new_nodes=[
BatchEndpoint(name="PERSON", type="PERSON"),
BatchEndpoint(name="Contract", type="DOCUMENT"),
],
relationships=[
BatchRelationship(
tail=BatchEndpoint(uuid="u1", name="Alice", type="PERSON"),
tip=BatchEndpoint(uuid="u2", name="Hiring", type="EVENT"),
name="MADE",
source_span="Alice hired Bob as a contractor",
),
BatchRelationship(
tail=BatchEndpoint(uuid="u2", name="Hiring", type="EVENT"),
tip=BatchEndpoint(uuid="missing", name="Carol", type="PERSON"),
name="TARGETED",
source_span="Alice hired Bob as a contractor in May 2024",
),
],
)
result = validate_batch_extract(
payload, source_text=source, scout_entities=entities
)
assert result.usable
assert not result.ok
assert len(result.accepted) == 2
assert {r.name for r in result.accepted} == {"MADE", "TARGETED"}
assert len(result.new_nodes) >= 2
assert any(n.name == "Contract" for n in result.new_nodes)
assert any(n.name == "Carol" for n in result.new_nodes)
reasons = {issue.reason for issue in result.rejected}
assert "unknown_tip_endpoint" not in reasons
assert any(r.startswith("new_node_type_named_placeholder:") for r in reasons)
|