Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
45 changes: 43 additions & 2 deletions nemo_retriever/src/nemo_retriever/service/mcp_server.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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,
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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(
Expand Down
34 changes: 34 additions & 0 deletions nemo_retriever/tests/test_service_mcp.py
Original file line number Diff line number Diff line change
Expand Up @@ -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""
Expand Down
Loading