server.py

"""A RAG-over-MCP server: exposes the rag-poc retrieval step as an MCP tool.

The idea is the canonical "RAG over MCP" pattern: this server does *retrieval*
only -- it embeds a query with the same local Ollama model that built the store,
finds the most similar chunks, and hands them back. The MCP client's model (e.g.
Claude) then reads those chunks and writes the grounded answer itself.

It reuses the sibling rag-poc project's code rather than reimplementing it: the
folder is put on sys.path (see RAG_POC_PATH) and its `rag` package is imported.

Run it:
    mcp dev server.py        # opens the MCP Inspector (needs the mcp[cli] extra)
    python server.py         # runs the raw stdio server (what a client launches)

Prerequisites:
    - Ollama running at localhost:11434 with the 'nomic-embed-text' model.
    - rag-poc already ingested (its store.npz exists): `python main.py ingest`.
"""

from __future__ import annotations

import os
import sys
from pathlib import Path

import httpx

from mcp.server.fastmcp import FastMCP

# Locate the sibling rag-poc project. Defaults to ../rag-poc relative to this
# file; override with the RAG_POC_PATH env var if it lives elsewhere.
RAG_POC_PATH = Path(
    os.environ.get("RAG_POC_PATH", Path(__file__).resolve().parent.parent / "rag-poc")
).resolve()
STORE_PATH = RAG_POC_PATH / "store.npz"

# Make rag-poc's `rag` package importable, then borrow exactly the pieces we need:
# the vector store, the query embedder, and the input-sanitising hook.
sys.path.insert(0, str(RAG_POC_PATH))
from rag import security  # noqa: E402
from rag.embed import embed  # noqa: E402
from rag.store import VectorStore  # noqa: E402

mcp = FastMCP("mcp-poc")


@mcp.tool()
def rag_search(query: str, k: int = 4) -> list[dict]:
    """Search the rag-poc document store and return the most relevant chunks.

    Retrieval only: this embeds your query with the same local model that built
    the store, ranks the stored chunks by cosine similarity, and returns the top
    matches. You (the model reading this) should answer from these chunks and
    cite the `source` of each one you use; if they don't contain the answer, say
    so rather than guessing.

    Args:
        query: the question or search phrase to look up.
        k: how many chunks to return (clamped to 1-20; default 4).

    Returns:
        A list of {source, score, text} dicts, most similar first. `score` is the
        cosine similarity (roughly 0..1, higher is more relevant).
    """
    query = security.sanitize_query(query)  # validate + bound the user input
    k = max(1, min(int(k), 20))

    if not STORE_PATH.exists():
        raise FileNotFoundError(
            f"No store at {STORE_PATH}. In rag-poc, run: python main.py ingest"
        )

    try:
        query_vector = embed(query)  # one Ollama embedding call
    except httpx.HTTPError as e:
        raise RuntimeError(
            "Could not reach the Ollama embedding model at localhost:11434. "
            "Start Ollama and make sure 'nomic-embed-text' is pulled. "
            f"(underlying error: {e})"
        ) from e

    store = VectorStore.load(str(STORE_PATH))
    hits = store.search(query_vector, k)
    return [
        {"source": hit.source, "score": round(hit.score, 4), "text": hit.text}
        for hit in hits
    ]


if __name__ == "__main__":
    # stdio transport: JSON-RPC over stdin/stdout, how a local client launches us.
    mcp.run()