feat(fase2): expose retriever as citation-bearing data source

This commit is contained in:
Kjell Tore Guttormsen 2026-06-24 13:34:53 +02:00
commit f981c51e0b
2 changed files with 131 additions and 0 deletions

View 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