Spaces:
Sleeping
Sleeping
File size: 2,802 Bytes
14184e3 | 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 | """Tests for the parallel-streaming curriculum building blocks (Phase 3)."""
from app.agents.brain_agent import (
BrainAgent,
_ExpansionChild,
_RootAndSections,
_SectionExpansion,
_SectionItem,
)
class _CapturingClient:
"""Captures the messages passed for the multi-doc prompt assertion."""
def __init__(self):
self.messages = None
def structured_complete(self, messages, output_model, model=None):
self.messages = messages
return _RootAndSections(root_label="X", sections=[_SectionItem(label="S", source_docs=[1])])
def test_multi_doc_prompt_lists_documents_and_section_tagging():
cap = _CapturingClient()
brain = BrainAgent(client=cap)
rs = brain.derive_root_and_sections("structure", "high_school", "", "", ["paperA.pdf", "paperB.pdf"])
blob = " ".join(m["content"] for m in cap.messages)
assert "paperA.pdf" in blob and "paperB.pdf" in blob
assert "source_docs" in blob
assert rs.sections[0].source_docs == [1]
def test_single_doc_prompt_omits_doc_listing():
cap = _CapturingClient()
BrainAgent(client=cap).derive_root_and_sections("structure", "high_school", "", "", ["only.pdf"])
blob = " ".join(m["content"] for m in cap.messages)
assert "source_docs" not in blob # no multi-doc tagging instruction for a single paper
class _FakeClient:
"""Returns canned structured output based on the requested model."""
def structured_complete(self, messages, output_model, model=None):
if output_model is _RootAndSections:
return _RootAndSections(
root_label="Optimization",
root_description="Methods to minimize loss",
sections=[_SectionItem(label="Gradient Descent"), _SectionItem(label="Adam")],
)
if output_model is _SectionExpansion:
return _SectionExpansion(
children=[_ExpansionChild(label="Learning Rate"), _ExpansionChild(label="Momentum")]
)
raise AssertionError(f"unexpected model {output_model}")
def test_derive_root_and_sections():
brain = BrainAgent(client=_FakeClient())
rs = brain.derive_root_and_sections("doc structure", "high_school", "Optimization")
assert rs.root_label == "Optimization"
assert [s.label for s in rs.sections] == ["Gradient Descent", "Adam"]
def test_expand_section():
brain = BrainAgent(client=_FakeClient())
doc_excerpts = [{"index": 0, "filename": "doc.pdf", "structure_text": "doc structure"}]
exp = brain.expand_section("Gradient Descent", doc_excerpts, "high_school")
assert len(exp.children) == 2
assert exp.children[0].label == "Learning Rate"
# relationship defaults are valid
assert exp.children[0].relationship in {"prerequisite", "related", "builds-on"}
|