Skip to content
Open
Show file tree
Hide file tree
Changes from 19 commits
Commits
Show all changes
23 commits
Select commit Hold shift + click to select a range
970dd29
feat(cosmos-retriever): vendor the Cosmos Retriever Python service
Jun 26, 2026
53db022
feat(agentic_search): add agentic_search MCP tool
Jun 26, 2026
9db5640
chore: remove benchmark scripts and dev artifacts from toolkit
aryan-410 Jun 30, 2026
905acc0
chore: remove generated datagen/splits folder
aryan-410 Jun 30, 2026
43eb203
refactor(cosmos-retriever): extract vLLM runtime policy, drop eval ha…
aryan-410 Jun 30, 2026
03aa2e0
chore(toolkit): polish agentic_search descriptions, align maxDocument…
aryan-410 Jun 30, 2026
dc0d8fc
chore(cosmos-retriever): remove datagen package and unit tests folder
aryan-410 Jun 30, 2026
a7d8b2f
feat(cosmos-retriever): expose per-query agent trajectory on Retrieva…
aryan-410 Jun 30, 2026
2bb468a
feat(cosmos-retriever): endpoint-only agentic search + schema-decoupl…
aryan-410 Jul 15, 2026
4e2499d
docs(cosmos-retriever): add detailed retrieval-system and agentic-wor…
aryan-410 Jul 15, 2026
6f4a799
chore(cosmos-retriever): remove benchmark scripts and foundry-harness
aryan-410 Jul 15, 2026
ab601aa
refactor(cosmos-retriever): strip all docstrings and comments
aryan-410 Jul 15, 2026
0be3253
refactor(cosmos-retriever): replace legacy.py with defaults.py (defau…
aryan-410 Jul 15, 2026
a45e09d
docs: remove all Harness-1/vLLM lineage; state agent is built from sc…
aryan-410 Jul 15, 2026
cda73a5
refactor(cosmos-retriever): rename inference/openai_chat.py -> agent_…
aryan-410 Jul 15, 2026
d25a17b
feat(cosmos-retriever): add anthropic_messages backend (Claude via Fo…
aryan-410 Jul 15, 2026
3f920b5
feat(cosmos-retriever): cross-collection search + live schema discovery
aryan-410 Jul 23, 2026
d1145f2
feat(cosmos-retriever): token budgeting, cross-turn dedup, and real t…
aryan-410 Jul 24, 2026
438c296
Tombstone chunk pruning, item-is-document read path, per-database ret…
aryan-410 Jul 26, 2026
b53242f
Address PR review: docstrings, compiler+anthropic tests, budget parit…
aryan-410 Aug 4, 2026
36b8553
test: add exhaustive unit + SKF live E2E suites; src docstrings
aryan-410 Aug 7, 2026
a4d71f3
chore(tests): move unit tests into tests/unit; drop retrieval-system doc
aryan-410 Aug 7, 2026
0351bdd
docs+tests: address PR review feedback
aryan-410 Aug 7, 2026
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
16 changes: 16 additions & 0 deletions .env.example
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,22 @@ ASPNETCORE_LOGGING__LOGLEVEL__DEFAULT=Information
# Optional: Server URLs (default: http://+:8080)
ASPNETCORE_URLS=http://+:8080

# ============================================================================
# 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

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

"The Harness-1 multi-turn retrieval agent"

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

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

# `python -m cosmos_retriever serve`. See docs/AGENTIC_SEARCH.md.
# Both vars below are optional with sensible defaults; if the service is not
# reachable, agentic_search simply returns a clean JSON error envelope to the
# caller.

# Base URL of the cosmos-retriever FastAPI service (default http://127.0.0.1:9000).
# COSMOS_RETRIEVER_URL=http://127.0.0.1:9000

# Per-request wall-clock cap in seconds (default 600).
# COSMOS_RETRIEVER_TIMEOUT_S=600

# ============================================================================
# DOCKER COMPOSE NOTES
# ============================================================================
Expand Down
2 changes: 2 additions & 0 deletions .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -5311,3 +5311,5 @@ node_modules/
.venv/Scripts/python.exe
.venv/Scripts/pythonw.exe
.venv/Scripts/tqdm.exe
foundry-harness/**/__pycache__/
*.pyc
26 changes: 26 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,32 @@ All notable changes to this project will be documented in this file.
The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).

## [1.2.0] - 2026-06-18

### 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

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@aryan-410 Aryan Saboo (aryan-410) Aug 7, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

issues hybrid (vector + full-text) RRF searches, optionally reranks with
Qwen3-Reranker-8B, reads full documents, and prunes its context across
multiple turns. Implemented as a subprocess call into the companion
[`cosmos-retriever`](https://github.com/your-org/cosmos-retriever)
Python package; see [`docs/AGENTIC_SEARCH.md`](docs/AGENTIC_SEARCH.md) for
the deployment story.
- Optional `database` and `container` arguments on `agentic_search` so a
single MCP server can target multiple Cosmos corpora at request time. When
the corpus registry (`CORPUS_REGISTRY` / `CORPUS_REGISTRY_FILE`) is set
in the host environment, the matching account, database, and embedding
model are picked automatically per call.
- New service: `AgenticSearchExecutor` (subprocess lifecycle, timeout, error
envelope generation).
- New env vars: `COSMOS_RETRIEVER_PYTHON`, `COSMOS_RETRIEVER_DIR`,
`COSMOS_RETRIEVER_TIMEOUT_S` — see [`.env.example`](.env.example).

### Changed
- `AppState` now also exposes `ILoggerFactory` so static `[McpServerTool]`
methods can obtain a properly-named logger.

## [1.1.2] - 2026-05-29

### Added
Expand Down
1 change: 1 addition & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,7 @@ This toolkit provides:
| `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. |

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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)


## Project Structure

Expand Down
58 changes: 58 additions & 0 deletions cosmos-retriever/.env.example
Original file line number Diff line number Diff line change
@@ -0,0 +1,58 @@
# ----- Inference backend -----

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There is an .env.example in cosmos-retriever and at the top level?

Also expose defaults like prune budget here -- up to you

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yes, there are two separate .env.example files because they configure two separate services. the top-level file is for the .net mcp server, while cosmos-retriever/.env.example is for the python retriever service. there is also a separate one for the sample client.

i updated the python .env.example to make it a complete reference for the service rather than only showing the minimum required variables. it now includes every configurable setting, grouped by llm, cosmos, embeddings, corpus registry, reranking, budgets, cache, and server settings, with the defaults documented alongside them.

this includes the threshold and token budgets, search limits, max turns, cache sizing, schema overrides, and the other settings that were previously missing. i also corrected a few stale examples that referenced variables the python service does not actually read.

the lower-level prune settings, such as the per-tool output budget and spillage fraction, were previously code-only constants. i have exposed those as config settings as well so the .env.example can serve as the single place people refer to for all defaults instead of having to inspect the implementation.

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yes!

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we have a repro/test files for these?

# "openai_responses" (default): any OpenAI-compatible /responses model
# (reasoning models such as gpt-5.x).
# "openai_chat": any OpenAI-compatible /chat/completions model (Azure AI Foundry
# deployment, OpenAI, local server, ...).
# "anthropic_messages": any Anthropic Messages API endpoint (e.g. Claude on
# Azure AI Foundry, served over the Messages API).
INFERENCE_BACKEND=openai_responses

# ----- LLM endpoint (openai_responses / openai_chat / anthropic_messages) -----
# For Azure AI Foundry, CHAT_BASE_URL is the endpoint URL and CHAT_MODEL is the
# deployment name. Set CHAT_API_VERSION to use the Azure OpenAI client.
# CHAT_BASE_URL=https://your-resource.services.ai.azure.com/openai/v1
# CHAT_API_KEY=
# CHAT_MODEL=gpt-5.2
# CHAT_API_VERSION=
# For anthropic_messages, optionally override the version / auth header:
# ANTHROPIC_VERSION=2023-06-01
# ANTHROPIC_AUTH_HEADER=x-api-key
# CHAT_TEMPERATURE=0.7
# CHAT_MAX_TOKENS=4096
# CHAT_MAX_TURNS=20

# ----- Cosmos DB (required) -----
# The corpus container must already be ingested with `id`, `docid`, `chunk_idx`,
# `text`, and `embedding` fields.
ACCOUNT_URI=https://your-cosmos-account.documents.azure.com:443/
COSMOS_DATABASE=your-database-name
COSMOS_CORPUS_CONTAINER=your-corpus-container
# Optional: leave unset to use AzureCliCredential / DefaultAzureCredential.
# COSMOS_KEY=

# ----- Embeddings for SearchCorpusTool (required) -----
# OpenAI by default; set AZURE_OPENAI_* to route through Azure OpenAI instead.
OPENAI_API_KEY=sk-...
OPENAI_EMBEDDING_MODEL=text-embedding-3-small
# AZURE_OPENAI_ENDPOINT=https://your-resource.openai.azure.com
# AZURE_OPENAI_API_KEY=
# Optional instruction prepended to embedding queries (used by some Qwen embedders).
# EMBED_QUERY_INSTRUCTION=

# ----- Reranker (optional) -----
# Pick at most one of these. Leave both unset to disable reranking.
# Baseten:
# BASETEN_API_KEY=
# BASETEN_MODEL_URL=https://model-xyz.api.baseten.co/environments/production/sync
# Local vLLM Qwen3-Reranker-8B (run on a separate port):
# VLLM_RERANKER_URL=http://127.0.0.1:8011

# ----- Retriever budget knobs (optional) -----
# COSMOS_RETRIEVER_MAX_TURNS=35
# COSMOS_RETRIEVER_THRESHOLD_BUDGET=16384
# COSMOS_RETRIEVER_TOKEN_BUDGET=32268

# ----- HTTP server -----
HOST=0.0.0.0
PORT=9000
LOG_LEVEL=info
37 changes: 37 additions & 0 deletions cosmos-retriever/.github/workflows/ci.yml

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Original file line number Diff line number Diff line change
@@ -0,0 +1,37 @@
name: ci

on:
push:
branches: [main]
pull_request:

concurrency:
group: ci-${{ github.ref }}
cancel-in-progress: true

jobs:
lint-and-test:
runs-on: ubuntu-latest
strategy:
matrix:
python-version: ["3.11", "3.12"]
steps:
- uses: actions/checkout@v4

- uses: actions/setup-python@v5
with:
python-version: ${{ matrix.python-version }}

- name: Install uv
uses: astral-sh/setup-uv@v3
with:
enable-cache: true

- name: Install package with dev extras
run: uv pip install --system -e ".[dev]"

- name: Ruff lint
run: ruff check src tests

- name: Pytest
run: pytest -q
39 changes: 39 additions & 0 deletions cosmos-retriever/.gitignore
Original file line number Diff line number Diff line change
@@ -0,0 +1,39 @@
# --- Python ---
__pycache__/
*.py[cod]
*$py.class
*.egg-info/
.eggs/
build/
dist/
.coverage
.coverage.*
htmlcov/
.pytest_cache/
.mypy_cache/
.ruff_cache/

# --- Virtual envs ---
.venv/
venv/
env/

# --- IDE ---
.vscode/
.idea/
*.swp

# --- Secrets / local config ---
.env
.env.local
.env.*.local
.env.*
!.env.example

# --- Logs / scratch ---
*.log
tmp/
runs/

# --- Build artefacts ---
src/*.egg-info/
1 change: 1 addition & 0 deletions cosmos-retriever/.python-version
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
3.11

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm a bit confused as to why this is pinned here. It seems like this PR is to be merged into an existing repo, in which case the PR should not pin versions and dependencies ?

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

good point. .python-version was only a local pyenv/uv hint for choosing the default interpreter in this directory. it was not used by ci, did not affect package consumers, and did not pin any dependencies.

the actual compatibility requirement is requires-python = ">=3.11" in pyproject.toml. ci also tests both python 3.11 and 3.12 explicitly through setup-python, so .python-version was redundant and could give the impression that the project only supports 3.11.

i removed .python-version. the authoritative version range remains in pyproject.toml, ci continues to cover both supported versions, and the readme still shows how to create a local 3.11 environment explicitly.

Loading