portfolio-optimiser/src/portfolio_optimiser/datasource.py

"""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 collections.abc import Sequence
from typing import Any

from agent_framework import FunctionTool, tool

from portfolio_optimiser.okf import Bundle
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 bundle_citations(bundle: Bundle) -> list[Citation]:
    """First-class citations for the OKF-navigated bundle context: one ``Citation`` per non-verdict
    concept file (exactly the files ``okf.bundle_context`` renders). Whole-file locators, exact by
    construction (``snippet == body[start:end]``). The verdict layer is deliberately uncited — it is
    not cost documentation; it enters only via the gated ExpeL fold."""
    return [
        Citation(
            file=f.name,
            locator=TextSpan(start_index=0, end_index=len(f.body)),
            snippet=f.body,
        )
        for f in bundle.context_files
    ]


def bundle_excerpt_citations(bundle: Bundle, concept_ids: Sequence[str]) -> list[Citation]:
    """Citations for exactly the concepts a pre-pass DELIVERED (order 20260907T080223Z).

    ``bundle_citations`` above cites the whole navigated base. Under a declared cut that would
    have the run's stamp claim every concept for a proposal that saw a handful — the same
    undeclared claim the pre-pass seam exists to remove, one artefact over.

    The body comes from the NAVIGATED ``BundleFile``, never from the payload's ``text``: the
    delivered text is NFC-normalised with trailing whitespace stripped, so a locator over it would
    not index the file it names. Reusing ``context_files`` also means this shares
    ``bundle_citations``' verdict-layer exclusion rather than restating it.

    An id the navigation did not reach raises: ``prepass.verify_against_bundle`` has already
    resolved every one of them on disk, so a miss here means the two disagree about the base, and
    a citation list that silently drops entries would under-report the very denominator this seam
    publishes.
    """
    by_name = {f.name: f for f in bundle.context_files}
    citations: list[Citation] = []
    for concept_id in concept_ids:
        name = concept_id + ".md"
        try:
            file = by_name[name]
        except KeyError as error:
            raise ValueError(
                f"the pre-pass delivered {concept_id!r}, which this bundle's navigation does not "
                "reach; the payload and the navigated base disagree"
            ) from error
        citations.append(
            Citation(
                file=name,
                locator=TextSpan(start_index=0, end_index=len(file.body)),
                snippet=file.body,
            )
        )
    return citations


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