55 lines
1.9 KiB
Python
55 lines
1.9 KiB
Python
"""Step 7 tests — the retriever exposed as a citation-bearing data source.
|
|
|
|
The in-process GA tool path and the (GA-present) MCP path both wrap the same retrieve()
|
|
core, so their structuredContent shape is identical and every chunk maps into a
|
|
provenance.Citation. Pattern: tests/test_retrieval.py + tests/test_backends.py.
|
|
"""
|
|
|
|
import pytest
|
|
|
|
from agent_framework import FunctionTool
|
|
|
|
from portfolio_optimiser.datasource import (
|
|
build_mcp_server,
|
|
chunk_dict_to_citation,
|
|
make_retrieval_tool,
|
|
retrieve_chunks,
|
|
)
|
|
from portfolio_optimiser.provenance import Citation
|
|
|
|
|
|
@pytest.fixture()
|
|
def docs(tmp_path):
|
|
d = tmp_path / "docs"
|
|
d.mkdir()
|
|
(d / "asphalt.txt").write_text(
|
|
"Asphalt Ab11 unit rate renegotiation reduced the paving cost on the school stretch.",
|
|
encoding="utf-8",
|
|
)
|
|
return d
|
|
|
|
|
|
def test_in_process_tool_returns_citation_ready_chunks(docs) -> None:
|
|
chunks = retrieve_chunks("asphalt paving cost", str(docs), top_k=3)
|
|
assert chunks # at least one citation-ready chunk
|
|
for d in chunks:
|
|
assert set(d) >= {"file", "locator", "snippet", "score"}
|
|
cit = chunk_dict_to_citation(d)
|
|
assert isinstance(cit, Citation)
|
|
assert cit.locator.start_index <= cit.locator.end_index
|
|
assert cit.snippet
|
|
|
|
|
|
def test_make_retrieval_tool_builds_ga_function_tool(docs) -> None:
|
|
t = make_retrieval_tool(str(docs))
|
|
assert isinstance(t, FunctionTool)
|
|
assert t.name == "retrieve_cost_docs"
|
|
|
|
|
|
async def test_mcp_wrapper_returns_same_structuredcontent_shape(docs) -> None:
|
|
query = "asphalt paving cost"
|
|
expected = retrieve_chunks(query, str(docs), top_k=3)
|
|
server = build_mcp_server(str(docs), top_k=3)
|
|
_content, structured = await server.call_tool("retrieve_cost_docs", {"query": query})
|
|
# FastMCP wraps a list return as {"result": [...]} structuredContent.
|
|
assert structured["result"] == expected
|