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:
@@ -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
@@ -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}"
|
||||
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user