Exposing RAG through MCP

"""Retrieval back ends for the docs-assistant RAG server. The MCP layer depends only on the Retriever protocol."""
import os
import re
from typing import Protocol

import httpx
from pydantic import BaseModel


class Passage(BaseModel):
    chunk_id: int
    page_id: str          # 'algebra/matrices': a readable ID that also fits docs://{component}/{page}
    component: str
    title: str
    section: str
    text: str
    score: float


class Retriever(Protocol):
    def search(self, query: str, groups: list[str], component: str | None, limit: int) -> list[Passage]: ...
    def page(self, component: str, page: str, groups: list[str]) -> str | None: ...


# --- pgvector: the DocsAssistant schema plus one column, documents.acl_groups ---------------------------------
ACL_DDL = """
ALTER TABLE documents ADD COLUMN acl_groups text[] NOT NULL DEFAULT ARRAY['staff'];
CREATE INDEX documents_acl_gin ON documents USING gin (acl_groups);
"""

SEARCH_SQL = r"""
SELECT c.id AS chunk_id,
       d.component,
       d.component || '/' || regexp_replace(d.source_path, '^.*/|\.adoc$', '', 'g') AS page_id,
       coalesce(d.title, '') AS title,
       coalesce(c.section_path, '') AS section,
       c.content AS text,
       1 - (e.embedding <=> %(q)s) AS score
FROM chunk_embeddings e
JOIN chunks c ON c.id = e.chunk_id
JOIN documents d ON d.id = c.document_id
WHERE e.model = %(model)s
  AND d.acl_groups && %(groups)s::text[]
  AND (%(component)s::text IS NULL OR d.component = %(component)s::text)
ORDER BY e.embedding <=> %(q)s
LIMIT %(limit)s
"""

PAGE_SQL = r"""
SELECT string_agg(c.content, E'\n\n' ORDER BY c.chunk_index)
FROM documents d
JOIN chunks c ON c.document_id = d.id
WHERE d.component = %(component)s
  AND regexp_replace(d.source_path, '^.*/|\.adoc$', '', 'g') = %(page)s
  AND d.acl_groups && %(groups)s::text[]
GROUP BY d.id
"""

EMBED_MODEL = "nomic-embed-text"


class PgVectorRetriever:
    def __init__(self) -> None:
        from pgvector.psycopg import register_vector
        from psycopg_pool import ConnectionPool

        # DATABASE_URL comes from the environment or a secret store, never from source code
        self._pool = ConnectionPool(os.environ["DATABASE_URL"], configure=register_vector, open=True)
        self._ollama = httpx.Client(base_url=os.environ.get("OLLAMA_URL", "http://127.0.0.1:11434"), timeout=30)

    def _embed(self, text: str) -> list[float]:
        r = self._ollama.post("/api/embed", json={"model": EMBED_MODEL, "input": text})
        r.raise_for_status()
        return r.json()["embeddings"][0]

    def search(self, query: str, groups: list[str], component: str | None, limit: int) -> list[Passage]:
        from pgvector import Vector

        params = {"q": Vector(self._embed(query)), "model": EMBED_MODEL, "groups": groups,
                  "component": component, "limit": limit}
        with self._pool.connection() as conn:
            rows = conn.execute(SEARCH_SQL, params).fetchall()
        return [Passage(chunk_id=r[0], component=r[1], page_id=r[2], title=r[3], section=r[4], text=r[5],
                        score=float(r[6])) for r in rows]

    def page(self, component: str, page: str, groups: list[str]) -> str | None:
        with self._pool.connection() as conn:
            row = conn.execute(PAGE_SQL, {"component": component, "page": page, "groups": groups}).fetchone()
        return row[0] if row else None

package dev.example.docsassistant;

import java.util.List;
import org.springframework.ai.document.Document;
import org.springframework.ai.mcp.annotation.McpTool;
import org.springframework.ai.mcp.annotation.McpToolParam;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.ai.vectorstore.filter.FilterExpressionBuilder;
import org.springframework.security.core.Authentication;
import org.springframework.security.core.GrantedAuthority;
import org.springframework.security.core.context.SecurityContextHolder;
import org.springframework.stereotype.Component;

@Component
public class RagTools {

    public record Hit(int rank, String id, String title, String snippet, double score) {}

    public record SearchResult(List<Hit> hits) {}

    private final VectorStore vectorStore;

    public RagTools(VectorStore vectorStore) {
        this.vectorStore = vectorStore;
    }

    @McpTool(name = "search_docs",
            description = "Semantic search over the documentation you are allowed to read. Returns short snippets.",
            generateOutputSchema = true,
            annotations = @McpTool.McpAnnotations(readOnlyHint = true, destructiveHint = false, openWorldHint = false))
    public SearchResult searchDocs(
            @McpToolParam(description = "Free-text query", required = true) String query,
            @McpToolParam(description = "algebra, units or statistics; omit to search all", required = false) String component,
            @McpToolParam(description = "Maximum number of hits, 1 to 10", required = false) Integer topK) {
        var groups = callerGroups();                                   // from the verified JWT, never from an argument
        var b = new FilterExpressionBuilder();
        var acl = b.in("acl", groups.toArray());
        Filter.Expression filter = (component == null ? acl : b.and(b.eq("component", component), acl)).build();
        int k = topK == null ? 5 : Math.max(1, Math.min(topK, 10));
        var docs = vectorStore.similaritySearch(
                SearchRequest.builder().query(query).topK(k).filterExpression(filter).build());
        var hits = new java.util.ArrayList<Hit>();
        for (int i = 0; i < docs.size(); i++) {
            Document d = docs.get(i);
            String text = d.getText() == null ? "" : d.getText();
            hits.add(new Hit(i + 1, (String) d.getMetadata().get("id"), (String) d.getMetadata().get("title"),
                    text.substring(0, Math.min(200, text.length())), d.getScore() == null ? 0 : d.getScore()));
        }
        return new SearchResult(hits);
    }

    private static List<String> callerGroups() {
        Authentication auth = SecurityContextHolder.getContext().getAuthentication();
        if (auth == null || !auth.isAuthenticated()) {
            throw new IllegalStateException("Authentication required");
        }
        return auth.getAuthorities().stream().map(GrantedAuthority::getAuthority)
                .filter(a -> a.startsWith("GROUP_")).map(a -> a.substring("GROUP_".length())).toList();
    }
}

    /** SCOPE_* authorities from the scope claim, plus GROUP_* authorities from the verified groups claim. */
    private static JwtAuthenticationConverter groupsConverter() {
        var scopes = new JwtGrantedAuthoritiesConverter();
        var converter = new JwtAuthenticationConverter();
        converter.setJwtGrantedAuthoritiesConverter(jwt -> {
            List<GrantedAuthority> authorities = new ArrayList<>(scopes.convert(jwt));
            List<String> groups = jwt.getClaimAsStringList("groups");
            if (groups != null) groups.forEach(g -> authorities.add(new SimpleGrantedAuthority("GROUP_" + g)));
            return authorities;
        });
        return converter;
    }

spring:
  datasource:
    url: ${DATABASE_URL}
    username: ${DATABASE_USER}
    password: ${DATABASE_PASSWORD}
  ai:
    vectorstore:
      pgvector:
        index-type: HNSW
        distance-type: COSINE_DISTANCE
        dimensions: 768
        table-name: docs_vector_store
    ollama:
      embedding:
        model: nomic-embed-text

{
  "mcpServers": {
    "docs-rag": {
      "type": "http",
      "url": "http://127.0.0.1:8000/mcp",
      "headers": { "Authorization": "Bearer ${DOCS_MCP_TOKEN}" }
    }
  }
}

$ claude mcp add --transport http docs-rag http://127.0.0.1:8000/mcp \
    --header 'Authorization: Bearer ${DOCS_MCP_TOKEN}'
$ DOCS_MCP_TOKEN="$(your-secret-tool get docs-mcp)" claude

import asyncio
import os

from langchain.agents import create_agent
from langchain_mcp_adapters.client import MultiServerMCPClient
from langchain_ollama import ChatOllama


async def main() -> None:
    client = MultiServerMCPClient({
        "docs": {
            "transport": "streamable_http",
            "url": os.environ.get("DOCS_MCP_URL", "http://127.0.0.1:8000/mcp"),
            "headers": {"Authorization": f"Bearer {os.environ['DOCS_MCP_TOKEN']}"},   # from the environment
        }
    })
    tools = await client.get_tools()               # search_docs and read_passages as LangChain tools
    agent = create_agent(ChatOllama(model="llama3.1:8b", base_url=os.environ.get("OLLAMA_URL")), tools,
                         system_prompt="Answer only from search_docs / read_passages results and cite page IDs.")
    out = await agent.ainvoke({"messages": [{"role": "user", "content": "How do I solve a linear system?"}]})
    for m in out["messages"]:
        print(type(m).__name__, str(m.content)[:150].replace("\n", " "))


asyncio.run(main())