Port docs-mcp-template upgrades (eval, citations, chunking, prefixes) (#16)

Co-authored-by: claude <[email protected]>
This commit was merged in pull request #16.
This commit is contained in:
2026-09-29 22:27:16 -04:00
committed by claude
parent 1a3963f306
commit 92219878d9
13 changed files with 803 additions and 74 deletions
+43
View File
@@ -0,0 +1,43 @@
"""Markdown formatters for MCP tool output.
No third-party imports — tests can run without mcp/httpx/chromadb.
Citation numbers are per-call (1-based, dense) and stateless.
"""
from __future__ import annotations
def format_search_hit(
n: int,
title: str,
url: str,
text: str,
extra: str = "",
) -> str:
"""One numbered hit. `n` is 1-based in final reranked/fused order."""
head = f"[{n}] **{title}**"
if url:
head += f" — {url}"
parts = [head]
if extra:
parts.append(extra)
body = (text or "").strip()
if body:
parts.append(body)
return "\n".join(parts)
def format_search_hits(
hits: list[tuple[str, str, str, str]],
) -> str:
"""Render hits as `[1] **title** — url` then text.
`hits` is a list of (title, url, text, extra) in display order.
Empty list → empty string (no invented citations).
"""
if not hits:
return ""
blocks = [
format_search_hit(n, title, url, text, extra)
for n, (title, url, text, extra) in enumerate(hits, start=1)
]
return "\n\n".join(blocks) + "\n"
+29 -17
View File
@@ -31,6 +31,7 @@ from mcp.server.mcpserver import MCPServer
from mcp.server.transport_security import TransportSecuritySettings
from pydantic import Field
from .format import format_search_hits
from .usage import TimedCall
log = logging.getLogger(__name__)
@@ -145,6 +146,17 @@ def _collection():
return _CHROMA
def _query_dense(col, query: str, n: int, where: dict | None = None, **extra):
"""Dense query with nomic prefixes applied at embed time only."""
from rag.embeddings import EMBED_QUERY_PREFIX, embed_texts
qvec = embed_texts([query], prefix=EMBED_QUERY_PREFIX)[0]
kwargs: dict[str, Any] = {"query_embeddings": [qvec], "n_results": n}
if where:
kwargs["where"] = where
kwargs.update(extra)
return col.query(**kwargs)
def _bm25():
"""Lazy BM25Index handle. None if the FTS5 db isn't built."""
global _BM25
@@ -258,9 +270,10 @@ def search_docs(
"""Search the HPE Morpheus VM Essentials (HVM) docs corpus.
Returns the top-k most relevant chunks (with full source page URLs)
given a natural-language query. Optional filters narrow the search
to one version, one platform, or one bundle. Use list_versions()
first if you need to discover the available facet values.
given a natural-language query. Hits are numbered [1]… for citation.
Optional filters narrow the search to one version, one platform, or
one bundle. Use list_versions() first if you need to discover the
available facet values.
Call this tool whenever the user asks anything that should be
answerable from the official product documentation — install,
@@ -296,7 +309,7 @@ def search_docs(
if HYBRID_SEARCH and bm is not None:
try:
dense_res = col.query(query_texts=[query], n_results=pool, where=where)
dense_res = _query_dense(col, query, pool, where)
dense_ids = (dense_res.get("ids") or [[]])[0]
bm_hits = bm.query(query, n=pool, where=bm25_where)
bm_ids = [cid for cid, _s in bm_hits]
@@ -326,7 +339,7 @@ def search_docs(
log.warning("BM25 retrieval failed, falling back to dense: %s", e)
if not docs:
res = col.query(query_texts=[query], n_results=k, where=where)
res = _query_dense(col, query, k, where)
docs = (res.get("documents") or [[]])[0]
metas = (res.get("metadatas") or [[]])[0]
dists = (res.get("distances") or [[]])[0]
@@ -342,7 +355,7 @@ def search_docs(
extra = bm.query(query, n=pool_size, where=bm25_where) if bm else []
extra_ids = [cid for cid, _s in extra]
else:
extra_res = col.query(query_texts=[query], n_results=pool_size, where=where)
extra_res = _query_dense(col, query, pool_size, where)
extra_ids = (extra_res.get("ids") or [[]])[0]
if extra_ids:
d2, m2, _ = _enrich_from_chroma(col, extra_ids, None)
@@ -371,22 +384,21 @@ def search_docs(
if not docs:
return f"_No matches for `{query}`._"
out = [f"# {len(docs)} result(s) for `{query}`", ""]
hits: list[tuple[str, str, str, str]] = []
for doc, meta, dist in zip(docs, metas, dists):
bid = meta.get("bundle_id", "")
pid = meta.get("page_id", "")
title = meta.get("title") or pid
ver = meta.get("version") or ""
url = _source_url(bid, pid)
header = f"## {title}"
extra = f"score={1 - dist:.3f}"
if ver:
header += f" _(v{ver})_"
out.append(header)
out.append(f"[{bid}/{pid}]({url}) · score={1 - dist:.3f}")
out.append("")
out.append(doc.strip())
out.append("")
return "\n".join(out)
extra = f"v{ver} " + extra
if pid:
extra += f" page_id: `{pid}`"
hits.append((title, url, doc, extra))
header = f"# {len(docs)} result(s) for `{query}`\n\n"
return header + format_search_hits(hits)
@mcp.tool()
@@ -1073,8 +1085,8 @@ def find_doc_inconsistencies(
return f"Couldn't open Chroma collection: {e}"
where = _build_where(version, platform, bundle_id)
try:
res = col.query(query_texts=[scope_query], n_results=max_pages * 3,
where=where, include=["metadatas"])
res = _query_dense(col, scope_query, max_pages * 3, where,
include=["metadatas"])
except Exception as e:
_call.set(error=f"query: {e}")
return f"Scope query failed: {e}"
+37
View File
@@ -0,0 +1,37 @@
"""Citation markdown renderer — no Chroma, no mcp."""
from __future__ import annotations
import unittest
from docs_mcp.format import format_search_hits
class FormatSearchHitsTests(unittest.TestCase):
def test_two_hits_numbered_with_urls(self) -> None:
out = format_search_hits([
("Install on Linux", "https://docs.example.com/install", "chunk one", ""),
("HA setup", "https://docs.example.com/ha", "chunk two", ""),
])
self.assertIn("[1] **Install on Linux** — https://docs.example.com/install", out)
self.assertIn("[2] **HA setup** — https://docs.example.com/ha", out)
self.assertIn("chunk one", out)
self.assertIn("chunk two", out)
self.assertNotIn("[3]", out)
# [1] before [2]
self.assertLess(out.index("[1]"), out.index("[2]"))
# each number is followed by its URL on the same logical hit
first = out.split("[2]")[0]
self.assertIn("https://docs.example.com/install", first)
self.assertNotIn("https://docs.example.com/ha", first)
def test_empty_hits_invents_nothing(self) -> None:
self.assertEqual(format_search_hits([]), "")
def test_missing_url_omits_emdash(self) -> None:
out = format_search_hits([("Title only", "", "body", "")])
self.assertIn("[1] **Title only**", out)
self.assertNotIn("—", out)
if __name__ == "__main__":
unittest.main()