diff --git a/.env.example b/.env.example index 2b9b14b..8a81e8d 100644 --- a/.env.example +++ b/.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 the trained Harness-1 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 # ============================================================================ diff --git a/.gitignore b/.gitignore index 93cdbfa..6295ca2 100644 --- a/.gitignore +++ b/.gitignore @@ -5311,3 +5311,5 @@ node_modules/ .venv/Scripts/python.exe .venv/Scripts/pythonw.exe .venv/Scripts/tqdm.exe +foundry-harness/**/__pycache__/ +*.pyc diff --git a/CHANGELOG.md b/CHANGELOG.md index 44bde2a..baa56cd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -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 + 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 diff --git a/README.md b/README.md index 2fd336c..79da466 100644 --- a/README.md +++ b/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` | 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. | ## Project Structure diff --git a/cosmos-retriever/.env.example b/cosmos-retriever/.env.example new file mode 100644 index 0000000..dbb992b --- /dev/null +++ b/cosmos-retriever/.env.example @@ -0,0 +1,58 @@ +# ----- Inference backend ----- +# "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 diff --git a/cosmos-retriever/.github/workflows/ci.yml b/cosmos-retriever/.github/workflows/ci.yml new file mode 100644 index 0000000..86d3c81 --- /dev/null +++ b/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 diff --git a/cosmos-retriever/.gitignore b/cosmos-retriever/.gitignore new file mode 100644 index 0000000..8133b6c --- /dev/null +++ b/cosmos-retriever/.gitignore @@ -0,0 +1,37 @@ +# --- 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 + +# --- Logs / scratch --- +*.log +tmp/ +runs/ + +# --- Build artefacts --- +src/*.egg-info/ diff --git a/cosmos-retriever/.python-version b/cosmos-retriever/.python-version new file mode 100644 index 0000000..2c07333 --- /dev/null +++ b/cosmos-retriever/.python-version @@ -0,0 +1 @@ +3.11 diff --git a/cosmos-retriever/LICENSE b/cosmos-retriever/LICENSE new file mode 100644 index 0000000..29f81d8 --- /dev/null +++ b/cosmos-retriever/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/cosmos-retriever/QUICKSTART.md b/cosmos-retriever/QUICKSTART.md new file mode 100644 index 0000000..c7a65de --- /dev/null +++ b/cosmos-retriever/QUICKSTART.md @@ -0,0 +1,99 @@ +# Cosmos Retriever — Quickstart + +Spin up the `agentic_search` retriever service and test it locally on Windows (PowerShell). + +## 1. Prerequisites + +- Python 3.11 (fetched automatically by `uv`) +- `uv`, Azure CLI (`az`) +- An Azure Cosmos DB corpus container (fields: `id`, `docid`, `chunk_idx`, `text`, `embedding`; vector + FTS indexes enabled) +- An OpenAI-compatible LLM endpoint (e.g. Azure AI Foundry) +- An Azure OpenAI embeddings deployment (`text-embedding-3-small`) + +## 2. Install (one time) + +```powershell +cd cosmos-retriever +uv venv --python 3.11 .venv +uv pip install --python .venv\Scripts\python.exe -e ".[dev]" +``` + +## 3. Configure + +Edit [`.env`](.env) and replace every ``: + +| Variable | What to put | +|---|---| +| `CHAT_BASE_URL` | LLM endpoint, e.g. `https://.services.ai.azure.com/openai/v1` | +| `CHAT_API_KEY` | LLM API key | +| `CHAT_MODEL` | Deployment/model name, e.g. `gpt-5.2` | +| `ACCOUNT_URI` | `https://.documents.azure.com:443/` | +| `COSMOS_DATABASE` | Database name | +| `COSMOS_CORPUS_CONTAINER` | Corpus container name | +| `EMBED_ENDPOINT` | Azure OpenAI v1 base URL | +| `OPENAI_API_KEY` | Azure OpenAI embeddings key | + +Defaults already set: `INFERENCE_BACKEND=openai_responses`, embeddings `text-embedding-3-small`, no reranker, Azure CLI credential for Cosmos, port `9000`. + +## 4. Authenticate to Cosmos + +```powershell +az login +``` + +## 5. Run the service + +```powershell +.\run-retriever.ps1 +``` + +It binds `0.0.0.0:9000` and keeps its Cosmos/LLM/embedding clients warm. + +## 6. Test + +**Health check** (new terminal): + +```powershell +Invoke-RestMethod http://127.0.0.1:9000/health +# -> status : ok +``` + +**One-shot search via HTTP:** + +```powershell +$body = '{"query":"Who discovered radium?","maxDocuments":5}' +Invoke-RestMethod -Uri http://127.0.0.1:9000/search -Method Post -ContentType application/json -Body $body +``` + +**Or via the CLI** (JSON to stdout, logs to stderr): + +```powershell +.venv\Scripts\python.exe -m cosmos_retriever search --query "Who discovered radium?" --max-documents 5 +``` + +**Target a specific corpus** from `corpus_registry.json`: + +```powershell +$body = '{"query":"...","container":"enterprise_ragbench_corpus"}' +Invoke-RestMethod -Uri http://127.0.0.1:9000/search -Method Post -ContentType application/json -Body $body +``` + +Expected response shape: + +```json +{ + "query": "Who discovered radium?", + "num_turns": 5, + "elapsed_s": 32.3, + "documents": [ + { "id": "96308__3", "rank": 0, "justification": "...", "text": "..." } + ] +} +``` + +## Troubleshooting + +- **`status` not returned / connection refused** — the service isn't running; check the `run-retriever.ps1` terminal for errors. +- **Cosmos auth errors** — run `az login`; or set `COSMOS_KEY` in `.env`; or `COSMOS_USE_DEFAULT_CREDENTIAL=1` for managed identity. +- **Placeholder still in `.env`** — every `<...>` value must be replaced. +- **Wire into the MCP server** — set `COSMOS_RETRIEVER_URL=http://127.0.0.1:9000` in the repo-root `.env` and start the .NET server with `..\run-mcp-server.ps1`. diff --git a/cosmos-retriever/README.md b/cosmos-retriever/README.md new file mode 100644 index 0000000..8f31319 --- /dev/null +++ b/cosmos-retriever/README.md @@ -0,0 +1,141 @@ +# Cosmos Retriever (Python helper) + +A Python library + FastAPI service that runs a multi-turn search agent +(a fine-tuned `openai/gpt-oss-20b` served by vLLM, or any OpenAI-compatible +model) against an Azure Cosmos DB corpus and returns the curated documents as +JSON. + +The [Azure Cosmos DB MCP Toolkit](../MCPToolKit/)'s `agentic_search` tool +calls this service's `POST /search` endpoint over HTTP. A one-shot CLI is also +provided for local testing. + +```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 | + +```bash +curl -s http://127.0.0.1:9000/search \ + -H 'content-type: application/json' \ + -d '{"query": "Who discovered radium?", "maxDocuments": 5}' +``` + +## CLI + +A one-shot CLI for local testing. JSON goes to **stdout**, logs go to **stderr**. + +```bash +python -m cosmos_retriever search \ + --query "Who discovered radium?" \ + --max-documents 5 +``` + +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 a `.env` / `.env.local` file +at the repo root). Required: + +| Variable | Purpose | +|---|---| +| `VLLM_BASE_URL` | OpenAI-compatible vLLM endpoint serving the local model | +| `ACCOUNT_URI` / `COSMOS_DATABASE` / `COSMOS_CORPUS_CONTAINER` | Cosmos target | +| `OPENAI_API_KEY` *(or `AZURE_OPENAI_*`)* | Embeddings backend | + +### 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 (pick at most one): +- `BASETEN_API_KEY` + `BASETEN_MODEL_URL` — Baseten Qwen3-Reranker-8B classify +- `VLLM_RERANKER_URL` — local vLLM `/score` endpoint with Qwen3-Reranker-8B + +## 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 + +Apache 2.0. diff --git a/cosmos-retriever/corpus_registry.json b/cosmos-retriever/corpus_registry.json new file mode 100644 index 0000000..3f608fb --- /dev/null +++ b/cosmos-retriever/corpus_registry.json @@ -0,0 +1,17 @@ +{ + "browsecomp_corpus_container": { + "account_uri": "https://aryans-internship-cosmos.documents.azure.com:443/", + "database": "search_retrieval_database", + "embed_base_url": "https://embedding-west-us-resource.services.ai.azure.com/openai/v1", + "embed_api_key_env": "AZURE_OPENAI_EMBED_API_KEY", + "embed_model": "text-embedding-3-small" + }, + "enterprise_ragbench_corpus": { + "account_uri": "https://aryans-internship-cosmos-prov.documents.azure.com:443/", + "database": "search_retrieval_database", + "embed_base_url": "http://172.17.0.2:8002/v1", + "embed_api_key_env": null, + "embed_model": "qwen3-embed", + "embed_query_instruction": "Given a question, retrieve documents that answer it" + } +} diff --git a/cosmos-retriever/docs/AGENTIC_WORKFLOW.md b/cosmos-retriever/docs/AGENTIC_WORKFLOW.md new file mode 100644 index 0000000..10e0909 --- /dev/null +++ b/cosmos-retriever/docs/AGENTIC_WORKFLOW.md @@ -0,0 +1,275 @@ +# The Agentic Search Workflow + +This document explains the **end-to-end agentic retrieval workflow**: how a +natural-language question becomes a curated, ranked set of documents. It covers +the network entry point, the multi-turn search agent, the four tools it drives, +the inference backends, budgets/pruning, and how everything is configured. + +Where [RETRIEVAL_SYSTEM.md](RETRIEVAL_SYSTEM.md) describes the *plumbing* (how a +single query becomes safe Cosmos SQL), this document describes the *brain* (how an +LLM agent plans, issues many searches, prunes, and decides when it's done). + +--- + +## 1. The big picture + +```mermaid +flowchart TD + subgraph dotnet[".NET MCP Toolkit"] + T[agentic_search MCP tool] + end + subgraph py["Python service (this repo)"] + S[FastAPI server
POST /search] + P[_RetrieverPool
one CosmosRetriever per corpus] + R[CosmosRetriever
multi-turn agent loop] + TS[ToolSet
4 Cosmos tools] + RL[Retrieval layer
CorpusRetriever] + end + LLM[[OpenAI-compatible model
/responses or /chat]] + DB[(Azure Cosmos DB
NoSQL corpus)] + EMB[[Embedding endpoint]] + + T -->|HTTP POST| S --> P --> R + R <-->|tool calls| LLM + R --> TS --> RL --> DB + RL --> EMB + R -->|ranked documents JSON| S --> T +``` + +1. The .NET toolkit's **`agentic_search`** tool makes an HTTP `POST /search` to a + long-lived instance of the Python FastAPI service (keeping the Cosmos SDK, + embedding client, and tokenizer warm across calls). +2. The server routes to a per-corpus **`CosmosRetriever`**, which runs a + **multi-turn agent loop** against an OpenAI-compatible model. +3. The model drives four **Cosmos tools** (search / grep / read / prune), each of + which delegates to the schema-decoupled retrieval layer. +4. When the model is satisfied it emits ranked `` blocks, which are + parsed into the JSON response. + +--- + +## 2. The network entry point (`server.py`) + +A FastAPI app exposes two routes: + +- **`GET /health`** → `{"status": "ok"}` (liveness; never touches Cosmos or the model). +- **`POST /search`** → runs the agent and returns curated documents. + +Request body (`SearchRequest`): `query`, `maxDocuments` (1–30), optional +`database` / `container` overrides. + +**`_RetrieverPool`** lazily builds and caches **one `CosmosRetriever` per corpus** +(keyed by `(database, container)`), so a single process serves many corpora while +keeping heavy clients warm. Because each retriever holds *synchronous* Cosmos/HTTP +clients and per-call agent state that are **not** thread-safe: + +- Every request runs the (sync) search on a worker thread via + `anyio.to_thread.run_sync`. +- Same-corpus requests are **serialised with a per-corpus `asyncio.Lock`**; + different corpora run concurrently. + +--- + +## 3. The agent (`CosmosRetriever`, `retriever.py`) + +Constructed once per corpus. On init it: + +1. Resolves the **`CorpusConfig`** for the target container (via + `RetrieverSettings.resolve_corpus`, which consults `corpus_registry.json`). +2. Builds the **Cosmos** database client and the **embedding** client. +3. Builds the **`ToolSet`** (the four tools) wired to a `CorpusRetriever`. +4. Builds the **inference client** (chat or responses) and an optional **reranker**. +5. Loads a **tiktoken** encoder for token accounting/budgets. + +Its public method is **`search(query, *, max_documents, max_turns, +threshold_budget, token_budget)`**, returning a **`RetrievalResult`**: + +```python +RetrievalResult( + query, documents=[RetrievedDocument(id, text, justification, rank)], + num_turns, final_text, pool_doc_ids, elapsed_s, usage, trajectory, metadata, +) +``` + +`search()` dispatches to one of two backends based on `INFERENCE_BACKEND`. + +--- + +## 4. Inference backends (`inference/agent_loop.py`) + +All three backends drive the **same four Cosmos tools** via function-calling; +they differ only in the API surface: + +| Backend | Function | API | Use for | +|---|---|---|---| +| `openai_responses` | `run_responses_search` | `/responses` | Reasoning models (gpt-5.x) — exposes turn-level trajectory + reasoning tokens. | +| `openai_chat` | `run_chat_search` | `/chat/completions` | Generic OpenAI-compatible chat models. | +| `anthropic_messages` | `run_anthropic_search` | `/v1/messages` | Anthropic Messages API models — e.g. Claude on Azure AI Foundry (tool-use blocks). | + +Each backend runs the **agent loop** (up to `max_turns`, default 20): + +```mermaid +sequenceDiagram + participant M as Model + participant A as Agent loop + participant TS as ToolSet + participant DB as Cosmos + + A->>M: system prompt + query + tool schemas + loop until final answer or max_turns + M-->>A: tool call(s) (search / grep / read / prune) + A->>TS: execute tool(s) (in parallel where possible) + TS->>DB: compiled Cosmos SQL + DB-->>TS: rows + TS-->>A: formatted observations (with token counts) + A->>A: accumulate usage; enforce token budget + A-->>M: tool observations (+ over-budget nudge if needed) + end + M-->>A: final blocks + A->>A: parse documents, attach cached text, rank +``` + +Responsibilities inside the loop: +- **Tool-argument parsing** (`_parse_tool_arguments`) tolerantly decodes model JSON. +- **Usage accounting** (`_acc_chat_usage` / `_acc_responses_usage`) tracks + input/output/reasoning tokens across turns. +- **Document text caching** (`_collect_doc_text`) remembers the text of every chunk + the agent saw, so the final answer's document ids can be rehydrated with content. +- **Document extraction** (`_extract_documents`) parses the final `…` blocks into ranked results. + +--- + +## 5. The system prompt (`prompts.py`) + +`get_retrieval_subagent_prompt(query, num_output_docs)` frames the model as a +**retrieval subagent** (it finds documents, it does *not* answer the question). +It instructs the model to: + +- decompose the query into distinct information needs, +- plan several **non-overlapping** search strategies and issue them **in parallel**, +- after each round, reflect: *what do I know / what to search next / what to prune / + do I have enough?*, +- prune proactively as the token budget approaches its limit, +- output only the ranked `` blocks (most to least relevant). + +When the soft token budget is crossed, +`get_retrieval_subagent_budget_exhausted_message()` is injected as a user turn, +forcing a decision: **prune chunks and continue**, or **conclude**. + +--- + +## 6. The four tools (`tools.py`) + +The agent is given exactly four tools. Each builds a *logical* request and +delegates to the retrieval layer — **no SQL or physical field names** live here. + +| Tool | Schema name | What it does | +|---|---|---| +| `SearchCorpusTool` | `search_corpus` | Hybrid/vector/full-text search; optional rerank; returns the relevant section of each hit. | +| `GrepCorpusTool` | `grep_corpus` | Fetches a full-text candidate pool, then applies a client-side **regex** filter. | +| `ReadDocumentTool` | `read_document` | Reconstructs a full document from its chunks via the configured resolver. | +| `PruneChunksTool` | `prune_chunks` | Records chunk ids whose content should be dropped from context to reclaim tokens. | + +Supporting types: `ToolSchema` (provider-agnostic → OpenAI / Harmony formats), +`ToolSet` (named collection + `build()` factory), `MultiToolUseTool` (wraps a +parallel tool-call bundle), and `ToolCallMetadata` (per-call telemetry such as +returned chunk ids). + +**`ToolSet.build()`** is the wiring point: pass either a pre-built `retriever` +(custom schema) *or* the `cosmos_database` + `container` + `openai_client` trio, +in which case the **default chunked-corpus retriever** is constructed +automatically. It also injects the schema's `agent_field_summary()` into the +search/grep tool descriptions so the model knows which fields it can target. + +--- + +## 7. Budgets, turns, and pruning + +The loop is bounded on three axes so it terminates and stays within context: + +- **`max_turns`** — hard cap on model round-trips (`CHAT_MAX_TURNS`, default 20). +- **`threshold_budget`** (soft) — when accumulated tokens cross it, the over-budget + message is injected, steering the model to prune or conclude. +- **`token_budget`** (hard) — the ceiling the agent must stay under. + +`PruneChunksTool` + the token counter (tiktoken) let the agent trade already-seen, +low-value chunks for fresh searches without blowing the context window. + +--- + +## 8. Reranking (optional, `rerank.py`) + +If configured, a `Reranker` re-scores `search_corpus` / `read_document` results +before they're shown to the model: + +- **`BasetenReranker`** — a hosted reranker endpoint (if `BASETEN_*` set), else +- **`VLLMReranker`** — a local vLLM reranker (if `VLLM_RERANKER_URL` set), else +- **None** — results are returned in the retrieval layer's native order. + +--- + +## 9. Configuration (`config.py`) + +`RetrieverSettings` (pydantic-settings; env vars + `.env`) is the single source of +truth. Highlights: + +- **Corpus targeting** — `ACCOUNT_URI`, `COSMOS_DATABASE`, + `COSMOS_CORPUS_CONTAINER`, plus a `corpus_registry.json` that maps a container + 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. +- **Embeddings** — per-corpus base URL / model / query instruction. +- **Auth** — Cosmos uses `AzureCliCredential` by default (opt into the broader + `DefaultAzureCredential` chain with `COSMOS_USE_DEFAULT_CREDENTIAL=1`); secrets + are read from env, never written to files. + +--- + +## 10. The response + +`POST /search` returns the agent's curated set: + +```json +{ + "query": "…", + "num_turns": 6, + "elapsed_s": 38.2, + "documents": [ + {"id": "doc_123", "text": "…", "justification": "why relevant", "rank": 0} + ] +} +``` + +For the `/responses` backend, a per-query **trajectory** (the search queries +issued, per-turn tool calls, and the final document set) is also captured on the +`RetrievalResult`, which is invaluable for debugging and evaluation. + +--- + +## 11. Concurrency & safety summary + +| Concern | Mechanism | +|---|---| +| Warm clients across requests | `_RetrieverPool` caches one retriever per corpus | +| Sync clients on an async server | `anyio.to_thread.run_sync` | +| 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) | + +--- + +## 12. File map + +| File | Role | +|---|---| +| `server.py` | FastAPI service, `/health`, `/search`, `_RetrieverPool` | +| `retriever.py` | `CosmosRetriever` agent façade + `RetrievalResult` | +| `inference/agent_loop.py` | `run_chat_search` / `run_responses_search` agent loops | +| `prompts.py` | System prompt + budget-exhausted message | +| `tools.py` | The four tools, `ToolSchema`, `ToolSet` | +| `rerank.py` | Optional Baseten / vLLM rerankers | +| `config.py` | `RetrieverSettings`, `CorpusConfig`, corpus registry | +| `cosmos_retriever/retrieval/` | The schema-decoupled retrieval layer (see [RETRIEVAL_SYSTEM.md](RETRIEVAL_SYSTEM.md)) | diff --git a/cosmos-retriever/docs/RETRIEVAL_SYSTEM.md b/cosmos-retriever/docs/RETRIEVAL_SYSTEM.md new file mode 100644 index 0000000..9432b65 --- /dev/null +++ b/cosmos-retriever/docs/RETRIEVAL_SYSTEM.md @@ -0,0 +1,297 @@ +# The Retrieval System + +This document explains how the **schema-decoupled retrieval layer** +(`cosmos_retriever.retrieval`) works: the layer that turns a *logical* retrieval +request ("search the corpus for X") into safe, parameterised Azure Cosmos DB for +NoSQL queries, executes them, and returns normalised results — **without any +agent-facing code ever touching Cosmos SQL or physical property paths**. + +If you only remember one thing: the agent tools speak in *logical fields* +(`text`, `embedding`, `docid`), and a single **`CorpusSchema`** maps those to the +*physical* Cosmos paths (`/text`, `/embedding`, `/docid`) for a given container. +Supporting a brand-new corpus shape is a matter of constructing a schema — no +edits to the tools, planner, or compiler. + +--- + +## 1. Design goals + +| Goal | How it's achieved | +|---|---| +| **No hardcoded paths/SQL in tools** | Tools build logical request models; a `CorpusRetriever` façade owns everything physical. | +| **Portability across corpus shapes** | Three separated concerns — *logical schema*, *physical capabilities*, *partition policy* — are supplied per corpus. | +| **Fail loudly, never silently degrade** | Typed errors (`errors.py`); the planner refuses to "try something cheaper" when an index is missing. | +| **Injection-safe queries** | Every value is a bound `@param`; every property path is validated/rendered by `CosmosPath`. | +| **Resilience under load** | The executor bounds concurrency, retries transient Cosmos errors, and logs slow queries. | + +--- + +## 2. The three inputs that describe a corpus + +Everything the layer does is driven by three objects you provide once per corpus. + +### 2.1 `CorpusSchema` — logical → physical mapping (`schema.py`) +The heart of the system. It says *where each logical field lives*: + +```python +CorpusSchema( + item_id_path="/id", # unique id of a Cosmos item (a chunk) + text_paths=["/text"], # one or more searchable text fields + primary_text_path="/text", # the default text field + vector_fields=[VectorFieldConfig(path="/embedding", dimensions=1536)], + document_id_path="/docid", # parent-document id (None => item *is* the document) + chunk_id_path="/id", + chunk_order_path="/chunk_idx", # used to re-assemble a document in order + partition_key_paths=["/docid"], + metadata_paths={...}, # optional extra projected fields +) +``` + +Key behaviours: +- **Named fields.** `text_field_map()` / `vector_field_map()` expose each field by + its last path segment (`text`, `embedding`), so the agent can pick a field *by + name* without knowing the path. Collisions fall back to the full path string. +- **`agent_field_summary()`** renders a human-readable list of queryable fields + (with optional descriptions) that is injected into the tool description shown to + the model. +- **`resolve_text_fields()` / `resolve_vector_config()`** turn requested field + *names* into physical `CosmosPath`s, raising `UnknownField` if a name is bogus. +- **`ChunkIdentityCodec`** (attached as `identity_codec`) converts a returned + chunk id into its parent document id. `LegacyDunderCodec` implements the + `"__" → ""` convention. +- **Validation.** A pydantic `model_validator` enforces invariants (e.g. + `primary_text_path` must be one of `text_paths`, vector dims > 0), raising + `InvalidCorpusSchema`. + +### 2.2 `RetrievalCapabilities` — what the container can *efficiently* do (`capabilities.py`) +The schema says what fields *exist*; capabilities say what the container is +*indexed* to do well: + +```python +RetrievalCapabilities( + vector_fields=[VectorCapability(path="/embedding", dimensions=1536, + support=SupportLevel.INDEXED)], + full_text_paths=["/text"], + native_hybrid_supported=True, # ORDER BY RANK RRF(...) available + full_text_supported=True, + vector_supported=True, + efficient_document_lookup_supported=True, +) +``` + +`SupportLevel` (`INDEXED` / `SCAN` / `UNSUPPORTED` / `UNKNOWN`) lets the planner +distinguish "indexed and fast" from "possible but a scan". The planner uses this +to choose a strategy and to **refuse** operations that would silently be slow. + +### 2.3 `PartitionQueryPolicy` — cross-partition guard-rails (`models.py`) +Controls what the layer is *allowed* to do when a partition key isn't supplied: +`allow_cross_partition_search`, `allow_cross_partition_document_read`, +`allow_bounded_scan`, `maximum_partitions`, etc. This keeps accidental +fan-out/full-scans behind an explicit opt-in. + +--- + +## 3. The request/response models (`models.py`) + +These carry **no** SQL or physical knowledge — they are the logical vocabulary +the tools speak: + +- **`SearchRequest`** — `query`, optional `query_vector`, `limit`, `text_fields`, + `vector_field`, `mode` (`auto|hybrid|vector|text`), `filters`, + `ignored_item_ids` (already-seen chunks to exclude), `partition_key`. +- **`GrepRequest`** — `pattern`, `candidate_limit`, `result_limit`, `text_field`. +- **`ReadDocumentRequest`** — `document_id`/`item_id`, `max_chunks`, `partition_key`. +- **Filters** — `EqualsFilter` / `RangeFilter` / `InFilter` (a discriminated + union), addressing *logical* field names. +- **`RetrievedItem`** — the normalised hit: `item_id`, `document_id`, `chunk_id`, + `chunk_order`, `text` (the display text), `text_fields` (every projected text + field keyed by name), `metadata`, `retrieval_strategy`, `retrieval_channels`, + `rank`. +- **`NormalizedDocument`** — a reconstructed document: ordered `chunk_texts` + + `chunk_ids`, with an `assembled` property that concatenates them. +- **`CompiledCosmosQuery`** — the compiler's output: `sql`, bound `parameters`, + `partition_key`, `enable_cross_partition_query`, `projected_aliases`. + +--- + +## 4. The pipeline + +A single `CorpusRetriever.search()` call flows through six stages: + +```mermaid +flowchart LR + A[SearchRequest
logical] --> B[RetrievalPlanner
pick strategy] + B --> C[SearchStrategy
resolve fields] + C --> D[CosmosQueryCompiler
build safe SQL] + D --> E[CosmosExecutor
run + retry] + E --> F[normalize_rows
rows -> items] + F --> G[list RetrievedItem] +``` + +### 4.1 `CorpusRetriever` — the façade (`retriever.py`) +The one object the agent tools depend on. It wires together the schema, +capabilities, planner, compiler, executor, strategies, and the document resolver, +and exposes exactly three methods: + +- **`search(SearchRequest) -> list[RetrievedItem]`** — validates any explicitly + requested field names up front (so an unknown field raises rather than silently + falling back), asks the planner for a strategy, lazily embeds the query if the + strategy needs a vector and none was supplied, then executes. +- **`grep_candidates(GrepRequest) -> list[RetrievedItem]`** — full-text candidate + fetch used as the pool for client-side regex filtering. +- **`read_document(ReadDocumentRequest) -> NormalizedDocument`** — reconstruct a + full document via the configured resolver. + +### 4.2 `RetrievalPlanner` — strategy selection (`planner.py`) +Turns *request + schema + capabilities + policy* into a concrete strategy. It +never equates "the query didn't throw" with "the operation is indexed": + +- `_vector_ok()` — vector search is viable only if the field exists, is + `vector_supported`, its capability is `INDEXED`, **and stored dimensions match + the schema's** (guards embedding-profile mismatches). +- `_fts_ok()` — full-text is viable only if every requested text path is a + declared `full_text_path`. +- **`plan_search()`** honours `mode`: + - `vector` / `text` → force that channel (raise `UnsupportedRetrievalCapability` + if unavailable); + - `hybrid` → `NativeHybridStrategy` if native RRF is supported, else + `ClientSideFusionStrategy`; + - `auto` → native hybrid → client-side fusion → vector-only → text-only → + bounded scan (if policy allows) → else raise. +- **`plan_grep()`** → `FullTextGrepCandidateStrategy` when full-text is available. + +### 4.3 Strategies — how each search actually runs (`strategies.py`) +All strategies share a `RetrievalContext` (schema, compiler, executor, +capabilities, policy) and return `list[RetrievedItem]`. + +| Strategy | What it emits | Notes | +|---|---|---| +| `NativeHybridStrategy` | `ORDER BY RANK RRF(VectorDistance(...), FullTextScore(...))` | Server-side Reciprocal Rank Fusion; multi-field FTS supported. | +| `VectorSearchStrategy` | `ORDER BY RANK VectorDistance(...)` | Pure semantic. | +| `FullTextSearchStrategy` | `ORDER BY RANK FullTextScore(...)` (RRF if multi-field) | Pure keyword/BM25. | +| `ClientSideFusionStrategy` | Runs vector + FTS, fuses with RRF (`k=60`) in Python | Fallback when the container lacks native RRF. | +| `BoundedScanStrategy` | Filter-only query, no ranking | Opt-in via `allow_bounded_scan`; last resort. | +| `FullTextGrepCandidateStrategy` | FTS candidate pool for grep | Feeds client-side regex. | + +Cross-partition safety: `_resolve_cross_partition()` decides whether a query may +fan out, raising `CrossPartitionQueryDisabled` when policy forbids it. + +### 4.4 `CosmosQueryCompiler` — safe SQL generation (`compiler.py`) +Builds the actual Cosmos SQL. Everything hostile-to-inject is parameterised or +path-validated: + +- **`_ParamBag`** collects bound `@p0, @p1, …` parameters; no value is ever string-interpolated. +- **`projection()`** builds the `SELECT`, projecting logical columns (`item_id`, + `text`, `document_id`, `chunk_order`, …), plus every text field as `txt_` and + every metadata field as `md_`. It returns an **alias legend** so the + normaliser can map columns back to field names. +- **`_where()`** compiles filters and appends the `ignored_item_ids` exclusion + (`NOT ARRAY_CONTAINS(@ids, c.)`). +- **`compile_hybrid` / `compile_vector` / `compile_full_text` / `compile_structured` + / `compile_document_read`** each return a `CompiledCosmosQuery`. +- Physical paths are rendered through `CosmosPath.render()` (see §5), never raw strings. + +### 4.5 `CosmosExecutor` — the single DB chokepoint (`executor.py`) +Every query in the system goes through `_query_items()`, which provides: + +- **Bounded concurrency** — a process-wide `BoundedSemaphore` + (`COSMOS_QUERY_MAX_CONCURRENCY`, default 8) caps simultaneous Cosmos calls. +- **Retries** — `@tenacity.retry` (5 attempts, exponential backoff 4–15 s) on + transient Cosmos statuses (408/429/449/500/502/503/504). +- **Partition routing** — passes `partition_key` when known, else + `enable_cross_partition_query=True`. +- **Slow-query logging** — warns when a query exceeds ~4.5 s. + +### 4.6 `normalize_rows` — raw rows → `RetrievedItem` (`normalization.py`) +Maps the projected aliases back into structured items: + +- `md_*` columns → `metadata`. +- `txt_*` columns → `text_fields` (keyed by the real field name via the alias legend). +- **Display text** (`_display_text`) returns the field(s) the caller actually + queried, and *also* appends the configured primary field if it wasn't already + included — so a keyword hit on a title still shows the body text too. + +--- + +## 5. Path safety (`paths.py`) + +`CosmosPath` parses a `"/a/b"` path into validated segments and rejects anything +unsafe (bad characters, injection attempts), raising `UnsafeCosmosPath`. It +renders to `c.a.b` against the query alias. Because **every** physical path in a +compiled query flows through `CosmosPath`, a malicious or malformed schema path +cannot produce injectable SQL. + +## 6. Text processing (`expressions.py`) + +`tokenize_for_fts()` lower-cases, strips a stopword list, de-duplicates, and caps +at 30 terms; `fts_literal_args()` escapes each term into the quoted argument list +that `FullTextScore(path, "t1", "t2", …)` expects. + +## 7. Document reconstruction (`document_resolvers.py`) + +`read_document` reassembles a full document from its chunks. The factory +`build_document_resolver()` picks the right resolver from the schema: + +| Resolver | When | How | +|---|---|---| +| `ItemIsDocumentResolver` | `document_id_path is None` (item *is* the doc) | Single-item lookup. | +| `ChunkedDocumentResolver` | partition key **is** the document id | Partition-scoped read, ordered by `chunk_order`. | +| `CrossPartitionChunkedDocumentResolver` | docid ≠ partition key | Cross-partition read (requires policy permission). | + +All derive the parent id via the schema's identity codec and sort chunks by +`chunk_order` before assembly (default cap `DEFAULT_MAX_CHUNKS = 300`). + +## 8. Typed errors (`errors.py`) + +`RetrievalError` subclasses make every failure mode explicit instead of degrading: +`InvalidCorpusSchema`, `UnsafeCosmosPath`, `UnsupportedRetrievalCapability`, +`UnknownField`, `EmbeddingProfileMismatch`, `CrossPartitionQueryDisabled`, +`UnboundedScanRejected`, `DocumentResolutionUnsupported`, `QueryCompilationError`, +`IndexNotReady`, `MissingPartitionKey`. + +## 9. The default profile (`legacy.py`) + +The MCP server must work out-of-the-box against the conventional chunked corpus +(`/id`, `/text`, `/embedding`, `/docid`, `/chunk_idx`, `__` +chunk ids, native RRF hybrid). `legacy.py` packages exactly that as a named +profile: + +- `build_legacy_schema()` → the standard `CorpusSchema` (+ `LegacyDunderCodec`), +- `legacy_capabilities_for()` → capabilities with native hybrid enabled, +- `build_legacy_retriever()` → a ready `CorpusRetriever`. + +`ToolSet.build()` calls `build_legacy_retriever()` when no explicit retriever is +supplied, so existing deployments keep working while custom corpora can pass their +own schema. + +--- + +## 10. Adding a new corpus (recipe) + +1. Write a `CorpusSchema` mapping your logical fields to physical paths. +2. Attach an identity codec if your chunk ids encode the parent doc id. +3. Declare a `RetrievalCapabilities` describing indexed vector/full-text support. +4. Build a `CorpusRetriever(container, schema, capabilities, query_embedder)`. +5. Pass it to `ToolSet.build(retriever=...)`. + +No changes to tools, planner, compiler, executor, or resolvers are required. + +## 11. Module map + +| File | Responsibility | +|---|---| +| `schema.py` | `CorpusSchema`, `VectorFieldConfig`, identity codec, field resolution | +| `capabilities.py` | `RetrievalCapabilities`, `VectorCapability`, `SupportLevel` | +| `models.py` | Request/response models, filters, `PartitionQueryPolicy`, `CompiledCosmosQuery` | +| `paths.py` | `CosmosPath` safe parse/render | +| `expressions.py` | FTS tokenisation + literal escaping | +| `planner.py` | Strategy selection from capabilities | +| `strategies.py` | Hybrid / vector / full-text / fusion / scan / grep execution | +| `compiler.py` | Logical plan → parameterised Cosmos SQL | +| `executor.py` | Concurrency, retries, slow-query logging | +| `normalization.py` | Raw rows → `RetrievedItem` | +| `document_resolvers.py` | Full-document reconstruction | +| `errors.py` | Typed failure modes | +| `retriever.py` | `CorpusRetriever` façade | +| `legacy.py` | Default chunked-corpus profile | +| `__init__.py` | Public API surface | diff --git a/cosmos-retriever/pyproject.toml b/cosmos-retriever/pyproject.toml new file mode 100644 index 0000000..7e15e3a --- /dev/null +++ b/cosmos-retriever/pyproject.toml @@ -0,0 +1,74 @@ +[project] +name = "cosmos-retriever" +version = "0.1.0" +description = "Multi-turn Cosmos DB search agent (driven by any OpenAI-compatible model) as a Python library + CLI, designed to be invoked by the Azure Cosmos DB MCP Toolkit's `agentic_search` tool." +readme = "README.md" +requires-python = ">=3.11" +license = { text = "Apache-2.0" } +authors = [{ name = "Cosmos Retriever Contributors" }] +keywords = ["retrieval", "rag", "cosmos-db", "vllm", "agent"] +classifiers = [ + "Development Status :: 4 - Beta", + "Intended Audience :: Developers", + "License :: OSI Approved :: Apache Software License", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Topic :: Scientific/Engineering :: Artificial Intelligence", +] +dependencies = [ + "anyio>=4.0,<5", + "azure-cosmos>=4.7,<5", + "azure-identity>=1.17,<2", + "fastapi>=0.110,<1", + "httpx>=0.27,<1", + "json-repair>=0.20,<1", + "openai>=1.40,<2", + "openai-harmony>=0.0.8,<1", + "pydantic>=2.7,<3", + "pydantic-settings>=2.4,<3", + "structlog>=24,<26", + "tenacity>=8.3,<10", + "tiktoken>=0.7,<1", + "uvicorn>=0.30,<1", +] + +[project.optional-dependencies] +baseten = ["baseten-performance-client>=0.4,<1"] +dev = [ + "mypy>=1.10,<2", + "ruff>=0.6,<1", +] + +[project.scripts] +cosmos-retriever = "cosmos_retriever.__main__:main" + +[project.urls] +Homepage = "https://github.com/your-org/cosmos-retriever" + +[build-system] +requires = ["hatchling>=1.24"] +build-backend = "hatchling.build" + +[tool.hatch.build.targets.wheel] +packages = ["src/cosmos_retriever"] + +[tool.ruff] +line-length = 100 +target-version = "py311" + +[tool.ruff.lint] +select = ["E", "F", "I", "B", "UP", "SIM", "N"] +ignore = ["E501"] # line length handled by formatter + +[tool.ruff.lint.per-file-ignores] +# Retrieval error names are part of the public error model and intentionally +# do not use an "Error" suffix. +"src/cosmos_retriever/retrieval/errors.py" = ["N818"] + +[tool.mypy] +python_version = "3.11" +strict = false +warn_unused_ignores = true +warn_redundant_casts = true +ignore_missing_imports = true +files = ["src/cosmos_retriever"] diff --git a/cosmos-retriever/run-retriever.ps1 b/cosmos-retriever/run-retriever.ps1 new file mode 100644 index 0000000..4a2db9b --- /dev/null +++ b/cosmos-retriever/run-retriever.ps1 @@ -0,0 +1,11 @@ +# Loads cosmos-retriever/.env is handled automatically by python-dotenv, +# so this just activates the venv interpreter and starts the FastAPI service. +# Usage: .\run-retriever.ps1 +$ErrorActionPreference = "Stop" +$here = Split-Path -Parent $MyInvocation.MyCommand.Path +$py = Join-Path $here ".venv\Scripts\python.exe" +if (-not (Test-Path $py)) { + throw "venv not found at $py. Run: uv venv --python 3.11 .venv; uv pip install --python .venv\Scripts\python.exe -e `".[dev]`"" +} +Write-Host "Starting cosmos-retriever FastAPI service (reads .env)..." -ForegroundColor Cyan +& $py -m cosmos_retriever serve diff --git a/cosmos-retriever/src/cosmos_retriever/__init__.py b/cosmos-retriever/src/cosmos_retriever/__init__.py new file mode 100644 index 0000000..35f3221 --- /dev/null +++ b/cosmos-retriever/src/cosmos_retriever/__init__.py @@ -0,0 +1,11 @@ + +from __future__ import annotations + +from cosmos_retriever.retriever import ( + CosmosRetriever, + RetrievalResult, + RetrievedDocument, +) + +__all__ = ["CosmosRetriever", "RetrievalResult", "RetrievedDocument"] +__version__ = "0.1.0" diff --git a/cosmos-retriever/src/cosmos_retriever/__main__.py b/cosmos-retriever/src/cosmos_retriever/__main__.py new file mode 100644 index 0000000..ace95d7 --- /dev/null +++ b/cosmos-retriever/src/cosmos_retriever/__main__.py @@ -0,0 +1,101 @@ + +from __future__ import annotations + +import argparse +import json +import sys +from dataclasses import asdict + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + prog="cosmos-retriever", + description="Run the multi-turn Cosmos retrieval agent and emit JSON.", + ) + sub = parser.add_subparsers(dest="cmd", required=True) + + search = sub.add_parser("search", help="Run one search end-to-end.") + search.add_argument("--query", required=True) + search.add_argument("--max-documents", type=int, default=20) + search.add_argument( + "--database", + default=None, + help="Override Cosmos database name (else COSMOS_DATABASE env var).", + ) + search.add_argument( + "--container", + default=None, + help="Override Cosmos corpus container name (else COSMOS_CORPUS_CONTAINER env var).", + ) + + serve = sub.add_parser( + "serve", + help="Run the FastAPI HTTP service the MCP Toolkit calls into.", + ) + serve.add_argument( + "--host", + default=None, + help="Bind address (else HOST env var, default 0.0.0.0).", + ) + serve.add_argument( + "--port", + type=int, + default=None, + help="Bind port (else PORT env var, default 9000).", + ) + return parser + + +def _cmd_search(args: argparse.Namespace) -> int: + from cosmos_retriever.config import get_settings + from cosmos_retriever.retriever import CosmosRetriever + + settings = get_settings() + if args.database: + settings.cosmos_database = args.database + + retriever = CosmosRetriever(settings=settings, corpus_name=args.container) + result = retriever.search(args.query, max_documents=args.max_documents) + json.dump(asdict(result), sys.stdout, default=str, ensure_ascii=False) + sys.stdout.write("\n") + sys.stdout.flush() + return 0 + + +def _cmd_serve(args: argparse.Namespace) -> int: + import uvicorn + + from cosmos_retriever.config import get_settings + from cosmos_retriever.server import create_app + + settings = get_settings() + host = args.host or settings.host + port = args.port or settings.port + app = create_app(settings) + uvicorn.run(app, host=host, port=port, log_level=settings.log_level.lower()) + return 0 + + +def main(argv: list[str] | None = None) -> int: + parser = _build_parser() + args = parser.parse_args(argv) + try: + if args.cmd == "search": + return _cmd_search(args) + if args.cmd == "serve": + return _cmd_serve(args) + except Exception as exc: + json.dump( + {"error": str(exc), "type": type(exc).__name__}, + sys.stdout, + ensure_ascii=False, + ) + sys.stdout.write("\n") + sys.stdout.flush() + return 1 + parser.error(f"Unknown command: {args.cmd}") + return 2 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/cosmos-retriever/src/cosmos_retriever/config.py b/cosmos-retriever/src/cosmos_retriever/config.py new file mode 100644 index 0000000..15d7aaf --- /dev/null +++ b/cosmos-retriever/src/cosmos_retriever/config.py @@ -0,0 +1,350 @@ + +from __future__ import annotations + +import json +import logging +import os +import sys +from dataclasses import dataclass +from functools import lru_cache +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import structlog +from azure.cosmos import CosmosClient, DatabaseProxy +from azure.identity import AzureCliCredential, DefaultAzureCredential +from dotenv import load_dotenv +from openai import OpenAI +from pydantic import Field, SecretStr, field_validator +from pydantic_settings import BaseSettings, SettingsConfigDict + +if TYPE_CHECKING: + from baseten_performance_client import PerformanceClient + +REPO_ROOT = Path(__file__).resolve().parents[2] +DEFAULT_ENV_FILES = (str(REPO_ROOT / ".env.local"), str(REPO_ROOT / ".env")) + +for _env_path in DEFAULT_ENV_FILES: + load_dotenv(_env_path, override=False) + + +def init_logging( + app_level: int = logging.INFO, + *, + lib_level: int = logging.WARNING, + colors: bool = True, +) -> None: + + logging.basicConfig(level=lib_level, format="%(message)s", stream=sys.stderr, force=True) + structlog.configure_once( + processors=[ + structlog.processors.TimeStamper(fmt="iso", utc=True), + structlog.processors.add_log_level, + structlog.processors.StackInfoRenderer(), + structlog.processors.format_exc_info, + structlog.dev.ConsoleRenderer(colors=colors), + ], + wrapper_class=structlog.make_filtering_bound_logger(app_level), + cache_logger_on_first_use=True, + logger_factory=structlog.PrintLoggerFactory(file=sys.stderr), + ) + + +@dataclass(frozen=True) +class CorpusConfig: + + container: str + account_uri: str + database: str + embed_base_url: str | None + + embed_api_key: SecretStr | None + embed_model: str + embed_query_instruction: str | None = None + + cosmos_key: SecretStr | None = None + + +class RetrieverSettings(BaseSettings): + + model_config = SettingsConfigDict( + env_file=DEFAULT_ENV_FILES, + env_file_encoding="utf-8", + extra="ignore", + case_sensitive=False, + ) + + inference_backend: str = Field( + default="openai_responses", + description='Inference backend: "openai_responses", "openai_chat", or "anthropic_messages".', + ) + + @field_validator("inference_backend") + @classmethod + def _validate_inference_backend(cls, v: str) -> str: + normalized = (v or "").strip().lower() + allowed = {"openai_chat", "openai_responses", "anthropic_messages"} + if normalized not in allowed: + raise ValueError( + f"INFERENCE_BACKEND must be one of {sorted(allowed)}, got {v!r}." + ) + return normalized + + chat_base_url: str | None = Field( + default=None, + description="Base URL of an OpenAI-compatible chat-completions endpoint.", + ) + chat_api_key: SecretStr | None = None + chat_model: str | None = Field( + default=None, description="Chat model / Foundry deployment name." + ) + chat_api_version: str | None = Field( + default=None, + description="Set for Azure OpenAI-style endpoints (uses the AzureOpenAI client).", + ) + chat_temperature: float = Field(default=0.7, ge=0.0, le=2.0) + chat_max_tokens: int = Field(default=4096, ge=256) + chat_max_turns: int = Field(default=20, ge=1, le=200) + chat_reasoning_effort: str | None = Field( + default=None, + description='Reasoning effort for reasoning models on the responses API (e.g. "low", "medium", "high").', + ) + anthropic_version: str = Field( + default="2023-06-01", + description="anthropic-version header for the anthropic_messages backend.", + ) + anthropic_auth_header: str = Field( + default="x-api-key", + description='Auth header name for the anthropic_messages endpoint (e.g. "x-api-key" or "api-key").', + ) + + account_uri: str = Field(description="Cosmos DB account URI (default corpus).") + cosmos_database: str + cosmos_corpus_container: str + cosmos_key: SecretStr | None = None + + openai_api_key: SecretStr | None = None + openai_embedding_model: str | None = None + embed_endpoint: str | None = Field( + default=None, + description=( + "Embedding endpoint base URL. Leave unset to use plain OpenAI " + "(api.openai.com). For Azure pass https://.../openai/v1; " + "for a local server pass http://host:port/v1." + ), + ) + embed_query_instruction: str | None = None + + corpus_registry: str | None = Field( + default=None, + description="JSON string mapping container name -> CorpusConfig overrides.", + ) + corpus_registry_file: str | None = Field( + default=None, + description="Path to a JSON file holding the corpus registry.", + ) + + baseten_api_key: SecretStr | None = None + baseten_model_url: str | None = None + vllm_reranker_url: str | None = None + + cosmos_retriever_max_turns: int = Field(default=35, ge=1, le=200, alias="COSMOS_RETRIEVER_MAX_TURNS") + cosmos_retriever_threshold_budget: int = Field( + default=16384, ge=1024, alias="COSMOS_RETRIEVER_THRESHOLD_BUDGET" + ) + cosmos_retriever_token_budget: int = Field( + default=32268, ge=4096, alias="COSMOS_RETRIEVER_TOKEN_BUDGET" + ) + cosmos_retriever_search_display_limit: int = Field(default=15, ge=1, le=50) + + host: str = Field(default="0.0.0.0") + port: int = Field(default=9000, ge=1, le=65535) + log_level: str = Field(default="info") + + def _load_registry(self) -> dict[str, dict[str, Any]]: + + if self.corpus_registry_file: + path = Path(self.corpus_registry_file) + if not path.is_file(): + raise FileNotFoundError(f"CORPUS_REGISTRY_FILE points at missing file: {path}") + raw = path.read_text(encoding="utf-8") + elif self.corpus_registry: + raw = self.corpus_registry + else: + return {} + + try: + data = json.loads(raw) + except json.JSONDecodeError as exc: + raise ValueError(f"corpus_registry is not valid JSON: {exc}") from exc + + if not isinstance(data, dict): + raise ValueError("corpus_registry must be a JSON object {container_name: {...}}") + return data + + def resolve_corpus(self, container: str | None = None) -> CorpusConfig: + + registry = self._load_registry() + target = container or self.cosmos_corpus_container + entry = registry.get(target) + + def _resolve_default_embed() -> tuple[str | None, SecretStr | None, str | None]: + return self.embed_endpoint, self.openai_api_key, self.openai_embedding_model + + if entry is None: + base, key, model = _resolve_default_embed() + return CorpusConfig( + container=target, + account_uri=self.account_uri, + database=self.cosmos_database, + embed_base_url=base, + embed_api_key=key, + embed_model=model, + embed_query_instruction=self.embed_query_instruction, + cosmos_key=self.cosmos_key, + ) + + api_key_env = entry.get("embed_api_key_env") + api_key_value: SecretStr | None = None + if api_key_env: + raw_key = os.environ.get(api_key_env) + if raw_key: + api_key_value = SecretStr(raw_key) + + cosmos_key_env = entry.get("cosmos_key_env") + cosmos_key_value: SecretStr | None = self.cosmos_key + if cosmos_key_env: + raw_ck = os.environ.get(cosmos_key_env) + if raw_ck: + cosmos_key_value = SecretStr(raw_ck) + + return CorpusConfig( + container=target, + account_uri=entry.get("account_uri") or self.account_uri, + database=entry.get("database") or self.cosmos_database, + embed_base_url=entry.get("embed_base_url"), + embed_api_key=api_key_value, + embed_model=entry.get("embed_model") or self.openai_embedding_model, + embed_query_instruction=entry.get("embed_query_instruction"), + cosmos_key=cosmos_key_value, + ) + + def _cosmos_credential(self): + + if os.environ.get("COSMOS_USE_DEFAULT_CREDENTIAL", "").lower() in {"1", "true", "yes"}: + return DefaultAzureCredential() + return AzureCliCredential() + + def build_cosmos_database(self, corpus: CorpusConfig) -> DatabaseProxy: + + if corpus.cosmos_key is not None: + client = CosmosClient(corpus.account_uri, credential=corpus.cosmos_key.get_secret_value()) + else: + client = CosmosClient(corpus.account_uri, credential=self._cosmos_credential()) + return client.get_database_client(corpus.database) + + def build_openai_client(self, corpus: CorpusConfig) -> OpenAI: + + kwargs: dict[str, Any] = {} + if corpus.embed_base_url: + kwargs["base_url"] = corpus.embed_base_url + kwargs["api_key"] = ( + corpus.embed_api_key.get_secret_value() if corpus.embed_api_key is not None else "EMPTY" + ) + return OpenAI(**kwargs) + + @property + def use_chat_backend(self) -> bool: + + return self.inference_backend.strip().lower() == "openai_chat" + + @property + def use_responses_backend(self) -> bool: + + return self.inference_backend.strip().lower() == "openai_responses" + + @property + def use_anthropic_backend(self) -> bool: + + return self.inference_backend.strip().lower() == "anthropic_messages" + + @property + def use_generic_llm_backend(self) -> bool: + + return self.use_chat_backend or self.use_responses_backend + + def build_chat_client(self) -> OpenAI: + + if not self.chat_base_url: + raise ValueError( + "CHAT_BASE_URL must be set when INFERENCE_BACKEND=openai_chat." + ) + if not self.chat_model: + raise ValueError( + "CHAT_MODEL (the deployment / model name) must be set when " + "INFERENCE_BACKEND=openai_chat." + ) + api_key = ( + self.chat_api_key.get_secret_value() if self.chat_api_key is not None else "EMPTY" + ) + if self.chat_api_version: + from openai import AzureOpenAI + + return AzureOpenAI( + azure_endpoint=self.chat_base_url, + api_key=api_key, + api_version=self.chat_api_version, + ) + return OpenAI(base_url=self.chat_base_url, api_key=api_key) + + def get_cosmos_client(self) -> CosmosClient: + corpus = self.resolve_corpus() + if corpus.cosmos_key is not None: + return CosmosClient(corpus.account_uri, credential=corpus.cosmos_key.get_secret_value()) + return CosmosClient(corpus.account_uri, credential=self._cosmos_credential()) + + def get_cosmos_database(self) -> DatabaseProxy: + return self.build_cosmos_database(self.resolve_corpus()) + + def get_openai_client(self) -> OpenAI: + return self.build_openai_client(self.resolve_corpus()) + + def get_baseten_client(self) -> PerformanceClient: + + if self.baseten_api_key is None or not self.baseten_model_url: + raise ValueError( + "BASETEN_API_KEY and BASETEN_MODEL_URL must both be set to use Baseten reranking." + ) + from baseten_performance_client import PerformanceClient + + return PerformanceClient( + base_url=self.baseten_model_url, + api_key=self.baseten_api_key.get_secret_value(), + ) + + +@lru_cache(maxsize=1) +def get_settings() -> RetrieverSettings: + + settings = RetrieverSettings() + init_logging(app_level=_log_level_to_int(settings.log_level)) + return settings + + +def get_config() -> "RetrieverSettings": + return get_settings() + + +def _log_level_to_int(level: str) -> int: + return getattr(logging, level.upper(), logging.INFO) + + +__all__ = [ + "CorpusConfig", + "DEFAULT_ENV_FILES", + "REPO_ROOT", + "RetrieverSettings", + "get_config", + "get_settings", + "init_logging", +] diff --git a/cosmos-retriever/src/cosmos_retriever/inference/__init__.py b/cosmos-retriever/src/cosmos_retriever/inference/__init__.py new file mode 100644 index 0000000..1206d27 --- /dev/null +++ b/cosmos-retriever/src/cosmos_retriever/inference/__init__.py @@ -0,0 +1,18 @@ + +from __future__ import annotations + +from cosmos_retriever.inference.agent_loop import ( + ChatDocument, + AgentSearchResult, + run_anthropic_search, + run_chat_search, + run_responses_search, +) + +__all__ = [ + "ChatDocument", + "AgentSearchResult", + "run_anthropic_search", + "run_chat_search", + "run_responses_search", +] diff --git a/cosmos-retriever/src/cosmos_retriever/inference/agent_loop.py b/cosmos-retriever/src/cosmos_retriever/inference/agent_loop.py new file mode 100644 index 0000000..8b6dd83 --- /dev/null +++ b/cosmos-retriever/src/cosmos_retriever/inference/agent_loop.py @@ -0,0 +1,520 @@ + +from __future__ import annotations + +import json +import re +from dataclasses import dataclass, field + +import json_repair +import openai +import requests +import structlog + +from cosmos_retriever.prompts import get_retrieval_subagent_prompt +from cosmos_retriever.tools import ToolSet +from cosmos_retriever.utils import ProviderFormat + +logger = structlog.get_logger("cosmos_retriever.inference.agent_loop") + +_CHAT_TOOL_NAMES = ("search_corpus", "grep_corpus", "read_document", "prune_chunks") + +_DOC_RESULT_RE = re.compile(r"#\s*DOCUMENT ID:\s*(?P\S+)(?:\s*\(\d+\s*tokens\))?") + +_FINAL_DOC_RE = re.compile( + r"[^\"'\s>]+)[\"']?\s*>\s*" + r"(?:\s*(?P.*?)\s*\s*)?" + r"", + re.IGNORECASE | re.DOTALL, +) + + +@dataclass +class ChatDocument: + + id: str + text: str = "" + justification: str | None = None + rank: int | None = None + + +@dataclass +class AgentSearchResult: + + documents: list[ChatDocument] + num_turns: int + final_text: str = "" + pool_doc_ids: list[str] = field(default_factory=list) + usage: dict[str, int] = field(default_factory=dict) + trajectory: dict[str, object] = field(default_factory=dict) + metadata: dict[str, str | int | float] = field(default_factory=dict) + + +def _empty_usage() -> dict[str, int]: + return { + "prompt_tokens": 0, + "completion_tokens": 0, + "reasoning_tokens": 0, + "total_tokens": 0, + "llm_calls": 0, + } + + +def _acc_chat_usage(usage: dict[str, int], resp) -> None: + u = getattr(resp, "usage", None) + if u is None: + return + usage["prompt_tokens"] += int(getattr(u, "prompt_tokens", 0) or 0) + usage["completion_tokens"] += int(getattr(u, "completion_tokens", 0) or 0) + usage["total_tokens"] += int(getattr(u, "total_tokens", 0) or 0) + usage["llm_calls"] += 1 + + +def _acc_responses_usage(usage: dict[str, int], resp) -> None: + u = getattr(resp, "usage", None) + if u is None: + return + usage["prompt_tokens"] += int(getattr(u, "input_tokens", 0) or 0) + usage["completion_tokens"] += int(getattr(u, "output_tokens", 0) or 0) + usage["total_tokens"] += int(getattr(u, "total_tokens", 0) or 0) + details = getattr(u, "output_tokens_details", None) + if details is not None: + usage["reasoning_tokens"] += int(getattr(details, "reasoning_tokens", 0) or 0) + usage["llm_calls"] += 1 + + +def _parse_tool_arguments(raw: str | None) -> dict: + + if not raw: + return {} + try: + parsed = json.loads(raw) + except json.JSONDecodeError: + try: + parsed = json_repair.loads(raw) + except Exception: + return {} + return parsed if isinstance(parsed, dict) else {} + + +def _collect_doc_text(observation: str, store: dict[str, str]) -> None: + + matches = list(_DOC_RESULT_RE.finditer(observation)) + for idx, match in enumerate(matches): + chunk_id = match.group("id") + start = match.end() + end = matches[idx + 1].start() if idx + 1 < len(matches) else len(observation) + body = observation[start:end].strip() + if body and not store.get(chunk_id): + store[chunk_id] = body + + +def _extract_documents( + final_text: str, doc_text: dict[str, str], max_documents: int +) -> list[ChatDocument]: + + documents: list[ChatDocument] = [] + seen: set[str] = set() + for match in _FINAL_DOC_RE.finditer(final_text): + doc_id = match.group("id") + if doc_id in seen: + continue + seen.add(doc_id) + justification = match.group("justification") + text = doc_text.get(doc_id) or doc_text.get(doc_id.split("__")[0]) or "" + documents.append( + ChatDocument( + id=doc_id, + text=text, + justification=justification.strip() if justification else None, + rank=len(documents), + ) + ) + if len(documents) >= max_documents: + break + return documents + + +def run_chat_search( + *, + toolset: ToolSet, + client: openai.OpenAI, + model: str, + query: str, + max_documents: int = 20, + max_turns: int = 20, + temperature: float = 0.7, + max_tokens: int = 4096, +) -> AgentSearchResult: + + tool_specs = [ + tool.get_format(ProviderFormat.OPENAI_HARMONY) + for name, tool in toolset.tools.items() + if name in _CHAT_TOOL_NAMES + ] + + messages: list[dict] = [ + {"role": "system", "content": get_retrieval_subagent_prompt(query, num_output_docs=max_documents)}, + { + "role": "user", + "content": ( + "Use the available tools to search the corpus, then return ONLY the " + "ranked blocks (with a ) for the most " + "relevant documents. Do not answer the question yourself." + ), + }, + ] + + doc_text: dict[str, str] = {} + tool_types_used: set[str] = set() + tool_call_count = 0 + final_text = "" + num_turns = 0 + usage = _empty_usage() + + for _ in range(max_turns): + response = client.chat.completions.create( + model=model, + messages=messages, + tools=tool_specs, + tool_choice="auto", + temperature=temperature, + max_tokens=max_tokens, + ) + num_turns += 1 + _acc_chat_usage(usage, response) + message = response.choices[0].message + tool_calls = message.tool_calls or [] + + assistant_entry: dict = {"role": "assistant", "content": message.content or ""} + if tool_calls: + assistant_entry["tool_calls"] = [ + { + "id": tc.id, + "type": "function", + "function": {"name": tc.function.name, "arguments": tc.function.arguments}, + } + for tc in tool_calls + ] + messages.append(assistant_entry) + + if not tool_calls: + final_text = message.content or "" + break + + for tc in tool_calls: + name = tc.function.name + tool_types_used.add(name) + tool_call_count += 1 + args = _parse_tool_arguments(tc.function.arguments) + tool = toolset.get_tool(name) + if tool is None: + output = f"Error: unknown tool '{name}'." + else: + try: + output, _metadata = tool(args) + _collect_doc_text(output, doc_text) + except Exception as exc: + logger.warning("chat_tool_error", tool=name, error=str(exc)) + output = f"Error executing '{name}': {exc}" + messages.append({"role": "tool", "tool_call_id": tc.id, "content": output}) + else: + for entry in reversed(messages): + if entry.get("role") == "assistant" and entry.get("content"): + final_text = entry["content"] + break + + documents = _extract_documents(final_text, doc_text, max_documents) + + logger.info( + "chat_search_complete", + model=model, + num_turns=num_turns, + num_documents=len(documents), + tool_calls=tool_call_count, + ) + + return AgentSearchResult( + documents=documents, + num_turns=num_turns, + final_text=final_text, + usage=usage, + metadata={ + "backend": "openai_chat", + "model": model, + "tool_calls": tool_call_count, + "tool_types_used": ",".join(sorted(tool_types_used)), + }, + ) + + +def run_responses_search( + *, + toolset: ToolSet, + client: openai.OpenAI, + model: str, + query: str, + max_documents: int = 20, + max_turns: int = 20, + max_tokens: int = 4096, + reasoning_effort: str | None = None, +) -> AgentSearchResult: + + tool_specs = [ + tool.get_format(ProviderFormat.OPENAI) + for name, tool in toolset.tools.items() + if name in _CHAT_TOOL_NAMES + ] + + prompt = ( + get_retrieval_subagent_prompt(query, num_output_docs=max_documents) + + "\n\nUse the available tools to search the corpus, then return ONLY the " + "ranked blocks (each with a ) for the most " + "relevant documents. Do not answer the question yourself." + ) + + common: dict = {"model": model, "tools": tool_specs, "max_output_tokens": max_tokens} + if reasoning_effort: + common["reasoning"] = {"effort": reasoning_effort} + + doc_text: dict[str, str] = {} + tool_types_used: set[str] = set() + tool_call_count = 0 + final_text = "" + usage = _empty_usage() + search_history: list[str] = [] + turn_tools: list[list[str]] = [] + + response = client.responses.create(input=prompt, **common) + num_turns = 1 + _acc_responses_usage(usage, response) + + while True: + function_calls = [o for o in response.output if getattr(o, "type", None) == "function_call"] + if not function_calls: + final_text = getattr(response, "output_text", "") or "" + break + if num_turns >= max_turns: + final_text = getattr(response, "output_text", "") or "" + break + + turn_tools.append([fc.name for fc in function_calls]) + outputs: list[dict] = [] + for fc in function_calls: + name = fc.name + tool_types_used.add(name) + tool_call_count += 1 + args = _parse_tool_arguments(fc.arguments) + if name in ("search_corpus", "grep_corpus"): + q = args.get("query") or args.get("pattern") or args.get("q") or "" + if q: + search_history.append(f"{name}: {str(q)[:100]}") + tool = toolset.get_tool(name) + if tool is None: + output = f"Error: unknown tool '{name}'." + else: + try: + output, _metadata = tool(args) + _collect_doc_text(output, doc_text) + except Exception as exc: + logger.warning("responses_tool_error", tool=name, error=str(exc)) + output = f"Error executing '{name}': {exc}" + outputs.append( + {"type": "function_call_output", "call_id": fc.call_id, "output": output} + ) + + response = client.responses.create( + previous_response_id=response.id, input=outputs, **common + ) + num_turns += 1 + _acc_responses_usage(usage, response) + + documents = _extract_documents(final_text, doc_text, max_documents) + + pool_doc_ids = sorted({cid.split("__")[0] for cid in doc_text}) + + logger.info( + "responses_search_complete", + model=model, + num_turns=num_turns, + num_documents=len(documents), + tool_calls=tool_call_count, + pool_size=len(pool_doc_ids), + ) + + return AgentSearchResult( + documents=documents, + num_turns=num_turns, + final_text=final_text, + pool_doc_ids=pool_doc_ids, + usage=usage, + trajectory={ + "search_history": search_history, + "turn_tools": turn_tools, + "final_docs": [d.id for d in documents], + }, + metadata={ + "backend": "openai_responses", + "model": model, + "tool_calls": tool_call_count, + "tool_types_used": ",".join(sorted(tool_types_used)), + }, + ) + + +def _acc_anthropic_usage(usage: dict[str, int], data: dict) -> None: + u = data.get("usage") or {} + inp = int(u.get("input_tokens") or 0) + out = int(u.get("output_tokens") or 0) + usage["prompt_tokens"] += inp + usage["completion_tokens"] += out + usage["total_tokens"] += inp + out + usage["llm_calls"] += 1 + + +def _anthropic_messages_url(base_url: str) -> str: + b = base_url.rstrip("/") + if b.endswith("/messages"): + return b + if b.endswith("/v1"): + return b + "/messages" + return b + "/v1/messages" + + +def run_anthropic_search( + *, + toolset: ToolSet, + base_url: str, + api_key: str, + model: str, + query: str, + max_documents: int = 20, + max_turns: int = 20, + max_tokens: int = 4096, + anthropic_version: str = "2023-06-01", + auth_header: str = "x-api-key", + timeout_s: int = 600, +) -> AgentSearchResult: + tools = [ + tool.get_format(ProviderFormat.ANTHROPIC) + for name, tool in toolset.tools.items() + if name in _CHAT_TOOL_NAMES + ] + system = get_retrieval_subagent_prompt(query, num_output_docs=max_documents) + messages: list[dict] = [ + { + "role": "user", + "content": ( + "Use the available tools to search the corpus, then return ONLY the " + "ranked blocks (each with a ) for the most " + "relevant documents. Do not answer the question yourself." + ), + } + ] + + url = _anthropic_messages_url(base_url) + headers = { + "content-type": "application/json", + "anthropic-version": anthropic_version, + auth_header: api_key, + } + + doc_text: dict[str, str] = {} + tool_types_used: set[str] = set() + tool_call_count = 0 + final_text = "" + num_turns = 0 + usage = _empty_usage() + search_history: list[str] = [] + turn_tools: list[list[str]] = [] + + for _ in range(max_turns): + payload = { + "model": model, + "max_tokens": max_tokens, + "system": system, + "messages": messages, + "tools": tools, + } + response = requests.post(url, json=payload, headers=headers, timeout=timeout_s) + response.raise_for_status() + data = response.json() + num_turns += 1 + _acc_anthropic_usage(usage, data) + + content = data.get("content") or [] + messages.append({"role": "assistant", "content": content}) + + tool_uses = [b for b in content if b.get("type") == "tool_use"] + if not tool_uses: + final_text = "".join( + b.get("text", "") for b in content if b.get("type") == "text" + ) + break + + turn_tools.append([tu.get("name", "") for tu in tool_uses]) + tool_results: list[dict] = [] + for tu in tool_uses: + name = tu.get("name", "") + tool_types_used.add(name) + tool_call_count += 1 + args = tu.get("input") or {} + if name in ("search_corpus", "grep_corpus"): + q = args.get("query") or args.get("pattern") or "" + if q: + search_history.append(f"{name}: {str(q)[:100]}") + tool = toolset.get_tool(name) + if tool is None: + output = f"Error: unknown tool '{name}'." + else: + try: + output, _metadata = tool(args) + _collect_doc_text(output, doc_text) + except Exception as exc: + logger.warning("anthropic_tool_error", tool=name, error=str(exc)) + output = f"Error executing '{name}': {exc}" + tool_results.append( + {"type": "tool_result", "tool_use_id": tu.get("id"), "content": output} + ) + messages.append({"role": "user", "content": tool_results}) + else: + final_text = "" + + documents = _extract_documents(final_text, doc_text, max_documents) + pool_doc_ids = sorted({cid.split("__")[0] for cid in doc_text}) + + logger.info( + "anthropic_search_complete", + model=model, + num_turns=num_turns, + num_documents=len(documents), + tool_calls=tool_call_count, + pool_size=len(pool_doc_ids), + ) + + return AgentSearchResult( + documents=documents, + num_turns=num_turns, + final_text=final_text, + pool_doc_ids=pool_doc_ids, + usage=usage, + trajectory={ + "search_history": search_history, + "turn_tools": turn_tools, + "final_docs": [d.id for d in documents], + }, + metadata={ + "backend": "anthropic_messages", + "model": model, + "tool_calls": tool_call_count, + "tool_types_used": ",".join(sorted(tool_types_used)), + }, + ) + + +__all__ = [ + "ChatDocument", + "AgentSearchResult", + "run_anthropic_search", + "run_chat_search", + "run_responses_search", +] diff --git a/cosmos-retriever/src/cosmos_retriever/prompts.py b/cosmos-retriever/src/cosmos_retriever/prompts.py new file mode 100644 index 0000000..6920df8 --- /dev/null +++ b/cosmos-retriever/src/cosmos_retriever/prompts.py @@ -0,0 +1,83 @@ + +from __future__ import annotations + + +def get_retrieval_subagent_prompt(query: str, *, num_output_docs: int = 30) -> str: + + return f""" + + You are a retrieval subagent in a multi-agent system. Your specific role is to identify and retrieve the most relevant documents from a large corpus to help another agent answer questions. You do NOT answer questions yourself - you only find and retrieve relevant documents. + + Here is the query you need to find documents for: + + + {query} + + + **Available Tools:** + - SearchTool: Hybrid semantic and keyword search + - GrepTool: Text pattern matching + - ReadDocument: Read specific document snippets that look promising but incomplete + - PruneChunksTool: Remove irrelevant chunks to free up context space + + **Your Process:** + - Break down the query into its key concepts and information needs (list each one explicitly) + - For each key concept, develop a specific search strategy that targets that concept + - Consider what types of documents and evidence would be most helpful for answering this query + - Plan several distinct, non-overlapping search strategies that approach the question from different angles + - Then execute your searches using multiple parallel tool calls. + + **Your Thinking:** + After each round of searches, in your thinking: + - Consider the following: + - **What do I know?**: List the key topics, themes, or aspects of the question that your currently retrieved documents address. What specific information do you have? + - **What should I search for next?**: Systematically consider what search approaches, keywords, or document types you haven't yet tried that might yield valuable information. + - **What should I prune?**: If you were to prune chunks, what would you remove and what new searches would you prioritize? Would this likely yield significantly better or more complete information than what you currently have? + - **Do I have enough information?**: Given the question's complexity and requirements, do you have sufficient information to help answer it, or are there critical gaps? + - Decide if additional searches are needed (and if so, ensure they use genuinely different approaches and do not duplicate or redundant searches) + - Avoid getting stuck on a single search strategy - if one approach isn't yielding results, prune and backtrack and try different approaches + + **Tactics to Consider:** + - When queries fail, try different approaches or keywords to improve the results + - Avoid duplicate or redundant searches + - Execute multiple tool calls in parallel when possible + - It's OK for this section to be quite long. + - If you notice your token budget is approaching the threshold, prune irrelevant chunks proactively to avoid running out of context. + - Focus on gathering as much relevant information as possible, it is useful to get multiple perspectives on the same topic or redundant information to confirm the information you have found is correct. + - Follow explicit textual evidence rather than speculation + + **Output Format:** + Present your final results in order from most relevant to least relevant using this structure: + + + + Brief explanation (1-3 sentences) of why this document is relevant to the query. + + + + Example: + + + This document contains detailed analysis of the specific topic mentioned in the query and provides quantitative data that directly supports answering the question. + + + + Your final output should consist only of the up to {num_output_docs} ranked document results in the specified format and should not duplicate or rehash any of the search planning or evaluation work you did in the thinking block. +` + """ + + +def get_retrieval_subagent_budget_exhausted_message( + current_token_usage: int, threshold_budget: int +) -> str: + + return ( + f"[Token usage: {current_token_usage}/{threshold_budget}] **OVER BUDGET.** \n" + "**CRITICAL CONSTRAINT:** You are currently at or near your token budget limit. " + "You CANNOT search, grep, or read any additional documents unless you prune chunks and reduce your token usage.\n" + "You must now make a strategic decision between two options:\n" + "**Option 1: Prune chunks** By using the PruneChunksTool and continue searching after.\n" + "Account for the tokens used by each chunk and the relevancy of the chunks to determine which chunks to prune.**\n" + "\n**Option 2: Conclude your search**\n" + "Before making your decision, work through your strategic analysis and if concluding your search ensure you have the final correct exhaustive set of documents to answer the question and all its subquestions." + ) diff --git a/cosmos-retriever/src/cosmos_retriever/rerank.py b/cosmos-retriever/src/cosmos_retriever/rerank.py new file mode 100644 index 0000000..5ccf5d7 --- /dev/null +++ b/cosmos-retriever/src/cosmos_retriever/rerank.py @@ -0,0 +1,391 @@ +from __future__ import annotations + +from abc import ABC, abstractmethod +from dataclasses import dataclass +import time +from typing import TYPE_CHECKING, Callable, List, Optional + +import requests +import structlog + +from cosmos_retriever.config import get_config + +if TYPE_CHECKING: + from baseten_performance_client import ClassificationResponse, PerformanceClient + +logger = structlog.get_logger("search_agent.rerank") + + +@dataclass +class RerankResult: + + document: str + score: float + original_index: int + tokens: Optional[int] = None + + +class Reranker(ABC): + + def __init__( + self, + token_counter: Optional[Callable[[str], int]] = None, + max_tokens: Optional[int] = None, + ): + if max_tokens is not None and token_counter is None: + raise ValueError("token_counter is required when max_tokens is specified") + self.token_counter = token_counter + self.max_tokens = max_tokens + + def _truncate_results( + self, results: List[RerankResult], max_tokens: Optional[int] = None + ) -> List[RerankResult]: + if self.token_counter is not None: + for result in results: + result.tokens = self.token_counter(result.document) + + effective_max_tokens = max_tokens if max_tokens is not None else self.max_tokens + if self.token_counter is None or effective_max_tokens is None: + return results + + truncated: List[RerankResult] = [] + total_tokens = 0 + for result in results: + doc_tokens = result.tokens + assert doc_tokens is not None + if total_tokens + doc_tokens > effective_max_tokens: + logger.info( + "truncating_results", + kept=len(truncated), + dropped=len(results) - len(truncated), + total_tokens=total_tokens, + max_tokens=effective_max_tokens, + ) + break + truncated.append(result) + total_tokens += doc_tokens + + return truncated + + @abstractmethod + def _rerank( + self, + query: str, + documents: List[str], + instruction: Optional[str] = None, + ) -> List[RerankResult]: + pass + + def __call__( + self, + query: str, + documents: List[str], + instruction: Optional[str] = None, + max_tokens: Optional[int] = None, + ) -> List[RerankResult]: + start = time.perf_counter() + results = self._rerank(query, documents, instruction) + elapsed_ms = (time.perf_counter() - start) * 1000 + if elapsed_ms > 1500: + logger.warning( + "Extremely slow reranking", + elapsed_ms=round(elapsed_ms, 1), + ) + return self._truncate_results(results, max_tokens=max_tokens) + + +class BasetenReranker(Reranker): + + PREFIX = '<|im_start|>system\nJudge whether the Document meets the requirements based on the Query and the Instruct provided. Note that the answer can only be "yes" or "no".<|im_end|>\n<|im_start|>user\n' + SUFFIX = "<|im_end|>\n<|im_start|>assistant\n\n\n\n\n" + DEFAULT_INSTRUCTION = ( + "Given a web search query, retrieve relevant passages that answer the query" + ) + + def __init__( + self, + client: Optional[PerformanceClient] = None, + token_counter: Optional[Callable[[str], int]] = None, + max_tokens: Optional[int] = None, + batch_size: int = 16, + max_concurrent_requests: int = 256, + timeout_s: int = 360, + ): + super().__init__(token_counter=token_counter, max_tokens=max_tokens) + if client is None: + config = get_config() + client = config.get_baseten_client() + + + + self.client = client + self.batch_size = batch_size + self.max_concurrent_requests = max_concurrent_requests + self.timeout_s = timeout_s + + def _format_input( + self, instruction: Optional[str], query: str, document: str + ) -> str: + if instruction is None: + instruction = self.DEFAULT_INSTRUCTION + return f"{self.PREFIX}: {instruction}\n: {query}\n: {document}{self.SUFFIX}" + + def _rerank( + self, + query: str, + documents: list[str], + instruction: Optional[str] = None, + ) -> list[RerankResult]: + if not documents: + return [] + + inputs = [self._format_input(instruction, query, doc) for doc in documents] + + response: ClassificationResponse = self.client.classify( + inputs=inputs, + truncate=True, + batch_size=self.batch_size, + max_concurrent_requests=self.max_concurrent_requests, + timeout_s=self.timeout_s, + ) + + results = [] + for idx, (doc, group) in enumerate(zip(documents, response.data)): + score = 0.0 + for result in group: + if result.label == "yes": + score = result.score + break + results.append(RerankResult(document=doc, score=score, original_index=idx)) + + results.sort(key=lambda x: x.score, reverse=True) + return results + + +class VLLMQwen3Reranker(Reranker): + + PREFIX = '<|im_start|>system\nJudge whether the Document meets the requirements based on the Query and the Instruct provided. Note that the answer can only be "yes" or "no".<|im_end|>\n<|im_start|>user\n' + SUFFIX = "<|im_end|>\n<|im_start|>assistant\n\n\n\n\n" + DEFAULT_INSTRUCTION = ( + "Given a web search query, retrieve relevant passages that answer the query" + ) + + def __init__( + self, + base_url: Optional[str] = None, + model: str = "Qwen/Qwen3-Reranker-8B", + token_counter: Optional[Callable[[str], int]] = None, + max_tokens: Optional[int] = None, + batch_size: int = 32, + timeout_s: int = 360, + ): + super().__init__(token_counter=token_counter, max_tokens=max_tokens) + import os + + self.base_url = ( + base_url or os.getenv("VLLM_RERANKER_URL", "http://127.0.0.1:8011") + ).rstrip("/") + self.model = model + self.batch_size = batch_size + self.timeout_s = timeout_s + + def _rerank( + self, + query: str, + documents: List[str], + instruction: Optional[str] = None, + ) -> List[RerankResult]: + if not documents: + return [] + if instruction is None: + instruction = self.DEFAULT_INSTRUCTION + + text_1 = f"{self.PREFIX}: {instruction}\n: {query}\n" + scores: List[float] = [] + for start in range(0, len(documents), self.batch_size): + batch = documents[start : start + self.batch_size] + payload = { + "model": self.model, + "text_1": text_1, + "text_2": [f": {doc}{self.SUFFIX}" for doc in batch], + "truncate_prompt_tokens": -1, + } + last_error: Optional[Exception] = None + for attempt in range(3): + try: + response = requests.post( + f"{self.base_url}/score", + json=payload, + timeout=self.timeout_s, + ) + response.raise_for_status() + data = response.json()["data"] + scores.extend(float(item["score"]) for item in data) + last_error = None + break + except requests.exceptions.RequestException as exc: + last_error = exc + logger.warning( + "vllm_rerank_retry", attempt=attempt + 1, error=str(exc) + ) + time.sleep(2**attempt) + if last_error is not None: + logger.error("vllm_rerank_failed", error=str(last_error)) + raise last_error + + results = [ + RerankResult(document=doc, score=score, original_index=idx) + for idx, (doc, score) in enumerate(zip(documents, scores)) + ] + results.sort(key=lambda x: x.score, reverse=True) + return results + + +class ContextualReranker(Reranker): + + API_URL = "https://api.contextual.ai/v1/rerank" + DEFAULT_MODEL = "ctxl-rerank-v2-instruct-multilingual" + DEFAULT_INSTRUCTION = "Prioritize results that most closely align with the criteria outlined in the query" + + def __init__( + self, + api_key: Optional[str] = None, + model: Optional[str] = None, + token_counter: Optional[Callable[[str], int]] = None, + max_tokens: Optional[int] = None, + top_n: Optional[int] = None, + timeout_s: int = 60, + ): + super().__init__(token_counter=token_counter, max_tokens=max_tokens) + if api_key is None: + config = get_config() + api_key = config.contextual_api_key.get_secret_value() + self.api_key = api_key + self.model = model or self.DEFAULT_MODEL + self.top_n = top_n + self.timeout_s = timeout_s + + def _rerank( + self, + query: str, + documents: list[str], + instruction: Optional[str] = None, + ) -> list[RerankResult]: + if not documents: + return [] + + payload: dict[str, str | list[str] | int] = { + "query": query, + "documents": documents, + "model": self.model, + } + + if self.top_n is not None: + payload["top_n"] = self.top_n + + if instruction is not None: + payload["instruction"] = instruction + else: + payload["instruction"] = self.DEFAULT_INSTRUCTION + + headers = { + "Authorization": f"Bearer {self.api_key}", + "Content-Type": "application/json", + } + + try: + response = requests.post( + self.API_URL, + json=payload, + headers=headers, + timeout=self.timeout_s, + ) + response.raise_for_status() + data = response.json() + except requests.exceptions.RequestException as e: + logger.error("contextual_rerank_failed", error=str(e)) + raise + + results = [] + for item in data.get("results", []): + idx = item["index"] + score = item["relevance_score"] + results.append( + RerankResult( + document=documents[idx], + score=score, + original_index=idx, + ) + ) + + results.sort(key=lambda x: x.score, reverse=True) + return results + + +if __name__ == "__main__": + import argparse + import tiktoken + + parser = argparse.ArgumentParser(description="Run reranker example") + parser.add_argument( + "--reranker", + choices=["baseten", "contextual"], + default="baseten", + help="Reranker to use (default: baseten)", + ) + parser.add_argument( + "--max-tokens", + type=int, + default=30, + help="Maximum tokens for output (default: 30)", + ) + args = parser.parse_args() + + logger.info( + "Running reranker example", reranker=args.reranker, max_tokens=args.max_tokens + ) + + enc = tiktoken.get_encoding("o200k_harmony") + token_counter = lambda text: len(enc.encode(text)) + + reranker: Reranker + if args.reranker == "contextual": + reranker = ContextualReranker( + token_counter=token_counter, + max_tokens=args.max_tokens, + ) + elif args.reranker == "baseten": + reranker = BasetenReranker( + token_counter=token_counter, + max_tokens=args.max_tokens, + ) + else: + raise ValueError(f"Invalid reranker: {args.reranker}") + + query = "What is the capital of China?" + documents = [ + "The capital of France is Paris.", + "The capital of China is Beijing.", + "The capital of Poland is Warsaw.", + "The capital of Germany is Berlin.", + "Chocolate is a delicious treat.", + "Pizza is a food", + "China has a population of 1.4 billion.", + "Germany has a population of 83 million.", + "Poland has a population of 38 million.", + "Warsaw is the capital of Poland.", + "Berlin is the capital of Germany.", + "Paris is the capital of France.", + "Beijing is the capital of China.", + "Warsaw is the capital of Poland.", + "Berlin is the capital of Germany.", + "Shanghai is not the capital of China.", + "Japan is closer to China than to the United States.", + "The capital of China has been Beijing for a long time.", + ] + results = reranker(query, documents) + logger.info("rerank_complete", num_results=len(results), max_tokens=args.max_tokens) + for result in results: + logger.info("result", score=result.score, document=result.document) + +VLLMReranker = VLLMQwen3Reranker diff --git a/cosmos-retriever/src/cosmos_retriever/retrieval/__init__.py b/cosmos-retriever/src/cosmos_retriever/retrieval/__init__.py new file mode 100644 index 0000000..cc0468e --- /dev/null +++ b/cosmos-retriever/src/cosmos_retriever/retrieval/__init__.py @@ -0,0 +1,58 @@ + +from __future__ import annotations + +from cosmos_retriever.retrieval.capabilities import ( + RetrievalCapabilities, + SupportLevel, + VectorCapability, +) +from cosmos_retriever.retrieval.embedding import QueryEmbedder +from cosmos_retriever.retrieval.defaults import ( + build_default_retriever, + default_capabilities_for, + default_chunked_schema, +) +from cosmos_retriever.retrieval.models import ( + EqualsFilter, + GrepRequest, + InFilter, + NormalizedDocument, + PartitionQueryPolicy, + RangeFilter, + ReadDocumentRequest, + RetrievedItem, + SearchRequest, +) +from cosmos_retriever.retrieval.paths import CosmosPath +from cosmos_retriever.retrieval.retriever import CorpusRetriever +from cosmos_retriever.retrieval.schema import ( + ChunkIdentityCodec, + CorpusSchema, + DunderChunkCodec, + VectorFieldConfig, +) + +__all__ = [ + "ChunkIdentityCodec", + "CorpusRetriever", + "CorpusSchema", + "CosmosPath", + "DunderChunkCodec", + "EqualsFilter", + "GrepRequest", + "InFilter", + "NormalizedDocument", + "PartitionQueryPolicy", + "QueryEmbedder", + "RangeFilter", + "ReadDocumentRequest", + "RetrievalCapabilities", + "RetrievedItem", + "SearchRequest", + "SupportLevel", + "VectorCapability", + "VectorFieldConfig", + "build_default_retriever", + "default_capabilities_for", + "default_chunked_schema", +] diff --git a/cosmos-retriever/src/cosmos_retriever/retrieval/capabilities.py b/cosmos-retriever/src/cosmos_retriever/retrieval/capabilities.py new file mode 100644 index 0000000..663b866 --- /dev/null +++ b/cosmos-retriever/src/cosmos_retriever/retrieval/capabilities.py @@ -0,0 +1,43 @@ +from __future__ import annotations + +from enum import StrEnum + +from pydantic import BaseModel + +from cosmos_retriever.retrieval.paths import CosmosPath +from cosmos_retriever.retrieval.schema import PathField + + +class SupportLevel(StrEnum): + INDEXED = "indexed" + SCAN = "scan" + UNSUPPORTED = "unsupported" + UNKNOWN = "unknown" + + +class VectorCapability(BaseModel): + path: PathField + dimensions: int + distance_function: str = "cosine" + index_type: str | None = None + support: SupportLevel = SupportLevel.UNKNOWN + + +class RetrievalCapabilities(BaseModel): + vector_fields: list[VectorCapability] = [] + full_text_paths: list[PathField] = [] + range_indexed_paths: list[PathField] = [] + partition_key_paths: list[PathField] = [] + native_hybrid_supported: bool = False + full_text_supported: bool = False + vector_supported: bool = False + efficient_document_lookup_supported: bool = False + + def vector_capability_for(self, path: CosmosPath) -> VectorCapability | None: + for v in self.vector_fields: + if str(v.path) == str(path): + return v + return None + + def has_full_text_path(self, path: CosmosPath) -> bool: + return any(str(p) == str(path) for p in self.full_text_paths) diff --git a/cosmos-retriever/src/cosmos_retriever/retrieval/compiler.py b/cosmos-retriever/src/cosmos_retriever/retrieval/compiler.py new file mode 100644 index 0000000..2d5a7a0 --- /dev/null +++ b/cosmos-retriever/src/cosmos_retriever/retrieval/compiler.py @@ -0,0 +1,259 @@ +from __future__ import annotations + +from typing import Any + +from cosmos_retriever.retrieval.errors import QueryCompilationError +from cosmos_retriever.retrieval.expressions import fts_literal_args, tokenize_for_fts +from cosmos_retriever.retrieval.models import ( + CompiledCosmosQuery, + EqualsFilter, + FilterExpression, + InFilter, + RangeFilter, +) +from cosmos_retriever.retrieval.paths import CosmosPath +from cosmos_retriever.retrieval.schema import CorpusSchema + +_ALIAS = "c" + + +class _ParamBag: + + def __init__(self) -> None: + self.params: list[dict[str, Any]] = [] + self._n = 0 + + def add(self, value: Any, prefix: str = "p") -> str: + name = f"@{prefix}{self._n}" + self._n += 1 + self.params.append({"name": name, "value": value}) + return name + + +class CosmosQueryCompiler: + + def __init__(self, schema: CorpusSchema) -> None: + self.schema = schema + + def _resolve_logical(self, name: str) -> CosmosPath: + s = self.schema + mapping: dict[str, CosmosPath | None] = { + "item_id": s.item_id_path, + "text": s.primary_text_path, + "primary_text": s.primary_text_path, + "document_id": s.document_id_path, + "chunk_id": s.chunk_id_path, + "chunk_order": s.chunk_order_path, + "title": s.title_path, + "source": s.source_path, + } + if name in mapping and mapping[name] is not None: + return mapping[name] + if name in s.metadata_paths: + return s.metadata_paths[name] + raise QueryCompilationError(f"unknown logical field {name!r}") + + + def projection(self, limit_param: str) -> tuple[str, dict[str, str]]: + + + s = self.schema + cols: list[str] = [] + aliases: dict[str, str] = {} + + def add(logical: str, path: CosmosPath | None) -> None: + if path is None: + return + cols.append(f"{path.render(_ALIAS)} AS {logical}") + aliases[logical] = logical + + add("item_id", s.item_id_path) + add("text", s.primary_text_path) + add("document_id", s.document_id_path) + add("chunk_id", s.chunk_id_path) + add("chunk_order", s.chunk_order_path) + add("title", s.title_path) + add("source", s.source_path) + + for i, (fname, fpath) in enumerate(s.text_field_map().items()): + alias = f"txt_{i}" + cols.append(f"{fpath.render(_ALIAS)} AS {alias}") + aliases[alias] = fname + for key, path in s.metadata_paths.items(): + cols.append(f"{path.render(_ALIAS)} AS md_{key}") + aliases[f"md_{key}"] = key + + select = f"SELECT TOP {limit_param} " + ", ".join(cols) + f" FROM {_ALIAS}" + return select, aliases + + + + def _compile_filter(self, f: FilterExpression, bag: _ParamBag) -> str: + path = self._resolve_logical(f.logical_field).render(_ALIAS) + if isinstance(f, EqualsFilter): + return f"{path} = {bag.add(f.value)}" + if isinstance(f, RangeFilter): + parts: list[str] = [] + if f.minimum is not None: + parts.append(f"{path} >= {bag.add(f.minimum)}") + if f.maximum is not None: + parts.append(f"{path} <= {bag.add(f.maximum)}") + return "(" + " AND ".join(parts) + ")" if parts else "true" + if isinstance(f, InFilter): + return f"ARRAY_CONTAINS({bag.add(list(f.values))}, {path})" + raise QueryCompilationError(f"unsupported filter {type(f).__name__}") + + def _where( + self, + filters: list[FilterExpression], + ignored_item_ids: list[str], + bag: _ParamBag, + ) -> str: + clauses = [self._compile_filter(f, bag) for f in filters] + if ignored_item_ids: + item_id = self.schema.item_id_path.render(_ALIAS) + clauses.append(f"NOT ARRAY_CONTAINS({bag.add(ignored_item_ids)}, {item_id})") + return (" WHERE " + " AND ".join(clauses)) if clauses else "" + + + def compile_hybrid( + self, + *, + query: str, + query_vector: list[float], + limit: int, + ignored_item_ids: list[str], + filters: list[FilterExpression], + partition_key: Any | None, + cross_partition: bool, + vector_path: CosmosPath, + text_paths: list[CosmosPath], + ) -> CompiledCosmosQuery: + bag = _ParamBag() + limit_p = bag.add(limit, prefix="k") + vec_p = bag.add(query_vector, prefix="qVec") + select, aliases = self.projection(limit_p) + where = self._where(filters, ignored_item_ids, bag) + terms = fts_literal_args(tokenize_for_fts(query)) + fts = ", ".join( + f"FullTextScore({tp.render(_ALIAS)}, {terms})" for tp in text_paths + ) + order = ( + " ORDER BY RANK RRF(" + f"VectorDistance({vector_path.render(_ALIAS)}, {vec_p}), {fts})" + ) + return CompiledCosmosQuery( + sql=select + where + order, + parameters=bag.params, + partition_key=partition_key, + enable_cross_partition_query=cross_partition, + strategy="native_hybrid", + projected_aliases=aliases, + ) + + def compile_vector( + self, + *, + query_vector: list[float], + limit: int, + ignored_item_ids: list[str], + filters: list[FilterExpression], + partition_key: Any | None, + cross_partition: bool, + vector_path: CosmosPath, + ) -> CompiledCosmosQuery: + bag = _ParamBag() + limit_p = bag.add(limit, prefix="k") + vec_p = bag.add(query_vector, prefix="qVec") + select, aliases = self.projection(limit_p) + where = self._where(filters, ignored_item_ids, bag) + order = f" ORDER BY RANK VectorDistance({vector_path.render(_ALIAS)}, {vec_p})" + return CompiledCosmosQuery( + sql=select + where + order, + parameters=bag.params, + partition_key=partition_key, + enable_cross_partition_query=cross_partition, + strategy="vector", + projected_aliases=aliases, + ) + + def compile_full_text( + self, + *, + query: str, + limit: int, + ignored_item_ids: list[str], + filters: list[FilterExpression], + partition_key: Any | None, + cross_partition: bool, + text_paths: list[CosmosPath], + strategy: str = "full_text", + ) -> CompiledCosmosQuery: + bag = _ParamBag() + limit_p = bag.add(limit, prefix="k") + select, aliases = self.projection(limit_p) + where = self._where(filters, ignored_item_ids, bag) + terms = fts_literal_args(tokenize_for_fts(query)) + if len(text_paths) == 1: + order = f" ORDER BY RANK FullTextScore({text_paths[0].render(_ALIAS)}, {terms})" + else: + fts = ", ".join( + f"FullTextScore({tp.render(_ALIAS)}, {terms})" for tp in text_paths + ) + order = f" ORDER BY RANK RRF({fts})" + return CompiledCosmosQuery( + sql=select + where + order, + parameters=bag.params, + partition_key=partition_key, + enable_cross_partition_query=cross_partition, + strategy=strategy, + projected_aliases=aliases, + ) + + def compile_structured( + self, + *, + limit: int, + filters: list[FilterExpression], + ignored_item_ids: list[str], + partition_key: Any | None, + cross_partition: bool, + ) -> CompiledCosmosQuery: + bag = _ParamBag() + limit_p = bag.add(limit, prefix="k") + select, aliases = self.projection(limit_p) + where = self._where(filters, ignored_item_ids, bag) + return CompiledCosmosQuery( + sql=select + where, + parameters=bag.params, + partition_key=partition_key, + enable_cross_partition_query=cross_partition, + strategy="structured", + projected_aliases=aliases, + ) + + + def compile_document_read( + self, + *, + document_id: str, + max_chunks: int, + partition_key: Any | None, + cross_partition: bool, + ) -> CompiledCosmosQuery: + s = self.schema + if s.document_id_path is None: + raise QueryCompilationError("document_id_path is not configured") + bag = _ParamBag() + limit_p = bag.add(max_chunks, prefix="k") + select, aliases = self.projection(limit_p) + doc_p = bag.add(document_id, prefix="doc") + where = f" WHERE {s.document_id_path.render(_ALIAS)} = {doc_p}" + return CompiledCosmosQuery( + sql=select + where, + parameters=bag.params, + partition_key=partition_key, + enable_cross_partition_query=cross_partition, + strategy="document_read", + projected_aliases=aliases, + ) diff --git a/cosmos-retriever/src/cosmos_retriever/retrieval/defaults.py b/cosmos-retriever/src/cosmos_retriever/retrieval/defaults.py new file mode 100644 index 0000000..d7e5454 --- /dev/null +++ b/cosmos-retriever/src/cosmos_retriever/retrieval/defaults.py @@ -0,0 +1,79 @@ +from __future__ import annotations + +from cosmos_retriever.retrieval.capabilities import ( + RetrievalCapabilities, + SupportLevel, + VectorCapability, +) +from cosmos_retriever.retrieval.embedding import QueryEmbedder +from cosmos_retriever.retrieval.models import PartitionQueryPolicy +from cosmos_retriever.retrieval.retriever import CorpusRetriever +from cosmos_retriever.retrieval.schema import ( + CorpusSchema, + DunderChunkCodec, + VectorFieldConfig, +) + +_DEFAULT_DIMENSIONS = 1536 + + +def default_chunked_schema( + embedding_model: str = "text-embedding-3-small", + dimensions: int = _DEFAULT_DIMENSIONS, +) -> CorpusSchema: + schema = CorpusSchema( + item_id_path="/id", + text_paths=["/text"], + primary_text_path="/text", + vector_fields=[ + VectorFieldConfig( + path="/embedding", + embedding_model=embedding_model, + dimensions=dimensions, + distance_function="cosine", + ) + ], + document_id_path="/docid", + chunk_id_path="/id", + chunk_order_path="/chunk_idx", + partition_key_paths=["/docid"], + ) + schema.identity_codec = DunderChunkCodec() + return schema + + +def default_capabilities_for(schema: CorpusSchema) -> RetrievalCapabilities: + field = schema.vector_fields[0] + return RetrievalCapabilities( + vector_fields=[ + VectorCapability( + path=field.path, + dimensions=field.dimensions, + distance_function=field.distance_function, + support=SupportLevel.INDEXED, + ) + ], + full_text_paths=[schema.primary_text_path], + partition_key_paths=list(schema.partition_key_paths), + native_hybrid_supported=True, + full_text_supported=True, + vector_supported=True, + efficient_document_lookup_supported=True, + ) + + +def build_default_retriever( + *, + container, + embedder: QueryEmbedder | None, + embedding_model: str = "text-embedding-3-small", + partition_policy: PartitionQueryPolicy | None = None, +) -> CorpusRetriever: + schema = default_chunked_schema(embedding_model=embedding_model) + return CorpusRetriever( + container=container, + schema=schema, + capabilities=default_capabilities_for(schema), + query_embedder=embedder, + partition_policy=partition_policy or PartitionQueryPolicy(), + ) diff --git a/cosmos-retriever/src/cosmos_retriever/retrieval/document_resolvers.py b/cosmos-retriever/src/cosmos_retriever/retrieval/document_resolvers.py new file mode 100644 index 0000000..a2ecced --- /dev/null +++ b/cosmos-retriever/src/cosmos_retriever/retrieval/document_resolvers.py @@ -0,0 +1,149 @@ + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import Any + +from cosmos_retriever.retrieval.compiler import CosmosQueryCompiler +from cosmos_retriever.retrieval.errors import ( + CrossPartitionQueryDisabled, + DocumentResolutionUnsupported, +) +from cosmos_retriever.retrieval.executor import CosmosExecutor +from cosmos_retriever.retrieval.models import ( + EqualsFilter, + NormalizedDocument, + PartitionQueryPolicy, + ReadDocumentRequest, +) +from cosmos_retriever.retrieval.schema import CorpusSchema + +DEFAULT_MAX_CHUNKS = 300 + + +class DocumentResolver(ABC): + def __init__( + self, + schema: CorpusSchema, + compiler: CosmosQueryCompiler, + executor: CosmosExecutor, + policy: PartitionQueryPolicy, + ) -> None: + self.schema = schema + self.compiler = compiler + self.executor = executor + self.policy = policy + + @abstractmethod + def resolve(self, request: ReadDocumentRequest) -> NormalizedDocument: ... + + def _derive_document_id(self, request: ReadDocumentRequest) -> str: + raw = request.document_id or request.item_id or "" + codec = self.schema.identity_codec + return codec.to_document_id(raw) if codec is not None else raw + + @staticmethod + def _sorted_rows(rows: list[dict[str, Any]]) -> list[dict[str, Any]]: + return sorted(rows, key=lambda r: r.get("chunk_order") or 0) + + +class ItemIsDocumentResolver(DocumentResolver): + + + def resolve(self, request: ReadDocumentRequest) -> NormalizedDocument: + item_id = request.item_id or request.document_id or "" + compiled = self.compiler.compile_structured( + + limit=1, + filters=[EqualsFilter(logical_field="item_id", value=item_id)], + ignored_item_ids=[], + + partition_key=request.partition_key, + cross_partition=request.partition_key is None, + ) + rows = self.executor.run(compiled) + return NormalizedDocument( + document_id=item_id, + + chunk_texts=[r.get("text", "") or "" for r in rows], + chunk_ids=[str(r.get("item_id")) for r in rows], + ) + + +class ChunkedDocumentResolver(DocumentResolver): + + def resolve(self, request: ReadDocumentRequest) -> NormalizedDocument: + doc_id = self._derive_document_id(request) + + max_chunks = request.max_chunks or DEFAULT_MAX_CHUNKS + + partition_key = request.partition_key or doc_id + compiled = self.compiler.compile_document_read( + + document_id=doc_id, + max_chunks=max_chunks, + + partition_key=partition_key, + cross_partition=False, + ) + rows = self._sorted_rows(self.executor.run(compiled)) + return NormalizedDocument( + + document_id=doc_id, + chunk_texts=[r.get("text", "") or "" for r in rows], + + chunk_ids=[str(r.get("item_id")) for r in rows], + ) + + +class CrossPartitionChunkedDocumentResolver(DocumentResolver): + def resolve(self, request: ReadDocumentRequest) -> NormalizedDocument: + if not self.policy.allow_cross_partition_document_read: + + raise CrossPartitionQueryDisabled( + "read_document requires cross-partition reconstruction, which is disabled" + ) + doc_id = self._derive_document_id(request) + + max_chunks = request.max_chunks or DEFAULT_MAX_CHUNKS + + compiled = self.compiler.compile_document_read( + document_id=doc_id, + + max_chunks=max_chunks, + partition_key=request.partition_key, + + cross_partition=request.partition_key is None, + ) + rows = self._sorted_rows(self.executor.run(compiled)) + + return NormalizedDocument( + document_id=doc_id, + + chunk_texts=[r.get("text", "") or "" for r in rows], + chunk_ids=[str(r.get("item_id")) for r in rows], + + warnings=["cross-partition document reconstruction"], + ) + + +def build_document_resolver( + schema: CorpusSchema, + + compiler: CosmosQueryCompiler, + executor: CosmosExecutor, + + + policy: PartitionQueryPolicy, +) -> DocumentResolver: + + if schema.is_item_document_mode: + return ItemIsDocumentResolver(schema, compiler, executor, policy) + if schema.document_id_path is None: + + raise DocumentResolutionUnsupported("no document reconstruction is possible") + if schema.partition_key_is_document_id: + return ChunkedDocumentResolver(schema, compiler, executor, policy) + + + return CrossPartitionChunkedDocumentResolver(schema, compiler, executor, policy) diff --git a/cosmos-retriever/src/cosmos_retriever/retrieval/embedding.py b/cosmos-retriever/src/cosmos_retriever/retrieval/embedding.py new file mode 100644 index 0000000..5d398b6 --- /dev/null +++ b/cosmos-retriever/src/cosmos_retriever/retrieval/embedding.py @@ -0,0 +1,24 @@ +from __future__ import annotations + +import openai + + +class QueryEmbedder: + + def __init__( + self, + client: openai.OpenAI, + model: str, + query_instruction: str | None = None, + ) -> None: + self._client = client + self._model = model + self._instruction = query_instruction + + def embed(self, text: str) -> list[float]: + if self._instruction: + text = f"Instruct: {self._instruction}\nQuery: {text}" + resp = self._client.embeddings.create( + model=self._model, input=[text], encoding_format="float" + ) + return resp.data[0].embedding diff --git a/cosmos-retriever/src/cosmos_retriever/retrieval/errors.py b/cosmos-retriever/src/cosmos_retriever/retrieval/errors.py new file mode 100644 index 0000000..4ceca8e --- /dev/null +++ b/cosmos-retriever/src/cosmos_retriever/retrieval/errors.py @@ -0,0 +1,50 @@ + +from __future__ import annotations + + +class RetrievalError(Exception): + pass + + +class InvalidCorpusSchema(RetrievalError): + pass + + +class UnsafeCosmosPath(RetrievalError): + pass + + +class UnsupportedRetrievalCapability(RetrievalError): + pass + + +class UnknownField(RetrievalError): + pass + + +class EmbeddingProfileMismatch(RetrievalError): + pass + + +class MissingPartitionKey(RetrievalError): + pass + + +class CrossPartitionQueryDisabled(RetrievalError): + pass + + +class UnboundedScanRejected(RetrievalError): + pass + + +class DocumentResolutionUnsupported(RetrievalError): + pass + + +class QueryCompilationError(RetrievalError): + pass + + +class IndexNotReady(RetrievalError): + pass diff --git a/cosmos-retriever/src/cosmos_retriever/retrieval/executor.py b/cosmos-retriever/src/cosmos_retriever/retrieval/executor.py new file mode 100644 index 0000000..64b628a --- /dev/null +++ b/cosmos-retriever/src/cosmos_retriever/retrieval/executor.py @@ -0,0 +1,93 @@ +from __future__ import annotations + +import os +import threading +import time +from typing import Any + +import structlog +import tenacity +from azure.cosmos import ContainerProxy +from azure.cosmos.exceptions import CosmosHttpResponseError + +from cosmos_retriever.retrieval.models import CompiledCosmosQuery + +logger = structlog.get_logger("cosmos_retriever.retrieval.executor") + + +def _read_positive_int_env(name: str, default: int) -> int: + raw = os.environ.get(name) + if raw is None: + return default + try: + value = int(raw) + except ValueError: + logger.warning("invalid_int_env", name=name, value=raw, default=default) + return default + if value < 1: + logger.warning("invalid_positive_int_env", name=name, value=raw, default=default) + return default + return value + + +COSMOS_QUERY_MAX_CONCURRENCY = _read_positive_int_env("COSMOS_QUERY_MAX_CONCURRENCY", 8) +_COSMOS_QUERY_SEMAPHORE = threading.BoundedSemaphore(COSMOS_QUERY_MAX_CONCURRENCY) + + +def _is_retryable_cosmos_error(exc: BaseException) -> bool: + if not isinstance(exc, CosmosHttpResponseError): + return False + status = getattr(exc, "status_code", None) + return status in (408, 429, 449, 500, 502, 503, 504) + + +@tenacity.retry( + stop=tenacity.stop_after_attempt(5), + wait=tenacity.wait_exponential(multiplier=1, min=4, max=15), + retry=tenacity.retry_if_exception(_is_retryable_cosmos_error), + before_sleep=lambda retry_state: logger.warning( + "retry_cosmos_query", + attempt=retry_state.attempt_number, + error=str(retry_state.outcome.exception()) if retry_state.outcome else None, + ), +) +def _query_items( + container: ContainerProxy, + query: str, + parameters: list[dict[str, Any]], + *, + partition_key: Any | None, + enable_cross_partition_query: bool, +) -> list[dict[str, Any]]: + start = time.perf_counter() + with _COSMOS_QUERY_SEMAPHORE: + kwargs: dict[str, Any] = {"query": query, "parameters": parameters} + if partition_key is not None: + kwargs["partition_key"] = partition_key + elif enable_cross_partition_query: + kwargs["enable_cross_partition_query"] = True + result = list(container.query_items(**kwargs)) + elapsed_ms = (time.perf_counter() - start) * 1000 + if elapsed_ms > 4500: + logger.warning( + "slow_cosmos_query", + elapsed_ms=round(elapsed_ms, 1), + cosmos_max_concurrency=COSMOS_QUERY_MAX_CONCURRENCY, + ) + return result + + +class CosmosExecutor: + + + def __init__(self, container: ContainerProxy) -> None: + self._container = container + + def run(self, compiled: CompiledCosmosQuery) -> list[dict[str, Any]]: + return _query_items( + self._container, + compiled.sql, + compiled.parameters, + partition_key=compiled.partition_key, + enable_cross_partition_query=compiled.enable_cross_partition_query, + ) diff --git a/cosmos-retriever/src/cosmos_retriever/retrieval/expressions.py b/cosmos-retriever/src/cosmos_retriever/retrieval/expressions.py new file mode 100644 index 0000000..ed1200d --- /dev/null +++ b/cosmos-retriever/src/cosmos_retriever/retrieval/expressions.py @@ -0,0 +1,34 @@ +from __future__ import annotations + +import re + +_TOKEN_RE = re.compile(r"\w+", re.UNICODE) +_STOPWORDS = frozenset( + ["a", "about", "above", "after", "again", "against", "all", "am", "an", "and", "any", "are", "as", "at", "be", "because", "been", "before", "being", "below", "between", "both", "but", "by", "can", "did", "do", "does", "doing", "don", "down", "during", "each", "few", "for", "from", "further", "had", "has", "have", "having", "he", "her", "here", "hers", "herself", "him", "himself", "his", "how", "i", "if", "in", "into", "is", "it", "its", "itself", "just", "like", "me", "more", "most", "my", "myself", "no", "nor", "not", "now", "of", "off", "on", "once", "only", "or", "other", "our", "ours", "ourselves", "out", "over", "own", "please", "same", "she", "should", "so", "some", "such", "tell", "than", "that", "the", "their", "theirs", "them", "themselves", "then", "there", "these", "they", "this", "those", "through", "to", "too", "under", "until", "up", "very", "was", "we", "were", "what", "when", "where", "which", "while", "who", "whom", "why", "will", "with", "would", "you", "your", "yours", "yourself", "yourselves"] +) + +_FTS_MAX_TERMS = 30 + + +def tokenize_for_fts(query: str) -> list[str]: + + out: list[str] = [] + seen: set[str] = set() + for raw in _TOKEN_RE.findall(query): + t = raw.lower() + if t in _STOPWORDS or t in seen: + continue + seen.add(t) + out.append(t) + if len(out) >= _FTS_MAX_TERMS: + break + return out + + +def fts_literal_args(terms: list[str]) -> str: + + + def esc(t: str) -> str: + return '"' + t.replace("\\", "\\\\").replace('"', '\\"') + '"' + + return ", ".join(esc(t) for t in terms) diff --git a/cosmos-retriever/src/cosmos_retriever/retrieval/formatting.py b/cosmos-retriever/src/cosmos_retriever/retrieval/formatting.py new file mode 100644 index 0000000..6c1834e --- /dev/null +++ b/cosmos-retriever/src/cosmos_retriever/retrieval/formatting.py @@ -0,0 +1,18 @@ +from __future__ import annotations + +DOC_TRUNCATION = 51_200_000 + + +def format_result_blocks( + triples: list[tuple[str, str, int | None]], +) -> str: + + blocks = [ + "\n# DOCUMENT ID: {}{} \n{}".format( + id_, + f" ({tokens} tokens)" if tokens is not None else "", + text[:DOC_TRUNCATION], + ) + for id_, text, tokens in triples + ] + return "\n".join(blocks) if blocks else "No results found" diff --git a/cosmos-retriever/src/cosmos_retriever/retrieval/models.py b/cosmos-retriever/src/cosmos_retriever/retrieval/models.py new file mode 100644 index 0000000..3a55673 --- /dev/null +++ b/cosmos-retriever/src/cosmos_retriever/retrieval/models.py @@ -0,0 +1,139 @@ + +from __future__ import annotations + +from typing import Annotated, Any, Literal + +from pydantic import BaseModel, Field + + +class EqualsFilter(BaseModel): + + kind: Literal["equals"] = "equals" + + + logical_field: str + value: Any + + +class RangeFilter(BaseModel): + kind: Literal["range"] = "range" + logical_field: str + minimum: Any | None = None + + maximum: Any | None = None + + +class InFilter(BaseModel): + kind: Literal["in"] = "in" + + + logical_field: str + values: list[Any] + + +FilterExpression = Annotated[ + EqualsFilter | RangeFilter | InFilter, Field(discriminator="kind") +] +class SearchRequest(BaseModel): + query: str + + + query_vector: list[float] | None = None + limit: int = 50 + ignored_item_ids: list[str] = Field(default_factory=list) + + + + filters: list[FilterExpression] = Field(default_factory=list) + partition_key: Any | None = None + + + + text_fields: list[str] | None = None + vector_field: str | None = None + mode: Literal["auto", "hybrid", "vector", "text"] = "auto" + + +class GrepRequest(BaseModel): + pattern: str + + candidate_limit: int = 50 + + result_limit: int = 5 + filters: list[FilterExpression] = Field(default_factory=list) + partition_key: Any | None = None + text_field: str | None = None + + +class ReadDocumentRequest(BaseModel): + document_id: str | None = None + item_id: str | None = None + + partition_key: Any | None = None + + + max_chunks: int | None = None + query: str | None = None + + +class RetrievedItem(BaseModel): + item_id: str + document_id: str | None = None + chunk_id: str | None = None + + chunk_order: int | None = None + text: str = "" + + + + text_fields: dict[str, str] = Field(default_factory=dict) + title: str | None = None + + source: str | None = None + metadata: dict[str, Any] = Field(default_factory=dict) + + + + partition_key: Any | None = None + retrieval_strategy: str = "" + retrieval_channels: list[str] = Field(default_factory=list) + + raw_scores: dict[str, float] = Field(default_factory=dict) + rank: int = 0 + + +class NormalizedDocument(BaseModel): + document_id: str | None = None + chunk_texts: list[str] = Field(default_factory=list) + chunk_ids: list[str] = Field(default_factory=list) + + warnings: list[str] = Field(default_factory=list) + + @property + def assembled(self) -> str: + return "".join(self.chunk_texts) + + +class CompiledCosmosQuery(BaseModel): + sql: str + parameters: list[dict[str, Any]] = Field(default_factory=list) + partition_key: Any | None = None + enable_cross_partition_query: bool = False + + + strategy: str = "" + projected_aliases: dict[str, str] = Field(default_factory=dict) + warnings: list[str] = Field(default_factory=list) + + + +class PartitionQueryPolicy(BaseModel): + allow_cross_partition_search: bool = True + allow_cross_partition_document_read: bool = False + + require_partition_filter_when_available: bool = False + + maximum_partitions: int | None = None + + + allow_bounded_scan: bool = False diff --git a/cosmos-retriever/src/cosmos_retriever/retrieval/normalization.py b/cosmos-retriever/src/cosmos_retriever/retrieval/normalization.py new file mode 100644 index 0000000..1e67196 --- /dev/null +++ b/cosmos-retriever/src/cosmos_retriever/retrieval/normalization.py @@ -0,0 +1,70 @@ +from __future__ import annotations + +from typing import Any + +from cosmos_retriever.retrieval.models import RetrievedItem + + +def _display_text( + text_fields: dict[str, str], + fallback: str, + queried: list[str] | None, + primary: str | None, +) -> str: + + names = [n for n in (queried or []) if n in text_fields] + if not names: + if primary and primary in text_fields: + return text_fields[primary] + return fallback + if primary and primary in text_fields and primary not in names: + names = [*names, primary] + if len(names) == 1: + return text_fields.get(names[0], "") or "" + return "\n\n".join(f"[{n}]\n{text_fields.get(n, '') or ''}" for n in names) + + +def normalize_rows( + rows: list[dict[str, Any]], + *, + strategy: str, + channels: list[str] | None = None, + start_rank: int = 0, + projected_aliases: dict[str, str] | None = None, + queried_text_fields: list[str] | None = None, + primary_text_field: str | None = None, +) -> list[RetrievedItem]: + aliases = projected_aliases or {} + items: list[RetrievedItem] = [] + for i, row in enumerate(rows): + metadata = { + key[len("md_") :]: value for key, value in row.items() if key.startswith("md_") + } + text_fields: dict[str, str] = {} + for key, value in row.items(): + if key.startswith("txt_") and key in aliases: + text_fields[aliases[key]] = value or "" + display = _display_text( + text_fields, + row.get("text", "") or "", + queried_text_fields, + primary_text_field, + ) + chunk_order = row.get("chunk_order") + items.append( + RetrievedItem( + item_id=str(row.get("item_id")), + document_id=(str(row["document_id"]) if row.get("document_id") is not None else None), + chunk_id=(str(row["chunk_id"]) if row.get("chunk_id") is not None else None), + chunk_order=chunk_order if isinstance(chunk_order, int) else None, + text=display, + text_fields=text_fields, + title=row.get("title"), + source=row.get("source"), + metadata=metadata, + retrieval_strategy=strategy, + retrieval_channels=list(channels or []), + rank=start_rank + i, + ) + ) + return items diff --git a/cosmos-retriever/src/cosmos_retriever/retrieval/paths.py b/cosmos-retriever/src/cosmos_retriever/retrieval/paths.py new file mode 100644 index 0000000..5481906 --- /dev/null +++ b/cosmos-retriever/src/cosmos_retriever/retrieval/paths.py @@ -0,0 +1,54 @@ + +from __future__ import annotations + +import re +from typing import Any + +from pydantic import BaseModel, ConfigDict + +from cosmos_retriever.retrieval.errors import UnsafeCosmosPath + +_ALLOWED_SEGMENT = re.compile(r"^[A-Za-z_][A-Za-z0-9_ .\-]*$") + + +class CosmosPath(BaseModel): + + model_config = ConfigDict(frozen=True) + + segments: tuple[str, ...] + + @classmethod + def parse(cls, raw: str | CosmosPath) -> CosmosPath: + + if isinstance(raw, CosmosPath): + return raw + if not isinstance(raw, str): + raise UnsafeCosmosPath(f"path must be a string, got {type(raw).__name__}") + if not raw.startswith("/"): + raise UnsafeCosmosPath(f"path must start with '/': {raw!r}") + if len(raw) < 2 or raw.endswith("/"): + raise UnsafeCosmosPath(f"path is empty or has a trailing '/': {raw!r}") + + segments = raw[1:].split("/") + for seg in segments: + if seg == "" or not _ALLOWED_SEGMENT.fullmatch(seg): + raise UnsafeCosmosPath(f"unsafe path segment {seg!r} in {raw!r}") + return cls(segments=tuple(segments)) + + def render(self, alias: str = "c") -> str: + + out = alias + for seg in self.segments: + escaped = seg.replace("\\", "\\\\").replace('"', '\\"') + out += f'["{escaped}"]' + return out + + def __str__(self) -> str: + return "/" + "/".join(self.segments) + + +def coerce_path(value: Any) -> CosmosPath: + + if isinstance(value, CosmosPath): + return value + return CosmosPath.parse(value) diff --git a/cosmos-retriever/src/cosmos_retriever/retrieval/planner.py b/cosmos-retriever/src/cosmos_retriever/retrieval/planner.py new file mode 100644 index 0000000..490cc3d --- /dev/null +++ b/cosmos-retriever/src/cosmos_retriever/retrieval/planner.py @@ -0,0 +1,115 @@ + +from __future__ import annotations + +import structlog + +from cosmos_retriever.retrieval.capabilities import RetrievalCapabilities, SupportLevel +from cosmos_retriever.retrieval.errors import UnsupportedRetrievalCapability +from cosmos_retriever.retrieval.models import GrepRequest, PartitionQueryPolicy, SearchRequest +from cosmos_retriever.retrieval.schema import CorpusSchema +from cosmos_retriever.retrieval.strategies import ( + BoundedScanStrategy, + ClientSideFusionStrategy, + FullTextGrepCandidateStrategy, + FullTextSearchStrategy, + GrepCandidateStrategy, + NativeHybridStrategy, + SearchStrategy, + VectorSearchStrategy, +) + +logger = structlog.get_logger("cosmos_retriever.retrieval.planner") + + +class RetrievalPlanner: + def __init__( + self, + schema: CorpusSchema, + capabilities: RetrievalCapabilities, + policy: PartitionQueryPolicy, + ) -> None: + self.schema = schema + self.capabilities = capabilities + self.policy = policy + + def _vector_ok(self, req: SearchRequest | None = None) -> bool: + if not self.schema.vector_fields or not self.capabilities.vector_supported: + return False + name = req.vector_field if req is not None else None + try: + field = self.schema.resolve_vector_config(name) + except Exception: + return False + cap = self.capabilities.vector_capability_for(field.path) + if cap is None or cap.support in (SupportLevel.UNSUPPORTED, SupportLevel.UNKNOWN): + return False + if cap.dimensions != field.dimensions: + logger.warning( + "embedding_dimension_mismatch", + schema_dims=field.dimensions, + capability_dims=cap.dimensions, + ) + return False + return True + + def _fts_ok(self, req: SearchRequest | None = None) -> bool: + if not self.capabilities.full_text_supported: + return False + names = req.text_fields if req is not None else None + try: + paths = self.schema.resolve_text_fields(names) + except Exception: + return False + return all(self.capabilities.has_full_text_path(p) for p in paths) + + def plan_search(self, req: SearchRequest) -> SearchStrategy: + vector_ok = self._vector_ok(req) + fts_ok = self._fts_ok(req) + mode = getattr(req, "mode", "auto") + + if mode == "vector": + if not vector_ok: + raise UnsupportedRetrievalCapability( + "vector mode requested but the selected vector field is unavailable " + "or embedding-incompatible" + ) + return VectorSearchStrategy() + if mode == "text": + if not fts_ok: + raise UnsupportedRetrievalCapability( + "text mode requested but full-text search is unavailable for the " + "selected field(s)" + ) + return FullTextSearchStrategy() + if mode == "hybrid": + if vector_ok and fts_ok: + return ( + NativeHybridStrategy() + if self.capabilities.native_hybrid_supported + else ClientSideFusionStrategy() + ) + raise UnsupportedRetrievalCapability( + "hybrid mode requested but vector and full-text are not both available " + "for the selected fields" + ) + + if self.capabilities.native_hybrid_supported and vector_ok and fts_ok: + return NativeHybridStrategy() + if vector_ok and fts_ok: + return ClientSideFusionStrategy() + if vector_ok: + return VectorSearchStrategy() + if fts_ok: + return FullTextSearchStrategy() + if self.policy.allow_bounded_scan: + return BoundedScanStrategy() + raise UnsupportedRetrievalCapability( + "no search strategy available for the configured container" + ) + + def plan_grep(self, req: GrepRequest) -> GrepCandidateStrategy: + if self.capabilities.full_text_supported: + return FullTextGrepCandidateStrategy() + raise UnsupportedRetrievalCapability( + "grep requires a full-text candidate source, which is unavailable" + ) diff --git a/cosmos-retriever/src/cosmos_retriever/retrieval/retriever.py b/cosmos-retriever/src/cosmos_retriever/retrieval/retriever.py new file mode 100644 index 0000000..c2e41d3 --- /dev/null +++ b/cosmos-retriever/src/cosmos_retriever/retrieval/retriever.py @@ -0,0 +1,81 @@ + + + +from __future__ import annotations + +import structlog + +from cosmos_retriever.retrieval.capabilities import RetrievalCapabilities +from cosmos_retriever.retrieval.compiler import CosmosQueryCompiler +from cosmos_retriever.retrieval.document_resolvers import build_document_resolver +from cosmos_retriever.retrieval.embedding import QueryEmbedder +from cosmos_retriever.retrieval.executor import CosmosExecutor +from cosmos_retriever.retrieval.models import ( + GrepRequest, + NormalizedDocument, + PartitionQueryPolicy, + ReadDocumentRequest, + RetrievedItem, + SearchRequest, +) +from cosmos_retriever.retrieval.planner import RetrievalPlanner +from cosmos_retriever.retrieval.schema import CorpusSchema +from cosmos_retriever.retrieval.strategies import RetrievalContext + +logger = structlog.get_logger("cosmos_retriever.retrieval.retriever") + + +class CorpusRetriever: + def __init__( + self, + *, + container, + schema: CorpusSchema, + capabilities: RetrievalCapabilities, + query_embedder: QueryEmbedder | None = None, + partition_policy: PartitionQueryPolicy | None = None, + ) -> None: + self.schema = schema + self.capabilities = capabilities + self.policy = partition_policy or PartitionQueryPolicy() + self._embedder = query_embedder + self._compiler = CosmosQueryCompiler(schema) + self._executor = CosmosExecutor(container) + self._planner = RetrievalPlanner(schema, capabilities, self.policy) + self._ctx = RetrievalContext( + schema=schema, + compiler=self._compiler, + executor=self._executor, + capabilities=capabilities, + policy=self.policy, + ) + self._resolver = build_document_resolver( + schema, self._compiler, self._executor, self.policy + ) + + def search(self, request: SearchRequest) -> list[RetrievedItem]: + if request.vector_field is not None: + self.schema.resolve_vector_config(request.vector_field) + if request.text_fields: + self.schema.resolve_text_fields(request.text_fields) + strategy = self._planner.plan_search(request) + if strategy.requires_embedding and request.query_vector is None: + if self._embedder is None: + from cosmos_retriever.retrieval.errors import EmbeddingProfileMismatch + + raise EmbeddingProfileMismatch( + "selected strategy requires a query embedding but no embedder is configured" + ) + request = request.model_copy( + update={"query_vector": self._embedder.embed(request.query)} + ) + return strategy.execute(request, self._ctx) + + def grep_candidates(self, request: GrepRequest) -> list[RetrievedItem]: + if request.text_field: + self.schema.resolve_text_fields([request.text_field]) + strategy = self._planner.plan_grep(request) + return strategy.candidates(request, self._ctx) + + def read_document(self, request: ReadDocumentRequest) -> NormalizedDocument: + return self._resolver.resolve(request) diff --git a/cosmos-retriever/src/cosmos_retriever/retrieval/schema.py b/cosmos-retriever/src/cosmos_retriever/retrieval/schema.py new file mode 100644 index 0000000..7a03591 --- /dev/null +++ b/cosmos-retriever/src/cosmos_retriever/retrieval/schema.py @@ -0,0 +1,160 @@ + +from __future__ import annotations + +from typing import Annotated, Literal, Protocol, runtime_checkable + +from pydantic import BaseModel, BeforeValidator, model_validator + +from cosmos_retriever.retrieval.errors import InvalidCorpusSchema, UnknownField +from cosmos_retriever.retrieval.paths import CosmosPath, coerce_path + +PathField = Annotated[CosmosPath, BeforeValidator(coerce_path)] + + +class VectorFieldConfig(BaseModel): + path: PathField + name: str | None = None + description: str | None = None + embedding_model: str | None = None + dimensions: int + distance_function: Literal["cosine", "dotproduct", "euclidean"] = "cosine" + data_type: str = "float32" + + +@runtime_checkable +class ChunkIdentityCodec(Protocol): + + def to_document_id(self, raw_id: str) -> str: ... + + +class DunderChunkCodec: + + def to_document_id(self, raw_id: str) -> str: + if isinstance(raw_id, str) and "__" in raw_id: + return raw_id.split("__", 1)[0] + return raw_id + + +class CorpusSchema(BaseModel): + item_id_path: PathField + text_paths: list[PathField] + primary_text_path: PathField + vector_fields: list[VectorFieldConfig] = [] + document_id_path: PathField | None = None + chunk_id_path: PathField | None = None + chunk_order_path: PathField | None = None + title_path: PathField | None = None + source_path: PathField | None = None + partition_key_paths: list[PathField] = [] + metadata_paths: dict[str, PathField] = {} + text_field_descriptions: dict[str, str] = {} + model_config = {"arbitrary_types_allowed": True} + + identity_codec: ChunkIdentityCodec | None = None + + @model_validator(mode="after") + def _check(self) -> CorpusSchema: + errors: list[str] = [] + primary = str(self.primary_text_path) + if primary not in {str(p) for p in self.text_paths}: + errors.append("primary_text_path must be one of text_paths") + for v in self.vector_fields: + if v.dimensions <= 0: + errors.append(f"vector field {v.path} has non-positive dimensions") + if errors: + raise InvalidCorpusSchema("; ".join(errors)) + return self + + @property + def is_item_document_mode(self) -> bool: + + return self.document_id_path is None + + @property + def partition_key_is_document_id(self) -> bool: + + if self.document_id_path is None or len(self.partition_key_paths) != 1: + return False + return str(self.partition_key_paths[0]) == str(self.document_id_path) + + @staticmethod + def _seg_name(path: CosmosPath) -> str: + return path.segments[-1] + + def text_field_map(self) -> dict[str, CosmosPath]: + + out: dict[str, CosmosPath] = {} + for p in self.text_paths: + name = self._seg_name(p) + if name in out and str(out[name]) != str(p): + name = str(p) + out[name] = p + return out + + def vector_field_map(self) -> dict[str, CosmosPath]: + + out: dict[str, CosmosPath] = {} + for i, vf in enumerate(self.vector_fields): + name = vf.name or self._seg_name(vf.path) + if name in out: + name = f"{name}_{i}" + out[name] = vf.path + return out + + def primary_text_field_name(self) -> str: + for name, p in self.text_field_map().items(): + if str(p) == str(self.primary_text_path): + return name + return self._seg_name(self.primary_text_path) + + def resolve_text_fields(self, names: list[str] | None) -> list[CosmosPath]: + if not names: + return [self.primary_text_path] + m = self.text_field_map() + paths: list[CosmosPath] = [] + for n in names: + if n not in m: + raise UnknownField( + f"unknown text field {n!r}; available: {sorted(m)}" + ) + paths.append(m[n]) + return paths + + def resolve_vector_config(self, name: str | None) -> VectorFieldConfig: + if not self.vector_fields: + raise UnknownField("no vector fields are configured") + if name is None: + return self.vector_fields[0] + for i, vf in enumerate(self.vector_fields): + vname = vf.name or self._seg_name(vf.path) + if vname == name or f"{vname}_{i}" == name: + return vf + available = sorted(self.vector_field_map()) + raise UnknownField(f"unknown vector field {name!r}; available: {available}") + + def resolve_vector_field(self, name: str | None) -> CosmosPath: + return self.resolve_vector_config(name).path + + def agent_field_summary(self) -> str: + + tm = self.text_field_map() + vm = self.vector_field_map() + lines: list[str] = [] + tparts = [] + for n, p in tm.items(): + d = self.text_field_descriptions.get(n) or self.text_field_descriptions.get(str(p)) + tparts.append(f"'{n}'" + (f" — {d}" if d else "")) + lines.append("Text fields (keyword / BM25): " + ", ".join(tparts)) + if vm: + vparts = [] + for n, p in vm.items(): + cfg = next((v for v in self.vector_fields if str(v.path) == str(p)), None) + d = cfg.description if cfg else None + vparts.append(f"'{n}'" + (f" — {d}" if d else "")) + lines.append("Vector fields (semantic): " + ", ".join(vparts)) + default_v = next(iter(vm), None) + default = f"hybrid over text='{self.primary_text_field_name()}'" + if default_v: + default += f" + vector='{default_v}'" + lines.append(f"Default when unspecified: {default}.") + return "\n".join(lines) diff --git a/cosmos-retriever/src/cosmos_retriever/retrieval/strategies.py b/cosmos-retriever/src/cosmos_retriever/retrieval/strategies.py new file mode 100644 index 0000000..362cb65 --- /dev/null +++ b/cosmos-retriever/src/cosmos_retriever/retrieval/strategies.py @@ -0,0 +1,228 @@ +from __future__ import annotations + +from abc import ABC, abstractmethod +from dataclasses import dataclass + +from cosmos_retriever.retrieval.capabilities import RetrievalCapabilities +from cosmos_retriever.retrieval.compiler import CosmosQueryCompiler +from cosmos_retriever.retrieval.errors import ( + CrossPartitionQueryDisabled, + UnboundedScanRejected, +) +from cosmos_retriever.retrieval.executor import CosmosExecutor +from cosmos_retriever.retrieval.models import ( + GrepRequest, + PartitionQueryPolicy, + RetrievedItem, + SearchRequest, +) +from cosmos_retriever.retrieval.normalization import normalize_rows +from cosmos_retriever.retrieval.schema import CorpusSchema + + +@dataclass +class RetrievalContext: + schema: CorpusSchema + compiler: CosmosQueryCompiler + executor: CosmosExecutor + capabilities: RetrievalCapabilities + policy: PartitionQueryPolicy + + +def _resolve_cross_partition(req_partition_key, policy: PartitionQueryPolicy) -> bool: + + if req_partition_key is not None: + return False + if not policy.allow_cross_partition_search: + raise CrossPartitionQueryDisabled( + "search requires a partition key or cross-partition permission" + ) + return True + + + +class SearchStrategy(ABC): + name: str = "" + requires_embedding: bool = False + + @abstractmethod + def execute(self, req: SearchRequest, ctx: RetrievalContext) -> list[RetrievedItem]: ... + + +class NativeHybridStrategy(SearchStrategy): + name = "native_hybrid" + requires_embedding = True + + def execute(self, req: SearchRequest, ctx: RetrievalContext) -> list[RetrievedItem]: + vector_path = ctx.schema.resolve_vector_field(req.vector_field) + text_paths = ctx.schema.resolve_text_fields(req.text_fields) + cross = _resolve_cross_partition(req.partition_key, ctx.policy) + compiled = ctx.compiler.compile_hybrid( + query=req.query, + query_vector=req.query_vector or [], + limit=req.limit, + ignored_item_ids=req.ignored_item_ids, + filters=req.filters, + partition_key=req.partition_key, + cross_partition=cross, + vector_path=vector_path, + text_paths=text_paths, + ) + rows = ctx.executor.run(compiled) + return normalize_rows( + rows, + strategy=self.name, + channels=["vector", "full_text"], + projected_aliases=compiled.projected_aliases, + queried_text_fields=req.text_fields, + primary_text_field=ctx.schema.primary_text_field_name(), + ) + + +class VectorSearchStrategy(SearchStrategy): + name = "vector" + requires_embedding = True + + def execute(self, req: SearchRequest, ctx: RetrievalContext) -> list[RetrievedItem]: + vector_path = ctx.schema.resolve_vector_field(req.vector_field) + cross = _resolve_cross_partition(req.partition_key, ctx.policy) + compiled = ctx.compiler.compile_vector( + query_vector=req.query_vector or [], + limit=req.limit, + ignored_item_ids=req.ignored_item_ids, + filters=req.filters, + partition_key=req.partition_key, + cross_partition=cross, + vector_path=vector_path, + ) + rows = ctx.executor.run(compiled) + return normalize_rows( + rows, + strategy=self.name, + channels=["vector"], + projected_aliases=compiled.projected_aliases, + primary_text_field=ctx.schema.primary_text_field_name(), + ) + + +class FullTextSearchStrategy(SearchStrategy): + name = "full_text" + requires_embedding = False + + def execute(self, req: SearchRequest, ctx: RetrievalContext) -> list[RetrievedItem]: + text_paths = ctx.schema.resolve_text_fields(req.text_fields) + cross = _resolve_cross_partition(req.partition_key, ctx.policy) + compiled = ctx.compiler.compile_full_text( + query=req.query, + limit=req.limit, + ignored_item_ids=req.ignored_item_ids, + filters=req.filters, + partition_key=req.partition_key, + cross_partition=cross, + text_paths=text_paths, + ) + rows = ctx.executor.run(compiled) + return normalize_rows( + rows, + strategy=self.name, + channels=["full_text"], + projected_aliases=compiled.projected_aliases, + queried_text_fields=req.text_fields, + primary_text_field=ctx.schema.primary_text_field_name(), + ) + + +class ClientSideFusionStrategy(SearchStrategy): + + name = "client_fusion" + requires_embedding = True + _RRF_K = 60 + + def execute(self, req: SearchRequest, ctx: RetrievalContext) -> list[RetrievedItem]: + vector_hits = VectorSearchStrategy().execute(req, ctx) + fts_hits = FullTextSearchStrategy().execute(req, ctx) + scores: dict[str, float] = {} + channels: dict[str, list[str]] = {} + item_by_id: dict[str, RetrievedItem] = {} + for hits, channel in ((vector_hits, "vector"), (fts_hits, "full_text")): + for rank, item in enumerate(hits): + scores[item.item_id] = scores.get(item.item_id, 0.0) + 1.0 / (self._RRF_K + rank) + channels.setdefault(item.item_id, []).append(channel) + item_by_id.setdefault(item.item_id, item) + ranked_ids = sorted(scores, key=lambda i: scores[i], reverse=True)[: req.limit] + out: list[RetrievedItem] = [] + for rank, item_id in enumerate(ranked_ids): + base = item_by_id[item_id] + out.append( + base.model_copy( + update={ + "rank": rank, + "retrieval_strategy": self.name, + "retrieval_channels": channels[item_id], + "raw_scores": {"rrf": scores[item_id]}, + } + ) + ) + return out + + +class BoundedScanStrategy(SearchStrategy): + + + name = "bounded_scan" + requires_embedding = False + + def execute(self, req: SearchRequest, ctx: RetrievalContext) -> list[RetrievedItem]: + if not ctx.policy.allow_bounded_scan: + raise UnboundedScanRejected("bounded scan is not enabled") + cross = _resolve_cross_partition(req.partition_key, ctx.policy) + compiled = ctx.compiler.compile_structured( + limit=req.limit, + filters=req.filters, + ignored_item_ids=req.ignored_item_ids, + partition_key=req.partition_key, + cross_partition=cross, + ) + compiled.warnings.append("bounded scan active") + rows = ctx.executor.run(compiled) + return normalize_rows(rows, strategy=self.name) + + + +class GrepCandidateStrategy(ABC): + @abstractmethod + + def candidates(self, req: GrepRequest, ctx: RetrievalContext) -> list[RetrievedItem]: ... + + +class FullTextGrepCandidateStrategy(GrepCandidateStrategy): + + + def candidates(self, req: GrepRequest, ctx: RetrievalContext) -> list[RetrievedItem]: + from cosmos_retriever.retrieval.expressions import tokenize_for_fts + + if not tokenize_for_fts(req.pattern): + return [] + text_paths = ctx.schema.resolve_text_fields( + [req.text_field] if req.text_field else None + ) + cross = _resolve_cross_partition(req.partition_key, ctx.policy) + compiled = ctx.compiler.compile_full_text( + query=req.pattern, + limit=req.candidate_limit, + ignored_item_ids=[], + filters=req.filters, + partition_key=req.partition_key, + cross_partition=cross, + text_paths=text_paths, + strategy="grep_full_text", + ) + rows = ctx.executor.run(compiled) + return normalize_rows( + rows, + strategy="grep_full_text", + channels=["full_text"], + projected_aliases=compiled.projected_aliases, + queried_text_fields=[req.text_field] if req.text_field else None, + primary_text_field=ctx.schema.primary_text_field_name(), + ) diff --git a/cosmos-retriever/src/cosmos_retriever/retriever.py b/cosmos-retriever/src/cosmos_retriever/retriever.py new file mode 100644 index 0000000..dde101d --- /dev/null +++ b/cosmos-retriever/src/cosmos_retriever/retriever.py @@ -0,0 +1,292 @@ + +from __future__ import annotations + +import time +from dataclasses import dataclass, field + +import structlog +import tiktoken + +from cosmos_retriever.config import CorpusConfig, RetrieverSettings, get_settings +from cosmos_retriever.rerank import BasetenReranker, Reranker, VLLMReranker +from cosmos_retriever.tools import ToolSet + +logger = structlog.get_logger("cosmos_retriever.retriever") + + +@dataclass +class RetrievedDocument: + + id: str + text: str = "" + justification: str | None = None + rank: int | None = None + + +@dataclass +class RetrievalResult: + + query: str + documents: list[RetrievedDocument] + num_turns: int + final_text: str = "" + pool_doc_ids: list[str] = field(default_factory=list) + elapsed_s: float = 0.0 + usage: dict[str, int] = field(default_factory=dict) + metadata: dict[str, str | int | float] = field(default_factory=dict) + trajectory: dict[str, object] = field(default_factory=dict) + + +class CosmosRetriever: + + def __init__( + self, + settings: RetrieverSettings | None = None, + *, + corpus_name: str | None = None, + reranker: Reranker | None = None, + ) -> None: + self.settings = settings or get_settings() + self.corpus: CorpusConfig = self.settings.resolve_corpus(corpus_name) + + self._tiktoken = tiktoken.get_encoding("o200k_harmony") + self._reranker = reranker or self._build_default_reranker() + + cosmos_db = self.settings.build_cosmos_database(self.corpus) + openai_client = self.settings.build_openai_client(self.corpus) + self._use_chat = self.settings.use_chat_backend + self._use_responses = self.settings.use_responses_backend + self._use_anthropic = self.settings.use_anthropic_backend + + self.toolset: ToolSet = ToolSet.build( + cosmos_database=cosmos_db, + cosmos_container_name=self.corpus.container, + openai_client=openai_client, + openai_embedding_model=self.corpus.embed_model, + embed_query_instruction=self.corpus.embed_query_instruction, + reranker=self._reranker, + token_counter=self._text_token_counter, + search_display_limit=self.settings.cosmos_retriever_search_display_limit, + ) + + self._chat_client = ( + self.settings.build_chat_client() + if self.settings.use_generic_llm_backend + else None + ) + self._chat_model: str | None = self.settings.chat_model + + logger.info( + "cosmos_retriever_initialized", + inference_backend=self.settings.inference_backend, + chat_base_url=self.settings.chat_base_url, + chat_model=self._chat_model, + cosmos_account=self.corpus.account_uri, + cosmos_db=self.corpus.database, + cosmos_container=self.corpus.container, + embed_base_url=self.corpus.embed_base_url, + embed_model=self.corpus.embed_model, + embed_query_instruction=self.corpus.embed_query_instruction, + reranker=type(self._reranker).__name__ if self._reranker is not None else None, + ) + + def search( + self, + query: str, + *, + max_documents: int = 20, + max_turns: int | None = None, + threshold_budget: int | None = None, + token_budget: int | None = None, + ) -> RetrievalResult: + + if not query or not query.strip(): + raise ValueError("query must be a non-empty string") + + return self._search_sync( + query, + max_documents, + max_turns or self.settings.cosmos_retriever_max_turns, + threshold_budget or self.settings.cosmos_retriever_threshold_budget, + token_budget or self.settings.cosmos_retriever_token_budget, + ) + + def _search_sync( + self, + query: str, + max_documents: int, + max_turns: int, + threshold_budget: int, + token_budget: int, + ) -> RetrievalResult: + + if self._use_chat: + return self._search_chat(query, max_documents) + if self._use_anthropic: + return self._search_anthropic(query, max_documents) + return self._search_responses(query, max_documents) + + def _search_chat(self, query: str, max_documents: int) -> RetrievalResult: + + from cosmos_retriever.inference.agent_loop import ( + run_chat_search, + ) + + if self._chat_client is None or self._chat_model is None: + raise RuntimeError("chat backend selected but chat client/model not initialised") + + start = time.perf_counter() + chat_result = run_chat_search( + toolset=self.toolset, + client=self._chat_client, + model=self._chat_model, + query=query, + max_documents=max_documents, + max_turns=self.settings.chat_max_turns, + temperature=self.settings.chat_temperature, + max_tokens=self.settings.chat_max_tokens, + ) + elapsed = time.perf_counter() - start + + documents = [ + RetrievedDocument(id=d.id, text=d.text, justification=d.justification, rank=d.rank) + for d in chat_result.documents + ] + result = RetrievalResult( + query=query, + documents=documents, + num_turns=chat_result.num_turns, + final_text=chat_result.final_text, + elapsed_s=round(elapsed, 3), + usage=chat_result.usage, + metadata=chat_result.metadata, + ) + logger.info( + "search_complete", + query=query[:200], + backend="openai_chat", + num_documents=len(result.documents), + num_turns=result.num_turns, + elapsed_s=result.elapsed_s, + ) + return result + + def _search_responses(self, query: str, max_documents: int) -> RetrievalResult: + + from cosmos_retriever.inference.agent_loop import ( + run_responses_search, + ) + + if self._chat_client is None or self._chat_model is None: + raise RuntimeError("responses backend selected but chat client/model not initialised") + + start = time.perf_counter() + chat_result = run_responses_search( + toolset=self.toolset, + client=self._chat_client, + model=self._chat_model, + query=query, + max_documents=max_documents, + max_turns=self.settings.chat_max_turns, + max_tokens=self.settings.chat_max_tokens, + reasoning_effort=self.settings.chat_reasoning_effort, + ) + elapsed = time.perf_counter() - start + + documents = [ + RetrievedDocument(id=d.id, text=d.text, justification=d.justification, rank=d.rank) + for d in chat_result.documents + ] + result = RetrievalResult( + query=query, + documents=documents, + num_turns=chat_result.num_turns, + final_text=chat_result.final_text, + elapsed_s=round(elapsed, 3), + pool_doc_ids=chat_result.pool_doc_ids, + usage=chat_result.usage, + trajectory=chat_result.trajectory, + metadata=chat_result.metadata, + ) + logger.info( + "search_complete", + query=query[:200], + backend="openai_responses", + num_documents=len(result.documents), + num_turns=result.num_turns, + elapsed_s=result.elapsed_s, + ) + return result + + def _search_anthropic(self, query: str, max_documents: int) -> RetrievalResult: + from cosmos_retriever.inference.agent_loop import ( + run_anthropic_search, + ) + + if ( + not self.settings.chat_base_url + or self.settings.chat_api_key is None + or self._chat_model is None + ): + raise RuntimeError( + "anthropic backend selected but CHAT_BASE_URL / CHAT_API_KEY / CHAT_MODEL not set" + ) + + start = time.perf_counter() + chat_result = run_anthropic_search( + toolset=self.toolset, + base_url=self.settings.chat_base_url, + api_key=self.settings.chat_api_key.get_secret_value(), + model=self._chat_model, + query=query, + max_documents=max_documents, + max_turns=self.settings.chat_max_turns, + max_tokens=self.settings.chat_max_tokens, + anthropic_version=self.settings.anthropic_version, + auth_header=self.settings.anthropic_auth_header, + ) + elapsed = time.perf_counter() - start + + documents = [ + RetrievedDocument(id=d.id, text=d.text, justification=d.justification, rank=d.rank) + for d in chat_result.documents + ] + result = RetrievalResult( + query=query, + documents=documents, + num_turns=chat_result.num_turns, + final_text=chat_result.final_text, + elapsed_s=round(elapsed, 3), + pool_doc_ids=chat_result.pool_doc_ids, + usage=chat_result.usage, + trajectory=chat_result.trajectory, + metadata=chat_result.metadata, + ) + logger.info( + "search_complete", + query=query[:200], + backend="anthropic_messages", + num_documents=len(result.documents), + num_turns=result.num_turns, + elapsed_s=result.elapsed_s, + ) + return result + + def _build_default_reranker(self) -> Reranker | None: + if self.settings.baseten_api_key and self.settings.baseten_model_url: + return BasetenReranker( + client=self.settings.get_baseten_client(), + token_counter=self._text_token_counter, + ) + if self.settings.vllm_reranker_url: + return VLLMReranker( + base_url=self.settings.vllm_reranker_url, + token_counter=self._text_token_counter, + ) + return None + + def _text_token_counter(self, text: str) -> int: + return len(self._tiktoken.encode(text)) + + +__all__ = ["CosmosRetriever", "RetrievalResult", "RetrievedDocument"] diff --git a/cosmos-retriever/src/cosmos_retriever/server.py b/cosmos-retriever/src/cosmos_retriever/server.py new file mode 100644 index 0000000..416c6de --- /dev/null +++ b/cosmos-retriever/src/cosmos_retriever/server.py @@ -0,0 +1,127 @@ + +from __future__ import annotations + +import asyncio +from collections import defaultdict +from contextlib import asynccontextmanager +from dataclasses import asdict +from typing import TYPE_CHECKING + +import anyio +import structlog +from fastapi import FastAPI +from fastapi.responses import JSONResponse +from pydantic import BaseModel, Field + +from cosmos_retriever.config import RetrieverSettings, get_settings +from cosmos_retriever.retriever import CosmosRetriever + +if TYPE_CHECKING: + from collections.abc import AsyncIterator + +logger = structlog.get_logger("cosmos_retriever.server") + + +class SearchRequest(BaseModel): + + query: str = Field(..., min_length=1, description="Natural-language information need.") + max_documents: int = Field( + default=20, + ge=1, + le=30, + alias="maxDocuments", + description="Cap on the number of curated documents to return.", + ) + database: str | None = Field( + default=None, + description="Override Cosmos database name (else COSMOS_DATABASE env var).", + ) + container: str | None = Field( + default=None, + description="Override Cosmos corpus container name (else COSMOS_CORPUS_CONTAINER).", + ) + + model_config = {"populate_by_name": True} + + +class _RetrieverPool: + + def __init__(self, settings: RetrieverSettings) -> None: + self._settings = settings + self._retrievers: dict[tuple[str | None, str | None], CosmosRetriever] = {} + self._locks: dict[tuple[str | None, str | None], asyncio.Lock] = defaultdict(asyncio.Lock) + self._build_lock = asyncio.Lock() + + async def get( + self, database: str | None, container: str | None + ) -> tuple[CosmosRetriever, asyncio.Lock]: + key = (database, container) + retriever = self._retrievers.get(key) + if retriever is None: + async with self._build_lock: + retriever = self._retrievers.get(key) + if retriever is None: + retriever = await anyio.to_thread.run_sync( + lambda: self._build(database, container) + ) + self._retrievers[key] = retriever + return retriever, self._locks[key] + + def _build(self, database: str | None, container: str | None) -> CosmosRetriever: + settings = self._settings.model_copy(deep=True) + if database: + settings.cosmos_database = database + return CosmosRetriever(settings=settings, corpus_name=container) + + +def create_app(settings: RetrieverSettings | None = None) -> FastAPI: + + resolved = settings or get_settings() + + @asynccontextmanager + async def lifespan(app: FastAPI) -> AsyncIterator[None]: + app.state.pool = _RetrieverPool(resolved) + logger.info( + "cosmos_retriever_server_started", + host=resolved.host, + port=resolved.port, + default_container=resolved.cosmos_corpus_container, + ) + yield + + app = FastAPI( + title="Cosmos Retriever", + version="0.1.0", + description="HTTP service running the multi-turn Cosmos search agent.", + lifespan=lifespan, + ) + + @app.get("/health") + async def health() -> dict[str, str]: + return {"status": "ok"} + + @app.post("/search") + async def search(request: SearchRequest) -> JSONResponse: + pool: _RetrieverPool = app.state.pool + try: + retriever, lock = await pool.get(request.database, request.container) + async with lock: + result = await anyio.to_thread.run_sync( + lambda: retriever.search( + request.query, max_documents=request.max_documents + ) + ) + except Exception as exc: + logger.error( + "search_failed", + query=request.query[:200], + error=str(exc), + error_type=type(exc).__name__, + ) + return JSONResponse( + status_code=500, + content={"error": str(exc), "type": type(exc).__name__}, + ) + return JSONResponse(content=asdict(result)) + + return app diff --git a/cosmos-retriever/src/cosmos_retriever/tools.py b/cosmos-retriever/src/cosmos_retriever/tools.py new file mode 100644 index 0000000..ec006e2 --- /dev/null +++ b/cosmos-retriever/src/cosmos_retriever/tools.py @@ -0,0 +1,685 @@ + +from __future__ import annotations + +import json +import re +from abc import ABC, abstractmethod +from collections.abc import Callable +from typing import Any, TypeAlias + +import openai +import structlog +from azure.cosmos import DatabaseProxy +from pydantic import BaseModel, Field + +from cosmos_retriever.rerank import Reranker +from cosmos_retriever.retrieval import ( + CorpusRetriever, + GrepRequest, + QueryEmbedder, + ReadDocumentRequest, + SearchRequest, + build_default_retriever, +) +from cosmos_retriever.retrieval.errors import UnknownField, UnsupportedRetrievalCapability +from cosmos_retriever.retrieval.executor import COSMOS_QUERY_MAX_CONCURRENCY +from cosmos_retriever.retrieval.formatting import DOC_TRUNCATION, format_result_blocks +from cosmos_retriever.retrieval.schema import CorpusSchema +from cosmos_retriever.utils import ProviderFormat + +logger = structlog.get_logger("cosmos_retriever.tools") + + + + +class ToolSchema(BaseModel): + + name: str + description: str + parameters: dict[str, Any] + required: list[str] = Field(default_factory=list) + + def _to_openai_format(self) -> dict[str, Any]: + return { + "type": "function", + "name": self.name, + "description": self.description, + "parameters": { + "type": "object", + "properties": self.parameters, + "required": self.required, + }, + } + + def _to_openai_harmony_format(self) -> dict[str, Any]: + return { + "type": "function", + "function": { + "name": self.name, + "description": self.description, + "parameters": { + "type": "object", + "properties": self.parameters, + "required": self.required, + }, + }, + } + + def _to_anthropic_format(self) -> dict[str, Any]: + return { + "name": self.name, + "description": self.description, + "input_schema": { + "type": "object", + "properties": self.parameters, + "required": self.required, + }, + } + + def to_provider_format(self, provider: ProviderFormat) -> dict[str, Any]: + if provider is ProviderFormat.OPENAI: + return self._to_openai_format() + if provider is ProviderFormat.OPENAI_HARMONY: + return self._to_openai_harmony_format() + if provider is ProviderFormat.ANTHROPIC: + return self._to_anthropic_format() + raise ValueError(f"Unsupported provider format: {provider}") + + + +SEARCH_CORPUS_SCHEMA = ToolSchema( + name="search_corpus", + description=( + "Searches the corpus for relevant documents based on the input query. " + "Returns a section of the document that is relevant to the query." + ), + parameters={ + "query": { + "type": "string", + "description": "The search query to find relevant documents in the corpus.", + } + }, + required=["query"], +) + +READ_DOCUMENT_SCHEMA = ToolSchema( + name="read_document", + description="Reads the content of a document based on its ID.", + parameters={ + "doc_id": { + "type": "string", + "description": "The unique identifier of the document to read.", + } + }, + required=["doc_id"], +) + +GREP_CORPUS_SCHEMA = ToolSchema( + name="grep_corpus", + description="Performs a regex search on the corpus to find documents matching the query.", + parameters={ + "pattern": { + "type": "string", + "description": "The regex query to search for in the corpus.", + } + }, + required=["pattern"], +) + +MULTI_TOOL_USE_SCHEMA = ToolSchema( + name="multi_tool_use", + description="Allows the agent to use multiple tools in parallel to gather information.", + parameters={ + "tool_calls": { + "type": "array", + "description": "List of tool calls to execute in parallel.", + "items": { + "type": "object", + "properties": { + "tool_name": {"type": "string"}, + "parameters": {"type": "object"}, + }, + "required": ["tool_name", "parameters"], + }, + } + }, + required=["tool_calls"], +) + +PRUNE_CHUNKS_SCHEMA = ToolSchema( + name="prune_chunks", + description=( + "Prunes the chunks by id that are not relevant to the main question from the " + "history of the conversation." + ), + parameters={"chunk_ids": {"type": "array", "items": {"type": "string"}}}, + required=["chunk_ids"], +) + + + + +class ToolCallMetadata(BaseModel): + pass + + +class Tool(ABC, BaseModel): + + tool_schema: ToolSchema + + @abstractmethod + def __call__( + self, + params: dict[Any, Any], + overrides: dict[Any, Any] | None = None, + ) -> tuple[str, ToolCallMetadata | None]: + pass + + def get_format(self, provider: ProviderFormat) -> dict[str, Any]: + return self.tool_schema.to_provider_format(provider) + + def __repr__(self) -> str: + return f"Tool(name={self.tool_schema.name!r})" + + +class SerializedTool(Tool): + + def __call__( + self, + params: dict[Any, Any], + overrides: dict[Any, Any] | None = None, + ) -> tuple[str, ToolCallMetadata | None]: + raise NotImplementedError("SerializedTool is a placeholder and cannot be executed.") + + + + +def _search_schema_for(schema: CorpusSchema) -> ToolSchema: + + text_names = list(schema.text_field_map()) + vector_names = list(schema.vector_field_map()) + params: dict[str, Any] = { + "query": { + "type": "string", + "description": "The search query to find relevant documents in the corpus.", + } + } + if len(text_names) > 1: + params["fields"] = { + "type": "array", + "items": {"type": "string", "enum": text_names}, + "description": ( + "Optional. Text field(s) to keyword-match against " + f"(available: {text_names}). Defaults to the primary text field." + ), + } + if len(vector_names) > 1: + params["vector_field"] = { + "type": "string", + "enum": vector_names, + "description": ( + "Optional. Vector field for semantic similarity " + f"(available: {vector_names}). Defaults to the first vector field." + ), + } + if text_names and vector_names: + params["mode"] = { + "type": "string", + "enum": ["auto", "hybrid", "vector", "text"], + "description": ( + "Optional retrieval method: 'hybrid' (semantic + keyword), " + "'vector' (semantic only), 'text' (keyword only), or 'auto' (default)." + ), + } + desc = ( + "Searches the corpus for relevant documents based on the input query. " + "Returns a section of the document that is relevant to the query.\n\n" + "Queryable schema:\n" + schema.agent_field_summary() + ) + return ToolSchema( + name="search_corpus", description=desc, parameters=params, required=["query"] + ) + + +def _grep_schema_for(schema: CorpusSchema) -> ToolSchema: + + text_names = list(schema.text_field_map()) + params: dict[str, Any] = { + "pattern": { + "type": "string", + "description": "The regex query to search for in the corpus.", + } + } + if len(text_names) > 1: + params["field"] = { + "type": "string", + "enum": text_names, + "description": ( + "Optional. Text field to search " + f"(available: {text_names}). Defaults to the primary text field." + ), + } + desc = ( + "Performs a regex search on the corpus to find documents matching the query.\n\n" + "Queryable text fields: " + ", ".join(f"'{n}'" for n in text_names) + ) + return ToolSchema( + name="grep_corpus", description=desc, parameters=params, required=["pattern"] + ) + + +class SearchCorpusToolCallMetadata(ToolCallMetadata): + + returned_chunk_ids: list[str] + pre_rerank_chunk_ids: list[str] | None = None + + +class SearchCorpusTool(Tool): + + tool_schema: ToolSchema + _retriever: CorpusRetriever + _reranker: Reranker | None + _search_limit: int + _display_limit: int + + def __init__( + self, + retriever: CorpusRetriever, + reranker: Reranker | None = None, + search_limit: int = 50, + display_limit: int = 10, + ) -> None: + super().__init__(tool_schema=_search_schema_for(retriever.schema)) + self._retriever = retriever + self._reranker = reranker + self._search_limit = search_limit + self._display_limit = display_limit + + def __call__( + self, + params: dict[Any, Any], + overrides: dict[Any, Any] | None = None, + ) -> tuple[str, SearchCorpusToolCallMetadata | None]: + log = logger.bind(tool=self.tool_schema.name) + if not isinstance(params, dict) or "query" not in params: + log.error("invalid_params", params_type=type(params).__name__) + raise ValueError(f"Invalid params type: {type(params)}") + + query = params["query"] + ignore_ids: list[str] = [] + if overrides is not None and "ignore_ids" in overrides: + ignore_ids = overrides["ignore_ids"] + + fields = params.get("fields") + if isinstance(fields, str): + fields = [fields] + vector_field = params.get("vector_field") + mode = params.get("mode") or "auto" + if mode not in ("auto", "hybrid", "vector", "text"): + mode = "auto" + log.info( + "search_corpus", + query=query, + ignore_ids=len(ignore_ids), + fields=fields, + vector_field=vector_field, + mode=mode, + ) + + request = SearchRequest( + query=query, + limit=self._search_limit, + ignored_item_ids=ignore_ids, + text_fields=fields, + vector_field=vector_field, + mode=mode, + ) + try: + items = self._retriever.search(request) + except (UnknownField, UnsupportedRetrievalCapability) as exc: + log.warning("search_field_error", error=str(exc)) + return ( + f"Search field/mode error: {exc}", + SearchCorpusToolCallMetadata(returned_chunk_ids=[]), + ) + ids = [it.item_id for it in items] + documents = [it.text for it in items] + + max_tokens_override = ( + overrides.get("max_tokens") if overrides and "max_tokens" in overrides else None + ) + + token_counts: list[int | None] = [None] * len(ids) + if self._reranker is not None and ids: + rerank_results = self._reranker(query, documents, max_tokens=max_tokens_override) + ids = [ids[r.original_index] for r in rerank_results] + documents = [r.document for r in rerank_results] + token_counts = [r.tokens for r in rerank_results] + log.info("reranked_results", num_results=len(ids)) + + triples = list(zip(ids, documents, token_counts, strict=True))[: self._display_limit] + text = format_result_blocks(triples) + returned = [t[0] for t in triples] + return text, SearchCorpusToolCallMetadata(returned_chunk_ids=returned) + + +class GrepCorpusToolCallMetadata(ToolCallMetadata): + + returned_chunk_ids: list[str] + + +class GrepCorpusTool(Tool): + + tool_schema: ToolSchema + _retriever: CorpusRetriever + _token_counter: Callable[[str], int] | None + + def __init__( + self, + retriever: CorpusRetriever, + token_counter: Callable[[str], int] | None = None, + ) -> None: + super().__init__(tool_schema=_grep_schema_for(retriever.schema)) + self._retriever = retriever + self._token_counter = token_counter + + def __call__( + self, + params: dict[Any, Any], + overrides: dict[Any, Any] | None = None, + ) -> tuple[str, ToolCallMetadata | None]: + log = logger.bind(tool=self.tool_schema.name) + if not isinstance(params, dict) or "pattern" not in params: + log.error("invalid_params", params_type=type(params).__name__) + raise ValueError(f"Invalid params type: {type(params)}") + + pattern = params["pattern"] + field = params.get("field") + log.info("grep_corpus", pattern=pattern, field=field) + + try: + candidates = self._retriever.grep_candidates( + GrepRequest( + pattern=pattern, candidate_limit=50, result_limit=5, text_field=field + ) + ) + except (UnknownField, UnsupportedRetrievalCapability) as exc: + log.warning("grep_field_error", error=str(exc)) + return ( + f"Grep field error: {exc}", + GrepCorpusToolCallMetadata(returned_chunk_ids=[]), + ) + if not candidates: + return "No results found", GrepCorpusToolCallMetadata(returned_chunk_ids=[]) + + try: + regex = re.compile(pattern, re.IGNORECASE) + matched = [it for it in candidates if regex.search(it.text)][:5] + except re.error: + matched = candidates[:5] + + ids = [it.item_id for it in matched] + documents = [it.text for it in matched] + token_counts: list[int | None] = ( + [self._token_counter(doc) for doc in documents] + if self._token_counter is not None + else [None] * len(documents) + ) + + triples = list(zip(ids, documents, token_counts, strict=True)) + text = format_result_blocks(triples) + return text, GrepCorpusToolCallMetadata(returned_chunk_ids=ids) + + +class ReadDocumentTool(Tool): + + tool_schema: ToolSchema + _retriever: CorpusRetriever + _reranker: Reranker | None + _token_counter: Callable[[str], int] | None + _max_tokens: int | None + + def __init__( + self, + retriever: CorpusRetriever, + reranker: Reranker | None = None, + token_counter: Callable[[str], int] | None = None, + max_tokens: int | None = None, + ) -> None: + if max_tokens is not None and token_counter is None: + raise ValueError("token_counter is required when max_tokens is specified") + super().__init__(tool_schema=READ_DOCUMENT_SCHEMA) + self._retriever = retriever + self._reranker = reranker + self._token_counter = token_counter + self._max_tokens = max_tokens + + def __call__( + self, + params: dict[Any, Any], + overrides: dict[Any, Any] | None = None, + ) -> tuple[str, ToolCallMetadata | None]: + log = logger.bind(tool=self.tool_schema.name) + if not isinstance(params, dict) or ("doc_id" not in params and "id" not in params): + log.error("invalid_params", params_type=type(params).__name__) + raise ValueError(f"Invalid params type: {type(params)}") + + doc_id = params.get("doc_id") or params.get("id") + log.info("read_document", doc_id=doc_id) + + query = overrides.get("query") if overrides else None + max_tokens = ( + overrides.get("max_tokens") if overrides and "max_tokens" in overrides else None + ) or self._max_tokens + + document = self._retriever.read_document( + ReadDocumentRequest(document_id=doc_id, query=query) + ) + documents = document.chunk_texts + assembled = document.assembled + + if self._reranker is not None and query is not None and max_tokens is not None: + rerank_results = self._reranker(query, documents, max_tokens=max_tokens) + kept_indices = {r.original_index for r in rerank_results} + kept_docs = [documents[i] for i in range(len(documents)) if i in kept_indices] + assembled = "".join(kept_docs) + log.info("reranked_and_filtered", original=len(documents), kept=len(kept_docs)) + elif self._token_counter is not None and max_tokens is not None: + total_tokens = self._token_counter(assembled) + if total_tokens > max_tokens: + truncated: list[str] = [] + running = 0 + for doc in documents: + n = self._token_counter(doc) + if running + n > max_tokens: + break + truncated.append(doc) + running += n + assembled = "".join(truncated) + log.info("truncated_by_tokens", original=len(documents), kept=len(truncated)) + + if self._token_counter is not None: + token_count = self._token_counter(assembled) + return f"# Document ({token_count} tokens)\n{assembled}", None + return assembled, None + + +class PruneChunksTool(Tool): + + tool_schema: ToolSchema + + def __init__(self) -> None: + super().__init__(tool_schema=PRUNE_CHUNKS_SCHEMA) + + def __call__( + self, + params: dict[Any, Any], + overrides: dict[Any, Any] | None = None, + ) -> tuple[str, ToolCallMetadata | None]: + log = logger.bind(tool=self.tool_schema.name) + if not isinstance(params, dict) or "chunk_ids" not in params: + log.error("invalid_params", params_type=type(params).__name__) + raise ValueError(f"Invalid params type: {type(params)}") + log.info("prune_chunks", chunk_ids=len(params["chunk_ids"])) + return "Pruned", None + + +_ToolSetT: TypeAlias = "ToolSet" + + +class MultiToolUseTool(Tool): + + tool_schema: ToolSchema + toolset: _ToolSetT + + def __init__(self, toolset: ToolSet) -> None: + super().__init__(tool_schema=MULTI_TOOL_USE_SCHEMA, toolset=toolset) + + def __call__( + self, + params: dict[Any, Any], + overrides: dict[Any, Any] | None = None, + ) -> tuple[str, ToolCallMetadata | None]: + results: list[str] = [] + for tool_call in params["tool_calls"]: + tool = self.toolset.get_tool(tool_call["tool_name"]) + if tool is None: + raise ValueError(f"Tool {tool_call['tool_name']} not found in toolset") + output, _ = tool(tool_call["parameters"]) + results.append(output) + return json.dumps(results), None + + +class UserTextTool(Tool): + + tool_schema: ToolSchema + + def __init__(self) -> None: + super().__init__( + tool_schema=ToolSchema( + name="user_text", + description="Produces text for the user.", + parameters={}, + required=[], + ) + ) + + def __call__( + self, + params: dict[Any, Any], + overrides: dict[Any, Any] | None = None, + ) -> tuple[str, ToolCallMetadata | None]: + raise ValueError("UserTextTool should not be called directly") + + + + +class ToolSet(BaseModel): + + tools: dict[str, Tool] = Field(default_factory=dict) + name: str | None = None + + def add_tool(self, tool: Tool) -> None: + if tool.tool_schema.name in self.tools: + raise ValueError(f"Tool with name {tool.tool_schema.name} already exists") + self.tools[tool.tool_schema.name] = tool + + def remove_tool(self, name: str) -> None: + self.tools.pop(name, None) + + def get_tool(self, name: str) -> Tool | None: + return self.tools.get(name) + + def get_formats(self, provider: ProviderFormat) -> list[dict[str, Any]]: + return [tool.get_format(provider) for tool in self.tools.values()] + + def __repr__(self) -> str: + names = ", ".join(sorted(self.tools.keys())) + suffix = f" ({self.name})" if self.name else "" + return f"ToolSet{suffix}[{len(self.tools)} tools: {names}]" + + @classmethod + def build( + cls, + *, + cosmos_database: DatabaseProxy | None = None, + cosmos_container_name: str | None = None, + openai_client: openai.OpenAI | None = None, + openai_embedding_model: str = "text-embedding-3-small", + embed_query_instruction: str | None = None, + retriever: CorpusRetriever | None = None, + reranker: Reranker | None = None, + token_counter: Callable[[str], int] | None = None, + max_tokens: int | None = None, + search_limit: int = 50, + search_display_limit: int = 10, + name: str | None = None, + ) -> ToolSet: + + if retriever is None: + if cosmos_database is None or cosmos_container_name is None or openai_client is None: + raise ValueError( + "ToolSet.build requires either 'retriever' or " + "'cosmos_database' + 'cosmos_container_name' + 'openai_client'" + ) + container = cosmos_database.get_container_client(cosmos_container_name) + embedder = QueryEmbedder( + client=openai_client, + model=openai_embedding_model, + query_instruction=embed_query_instruction, + ) + retriever = build_default_retriever( + container=container, + embedder=embedder, + embedding_model=openai_embedding_model, + ) + + toolset = cls(name=name) + toolset.add_tool( + SearchCorpusTool( + retriever=retriever, + reranker=reranker, + search_limit=search_limit, + display_limit=search_display_limit, + ) + ) + toolset.add_tool( + GrepCorpusTool( + retriever=retriever, + token_counter=token_counter, + ) + ) + toolset.add_tool( + ReadDocumentTool( + retriever=retriever, + reranker=reranker, + token_counter=token_counter, + max_tokens=max_tokens, + ) + ) + toolset.add_tool(PruneChunksTool()) + return toolset + + +__all__ = [ + "COSMOS_QUERY_MAX_CONCURRENCY", + "DOC_TRUNCATION", + "GREP_CORPUS_SCHEMA", + "GrepCorpusTool", + "GrepCorpusToolCallMetadata", + "MULTI_TOOL_USE_SCHEMA", + "MultiToolUseTool", + "PRUNE_CHUNKS_SCHEMA", + "PruneChunksTool", + "READ_DOCUMENT_SCHEMA", + "ReadDocumentTool", + "SEARCH_CORPUS_SCHEMA", + "SearchCorpusTool", + "SearchCorpusToolCallMetadata", + "SerializedTool", + "Tool", + "ToolCallMetadata", + "ToolSchema", + "ToolSet", + "UserTextTool", +] diff --git a/cosmos-retriever/src/cosmos_retriever/utils.py b/cosmos-retriever/src/cosmos_retriever/utils.py new file mode 100644 index 0000000..53f1229 --- /dev/null +++ b/cosmos-retriever/src/cosmos_retriever/utils.py @@ -0,0 +1,14 @@ + +from __future__ import annotations + +from enum import StrEnum + + +class ProviderFormat(StrEnum): + + OPENAI = "openai" + OPENAI_HARMONY = "openai_harmony" + ANTHROPIC = "anthropic" + + +__all__ = ["ProviderFormat"] diff --git a/docs/AGENTIC_SEARCH.md b/docs/AGENTIC_SEARCH.md new file mode 100644 index 0000000..067dfb0 --- /dev/null +++ b/docs/AGENTIC_SEARCH.md @@ -0,0 +1,247 @@ +# `agentic_search` — multi-turn retrieval as an MCP tool + +`agentic_search` runs a multi-turn search agent — built from scratch for this +toolkit — against an Azure Cosmos DB corpus and returns the ranked, curated set of documents +that best answer a natural-language query. The agent issues hybrid (vector + +full-text) RRF searches, optionally reranks with Qwen3-Reranker-8B, fetches +full documents, and prunes its working context across multiple turns. From +the MCP client's perspective it's a single tool call; under the hood the +agent can take 20–40 turns and 30–60 s of wall-clock time. + +## Architecture + +```text + MCP client MCPToolKit (.NET) cosmos-retriever (Python, FastAPI) + ────────── ───────────────── ─────────────────────────────────── + Claude Desktop ┌─ TokenBudgetRetrievalSubagent + AI Foundry ─── MCP HTTP ───► [McpServerTool] AgenticSearch │ ├─ SearchCorpus / Grep / Read / Prune + VS Code Copilot │ │ └─ OpenAI-compatible inference + ▼ │ + AgenticSearchExecutor ── HTTP POST ───► POST /search (uvicorn, kept warm) + │ │ + │ ◄────── JSON body ──────────────┤ + │ └─► LLM endpoint + Cosmos DB + embeddings + ▼ + MCP tool response +``` + +The .NET server and the Python retriever are now **two long-lived +processes**. The retriever is started once (`python -m cosmos_retriever +serve`) and keeps its Cosmos/embedding/LLM clients warm; the .NET server +calls its `POST /search` endpoint per MCP tool call and passes the JSON +response through verbatim. + +## Prerequisites + +You need three things running on the same host (or reachable from it): + +| Component | What it is | +|---|---| +| **An LLM endpoint** | Any OpenAI-compatible model — an Azure AI Foundry deployment, OpenAI, or a local server — via `INFERENCE_BACKEND=openai_responses` (default) or `openai_chat` (see below). | +| **Azure Cosmos DB for NoSQL** | Container populated with the standard chunked-corpus schema (`id`, `docid`, `chunk_idx`, `text`, `embedding`), vector + FTS indexes enabled. | +| **Embeddings backend** | Whatever model your corpus was ingested with — Azure OpenAI `text-embedding-3-small`, OpenAI native, or a local vLLM embedding server. | + +### Inference backend + +The retriever supports two backends, selected by `INFERENCE_BACKEND`: + +- `openai_responses` *(default)* — any OpenAI-compatible `/responses` model + (e.g. a reasoning model such as gpt-5.x), driven with standard tool calling. +- `openai_chat` — any OpenAI-compatible `/chat/completions` model (an Azure AI + Foundry deployment, OpenAI, a local server, ...). Set `CHAT_BASE_URL`, + `CHAT_API_KEY`, `CHAT_MODEL` (and `CHAT_API_VERSION` for Azure OpenAI-style + endpoints). The agent uses the same Cosmos tools, so retrieval quality tracks + the chosen model's tool-use ability. + +The Python helper is **bundled in this repository** at +[`cosmos-retriever/`](../cosmos-retriever/) — no separate clone needed. +Install it into a virtualenv: + +```bash +cd cosmos-retriever +uv venv --python 3.11 .venv +uv pip install --python .venv/bin/python -e . +``` + +Confirm it works: + +```bash +.venv/bin/python -m cosmos_retriever serve --help +``` + +Then start the service (it reads its own `.env` / `.env.local` for +`CHAT_BASE_URL`, `ACCOUNT_URI`, `COSMOS_*`, `AZURE_OPENAI_*`, `HOST`, `PORT`): + +```bash +.venv/bin/python -m cosmos_retriever serve # binds HOST:PORT (default 0.0.0.0:9000) +curl -s http://127.0.0.1:9000/health # -> {"status":"ok"} +``` + +## Server configuration + +Two env vars are read by the `AgenticSearchExecutor` service; both optional. +If `COSMOS_RETRIEVER_URL` doesn't point at a running retriever service, the +tool returns a clean JSON `{"error":"...","hint":"..."}` envelope rather than +crashing the server. + +| Variable | Default | Purpose | +|---|---|---| +| `COSMOS_RETRIEVER_URL` | `http://127.0.0.1:9000` | Base URL of the cosmos-retriever FastAPI service. | +| `COSMOS_RETRIEVER_TIMEOUT_S` | `600` | Per-request wall-clock cap; the request is abandoned if it exceeds this. | + +Unlike the previous subprocess design, the retriever service has its **own** +environment. Everything it needs (`CHAT_BASE_URL`, `ACCOUNT_URI`, +`COSMOS_DATABASE`, `COSMOS_CORPUS_CONTAINER`, `AZURE_OPENAI_*`, +`CORPUS_REGISTRY_FILE`, …) is read from the retriever process's environment / +`.env` file, **not** inherited from the .NET server. + +## Tool schema + +```jsonc +{ + "name": "agentic_search", + "description": "Runs a multi-turn retrieval agent against a Cosmos DB corpus and returns ranked, curated documents.", + "inputSchema": { + "type": "object", + "properties": { + "query": { "type": "string", "maxLength": 4096 }, + "maxDocuments": { "type": "integer", "minimum": 1, "maximum": 30, "default": 20 }, + "database": { "type": "string", "maxLength": 256 }, + "container": { "type": "string", "maxLength": 256 } + }, + "required": ["query"], + "additionalProperties": false + } +} +``` + +Tool result (the retriever service's `POST /search` JSON body, passed through verbatim): + +```jsonc +{ + "query": "Who discovered radium and when did she win her second Nobel?", + "num_turns": 5, + "elapsed_s": 32.3, + "documents": [ + { + "id": "96308__3", + "rank": 0, + "justification": "This biography directly states that Marie Curie ...", + "text": "..." + } + ] +} +``` + +On failure the helper (or the C# executor) returns a JSON error envelope: + +```jsonc +{ "error": "agentic_search timed out after 600s.", "stderr": "..." } +``` + +## Multi-corpus targeting + +`agentic_search` accepts optional `database` and `container` arguments so a +single MCP server can be aimed at multiple Cosmos corpora at request time. +For per-corpus *embedding-model* selection (e.g. one corpus ingested with +`text-embedding-3-small`, another with `qwen3-embed`), point +`CORPUS_REGISTRY_FILE` at a JSON file in the cosmos-retriever package: + +```jsonc +{ + "browsecomp_corpus_container": { + "account_uri": "https://acct-a.documents.azure.com:443/", + "database": "search_retrieval_database", + "embed_base_url": "https://embedding.services.ai.azure.com/openai/v1", + "embed_api_key_env": "AZURE_OPENAI_API_KEY", + "embed_model": "text-embedding-3-small" + }, + "enterprise_ragbench_corpus": { + "account_uri": "https://acct-b.documents.azure.com:443/", + "database": "search_retrieval_database", + "embed_base_url": "http://localhost:8002/v1", + "embed_api_key_env": null, + "embed_model": "qwen3-embed", + "embed_query_instruction": "Given a question, retrieve documents that answer it" + } +} +``` + +Then call: + +```jsonc +{ "name": "agentic_search", + "arguments": { + "query": "What was the temporary mitigation applied to the internal load balancer ...", + "container": "enterprise_ragbench_corpus" + } } +``` + +The matching account, database, embedding URL, model, and optional +`Instruct:` prefix all get picked automatically per call. Adding a third +corpus is a one-line registry edit — no rebuild, no restart. + +## Local demo + +End-to-end against any OpenAI-compatible endpoint plus Cosmos DB and embeddings. + +**1. Start the retriever service** (the bundled `cosmos-retriever/` folder; it +reads its own `.env`): + +```bash +cd cosmos-retriever +INFERENCE_BACKEND=openai_responses \ +CHAT_BASE_URL=https://your-resource.services.ai.azure.com/openai/v1 \ +CHAT_API_KEY=... \ +CHAT_MODEL=gpt-5.2 \ +VLLM_RERANKER_URL=http://localhost:8011 \ +CORPUS_REGISTRY_FILE=$PWD/corpus_registry.json \ +PORT=9000 \ +.venv/bin/python -m cosmos_retriever serve +``` + +**2. Start the .NET MCP server** (from the repo root), pointing it at the retriever URL: + +```bash +DEV_BYPASS_AUTH=true \ +COSMOS_RETRIEVER_URL=http://127.0.0.1:9000 \ +OPENAI_ENDPOINT="$AZURE_OPENAI_ENDPOINT" \ +OPENAI_EMBEDDING_DEPLOYMENT="$AZURE_OPENAI_EMBED_DEPLOYMENT" \ +dotnet run --project src/AzureCosmosDB.MCP.Toolkit +``` + +Then point any MCP client at `http://127.0.0.1:8080/mcp/`. + +## Operational notes + +- **The retriever service has its own environment.** Configure + `CHAT_BASE_URL`, `ACCOUNT_URI`, `COSMOS_*`, `AZURE_OPENAI_*`, + `CORPUS_REGISTRY_FILE`, etc. where you launch `cosmos_retriever serve` + (env or its `.env` file) — the .NET server no longer forwards them. +- **`COSMOS_USE_DEFAULT_CREDENTIAL`** controls the retriever's Cosmos auth. + By default it uses `AzureCliCredential`; set it to `1` to opt into the + broader `DefaultAzureCredential` chain (managed identity, etc.). +- **Warm process, no cold start.** Because the service stays up, the heavy + client init happens once. Per-call latency is dominated by Cosmos + round-trips + LLM generation; don't expect sub-second latency. +- **Retrieval quality is corpus-dependent.** Cosmos's hybrid RRF puts gold + docs in the top 5 reliably; the Qwen3-Reranker step on top can over- or + under-shoot depending on how close the corpus distribution is to the + reranker's training data. If you see recall regressions, try disabling the + reranker for that corpus (omit `VLLM_RERANKER_URL` / `BASETEN_API_KEY`). +- **The tool always returns parseable JSON.** Unreachable service, request + timeouts, and non-2xx responses all yield + `{"error": "...", "hint"?: "...", "body"?: "..."}` envelopes rather than + HTTP 500s to the MCP client. + +## Implementation pointers + +| File | Role | +|---|---| +| [`Services/AgenticSearchExecutor.cs`](../src/AzureCosmosDB.MCP.Toolkit/Services/AgenticSearchExecutor.cs) | HTTP call to the retriever service, timeout, error-envelope generation. | +| [`Services/CosmosDbToolsService.cs`](../src/AzureCosmosDB.MCP.Toolkit/Services/CosmosDbToolsService.cs) | `AgenticSearch` instance method called by both controllers. | +| [`Program.cs`](../src/AzureCosmosDB.MCP.Toolkit/Program.cs) | `[McpServerTool] AgenticSearch` static method discovered by the MCP SDK. | +| [`Controllers/MCPProtocolController.cs`](../src/AzureCosmosDB.MCP.Toolkit/Controllers/MCPProtocolController.cs) | JSON-RPC `tools/list` + `tools/call` dispatch for the custom `/mcp/http` transport. | +| [`Controllers/MCPTestController.cs`](../src/AzureCosmosDB.MCP.Toolkit/Controllers/MCPTestController.cs) | REST sibling at `POST /api/mcp/tools/agentic_search`. | +| [`Services/McpToolRequestValidator.cs`](../src/AzureCosmosDB.MCP.Toolkit/Services/McpToolRequestValidator.cs) | Strict input validation schema. | +| [`cosmos-retriever/`](../cosmos-retriever/) | The bundled Python FastAPI service (`POST /search`) the executor calls; run with `python -m cosmos_retriever serve`. | diff --git a/run-mcp-server.ps1 b/run-mcp-server.ps1 new file mode 100644 index 0000000..6bf22e6 --- /dev/null +++ b/run-mcp-server.ps1 @@ -0,0 +1,28 @@ +# Loads the repo-root .env into the current process environment (the .NET +# server does NOT read .env on its own), then starts the MCP Toolkit server. +# Usage: .\run-mcp-server.ps1 +$ErrorActionPreference = "Stop" +$here = Split-Path -Parent $MyInvocation.MyCommand.Path +$envFile = Join-Path $here ".env" + +if (Test-Path $envFile) { + Get-Content $envFile | ForEach-Object { + $line = $_.Trim() + if ($line -and -not $line.StartsWith("#") -and $line.Contains("=")) { + $idx = $line.IndexOf("=") + $name = $line.Substring(0, $idx).Trim() + $value = $line.Substring($idx + 1).Trim() + if ($value -match '^<.*>$') { + Write-Warning "Env var '$name' still has a placeholder value; skipping. Edit .env." + } else { + [Environment]::SetEnvironmentVariable($name, $value, "Process") + Write-Host " set $name" -ForegroundColor DarkGray + } + } + } +} else { + Write-Warning "No .env file found at $envFile" +} + +Write-Host "Starting Azure Cosmos DB MCP Toolkit on http://localhost:8080 ..." -ForegroundColor Cyan +dotnet run --project (Join-Path $here "src\AzureCosmosDB.MCP.Toolkit") diff --git a/src/AzureCosmosDB.MCP.Toolkit/AzureCosmosDB.MCP.Toolkit.csproj b/src/AzureCosmosDB.MCP.Toolkit/AzureCosmosDB.MCP.Toolkit.csproj index f7a1e6a..a09fdb3 100644 --- a/src/AzureCosmosDB.MCP.Toolkit/AzureCosmosDB.MCP.Toolkit.csproj +++ b/src/AzureCosmosDB.MCP.Toolkit/AzureCosmosDB.MCP.Toolkit.csproj @@ -7,7 +7,7 @@ false true AzureCosmosDB.MCP.Toolkit - 1.1.2 + 1.2.0 Azure Cosmos DB Team Microsoft Azure Cosmos DB MCP Toolkit diff --git a/src/AzureCosmosDB.MCP.Toolkit/Controllers/MCPProtocolController.cs b/src/AzureCosmosDB.MCP.Toolkit/Controllers/MCPProtocolController.cs index fd51717..4bca534 100644 --- a/src/AzureCosmosDB.MCP.Toolkit/Controllers/MCPProtocolController.cs +++ b/src/AzureCosmosDB.MCP.Toolkit/Controllers/MCPProtocolController.cs @@ -259,6 +259,21 @@ public async Task HandleMCPRequest([FromBody] JsonElement request required = new string[] { "databaseId", "containerId", "searchText", "textProperty", "vectorProperty", "selectProperties" }, additionalProperties = false } + }, + new { + name = "agentic_search", + description = "PREFERRED tool for answering knowledge questions from a Cosmos DB corpus. Runs an autonomous multi-turn retrieval agent that plans sub-queries, issues several vector/keyword searches, follows leads across documents, reranks candidates, and returns a curated, ranked set of the most relevant documents with their content. Use this for anything beyond a trivial lookup: complex, ambiguous, multi-part, or multi-hop questions; or whenever one-shot vector_search/text_search might miss relevant context. It is more thorough (but slower) than the single-shot search tools, so prefer it when answer quality matters more than latency. Just pass a natural-language `query`; the agent handles query planning and ranking for you. Optionally pass `container=` to target a registered corpus (see the CORPUS_REGISTRY env var on the host): the matching Cosmos account + database + embedding model is selected automatically per call. With no `container`, the default-corpus env vars are used. Use `maxDocuments` to cap how many curated documents are returned.", + inputSchema = new { + type = "object", + properties = new { + query = new { type = "string", description = "Natural-language information need to retrieve documents for", maxLength = 4096 }, + maxDocuments = new { type = "integer", description = "Maximum number of curated documents to return (1-50, default 20)", minimum = 1, maximum = 50, @default = 20 }, + database = new { type = "string", description = "Optional Cosmos database override (else COSMOS_DATABASE env var)", maxLength = 256 }, + container = new { type = "string", description = "Optional Cosmos corpus container override (else COSMOS_CORPUS_CONTAINER env var)", maxLength = 256 } + }, + required = new string[] { "query" }, + additionalProperties = false + } } } } @@ -452,6 +467,12 @@ private async Task ExecuteTool(string toolName, Dictionary await _cosmosDbTools.AgenticSearch( + GetStringArg(args, "query"), + GetOptionalIntArg(args, "maxDocuments", 20), + GetOptionalStringArg(args, "database"), + GetOptionalStringArg(args, "container"), + cancellationToken), _ => throw new ArgumentException($"Unknown tool: {toolName}") }; } @@ -461,6 +482,13 @@ private static string GetStringArg(Dictionary args, string key) return args.TryGetValue(key, out var value) ? value?.ToString() ?? "" : ""; } + private static string? GetOptionalStringArg(Dictionary args, string key) + { + if (!args.TryGetValue(key, out var value)) return null; + var s = value?.ToString(); + return string.IsNullOrWhiteSpace(s) ? null : s; + } + private static int GetRequiredIntArg(Dictionary args, string key) { if (!args.TryGetValue(key, out var value)) diff --git a/src/AzureCosmosDB.MCP.Toolkit/Controllers/MCPTestController.cs b/src/AzureCosmosDB.MCP.Toolkit/Controllers/MCPTestController.cs index 7696874..9eb5e37 100644 --- a/src/AzureCosmosDB.MCP.Toolkit/Controllers/MCPTestController.cs +++ b/src/AzureCosmosDB.MCP.Toolkit/Controllers/MCPTestController.cs @@ -34,6 +34,7 @@ public async Task CallTool(string toolName, [FromBody] MCPToolReq "text_search" => await CallTextSearch(request.Parameters), "vector_search" => await CallVectorSearch(request.Parameters), "get_approximate_schema" => await CallGetApproximateSchema(request.Parameters), + "agentic_search" => await CallAgenticSearch(request.Parameters), _ => throw new ArgumentException($"Unknown tool: {toolName}") }; @@ -71,7 +72,8 @@ public IActionResult ListTools() new { name = "find_document_by_id", description = "Finds a document by its ID in the specified database/container" }, new { name = "text_search", description = "Select TOP N documents where a given property contains the provided search string. N must be between 1-20" }, new { name = "vector_search", description = "Performs vector search on Cosmos DB using Azure OpenAI embeddings" }, - new { name = "get_approximate_schema", description = "Approximates a container schema by sampling up to 10 documents" } + new { name = "get_approximate_schema", description = "Approximates a container schema by sampling up to 10 documents" }, + new { name = "agentic_search", description = "Runs an autonomous multi-turn retrieval agent against a Cosmos DB corpus and returns ranked, curated documents that best answer the query." } }; return Ok(new { tools, count = tools.Length, timestamp = DateTime.UtcNow }); @@ -127,6 +129,17 @@ private async Task CallGetApproximateSchema(Dictionary p return await _cosmosDbTools.GetApproximateSchema(databaseId, containerId); } + private async Task CallAgenticSearch(Dictionary parameters) + { + var query = GetRequiredParameter(parameters, "query"); + var maxDocuments = parameters.ContainsKey("maxDocuments") + ? GetRequiredParameter(parameters, "maxDocuments") + : 20; + string? database = parameters.ContainsKey("database") ? GetRequiredParameter(parameters, "database") : null; + string? container = parameters.ContainsKey("container") ? GetRequiredParameter(parameters, "container") : null; + return await _cosmosDbTools.AgenticSearch(query, maxDocuments, database, container); + } + private T GetRequiredParameter(Dictionary parameters, string paramName) { if (!parameters.TryGetValue(paramName, out var value)) diff --git a/src/AzureCosmosDB.MCP.Toolkit/Program.cs b/src/AzureCosmosDB.MCP.Toolkit/Program.cs index ffa178c..7c9a4ed 100644 --- a/src/AzureCosmosDB.MCP.Toolkit/Program.cs +++ b/src/AzureCosmosDB.MCP.Toolkit/Program.cs @@ -243,6 +243,7 @@ // Store configuration in static state for access by static tool methods AppState.Configuration = builder.Configuration; +AppState.LoggerFactory = app.Services.GetRequiredService(); // Add security headers middleware to allow MSAL authentication app.Use(async (context, next) => @@ -345,6 +346,7 @@ internal static class AppState { public static IConfiguration? Configuration { get; set; } + public static ILoggerFactory? LoggerFactory { get; set; } } public partial class Program @@ -1093,4 +1095,26 @@ FROM c return JsonSerializer.Serialize(new { error = ex.Message }); } } + + [McpServerTool, Description("PREFERRED tool for answering knowledge questions from a Cosmos DB corpus. Runs an autonomous multi-turn retrieval agent that plans sub-queries, issues several vector/keyword searches, follows leads across documents, reranks candidates, and returns a curated, ranked set of the most relevant documents with their content. Use this for anything beyond a trivial lookup: complex, ambiguous, multi-part, or multi-hop questions; or whenever one-shot vector_search/text_search might miss relevant context. It is more thorough (but slower) than the single-shot search tools, so prefer it when answer quality matters more than latency. Just pass a natural-language `query`; the agent handles query planning and ranking for you. Optionally pass `container=` to target a registered corpus (see the CORPUS_REGISTRY env var on the host): the matching Cosmos account + database + embedding model is selected automatically per call. With no `container` the default-corpus env vars are used. Use `maxDocuments` to cap how many curated documents are returned.")] + public static async Task AgenticSearch( + [Description("Natural-language information need to retrieve documents for.")] string query, + [Description("Maximum number of curated documents to return (1-50, default 20).")] int maxDocuments = 20, + [Description("Optional Cosmos database name override (else COSMOS_DATABASE env var).")] string? database = null, + [Description("Optional Cosmos corpus container name override (else COSMOS_CORPUS_CONTAINER env var).")] string? container = null) + { + var logger = (AppState.LoggerFactory ?? Microsoft.Extensions.Logging.Abstractions.NullLoggerFactory.Instance) + .CreateLogger("AzureCosmosDB.MCP.Toolkit.CosmosDbTools.AgenticSearch"); + + if (string.IsNullOrWhiteSpace(query)) + { + return JsonSerializer.Serialize(new { error = "Parameter 'query' is required and must be non-empty." }); + } + if (maxDocuments < 1 || maxDocuments > 50) + { + return JsonSerializer.Serialize(new { error = "Parameter 'maxDocuments' must be between 1 and 50." }); + } + + return await AgenticSearchExecutor.RunAsync(query, maxDocuments, logger, database, container); + } } diff --git a/src/AzureCosmosDB.MCP.Toolkit/Services/AgenticSearchExecutor.cs b/src/AzureCosmosDB.MCP.Toolkit/Services/AgenticSearchExecutor.cs new file mode 100644 index 0000000..25dc83b --- /dev/null +++ b/src/AzureCosmosDB.MCP.Toolkit/Services/AgenticSearchExecutor.cs @@ -0,0 +1,191 @@ +using System.Globalization; +using System.Net.Http.Json; +using System.Text.Json; + +namespace AzureCosmosDB.MCP.Toolkit.Services; + +/// +/// Calls the long-lived cosmos-retriever FastAPI service over HTTP and +/// returns its response body (a single JSON document) verbatim. +/// +/// +/// +/// The Python helper runs a multi-turn retrieval agent against an Azure Cosmos +/// DB corpus and returns a JSON document of curated, ranked results. It is +/// started once (python -m cosmos_retriever serve) and kept warm so the +/// heavy clients (Cosmos SDK, embeddings, model encoder) are not re-initialised +/// on every call. +/// +/// +/// Host environment variables (read on every call): +/// +/// VariableDefault / purpose +/// +/// (COSMOS_RETRIEVER_URL) +/// Base URL of the cosmos-retriever FastAPI service. +/// Defaults to . +/// +/// +/// (COSMOS_RETRIEVER_TIMEOUT_S) +/// Per-request wall-clock cap in seconds; the request is +/// abandoned if it exceeds the timeout. Defaults to +/// . +/// +/// +/// +/// +/// The retriever service owns its own configuration (model endpoint, +/// ACCOUNT_URI, COSMOS_DATABASE, COSMOS_CORPUS_CONTAINER, +/// CORPUS_REGISTRY_FILE, AZURE_OPENAI_*, etc.) read from its own +/// environment / .env file; none of it flows through this process. +/// +/// +public static class AgenticSearchExecutor +{ + public const string BaseUrlEnvVar = "COSMOS_RETRIEVER_URL"; + + public const string TimeoutEnvVar = "COSMOS_RETRIEVER_TIMEOUT_S"; + + public const string DefaultBaseUrl = "http://127.0.0.1:9000"; + + public const int DefaultTimeoutSeconds = 600; + + private const int BodyTruncateBytes = 4096; + + // A single shared HttpClient with no built-in timeout — each call drives + // its own deadline via a linked CancellationTokenSource. + private static readonly HttpClient HttpClient = new() + { + Timeout = Timeout.InfiniteTimeSpan, + }; + + /// + /// Run a single cosmos-retriever search by calling the FastAPI + /// POST /search endpoint. + /// + /// Natural-language information need. + /// Cap on the number of curated docs returned (1–50). + /// Logger for request lifecycle events. + /// Optional Cosmos database override. + /// Optional Cosmos container override. + /// Cooperative cancellation. + /// + /// The service's response body, expected to be a single JSON document. On + /// any failure (service unreachable, timed out, non-success status, empty + /// body) returns a serialised { "error": "...", ... } envelope so + /// the MCP tool always returns parseable JSON to the caller. + /// + public static async Task RunAsync( + string query, + int maxDocuments, + ILogger logger, + string? database = null, + string? container = null, + CancellationToken cancellationToken = default) + { + var baseUrl = ResolveString(BaseUrlEnvVar, defaultValue: DefaultBaseUrl).TrimEnd('/'); + var timeoutSeconds = ResolveInt(TimeoutEnvVar, DefaultTimeoutSeconds); + var requestUri = $"{baseUrl}/search"; + + var payload = new Dictionary + { + ["query"] = query, + ["maxDocuments"] = maxDocuments, + }; + if (!string.IsNullOrWhiteSpace(database)) payload["database"] = database; + if (!string.IsNullOrWhiteSpace(container)) payload["container"] = container; + + logger.LogInformation( + "agentic_search: POST {RequestUri} (database={Database} container={Container} timeout={Timeout}s)", + requestUri, database ?? "", container ?? "", timeoutSeconds); + + using var timeoutCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken); + timeoutCts.CancelAfter(TimeSpan.FromSeconds(timeoutSeconds)); + + HttpResponseMessage response; + try + { + using var content = JsonContent.Create(payload); + response = await HttpClient + .PostAsync(requestUri, content, timeoutCts.Token) + .ConfigureAwait(false); + } + catch (OperationCanceledException) when (timeoutCts.IsCancellationRequested && !cancellationToken.IsCancellationRequested) + { + logger.LogWarning("agentic_search: request exceeded {Timeout}s.", timeoutSeconds); + return ErrorEnvelope( + $"agentic_search timed out after {timeoutSeconds}s.", + hint: $"Increase {TimeoutEnvVar} or check that the cosmos-retriever service at {baseUrl} is responsive."); + } + catch (HttpRequestException ex) + { + logger.LogError(ex, + "agentic_search: failed to reach the cosmos-retriever service at {BaseUrl}.", baseUrl); + return ErrorEnvelope( + $"Failed to reach the cosmos-retriever service: {ex.Message}", + hint: $"Start it with 'python -m cosmos_retriever serve' and set {BaseUrlEnvVar} to its base URL (default {DefaultBaseUrl})."); + } + + using (response) + { + var body = (await response.Content.ReadAsStringAsync(cancellationToken).ConfigureAwait(false)).Trim(); + + if (!response.IsSuccessStatusCode) + { + logger.LogWarning( + "agentic_search: service returned {StatusCode}. body tail: {Body}", + (int)response.StatusCode, TruncateTail(body, 512)); + + // The FastAPI service emits its own JSON error envelope on most + // failures; pass it through verbatim if so, otherwise wrap it. + if (LooksLikeJson(body)) + { + return body; + } + return ErrorEnvelope( + $"agentic_search service returned HTTP {(int)response.StatusCode}.", + bodyTail: TruncateTail(body, BodyTruncateBytes)); + } + + if (string.IsNullOrWhiteSpace(body)) + { + return ErrorEnvelope("agentic_search service produced no output."); + } + + return body; + } + } + + private static string ResolveString(string envVar, string defaultValue) + { + var value = Environment.GetEnvironmentVariable(envVar); + return string.IsNullOrWhiteSpace(value) ? defaultValue : value; + } + + private static int ResolveInt(string envVar, int defaultValue) + { + var raw = Environment.GetEnvironmentVariable(envVar); + if (!string.IsNullOrWhiteSpace(raw) && int.TryParse(raw, NumberStyles.Integer, CultureInfo.InvariantCulture, out var parsed) && parsed > 0) + { + return parsed; + } + return defaultValue; + } + + private static bool LooksLikeJson(string s) => + s.Length > 0 && (s[0] == '{' || s[0] == '['); + + private static string ErrorEnvelope(string error, string? hint = null, string? bodyTail = null) + { + var payload = new Dictionary { ["error"] = error }; + if (hint is not null) payload["hint"] = hint; + if (bodyTail is not null) payload["body"] = bodyTail; + return JsonSerializer.Serialize(payload); + } + + private static string TruncateTail(string s, int maxChars) + { + if (string.IsNullOrEmpty(s) || s.Length <= maxChars) return s ?? string.Empty; + return "..." + s[^maxChars..]; + } +} diff --git a/src/AzureCosmosDB.MCP.Toolkit/Services/CosmosClientFactory.cs b/src/AzureCosmosDB.MCP.Toolkit/Services/CosmosClientFactory.cs index 82958b3..ad2cc33 100644 --- a/src/AzureCosmosDB.MCP.Toolkit/Services/CosmosClientFactory.cs +++ b/src/AzureCosmosDB.MCP.Toolkit/Services/CosmosClientFactory.cs @@ -74,7 +74,14 @@ public static CosmosClient CreateCosmosClient(IConfiguration configuration, ILog } logger.LogInformation("Creating CosmosClient using Azure credentials (cloud mode)"); - var credential = new DefaultAzureCredential(); + // Exclude ManagedIdentityCredential: on Azure VMs MSI_ENDPOINT/IMDS is present + // but the managed identity often lacks Cosmos RBAC (SSO failure). Skipping it + // lets the chain fall through to the Azure CLI login (az login), which the + // Python retriever uses successfully. + var credential = new DefaultAzureCredential(new DefaultAzureCredentialOptions + { + ExcludeManagedIdentityCredential = true, + }); return new CosmosClient(endpoint, credential, BuildClientOptions(configuration, logger, useGatewayMode: false)); } diff --git a/src/AzureCosmosDB.MCP.Toolkit/Services/CosmosDbToolsService.cs b/src/AzureCosmosDB.MCP.Toolkit/Services/CosmosDbToolsService.cs index 46bb07f..da77bc2 100644 --- a/src/AzureCosmosDB.MCP.Toolkit/Services/CosmosDbToolsService.cs +++ b/src/AzureCosmosDB.MCP.Toolkit/Services/CosmosDbToolsService.cs @@ -700,4 +700,22 @@ public async Task GetApproximateSchema(string databaseId, string contain return new { error = ex.Message }; } } + + /// + /// Calls the cosmos-retriever FastAPI service and returns its raw response + /// body (a single JSON document). See + /// for the environment-variable contract and timeout knobs. + /// + public async Task AgenticSearch( + string query, + int maxDocuments = 20, + string? database = null, + string? container = null, + CancellationToken cancellationToken = default) + { + var raw = await AgenticSearchExecutor.RunAsync(query, maxDocuments, _logger, database, container, cancellationToken); + // Pass the JSON string through verbatim so the MCP envelope serialises it + // as a single string (matching the other tools, which also return JSON strings). + return raw; + } } diff --git a/src/AzureCosmosDB.MCP.Toolkit/Services/McpToolRequestValidator.cs b/src/AzureCosmosDB.MCP.Toolkit/Services/McpToolRequestValidator.cs index 859c566..3875642 100644 --- a/src/AzureCosmosDB.MCP.Toolkit/Services/McpToolRequestValidator.cs +++ b/src/AzureCosmosDB.MCP.Toolkit/Services/McpToolRequestValidator.cs @@ -58,6 +58,13 @@ public sealed class McpToolRequestValidator ["vectorProperty"] = ToolArgumentSchema.String(required: true, maxLength: 256), ["selectProperties"] = ToolArgumentSchema.String(required: true, maxLength: 512), ["topN"] = ToolArgumentSchema.Integer(required: false, minValue: 1, maxValue: 50) + }), + ["agentic_search"] = new(new Dictionary(StringComparer.Ordinal) + { + ["query"] = ToolArgumentSchema.String(required: true, maxLength: 4096), + ["maxDocuments"] = ToolArgumentSchema.Integer(required: false, minValue: 1, maxValue: 50), + ["database"] = ToolArgumentSchema.String(required: false, maxLength: 256), + ["container"] = ToolArgumentSchema.String(required: false, maxLength: 256) }) }; @@ -68,7 +75,9 @@ public ToolValidationResult ValidateToolCall(JsonElement paramsElement) throw new ToolInputValidationException("'params' must be a JSON object."); } - RejectUnknownProperties(paramsElement, ["name", "arguments"], "params"); + // `_meta` is a standard MCP field clients may attach to params (e.g. progress + // tokens); accept and ignore it rather than rejecting the request. + RejectUnknownProperties(paramsElement, ["name", "arguments", "_meta"], "params"); if (!paramsElement.TryGetProperty("name", out var toolNameElement) || toolNameElement.ValueKind != JsonValueKind.String) { diff --git a/tests/AzureCosmosDB.MCP.Toolkit.Tests/AgenticSearchExecutorTests.cs b/tests/AzureCosmosDB.MCP.Toolkit.Tests/AgenticSearchExecutorTests.cs new file mode 100644 index 0000000..daa7b0b --- /dev/null +++ b/tests/AzureCosmosDB.MCP.Toolkit.Tests/AgenticSearchExecutorTests.cs @@ -0,0 +1,227 @@ +using System.Net; +using System.Net.Sockets; +using System.Text; +using System.Text.Json; +using AzureCosmosDB.MCP.Toolkit.Services; +using FluentAssertions; +using Microsoft.Extensions.Logging.Abstractions; +using Xunit; + +namespace AzureCosmosDB.MCP.Toolkit.Tests; + +/// +/// Tests for . Stands in for the +/// cosmos-retriever FastAPI service with a tiny in-process +/// so we can verify the executor's response +/// pass-through, timeout behaviour, and error-envelope generation without +/// needing the real retriever service running. +/// +public sealed class AgenticSearchExecutorTests : IDisposable +{ + private readonly Dictionary _savedEnv = new(); + private static readonly NullLogger _logger = NullLogger.Instance; + + private void SetEnv(string name, string? value) + { + if (!_savedEnv.ContainsKey(name)) + { + _savedEnv[name] = Environment.GetEnvironmentVariable(name); + } + Environment.SetEnvironmentVariable(name, value); + } + + public void Dispose() + { + foreach (var (k, v) in _savedEnv) + { + Environment.SetEnvironmentVariable(k, v); + } + } + + [Fact] + public async Task RunAsync_passes_through_service_response_body() + { + const string body = + "{\"query\":\"hi\",\"documents\":[{\"id\":\"doc_a\",\"rank\":0}],\"num_turns\":1,\"elapsed_s\":0.01}"; + + using var server = StubServer.Start((ctx, _) => + { + ctx.Response.StatusCode = 200; + ctx.Response.ContentType = "application/json"; + return body; + }); + + SetEnv(AgenticSearchExecutor.BaseUrlEnvVar, server.BaseUrl); + SetEnv(AgenticSearchExecutor.TimeoutEnvVar, "30"); + + var raw = await AgenticSearchExecutor.RunAsync("hi", maxDocuments: 5, logger: _logger); + + using var doc = JsonDocument.Parse(raw); + doc.RootElement.GetProperty("query").GetString().Should().Be("hi"); + doc.RootElement.GetProperty("num_turns").GetInt32().Should().Be(1); + doc.RootElement.GetProperty("documents")[0].GetProperty("id").GetString().Should().Be("doc_a"); + } + + [Fact] + public async Task RunAsync_forwards_request_payload_to_service() + { + string? capturedBody = null; + using var server = StubServer.Start((ctx, reqBody) => + { + capturedBody = reqBody; + ctx.Response.StatusCode = 200; + return "{\"query\":\"q\",\"documents\":[],\"num_turns\":0,\"elapsed_s\":0.0}"; + }); + + SetEnv(AgenticSearchExecutor.BaseUrlEnvVar, server.BaseUrl); + SetEnv(AgenticSearchExecutor.TimeoutEnvVar, "30"); + + await AgenticSearchExecutor.RunAsync( + "find me docs", maxDocuments: 7, logger: _logger, database: "db1", container: "corpus-x"); + + capturedBody.Should().NotBeNull(); + using var doc = JsonDocument.Parse(capturedBody!); + doc.RootElement.GetProperty("query").GetString().Should().Be("find me docs"); + doc.RootElement.GetProperty("maxDocuments").GetInt32().Should().Be(7); + doc.RootElement.GetProperty("database").GetString().Should().Be("db1"); + doc.RootElement.GetProperty("container").GetString().Should().Be("corpus-x"); + } + + [Fact] + public async Task RunAsync_passes_through_service_error_envelope_on_non_success() + { + using var server = StubServer.Start((ctx, _) => + { + ctx.Response.StatusCode = 500; + ctx.Response.ContentType = "application/json"; + return "{\"error\":\"vllm unreachable\",\"type\":\"RuntimeError\"}"; + }); + + SetEnv(AgenticSearchExecutor.BaseUrlEnvVar, server.BaseUrl); + SetEnv(AgenticSearchExecutor.TimeoutEnvVar, "30"); + + var raw = await AgenticSearchExecutor.RunAsync("hi", maxDocuments: 5, logger: _logger); + + using var doc = JsonDocument.Parse(raw); + doc.RootElement.GetProperty("error").GetString().Should().Be("vllm unreachable"); + } + + [Fact] + public async Task RunAsync_returns_error_envelope_when_service_unreachable() + { + // Reserve+release a port so nothing is listening on it. + var port = GetFreePort(); + SetEnv(AgenticSearchExecutor.BaseUrlEnvVar, $"http://127.0.0.1:{port}"); + SetEnv(AgenticSearchExecutor.TimeoutEnvVar, "5"); + + var raw = await AgenticSearchExecutor.RunAsync("hi", maxDocuments: 5, logger: _logger); + + using var doc = JsonDocument.Parse(raw); + doc.RootElement.GetProperty("error").GetString().Should().Contain("Failed to reach"); + doc.RootElement.TryGetProperty("hint", out var hint).Should().BeTrue(); + hint.GetString().Should().Contain(AgenticSearchExecutor.BaseUrlEnvVar); + } + + [Fact] + public async Task RunAsync_returns_error_envelope_when_service_times_out() + { + using var server = StubServer.Start((ctx, _) => + { + // Sleep for longer than the 1s timeout we're about to set. + Thread.Sleep(5000); + ctx.Response.StatusCode = 200; + return "{}"; + }); + + SetEnv(AgenticSearchExecutor.BaseUrlEnvVar, server.BaseUrl); + SetEnv(AgenticSearchExecutor.TimeoutEnvVar, "1"); + + var raw = await AgenticSearchExecutor.RunAsync("hi", maxDocuments: 5, logger: _logger); + + using var doc = JsonDocument.Parse(raw); + doc.RootElement.GetProperty("error").GetString().Should().Contain("timed out after 1s"); + } + + private static int GetFreePort() + { + var listener = new TcpListener(IPAddress.Loopback, 0); + listener.Start(); + var port = ((IPEndPoint)listener.LocalEndpoint).Port; + listener.Stop(); + return port; + } + + /// + /// Minimal in-process HTTP server backed by . + /// The handler receives the request context plus the request body and + /// returns the response body string. + /// + private sealed class StubServer : IDisposable + { + private readonly HttpListener _listener; + private readonly CancellationTokenSource _cts = new(); + + public string BaseUrl { get; } + + private StubServer(HttpListener listener, string baseUrl) + { + _listener = listener; + BaseUrl = baseUrl; + } + + public static StubServer Start(Func handler) + { + var port = GetFreePort(); + var baseUrl = $"http://127.0.0.1:{port}"; + var listener = new HttpListener(); + listener.Prefixes.Add($"{baseUrl}/"); + listener.Start(); + var server = new StubServer(listener, baseUrl); + _ = Task.Run(() => server.LoopAsync(handler)); + return server; + } + + private async Task LoopAsync(Func handler) + { + while (!_cts.IsCancellationRequested) + { + HttpListenerContext ctx; + try + { + ctx = await _listener.GetContextAsync().ConfigureAwait(false); + } + catch + { + return; // listener stopped + } + + try + { + string reqBody; + using (var reader = new StreamReader(ctx.Request.InputStream, Encoding.UTF8)) + { + reqBody = await reader.ReadToEndAsync().ConfigureAwait(false); + } + + var responseBody = handler(ctx, reqBody); + var buffer = Encoding.UTF8.GetBytes(responseBody); + ctx.Response.ContentLength64 = buffer.Length; + await ctx.Response.OutputStream.WriteAsync(buffer).ConfigureAwait(false); + ctx.Response.OutputStream.Close(); + } + catch + { + try { ctx.Response.Abort(); } catch { /* best effort */ } + } + } + } + + public void Dispose() + { + _cts.Cancel(); + try { _listener.Stop(); } catch { /* best effort */ } + try { _listener.Close(); } catch { /* best effort */ } + _cts.Dispose(); + } + } +}