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

Aryan Saboo (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 Saboo (aryan-410) · 2026-07-27T21:16:50Z

Sajeetharan (@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 Saboo (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

yes I agree.

i replaced the committed corpus_registry.json with a safe example file and removed the personal account urls, internal embedding host, regional endpoints, and specific model names.

it now contains two placeholder entries:

the first is a full example showing every supported option, including the cosmos location, embedding configuration, env var references for credentials, dimensions, query instructions, and a nested schema override.

the second is a minimal example showing only the fields needed to identify the corpus and embedding model, with everything else inherited from the normal environment defaults.

i also added a comment explaining the database/container key format, that the values are placeholders, and that secrets should never be stored in this file. the registry only names the environment variables that contain the keys.

the committed file is therefore now documentation rather than a copy of anyone’s local setup. real development values should live in a gitignored corpus_registry.local.json, with corpus_registry_file pointed at that file locally.

Magdalen Dobson Manohar (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.

Aryan Saboo (@aryan-410) I think this comment was missed earlier.

Sorry just saw this. Yes the stop words are currently only in english. If a non-english phrase was passed to be grepped it would just result in a FTS across the path.

Okay but is this a problem? Why don't we need stop words for other languages?

Aryan Saboo (@aryan-410)

Magdalen Dobson Manohar (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"

i replaced the committed corpus_registry.json with a safe example file and removed the personal account urls, internal embedding host, regional endpoints, and specific model names.

it now contains two placeholder entries:

the first is a full example showing every supported option, including the cosmos location, embedding configuration, env var references for credentials, dimensions, query instructions, and a nested schema override.

the second is a minimal example showing only the fields needed to identify the corpus and embedding model, with everything else inherited from the normal environment defaults.

i also added a comment explaining the database/container key format, that the values are placeholders, and that secrets should never be stored in this file. the registry only names the environment variables that contain the keys.

the committed file is therefore now documentation rather than a copy of anyone’s local setup. real development values should live in a gitignored corpus_registry.local.json, with corpus_registry_file pointed at that file locally.

Magdalen Dobson Manohar (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?

yes, the apache license in cosmos-retriever was a leftover from when it was maintained as a separate project. it was not intended to create a differently licensed subtree inside the mit-licensed repository.

i removed the separate license file and updated the package metadata and readme to use mit consistently. the top-level license now governs the package along with the rest of the repo, and there are no remaining apache references.

Magdalen Dobson Manohar (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.

the Qwen3-Reranker is a yes/no classifier, not a text generator. the prompt asks the model to judge "yes" or "no", but the serving layer never lets it write a word. instead it reads the model's probability for the "yes" token at the single judgment position and returns that as the relevance score (a float in [0, 1]).

two quick details:

"<|im_start|>assistant" : the empty block forces the model to skip reasoning and land directly on the yes/no judgment token. So the very next token's logits are a clean relevance signal.

the server (Baseten [classify] / vLLM [/score]) applies softmax over the yes/no token logits and hands back P("yes"). that's is the score

so there's not really any dropped response, the yes/no is the score. "Would the model say yes?" collapses to "how much probability does it put on the yes token?", which is a better ranking signal than a hard yes/no anyway (it gives a continuous ordering).

Is re-rank configurable? Where is the re-ranker contract (that it will look at probability of yes/no) enforced?

Okay, thanks for clarifying. Might be better to use something other than "score" because to me that implies a numeric range is being used.

Magdalen Dobson Manohar (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?

It can't because the model is never in free generation mode. The yes/no wording in PREFIX is the Qwen3-Reranker scoring template:

BasetenReranker calls client.classify(), which returns a structured {label, score} per document so the code takes the probability score on the "yes" label that is set to 0 if absent.

VLLMQwen3Reranker hits /score which returns a numeric score directly.
ContextualReranker returns relevance_score.

So there's no brittle string matching on generated text, the only match is against the classifier's fixed label set.

Got it, makes sense.

Magdalen Dobson Manohar (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.

Added tests.

Don Dennis (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?

I agree, the design doc will definitely not stay true to the code over time especially once it is translated to the native language of the Cosmos AI SDK. I have now removed it from the pr

Is this meant to be removed?

Aryan Saboo (@aryan-410)

Don Dennis (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

added doc strings!

Don Dennis (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?

i agree that there is some repeated tool execution code here, but the three run functions are intentionally kept separate.

consolidating them right now would not really produce one clean shared loop. it would mostly turn the current linear functions into a large sequence of backend checks: if chat, extract calls this way and append this message shape; if responses, use a different id and update the transcript before execution; if anthropic, collect tool results and append them together afterward. usage tracking, termination conditions, transcript mutation, and tool-call extraction would all still need backend-specific branches.

the alternative would be introducing a backend adapter with several methods, but that creates more code and indirection than the duplication it removes. understanding one backend would then require jumping between the shared loop and its adapter instead of reading one function from top to bottom.

the separation is also deliberate because more backends may be added later. each provider can have its own transcript format, tool-call representation, usage model, and execution behavior without forcing those differences into an increasingly branch-heavy shared function. adding a backend should mean adding another self-contained runner, rather than modifying one central loop and increasing the blast radius for every existing backend.

the repeated section is mostly orchestration around shared helpers and the budget controller. the actual budget, deduplication, rejection, clamping, and pruning logic already lives in one place, so we are not maintaining three separate implementations of that policy.

Don Dennis (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.

Added doc strings to all functions!

Don Dennis (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?

yes, users can override these.

the default threshold and token budget constants in agent_loop.py are only fallbacks for anyone calling the run functions directly. when requests go through the retriever, those values are replaced by the server config.

they can currently be changed through the cosmos_retriever_threshold_budget and cosmos_retriever_token_budget environment variables, which default to 16384 and 32268. they can also be updated through patch /config without restarting the server.

the search paths pass the configured values into the run functions, so the module constants are not the only way to control them. i also added a comment next to the constants to make that clearer.

the one limitation is that these are currently server-level settings. they are not part of the per-request runtimeconfig in the /search body. adding per-request overrides would mean adding token_budget and threshold_budget to runtimeconfig and passing them through the config resolution and search paths.

i can add that if per-request control is useful, otherwise i’ll leave them as server-level settings.

Leave them as is.

Don Dennis (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 ?

the /config endpoints are operator-facing control-plane endpoints.

get /config returns the current server settings with secrets redacted, along with retriever pool stats.

patch /config allows an operator to update server-level defaults without restarting the service. after an update, the retriever pool is cleared and rebuilt so the new settings are used.

the fields that can be changed include the inference backend and llm defaults, chat and anthropic settings, embedding configuration, the default cosmos account and corpus, retrieval budgets and limits, reranker credentials, and cache sizing.

these are server-wide defaults and are separate from the per-request runtimeconfig passed in the /search body. request-level config only affects that individual search, while patch /config changes the defaults used by the running service.

i added docstrings to both get_config and patch_config to make the intended audience and behavior explicit.

Don Dennis (metastableB) · 2026-07-28T18:17:25Z

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


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

Don Dennis (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.

it is both user-facing and internal, with a fairly clear boundary between the two.

the fields on retrieversettings are user-facing. they define the env and .env configuration that an operator can set, such as the inference backend, base urls, account uri, budgets, and cache limits. each field is also documented through its description.

the methods on the class are internal. things like resolving the corpus, loading the registry, building clients, applying structural overrides, and selecting a backend are only used by the service to turn those settings into concrete clients and a resolved corpus config.

i added a module docstring in config.py to make that distinction explicit and to note that, if the module grows further, the natural split would be separating the settings fields from the resolution logic.

Don Dennis (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?

this is server-side caching, but it is not query or result caching.

boundedttlcache is a generic thread-safe lru cache with a ttl. in this service, the retriever pool uses it to cache retriever engine instances based on the corpus scope and structural overrides.

what is being cached is the expensive engine setup. creating a new retriever initializes the cosmos clients, embedding client, reranker, and performs live schema discovery. caching the constructed engine lets us reuse that setup for repeated searches against the same scope.

the actual searches are not cached. every /search request still executes a fresh query against cosmos and returns current results.

the cache is also bounded. max_entries, which defaults to 32, prevents memory from growing indefinitely as more corpora are accessed. the ttl, which defaults to 900 seconds, periodically rebuilds engines so schema or credential changes are eventually picked up. patch /config also clears the pool immediately.

i added docstrings in cache.py and server.py to make the distinction explicit: this caches retriever engines, not queries or search results. no additional code change should be needed for this comment.

Don Dennis (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?

These paths are not canonical Cosmos fields and are not guaranteed to exist. These are only optionally configurable if the user wants to manually provide a schema for a certain collection instead of discovering it in runtime

Magdalen Dobson Manohar (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?

INFERENCE_BACKEND selects which wire protocol / API shape the agent loop uses to talk to the LLM

Magdalen Dobson Manohar (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

yeah so there are two separate injection surfaces:

Values are passed as bound parameters.

Filter values, document IDs, ignored-ID lists, and vectors are never interpolated into the SQL text. The compiler stores them separately as Pavel Avgustinov (@p0), @qVec1, etc., and Cosmos binds them as literal values. test_filter_values_are_bound_never_inlined verifies that attacker-controlled input appears only in the parameter list and never in the generated query.

Identifiers and field paths are validated.
Paths cannot be parameterized because they are part of the query structure. Every physical path therefore passes through CosmosPath, which applies a strict segment allowlist, rejects unsafe characters, and escapes quotes and backslashes during rendering.

FTS terms follow a separate safe path: they are tokenized to \w+ terms and escaped before being rendered as quoted literals.

I updated AGENTIC_WORKFLOW.md to make this distinction explicit: bound parameters protect values, while the CosmosPath allowlist protects structural identifiers.

Aryan Saboo (@aryan-410)
see comment here which has accidentally tagged :D

Magdalen Dobson Manohar (magdalendobson) · 2026-07-28T18:43:44Z

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


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

Do we have a repro/test files for these?

Magdalen Dobson Manohar (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?

not really, o200k_harmony is not exact for every allowed model but I felt like it doesn't need to be. is IS exact for gpt-oss/Harmony models, within a few percent for other o200k-family OpenAI models (gpt-4o, gpt-5.x), and approximate for non-OpenAI models (Claude, Qwen) that use different tokenizers. I felt like that's was acceptable because the count only drives budget/truncation heuristics and not really ever billing or hard context limit enforcement, not to mention that the budgets carry headroom to absorb the drift.

using one cheap local tokenizer helps to avoid needing per model tokenizers, several of which aren't available as tiktoken encodings. However if you feel we need an exact token count, I can implement a different method for it.

I think this is okay as long as you make it clear in the documentation that the token limits are not completely guaranteed to be respected.

I think the cheap method is fine, but I would like some quantification on the 'few percent points' bit. Documented benchmarks somewhere or external reference -- just need rough, but real numbers.

Magdalen Dobson Manohar (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

Magdalen Dobson Manohar (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?

the different license was a leftover from when cosmos-retriever was a standalone project. i’ve removed the separate apache license and updated the package metadata and readme to use the repo’s top-level mit license, so they are now consistent.

…y, hygiene - Docstrings: run_* backends, Tool ABC, Reranker/ContextualReranker, BoundedTTLCache/_RetrieverPool, config module+classes, GET/PATCH /config - Tests: add test_compiler.py (strategies/filters/injection-safety) and test_anthropic_budget.py; add pytest to dev extras - agent_loop: bring anthropic backend to budget/prune/reject/clamp/timing parity with chat/responses - config: add per-corpus embed_dimensions (MRL output dims) threaded through the query embedder - docs: English-only stopwords note, schema_override canonical-fields clarification, query-injection explanation, token-counter approximation note - .env.example: expose all settings with defaults; correct EMBED_ENDPOINT vs AZURE_OPENAI_* - corpus_registry.json: replace personal values with labelled full + minimal examples - license: align sub-package to repo MIT (remove Apache LICENSE, update pyproject + README); remove redundant .python-version

Unit tests: config, tools, retriever, planner, strategies, document_resolvers, normalization, embeddings, paths, expressions, security, server, rerank, token counting, cross-tokenizer comparison. Opt-in (RUN_SKF_LIVE) live E2E suites against skf-rag-test: agentic search, live schema discovery, cross-collection RRF. tests/conftest.py isolates unit tests from ambient .env config. Module docstring updates across src/.

Don Dennis (metastableB) · 2026-08-07T19:17:30Z

+### Added
+- **`agentic_search` tool**: Runs a multi-turn retrieval agent — built from
+  scratch for this toolkit — against a Cosmos DB corpus and
+  returns ranked, curated documents that best answer the query. The agent


Your-org? > The agent issues hybrid (vector + full-text) RRF searches, optionally reranks with Qwen3-Reranker-8B, reads full documents, and prunes its context across multiple turns Also can you simplify this sentence? Follow the 'pseudo-code pattern' "Given a query, the agentic search tool will (1) do step A, (2) do step b, (3) do step c --" The content is fine as is. Just hard to read.

Apologies, the readme and changelog are an extremely old version from the research side of things so needed to be updated. Pushed all the updates necessary to make sure it is up to date.

Don Dennis (metastableB) · 2026-08-07T19:20:16Z

 | `text_search` | Search for documents where a property contains a search phrase |
 | `vector_search` | Perform vector search using Azure OpenAI embeddings |
 | `hybrid_search` | Perform hybrid search combining vector similarity and full-text keyword search using Reciprocal Rank Fusion (RRF) |
+| `agentic_search` | Run a multi-turn retrieval agent (built from scratch for this toolkit) against a Cosmos DB corpus. Backed by the bundled [`cosmos-retriever/`](cosmos-retriever/) FastAPI service; see [docs/AGENTIC_SEARCH.md](docs/AGENTIC_SEARCH.md) for setup and per-corpus configuration. |


"Perform multi-turn retrieval with the help of a configurable agent: An agent rewrites the queries, issues tool calls against the configured corpus/containers and returns responses. See docs[] for config"

We don't need to have sentences like "build from scratch for this toolkit", "backed by bundled fast API service" etc. This is just llm bleeding context.

A readme at the repo root should not assume context -- in fact it should be defining/providing context.

The additional details added here can go into the chagelog for instance, as changelog is typically for people who have context of the repo. (Not that you need to add them -- just trying to scope out README vs changelog)

Don Dennis (metastableB) · 2026-08-07T19:23:21Z

+# OPTIONAL: agentic_search TOOL (Cosmos retriever HTTP service)
+# ============================================================================
+# The `agentic_search` MCP tool calls the trained Harness-1 multi-turn
+# retrieval agent, which runs as a long-lived FastAPI service started with


"The Harness-1 multi-turn retrieval agent"

What is "the" harness-1 here? This is an example env config -- please don't assume context.

Otherwise this doc string is good. Explains clearly what the config options are, what the defaults are and what they do.

Don Dennis (metastableB) · 2026-08-07T19:24:51Z

+/// <summary>
+/// Tests for <see cref="AgenticSearchExecutor"/>. Stands in for the
+/// cosmos-retriever FastAPI service with a tiny in-process
+/// <see cref="HttpListener"/> so we can verify the executor's response


What is this <see cref=".."> syntax?

Aryan Saboo (@aryan-410)

Don Dennis (metastableB) · 2026-08-07T19:27:26Z

@@ -0,0 +1,143 @@
+# Cosmos Retriever (Python helper)
+
+A Python library + FastAPI service that runs a multi-turn search agent


What is the library part here? is there a reusable library ? or is this a feature implementing (a) a search call and (b) the service backing up the search call?

Don Dennis (metastableB) · 2026-08-07T20:19:43Z

+The number of queries allowed to run at the same time is capped, 
+so a burst of searches can't overwhelm the account. 
+
+the cap defaults to a sensible value 


Fix formatting. o/w good doc string

removed and fixed!

Don Dennis (metastableB) · 2026-08-07T20:22:38Z

+"""Comprehensive tests for the token-budget system.
+
+Covers, against the REAL agent loops (with a scripted fake LLM client + fake
+tools), every feature ported from upstream harness/agent.py:


Do we need to refer to upstream harness/agent.py

Removed the reference!

Don Dennis (metastableB) · 2026-08-07T20:24:16Z

@@ -0,0 +1,246 @@
+"""Security-focused tests for cosmos-retriever.
+
+Consolidates the safety-critical behaviours that are easy to regress:


Don Dennis (metastableB) · 2026-08-07T20:25:22Z

@@ -0,0 +1,279 @@
+"""Exhaustive tests for `cosmos_retriever.retrieval.planner`.
+
+RetrievalPlanner only *decides*: it inspects schema + capabilities (+ policy)


Simplify this sentence please. "RetrievalPlanner only decides" <-- What does this mean?

Don Dennis (metastableB) · 2026-08-07T20:29:30Z

+where they live. the capabilities (see the capabilities file) say which of those
+fields are truly searchable, and how. 
+
+A request may also pin a mode outright, 


fix formatting here. O/w good doc string

Magdalen Dobson Manohar (magdalendobson)

Thank you! The documentation and testing look much better now. Please take a look and make sure all comments are properly resolved before asking for Don's approval, but no outstanding strict blockers.

Docstring formatting/clarity (planner, executor, inference, test_planner, test_token_budget). Removed internal upstream/Harness references. Added adversarial reranker tests (degenerate/malformed/mismatched outputs) and a token-count result table. Added How-to-run steps to end-to-end tests plus tests/README.md. Reworked repo README, CHANGELOG, .env.example, cosmos-retriever README, and the C# test docstring for context-free clarity.

Don Dennis (metastableB)

Overall looks good. There are couple minor comments not addressed in the README, but otherwise this is in a good place.

Don Dennis (metastableB) · 2026-08-07T22:38:23Z

+2. Grant your identity the **Cosmos DB Built-in Data Reader** role on the account:
+   ```bash
+   az cosmosdb sql role assignment create \
+     --account-name skf-rag-test --resource-group DiskANN_development \


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 Saboo (aryan-410) mentioned this pull request Jul 23, 2026

Feat/agentic search cosmos retriever #148

Closed

Aryan Saboo (aryan-410) added 2 commits July 24, 2026 19:56

Aryan Saboo (aryan-410) commented Jul 27, 2026

View reviewed changes

Magdalen Dobson Manohar (magdalendobson) reviewed Jul 28, 2026

View reviewed changes

Don Dennis (metastableB) reviewed Jul 28, 2026

View reviewed changes

Magdalen Dobson Manohar (magdalendobson) reviewed Jul 28, 2026

View reviewed changes

Magdalen Dobson Manohar (magdalendobson) suggested changes Jul 28, 2026

View reviewed changes

Magdalen Dobson Manohar (magdalendobson) reviewed Jul 28, 2026

View reviewed changes

Aryan Saboo (aryan-410) added 3 commits August 4, 2026 00:01

chore(tests): move unit tests into tests/unit; drop retrieval-system doc

a4d71f3

Don Dennis (metastableB) suggested changes Aug 7, 2026

View reviewed changes

Magdalen Dobson Manohar (magdalendobson) approved these changes Aug 7, 2026

View reviewed changes

Don Dennis (metastableB) approved these changes Aug 7, 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

		@@ -0,0 +1,143 @@
		# Cosmos Retriever (Python helper)

		A Python library + FastAPI service that runs a multi-turn search agent

		@@ -0,0 +1,246 @@
		"""Security-focused tests for cosmos-retriever.

		Consolidates the safety-critical behaviours that are easy to regress:

		@@ -0,0 +1,279 @@
		"""Exhaustive tests for `cosmos_retriever.retrieval.planner`.

		RetrievalPlanner only decides: it inspects schema + capabilities (+ policy)

Uh oh!

Conversation

Aryan Saboo (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!

Aryan Saboo (aryan-410) Aug 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

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!

Aryan Saboo (aryan-410) Aug 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Aryan Saboo (aryan-410) Aug 3, 2026 •

edited

Loading

Aryan Saboo (aryan-410) Aug 3, 2026 •

edited

Loading

Aryan Saboo (aryan-410) Aug 3, 2026 •

edited

Loading