Cross-collection agentic retrieval + live schema discovery by aryan-410 · Pull Request #150 · AzureCosmosDB/MCPToolKit

aryan-410 · 2026-07-23T04:52:39Z

Summary

Upgrades the cosmos-retriever agentic search service to (1) search across all collections in a database and (2) discover each container's schema/capabilities live instead of assuming a hardcoded chunked schema.

Key changes

Live discovery (retrieval/discovery/, retrieval/binding.py): infers vector/full-text indexes, partition keys, and dimensions from container metadata. Removes the hardcoded default_chunked_schema (defaults.py deleted).
Cross-collection retrieval (retrieval/orchestration.py): container="*" fans out across every searchable collection in a database and fuses hits with RRF. Per-collection depth tracks the fused limit so a gold doc ranked deep within its own collection isn't truncated before fusion.
Structured schema overrides (retrieval/schema_override.py): user-supplied JSON override (document_id_path, chunk_order_path, dunder codec, …) via the corpus registry, replacing the none/legacy_chunked enum.
Database-aware corpus registry: per-database embedding endpoint/model, with fallback to the default endpoint+key when an entry lists only a model.
container is now optional — omit to search the whole database; database is required (no silent server defaults).
Admin + tuning: GET/PATCH /config, per-request RuntimeConfig overrides, per-turn timing instrumentation, bounded-TTL retriever cache.
.NET agentic_search: optional container, JSON schemaOverride, and a Managed-Identity-excluded Cosmos credential so the native tools fall through to developer CLI auth locally.

Tests

Adds unit tests for binding, discovery, orchestration, cache, and runtime config. All pass (pytest, 38).

Notes

Retrieval is token-neutral vs. the previous per-collection depth; the deeper fusion improved paired recall by ~+0.07 on a 9-collection enterprise benchmark with no change in turns/tokens.

FastAPI service (POST /search, GET /health) wrapping CosmosRetriever, which runs a multi-turn retrieval agent over a Cosmos DB corpus. Pluggable inference backend: harmony_vllm (fine-tuned pat-jj/harness-1), openai_chat (any OpenAI-compatible chat model), or openai_responses (reasoning models such as gpt-5.4 on Azure AI Foundry). Includes tests for the server and agent loops. The .NET agentic_search tool calls this service over HTTP.

Add a 9th MCP tool, agentic_search, that runs the Cosmos Retriever agent over a Cosmos DB corpus and returns ranked, curated documents. - AgenticSearchExecutor: calls the cosmos-retriever service over HTTP (COSMOS_RETRIEVER_URL, COSMOS_RETRIEVER_TIMEOUT_S); always returns parseable JSON (error envelope on failure). - Wire into Program.cs, MCPProtocolController (tools/list + tools/call), MCPTestController, and McpToolRequestValidator. - CosmosClientFactory: exclude ManagedIdentityCredential (fall through to az login); accept the standard MCP _meta params field. - Docs: docs/AGENTIC_SEARCH.md, README + CHANGELOG + .env.example.

…rness - Move VllmTokenCompleter + run_single_episode into inference/vllm_policy.py - Delete inference/evaluate_harness1_vllm.py (eval/benchmark code) - Repoint retriever.py to vllm_policy; update env_rl docstring - Include pre-existing: pool_doc_ids trajectory pooling (openai_chat), optional baseten import (rerank)

…s to 1-50, de-brand harness references

The datagen/ package deletion was previously only staged, never committed, so it still appeared in the PR. Actually remove it (search_dataset.py, generate_sft_rl_splits.py, BrowseComp-Plus, README, __init__) along with the unit tests folder, the datagen TYPE_CHECKING import in tasks.py, the stale datagen comment in config.py, and the now-dangling pytest/respx dev deps and pytest/ruff test config in pyproject.toml.

…lResult Add a trajectory field to RetrievalResult populated by the harmony_vllm backend: the search queries issued (search_history), per-turn tool calls (turn_tools), programmatic per-turn status summaries (turn_summaries), and the final per-doc importance tags from curation (curated_importance).

…ed retrieval layer - Add schema-driven retrieval package (planner/compiler/executor/strategies/ document-resolvers) decoupled from physical Cosmos paths - Endpoint-only inference (OpenAI chat/responses); drop vLLM/RL runtime - Add foundry-harness and retrieval/agent benchmark scripts

…kflow guides

…lt corpus profile); rename LegacyDunderCodec -> DunderChunkCodec

…ratch, endpoint-only backends

…loop.py and ChatSearchResult -> AgentSearchResult (both backends live there); keep backend-id strings

…undry Messages API); refresh backend docs/.env

- Live discovery of container schema/capabilities (retrieval/discovery/, binding.py), replacing the hardcoded default_chunked_schema (defaults.py removed) - Cross-collection retrieval: container="*" fans out across all searchable collections in a database and fuses with RRF (orchestration.py); per-collection depth tracks the fused limit so gold docs ranked deep in a collection aren't dropped - Structured, user-supplied schema overrides (schema_override.py) via the corpus registry, replacing the none/legacy_chunked enum - Database-aware corpus registry: per-database embedding endpoint/model with default-endpoint fallback - container is now optional (defaults to whole-database); database required, no silent server defaults - /config admin endpoints, per-request RuntimeConfig overrides, timing instrumentation, bounded-TTL retriever cache - .NET agentic_search: optional container, JSON schemaOverride, Managed-Identity- excluded Cosmos credential for the native tools - Tests: binding, discovery, orchestration, cache, runtime config

…rajectory pruning Port the upstream token-budget context-management system into the agent loop: - _BudgetController: threshold/token/tool-output budgets, spillage-based rejection control, per-step token accounting, and usage annotation - Cross-turn dedup via ignore_ids + first-query-wins doc->query mapping - Real trajectory pruning (excise DOCUMENT ID blocks) using local transcript instead of previous_response_id, for both responses and chat backends - Tool-output clamping when remaining budget is low - Wire budgets from settings through retriever into both run loops - Add COSMOS_RETRIEVER_THRESHOLD_BUDGET / _TOKEN_BUDGET settings - 14 new tests (tests/test_token_budget.py); ignore .env.* secrets

…riever URL map - agent_loop: replace hard-delete pruning with idempotent [pruned] tombstones so the model retains a trace of its own prune actions - tools: add distinct item-is-document read path that skips the reranker - corpus_registry: drop schema_override for the browsecomp corpus - AgenticSearchExecutor: add COSMOS_RETRIEVER_URL_MAP for per-database retriever endpoints - tests: update prune assertion for tombstone output

aryan-410 · 2026-07-27T21:16:50Z

@sajeetharan not sure if this github workflows file needs to be present in the MCP repo. Usually, if this was a monolithic repo I would include it as it is needed for replication of venvs but if it is being merged to this MCP Toolkit I am not sure if it should be included.

aryan-410 · 2026-07-27T21:30:35Z

Walk the user through all the configs they can + need to provide

There is a .env file as well. Is this a config file? aryans-internship-cosmos should not leak into the repo

magdalendobson · 2026-07-28T17:19:13Z

+import re
+
+_TOKEN_RE = re.compile(r"\w+", re.UNICODE)
+_STOPWORDS = frozenset(


Do we need stop words in other languages too? Or are we just targeting English here? If so is that documented anywhere?

test comment.

magdalendobson · 2026-07-28T17:55:47Z

+    "embed_model":  "qwen3-embed",
+    "embed_query_instruction": "Given a question, retrieve documents that answer it"
+  },
+  "search_retrieval_database/browsecomp_corpus_container": {


Is this meant to be some kind of example? Or is this something that anyone can run? It seems a little strange to have personal account URIs in the public repo. If it's just an example, maybe change the URIs to something more generic like "my-account-uri.documents.azure.com"

magdalendobson · 2026-07-28T17:56:44Z

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/


Is there a license here because code under other folders has a different license, or does this apply to every file in the repo/PR?

magdalendobson · 2026-07-28T18:07:56Z

+                        json=payload,
+                        timeout=self.timeout_s,
+                    )
+                    response.raise_for_status()


I'm a little confused about what's going on here. The prompt above seems to ask for a "yes" or "no" response, but then that response doesn't seem to be taken into account anywhere, and instead some kind of "score" is referenced.

magdalendobson · 2026-07-28T18:09:00Z

+        for idx, (doc, group) in enumerate(zip(documents, response.data)):
+            score = 0.0
+            for result in group:
+                if result.label == "yes":


What if the LLM doesn't respond "yes" or "no" precisely? It's not guaranteed that it will even if told to, right?

magdalendobson · 2026-07-28T18:28:55Z

@@ -0,0 +1,256 @@
+from __future__ import annotations


It would be good to have some tests for the compiler. You could probably just commit a small text file with different types of queries and check that they compile as expected.

metastableB · 2026-07-28T17:26:35Z

@@ -0,0 +1,297 @@
+# The Retrieval System


this is a good description, but do you want to commit the design doc into the repo? This doc might drift away from the actual code over time right?

metastableB · 2026-07-28T17:32:22Z

+    return b + "/v1/messages"
+
+
+def run_anthropic_search(


Add doc strings for these public methods? Its unclear how run_anthropic_search vs run_chat_search vs run_responses_search differ

metastableB · 2026-07-28T17:37:15Z

+        turn_tools.append([fc.name for fc in function_calls])
+        # ── act: execute tools with dedup + reject + clamp; observe: annotate
+        budget.reset_step()
+        for fc in function_calls:


This loop injecting function call results seem common to all the public api methods? they only differ in how they handle the model end-point and responses right? Or is there enough differences between the methods to warrant duplication in each of the run_* methods?

The tool handling should not care about the model/inference end-point right?

metastableB · 2026-07-28T17:41:23Z

+    text_token_counter=None,
+    threshold_budget: int = _DEFAULT_THRESHOLD_BUDGET,
+    token_budget: int = _DEFAULT_TOKEN_BUDGET,
+) -> AgentSearchResult:


Doc strings please. On a first pass, it seems like a good chunk of the three run_* methods have same code/pattern.

metastableB · 2026-07-28T17:45:56Z

+logger = structlog.get_logger("cosmos_retriever.inference.agent_loop")
+
+
+_DEFAULT_THRESHOLD_BUDGET = 16384   # soft cap: prompt prune/conclude + restrict to prune


Can users override these?

metastableB · 2026-07-28T18:16:00Z

+            content={"config": pool.settings.redacted_config(), "pool": pool.stats()}
+        )
+
+    @app.patch("/config")


What are things that are configurable on the fly? is this user facing end point or something internal ?

metastableB · 2026-07-28T18:17:25Z

+        return results
+
+
+class ContextualReranker(Reranker):


This is a fundamental component right? add a doc-string?

metastableB · 2026-07-28T18:20:25Z

@@ -0,0 +1,710 @@
+


can you add a doc string for this config file? is this user facing or something for internal consumption. If its user facing, I'd separate the config defaults from the config resolution/parsing stuff -- 500+ lines in a config file is not easy to parse for people without context. However, if its only consumed by the code, then its fine I think.

metastableB · 2026-07-28T18:22:08Z

+    expirations: int
+
+
+class BoundedTTLCache(Generic[K, V]):


What are we caching here? How is this caching used? I'm guessing this is client/user side caching for queries or results?

metastableB · 2026-07-28T18:24:46Z

+    paths, partition keys) but cannot infer the *semantic role* of fields. This
+    override names them so chunks can be grouped back into parent documents:
+
+      - ``document_id_path``  path of the parent-document id (groups chunks)


I don't understand the cosmos db storage schema well enough; so this bit confuses me. Are these document_id_path, chunk_id_path etc canonical and always available on all databases?

magdalendobson · 2026-07-28T18:34:27Z

+  name to its account/database/embedding endpoint/model. `resolve_corpus()`
+  returns a fully-resolved `CorpusConfig`.
+- **Inference** — `INFERENCE_BACKEND` (`openai_responses` | `openai_chat`),
+  `CHAT_BASE_URL`, `CHAT_MODEL`, `CHAT_MAX_TURNS`, `CHAT_REASONING_EFFORT`, etc.


What is this parameter supposed to mean?

magdalendobson · 2026-07-28T18:35:42Z

+| Same-corpus thread-safety | per-corpus `asyncio.Lock` |
+| Cosmos overload / throttling | executor `BoundedSemaphore` + tenacity retries |
+| Runaway agents | `max_turns` + token budgets + pruning |
+| Query injection | bound `@params` + `CosmosPath` validation (retrieval layer) |


Not sure I understand how this helps with safety

magdalendobson · 2026-07-28T18:43:44Z

@@ -0,0 +1,58 @@
+# ----- Inference backend -----


Have you done end-to-end testing with all three APIs?

magdalendobson · 2026-07-28T19:06:55Z

+        "Running reranker example", reranker=args.reranker, max_tokens=args.max_tokens
+    )
+
+    enc = tiktoken.get_encoding("o200k_harmony")


Will "o200k_harmony" give an accurate token count for every allowed model?

magdalendobson

I left some comments in the doc, but I have a few overall concerns with this PR:

-- I would like to see much, much more testing on the Python side. It would be ideal to get as close to 100% coverage as possible, although in some cases this may not be possible without model access. Nevertheless, every function that can be tested without an endpoint should be tested. This is my #1 concern. I'm not sure how much testing is practical on the C# side but as much test coverage as possible should be done there too.

-- A lot of the filenames are not that descriptive, and the organization could be better. Docstrings at the top of every file describing their purpose would be good

-- I think there are a lot of baked-in assumptions and edge cases for the agentic API. For example, I worry it might not perform well for non-English documents. It would be good to put all of these out in the open.

-- Are there end-to-end testing results that we can put out in the public? It would be good to see those as part of this PR

magdalendobson · 2026-07-28T19:19:34Z

+classifiers = [
+  "Development Status :: 4 - Beta",
+  "Intended Audience :: Developers",
+  "License :: OSI Approved :: Apache Software License",


Why is it a different license from the top-level license?

cosmos-dev and others added 17 commits June 26, 2026 21:57

chore: remove benchmark scripts and dev artifacts from toolkit

9db5640

chore: remove generated datagen/splits folder

905acc0

chore(toolkit): polish agentic_search descriptions, align maxDocument…

03aa2e0

…s to 1-50, de-brand harness references

docs(cosmos-retriever): add detailed retrieval-system and agentic-wor…

4e2499d

…kflow guides

chore(cosmos-retriever): remove benchmark scripts and foundry-harness

6f4a799

refactor(cosmos-retriever): strip all docstrings and comments

ab601aa

refactor(cosmos-retriever): replace legacy.py with defaults.py (defau…

0be3253

…lt corpus profile); rename LegacyDunderCodec -> DunderChunkCodec

docs: remove all Harness-1/vLLM lineage; state agent is built from sc…

a45e09d

…ratch, endpoint-only backends

refactor(cosmos-retriever): rename inference/openai_chat.py -> agent_…

cda73a5

…loop.py and ChatSearchResult -> AgentSearchResult (both backends live there); keep backend-id strings

feat(cosmos-retriever): add anthropic_messages backend (Claude via Fo…

d25a17b

…undry Messages API); refresh backend docs/.env

aryan-410 mentioned this pull request Jul 23, 2026

Feat/agentic search cosmos retriever #148

Closed

aryan-410 added 2 commits July 24, 2026 19:56

aryan-410 commented Jul 27, 2026

View reviewed changes

magdalendobson reviewed Jul 28, 2026

View reviewed changes

metastableB reviewed Jul 28, 2026

View reviewed changes

magdalendobson reviewed Jul 28, 2026

View reviewed changes

magdalendobson suggested changes Jul 28, 2026

View reviewed changes

magdalendobson reviewed Jul 28, 2026

View reviewed changes

		logger = structlog.get_logger("cosmos_retriever.inference.agent_loop")


		_DEFAULT_THRESHOLD_BUDGET = 16384 # soft cap: prompt prune/conclude + restrict to prune

Uh oh!

Conversation

aryan-410 commented Jul 23, 2026

Summary

Key changes

Tests

Notes

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

magdalendobson left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants