feat(fase2): expose retriever as citation-bearing data source
This commit is contained in:
parent
3104e18b07
commit
f981c51e0b
2 changed files with 131 additions and 0 deletions
76
src/portfolio_optimiser/datasource.py
Normal file
76
src/portfolio_optimiser/datasource.py
Normal file
|
|
@ -0,0 +1,76 @@
|
|||
"""Expose the framework-agnostic retriever to the agents as a citation-bearing data source.
|
||||
|
||||
**MVP (GA-safe) path:** an in-process GA ``FunctionTool`` over ``retrieval.retrieve()`` whose
|
||||
chunks the orchestrator maps into ``provenance.Citation`` — zero new runtime dependency,
|
||||
D7-portable. Path-security lives in ``retrieval.py`` (Step 5).
|
||||
|
||||
Because ``mcp`` resolved as a GA release in Step 1, this module ALSO exposes a thin custom
|
||||
**stdio MCP server** (FastMCP) returning the SAME chunks as ``structuredContent`` — honoring
|
||||
the CLAUDE.md "data access via MCP" convention. Both paths wrap the identical ``retrieve()``
|
||||
core via ``retrieve_chunks``, so the citation seam (``{file, locator, snippet, score}``) is
|
||||
byte-identical whether the agents reach it in-process or over stdio.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from agent_framework import FunctionTool, tool
|
||||
|
||||
from portfolio_optimiser.provenance import Citation
|
||||
from portfolio_optimiser.retrieval import RetrievedChunk, TextSpan, retrieve
|
||||
|
||||
|
||||
def _chunk_to_dict(c: RetrievedChunk) -> dict[str, Any]:
|
||||
return {
|
||||
"file": c.file,
|
||||
"locator": {"start_index": c.locator.start_index, "end_index": c.locator.end_index},
|
||||
"snippet": c.snippet,
|
||||
"score": c.score,
|
||||
}
|
||||
|
||||
|
||||
def retrieve_chunks(query: str, docs_dir: str, top_k: int = 3) -> list[dict[str, Any]]:
|
||||
"""The shared data-source call: retrieve citation-ready chunks as plain dicts (the
|
||||
``structuredContent`` shape). Identical on the in-process tool path and the MCP path."""
|
||||
return [_chunk_to_dict(c) for c in retrieve(query, docs_dir, top_k)]
|
||||
|
||||
|
||||
def chunk_dict_to_citation(d: dict[str, Any]) -> Citation:
|
||||
"""Map a structuredContent chunk dict into a first-class ``provenance.Citation``."""
|
||||
loc = d["locator"]
|
||||
return Citation(
|
||||
file=d["file"],
|
||||
locator=TextSpan(start_index=loc["start_index"], end_index=loc["end_index"]),
|
||||
snippet=d["snippet"],
|
||||
)
|
||||
|
||||
|
||||
def make_retrieval_tool(docs_dir: str, *, top_k: int = 3) -> FunctionTool:
|
||||
"""Build the GA in-process data-source tool bound to a docs folder. The agents call it;
|
||||
the orchestrator maps the returned chunks into ``provenance.Citation``."""
|
||||
|
||||
@tool(
|
||||
name="retrieve_cost_docs",
|
||||
description="Retrieve cited snippets from the project's cost documentation.",
|
||||
)
|
||||
def retrieve_cost_docs(query: str) -> list[dict[str, Any]]:
|
||||
return retrieve_chunks(query, docs_dir, top_k)
|
||||
|
||||
return retrieve_cost_docs
|
||||
|
||||
|
||||
def build_mcp_server(docs_dir: str, *, top_k: int = 3) -> Any:
|
||||
"""Thin custom stdio MCP server (FastMCP) exposing the same ``retrieve()`` core as
|
||||
``structuredContent``. Run via ``server.run()`` for stdio; consumed by an
|
||||
``MCPStdioTool``. The tool delegates to ``retrieve_chunks`` so its shape matches the
|
||||
in-process path exactly."""
|
||||
from mcp.server.fastmcp import FastMCP
|
||||
|
||||
server = FastMCP("portfolio-optimiser-docs")
|
||||
|
||||
@server.tool()
|
||||
def retrieve_cost_docs(query: str) -> list[dict[str, Any]]:
|
||||
return retrieve_chunks(query, docs_dir, top_k)
|
||||
|
||||
return server
|
||||
55
tests/test_datasource.py
Normal file
55
tests/test_datasource.py
Normal file
|
|
@ -0,0 +1,55 @@
|
|||
"""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
|
||||
Loading…
Add table
Add a link
Reference in a new issue