Cross-collection agentic retrieval + live schema discovery

.env.example

@@ -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 a multi-turn retrieval agent, which runs
+# as a long-lived FastAPI service started with
+# `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
 # ============================================================================

.gitignore

@@ -5311,3 +5311,5 @@ node_modules/
 .venv/Scripts/python.exe
 .venv/Scripts/pythonw.exe
 .venv/Scripts/tqdm.exe
+foundry-harness/**/__pycache__/
+*.pyc

CHANGELOG.md

@@ -5,6 +5,31 @@ 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 against a Cosmos
+  DB corpus and returns ranked, curated documents that best answer the query.
+  Given a query, the agent will (1) issue hybrid (vector + full-text) RRF
+  searches, (2) optionally rerank the hits with Qwen3-Reranker-8B, (3) read full
+  documents, and (4) prune its context across multiple turns. Implemented as a
+  subprocess call into the companion [`cosmos-retriever/`](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

README.md

@@ -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` | Perform multi-turn retrieval with a configurable agent. The agent rewrites the query, issues tool calls against the configured corpus and containers, and returns the most relevant documents. See [docs/AGENTIC_SEARCH.md](docs/AGENTIC_SEARCH.md) for setup and configuration. |
 
 ## Project Structure
 

cosmos-retriever/.env.example

@@ -0,0 +1,84 @@
+# =============================================================================
+# Cosmos Retriever configuration (Python service)
+# =============================================================================
+# Every setting read by `RetrieverSettings` (config.py) is listed here with its
+# default. Values load from environment variables or a `.env` / `.env.local` file
+# at the repo root. Required keys are uncommented with placeholders; optional keys
+# are commented out showing their default. Variable names are case-insensitive.
+#
+# NOTE: this file configures the *Python retriever service*. The .NET MCP server
+# uses the separate top-level `../.env.example`.
+
+# ----- Inference backend -----
+# "openai_responses" (default): OpenAI-compatible /responses model (reasoning
+#   models such as gpt-5.x).
+# "openai_chat": OpenAI-compatible /chat/completions model (Azure AI Foundry
+#   deployment, OpenAI, local server, ...).
+# "anthropic_messages": Anthropic Messages API (e.g. Claude on Azure AI Foundry).
+INFERENCE_BACKEND=openai_responses
+
+# ----- LLM endpoint (drives the retrieval agent) -----
+# For Azure AI Foundry: CHAT_BASE_URL is the endpoint URL, CHAT_MODEL 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=
+# CHAT_TEMPERATURE=0.7          # sampling temperature (chat backend)
+# CHAT_MAX_TOKENS=4096          # max output tokens per model turn
+# CHAT_MAX_TURNS=20             # max model<->tool round-trips per search
+# CHAT_REASONING_EFFORT=        # low|medium|high (openai_responses reasoning models only)
+# anthropic_messages only:
+# ANTHROPIC_VERSION=2023-06-01
+# ANTHROPIC_AUTH_HEADER=x-api-key
+
+# ----- Cosmos DB target (required) -----
+ACCOUNT_URI=https://your-cosmos-account.documents.azure.com:443/
+COSMOS_DATABASE=your-database-name
+COSMOS_CORPUS_CONTAINER=your-corpus-container
+# COSMOS_KEY=                          # unset -> AzureCliCredential (default)
+# COSMOS_USE_DEFAULT_CREDENTIAL=false  # true -> use the DefaultAzureCredential chain
+
+# ----- Embeddings for SearchCorpusTool (required) -----
+# Default embedding endpoint/key/model, used when a corpus is NOT in the registry.
+OPENAI_API_KEY=sk-...
+OPENAI_EMBEDDING_MODEL=text-embedding-3-small
+# EMBED_ENDPOINT=               # OpenAI (api.openai.com) if unset. For Azure pass
+#                               #   https://<resource>.services.ai.azure.com/openai/v1;
+#                               #   for a local server pass http://host:port/v1
+# OPENAI_EMBEDDING_DIMENSIONS=  # request truncated (MRL) output dims, e.g. 2560 to
+#                               #   match a Qwen3-Embedding corpus. Unset = model native.
+# EMBED_QUERY_INSTRUCTION=      # optional "Instruct:" prefix (some Qwen embedders)
+
+# ----- Per-corpus embedding registry (optional) -----
+# Map a container to its own account / database / embedding endpoint+model+dims.
+# Provide ONE of these. A registry entry references its key via `embed_api_key_env`
+# (any env var name you choose, e.g. AZURE_OPENAI_EMBED_API_KEY below).
+# CORPUS_REGISTRY_FILE=corpus_registry.json
+# CORPUS_REGISTRY={"db/container": {"account_uri": "...", "embed_model": "..."}}
+# AZURE_OPENAI_EMBED_API_KEY=   # example key referenced by a registry entry
+
+# ----- Reranker (optional; pick at most one) -----
+# BASETEN_API_KEY=              # Baseten Qwen3-Reranker-8B classify
+# BASETEN_MODEL_URL=https://model-xyz.api.baseten.co/environments/production/sync
+# VLLM_RERANKER_URL=http://127.0.0.1:8011   # local vLLM Qwen3-Reranker /score
+
+# ----- Retriever budgets & limits (optional) -----
+# COSMOS_RETRIEVER_MAX_TURNS=35              # hard cap on agent turns
+# COSMOS_RETRIEVER_THRESHOLD_BUDGET=16384    # soft cap: prune-or-conclude kicks in
+# COSMOS_RETRIEVER_TOKEN_BUDGET=32268        # hard cap on transcript tokens
+# COSMOS_RETRIEVER_SEARCH_DISPLAY_LIMIT=15   # rows shown per search result
+# COSMOS_RETRIEVER_RAW_QUERY_ENABLED=true    # expose the read-only execute_query tool
+# COSMOS_RETRIEVER_SCHEMA_OVERRIDE=          # JSON: document_id_path, chunk_order_path, ...
+# Note: the per-tool output clamp (~4096) and spillage fraction (0.5) are code-level
+# constants in agent_loop.py (_DEFAULT_TOOL_OUTPUT_BUDGET / _DEFAULT_SPILLAGE_FRACTION),
+# not env-configurable.
+
+# ----- Retriever pool cache (optional) -----
+# COSMOS_RETRIEVER_CACHE_MAX_ENTRIES=32      # max pooled retriever engines (LRU)
+# COSMOS_RETRIEVER_CACHE_TTL_SECONDS=900.0   # engine TTL (seconds) before rebuild
+
+# ----- HTTP server -----
+HOST=0.0.0.0
+PORT=9000
+LOG_LEVEL=info

cosmos-retriever/.github/workflows/ci.yml

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

cosmos-retriever/.gitignore

@@ -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/

cosmos-retriever/README.md

@@ -0,0 +1,158 @@
+# Cosmos Retriever (Python helper)
+
+This package runs a multi-turn search agent against an Azure Cosmos DB corpus and
+returns the curated documents as JSON. The agent is model-agnostic: it drives any
+OpenAI-compatible endpoint (the `/responses` or `/chat/completions` APIs) or an
+Anthropic Messages endpoint. The same code is available three ways, an importable
+Python package (`CosmosRetriever`), a FastAPI service
+(`python -m cosmos_retriever serve`), and a one-shot CLI
+(`python -m cosmos_retriever search`).
+
+The [Azure Cosmos DB MCP Toolkit](../MCPToolKit/)'s `agentic_search` tool calls
+this service's `POST /search` endpoint over HTTP.
+
+```text
+  Claude Desktop / AI Foundry / VS Code
+            │
+            │   MCP streamable-HTTP
+            ▼
+  Azure Cosmos DB MCP Toolkit  (.NET)
+   ├─ list_databases / list_collections / ...   (8 native tools)
+   └─ agentic_search                            ◀─── 9th tool
+            │
+            │   HTTP: POST http://127.0.0.1:9000/search
+            ▼
+  cosmos_retriever  (this package, FastAPI + uvicorn)
+   ├─ TokenBudgetRetrievalSubagent
+   ├─ SearchCorpus / Grep / ReadDocument / PruneChunks tools
+   └─ VLLMHarmonyInferenceModel  ──► vLLM /v1/completions (token-IDs)
+                                     Cosmos DB hybrid RRF
+                                     Azure OpenAI embeddings
+                                     Qwen3-Reranker (Baseten or local vLLM)
+```
+
+## Install
+
+```bash
+cd cosmos-retriever
+uv venv --python 3.11 .venv
+uv pip install --python .venv/bin/python -e ".[dev]"
+```
+
+## HTTP service
+
+The MCP Toolkit talks to a long-lived FastAPI service. Start it with:
+
+```bash
+python -m cosmos_retriever serve            # binds HOST:PORT (default 0.0.0.0:9000)
+```
+
+Endpoints:
+
+| Method & path | Body / response |
+|---|---|
+| `GET /health` | `{"status": "ok"}` |
+| `POST /search` | request `{"query": str, "maxDocuments": int, "database": str?, "container": str?}` → the JSON result below |
+
+Example request to test a running service (the query and its answer depend on the corpus you configured):
+
+```bash
+curl -s http://127.0.0.1:9000/search \
+  -H 'content-type: application/json' \
+  -d '{"query": "Who discovered radium?", "maxDocuments": 5}'
+```
+
+## CLI
+
+To smoke-test locally, use the command below to query the service with a single
+question and print the answer documents. JSON goes to **stdout**, logs go to
+**stderr**.
+
+```bash
+python -m cosmos_retriever search \
+  --query "Who discovered radium?" \
+  --max-documents 5
+```
+
+Expected output (same schema returned by `POST /search`):
+```json
+{
+  "query": "Who discovered radium?",
+  "num_turns": 5,
+  "elapsed_s": 32.3,
+  "documents": [
+    { "id": "96308__3", "rank": 0, "justification": "...", "text": "..." }
+  ]
+}
+```
+
+## Configuration
+
+All settings come from environment variables, or from a `.env` / `.env.local`
+file in the `cosmos-retriever/` directory. Precedence is real environment
+variables first, then `.env.local`, then `.env`. Use `.env.local` for local
+secrets and overrides, it is gitignored. Required settings:
+
+| Variables | Purpose |
+|---|---|
+| `INFERENCE_BACKEND`, `CHAT_BASE_URL`, `CHAT_API_KEY`, `CHAT_MODEL` | The backend, endpoint, key, and model for the LLM that drives the agent (see Inference backend below) |
+| `ACCOUNT_URI`, `COSMOS_DATABASE`, `COSMOS_CORPUS_CONTAINER` | The Cosmos account, database, and container to search |
+| `OPENAI_API_KEY`, `OPENAI_EMBEDDING_MODEL` | The embeddings key and model (set `EMBED_ENDPOINT` for Azure or a local server) |
+
+Each row is a group of related settings, not alternatives. See
+[`.env.example`](.env.example) for the complete list and defaults.
+
+### Inference backend
+
+`INFERENCE_BACKEND` selects what drives the retrieval agent:
+
+| Value | Model | Endpoint vars |
+|---|---|---|
+| `openai_responses` *(default)* | Any OpenAI-compatible `/responses` model (reasoning models such as gpt-5.x). | `CHAT_BASE_URL`, `CHAT_API_KEY`, `CHAT_MODEL`, optional `CHAT_API_VERSION` |
+| `openai_chat` | Any OpenAI-compatible `/chat/completions` model (Azure AI Foundry deployment, OpenAI, local server, ...). | `CHAT_BASE_URL`, `CHAT_API_KEY`, `CHAT_MODEL`, optional `CHAT_API_VERSION` |
+| `anthropic_messages` | Any Anthropic Messages API endpoint — e.g. Claude on Azure AI Foundry (served over the Messages API, not OpenAI-shaped). | `CHAT_BASE_URL`, `CHAT_API_KEY`, `CHAT_MODEL`, optional `ANTHROPIC_VERSION`, `ANTHROPIC_AUTH_HEADER` |
+
+All backends drive the same Cosmos tools, so retrieval quality depends on the
+chosen model's tool-use ability. Example (Azure AI Foundry):
+
+```bash
+INFERENCE_BACKEND=openai_chat \
+CHAT_BASE_URL=https://your-resource.services.ai.azure.com/openai/v1 \
+CHAT_API_KEY=... \
+CHAT_MODEL=gpt-4o \
+python -m cosmos_retriever serve
+```
+
+### Optional reranker
+
+An independent reranker model can be configured to reorder the retrieved
+documents by relevance before they are returned, which improves the quality of
+the final ranking. It is optional. Without it, the agent keeps the raw retrieval
+order. Configure at most one of:
+
+- `VLLM_RERANKER_URL`, a local vLLM `/score` endpoint serving Qwen3-Reranker-8B.
+- `BASETEN_API_KEY` and `BASETEN_MODEL_URL`, a Baseten Qwen3-Reranker-8B deployment.
+
+## Layout
+
+```text
+src/cosmos_retriever/
+  __init__.py        # CosmosRetriever, RetrievalResult, RetrievedDocument
+  __main__.py        # `python -m cosmos_retriever {search,serve}`
+  server.py          # FastAPI app: GET /health + POST /search
+  retriever.py       # CosmosRetriever facade
+  agent.py           # 3 agent classes + prune_chunks_from_trajectory
+  tools.py           # SearchCorpus / Grep / ReadDocument / PruneChunks
+  trajectory.py      # Action / Observation / Trajectory + Harmony rendering
+  rerank.py          # Reranker ABC + Baseten + local-vLLM
+  inference/
+    base.py          # AgentInferenceModel ABC
+    vllm.py          # VLLMHarmonyInferenceModel (httpx → /v1/completions)
+  prompts.py         # retrieval subagent system prompt
+  config.py          # RetrieverSettings (pydantic-settings)
+  utils.py
+```
+
+## License
+
+MIT — this package is covered by the repository's top-level [LICENSE](../LICENSE).