From 4e0b68daa175d5247d4bfa9d4cda3b6ce85140cc Mon Sep 17 00:00:00 2001
From: Mahika Wason <mwason@nvidia.com>
Date: Thu, 9 Jul 2026 12:31:25 -0700
Subject: [PATCH] Compact MCP evidence query responses

---
 .../src/nemo_retriever/service/mcp_server.py  | 45 ++++++++++++++++++-
 nemo_retriever/tests/test_service_mcp.py      | 34 ++++++++++++++
 2 files changed, 77 insertions(+), 2 deletions(-)

diff --git a/nemo_retriever/src/nemo_retriever/service/mcp_server.py b/nemo_retriever/src/nemo_retriever/service/mcp_server.py
index 8f522eab83..0e742c1f31 100644
--- a/nemo_retriever/src/nemo_retriever/service/mcp_server.py
+++ b/nemo_retriever/src/nemo_retriever/service/mcp_server.py
@@ -202,6 +202,7 @@ async def query(
         top_k: int = 5,
         format: QueryFormat = "hits",
         payload: dict[str, Any] | None = None,
+        max_text_chars: int = 500,
     ) -> dict[str, Any]:
         body = dict(payload or {})
         body.setdefault("query", query)
@@ -210,7 +211,10 @@ async def query(
         async with self._client() as client:
             resp = await client.post("/v1/query", json=body)
         self._raise_for_status(resp)
-        return dict(self._json_or_text(resp))
+        result = dict(self._json_or_text(resp))
+        if body.get("format") == "evidence":
+            return _compact_evidence_response(result, max_text_chars=min(max_text_chars, 500))
+        return result
 
     async def answer(
         self,
@@ -438,6 +442,36 @@ def _document_bytes(doc: MCPDocumentInput) -> tuple[str, bytes]:
     return doc.filename or "document.bin", content
 
 
+def _compact_text(value: Any, *, max_chars: int) -> str:
+    text = str(value or "")
+    if max_chars <= 0 or len(text) <= max_chars:
+        return text
+    return text[:max_chars].rstrip() + "..."
+
+
+def _compact_evidence_response(payload: dict[str, Any], *, max_text_chars: int) -> dict[str, Any]:
+    """Trim verbose evidence text so MCP tool output preserves all top-k hits."""
+    compact = dict(payload)
+    compact_results: list[Any] = []
+    for result in payload.get("results") or []:
+        if not isinstance(result, dict):
+            compact_results.append(result)
+            continue
+        compact_result = dict(result)
+        evidence_items: list[Any] = []
+        for item in result.get("evidence") or []:
+            if not isinstance(item, dict):
+                evidence_items.append(item)
+                continue
+            compact_item = dict(item)
+            compact_item["text"] = _compact_text(compact_item.get("text"), max_chars=max_text_chars)
+            evidence_items.append(compact_item)
+        compact_result["evidence"] = evidence_items
+        compact_results.append(compact_result)
+    compact["results"] = compact_results
+    return compact
+
+
 def build_mcp(settings: ServiceMCPSettings | None = None) -> FastMCP:
     """Create a FastMCP server for a retriever service endpoint."""
     settings = settings or ServiceMCPSettings()
@@ -499,8 +533,15 @@ async def query(
         top_k: int = 5,
         format: QueryFormat = "hits",
         payload: dict[str, Any] | None = None,
+        max_text_chars: int = 500,
     ) -> dict[str, Any]:
-        return await service.query(query, top_k=top_k, format=format, payload=payload)
+        return await service.query(
+            query,
+            top_k=top_k,
+            format=format,
+            payload=payload,
+            max_text_chars=max_text_chars,
+        )
 
     @mcp.tool(name="answer", description="Search ingested documents and generate an answer.")
     async def answer(
diff --git a/nemo_retriever/tests/test_service_mcp.py b/nemo_retriever/tests/test_service_mcp.py
index 59d54261ec..a7df7de8f1 100644
--- a/nemo_retriever/tests/test_service_mcp.py
+++ b/nemo_retriever/tests/test_service_mcp.py
@@ -80,6 +80,40 @@ def _handler(request: httpx.Request) -> httpx.Response:
     assert result["results"][0]["hits"][0]["text"] == "match"
 
 
+def test_query_tool_compacts_evidence_text_without_dropping_hits() -> None:
+    seen: dict[str, Any] = {}
+    evidence = [
+        {
+            "text": f"chunk-{index} " + ("x" * 2000),
+            "source": "doc",
+            "locator": {"kind": "page", "value": index + 1},
+            "score": 1.0 / (index + 1),
+            "citation": f"doc p.{index + 1}",
+        }
+        for index in range(10)
+    ]
+
+    def _handler(request: httpx.Request) -> httpx.Response:
+        seen["body"] = json.loads(request.content)
+        return httpx.Response(200, json={"results": [{"evidence": evidence}]})
+
+    client = ServiceMCPClient(
+        ServiceMCPSettings(base_url="http://service:7670"),
+        transport=httpx.MockTransport(_handler),
+    )
+
+    result = _run(client.query("What is indexed?", top_k=10, format="evidence", max_text_chars=20_000))
+
+    compact_evidence = result["results"][0]["evidence"]
+    assert seen["body"] == {"query": "What is indexed?", "top_k": 10, "format": "evidence"}
+    assert len(compact_evidence) == 10
+    assert compact_evidence[0]["citation"] == "doc p.1"
+    assert compact_evidence[0]["locator"] == {"kind": "page", "value": 1}
+    assert compact_evidence[0]["text"].startswith("chunk-0 ")
+    assert compact_evidence[0]["text"].endswith("...")
+    assert len(compact_evidence[0]["text"]) == 503
+
+
 def test_ingest_documents_accepts_inline_base64_upload() -> None:
     calls: list[tuple[str, str]] = []
     upload_body = b""