| 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) |
|
|
