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

Co-authored-by: claude <[email protected]>

docs_mcp/format.py

@@ -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"

docs_mcp/server.py

@@ -29,6 +29,7 @@ from mcp.server.mcpserver import MCPServer
 from mcp.server.transport_security import TransportSecuritySettings
 from pydantic import Field
 
+from .format import format_search_hit
 from .usage import TimedCall
 
 log = logging.getLogger(__name__)
@@ -182,6 +183,17 @@ def _collection():
     return _chroma_collection
 
 
+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_index():
     """Return the BM25 index, or None if it doesn't exist on disk."""
     global _bm25
@@ -257,10 +269,8 @@ def _read_markdown(source: str, source_key: str) -> str | None:
         return None
 
 
-def _format_hit(doc: str, meta: dict, distance: float | None = None) -> str:
-    """Render one retrieval hit as a fenced markdown block with full
-    provenance attached. ``doc`` is the chunk text; ``meta`` is the
-    chunk's metadata dict."""
+def _format_hit(n: int, doc: str, meta: dict, distance: float | None = None) -> str:
+    """Render one numbered retrieval hit. ``n`` is 1-based display order."""
     src_url = meta.get("source_url") or ""
     src_key = meta.get("source_key") or ""
     src = meta.get("source") or ""
@@ -268,15 +278,10 @@ def _format_hit(doc: str, meta: dict, distance: float | None = None) -> str:
     brand = meta.get("brand") or ""
     crop = meta.get("crop") or ""
     name = meta.get("product_name") or src_key
-
-    header = (
-        f"### {name}  \n"
-        f"`{src}::{src_key}` — {vendor} / {brand} / {crop}  \n"
-        f"<{src_url}>"
-    )
+    extra = f"`{src}::{src_key}` — {vendor} / {brand} / {crop}"
     if distance is not None:
-        header += f"  \n_(distance={distance:.4f})_"
-    return f"{header}\n\n{doc.strip()}\n"
+        extra += f"  distance={distance:.4f}"
+    return format_search_hit(n, name, src_url, doc, extra)
 
 
 def _rrf_fuse(rankings: list[list[str]], k: int = RRF_K) -> list[str]:
@@ -548,7 +553,8 @@ def search_docs(
     """Search the seed-variety corpus for hybrids/varieties matching a query.
 
     Returns the top-k variety chunks with their full source URLs, ratings,
-    maturity, traits, and regional listings. Optional filters narrow to one
+    maturity, traits, and regional listings. Hits are numbered [1]… for
+    citation. Optional filters narrow to one
     crop, brand, vendor, or scraper source. Use **list_versions()** first
     to discover valid facet values. Use **lookup_variety()** on any
     candidate the user is serious about — that returns the canonical
@@ -593,11 +599,7 @@ def search_docs(
             )
 
         try:
-            dense = col.query(
-                query_texts=[query],
-                n_results=pool_size,
-                where=where,
-            )
+            dense = _query_dense(col, query, pool_size, where)
         except Exception as exc:  # noqa: BLE001
             _call.set(error_dense=str(exc), hits_returned=0)
             return f"_(retrieval failed: {exc})_"
@@ -702,20 +704,21 @@ def search_docs(
             )
 
         blocks: list[str] = []
-        for cid in final_ids:
+        for n, cid in enumerate(final_ids, start=1):
             doc = id_to_doc.get(cid, "")
             meta = id_to_meta.get(cid, {})
             dist = id_to_dist.get(cid) if not used_hybrid else None
-            blocks.append(_format_hit(doc, meta, dist))
+            blocks.append(_format_hit(n, doc, meta, dist))
 
         header = (
             f"# Search results — {len(final_ids)} variety chunk"
             f"{'s' if len(final_ids) != 1 else ''}"
             f"{' (dense + BM25 hybrid)' if used_hybrid else ' (dense only)'}\n"
-            f"_Use `lookup_variety(source_key=...)` on any candidate "
-            f"to fact-check ratings from the canonical sidecar._\n\n---\n\n"
+            f"_Hits are numbered [1]… for citation. "
+            f"Use `lookup_variety(source_key=...)` on any candidate "
+            f"to fact-check ratings from the canonical sidecar._\n\n"
         )
-        return header + "\n---\n\n".join(blocks)
+        return header + "\n\n".join(blocks) + "\n"
 
 
 @mcp.tool()
@@ -975,11 +978,7 @@ def search_trials(
             full_query = f"{query} {product}"
 
         try:
-            dense = col.query(
-                query_texts=[full_query],
-                n_results=pool_size,
-                where=where,
-            )
+            dense = _query_dense(col, full_query, pool_size, where)
         except Exception as exc:  # noqa: BLE001
             _call.set(error_dense=str(exc), hits_returned=0)
             return f"_(trial retrieval failed: {exc})_"
@@ -1093,28 +1092,27 @@ def search_trials(
             )
 
         blocks: list[str] = []
-        for cid in final_ids:
+        for n, cid in enumerate(final_ids, start=1):
             doc = id_to_doc.get(cid, "")
             meta = id_to_meta.get(cid, {})
             dist = id_to_dist.get(cid) if not used_hybrid else None
-            blocks.append(_format_trial_hit(doc, meta, dist))
+            blocks.append(_format_trial_hit(n, doc, meta, dist))
 
         header = (
             f"# Trial search results — {len(final_ids)} trial document"
             f"{'s' if len(final_ids) != 1 else ''}"
             f"{' (dense + BM25 hybrid)' if used_hybrid else ' (dense only)'}\n"
-            f"_Use `get_page(source=..., source_key=...)` to read the "
+            f"_Hits are numbered [1]… for citation. "
+            f"Use `get_page(source=..., source_key=...)` to read the "
             f"full trial body. Use `lookup_variety(source_key=...)` on "
             f"any product code to verify its identity (RM, traits, "
-            f"disease ratings)._\n\n---\n\n"
+            f"disease ratings)._\n\n"
         )
-        return header + "\n---\n\n".join(blocks)
+        return header + "\n\n".join(blocks) + "\n"
 
 
-def _format_trial_hit(doc: str, meta: dict, distance: float | None = None) -> str:
-    """Trial-specific result header. Highlights crop/state/year and
-    sources URL (vs variety hits which emphasize brand + product
-    identity)."""
+def _format_trial_hit(n: int, doc: str, meta: dict, distance: float | None = None) -> str:
+    """Trial-specific numbered hit. Highlights crop/state/year."""
     src_url = meta.get("source_url") or ""
     src_key = meta.get("source_key") or ""
     src = meta.get("source") or ""
@@ -1125,15 +1123,10 @@ def _format_trial_hit(doc: str, meta: dict, distance: float | None = None) -> st
 
     title_bits = [b for b in [crop.title(), region or state, str(year) if year else ""] if b]
     title = " · ".join(title_bits) if title_bits else src_key
-
-    header = (
-        f"### Trial: {title}  \n"
-        f"`{src}::{src_key}` — {meta.get('vendor', '')} / {meta.get('brand', '')}  \n"
-        f"<{src_url}>"
-    )
+    extra = f"`{src}::{src_key}` — {meta.get('vendor', '')} / {meta.get('brand', '')}"
     if distance is not None:
-        header += f"  \n_(distance={distance:.4f})_"
-    return f"{header}\n\n{doc.strip()}\n"
+        extra += f"  distance={distance:.4f}"
+    return format_search_hit(n, f"Trial: {title}", src_url, doc, extra)
 
 
 @mcp.tool()

docs_mcp/test_format.py

@@ -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()