Spaces:
Sleeping
Sleeping
File size: 2,311 Bytes
2e818da | 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 | from app.agents.d3.extractor import D3DataExtractor
from app.agents.d3.registry import TEMPLATES
from app.agents.d3 import templates # noqa: F401
from app.agents.d3.templates._hierarchy import HierNode, HierarchyData
class _FakeClient:
def __init__(self, obj): self._obj = obj
def structured_complete(self, messages, schema, **kw): return self._obj
def test_fill_returns_schema_instance():
t = TEMPLATES["bar_chart"]
filled = t.schema(title="GDP", x_label="Country", y_label="T$", data=[{"name": "US", "value": 21}])
ex = D3DataExtractor(client=_FakeClient(filled))
out = ex.fill(t, "GDP", [{"source": "s", "text": "US 21 trillion"}], "graduate")
assert out is not None and out.data[0].name == "US"
def test_fill_never_declines_when_model_returns_empty_data():
# D3DataExtractor.fill() no longer declines on an empty schema instance -- the
# extraction prompt now asks the model to synthesize plausible illustrative data
# itself rather than leave the schema empty, so fill() always returns what the
# model gave back, whatever its shape.
t = TEMPLATES["bar_chart"]
empty = t.schema(data=[])
ex = D3DataExtractor(client=_FakeClient(empty))
out = ex.fill(t, "x", [{"source": "s", "text": "no numbers here"}], "graduate")
assert out is empty
def test_fill_never_declines_hierarchy_root_with_no_children_and_no_value():
empty_root = HierNode(name="Root")
filled = HierarchyData(title="Empty", root=empty_root)
ex = D3DataExtractor(client=_FakeClient(filled))
out = ex.fill(
TEMPLATES["bar_chart"],
"x",
[{"source": "s", "text": "no hierarchy here"}],
"graduate",
)
assert out is filled
def test_fill_accepts_hierarchy_root_with_value_as_valid():
# A single-node root WITH a value (or with children) must NOT be declined —
# per the plan, decline only when root.children == [] AND root.value is None.
valued_root = HierNode(name="Root", value=42)
filled = HierarchyData(title="Single node", root=valued_root)
ex = D3DataExtractor(client=_FakeClient(filled))
out = ex.fill(
TEMPLATES["bar_chart"],
"x",
[{"source": "s", "text": "root has a value"}],
"graduate",
)
assert out is not None
assert out.root.value == 42
|