From f981c51e0b910f2bac7df4646514bf988d371c7b Mon Sep 17 00:00:00 2001
From: Kjell Tore Guttormsen <hello@fromaitochitta.com>
Date: Wed, 24 Jun 2026 13:34:53 +0200
Subject: [PATCH] feat(fase2): expose retriever as citation-bearing data source

---
 src/portfolio_optimiser/datasource.py | 76 +++++++++++++++++++++++++++
 tests/test_datasource.py              | 55 +++++++++++++++++++
 2 files changed, 131 insertions(+)
 create mode 100644 src/portfolio_optimiser/datasource.py
 create mode 100644 tests/test_datasource.py

diff --git a/src/portfolio_optimiser/datasource.py b/src/portfolio_optimiser/datasource.py
new file mode 100644
index 0000000..c145f5f
--- /dev/null
+++ b/src/portfolio_optimiser/datasource.py
@@ -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
diff --git a/tests/test_datasource.py b/tests/test_datasource.py
new file mode 100644
index 0000000..124f8f6
--- /dev/null
+++ b/tests/test_datasource.py
@@ -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