File size: 3,227 Bytes
66ad25b | 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 | from pluto.stages.route import run_route
from pluto.tracer import Tracer
class _FakeTools:
def __init__(self):
self.doc_index = None
self._chunks = {
"doc_a": [
"Introduction to the system architecture and design.",
"Method details for the proposed coordinator pipeline.",
"Results show 0% ASR across defended scenarios.",
"Additional background information with little relevance.",
"Conclusion and limitations of the coordinator pipeline.",
"Appendix references and citations.",
],
"doc_b": [
"Dataset description and benchmark setup.",
"Performance comparison against baselines.",
"More benchmark discussion and analysis.",
"Implementation details and ablations.",
"Future work and limitations.",
"References section.",
],
}
def list_docs(self):
return [{"doc_id": doc_id, "filename": f"{doc_id}.md"} for doc_id in self._chunks]
def search(self, query, filters=None):
del query
results = [{"doc_id": "doc_a", "score": 0.9}, {"doc_id": "doc_b", "score": 0.7}]
if filters and filters.get("doc_ids"):
allowed = set(filters["doc_ids"])
results = [item for item in results if item["doc_id"] in allowed]
return results
def get_all_chunks(self, doc_id):
return self._chunks[doc_id]
def test_route_limits_extractions_to_relevant_subset():
tools = _FakeTools()
tracer = Tracer()
route_out = run_route("coordinator pipeline architecture results", tools, tracer)
total_chunks = sum(len(chunks) for chunks in tools._chunks.values())
assert len(route_out.chunk_plan) < total_chunks
assert len(route_out.chunk_plan) <= route_out.budgets.max_extractions
assert len(route_out.chunk_plan) <= 12
assert any(chunk.doc_id == "doc_a" for chunk in route_out.chunk_plan)
assert any(chunk.doc_id == "doc_b" for chunk in route_out.chunk_plan)
def test_route_prefers_directly_named_document():
tools = _FakeTools()
tracer = Tracer()
route_out = run_route("Summarize doc_a and its architecture contributions", tools, tracer)
assert {doc.doc_id for doc in route_out.doc_scope} == {"doc_a"}
def test_route_respects_selected_doc_scope():
tools = _FakeTools()
tracer = Tracer()
route_out = run_route(
"Compare methodology and results",
tools,
tracer,
selected_doc_ids=["doc_b"],
)
assert [doc.doc_id for doc in route_out.doc_scope] == ["doc_b"]
assert all(chunk.doc_id == "doc_b" for chunk in route_out.chunk_plan)
def test_route_detailed_mode_expands_chunk_budget():
tools = _FakeTools()
tracer = Tracer()
standard = run_route("Compare methodology and results", tools, tracer, selected_doc_ids=["doc_a", "doc_b"])
detailed = run_route(
"Compare methodology and results",
tools,
tracer,
selected_doc_ids=["doc_a", "doc_b"],
detail_level="detailed",
)
assert len(detailed.chunk_plan) >= len(standard.chunk_plan)
|