diff --git a/.env.example b/.env.example index 2b9b14b..a8a198b 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 a multi-turn retrieval agent, which runs +# as a long-lived FastAPI service started with +# `python -m cosmos_retriever serve`. See docs/AGENTIC_SEARCH.md. +# Both vars below are optional with sensible defaults; if the service is not +# reachable, agentic_search simply returns a clean JSON error envelope to the +# caller. + +# Base URL of the cosmos-retriever FastAPI service (default http://127.0.0.1:9000). +# COSMOS_RETRIEVER_URL=http://127.0.0.1:9000 + +# Per-request wall-clock cap in seconds (default 600). +# COSMOS_RETRIEVER_TIMEOUT_S=600 + # ============================================================================ # DOCKER COMPOSE NOTES # ============================================================================ 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..04f1d38 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,31 @@ All notable changes to this project will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## [1.2.0] - 2026-06-18 + +### Added +- **`agentic_search` tool**: Runs a multi-turn retrieval agent against a Cosmos + DB corpus and returns ranked, curated documents that best answer the query. + Given a query, the agent will (1) issue hybrid (vector + full-text) RRF + searches, (2) optionally rerank the hits with Qwen3-Reranker-8B, (3) read full + documents, and (4) prune its context across multiple turns. Implemented as a + subprocess call into the companion [`cosmos-retriever/`](cosmos-retriever/) + Python package; see [`docs/AGENTIC_SEARCH.md`](docs/AGENTIC_SEARCH.md) for the + deployment story. +- Optional `database` and `container` arguments on `agentic_search` so a + single MCP server can target multiple Cosmos corpora at request time. When + the corpus registry (`CORPUS_REGISTRY` / `CORPUS_REGISTRY_FILE`) is set + in the host environment, the matching account, database, and embedding + model are picked automatically per call. +- New service: `AgenticSearchExecutor` (subprocess lifecycle, timeout, error + envelope generation). +- New env vars: `COSMOS_RETRIEVER_PYTHON`, `COSMOS_RETRIEVER_DIR`, + `COSMOS_RETRIEVER_TIMEOUT_S` — see [`.env.example`](.env.example). + +### Changed +- `AppState` now also exposes `ILoggerFactory` so static `[McpServerTool]` + methods can obtain a properly-named logger. + ## [1.1.2] - 2026-05-29 ### Added diff --git a/README.md b/README.md index 2fd336c..90d153c 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` | Perform multi-turn retrieval with a configurable agent. The agent rewrites the query, issues tool calls against the configured corpus and containers, and returns the most relevant documents. See [docs/AGENTIC_SEARCH.md](docs/AGENTIC_SEARCH.md) for setup and configuration. | ## Project Structure diff --git a/cosmos-retriever/.env.example b/cosmos-retriever/.env.example new file mode 100644 index 0000000..e746d2b --- /dev/null +++ b/cosmos-retriever/.env.example @@ -0,0 +1,84 @@ +# ============================================================================= +# Cosmos Retriever configuration (Python service) +# ============================================================================= +# Every setting read by `RetrieverSettings` (config.py) is listed here with its +# default. Values load from environment variables or a `.env` / `.env.local` file +# at the repo root. Required keys are uncommented with placeholders; optional keys +# are commented out showing their default. Variable names are case-insensitive. +# +# NOTE: this file configures the *Python retriever service*. The .NET MCP server +# uses the separate top-level `../.env.example`. + +# ----- Inference backend ----- +# "openai_responses" (default): OpenAI-compatible /responses model (reasoning +# models such as gpt-5.x). +# "openai_chat": OpenAI-compatible /chat/completions model (Azure AI Foundry +# deployment, OpenAI, local server, ...). +# "anthropic_messages": Anthropic Messages API (e.g. Claude on Azure AI Foundry). +INFERENCE_BACKEND=openai_responses + +# ----- LLM endpoint (drives the retrieval agent) ----- +# For Azure AI Foundry: CHAT_BASE_URL is the endpoint URL, CHAT_MODEL the +# deployment name. Set CHAT_API_VERSION to use the Azure OpenAI client. +CHAT_BASE_URL=https://your-resource.services.ai.azure.com/openai/v1 +CHAT_API_KEY= +CHAT_MODEL=gpt-5.2 +# CHAT_API_VERSION= +# CHAT_TEMPERATURE=0.7 # sampling temperature (chat backend) +# CHAT_MAX_TOKENS=4096 # max output tokens per model turn +# CHAT_MAX_TURNS=20 # max model<->tool round-trips per search +# CHAT_REASONING_EFFORT= # low|medium|high (openai_responses reasoning models only) +# anthropic_messages only: +# ANTHROPIC_VERSION=2023-06-01 +# ANTHROPIC_AUTH_HEADER=x-api-key + +# ----- Cosmos DB target (required) ----- +ACCOUNT_URI=https://your-cosmos-account.documents.azure.com:443/ +COSMOS_DATABASE=your-database-name +COSMOS_CORPUS_CONTAINER=your-corpus-container +# COSMOS_KEY= # unset -> AzureCliCredential (default) +# COSMOS_USE_DEFAULT_CREDENTIAL=false # true -> use the DefaultAzureCredential chain + +# ----- Embeddings for SearchCorpusTool (required) ----- +# Default embedding endpoint/key/model, used when a corpus is NOT in the registry. +OPENAI_API_KEY=sk-... +OPENAI_EMBEDDING_MODEL=text-embedding-3-small +# EMBED_ENDPOINT= # OpenAI (api.openai.com) if unset. For Azure pass +# # https://.services.ai.azure.com/openai/v1; +# # for a local server pass http://host:port/v1 +# OPENAI_EMBEDDING_DIMENSIONS= # request truncated (MRL) output dims, e.g. 2560 to +# # match a Qwen3-Embedding corpus. Unset = model native. +# EMBED_QUERY_INSTRUCTION= # optional "Instruct:" prefix (some Qwen embedders) + +# ----- Per-corpus embedding registry (optional) ----- +# Map a container to its own account / database / embedding endpoint+model+dims. +# Provide ONE of these. A registry entry references its key via `embed_api_key_env` +# (any env var name you choose, e.g. AZURE_OPENAI_EMBED_API_KEY below). +# CORPUS_REGISTRY_FILE=corpus_registry.json +# CORPUS_REGISTRY={"db/container": {"account_uri": "...", "embed_model": "..."}} +# AZURE_OPENAI_EMBED_API_KEY= # example key referenced by a registry entry + +# ----- Reranker (optional; pick at most one) ----- +# BASETEN_API_KEY= # Baseten Qwen3-Reranker-8B classify +# BASETEN_MODEL_URL=https://model-xyz.api.baseten.co/environments/production/sync +# VLLM_RERANKER_URL=http://127.0.0.1:8011 # local vLLM Qwen3-Reranker /score + +# ----- Retriever budgets & limits (optional) ----- +# COSMOS_RETRIEVER_MAX_TURNS=35 # hard cap on agent turns +# COSMOS_RETRIEVER_THRESHOLD_BUDGET=16384 # soft cap: prune-or-conclude kicks in +# COSMOS_RETRIEVER_TOKEN_BUDGET=32268 # hard cap on transcript tokens +# COSMOS_RETRIEVER_SEARCH_DISPLAY_LIMIT=15 # rows shown per search result +# COSMOS_RETRIEVER_RAW_QUERY_ENABLED=true # expose the read-only execute_query tool +# COSMOS_RETRIEVER_SCHEMA_OVERRIDE= # JSON: document_id_path, chunk_order_path, ... +# Note: the per-tool output clamp (~4096) and spillage fraction (0.5) are code-level +# constants in agent_loop.py (_DEFAULT_TOOL_OUTPUT_BUDGET / _DEFAULT_SPILLAGE_FRACTION), +# not env-configurable. + +# ----- Retriever pool cache (optional) ----- +# COSMOS_RETRIEVER_CACHE_MAX_ENTRIES=32 # max pooled retriever engines (LRU) +# COSMOS_RETRIEVER_CACHE_TTL_SECONDS=900.0 # engine TTL (seconds) before rebuild + +# ----- HTTP server ----- +HOST=0.0.0.0 +PORT=9000 +LOG_LEVEL=info 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..0ec8851 --- /dev/null +++ b/cosmos-retriever/.gitignore @@ -0,0 +1,39 @@ +# --- Python --- +__pycache__/ +*.py[cod] +*$py.class +*.egg-info/ +.eggs/ +build/ +dist/ +.coverage +.coverage.* +htmlcov/ +.pytest_cache/ +.mypy_cache/ +.ruff_cache/ + +# --- Virtual envs --- +.venv/ +venv/ +env/ + +# --- IDE --- +.vscode/ +.idea/ +*.swp + +# --- Secrets / local config --- +.env +.env.local +.env.*.local +.env.* +!.env.example + +# --- Logs / scratch --- +*.log +tmp/ +runs/ + +# --- Build artefacts --- +src/*.egg-info/ diff --git a/cosmos-retriever/README.md b/cosmos-retriever/README.md new file mode 100644 index 0000000..48c5831 --- /dev/null +++ b/cosmos-retriever/README.md @@ -0,0 +1,158 @@ +# Cosmos Retriever (Python helper) + +This package runs a multi-turn search agent against an Azure Cosmos DB corpus and +returns the curated documents as JSON. The agent is model-agnostic: it drives any +OpenAI-compatible endpoint (the `/responses` or `/chat/completions` APIs) or an +Anthropic Messages endpoint. The same code is available three ways, an importable +Python package (`CosmosRetriever`), a FastAPI service +(`python -m cosmos_retriever serve`), and a one-shot CLI +(`python -m cosmos_retriever search`). + +The [Azure Cosmos DB MCP Toolkit](../MCPToolKit/)'s `agentic_search` tool calls +this service's `POST /search` endpoint over HTTP. + +```text + Claude Desktop / AI Foundry / VS Code + │ + │ MCP streamable-HTTP + ▼ + Azure Cosmos DB MCP Toolkit (.NET) + ├─ list_databases / list_collections / ... (8 native tools) + └─ agentic_search ◀─── 9th tool + │ + │ HTTP: POST http://127.0.0.1:9000/search + ▼ + cosmos_retriever (this package, FastAPI + uvicorn) + ├─ TokenBudgetRetrievalSubagent + ├─ SearchCorpus / Grep / ReadDocument / PruneChunks tools + └─ VLLMHarmonyInferenceModel ──► vLLM /v1/completions (token-IDs) + Cosmos DB hybrid RRF + Azure OpenAI embeddings + Qwen3-Reranker (Baseten or local vLLM) +``` + +## Install + +```bash +cd cosmos-retriever +uv venv --python 3.11 .venv +uv pip install --python .venv/bin/python -e ".[dev]" +``` + +## HTTP service + +The MCP Toolkit talks to a long-lived FastAPI service. Start it with: + +```bash +python -m cosmos_retriever serve # binds HOST:PORT (default 0.0.0.0:9000) +``` + +Endpoints: + +| Method & path | Body / response | +|---|---| +| `GET /health` | `{"status": "ok"}` | +| `POST /search` | request `{"query": str, "maxDocuments": int, "database": str?, "container": str?}` → the JSON result below | + +Example request to test a running service (the query and its answer depend on the corpus you configured): + +```bash +curl -s http://127.0.0.1:9000/search \ + -H 'content-type: application/json' \ + -d '{"query": "Who discovered radium?", "maxDocuments": 5}' +``` + +## CLI + +To smoke-test locally, use the command below to query the service with a single +question and print the answer documents. JSON goes to **stdout**, logs go to +**stderr**. + +```bash +python -m cosmos_retriever search \ + --query "Who discovered radium?" \ + --max-documents 5 +``` + +Expected output (same schema returned by `POST /search`): +```json +{ + "query": "Who discovered radium?", + "num_turns": 5, + "elapsed_s": 32.3, + "documents": [ + { "id": "96308__3", "rank": 0, "justification": "...", "text": "..." } + ] +} +``` + +## Configuration + +All settings come from environment variables, or from a `.env` / `.env.local` +file in the `cosmos-retriever/` directory. Precedence is real environment +variables first, then `.env.local`, then `.env`. Use `.env.local` for local +secrets and overrides, it is gitignored. Required settings: + +| Variables | Purpose | +|---|---| +| `INFERENCE_BACKEND`, `CHAT_BASE_URL`, `CHAT_API_KEY`, `CHAT_MODEL` | The backend, endpoint, key, and model for the LLM that drives the agent (see Inference backend below) | +| `ACCOUNT_URI`, `COSMOS_DATABASE`, `COSMOS_CORPUS_CONTAINER` | The Cosmos account, database, and container to search | +| `OPENAI_API_KEY`, `OPENAI_EMBEDDING_MODEL` | The embeddings key and model (set `EMBED_ENDPOINT` for Azure or a local server) | + +Each row is a group of related settings, not alternatives. See +[`.env.example`](.env.example) for the complete list and defaults. + +### Inference backend + +`INFERENCE_BACKEND` selects what drives the retrieval agent: + +| Value | Model | Endpoint vars | +|---|---|---| +| `openai_responses` *(default)* | Any OpenAI-compatible `/responses` model (reasoning models such as gpt-5.x). | `CHAT_BASE_URL`, `CHAT_API_KEY`, `CHAT_MODEL`, optional `CHAT_API_VERSION` | +| `openai_chat` | Any OpenAI-compatible `/chat/completions` model (Azure AI Foundry deployment, OpenAI, local server, ...). | `CHAT_BASE_URL`, `CHAT_API_KEY`, `CHAT_MODEL`, optional `CHAT_API_VERSION` | +| `anthropic_messages` | Any Anthropic Messages API endpoint — e.g. Claude on Azure AI Foundry (served over the Messages API, not OpenAI-shaped). | `CHAT_BASE_URL`, `CHAT_API_KEY`, `CHAT_MODEL`, optional `ANTHROPIC_VERSION`, `ANTHROPIC_AUTH_HEADER` | + +All backends drive the same Cosmos tools, so retrieval quality depends on the +chosen model's tool-use ability. Example (Azure AI Foundry): + +```bash +INFERENCE_BACKEND=openai_chat \ +CHAT_BASE_URL=https://your-resource.services.ai.azure.com/openai/v1 \ +CHAT_API_KEY=... \ +CHAT_MODEL=gpt-4o \ +python -m cosmos_retriever serve +``` + +### Optional reranker + +An independent reranker model can be configured to reorder the retrieved +documents by relevance before they are returned, which improves the quality of +the final ranking. It is optional. Without it, the agent keeps the raw retrieval +order. Configure at most one of: + +- `VLLM_RERANKER_URL`, a local vLLM `/score` endpoint serving Qwen3-Reranker-8B. +- `BASETEN_API_KEY` and `BASETEN_MODEL_URL`, a Baseten Qwen3-Reranker-8B deployment. + +## Layout + +```text +src/cosmos_retriever/ + __init__.py # CosmosRetriever, RetrievalResult, RetrievedDocument + __main__.py # `python -m cosmos_retriever {search,serve}` + server.py # FastAPI app: GET /health + POST /search + retriever.py # CosmosRetriever facade + agent.py # 3 agent classes + prune_chunks_from_trajectory + tools.py # SearchCorpus / Grep / ReadDocument / PruneChunks + trajectory.py # Action / Observation / Trajectory + Harmony rendering + rerank.py # Reranker ABC + Baseten + local-vLLM + inference/ + base.py # AgentInferenceModel ABC + vllm.py # VLLMHarmonyInferenceModel (httpx → /v1/completions) + prompts.py # retrieval subagent system prompt + config.py # RetrieverSettings (pydantic-settings) + utils.py +``` + +## License + +MIT — this package is covered by the repository's top-level [LICENSE](../LICENSE). diff --git a/cosmos-retriever/corpus_registry.json b/cosmos-retriever/corpus_registry.json new file mode 100644 index 0000000..db5ba35 --- /dev/null +++ b/cosmos-retriever/corpus_registry.json @@ -0,0 +1,28 @@ +{ + "_comment": "EXAMPLE corpus registry. Each key is '/' (most specific), or just '', or '' (a DB-wide default). Copy this file, replace every with your own values, and point CORPUS_REGISTRY_FILE at your copy. API keys are NEVER stored here: 'embed_api_key_env' / 'cosmos_key_env' name the environment variables that hold them. Two labelled entries follow: a FULL end-to-end example (every supported field) and a MINIMAL one (only what's needed).", + + "your_database/full_example_container": { + "account_uri": "https://.documents.azure.com:443/", + "database": "your_database", + "embed_base_url": "https://.services.ai.azure.com/openai/v1", + "embed_api_key_env": "MY_EMBED_API_KEY", + "embed_model": "text-embedding-3-small", + "embed_dimensions": 1536, + "embed_query_instruction": "Given a question, retrieve documents that answer it", + "cosmos_key_env": "MY_COSMOS_KEY", + "schema_override": { + "document_id_path": "/docid", + "chunk_id_path": "/id", + "chunk_order_path": "/chunk_idx", + "title_path": "/title", + "source_path": "/source_type", + "use_dunder_codec": true + } + }, + + "minimal_example_container": { + "account_uri": "https://.documents.azure.com:443/", + "embed_model": "text-embedding-3-small" + } +} + diff --git a/cosmos-retriever/docs/AGENTIC_WORKFLOW.md b/cosmos-retriever/docs/AGENTIC_WORKFLOW.md new file mode 100644 index 0000000..6b4363d --- /dev/null +++ b/cosmos-retriever/docs/AGENTIC_WORKFLOW.md @@ -0,0 +1,285 @@ +# 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` (values) + `CosmosPath` allowlist (identifiers) | + +Because the agent is LLM-driven, tool-call arguments are untrusted. Query +injection is blocked on two fronts: user/LLM-supplied **values** (filter values, +ids, vectors) are never concatenated into SQL — they go into bound `@params` +that Cosmos binds separately, so `2020'; DROP TABLE…` is treated as a literal +string. Field **paths/identifiers** can't be parameterized, so every path flows +through `CosmosPath`, which allowlists each segment +(`^[A-Za-z_][A-Za-z0-9_ .\-]*$`), rejects unsafe characters, and escapes on +render. Values are parameterized, identifiers are validated — neither can escape +into executable SQL. + +--- + +## 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/pyproject.toml b/cosmos-retriever/pyproject.toml new file mode 100644 index 0000000..46ea114 --- /dev/null +++ b/cosmos-retriever/pyproject.toml @@ -0,0 +1,75 @@ +[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 = "MIT" } +authors = [{ name = "Cosmos Retriever Contributors" }] +keywords = ["retrieval", "rag", "cosmos-db", "vllm", "agent"] +classifiers = [ + "Development Status :: 4 - Beta", + "Intended Audience :: Developers", + "License :: OSI Approved :: MIT 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", + "pytest>=8,<9", +] + +[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/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..d54a61d --- /dev/null +++ b/cosmos-retriever/src/cosmos_retriever/__main__.py @@ -0,0 +1,103 @@ + +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="Cosmos database name to query (required).", + ) + search.add_argument( + "--container", + default=None, + help="Optional Cosmos container to narrow to; omit to search the whole database.", + ) + + 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 + + # Container is optional: default to the whole database (cross-collection). + container = args.container or "*" + retriever = CosmosRetriever(settings=settings, corpus_name=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/cache.py b/cosmos-retriever/src/cosmos_retriever/cache.py new file mode 100644 index 0000000..e31a7df --- /dev/null +++ b/cosmos-retriever/src/cosmos_retriever/cache.py @@ -0,0 +1,133 @@ +from __future__ import annotations + +import threading +import time +from collections import OrderedDict +from collections.abc import Callable +from dataclasses import dataclass +from typing import Generic, TypeVar + +K = TypeVar("K") +V = TypeVar("V") + + +@dataclass(frozen=True) +class CacheStats: + entries: int + max_entries: int + ttl_seconds: float + hits: int + misses: int + evictions: int + expirations: int + + +class BoundedTTLCache(Generic[K, V]): + """Generic, thread-safe LRU cache with a per-entry time-to-live. + + A reusable utility (it knows nothing about retrieval): entries expire after + ``ttl_seconds`` and the least-recently-used entry is evicted once the cache + exceeds ``max_entries``. An optional ``on_evict`` callback disposes of + values that are dropped (evicted, expired, invalidated, or cleared). + + In this service it is used server-side by ``server._RetrieverPool`` to pool + expensive-to-build ``CosmosRetriever`` engines keyed by corpus scope; it does + **not** cache queries or search results. + """ + + def __init__( self, + *, + max_entries: int = 128, + ttl_seconds: float = 900.0, + time_source: Callable[[], float] = time.monotonic, + on_evict: Callable[[K, V], None] | None = None, + ) -> None: + if max_entries < 1: + raise ValueError("max_entries must be >= 1") + if ttl_seconds <= 0: + raise ValueError("ttl_seconds must be > 0") + self._max = max_entries + self._ttl = ttl_seconds + self._now = time_source + self._on_evict = on_evict + self._lock = threading.RLock() + self._data: OrderedDict[K, tuple[float, V]] = OrderedDict() + self._hits = 0 + self._misses = 0 + self._evictions = 0 + self._expirations = 0 + + def get(self, key: K) -> V | None: + with self._lock: + entry = self._data.get(key) + if entry is None: + self._misses += 1 + return None + stamp, value = entry + if self._now() - stamp >= self._ttl: + del self._data[key] + self._expirations += 1 + self._misses += 1 + self._dispose(key, value) + return None + self._data.move_to_end(key) + self._hits += 1 + return value + + def put(self, key: K, value: V) -> None: + with self._lock: + existing = self._data.get(key) + if existing is not None: + self._data[key] = (self._now(), value) + self._data.move_to_end(key) + if existing[1] is not value: + self._dispose(key, existing[1]) + return + self._data[key] = (self._now(), value) + self._data.move_to_end(key) + while len(self._data) > self._max: + old_key, (_, old_value) = self._data.popitem(last=False) + self._evictions += 1 + self._dispose(old_key, old_value) + + def invalidate(self, key: K) -> bool: + with self._lock: + entry = self._data.pop(key, None) + if entry is None: + return False + self._dispose(key, entry[1]) + return True + + def clear(self) -> None: + with self._lock: + items = list(self._data.items()) + self._data.clear() + for key, (_, value) in items: + self._dispose(key, value) + + def stats(self) -> CacheStats: + with self._lock: + return CacheStats( + entries=len(self._data), + max_entries=self._max, + ttl_seconds=self._ttl, + hits=self._hits, + misses=self._misses, + evictions=self._evictions, + expirations=self._expirations, + ) + + def __len__(self) -> int: + with self._lock: + return len(self._data) + + def _dispose(self, key: K, value: V) -> None: + if self._on_evict is None: + return + try: + self._on_evict(key, value) + except Exception: # noqa: BLE001 + pass + + +__all__ = ["BoundedTTLCache", "CacheStats"] diff --git a/cosmos-retriever/src/cosmos_retriever/config.py b/cosmos-retriever/src/cosmos_retriever/config.py new file mode 100644 index 0000000..40d280f --- /dev/null +++ b/cosmos-retriever/src/cosmos_retriever/config.py @@ -0,0 +1,766 @@ +"""All the settings the service runs on, in one place. + +This module is where the Cosmos Retriever gets its configuration. An operator +sets values through environment variables or a .env file: which inference +backend to use, how to reach the chat and embedding endpoints, which Cosmos +account to query, token budgets, cache sizes, and so on. + +Everything the service +needs to know about its environment is gathered here and validated on the way in, +so a bad value is caught at startup rather than mid request. + +There are two things to read the module as. The settings object is the list of +knobs an operator can turn; each one carries a default and a short description of +what it does, and that set of fields is the whole public surface. The rest of the +module is the quiet machinery that turns those raw values into things the service +can actually use: live connections to Cosmos and the model endpoints, and a +resolved, per-corpus view that pins down the exact account, database, container, +and embedding details for one corpus. + +Callers ask for that resolved view and the +clients, they never touch the wiring behind it. + +A single corpus can also override the shared defaults through a registry, so one +deployment can serve several corpora that live in different places or use +different embedding models. +""" + +from __future__ import annotations + +import hashlib +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 BaseModel, Field, SecretStr, field_validator +from pydantic_settings import BaseSettings, SettingsConfigDict + +from cosmos_retriever.retrieval.schema_override import SchemaOverride + +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: + + """Fully-resolved configuration for a single corpus (internal). + + Produced by :meth:`RetrieverSettings.resolve_corpus` by merging a corpus + registry entry with the server defaults: the concrete Cosmos account / + database / container plus the embedding endpoint, model, and (optional) + output dimensionality and query instruction used to search it. + """ + + 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 + embed_dimensions: int | None = None + + cosmos_key: SecretStr | None = None + schema_override: SchemaOverride | None = None + + +class RetrieverSettings(BaseSettings): + + """Service settings (pydantic-settings), sourced from env vars + ``.env``. + + The *fields* below are the user-facing configuration schema; the *methods* + (``resolve_corpus``, ``build_*_client``, ``apply_structural_overrides``, the + ``use_*_backend`` properties) are internal resolution used by the server. See + the module docstring for the split. + """ + + 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 | None = Field( + default=None, + description=( + "Fallback Cosmos account URI for corpora not found in the registry. " + "Registry entries provide their own account_uri (which wins), so this " + "is only needed when querying an unregistered database/container." + ), + ) + cosmos_database: str | None = Field( + default=None, + description=( + "Cosmos database to query. There is no default: every request must " + "specify the database (the MCP client selects it dynamically)." + ), + ) + cosmos_corpus_container: str | None = Field( + default=None, + description=( + "Cosmos container to query. There is no default: every request must " + "specify the container (the MCP client selects it dynamically)." + ), + ) + cosmos_key: SecretStr | None = None + + openai_api_key: SecretStr | None = None + openai_embedding_model: str | None = None + openai_embedding_dimensions: int | None = Field( + default=None, + description=( + "Requested embedding output dimensionality. Set when the query " + "embedder supports Matryoshka/`dimensions` truncation (e.g. serving " + "Qwen3-Embedding-8B at 2560 dims to match a corpus). Corpus-registry " + "entries may override this per corpus via 'embed_dimensions'." + ), + ) + 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) + + cosmos_retriever_cache_max_entries: int = Field(default=32, ge=1, le=1024) + cosmos_retriever_cache_ttl_seconds: float = Field(default=900.0, gt=0.0) + + cosmos_retriever_schema_override: SchemaOverride | None = Field( + default=None, + description=( + "Default schema override applied to any corpus that does not define " + "its own in the corpus registry. Accepts a JSON object with keys like " + "document_id_path, chunk_id_path, chunk_order_path, title_path, " + "source_path, item_id_path, use_dunder_codec. Omit for pure discovery." + ), + ) + + cosmos_retriever_raw_query_enabled: bool = Field( + default=True, + description=( + "Expose the read-only custom Cosmos SQL query tool (execute_query) to " + "the agent. Set false to remove the tool entirely." + ), + ) + + @field_validator("cosmos_retriever_schema_override", mode="before") + @classmethod + def _coerce_schema_override(cls, v: Any) -> SchemaOverride | None: + return SchemaOverride.coerce(v) + + 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 + + @staticmethod + def _lookup_registry_entry( + registry: dict[str, dict[str, Any]], + database: str | None, + container: str, + ) -> dict[str, Any] | None: + """Resolve a corpus-registry entry for a (database, container) pair. + + Keys are tried from most to least specific so embedding config can be + pinned per database, per container, or per database-wide default: + + 1. "/" — exact collection + 2. "" — same container name in any database + 3. "" — database-wide default (all containers) + + Returns the first matching entry, or ``None`` if nothing matches. + """ + candidates: list[str] = [] + if database: + candidates.append(f"{database}/{container}") + candidates.append(container) + if database: + candidates.append(database) + for key in candidates: + entry = registry.get(key) + if entry is not None: + return entry + return None + + def resolve_corpus(self, container: str | None = None) -> CorpusConfig: + + registry = self._load_registry() + target = container or self.cosmos_corpus_container + if not target: + raise ValueError( + "No Cosmos container specified. Pass 'container' on the request " + "(the MCP client selects it per call); there is no server default." + ) + entry = self._lookup_registry_entry(registry, self.cosmos_database, 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: + database = self.cosmos_database + if not database: + raise ValueError( + "No Cosmos database specified. Pass 'database' on the request " + "(the MCP client selects it per call); there is no server default." + ) + base, key, model = _resolve_default_embed() + if not self.account_uri: + raise ValueError( + f"No Cosmos account for corpus '{target}': it is not in the " + "registry and no fallback ACCOUNT_URI is configured." + ) + return CorpusConfig( + container=target, + account_uri=self.account_uri, + database=database, + embed_base_url=base, + embed_api_key=key, + embed_model=model, + embed_query_instruction=self.embed_query_instruction, + embed_dimensions=self.openai_embedding_dimensions, + cosmos_key=self.cosmos_key, + schema_override=self.cosmos_retriever_schema_override, + ) + + 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) + + database = entry.get("database") or self.cosmos_database + if not database: + raise ValueError( + f"Corpus '{target}' has no database configured and none was passed " + "on the request." + ) + + # Endpoint fallback: an entry may list only a model (no endpoint). In + # that case reuse the server-default embed endpoint + key and simply + # query it for the entry's model. An entry that brings its own endpoint + # keeps its own key (never inherits the default key for a different host). + entry_base = entry.get("embed_base_url") + if entry_base: + embed_base_url = entry_base + embed_api_key = api_key_value + else: + embed_base_url = self.embed_endpoint + embed_api_key = api_key_value or self.openai_api_key + + # Schema override: an entry may define its own; otherwise fall back to + # the server-level default. Absent both -> pure discovery (None). + schema_override = SchemaOverride.coerce(entry.get("schema_override")) + if schema_override is None: + schema_override = self.cosmos_retriever_schema_override + + account_uri = entry.get("account_uri") or self.account_uri + if not account_uri: + raise ValueError( + f"Corpus '{target}' has no account_uri configured (registry entry " + "omits it and no fallback ACCOUNT_URI is set)." + ) + + return CorpusConfig( + container=target, + account_uri=account_uri, + database=database, + embed_base_url=embed_base_url, + embed_api_key=embed_api_key, + embed_model=entry.get("embed_model") or self.openai_embedding_model, + embed_query_instruction=entry.get("embed_query_instruction") + or self.embed_query_instruction, + embed_dimensions=entry.get("embed_dimensions") + if entry.get("embed_dimensions") is not None + else self.openai_embedding_dimensions, + cosmos_key=cosmos_key_value, + schema_override=schema_override, + ) + + def _cosmos_credential(self): + + if os.environ.get("COSMOS_USE_DEFAULT_CREDENTIAL", "").lower() in {"1", "true", "yes"}: + return DefaultAzureCredential() + return AzureCliCredential() + + def build_cosmos_client(self, corpus: CorpusConfig) -> CosmosClient: + + 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 build_cosmos_database(self, corpus: CorpusConfig) -> DatabaseProxy: + + return self.build_cosmos_client(corpus).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 apply_structural_overrides(self, rc: "RuntimeConfig | None") -> "RetrieverSettings": + if rc is None: + return self + updated = self.model_copy(deep=True) + mapping = { + "inference_backend": "inference_backend", + "chat_base_url": "chat_base_url", + "chat_model": "chat_model", + "chat_api_version": "chat_api_version", + "embed_endpoint": "embed_endpoint", + "openai_embedding_model": "openai_embedding_model", + "account_uri": "account_uri", + "embed_query_instruction": "embed_query_instruction", + "schema_override": "cosmos_retriever_schema_override", + "search_display_limit": "cosmos_retriever_search_display_limit", + } + for src, dst in mapping.items(): + value = getattr(rc, src) + if value is not None: + setattr(updated, dst, value) + if rc.chat_api_key is not None: + updated.chat_api_key = SecretStr(rc.chat_api_key) + if rc.openai_api_key is not None: + updated.openai_api_key = SecretStr(rc.openai_api_key) + return updated + + # Maps ServerConfigUpdate field names -> RetrieverSettings attribute names. + _SERVER_UPDATE_MAP = { + "schema_override": "cosmos_retriever_schema_override", + "search_display_limit": "cosmos_retriever_search_display_limit", + "token_budget": "cosmos_retriever_token_budget", + "threshold_budget": "cosmos_retriever_threshold_budget", + "max_turns": "cosmos_retriever_max_turns", + "cache_max_entries": "cosmos_retriever_cache_max_entries", + "cache_ttl_seconds": "cosmos_retriever_cache_ttl_seconds", + } + _SERVER_SECRET_FIELDS = frozenset( + {"chat_api_key", "openai_api_key", "cosmos_key", "baseten_api_key"} + ) + + def apply_server_updates(self, update: "ServerConfigUpdate") -> "RetrieverSettings": + """Return a new settings object with the runtime-mutable server-level + fields overridden. Only fields explicitly set on ``update`` are applied; + secrets are wrapped in ``SecretStr``.""" + updated = self.model_copy(deep=True) + for src, value in update.model_dump(exclude_none=True).items(): + dst = self._SERVER_UPDATE_MAP.get(src, src) + if src in self._SERVER_SECRET_FIELDS: + value = SecretStr(value) + elif src == "schema_override": + value = SchemaOverride.coerce(value) + setattr(updated, dst, value) + return updated + + def redacted_config(self) -> dict[str, Any]: + """Current server-level config for GET /config, with secrets masked.""" + + def _mask(v: SecretStr | None) -> str | None: + return "***set***" if v is not None else None + + return { + "inference_backend": self.inference_backend, + "chat_base_url": self.chat_base_url, + "chat_model": self.chat_model, + "chat_api_version": self.chat_api_version, + "chat_api_key": _mask(self.chat_api_key), + "chat_temperature": self.chat_temperature, + "chat_max_tokens": self.chat_max_tokens, + "chat_max_turns": self.chat_max_turns, + "chat_reasoning_effort": self.chat_reasoning_effort, + "anthropic_version": self.anthropic_version, + "anthropic_auth_header": self.anthropic_auth_header, + "embed_endpoint": self.embed_endpoint, + "openai_embedding_model": self.openai_embedding_model, + "openai_api_key": _mask(self.openai_api_key), + "embed_query_instruction": self.embed_query_instruction, + "account_uri": self.account_uri, + "cosmos_database": self.cosmos_database, + "cosmos_corpus_container": self.cosmos_corpus_container, + "cosmos_key": _mask(self.cosmos_key), + "schema_override": ( + self.cosmos_retriever_schema_override.model_dump() + if self.cosmos_retriever_schema_override is not None + else None + ), + "search_display_limit": self.cosmos_retriever_search_display_limit, + "token_budget": self.cosmos_retriever_token_budget, + "threshold_budget": self.cosmos_retriever_threshold_budget, + "max_turns": self.cosmos_retriever_max_turns, + "baseten_model_url": self.baseten_model_url, + "baseten_api_key": _mask(self.baseten_api_key), + "vllm_reranker_url": self.vllm_reranker_url, + "cache_max_entries": self.cosmos_retriever_cache_max_entries, + "cache_ttl_seconds": self.cosmos_retriever_cache_ttl_seconds, + "log_level": self.log_level, + "host": self.host, + "port": self.port, + } + + 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(), + ) + + +class RuntimeConfig(BaseModel): + model_config = {"extra": "forbid"} + + inference_backend: str | None = None + chat_base_url: str | None = None + chat_api_key: str | None = None + chat_model: str | None = None + chat_api_version: str | None = None + embed_endpoint: str | None = None + openai_api_key: str | None = None + openai_embedding_model: str | None = None + account_uri: str | None = None + embed_query_instruction: str | None = None + schema_override: SchemaOverride | None = None + search_display_limit: int | None = None + + chat_temperature: float | None = None + chat_max_tokens: int | None = None + chat_max_turns: int | None = None + chat_reasoning_effort: str | None = None + anthropic_version: str | None = None + anthropic_auth_header: str | None = None + max_documents: int | None = None + + @field_validator("inference_backend") + @classmethod + def _validate_backend(cls, v: str | None) -> str | None: + if v is None: + return None + normalized = v.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 + + @field_validator("schema_override", mode="before") + @classmethod + def _validate_schema_override(cls, v: Any) -> SchemaOverride | None: + return SchemaOverride.coerce(v) + + def structural_key(self) -> tuple: + def _h(value: str | None) -> str | None: + return hashlib.sha256(value.encode()).hexdigest()[:16] if value else None + + return ( + self.inference_backend, + self.chat_base_url, + _h(self.chat_api_key), + self.chat_model, + self.chat_api_version, + self.embed_endpoint, + _h(self.openai_api_key), + self.openai_embedding_model, + self.account_uri, + self.embed_query_instruction, + self.schema_override.stable_key() if self.schema_override else None, + self.search_display_limit, + ) + + +class ServerConfigUpdate(BaseModel): + """Partial, runtime-mutable server-level configuration accepted by the + PATCH /config admin endpoint. Every field is optional; only provided fields + are applied. Unknown fields are rejected.""" + + model_config = {"extra": "forbid"} + + # inference / LLM defaults + inference_backend: str | None = None + chat_base_url: str | None = None + chat_api_key: str | None = None + chat_model: str | None = None + chat_api_version: str | None = None + chat_temperature: float | None = Field(default=None, ge=0.0, le=2.0) + chat_max_tokens: int | None = Field(default=None, ge=256) + chat_max_turns: int | None = Field(default=None, ge=1, le=200) + chat_reasoning_effort: str | None = None + anthropic_version: str | None = None + anthropic_auth_header: str | None = None + + # embeddings + embed_endpoint: str | None = None + openai_api_key: str | None = None + openai_embedding_model: str | None = None + embed_query_instruction: str | None = None + + # default corpus / Cosmos account + account_uri: str | None = None + cosmos_database: str | None = None + cosmos_corpus_container: str | None = None + cosmos_key: str | None = None + + # retrieval defaults + schema_override: SchemaOverride | None = None + search_display_limit: int | None = Field(default=None, ge=1, le=50) + token_budget: int | None = Field(default=None, ge=4096) + threshold_budget: int | None = Field(default=None, ge=1024) + max_turns: int | None = Field(default=None, ge=1, le=200) + + # reranker + baseten_api_key: str | None = None + baseten_model_url: str | None = None + vllm_reranker_url: str | None = None + + # pool / logging + cache_max_entries: int | None = Field(default=None, ge=1, le=1024) + cache_ttl_seconds: float | None = Field(default=None, gt=0.0) + log_level: str | None = None + + @field_validator("inference_backend") + @classmethod + def _validate_backend(cls, v: str | None) -> str | None: + if v is None: + return None + normalized = v.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 + + @field_validator("schema_override", mode="before") + @classmethod + def _validate_schema_override(cls, v: Any) -> SchemaOverride | None: + return SchemaOverride.coerce(v) + + +@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", + "RuntimeConfig", + "ServerConfigUpdate", + "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..5168baa --- /dev/null +++ b/cosmos-retriever/src/cosmos_retriever/inference/__init__.py @@ -0,0 +1,23 @@ +"""All the code that interacts with the LLMs — plus the token budgeting system +for when it does — lives in this ``inference`` package. The rest of the codebase +is model-agnostic and does not need to know about the LLMs or how they are used. +""" + + +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..531f0fa --- /dev/null +++ b/cosmos-retriever/src/cosmos_retriever/inference/agent_loop.py @@ -0,0 +1,1059 @@ +"""Run a language model as a document-retrieval agent. + +This is the only part of the entire folder that interacts with the LLM endpoint, +everything else is built on top of the interaction here. + +Given a user query and a set of tools, the functions here let a language model +search a corpus over several turns and hand back the documents it judged most +relevant. The model works in a loop: it calls tools (to search, read, or discard +text), reads the results, and repeats until it has gathered enough to answer, at +which point it emits a ranked list of documents. Each result is returned as an +AgentSearchResult. + +The same loop is offered against three different model APIs, one function per +API. Callers pick whichever matches the model they are talking to. all three take +a query and return the same result type, so they are interchangeable from the +outside. + +A running loop keeps its own token budget so a conversation cannot grow without +bound. As the transcript fills up, old search results are trimmed, duplicate +documents are skipped, and oversized tool outputs are shortened. If the budget is +nearly spent, the model is asked to either discard material or give its final +answer. if it is fully spent, further searching is refused. +""" + + +from __future__ import annotations + +import json +import re +import time +from dataclasses import dataclass, field + +import json_repair +import openai +import requests +import structlog + +from cosmos_retriever.prompts import ( + get_retrieval_subagent_budget_exhausted_message, + 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") + + +# Library-level fallback budgets, used only when run_* is called directly without +# explicit values. +_DEFAULT_THRESHOLD_BUDGET = 16384 # soft cap: prompt prune/conclude + restrict to prune +_DEFAULT_TOKEN_BUDGET = 32268 # hard cap +_DEFAULT_TOOL_OUTPUT_BUDGET = 4096 # clamp search/read output when remaining < this +_DEFAULT_SPILLAGE_FRACTION = 0.5 # allowed spillage past threshold before hard reject + +# Marker appended to each observation so pruning never removes past it (upstream parity). +_TOKEN_MARKER_RE = re.compile(r"\n\n\[Token usage:") + + +def _remove_chunks_from_text(text: str, chunk_ids: set[str]) -> str: + """Replace pruned ``# DOCUMENT ID: `` blocks with a tombstone marker. + """ + if not text or not chunk_ids: + return text + matches = list(_DOC_RESULT_RE.finditer(text)) + if not matches: + return text + marker = _TOKEN_MARKER_RE.search(text) + text_end = marker.start() if marker else len(text) + + prune_ranges: list[tuple[int, int, str]] = [] + for idx, match in enumerate(matches): + doc_id = match.group("id") + if doc_id in chunk_ids: + start = match.start() + end = matches[idx + 1].start() if idx + 1 < len(matches) else text_end + prune_ranges.append((start, end, doc_id)) + if not prune_ranges: + return text + + parts: list[str] = [] + last = 0 + for start, end, doc_id in prune_ranges: + parts.append(text[last:start]) + parts.append(f"# DOCUMENT ID: {doc_id} [pruned]\n\n") + last = end + parts.append(text[last:]) + pruned = re.sub(r"\n{3,}", "\n\n", "".join(parts)) + return pruned.strip() + + +class _BudgetController: + """Per-query token-budget + dedup enforcer. + """ + + def __init__( + self, + *, + text_token_counter, + threshold_budget: int = _DEFAULT_THRESHOLD_BUDGET, + token_budget: int = _DEFAULT_TOKEN_BUDGET, + tool_output_budget: int = _DEFAULT_TOOL_OUTPUT_BUDGET, + spillage_fraction: float = _DEFAULT_SPILLAGE_FRACTION, + ) -> None: + self._count = text_token_counter or (lambda s: len(s) // 4) + self.threshold_budget = threshold_budget + self.token_budget = token_budget + self.tool_output_budget = tool_output_budget + spillage = int((token_budget - threshold_budget) * spillage_fraction) + self.rejection_budget = threshold_budget + spillage + # dedup state + self._ids_seen: set[str] = set() + self._doc_id_to_query: dict[str, str] = {} + + # prune state + self._pruned_chunk_ids: set[str] = set() + + # per-step (parallel tool calls in one turn) token tracking + self._step_tokens_used: int = 0 + + def search_overrides(self) -> dict: + return {"ignore_ids": list(self._ids_seen)} + + def record_search(self, returned_chunk_ids, query: str) -> None: + self._ids_seen.update(returned_chunk_ids) + for chunk_id in returned_chunk_ids: + doc_id = chunk_id.split("_")[0] if "_" in chunk_id else chunk_id + self._doc_id_to_query.setdefault(doc_id, query) + + def read_overrides(self, params: dict) -> dict: + doc_id = params.get("doc_id") or params.get("id", "") + if "_" in doc_id: + doc_id = doc_id.split("_")[0] + if doc_id in self._doc_id_to_query: + return {"query": self._doc_id_to_query[doc_id]} + return {} + + def record_prune(self, chunk_ids) -> None: + if isinstance(chunk_ids, (list, tuple, set)): + self._pruned_chunk_ids.update(str(c) for c in chunk_ids) + + def prune_text(self, text: str) -> str: + return _remove_chunks_from_text(text, self._pruned_chunk_ids) + + # ── token accounting (TokenBudgetRetrievalSubagent) ────────────────── + def reset_step(self) -> None: + self._step_tokens_used = 0 + + def add_step_tokens(self, text: str) -> None: + self._step_tokens_used += self._count(text) + + def annotate(self, text: str, current_usage: int) -> str: + return f"{text}\n\n[Token usage: {current_usage}/{self.threshold_budget}]" + + def tool_max_tokens(self, tool_name: str, current_usage: int) -> int | None: + """_call_tool(): clamp search/read output when budget is tight.""" + if tool_name in ("search_corpus", "read_document"): + remaining = self.token_budget - current_usage - self._step_tokens_used + if remaining < self.tool_output_budget: + return max(512, remaining // 2) + return None + + def should_reject(self, tool_name: str, current_usage: int) -> bool: + """_call_tool(): hard-reject non-prune tools past the rejection budget.""" + effective = current_usage + self._step_tokens_used + return effective > self.rejection_budget and tool_name != "prune_chunks" + + def rejection_message(self, current_usage: int) -> str: + effective = current_usage + self._step_tokens_used + return ( + f"Error: Token budget exceeded ({effective}/{self.threshold_budget} tokens). " + "You must use prune_chunks to reduce context size or provide your final answer." + ) + + def over_threshold(self, current_usage: int) -> bool: + return current_usage > self.threshold_budget + + def over_token_budget(self, current_usage: int) -> bool: + return current_usage > self.token_budget + +_CHAT_TOOL_NAMES = ( + "search_corpus", + "grep_corpus", + "read_document", + "prune_chunks", + "execute_query", +) + +_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) + timing: dict[str, 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 {} + parsed: object = raw + # Some models (e.g. gpt-oss via vLLM/Harmony) double-encode tool arguments, + # yielding a JSON string that itself contains JSON. Unwrap up to a few levels. + for _ in range(3): + if isinstance(parsed, dict): + return parsed + if not isinstance(parsed, str): + break + try: + parsed = json.loads(parsed) + except json.JSONDecodeError: + try: + parsed = json_repair.loads(parsed) + 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 _count_messages(messages: list[dict], counter) -> int: + """Token count of a chat-completions transcript (post-prune).""" + total = 0 + for m in messages: + c = m.get("content") + if isinstance(c, str): + total += counter(c) + for tc in m.get("tool_calls", []) or []: + fn = tc.get("function", {}) + for key in ("name", "arguments"): + v = fn.get(key) + if isinstance(v, str): + total += counter(v) + return total + + +# The three sibling run_* entry points (chat / responses / anthropic) deliberately +# duplicate the agent-loop scaffolding rather than sharing one parametrized function +# because each targets a different provider wire protocol: they differ in tool-call +# serialization (Harmony vs. OpenAI vs. Anthropic formats), transcript shape (a +# `messages` list vs. a locally-pruned `input_items` list vs. system-separate +# messages), and per-turn response bookkeeping. Folding them together would produce a +# function dominated by per-backend `if` branches that is hard to read and to test in +# isolation. Keeping them separate trades a little duplicated setup for three linear, +# independently-testable control flows; `CosmosRetriever` picks the right one from the +# configured `inference_backend`. +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, + text_token_counter=None, + threshold_budget: int = _DEFAULT_THRESHOLD_BUDGET, + token_budget: int = _DEFAULT_TOKEN_BUDGET, +) -> AgentSearchResult: + + """Run the multi-turn retrieval agent against an OpenAI-compatible **Chat + Completions** endpoint. + + Drives the search loop using the ``/chat/completions`` wire format: tools are + serialized as Harmony-style function specs, the transcript is a plain + ``messages`` list, and sampling is controlled by ``temperature``. Targets + standard (non-reasoning) chat deployments such as Azure AI Foundry, OpenAI, or + a local vLLM server. Returns an :class:`AgentSearchResult` carrying the ranked + documents, token usage, and trajectory. + """ + + tool_specs = [ + tool.get_format(ProviderFormat.OPENAI_HARMONY) + for name, tool in toolset.tools.items() + if name in _CHAT_TOOL_NAMES + ] + prune_specs = [ + tool.get_format(ProviderFormat.OPENAI_HARMONY) + for name, tool in toolset.tools.items() + if name == "prune_chunks" + ] + + budget = _BudgetController( + text_token_counter=text_token_counter, + threshold_budget=threshold_budget, + token_budget=token_budget, + ) + + 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() + + import collections as _collections + timing = {"llm_s": 0.0, "tools_s": 0.0, "retrieval_s": 0.0, "rerank_s": 0.0} + tool_s = _collections.defaultdict(float) + + for _ in range(max_turns): + # ── prepare_for_inference: prune tool messages, count, decide state + for m in messages: + if m.get("role") == "tool" and isinstance(m.get("content"), str): + m["content"] = budget.prune_text(m["content"]) + current_usage = _count_messages(messages, budget._count) + + turn_messages = messages + turn_specs = tool_specs + if budget.over_threshold(current_usage) and not budget.over_token_budget(current_usage): + turn_messages = messages + [ + { + "role": "user", + "content": get_retrieval_subagent_budget_exhausted_message( + current_usage, budget.threshold_budget + ), + } + ] + turn_specs = prune_specs or tool_specs + + out_cap = max(256, min(max_tokens, budget.token_budget - current_usage)) + + _t = time.perf_counter() + response = client.chat.completions.create( + model=model, + messages=turn_messages, + tools=turn_specs, + tool_choice="auto", + temperature=temperature, + max_tokens=out_cap, + ) + timing["llm_s"] += time.perf_counter() - _t + 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 + + budget.reset_step() + 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}'." + elif budget.should_reject(name, current_usage): + output = budget.rejection_message(current_usage) + logger.warning("tool_rejected_over_budget", tool=name, usage=current_usage) + else: + overrides: dict = {} + if name == "search_corpus": + overrides.update(budget.search_overrides()) + elif name == "read_document": + overrides.update(budget.read_overrides(args)) + clamp = budget.tool_max_tokens(name, current_usage) + if clamp is not None: + overrides["max_tokens"] = clamp + try: + _tt = time.perf_counter() + output, _metadata = tool(args, overrides or None) + _dt = time.perf_counter() - _tt + timing["tools_s"] += _dt + tool_s[name] += _dt + _rs = getattr(_metadata, "retrieval_s", None) + if _rs is not None: + timing["retrieval_s"] += _rs + timing["rerank_s"] += getattr(_metadata, "rerank_s", 0.0) or 0.0 + _collect_doc_text(output, doc_text) + if name == "search_corpus" and _metadata is not None: + budget.record_search( + getattr(_metadata, "returned_chunk_ids", []) or [], str(args.get("query", "")) + ) + elif name == "prune_chunks": + budget.record_prune(args.get("chunk_ids")) + budget.add_step_tokens(output) + except Exception as exc: # noqa: BLE001 — surface tool errors to the model + logger.warning("chat_tool_error", tool=name, error=str(exc)) + output = f"Error executing '{name}': {exc}" + + output = budget.annotate(output, current_usage) + 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) + pool_doc_ids = sorted({cid.split("__")[0] for cid in doc_text}) + + logger.info( + "chat_search_complete", + model=model, + num_turns=num_turns, + num_documents=len(documents), + tool_calls=tool_call_count, + pruned_chunks=len(budget._pruned_chunk_ids), + ) + + return AgentSearchResult( + documents=documents, + num_turns=num_turns, + final_text=final_text, + pool_doc_ids=pool_doc_ids, + usage=usage, + trajectory={"final_docs": [d.id for d in documents]}, + metadata={ + "backend": "openai_chat", + "model": model, + "tool_calls": tool_call_count, + "tool_types_used": ",".join(sorted(tool_types_used)), + }, + timing={ + "llm_s": round(timing["llm_s"], 2), + "tools_s": round(timing["tools_s"], 2), + "retrieval_s": round(timing["retrieval_s"], 2), + "rerank_s": round(timing["rerank_s"], 2), + **{f"tool.{k}_s": round(v, 2) for k, v in tool_s.items()}, + }, + ) + + +def _responses_output_to_call_item(fc) -> dict: + """Render a model function_call output item back into an input item so the + transcript can be resent (we drive the /responses API without + previous_response_id, which is what makes real pruning possible).""" + return { + "type": "function_call", + "call_id": fc.call_id, + "name": fc.name, + "arguments": fc.arguments, + } + + +def _count_items(items: list[dict], counter) -> int: + """Token count of a /responses input-items transcript (post-prune).""" + total = 0 + for it in items: + if not isinstance(it, dict): + continue + for key in ("content", "output", "arguments", "name"): + val = it.get(key) + if isinstance(val, str): + total += counter(val) + return total + + +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, + text_token_counter=None, + threshold_budget: int = _DEFAULT_THRESHOLD_BUDGET, + token_budget: int = _DEFAULT_TOKEN_BUDGET, +) -> AgentSearchResult: + + """Run the multi-turn retrieval agent against an OpenAI **Responses** endpoint. + + Same search loop as :func:`run_chat_search` but speaks the ``/responses`` API + used by reasoning models (gpt-5.x, o-series): tools use the plain OpenAI + function format, the transcript is a local ``input_items`` list (kept + client-side with no ``previous_response_id`` so it can be pruned in place), and + behaviour is tuned via ``reasoning_effort`` instead of ``temperature``. Returns + an :class:`AgentSearchResult`. + """ + + tool_specs = [ + tool.get_format(ProviderFormat.OPENAI) + for name, tool in toolset.tools.items() + if name in _CHAT_TOOL_NAMES + ] + # Prune-only toolset for the over-threshold "prune or conclude" state. + prune_specs = [ + tool.get_format(ProviderFormat.OPENAI) + for name, tool in toolset.tools.items() + if name == "prune_chunks" + ] + + budget = _BudgetController( + text_token_counter=text_token_counter, + threshold_budget=threshold_budget, + token_budget=token_budget, + ) + + 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} + if reasoning_effort: + common["reasoning"] = {"effort": reasoning_effort} + + # Local transcript (no previous_response_id) so we can prune it in place. + input_items: list[dict] = [{"role": "user", "content": prompt}] + + 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]] = [] + + import collections as _collections + timing = {"llm_s": 0.0, "tools_s": 0.0, "retrieval_s": 0.0, "rerank_s": 0.0} + tool_s = _collections.defaultdict(float) + + num_turns = 0 + while True: + # ── prepare_for_inference: prune the transcript, count tokens, decide state + for it in input_items: + if isinstance(it, dict) and it.get("type") == "function_call_output": + it["output"] = budget.prune_text(it["output"]) + current_usage = _count_items(input_items, budget._count) + + turn_input = input_items + turn_specs = tool_specs + if budget.over_threshold(current_usage) and not budget.over_token_budget(current_usage): + # Force prune-or-conclude: inject the budget message and restrict to prune. + turn_input = input_items + [ + { + "role": "user", + "content": get_retrieval_subagent_budget_exhausted_message( + current_usage, budget.threshold_budget + ), + } + ] + turn_specs = prune_specs or tool_specs + + out_cap = max(256, min(max_tokens, budget.token_budget - current_usage)) + call_kwargs = {**common, "tools": turn_specs, "max_output_tokens": out_cap} + + _t = time.perf_counter() + response = client.responses.create(input=turn_input, **call_kwargs) + timing["llm_s"] += time.perf_counter() - _t + num_turns += 1 + _acc_responses_usage(usage, response) + + 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]) + # ── act: execute tools with dedup + reject + clamp; observe: annotate + budget.reset_step() + 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]}") + + input_items.append(_responses_output_to_call_item(fc)) + + tool = toolset.get_tool(name) + if tool is None: + output = f"Error: unknown tool '{name}'." + elif budget.should_reject(name, current_usage): + output = budget.rejection_message(current_usage) + logger.warning("tool_rejected_over_budget", tool=name, usage=current_usage) + else: + overrides: dict = {} + if name == "search_corpus": + overrides.update(budget.search_overrides()) + elif name == "read_document": + overrides.update(budget.read_overrides(args)) + clamp = budget.tool_max_tokens(name, current_usage) + if clamp is not None: + overrides["max_tokens"] = clamp + try: + _tt = time.perf_counter() + output, _metadata = tool(args, overrides or None) + _dt = time.perf_counter() - _tt + timing["tools_s"] += _dt + tool_s[name] += _dt + _rs = getattr(_metadata, "retrieval_s", None) + if _rs is not None: + timing["retrieval_s"] += _rs + timing["rerank_s"] += getattr(_metadata, "rerank_s", 0.0) or 0.0 + _collect_doc_text(output, doc_text) + if name == "search_corpus" and _metadata is not None: + budget.record_search( + getattr(_metadata, "returned_chunk_ids", []) or [], str(args.get("query", "")) + ) + elif name == "prune_chunks": + budget.record_prune(args.get("chunk_ids")) + budget.add_step_tokens(output) + except Exception as exc: # noqa: BLE001 — surface tool errors to the model + logger.warning("responses_tool_error", tool=name, error=str(exc)) + output = f"Error executing '{name}': {exc}" + + output = budget.annotate(output, current_usage) + input_items.append( + {"type": "function_call_output", "call_id": fc.call_id, "output": output} + ) + + 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), + pruned_chunks=len(budget._pruned_chunk_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)), + }, + timing={ + "llm_s": round(timing["llm_s"], 2), + "tools_s": round(timing["tools_s"], 2), + "retrieval_s": round(timing["retrieval_s"], 2), + "rerank_s": round(timing["rerank_s"], 2), + **{f"tool.{k}_s": round(v, 2) for k, v in tool_s.items()}, + }, + ) + + +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 _count_anthropic_messages(messages: list[dict], counter) -> int: + """Token count of an Anthropic Messages transcript (post-prune).""" + total = 0 + for m in messages: + c = m.get("content") + if isinstance(c, str): + total += counter(c) + continue + if isinstance(c, list): + for block in c: + if not isinstance(block, dict): + continue + btype = block.get("type") + if btype == "text" and isinstance(block.get("text"), str): + total += counter(block["text"]) + elif btype == "tool_use": + inp = block.get("input") + if inp is not None: + total += counter(json.dumps(inp)) + if isinstance(block.get("name"), str): + total += counter(block["name"]) + elif btype == "tool_result" and isinstance(block.get("content"), str): + total += counter(block["content"]) + return total + + +def _with_appended_text(message: dict, text: str) -> dict: + """Return a copy of a user message with an extra text block appended. + + Used to inject the budget-exhausted 'prune or conclude' instruction without + adding a second consecutive user message (which the Anthropic API rejects). + """ + content = message.get("content") + if isinstance(content, str): + new_content: list[dict] = [ + {"type": "text", "text": content}, + {"type": "text", "text": text}, + ] + elif isinstance(content, list): + new_content = content + [{"type": "text", "text": text}] + else: + new_content = [{"type": "text", "text": text}] + return {**message, "content": new_content} + + +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, + text_token_counter=None, + threshold_budget: int = _DEFAULT_THRESHOLD_BUDGET, + token_budget: int = _DEFAULT_TOKEN_BUDGET, +) -> AgentSearchResult: + """Run the multi-turn retrieval agent against an **Anthropic Messages** endpoint. + + Same search loop as the other backends but targets the Anthropic Messages API + (e.g. Claude served on Azure AI Foundry) over raw HTTP: it builds the request + URL and auth headers itself, serializes tools in Anthropic format, and passes + the system prompt separately from the ``messages`` list. Returns an + :class:`AgentSearchResult`. + """ + tools = [ + tool.get_format(ProviderFormat.ANTHROPIC) + for name, tool in toolset.tools.items() + if name in _CHAT_TOOL_NAMES + ] + prune_specs = [ + tool.get_format(ProviderFormat.ANTHROPIC) + for name, tool in toolset.tools.items() + if name == "prune_chunks" + ] + budget = _BudgetController( + text_token_counter=text_token_counter, + threshold_budget=threshold_budget, + token_budget=token_budget, + ) + 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]] = [] + + import collections as _collections + timing = {"llm_s": 0.0, "tools_s": 0.0, "retrieval_s": 0.0, "rerank_s": 0.0} + tool_s = _collections.defaultdict(float) + + for _ in range(max_turns): + # ── prepare_for_inference: prune tool results, count tokens, decide state + for m in messages: + if m.get("role") == "user" and isinstance(m.get("content"), list): + for block in m["content"]: + if ( + isinstance(block, dict) + and block.get("type") == "tool_result" + and isinstance(block.get("content"), str) + ): + block["content"] = budget.prune_text(block["content"]) + current_usage = _count_anthropic_messages(messages, budget._count) + + turn_messages = messages + turn_tool_specs = tools + if budget.over_threshold(current_usage) and not budget.over_token_budget(current_usage): + budget_msg = get_retrieval_subagent_budget_exhausted_message( + current_usage, budget.threshold_budget + ) + turn_messages = messages[:-1] + [_with_appended_text(messages[-1], budget_msg)] + turn_tool_specs = prune_specs or tools + + out_cap = max(256, min(max_tokens, budget.token_budget - current_usage)) + payload = { + "model": model, + "max_tokens": out_cap, + "system": system, + "messages": turn_messages, + "tools": turn_tool_specs, + } + _t = time.perf_counter() + response = requests.post(url, json=payload, headers=headers, timeout=timeout_s) + response.raise_for_status() + data = response.json() + timing["llm_s"] += time.perf_counter() - _t + 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]) + budget.reset_step() + 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}'." + elif budget.should_reject(name, current_usage): + output = budget.rejection_message(current_usage) + logger.warning("tool_rejected_over_budget", tool=name, usage=current_usage) + else: + overrides: dict = {} + if name == "search_corpus": + overrides.update(budget.search_overrides()) + elif name == "read_document": + overrides.update(budget.read_overrides(args)) + clamp = budget.tool_max_tokens(name, current_usage) + if clamp is not None: + overrides["max_tokens"] = clamp + try: + _tt = time.perf_counter() + output, _metadata = tool(args, overrides or None) + _dt = time.perf_counter() - _tt + timing["tools_s"] += _dt + tool_s[name] += _dt + _rs = getattr(_metadata, "retrieval_s", None) + if _rs is not None: + timing["retrieval_s"] += _rs + timing["rerank_s"] += getattr(_metadata, "rerank_s", 0.0) or 0.0 + _collect_doc_text(output, doc_text) + if name == "search_corpus" and _metadata is not None: + budget.record_search( + getattr(_metadata, "returned_chunk_ids", []) or [], str(args.get("query", "")) + ) + elif name == "prune_chunks": + budget.record_prune(args.get("chunk_ids")) + budget.add_step_tokens(output) + except Exception as exc: # noqa: BLE001 — surface tool errors to the model + logger.warning("anthropic_tool_error", tool=name, error=str(exc)) + output = f"Error executing '{name}': {exc}" + output = budget.annotate(output, current_usage) + 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)), + }, + timing={ + "llm_s": round(timing["llm_s"], 2), + "tools_s": round(timing["tools_s"], 2), + "retrieval_s": round(timing["retrieval_s"], 2), + "rerank_s": round(timing["rerank_s"], 2), + **{f"tool.{k}_s": round(v, 2) for k, v in tool_s.items()}, + }, + ) + + +__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..84eab25 --- /dev/null +++ b/cosmos-retriever/src/cosmos_retriever/prompts.py @@ -0,0 +1,101 @@ +"""The words the retrieval agent is given to work with. + +This module holds the fixed text that steers the retrieval agent: the main +instructions that tell it what its job is and how to behave, and a short warning +sent when it is running low on room to work. Both are produced by small helper +functions that drop the current query or budget numbers into a ready-made +template, so the surrounding code never has to build these long strings by hand. + +These templates are consumed by the agent loop (see agent_loop.py), which feeds +the main prompt in at the start of a run and injects the budget message when the +transcript grows too large. +""" + +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 + - ExecuteQuery: Author and run a read-only Cosmos DB SELECT query on a chosen database/container. RESTRICTED — see rules below. + + **When to use ExecuteQuery (restricted):** + - Use ExecuteQuery ONLY when the information need cannot be satisfied by semantic/keyword search and specifically requires precise, structured access to the data: exact field-value filters, numeric or date range filters, counts/aggregations (COUNT, SUM, AVG, MIN, MAX), DISTINCT values, GROUP BY, or deterministic ordering by a specific field. + - Do NOT use ExecuteQuery for ordinary topical, conceptual, or semantic questions — use SearchTool, GrepTool, and ReadDocument for those. When in doubt, prefer the search tools. + - It is read-only (SELECT only) and results are truncated. Typical pattern: use it to pinpoint document ids by structured criteria, then ReadDocument those ids or include them in your final ranked output. + + **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..04cd1a1 --- /dev/null +++ b/cosmos-retriever/src/cosmos_retriever/rerank.py @@ -0,0 +1,437 @@ +"""Re-score a shortlist of documents by how well they match the query. + +Search gets you a rough shortlist; reranking sharpens it. Given a query and a +list of candidate documents, the classes here ask a relevance model to score each +one and hand the list back sorted best-first. They can also stop once a token +budget is reached, so a caller only keeps as much text as it has room for. + +Connection details and API keys are pulled from the service config (see +config.py) when not passed in explicitly. The module can also be run directly as +a small command-line demo that reranks a sample query. +""" + +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): + """Abstract base for rerankers: score query–document relevance, then rank. + + Concrete subclasses implement :meth:`_rerank` — call their scoring backend and + return :class:`RerankResult` objects sorted by descending score. The base + supplies the public ``__call__`` template: run ``_rerank``, warn if it is + slow, then apply optional token-budget truncation (``_truncate_results``) so + the returned set fits within ``max_tokens`` (which requires a + ``token_counter``). Implementations in this module are ``BasetenReranker`` and + ``VLLMQwen3Reranker`` (Qwen3-Reranker) and ``ContextualReranker``. + """ + + 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): + + """Qwen3-Reranker served via Baseten's ``classify`` endpoint. + + The yes/no framing in ``PREFIX`` is the reranker's scoring template, not a + free-text answer the code parses. ``classify`` returns a structured + ``{label, score}`` per document; the relevance score is the probability mass + on the ``"yes"`` label (0.0 if absent). The model is never in a generative + mode where it could reply with anything other than the fixed classifier + labels, so there is no brittle string-matching on model prose. + """ + + 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 + # The reranker is a yes/no classifier, not a text generator: the + # relevance score is the probability the judgment token is "yes" + # (softmax over the yes/no logits, computed server-side). We take that + # P("yes") as the score; 0.0 if the label is absent. + 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): + + """Reranker backed by Contextual AI's hosted ``/rerank`` API. + + Sends the query, candidate documents, and an optional ``instruction`` to the + Contextual endpoint and maps each returned ``relevance_score`` onto a + :class:`RerankResult`, sorted by descending score. Unlike the Qwen3 rerankers + this is a managed HTTP service with no local model or logits — the API returns + relevance scores directly. The API key is taken from the constructor argument + or, if omitted, from ``get_config()``. + """ + + 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..d0cd34a --- /dev/null +++ b/cosmos-retriever/src/cosmos_retriever/retrieval/__init__.py @@ -0,0 +1,76 @@ + +from __future__ import annotations + +from cosmos_retriever.retrieval.binding import ( + SchemaOverride, + build_capability_retriever, + build_capability_retriever_from_live, + capabilities_from_metadata, + schema_from_metadata, +) +from cosmos_retriever.retrieval.capabilities import ( + RetrievalCapabilities, + SupportLevel, + VectorCapability, +) +from cosmos_retriever.retrieval.embedding import QueryEmbedder +from cosmos_retriever.retrieval.models import ( + EqualsFilter, + GrepRequest, + InFilter, + NormalizedDocument, + PartitionQueryPolicy, + RangeFilter, + ReadDocumentRequest, + RetrievedItem, + SearchRequest, +) +from cosmos_retriever.retrieval.orchestration import ( + ContainerTarget, + CrossCollectionRetriever, + MultiContainerRetriever, + MultiSearchResult, + fuse_rrf, + select_search_targets, +) +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", + "ContainerTarget", + "CrossCollectionRetriever", + "DunderChunkCodec", + "EqualsFilter", + "GrepRequest", + "InFilter", + "MultiContainerRetriever", + "MultiSearchResult", + "NormalizedDocument", + "PartitionQueryPolicy", + "QueryEmbedder", + "RangeFilter", + "ReadDocumentRequest", + "RetrievalCapabilities", + "RetrievedItem", + "SchemaOverride", + "SearchRequest", + "SupportLevel", + "VectorCapability", + "VectorFieldConfig", + "build_capability_retriever", + "build_capability_retriever_from_live", + "capabilities_from_metadata", + "fuse_rrf", + "schema_from_metadata", + "select_search_targets", +] diff --git a/cosmos-retriever/src/cosmos_retriever/retrieval/binding.py b/cosmos-retriever/src/cosmos_retriever/retrieval/binding.py new file mode 100644 index 0000000..ea11b4c --- /dev/null +++ b/cosmos-retriever/src/cosmos_retriever/retrieval/binding.py @@ -0,0 +1,136 @@ +"""Turn a container's description into a ready to use retriever. + +This module is the bridge between discovery and search. Discovery tells us what a +Cosmos DB container looks like. the functions here take that description and +produce the two things a retriever needs to run against it: a statement of what +the container can do (its capabilities) and a map of where the useful fields live +(its schema). With those in hand, they assemble a fully wired CorpusRetriever. + +The mapping can be nudged by hand. A container's raw settings don't always name +fields the way this system expects, so a SchemaOverride may be supplied to +point at the id, chunk, title, or source fields explicitly. + +Two builders are offered. One takes a container description you already have; the +other reads the description straight from a live container first, then builds on +top of it, so callers with only a connection can get a retriever in one step. +""" + +from __future__ import annotations + +from cosmos_retriever.retrieval.capabilities import ( + RetrievalCapabilities, + SupportLevel, + VectorCapability, +) +from cosmos_retriever.retrieval.discovery.models import ContainerMetadata +from cosmos_retriever.retrieval.discovery.profiler import parse_container_metadata +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, +) +from cosmos_retriever.retrieval.schema_override import SchemaOverride + +__all__ = [ + "SchemaOverride", + "capabilities_from_metadata", + "schema_from_metadata", + "build_capability_retriever", + "build_capability_retriever_from_live", +] + + +def capabilities_from_metadata(metadata: ContainerMetadata) -> RetrievalCapabilities: + indexed = [v for v in metadata.vector_fields if v.indexed and (v.dimensions or 0) > 0] + has_fts = bool(metadata.full_text_paths) + has_vec = bool(indexed) + return RetrievalCapabilities( + vector_fields=[ + VectorCapability( + path=v.path, + dimensions=v.dimensions or 0, + distance_function=v.distance_function or "cosine", + support=SupportLevel.INDEXED, + ) + for v in indexed + ], + full_text_paths=list(metadata.full_text_paths), + partition_key_paths=list(metadata.partition_key_paths), + native_hybrid_supported=has_vec and has_fts, + full_text_supported=has_fts, + vector_supported=has_vec, + efficient_document_lookup_supported=True, + ) + + +def schema_from_metadata( + metadata: ContainerMetadata, + override: SchemaOverride | None = None, +) -> CorpusSchema: + o = override or SchemaOverride() + text_paths = list(metadata.full_text_paths) + vector_fields = [ + VectorFieldConfig( + path=v.path, + dimensions=v.dimensions or 0, + distance_function=v.distance_function or "cosine", + ) + for v in metadata.vector_fields + if v.indexed and (v.dimensions or 0) > 0 + ] + + schema = CorpusSchema( + item_id_path=o.item_id_path or "/id", + text_paths=text_paths, + vector_fields=vector_fields, + document_id_path=o.document_id_path, + chunk_id_path=o.chunk_id_path, + chunk_order_path=o.chunk_order_path, + title_path=o.title_path, + source_path=o.source_path, + partition_key_paths=list(metadata.partition_key_paths), + ) + if o.use_dunder_codec: + schema.identity_codec = DunderChunkCodec() + return schema + + +def build_capability_retriever( + *, + container, + metadata: ContainerMetadata, + embedder: QueryEmbedder | None = None, + override: SchemaOverride | None = None, + partition_policy: PartitionQueryPolicy | None = None, +) -> CorpusRetriever: + return CorpusRetriever( + container=container, + schema=schema_from_metadata(metadata, override), + capabilities=capabilities_from_metadata(metadata), + query_embedder=embedder, + partition_policy=partition_policy or PartitionQueryPolicy(), + ) + + +def build_capability_retriever_from_live( + *, + container, + database: str, + embedder: QueryEmbedder | None = None, + override: SchemaOverride | None = None, + partition_policy: PartitionQueryPolicy | None = None, +) -> CorpusRetriever: + props = container.read() + metadata = parse_container_metadata( + database, container.id, props, props.get("_etag") + ) + return build_capability_retriever( + container=container, + metadata=metadata, + embedder=embedder, + override=override, + partition_policy=partition_policy, + ) 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..8231b7b --- /dev/null +++ b/cosmos-retriever/src/cosmos_retriever/retrieval/compiler.py @@ -0,0 +1,280 @@ +"""Build Cosmos DB SQL queries for the retriever. + +A search request in this system is described in high-level terms: find the +nearest vectors, match this text, keep only rows that pass these filters, skip +these ids, return this many. This module translates such a request into the exact +Cosmos DB SQL string and parameter list needed to run it. + +Requests speak in logical field names: "title", "document_id", "chunk_id" and +so on rather than raw document paths. The compiler looks each name up in the +corpus schema to find where that field actually lives, so the same request works +against containers that store their data differently. + +One method is provided per search style: nearest-vector search, full-text search, +the two combined into a single ranked query, a plain filter-only lookup, and +fetching every chunk of one document. + +They share the same building blocks: the +list of columns to return, the WHERE clause, and the id-exclusion list and all +put user-supplied values into query parameters rather than into the SQL text, so +untrusted input can never alter the query's structure. + +To follow the query's flow through the system, follow the executor file to see the aftermath. +""" + +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, + "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("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 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/discovery/__init__.py b/cosmos-retriever/src/cosmos_retriever/retrieval/discovery/__init__.py new file mode 100644 index 0000000..34fe327 --- /dev/null +++ b/cosmos-retriever/src/cosmos_retriever/retrieval/discovery/__init__.py @@ -0,0 +1,31 @@ +from __future__ import annotations + +from cosmos_retriever.retrieval.discovery.catalog import ResourceCatalog +from cosmos_retriever.retrieval.discovery.connection import ( + CosmosAccountConnection, + CredentialProvider, + DefaultCredentialProvider, +) +from cosmos_retriever.retrieval.discovery.models import ( + CapabilityFlag, + CapabilityProfile, + ContainerMetadata, + VectorIndexInfo, +) +from cosmos_retriever.retrieval.discovery.profiler import ( + CapabilityProfiler, + parse_container_metadata, +) + +__all__ = [ + "CapabilityFlag", + "CapabilityProfile", + "CapabilityProfiler", + "ContainerMetadata", + "CosmosAccountConnection", + "CredentialProvider", + "DefaultCredentialProvider", + "ResourceCatalog", + "VectorIndexInfo", + "parse_container_metadata", +] diff --git a/cosmos-retriever/src/cosmos_retriever/retrieval/discovery/catalog.py b/cosmos-retriever/src/cosmos_retriever/retrieval/discovery/catalog.py new file mode 100644 index 0000000..ee66a42 --- /dev/null +++ b/cosmos-retriever/src/cosmos_retriever/retrieval/discovery/catalog.py @@ -0,0 +1,109 @@ +"""Browse a Cosmos DB account and remember what each container can do. + +This module is the single place the rest of the system asks "what databases and +containers exist, and what kind of search does this one support?". It talks to +Cosmos DB to list databases and containers, reads a container's settings when +asked, and hands back either the raw description or a ready made judgement of +which search strategies the container supports. + +Because reading a container's settings is a network round trip, answers are +cached and shared. Each cached entry expires after a set time, and the cache +keeps only a fixed number of the most recently used containers so it never grows +without bound. The cache can be cleared or refreshed at any time, and every +operation on it is safe to call from multiple threads at once. +""" + +from __future__ import annotations + +import threading +import time +from collections import OrderedDict +from typing import Any, Protocol + +from cosmos_retriever.retrieval.discovery.models import CapabilityProfile, ContainerMetadata +from cosmos_retriever.retrieval.discovery.profiler import ( + CapabilityProfiler, + parse_container_metadata, +) + + +class _Connection(Protocol): + def client(self) -> Any: ... + + +class ResourceCatalog: + def __init__( + self, + connection: _Connection, + *, + ttl_seconds: float = 300.0, + max_entries: int = 256, + profiler: CapabilityProfiler | None = None, + ) -> None: + self._conn = connection + self._ttl = ttl_seconds + self._max = max_entries + self._profiler = profiler or CapabilityProfiler() + self._meta: OrderedDict[tuple[str, str], ContainerMetadata] = OrderedDict() + self._lock = threading.RLock() + + def databases(self) -> list[str]: + return [db["id"] for db in self._conn.client().list_databases()] + + def containers(self, database: str) -> list[str]: + db = self._conn.client().get_database_client(database) + return [c["id"] for c in db.list_containers()] + + def container_metadata( + self, database: str, container: str, *, force: bool = False + ) -> ContainerMetadata: + key = (database, container) + now = time.time() + with self._lock: + hit = self._meta.get(key) + if hit is not None and not force and (now - hit.fetched_at) < self._ttl: + self._meta.move_to_end(key) + return hit + + props, etag = self._read_props(database, container) + meta = parse_container_metadata(database, container, props, etag) + + with self._lock: + self._meta[key] = meta + self._meta.move_to_end(key) + while len(self._meta) > self._max: + self._meta.popitem(last=False) + return meta + + def profile(self, database: str, container: str, *, force: bool = False) -> CapabilityProfile: + return self._profiler.profile( + self.container_metadata(database, container, force=force) + ) + + def refresh(self, database: str | None = None, container: str | None = None) -> None: + with self._lock: + if database is None: + self._meta.clear() + elif container is None: + for k in [k for k in self._meta if k[0] == database]: + self._meta.pop(k, None) + else: + self._meta.pop((database, container), None) + + def invalidate(self, database: str, container: str) -> None: + with self._lock: + self._meta.pop((database, container), None) + + def cached_container_count(self) -> int: + with self._lock: + return len(self._meta) + + def _read_props(self, database: str, container: str) -> tuple[dict[str, Any], str | None]: + c = self._conn.client().get_database_client(database).get_container_client(container) + props = c.read() + etag: str | None = None + try: + etag = c.client_connection.last_response_headers.get("etag") + except Exception: + etag = None + return dict(props), etag diff --git a/cosmos-retriever/src/cosmos_retriever/retrieval/discovery/connection.py b/cosmos-retriever/src/cosmos_retriever/retrieval/discovery/connection.py new file mode 100644 index 0000000..a696ea4 --- /dev/null +++ b/cosmos-retriever/src/cosmos_retriever/retrieval/discovery/connection.py @@ -0,0 +1,48 @@ +from __future__ import annotations + +import os +import threading +from typing import Protocol, runtime_checkable + +from azure.cosmos import CosmosClient +from azure.identity import AzureCliCredential, DefaultAzureCredential + + +@runtime_checkable +class CredentialProvider(Protocol): + def credential(self) -> object: ... + + +class DefaultCredentialProvider: + def credential(self) -> object: + opt_in = os.environ.get("COSMOS_USE_DEFAULT_CREDENTIAL", "").strip().lower() + if opt_in in {"1", "true", "yes"}: + return DefaultAzureCredential() + return AzureCliCredential() + + +class CosmosAccountConnection: + def __init__( + self, + endpoint: str, + credential_provider: CredentialProvider | None = None, + ) -> None: + self._endpoint = endpoint + self._provider = credential_provider or DefaultCredentialProvider() + self._client: CosmosClient | None = None + self._lock = threading.Lock() + + @property + def endpoint(self) -> str: + return self._endpoint + + def client(self) -> CosmosClient: + if self._client is None: + with self._lock: + if self._client is None: + self._client = CosmosClient(self._endpoint, self._provider.credential()) + return self._client + + def close(self) -> None: + with self._lock: + self._client = None diff --git a/cosmos-retriever/src/cosmos_retriever/retrieval/discovery/models.py b/cosmos-retriever/src/cosmos_retriever/retrieval/discovery/models.py new file mode 100644 index 0000000..db0e326 --- /dev/null +++ b/cosmos-retriever/src/cosmos_retriever/retrieval/discovery/models.py @@ -0,0 +1,64 @@ +from __future__ import annotations + +from typing import Any, Literal + +from pydantic import BaseModel, Field + +Provenance = Literal["authoritative", "inferred", "fallback", "unknown"] + + +class VectorIndexInfo(BaseModel): + path: str + dimensions: int | None = None + distance_function: str | None = None + data_type: str | None = None + index_type: str | None = None + indexed: bool = False + + +class ContainerMetadata(BaseModel): + database: str + container: str + etag: str | None = None + fetched_at: float + partition_key_paths: list[str] = Field(default_factory=list) + included_paths: list[str] = Field(default_factory=list) + excluded_paths: list[str] = Field(default_factory=list) + full_text_paths: list[str] = Field(default_factory=list) + full_text_policy_paths: list[str] = Field(default_factory=list) + vector_fields: list[VectorIndexInfo] = Field(default_factory=list) + + +class CapabilityFlag(BaseModel): + value: bool + provenance: Provenance = "authoritative" + + +class CapabilityProfile(BaseModel): + database: str + container: str + fetched_at: float + partition_key_paths: list[str] = Field(default_factory=list) + full_text_paths: list[str] = Field(default_factory=list) + vector_fields: list[VectorIndexInfo] = Field(default_factory=list) + can_full_text: CapabilityFlag + can_vector: CapabilityFlag + can_native_hybrid: CapabilityFlag + can_item_lookup: CapabilityFlag + recommended_strategies: list[str] = Field(default_factory=list) + confidence: float = 1.0 + + def summary(self) -> dict[str, Any]: + return { + "database": self.database, + "container": self.container, + "full_text": self.can_full_text.value, + "vector": self.can_vector.value, + "hybrid": self.can_native_hybrid.value, + "item_lookup": self.can_item_lookup.value, + "full_text_paths": list(self.full_text_paths), + "vector_paths": [v.path for v in self.vector_fields], + "partition_key_paths": list(self.partition_key_paths), + "recommended_strategies": list(self.recommended_strategies), + "confidence": self.confidence, + } diff --git a/cosmos-retriever/src/cosmos_retriever/retrieval/discovery/profiler.py b/cosmos-retriever/src/cosmos_retriever/retrieval/discovery/profiler.py new file mode 100644 index 0000000..556aa09 --- /dev/null +++ b/cosmos-retriever/src/cosmos_retriever/retrieval/discovery/profiler.py @@ -0,0 +1,107 @@ +"""Work out what kinds of search a Cosmos DB container supports. + +Cosmos DB describes each container with a bundle of raw settings: how it is +partitioned, which fields are indexed, and whether it is set up for full-text or +vector search. This module turns that raw description into two tidy answers: +what the container looks like, and what you can do with it. + +The work happens in two steps. First the raw settings are read and organised +into a plain ContainerMetadata record. Then that record is inspected to +decide which search strategies the container can actually serve: vector search, +full-text search, a native combination of the two, or a simple lookup by id +and the answer is returned as a CapabilityProfile. +""" + +from __future__ import annotations + +import time +from typing import Any + +from cosmos_retriever.retrieval.discovery.models import ( + CapabilityFlag, + CapabilityProfile, + ContainerMetadata, + VectorIndexInfo, +) + + +def parse_container_metadata( + database: str, + container: str, + props: dict[str, Any], + etag: str | None = None, + *, + fetched_at: float | None = None, +) -> ContainerMetadata: + pk_paths = list((props.get("partitionKey") or {}).get("paths") or []) + + idx = props.get("indexingPolicy") or {} + included = [p.get("path") for p in (idx.get("includedPaths") or []) if p.get("path")] + excluded = [p.get("path") for p in (idx.get("excludedPaths") or []) if p.get("path")] + ft_index_paths = [p.get("path") for p in (idx.get("fullTextIndexes") or []) if p.get("path")] + vec_index = {p.get("path"): p for p in (idx.get("vectorIndexes") or []) if p.get("path")} + + ftp = props.get("fullTextPolicy") or {} + ft_policy_paths = [p.get("path") for p in (ftp.get("fullTextPaths") or []) if p.get("path")] + + vep = props.get("vectorEmbeddingPolicy") or {} + vectors: list[VectorIndexInfo] = [] + for emb in vep.get("vectorEmbeddings") or []: + path = emb.get("path") + if not path: + continue + vi = vec_index.get(path) + vectors.append( + VectorIndexInfo( + path=path, + dimensions=emb.get("dimensions"), + distance_function=emb.get("distanceFunction"), + data_type=emb.get("dataType"), + index_type=(vi or {}).get("type"), + indexed=vi is not None, + ) + ) + + return ContainerMetadata( + database=database, + container=container, + etag=etag, + fetched_at=fetched_at if fetched_at is not None else time.time(), + partition_key_paths=pk_paths, + included_paths=included, + excluded_paths=excluded, + full_text_paths=ft_index_paths, + full_text_policy_paths=ft_policy_paths, + vector_fields=vectors, + ) + + +class CapabilityProfiler: + def profile(self, metadata: ContainerMetadata) -> CapabilityProfile: + indexed_vectors = [v for v in metadata.vector_fields if v.indexed] + has_fts = bool(metadata.full_text_paths) + has_vec = bool(indexed_vectors) + + strategies: list[str] = [] + if has_vec and has_fts: + strategies.append("native_hybrid") + if has_vec: + strategies.append("vector") + if has_fts: + strategies.append("full_text") + strategies.append("item_lookup") + + return CapabilityProfile( + database=metadata.database, + container=metadata.container, + fetched_at=metadata.fetched_at, + partition_key_paths=metadata.partition_key_paths, + full_text_paths=metadata.full_text_paths, + vector_fields=indexed_vectors, + can_full_text=CapabilityFlag(value=has_fts), + can_vector=CapabilityFlag(value=has_vec), + can_native_hybrid=CapabilityFlag(value=has_vec and has_fts), + can_item_lookup=CapabilityFlag(value=True), + recommended_strategies=strategies, + confidence=1.0, + ) 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..e191bf4 --- /dev/null +++ b/cosmos-retriever/src/cosmos_retriever/retrieval/document_resolvers.py @@ -0,0 +1,153 @@ + +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.normalization import assemble_text, row_text_fields +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) + aliases = compiled.projected_aliases + return NormalizedDocument( + document_id=item_id, + + chunk_texts=[assemble_text(row_text_fields(r, aliases)) 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)) + aliases = compiled.projected_aliases + return NormalizedDocument( + + document_id=doc_id, + chunk_texts=[assemble_text(row_text_fields(r, aliases)) 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)) + aliases = compiled.projected_aliases + + return NormalizedDocument( + document_id=doc_id, + + chunk_texts=[assemble_text(row_text_fields(r, aliases)) 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..4cb8c79 --- /dev/null +++ b/cosmos-retriever/src/cosmos_retriever/retrieval/embedding.py @@ -0,0 +1,29 @@ +from __future__ import annotations + +import openai + + +class QueryEmbedder: + + def __init__( + self, + client: openai.OpenAI, + model: str, + query_instruction: str | None = None, + dimensions: int | None = None, + ) -> None: + self._client = client + self._model = model + self._instruction = query_instruction + self._dimensions = dimensions + + def embed(self, text: str) -> list[float]: + if self._instruction: + text = f"Instruct: {self._instruction}\nQuery: {text}" + kwargs: dict[str, object] = {} + if self._dimensions is not None: + kwargs["dimensions"] = self._dimensions + resp = self._client.embeddings.create( + model=self._model, input=[text], encoding_format="float", **kwargs + ) + 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..a694cca --- /dev/null +++ b/cosmos-retriever/src/cosmos_retriever/retrieval/executor.py @@ -0,0 +1,107 @@ +"""Run a compiled query against Cosmos DB and hand back the rows. + +Once a search has been turned into SQL (for details on how that is done, refer to +the compiler), this module is what actually sends it to Cosmos DB and collects the +results. It is the last step before raw rows flow back into the retriever. + +Running a query here comes with three safeguards. Transient failures are retried +automatically with growing pauses between attempts, so a momentary hiccup doesn't +sink a request. The number of queries allowed to run at the same time is capped, +so a burst of searches can't overwhelm the account; the cap defaults to a sensible +value and can be raised or lowered through an environment variable. And any query +that takes unusually long is logged, to make slow spots easy to spot. +""" + +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..3954f89 --- /dev/null +++ b/cosmos-retriever/src/cosmos_retriever/retrieval/expressions.py @@ -0,0 +1,41 @@ +from __future__ import annotations + +import re + +_TOKEN_RE = re.compile(r"\w+", re.UNICODE) +# English-only stopword list. Full-text queries are assumed to be English; this +# list only removes English function words so they don't dominate FullTextScore. +# Non-English queries still work — `_TOKEN_RE` is Unicode-aware, so their tokens +# are tokenized/lower-cased/de-duplicated normally; their function words simply +# aren't stripped (mildly noisier, but RRF/FullTextScore tolerate it). The only +# degenerate case is an all-English-stopword query, which can reduce to zero +# terms. Add per-language lists here if broader language support is needed. +_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..9c467e3 --- /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 row_text_fields(row: dict[str, Any], aliases: dict[str, str]) -> dict[str, str]: + + out: dict[str, str] = {} + for key, value in row.items(): + if key.startswith("txt_") and key in aliases: + out[aliases[key]] = value or "" + return out + + +def assemble_text(text_fields: dict[str, str], names: list[str] | None = None) -> str: + + selected = [n for n in names if n in text_fields] if names else list(text_fields) + if not selected: + return "" + if len(selected) == 1: + return text_fields.get(selected[0], "") or "" + return "\n\n".join(f"[{n}]\n{text_fields.get(n, '') or ''}" for n in selected) + + +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, +) -> 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 = row_text_fields(row, aliases) + display = assemble_text(text_fields, queried_text_fields) + 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/orchestration.py b/cosmos-retriever/src/cosmos_retriever/retrieval/orchestration.py new file mode 100644 index 0000000..820eb5b --- /dev/null +++ b/cosmos-retriever/src/cosmos_retriever/retrieval/orchestration.py @@ -0,0 +1,245 @@ +from __future__ import annotations + +import time +from collections.abc import Callable, Sequence +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass, field +from typing import NamedTuple + +import structlog + +from cosmos_retriever.retrieval.models import ( + GrepRequest, + NormalizedDocument, + ReadDocumentRequest, + RetrievedItem, + SearchRequest, +) +from cosmos_retriever.retrieval.retriever import CorpusRetriever + +logger = structlog.get_logger("cosmos_retriever.orchestration") + +RRF_K = 60 + + +class ContainerTarget(NamedTuple): + database: str + container: str + + +RetrieverResolver = Callable[[ContainerTarget], CorpusRetriever] + + +@dataclass +class MultiSearchResult: + items: list[RetrievedItem] + searched: list[ContainerTarget] = field(default_factory=list) + errors: dict[str, str] = field(default_factory=dict) + per_container_counts: dict[str, int] = field(default_factory=dict) + elapsed_s: float = 0.0 + + +def _qualify(target: ContainerTarget, item: RetrievedItem) -> str: + return f"{target.database}/{target.container}:{item.item_id}" + + +def fuse_rrf( + ranked_lists: Sequence[tuple[ContainerTarget, list[RetrievedItem]]], + *, + k: int = RRF_K, + limit: int | None = None, +) -> list[RetrievedItem]: + scores: dict[str, float] = {} + chosen: dict[str, RetrievedItem] = {} + for target, items in ranked_lists: + for position, item in enumerate(items): + key = _qualify(target, item) + scores[key] = scores.get(key, 0.0) + 1.0 / (k + position) + if key not in chosen: + tagged = item.model_copy(deep=True) + tagged.metadata = { + **tagged.metadata, + "container": target.container, + "database": target.database, + } + chosen[key] = tagged + ordered = sorted(scores.items(), key=lambda kv: kv[1], reverse=True) + fused: list[RetrievedItem] = [] + for rank, (key, score) in enumerate(ordered): + item = chosen[key] + item.rank = rank + item.raw_scores = {**item.raw_scores, "rrf": score} + fused.append(item) + if limit is not None and len(fused) >= limit: + break + return fused + + +class MultiContainerRetriever: + def __init__( + self, + resolver: RetrieverResolver, + *, + max_workers: int = 8, + ) -> None: + if max_workers < 1: + raise ValueError("max_workers must be >= 1") + self._resolve = resolver + self._max_workers = max_workers + + def search( + self, + targets: Sequence[ContainerTarget], + request: SearchRequest, + *, + per_container_limit: int | None = None, + final_limit: int | None = None, + ) -> MultiSearchResult: + targets = list(dict.fromkeys(targets)) + if not targets: + return MultiSearchResult(items=[]) + + per_request = request + if per_container_limit is not None: + per_request = request.model_copy(update={"limit": per_container_limit}) + + start = time.perf_counter() + result = MultiSearchResult(items=[]) + ranked_lists: list[tuple[ContainerTarget, list[RetrievedItem]]] = [] + + workers = min(self._max_workers, len(targets)) + with ThreadPoolExecutor(max_workers=workers) as pool: + futures = { + pool.submit(self._search_one, target, per_request): target + for target in targets + } + for future in futures: + target = futures[future] + label = f"{target.database}/{target.container}" + try: + items = future.result() + except Exception as exc: # noqa: BLE001 + result.errors[label] = f"{type(exc).__name__}: {exc}" + logger.warning( + "container_search_failed", + database=target.database, + container=target.container, + error=str(exc), + ) + continue + result.searched.append(target) + result.per_container_counts[label] = len(items) + ranked_lists.append((target, items)) + + result.items = fuse_rrf(ranked_lists, limit=final_limit) + result.elapsed_s = round(time.perf_counter() - start, 3) + return result + + def _search_one( + self, target: ContainerTarget, request: SearchRequest + ) -> list[RetrievedItem]: + retriever = self._resolve(target) + return retriever.search(request) + + +def select_search_targets( + catalog, + database: str, + *, + containers: Sequence[str] | None = None, + require_capability: bool = True, +) -> list[ContainerTarget]: + names = list(containers) if containers is not None else catalog.containers(database) + targets: list[ContainerTarget] = [] + for name in names: + if require_capability: + profile = catalog.profile(database, name) + if not (profile.can_full_text.value or profile.can_vector.value): + continue + targets.append(ContainerTarget(database=database, container=name)) + return targets + + +class CrossCollectionRetriever: + """Duck-types the ``CorpusRetriever`` interface (``schema`` / ``search`` / + ``grep_candidates`` / ``read_document``) but fans every operation out across + all target collections of a database and fuses the hits with RRF. Drops in + wherever a single-container ``CorpusRetriever`` is expected (e.g. the agent + ``ToolSet``).""" + + def __init__( + self, + targets: Sequence[ContainerTarget], + retrievers: dict[ContainerTarget, CorpusRetriever], + *, + per_container_limit: int | None = None, + max_workers: int = 16, + ) -> None: + targets = list(dict.fromkeys(targets)) + if not targets: + raise ValueError("CrossCollectionRetriever requires at least one target") + self._targets = targets + self._retrievers = retrievers + self._per_container = per_container_limit + self._max_workers = max(1, min(max_workers, len(targets))) + self._mcr = MultiContainerRetriever( + lambda t: self._retrievers[t], max_workers=self._max_workers + ) + # Representative schema so the tools can build their JSON tool schema. + self.schema = retrievers[targets[0]].schema + + def search(self, request: SearchRequest) -> list[RetrievedItem]: + # Track the fused output depth (final_limit) by default: truncating a + # collection below what fusion can keep silently drops recoverable gold + # docs that rank between per_container_limit and final_limit within + # their own collection. An explicit per_container_limit still overrides. + per_container = self._per_container if self._per_container is not None else request.limit + result = self._mcr.search( + self._targets, + request, + per_container_limit=per_container, + final_limit=request.limit, + ) + return result.items + + def grep_candidates(self, request: GrepRequest) -> list[RetrievedItem]: + out: list[RetrievedItem] = [] + with ThreadPoolExecutor(max_workers=self._max_workers) as pool: + futures = [ + pool.submit(self._retrievers[t].grep_candidates, request) + for t in self._targets + ] + for fut in futures: + try: + out.extend(fut.result()) + except Exception as exc: # noqa: BLE001 + logger.warning("cross_grep_failed", error=str(exc)) + return out[: request.candidate_limit] + + def read_document(self, request: ReadDocumentRequest) -> NormalizedDocument: + # A document id can live in any collection; probe until one yields text. + first: NormalizedDocument | None = None + for target in self._targets: + try: + doc = self._retrievers[target].read_document(request) + except Exception: # noqa: BLE001 + continue + if first is None: + first = doc + if doc.assembled.strip(): + return doc + if first is not None: + return first + return self._retrievers[self._targets[0]].read_document(request) + + +__all__ = [ + "ContainerTarget", + "CrossCollectionRetriever", + "MultiContainerRetriever", + "MultiSearchResult", + "RRF_K", + "RetrieverResolver", + "fuse_rrf", + "select_search_targets", +] 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..ed0f009 --- /dev/null +++ b/cosmos-retriever/src/cosmos_retriever/retrieval/planner.py @@ -0,0 +1,145 @@ + +"""Pick the search strategy that best fits a container. + +Different containers can do different things: some are set up for vector search, +some for full text search, some for both, some for neither. This module looks at +what a container actually supports and chooses how a given request should be run, +handing back a ready-to-use strategy object. + +The choice is driven by two inputs. The corpus schema says which fields exist and +where they live; the capabilities (see the capabilities file) say which of those +fields are truly searchable, and how. + +A request may also pin a mode outright, in which case the planner either follows +it or raises if the container can't satisfy it. Left on auto, it prefers the +richest option available and falls back gracefully: native hybrid if the +container ranks vector and text together for you, otherwise combining the two +results itself, then whichever single mode is available, and finally a bounded +scan if the policy permits one. + +The planner only decides, it does not run anything. The strategy it returns (one +of the classes defined in the strategies file) is what goes on to build SQL +through the compiler and run it through the executor. +""" + +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 + if not names: + # No specific field requested. Full-text is available as long as the + # schema exposes at least one full-text-capable field; the concrete + # field(s) must be chosen by the caller at execution time. + return any( + self.capabilities.has_full_text_path(p) for p in self.schema.text_paths + ) + 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..3fde035 --- /dev/null +++ b/cosmos-retriever/src/cosmos_retriever/retrieval/retriever.py @@ -0,0 +1,104 @@ +""" + +the front door to corpus search. + +This module ties the retrieval pieces together into one object that the rest of +the system talks to. + +Give it a container plus a description of that container (its +schema and capabilities) and it can answer three kinds of request: search for the +most relevant items, find candidate rows for a grep, and read back a whole +document by id. + +Under the hood it wires up and drives the components: the planner picks a +strategy for each request, the compiler turns that into SQL, and the executor +runs it (each covered in its own file). + +When a request needs a query embedding +and none was supplied, it asks the configured embedder to produce one first, +document reads are handed off to a resolver built for the container's layout. All +of that is hidden behind the three methods, so callers never touch the moving +parts directly. + +This is the object that binding.py assembles and hands out, and it is where a +search request begins its journey through the system. +""" + +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..230706f --- /dev/null +++ b/cosmos-retriever/src/cosmos_retriever/retrieval/schema.py @@ -0,0 +1,173 @@ + +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] = [] + 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] = [] + for v in self.vector_fields: + if v.dimensions <= 0: + errors.append(f"vector field {v.path} has non-positive dimensions") + if not self.text_paths and not self.vector_fields: + errors.append("schema must declare at least one text or vector field") + 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 resolve_text_fields(self, names: list[str] | None) -> list[CosmosPath]: + m = self.text_field_map() + if not names: + if len(m) == 1: + return [next(iter(m.values()))] + if not m: + return [] + raise UnknownField( + "multiple text fields are available; specify one or more of " + f"{sorted(m)}" + ) + 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] = [] + if tm: + 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)) + text_names = list(tm) + default_v = next(iter(vm), None) + if len(text_names) > 1: + lines.append( + "You must choose which text field(s) to search on each call " + f"(available: {text_names})." + ) + elif len(text_names) == 1 and default_v: + lines.append( + f"Default when unspecified: hybrid over text='{text_names[0]}' " + f"+ vector='{default_v}'." + ) + elif len(text_names) == 1: + lines.append(f"Default when unspecified: full-text over '{text_names[0]}'.") + elif default_v: + lines.append(f"Default when unspecified: vector search over '{default_v}'.") + else: + lines.append("Default when unspecified: structured item lookup.") + return "\n".join(lines) diff --git a/cosmos-retriever/src/cosmos_retriever/retrieval/schema_override.py b/cosmos-retriever/src/cosmos_retriever/retrieval/schema_override.py new file mode 100644 index 0000000..a66e249 --- /dev/null +++ b/cosmos-retriever/src/cosmos_retriever/retrieval/schema_override.py @@ -0,0 +1,73 @@ +from __future__ import annotations + +import json +from typing import Any + +from pydantic import BaseModel + + +class SchemaOverride(BaseModel): + """User-supplied hints layered on top of live schema discovery. + + Discovery reads the container's physical schema (vector fields, full-text + paths, partition keys) but cannot infer the *semantic role* of fields. This + override names them so chunks can be grouped back into parent documents: + + - ``document_id_path`` path of the parent-document id (groups chunks) + - ``chunk_id_path`` path of the per-chunk id + - ``chunk_order_path`` path of the chunk ordinal (orders chunks) + - ``title_path`` / ``source_path`` optional display fields + - ``item_id_path`` path of the item id (defaults to ``/id``) + - ``use_dunder_codec`` chunk ids are encoded ``__`` + + These paths are **not** canonical Cosmos fields and are not guaranteed to + exist. Cosmos DB is schema-agnostic: the only field present on every + document is ``/id`` (plus system props like ``_rid``/``_ts``), which is why + ``item_id_path`` defaults to ``/id``. The remaining paths are + application-specific — they exist only if the ingestion pipeline created + them — and are only meaningful for *chunked* (RAG) corpora where one logical + document is split across many records. For one-record-per-document data, + omit them: each item is then treated as its own document + (see ``CorpusSchema.is_item_document_mode``). ``use_dunder_codec`` is a + naming convention, not a Cosmos feature. + + All fields are optional; omit the object entirely for pure discovery. + """ + + model_config = {"extra": "forbid"} + + item_id_path: str | None = None + document_id_path: str | None = None + chunk_id_path: str | None = None + chunk_order_path: str | None = None + title_path: str | None = None + source_path: str | None = None + use_dunder_codec: bool = False + + @classmethod + def coerce(cls, value: Any) -> "SchemaOverride | None": + """Build a SchemaOverride from ``None``, a dict, a JSON string, or an + existing instance. Returns ``None`` for empty/blank input.""" + if value is None: + return None + if isinstance(value, SchemaOverride): + return value + if isinstance(value, str): + text = value.strip() + if not text or text.lower() in {"none", "null"}: + return None + try: + value = json.loads(text) + except json.JSONDecodeError as exc: + raise ValueError(f"schema_override is not valid JSON: {exc}") from exc + if isinstance(value, dict): + if not value: + return None + return cls(**value) + raise ValueError( + f"schema_override must be a JSON object (or null), got {type(value).__name__}." + ) + + def stable_key(self) -> str: + """Deterministic string form for cache keys.""" + return json.dumps(self.model_dump(), sort_keys=True) 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..5a9b3e6 --- /dev/null +++ b/cosmos-retriever/src/cosmos_retriever/retrieval/strategies.py @@ -0,0 +1,227 @@ +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, + ) + + +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, + ) + +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, + ) + + +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, + projected_aliases=compiled.projected_aliases, + ) + + + +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, + ) diff --git a/cosmos-retriever/src/cosmos_retriever/retriever.py b/cosmos-retriever/src/cosmos_retriever/retriever.py new file mode 100644 index 0000000..87d5ed5 --- /dev/null +++ b/cosmos-retriever/src/cosmos_retriever/retriever.py @@ -0,0 +1,434 @@ + +"""The top-level entry point for running a search. + +This module is what an application actually holds onto: a single CosmosRetriever +that you point at a corpus and ask questions. It gathers everything a run needs +from the service settings (see config.py), which corpus to search, how to reach +Cosmos and the models, which reranker to use, wires those pieces together once at +construction time, and then answers searches through one search method. + +Each search is carried out by the multi-turn retrieval agent, which this class +hands off to depending on the configured model backend (chat, responses, or +Anthropic. + +The loops themselves live in agent_loop.py). It gives the agent the +set of tools it can call (the ToolSet), an approximate token counter used for +budgeting, and a reranker, then packages whatever the agent returns into a plain +RetrievalResult: the ranked documents plus timing, token usage, and a trace +of what happened. + +It also handles two shapes of corpus. A normal corpus is one Cosmos container, a +corpus named "*" means search every searchable container in the database at once, +in which case a cross-collection retriever is built across all of them. + +Callers only ever touch three things from here: CosmosRetriever and the two +result records, RetrievalResult and RetrievedDocument. +""" + +from __future__ import annotations + +import time +from dataclasses import dataclass, field + +import structlog +import tiktoken + +from cosmos_retriever.config import CorpusConfig, RetrieverSettings, RuntimeConfig, get_settings +from cosmos_retriever.rerank import BasetenReranker, Reranker, VLLMReranker +from cosmos_retriever.retrieval import ( + CrossCollectionRetriever, + QueryEmbedder, + build_capability_retriever_from_live, + select_search_targets, +) +from cosmos_retriever.retrieval.discovery import ResourceCatalog +from cosmos_retriever.tools import ToolSet + +logger = structlog.get_logger("cosmos_retriever.retriever") + + +class _StaticCosmosConnection: + """Adapts an already-built CosmosClient to the ResourceCatalog connection + protocol (a single ``client()`` accessor), so catalog reads reuse the same + client instead of opening another.""" + + def __init__(self, client) -> None: + self._client = client + + def client(self): + return self._client + + +@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) + + +@dataclass(frozen=True) +class _Exec: + turns: int + temperature: float + max_tokens: int + reasoning_effort: str | None + anthropic_version: str + anthropic_auth_header: str + + +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) + + # Single approximate tokenizer used only for budget/truncation heuristics + # (not billing or context-limit enforcement). o200k_harmony is exact for + # gpt-oss/Harmony models and within a few percent for other o200k-family + # models; non-OpenAI models (Claude, Qwen) are approximated. Budgets are + # set with headroom to absorb the drift. + self._tiktoken = tiktoken.get_encoding("o200k_harmony") + self._reranker = reranker or self._build_default_reranker() + + cosmos_client = self.settings.build_cosmos_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 + + schema_override = self.corpus.schema_override + + self.database_wide = self.corpus.container.strip() == "*" + if self.database_wide: + self.toolset: ToolSet = ToolSet.build( + retriever=self._build_cross_collection_retriever(cosmos_client, schema_override), + reranker=self._reranker, + token_counter=self._text_token_counter, + search_display_limit=self.settings.cosmos_retriever_search_display_limit, + cosmos_client=cosmos_client, + enable_raw_query=self.settings.cosmos_retriever_raw_query_enabled, + ) + else: + self.toolset = ToolSet.build( + cosmos_database=cosmos_client.get_database_client(self.corpus.database), + cosmos_container_name=self.corpus.container, + openai_client=self.settings.build_openai_client(self.corpus), + openai_embedding_model=self.corpus.embed_model, + embed_query_instruction=self.corpus.embed_query_instruction, + embed_dimensions=self.corpus.embed_dimensions, + reranker=self._reranker, + token_counter=self._text_token_counter, + search_display_limit=self.settings.cosmos_retriever_search_display_limit, + schema_override=schema_override, + cosmos_client=cosmos_client, + enable_raw_query=self.settings.cosmos_retriever_raw_query_enabled, + ) + + 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, + embed_dimensions=self.corpus.embed_dimensions, + 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, + overrides: RuntimeConfig | None = None, + ) -> RetrievalResult: + + if not query or not query.strip(): + raise ValueError("query must be a non-empty string") + + effective_docs = ( + overrides.max_documents + if overrides is not None and overrides.max_documents is not None + else max_documents + ) + return self._search_sync(query, effective_docs, self._resolve_exec(overrides)) + + def _resolve_exec(self, overrides: RuntimeConfig | None) -> _Exec: + def pick(attr: str, default): + value = getattr(overrides, attr, None) if overrides is not None else None + return value if value is not None else default + + return _Exec( + turns=pick("chat_max_turns", self.settings.chat_max_turns), + temperature=pick("chat_temperature", self.settings.chat_temperature), + max_tokens=pick("chat_max_tokens", self.settings.chat_max_tokens), + reasoning_effort=pick("chat_reasoning_effort", self.settings.chat_reasoning_effort), + anthropic_version=pick("anthropic_version", self.settings.anthropic_version), + anthropic_auth_header=pick( + "anthropic_auth_header", self.settings.anthropic_auth_header + ), + ) + + def _search_sync( + self, + query: str, + max_documents: int, + exec_params: _Exec, + ) -> RetrievalResult: + + if self._use_chat: + return self._search_chat(query, max_documents, exec_params) + if self._use_anthropic: + return self._search_anthropic(query, max_documents, exec_params) + return self._search_responses(query, max_documents, exec_params) + + def _search_chat( + self, query: str, max_documents: int, exec_params: _Exec + ) -> 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=exec_params.turns, + temperature=exec_params.temperature, + max_tokens=exec_params.max_tokens, + text_token_counter=self._text_token_counter, + threshold_budget=self.settings.cosmos_retriever_threshold_budget, + token_budget=self.settings.cosmos_retriever_token_budget, + ) + 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, "timing": chat_result.timing}, + 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, exec_params: _Exec + ) -> 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=exec_params.turns, + max_tokens=exec_params.max_tokens, + reasoning_effort=exec_params.reasoning_effort, + text_token_counter=self._text_token_counter, + threshold_budget=self.settings.cosmos_retriever_threshold_budget, + token_budget=self.settings.cosmos_retriever_token_budget, + ) + 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, "timing": chat_result.timing}, + 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, exec_params: _Exec + ) -> 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=exec_params.turns, + max_tokens=exec_params.max_tokens, + anthropic_version=exec_params.anthropic_version, + auth_header=exec_params.anthropic_auth_header, + text_token_counter=self._text_token_counter, + threshold_budget=self.settings.cosmos_retriever_threshold_budget, + token_budget=self.settings.cosmos_retriever_token_budget, + ) + 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, "timing": chat_result.timing}, + 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_cross_collection_retriever( + self, cosmos_client, schema_override + ) -> CrossCollectionRetriever: + db_client = cosmos_client.get_database_client(self.corpus.database) + embedder = QueryEmbedder( + client=self.settings.build_openai_client(self.corpus), + model=self.corpus.embed_model, + query_instruction=self.corpus.embed_query_instruction, + dimensions=self.corpus.embed_dimensions, + ) + catalog = ResourceCatalog(_StaticCosmosConnection(cosmos_client)) + targets = select_search_targets(catalog, self.corpus.database) + if not targets: + raise RuntimeError( + f"No searchable collections found in database {self.corpus.database!r} " + "(a container needs a full-text or vector index to be searchable)." + ) + retrievers = { + target: build_capability_retriever_from_live( + container=db_client.get_container_client(target.container), + database=target.database, + embedder=embedder, + override=schema_override, + ) + for target in targets + } + logger.info( + "cross_collection_retriever_built", + database=self.corpus.database, + collections=[t.container for t in targets], + count=len(targets), + ) + return CrossCollectionRetriever(targets, retrievers) + + 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..96c336d --- /dev/null +++ b/cosmos-retriever/src/cosmos_retriever/server.py @@ -0,0 +1,285 @@ + +"""The HTTP service that puts the retriever behind an API. + +This module wraps the search agent in a small web service so other systems can +call it over HTTP instead of importing the code. It exposes a /search +endpoint that takes a natural-language query (plus an optional database and +container to aim at) and returns the ranked documents, along with a couple of +operator endpoints for checking health and viewing or adjusting server settings +while it runs. + +Each request may also pin which database and container to search: a database is +always required, while leaving the container off means "search the whole +database". Per-request overrides can further tweak model and budget settings for +that one call without disturbing the server's defaults. + +The actual search work is delegated to CosmosRetriever (see retriever.py); +this module only handles the web layer, caching, and request wiring. +""" + +from __future__ import annotations + +import asyncio +from collections import defaultdict +from contextlib import asynccontextmanager +from dataclasses import asdict +from typing import TYPE_CHECKING, NamedTuple + +import anyio +import structlog +from fastapi import FastAPI +from fastapi.responses import JSONResponse +from pydantic import BaseModel, Field + +from cosmos_retriever.cache import BoundedTTLCache +from cosmos_retriever.config import ( + RetrieverSettings, + RuntimeConfig, + ServerConfigUpdate, + get_settings, +) +from cosmos_retriever.retriever import CosmosRetriever + +if TYPE_CHECKING: + from collections.abc import AsyncIterator + +logger = structlog.get_logger("cosmos_retriever.server") + + +class RetrievalScope(NamedTuple): + database: str | None + container: str | None + + @classmethod + def resolve(cls, settings: RetrieverSettings, database: str | None, container: str | None): + # Container is optional: when unspecified, search the WHOLE database + # (every searchable collection) via cross-collection mode ("*"). The + # database is what the agent reasons about; a specific container is only + # an optional narrowing override. Database has no default and is required. + return cls( + database=database or settings.cosmos_database, + container=container or settings.cosmos_corpus_container or "*", + ) + + +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="Cosmos database name to query (required; selected per request).", + ) + container: str | None = Field( + default=None, + description="Optional Cosmos container to narrow to; omit to search the whole database.", + ) + overrides: RuntimeConfig | None = Field( + default=None, + description="Per-request runtime overrides (model endpoint, backend, turns, budgets, etc.).", + ) + + model_config = {"populate_by_name": True} + + +class _RetrieverPool: + + """Server-side pool of built ``CosmosRetriever`` engines. + + Constructing a retriever is expensive (opens Cosmos clients, builds the + embedding client and reranker, runs live schema discovery), so one engine is + built per unique ``(RetrievalScope, structural-overrides)`` key and reused + across requests via a :class:`BoundedTTLCache`. Only the engine is cached — + every search still runs fresh against Cosmos. A per-key ``asyncio.Lock`` + serializes requests that share an engine, and a single ``_build_lock`` gives + double-checked building so concurrent first-hits don't construct duplicates. + """ + + def __init__(self, settings: RetrieverSettings) -> None: + self._settings = settings + self._cache: BoundedTTLCache[tuple, CosmosRetriever] = BoundedTTLCache( + max_entries=settings.cosmos_retriever_cache_max_entries, + ttl_seconds=settings.cosmos_retriever_cache_ttl_seconds, + ) + self._locks: dict[tuple, asyncio.Lock] = defaultdict(asyncio.Lock) + self._build_lock = asyncio.Lock() + + async def get( + self, + database: str | None, + container: str | None, + overrides: RuntimeConfig | None = None, + ) -> tuple[CosmosRetriever, asyncio.Lock]: + scope = RetrievalScope.resolve(self._settings, database, container) + key: tuple = (scope, overrides.structural_key() if overrides is not None else None) + retriever = self._cache.get(key) + if retriever is None: + async with self._build_lock: + retriever = self._cache.get(key) + if retriever is None: + retriever = await anyio.to_thread.run_sync( + lambda: self._build(scope, overrides) + ) + self._cache.put(key, retriever) + return retriever, self._locks[key] + + def stats(self) -> dict[str, object]: + s = self._cache.stats() + return { + "entries": s.entries, + "max_entries": s.max_entries, + "ttl_seconds": s.ttl_seconds, + "hits": s.hits, + "misses": s.misses, + "evictions": s.evictions, + "expirations": s.expirations, + } + + @property + def settings(self) -> RetrieverSettings: + return self._settings + + async def update_settings(self, new_settings: RetrieverSettings) -> None: + """Swap the server-level settings and drop all cached retrievers so the + next request rebuilds them with the new configuration. Rebuilds the + cache itself if its sizing changed.""" + async with self._build_lock: + size_changed = ( + new_settings.cosmos_retriever_cache_max_entries + != self._settings.cosmos_retriever_cache_max_entries + or new_settings.cosmos_retriever_cache_ttl_seconds + != self._settings.cosmos_retriever_cache_ttl_seconds + ) + self._settings = new_settings + if size_changed: + self._cache = BoundedTTLCache( + max_entries=new_settings.cosmos_retriever_cache_max_entries, + ttl_seconds=new_settings.cosmos_retriever_cache_ttl_seconds, + ) + else: + self._cache.clear() + self._locks.clear() + + def _build( + self, scope: RetrievalScope, overrides: RuntimeConfig | None + ) -> CosmosRetriever: + settings = self._settings.apply_structural_overrides(overrides) + if settings is self._settings: + settings = settings.model_copy(deep=True) + if scope.database: + settings.cosmos_database = scope.database + return CosmosRetriever(settings=settings, corpus_name=scope.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, object]: + pool: _RetrieverPool | None = getattr(app.state, "pool", None) + return {"status": "ok", "retriever_cache": pool.stats() if pool is not None else {}} + + @app.get("/config") + async def get_config() -> JSONResponse: + """Admin/operator endpoint: return the current server-level settings + (secrets redacted) plus retriever-pool stats. Read-only; not part of the + search request path and not end-user facing.""" + pool: _RetrieverPool = app.state.pool + return JSONResponse( + content={"config": pool.settings.redacted_config(), "pool": pool.stats()} + ) + + @app.patch("/config") + async def patch_config(update: ServerConfigUpdate) -> JSONResponse: + """Admin/operator endpoint: mutate server-level *default* settings at + runtime (not part of the search request path, not end-user facing). + Applies only the provided fields, then rebuilds the retriever pool so + subsequent requests pick them up. Per-request tweaks instead go through + ``RuntimeConfig`` in the ``/search`` body.""" + pool: _RetrieverPool = app.state.pool + try: + new_settings = pool.settings.apply_server_updates(update) + except Exception as exc: + return JSONResponse( + status_code=400, + content={"error": str(exc), "type": type(exc).__name__}, + ) + await pool.update_settings(new_settings) + changed = sorted(update.model_dump(exclude_none=True).keys()) + logger.info("server_config_updated", changed=changed) + return JSONResponse( + content={ + "status": "ok", + "changed": changed, + "config": new_settings.redacted_config(), + "pool": pool.stats(), + } + ) + + @app.post("/search") + async def search(request: SearchRequest) -> JSONResponse: + pool: _RetrieverPool = app.state.pool + scope = RetrievalScope.resolve(pool.settings, request.database, request.container) + if not scope.database: + return JSONResponse( + status_code=400, + content={ + "error": ( + "Missing required field: database. The database must be " + "specified on each request (container is optional — omit " + "it to search the whole database)." + ), + "type": "ValueError", + }, + ) + try: + retriever, lock = await pool.get( + request.database, request.container, request.overrides + ) + async with lock: + result = await anyio.to_thread.run_sync( + lambda: retriever.search( + request.query, + max_documents=request.max_documents, + overrides=request.overrides, + ) + ) + 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..eb9c571 --- /dev/null +++ b/cosmos-retriever/src/cosmos_retriever/tools.py @@ -0,0 +1,897 @@ +"""The actions the retrieval agent is allowed to take. + +Every move the agent can make during a search is a tool defined here: search the +corpus, grep it for a pattern, read a whole document, discard chunks it no longer +needs, and, when explicitly enabled, run a read-only query directly. Each tool +carries a self-describing schema (its name, what it does, and the arguments it +accepts) and knows how to run itself, returning both the text the agent sees and +a little metadata about what happened. + +The same tool has to be described in whichever shape the chosen model expects, so +each schema can render itself into the OpenAI, OpenAI-Harmony, or Anthropic wire +format on demand. That is what lets one set of tools work across all three +backends without change. + +Tools are gathered into a ToolSet, the bundle handed to the agent loop (see +agent_loop.py), which is what actually calls them turn by turn. A build helper +assembles the standard set for a corpus, wiring each tool to the underlying +corpus retriever and reranker so callers don't construct them piece by piece. + + +""" + +from __future__ import annotations + +import json +import re +import time +from abc import ABC, abstractmethod +from collections.abc import Callable +from typing import Any, TypeAlias + +import openai +import structlog +from azure.cosmos import CosmosClient, DatabaseProxy +from pydantic import BaseModel, Field + +from cosmos_retriever.rerank import Reranker +from cosmos_retriever.retrieval import ( + CorpusRetriever, + GrepRequest, + QueryEmbedder, + ReadDocumentRequest, + SchemaOverride, + SearchRequest, + build_capability_retriever_from_live, +) +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): + + """Abstract base for every agent tool in this package. + + Concrete tools (search / grep / read / prune / …) implement ``__call__`` to + run the action and return ``(text_output, metadata)``; ``get_format`` + serializes the tool's schema into the requested provider wire format. This is + the Python service's own tool abstraction — the .NET MCP Toolkit's tools are a + separate C# hierarchy and are not shared here. + """ + + 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": ( + "Required. Text field(s) to keyword-match against " + f"(available: {text_names}). Choose one or more on every call." + ), + } + 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() + ) + required = ["query"] + (["fields"] if len(text_names) > 1 else []) + return ToolSchema( + name="search_corpus", description=desc, parameters=params, required=required + ) + + +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": ( + "Required. Text field to search " + f"(available: {text_names}). Choose one on every call." + ), + } + 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) + ) + required = ["pattern"] + (["field"] if len(text_names) > 1 else []) + return ToolSchema( + name="grep_corpus", description=desc, parameters=params, required=required + ) + + +class SearchCorpusToolCallMetadata(ToolCallMetadata): + + returned_chunk_ids: list[str] + pre_rerank_chunk_ids: list[str] | None = None + retrieval_s: float = 0.0 + rerank_s: float = 0.0 + + +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: + _t = time.perf_counter() + items = self._retriever.search(request) + retrieval_s = time.perf_counter() - _t + 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) + rerank_s = 0.0 + if self._reranker is not None and ids: + _t = time.perf_counter() + rerank_results = self._reranker(query, documents, max_tokens=max_tokens_override) + rerank_s = time.perf_counter() - _t + 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, + retrieval_s=round(retrieval_s, 3), + rerank_s=round(rerank_s, 3), + ) + + +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 + + # Item-is-document mode: each item is a whole, single-chunk document, so + # there are no sub-chunks to rerank or select. Read the item directly and + # return it verbatim — no reranker call and no query-based filtering. + if self._retriever.schema.is_item_document_mode: + document = self._retriever.read_document( + ReadDocumentRequest(document_id=doc_id) + ) + assembled = document.assembled + log.info("read_item_document", doc_id=doc_id) + 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 + + 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 + + +RUN_QUERY_SCHEMA = ToolSchema( + name="execute_query", + description=( + "Author and execute a READ-ONLY Azure Cosmos DB NoSQL (SQL) SELECT query " + "against a chosen database/container in the account.\n\n" + "RESTRICTED USE — call this ONLY when the user's request cannot be answered " + "by semantic/keyword retrieval and specifically requires precise structured " + "access to the data, such as:\n" + " - exact field-value filters (e.g. status = 'closed', author = 'X'),\n" + " - numeric or date range filters,\n" + " - counts / aggregations (COUNT, SUM, AVG, MIN, MAX),\n" + " - DISTINCT values or GROUP BY,\n" + " - deterministic ordering by a specific field.\n" + "Do NOT use it for ordinary topical, conceptual, or semantic questions — use " + "search_corpus, grep_corpus, and read_document for those. Only SELECT queries " + "are permitted (writes are impossible via this API); results are truncated. " + "Typical pattern: use it to pinpoint document ids by structured criteria, then " + "read_document those ids or include them in your final ranked output." + ), + parameters={ + "query": { + "type": "string", + "description": ( + "A single read-only Cosmos DB NoSQL SELECT query. The container is " + "aliased as 'c' (e.g. SELECT c.id, c.title FROM c WHERE c.status = 'open')." + ), + }, + "database": { + "type": "string", + "description": "Target database id. Defaults to the current corpus database.", + }, + "container": { + "type": "string", + "description": "Target container/collection id. Defaults to the current corpus container.", + }, + }, + required=["query"], +) + + +_SELECT_RE = re.compile(r"^\s*\(*\s*select\b", re.IGNORECASE | re.DOTALL) + + +def _sanitize_query_value(value: Any) -> Any: + + if isinstance(value, list): + # Collapse embedding-like numeric vectors so they don't flood context. + if len(value) > 32 and all(isinstance(x, (int, float)) for x in value[:8]): + return f"" + return [_sanitize_query_value(x) for x in value] + if isinstance(value, dict): + return {k: _sanitize_query_value(v) for k, v in value.items()} + if isinstance(value, str) and len(value) > 4000: + return value[:4000] + "\u2026" + return value + + +class RunQueryTool(Tool): + + tool_schema: ToolSchema + _client: CosmosClient + _default_database: str + _default_container: str + _max_rows: int + _max_chars: int + + def __init__( + self, + client: CosmosClient, + default_database: str = "", + default_container: str = "", + max_rows: int = 20, + max_chars: int = 20_000, + ) -> None: + super().__init__(tool_schema=RUN_QUERY_SCHEMA) + self._client = client + self._default_database = default_database + self._default_container = default_container + self._max_rows = max_rows + self._max_chars = max_chars + + 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 "query" not in params: + log.error("invalid_params", params_type=type(params).__name__) + raise ValueError(f"Invalid params type: {type(params)}") + + query = str(params["query"] or "").strip() + database = str(params.get("database") or self._default_database or "").strip() + container = str(params.get("container") or self._default_container or "").strip() + + if not query: + return "execute_query error: 'query' must be a non-empty SELECT query.", None + if not _SELECT_RE.match(query): + return ( + "execute_query error: only read-only SELECT queries are allowed.", + None, + ) + if not database or not container: + return ( + "execute_query error: specify both 'database' and 'container'.", + None, + ) + + log.info("execute_query", database=database, container=container, query=query[:200]) + try: + cont = self._client.get_database_client(database).get_container_client(container) + rows: list[Any] = [] + for item in cont.query_items( + query=query, + enable_cross_partition_query=True, + max_item_count=self._max_rows, + ): + rows.append(item) + if len(rows) >= self._max_rows: + break + except Exception as exc: # noqa: BLE001 + log.warning("execute_query_error", error=str(exc)) + return f"execute_query error: {type(exc).__name__}: {exc}", None + + sanitized = [_sanitize_query_value(r) for r in rows] + body = json.dumps(sanitized, ensure_ascii=False, default=str) + truncated = len(body) > self._max_chars + if truncated: + body = body[: self._max_chars] + "\u2026" + header = ( + f"# execute_query: {len(rows)} row(s) from {database}/{container}" + + (f" (capped at {self._max_rows})" if len(rows) >= self._max_rows else "") + + (" [output truncated]" if truncated else "") + ) + return f"{header}\n{body}", 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, + embed_dimensions: int | 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, + schema_override: SchemaOverride | None = None, + cosmos_client: CosmosClient | None = None, + enable_raw_query: bool = False, + ) -> 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, + dimensions=embed_dimensions, + ) + retriever = build_capability_retriever_from_live( + container=container, + database=getattr(cosmos_database, "id", "") or "", + embedder=embedder, + override=schema_override, + ) + + 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()) + if enable_raw_query and cosmos_client is not None: + toolset.add_tool( + RunQueryTool( + client=cosmos_client, + default_database=getattr(cosmos_database, "id", "") or "", + default_container=cosmos_container_name or "", + ) + ) + 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", + "RUN_QUERY_SCHEMA", + "RunQueryTool", + "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..eecabcb --- /dev/null +++ b/cosmos-retriever/src/cosmos_retriever/utils.py @@ -0,0 +1,23 @@ + +from __future__ import annotations + +from enum import StrEnum + + +class ProviderFormat(StrEnum): + + """Tool-call wire format a tool schema is serialized into. + + Selected by the caller — each ``run_*`` agent loop passes the value matching + the configured ``INFERENCE_BACKEND`` — and is *not* auto-detected from model + responses: ``OPENAI`` is the ``/responses`` function format, + ``OPENAI_HARMONY`` the ``/chat/completions`` Harmony format, and + ``ANTHROPIC`` the Anthropic tool format. + """ + + OPENAI = "openai" + OPENAI_HARMONY = "openai_harmony" + ANTHROPIC = "anthropic" + + +__all__ = ["ProviderFormat"] diff --git a/cosmos-retriever/tests/README.md b/cosmos-retriever/tests/README.md new file mode 100644 index 0000000..a2f422d --- /dev/null +++ b/cosmos-retriever/tests/README.md @@ -0,0 +1,104 @@ +# Tests + +This suite has two kinds of tests. + +## Fake tests vs real tests + +**Fake tests (`tests/unit/`)** — the default suite. Every external dependency +(Cosmos DB, OpenAI/Azure embeddings, the chat model, the reranker, the network) +is replaced with an in-process fake or monkeypatch. They are fast, deterministic, +need no credentials, and run everywhere. This is the bulk of the coverage: config, +tools, retriever, planner, strategies, document resolvers, normalization, +embeddings, paths, expressions, security, server, rerank, token counting, and the +token-budget agent loops. + +**Real tests (`tests/end_to_end/`)** — opt-in integration tests that talk to +actual services and real model tokenizers. They are skipped by default so a plain +`pytest` run stays green offline. Each file documents its own `How to run` steps +in its module docstring. They fall into three groups: + +| File | Hits | Opt-in trigger | +|---|---|---| +| `test_skf_live.py` | Live Cosmos + embeddings + reranker + gpt-5.4 agent | `RUN_SKF_LIVE=1` | +| `test_skf_discovery_live.py` | Live Cosmos container metadata + embeddings | `RUN_SKF_LIVE=1` | +| `test_skf_cross_collection_live.py` | Live cross-collection RRF over Cosmos + embeddings | `RUN_SKF_LIVE=1` | +| `test_rerank_live.py` | A running Qwen3-Reranker `/score` server | server reachable | +| `test_tokenizer_comparison.py` | Real model tokenizers (tiktoken + HF ports) | libraries + tokenizers available | + +> Env isolation: `tests/conftest.py` strips the ambient service configuration +> (from `.env.local`) for the fake tests so they always see clean defaults, and +> deliberately leaves it in place for `tests/end_to_end/`. + +## How to run the fake tests + +From the `cosmos-retriever/` directory: + +```bash +uv venv --python 3.11 .venv +uv pip install --python .venv/bin/python -e ".[dev]" +source .venv/bin/activate + +pytest # runs all fake tests; real tests skip +``` + +You should see the unit tests pass and the `end_to_end` tests reported as +`skipped`. + +## How to set up the real tests + +### Reranker (`test_rerank_live.py`) + +1. Serve `Qwen/Qwen3-Reranker-8B` with a vLLM-compatible `/score` endpoint on + `http://127.0.0.1:8011`, for example: + ```bash + vllm serve Qwen/Qwen3-Reranker-8B --port 8011 \ + --hf-overrides '{"architectures":["Qwen3ForSequenceClassification"],"classifier_from_token":["no","yes"],"is_original_qwen3_reranker":true}' + ``` + (Needs vLLM >= 0.10 for the score conversion.) +2. If it runs elsewhere, set `VLLM_RERANKER_URL` to its base URL. +3. `pytest tests/end_to_end/test_rerank_live.py` + +### Tokenizer comparison (`test_tokenizer_comparison.py`) + +1. `uv pip install --python .venv/bin/python tiktoken tokenizers huggingface_hub` +2. Allow network access on first run so the model tokenizers download from + Hugging Face (or prime the HF cache first). Anything unavailable is skipped. +3. `pytest tests/end_to_end/test_tokenizer_comparison.py` +4. Regenerate the pinned numbers any time with + `python tests/end_to_end/tokenizer_panel.py`. + +### SKF live suites (`test_skf_*_live.py`) + +These run against the `skf-rag-test` Cosmos DB account. + +1. `az login`, and select the subscription that owns the `skf-rag-test` account. +2. Grant your identity the **Cosmos DB Built-in Data Reader** role on the account: + ```bash + az cosmosdb sql role assignment create \ + --account-name skf-rag-test --resource-group DiskANN_development \ + --role-definition-id 00000000-0000-0000-0000-000000000001 \ + --principal-id --scope "/" + ``` +3. Create `cosmos-retriever/.env.local` (copy from `.env.example`) and fill: + - `ACCOUNT_URI=https://skf-rag-test.documents.azure.com:443/` + - `COSMOS_DATABASE=skf-database`, `COSMOS_CORPUS_CONTAINER=skf-unstructured` + - Embeddings: `EMBED_ENDPOINT` (`.../openai/v1`), `OPENAI_API_KEY`, + `OPENAI_EMBEDDING_MODEL=text-embedding-3-small`, `OPENAI_EMBEDDING_DIMENSIONS=1536` + (the embedding model must match the container's vectors: 1536 for + `skf-unstructured`/`skf-structured`, 3072/`text-embedding-3-large` for + `skf-unstructured-text-large`). + - Chat (agent, `test_skf_live.py` only): `INFERENCE_BACKEND`, `CHAT_BASE_URL` + (`.../openai/v1`), `CHAT_API_KEY`, `CHAT_MODEL`. + - Reranker (agent assertions in `test_skf_live.py`): `VLLM_RERANKER_URL`, with + the reranker server from above running. +4. `export RUN_SKF_LIVE=1` +5. Run a suite: + ```bash + RUN_SKF_LIVE=1 pytest tests/end_to_end/test_skf_discovery_live.py # metadata + embeddings only + RUN_SKF_LIVE=1 pytest tests/end_to_end/test_skf_cross_collection_live.py # + cross-collection RRF + RUN_SKF_LIVE=1 pytest tests/end_to_end/test_skf_live.py # + full agent (chat + reranker) + ``` + +`test_skf_discovery_live.py` and `test_skf_cross_collection_live.py` need only +Cosmos read access and embeddings. `test_skf_live.py` additionally needs the chat +model and the reranker. `.env.local` holds secrets, keep it gitignored. diff --git a/cosmos-retriever/tests/conftest.py b/cosmos-retriever/tests/conftest.py new file mode 100644 index 0000000..3c3b159 --- /dev/null +++ b/cosmos-retriever/tests/conftest.py @@ -0,0 +1,59 @@ +"""Shared pytest fixtures. + +``config.py`` calls ``load_dotenv(.env.local)`` at import, which injects the +live service configuration into ``os.environ``. Unit tests construct +``RetrieverSettings`` expecting a clean environment (env vars outrank +``_env_file=None``), so this autouse fixture strips those keys for everything +except the live tests under ``tests/end_to_end/``, which intentionally use the +real ``.env.local``. +""" +from __future__ import annotations + +import pytest + +# Env vars that .env.local / a real deployment may set and that map onto +# RetrieverSettings fields — cleared so offline unit tests see defaults. +_CONFIG_ENV_VARS = [ + "INFERENCE_BACKEND", + "CHAT_BASE_URL", + "CHAT_API_KEY", + "CHAT_MODEL", + "CHAT_API_VERSION", + "CHAT_TEMPERATURE", + "CHAT_MAX_TOKENS", + "CHAT_MAX_TURNS", + "CHAT_REASONING_EFFORT", + "ANTHROPIC_VERSION", + "ANTHROPIC_AUTH_HEADER", + "ACCOUNT_URI", + "COSMOS_DATABASE", + "COSMOS_CORPUS_CONTAINER", + "COSMOS_KEY", + "OPENAI_API_KEY", + "OPENAI_EMBEDDING_MODEL", + "OPENAI_EMBEDDING_DIMENSIONS", + "EMBED_ENDPOINT", + "EMBED_QUERY_INSTRUCTION", + "CORPUS_REGISTRY", + "CORPUS_REGISTRY_FILE", + "BASETEN_API_KEY", + "BASETEN_MODEL_URL", + "VLLM_RERANKER_URL", + "LOG_LEVEL", + "HOST", + "PORT", +] + + +@pytest.fixture(autouse=True) +def _isolate_config_env(request, monkeypatch): + # Live end-to-end tests rely on the real .env.local; leave their env intact. + if "end_to_end" in str(request.fspath): + return + for var in _CONFIG_ENV_VARS: + monkeypatch.delenv(var, raising=False) + # Also stop pydantic-settings from auto-loading the repo-root .env.local file + # (it is read via model_config regardless of os.environ). + from cosmos_retriever.config import RetrieverSettings + + monkeypatch.setitem(RetrieverSettings.model_config, "env_file", None) diff --git a/cosmos-retriever/tests/end_to_end/test_rerank_live.py b/cosmos-retriever/tests/end_to_end/test_rerank_live.py new file mode 100644 index 0000000..fd36b6e --- /dev/null +++ b/cosmos-retriever/tests/end_to_end/test_rerank_live.py @@ -0,0 +1,75 @@ +"""Live integration test for VLLMQwen3Reranker against a REAL reranker server. + +Unlike test_rerank.py (which fakes all HTTP), this exercises the full +`VLLMQwen3Reranker` code path against an actual running Qwen3-Reranker server +that exposes a vLLM-compatible ``/score`` endpoint returning +``{"data": [{"score": float}, ...]}``. + +The test is skipped unless a reachable server is found. Point it at a server +with ``VLLM_RERANKER_URL`` (default ``http://127.0.0.1:8011``). It was validated +against the real Qwen3-Reranker-8B weights served locally. + +How to run. +1. Start a Qwen3 reranker that serves the vLLM score endpoint on port 8011. +2. If the server runs elsewhere, set VLLM_RERANKER_URL to its base url. +3. Run pytest on tests/end_to_end/test_rerank_live.py. + +See tests/README.md for the full setup. +""" +from __future__ import annotations + +import os + +import pytest +import requests + +from cosmos_retriever.rerank import RerankResult, VLLMQwen3Reranker + +RERANKER_URL = os.getenv("VLLM_RERANKER_URL", "http://127.0.0.1:8011") + + +def _server_reachable(url: str) -> bool: + try: + requests.get(f"{url}/health", timeout=2) + return True + except requests.exceptions.RequestException: + return False + + +pytestmark = pytest.mark.skipif( + not _server_reachable(RERANKER_URL), + reason=f"no reranker server reachable at {RERANKER_URL} (set VLLM_RERANKER_URL)", +) + + +def test_live_reranker_orders_relevant_documents_first() -> None: + reranker = VLLMQwen3Reranker(base_url=RERANKER_URL) + query = "What is the capital of China?" + documents = [ + "The capital of France is Paris.", + "The capital of China is Beijing.", + "Chocolate is a delicious treat.", + "Beijing has been the capital of China for a long time.", + ] + + results = reranker(query, documents) + + # Same count, all wrapped as RerankResult, original indices preserved as a set. + assert len(results) == len(documents) + assert all(isinstance(r, RerankResult) for r in results) + assert {r.original_index for r in results} == set(range(len(documents))) + + # Sorted by descending score (contract of _rerank). + scores = [r.score for r in results] + assert scores == sorted(scores, reverse=True) + + # The two China/Beijing documents (indices 1 and 3) must outrank the + # irrelevant France/chocolate ones (indices 0 and 2). + top_two = {results[0].original_index, results[1].original_index} + assert top_two == {1, 3} + assert results[1].score > results[2].score # clear relevance gap + + +def test_live_reranker_empty_documents_returns_empty() -> None: + reranker = VLLMQwen3Reranker(base_url=RERANKER_URL) + assert reranker("any query", []) == [] diff --git a/cosmos-retriever/tests/end_to_end/test_skf_cross_collection_live.py b/cosmos-retriever/tests/end_to_end/test_skf_cross_collection_live.py new file mode 100644 index 0000000..f0548c8 --- /dev/null +++ b/cosmos-retriever/tests/end_to_end/test_skf_cross_collection_live.py @@ -0,0 +1,344 @@ +"""End-to-end live tests for cross-collection RRF fusion over the real +``skf-rag-test`` / ``skf-database`` collections. + +Database-wide search fans each request out across every searchable collection +and fuses the per-collection ranked lists with Reciprocal Rank Fusion. These +tests exercise that path over real data at every layer: + + select_search_targets -> MultiContainerRetriever -> fuse_rrf + -> CrossCollectionRetriever (search / grep / read) + -> CosmosRetriever(corpus_name="*") + +``skf-database`` is a realistic stress case: its collections have mixed +embedding dimensions (1536 x2, 3072 x1) and different text-field sets, so the +fan-out's per-collection error tolerance (a failing collection is skipped, the +rest still fuse) is exercised with genuine dimension/field mismatches — not +mocks. + +Opt-in: skipped unless ``RUN_SKF_LIVE`` is truthy and ``.env.local`` points at +``skf-rag-test``. Requires ``az login`` with the Cosmos Data Reader role. + +How to run. +1. Run az login for the subscription that owns the skf-rag-test account. +2. Give your identity the Cosmos DB Data Reader role on that account. +3. Fill cosmos-retriever/.env.local with ACCOUNT_URI and the embeddings endpoint and key. +4. Set the environment variable RUN_SKF_LIVE to 1. +5. Run pytest on tests/end_to_end/test_skf_cross_collection_live.py. + +See tests/README.md for the full setup. +""" +from __future__ import annotations + +import os + +import pytest + + +def _live_enabled() -> bool: + if os.getenv("RUN_SKF_LIVE", "").strip().lower() not in {"1", "true", "yes"}: + return False + try: + from cosmos_retriever.config import get_settings + + s = get_settings() + return bool(s.account_uri and "skf-rag-test" in s.account_uri) + except Exception: + return False + + +pytestmark = pytest.mark.skipif( + not _live_enabled(), + reason="set RUN_SKF_LIVE=1 and configure .env.local for skf-rag-test", +) + +DATABASE = "skf-database" +C_UNSTRUCT = "skf-unstructured" # 1536-dim, 9 text fields +C_STRUCT = "skf-structured" # 1536-dim, 6 text fields (incl. description) +C_LARGE = "skf-unstructured-text-large" # 3072-dim -> vector-incompatible with 1536 embedder +COMMON_TEXT_FIELD = "description" # present in both 1536 collections + + +class _Conn: + def __init__(self, client): + self._client = client + + def client(self): + return self._client + + +def _target(container: str): + from cosmos_retriever.retrieval.orchestration import ContainerTarget + + return ContainerTarget(database=DATABASE, container=container) + + +# ─────────────────────────────── fixtures ───────────────────────────────── + + +@pytest.fixture(scope="session") +def settings(): + from cosmos_retriever.config import get_settings + + return get_settings() + + +@pytest.fixture(scope="session") +def cosmos_client(settings): + corpus = settings.resolve_corpus(C_UNSTRUCT) + return settings.build_cosmos_client(corpus) + + +@pytest.fixture(scope="session") +def embedder(settings): + from cosmos_retriever.retrieval import QueryEmbedder + + corpus = settings.resolve_corpus(C_UNSTRUCT) # 1536-dim + return QueryEmbedder( + client=settings.build_openai_client(corpus), + model=corpus.embed_model, + query_instruction=corpus.embed_query_instruction, + dimensions=corpus.embed_dimensions, + ) + + +@pytest.fixture(scope="session") +def retrievers(settings, cosmos_client, embedder): + """Per-collection CorpusRetrievers keyed by ContainerTarget (built once).""" + from cosmos_retriever.retrieval import build_capability_retriever_from_live + + db = cosmos_client.get_database_client(DATABASE) + out = {} + for name in (C_UNSTRUCT, C_STRUCT, C_LARGE): + try: + out[_target(name)] = build_capability_retriever_from_live( + container=db.get_container_client(name), + database=DATABASE, + embedder=embedder, + ) + except Exception as exc: + pytest.skip(f"could not build retriever for {name}: {exc}") + return out + + +@pytest.fixture(scope="session") +def catalog(cosmos_client): + from cosmos_retriever.retrieval.discovery import ResourceCatalog + + return ResourceCatalog(_Conn(cosmos_client)) + + +@pytest.fixture(scope="session") +def cross(retrievers): + """CrossCollectionRetriever over the two dimension-compatible collections.""" + from cosmos_retriever.retrieval.orchestration import CrossCollectionRetriever + + targets = [_target(C_UNSTRUCT), _target(C_STRUCT)] + subset = {t: retrievers[t] for t in targets} + return CrossCollectionRetriever(targets, subset) + + +def _vsearch(retriever, query="bearing steel", limit=5): + from cosmos_retriever.retrieval.models import SearchRequest + + return retriever.search(SearchRequest(query=query, limit=limit, mode="vector")) + + +# ═══════════════════════ select_search_targets ════════════════════════════ + + +def test_select_targets_finds_searchable_collections(catalog) -> None: + from cosmos_retriever.retrieval.orchestration import select_search_targets + + targets = select_search_targets(catalog, DATABASE) + names = {t.container for t in targets} + assert {C_UNSTRUCT, C_STRUCT, C_LARGE} <= names + assert all(t.database == DATABASE for t in targets) + for t in targets: + p = catalog.profile(DATABASE, t.container) + assert p.can_full_text.value or p.can_vector.value + + +def test_select_targets_capability_filter(catalog) -> None: + from cosmos_retriever.retrieval.orchestration import select_search_targets + + filtered = select_search_targets(catalog, DATABASE, require_capability=True) + unfiltered = select_search_targets(catalog, DATABASE, require_capability=False) + assert len(unfiltered) >= len(filtered) # unfiltered may include non-indexed containers + + +def test_select_targets_explicit_subset(catalog) -> None: + from cosmos_retriever.retrieval.orchestration import select_search_targets + + targets = select_search_targets(catalog, DATABASE, containers=[C_UNSTRUCT]) + assert [t.container for t in targets] == [C_UNSTRUCT] + + +# ═══════════════════════ fuse_rrf over real items ═════════════════════════ + + +def test_fuse_rrf_math_and_tagging(retrievers) -> None: + from cosmos_retriever.retrieval.orchestration import RRF_K, fuse_rrf + + t_u, t_s = _target(C_UNSTRUCT), _target(C_STRUCT) + a = _vsearch(retrievers[t_u], limit=5) + b = _vsearch(retrievers[t_s], limit=5) + assert a and b + + fused = fuse_rrf([(t_u, a), (t_s, b)]) + # Different collections -> qualified keys never collide -> all items survive. + assert len(fused) == len(a) + len(b) + assert [it.rank for it in fused] == list(range(len(fused))) + scores = [it.raw_scores["rrf"] for it in fused] + assert scores == sorted(scores, reverse=True) + assert fused[0].raw_scores["rrf"] == pytest.approx(1.0 / (RRF_K + 0)) # pos 0 -> 1/60 + tags = {(it.metadata["database"], it.metadata["container"]) for it in fused} + assert tags == {(DATABASE, C_UNSTRUCT), (DATABASE, C_STRUCT)} + + +def test_fuse_rrf_respects_limit(retrievers) -> None: + from cosmos_retriever.retrieval.orchestration import fuse_rrf + + t_u, t_s = _target(C_UNSTRUCT), _target(C_STRUCT) + fused = fuse_rrf( + [(t_u, _vsearch(retrievers[t_u], limit=5)), (t_s, _vsearch(retrievers[t_s], limit=5))], + limit=3, + ) + assert len(fused) == 3 + + +# ═══════════════════════ MultiContainerRetriever ══════════════════════════ + + +def _mcr(retrievers): + from cosmos_retriever.retrieval.orchestration import MultiContainerRetriever + + return MultiContainerRetriever(lambda t: retrievers[t], max_workers=4) + + +def test_multi_search_fuses_across_collections(retrievers) -> None: + from cosmos_retriever.retrieval.models import SearchRequest + + targets = [_target(C_UNSTRUCT), _target(C_STRUCT)] + res = _mcr(retrievers).search( + targets, SearchRequest(query="bearing steel", limit=5, mode="vector") + ) + assert res.items + assert set(res.searched) == set(targets) + assert set(res.per_container_counts) == {f"{DATABASE}/{C_UNSTRUCT}", f"{DATABASE}/{C_STRUCT}"} + assert res.errors == {} + assert res.elapsed_s > 0.0 + + +def test_multi_search_deduplicates_targets(retrievers) -> None: + from cosmos_retriever.retrieval.models import SearchRequest + + t_u = _target(C_UNSTRUCT) + res = _mcr(retrievers).search([t_u, t_u], SearchRequest(query="bearing", limit=5, mode="vector")) + assert res.searched == [t_u] + + +def test_multi_search_per_container_limit(retrievers) -> None: + from cosmos_retriever.retrieval.models import SearchRequest + + targets = [_target(C_UNSTRUCT), _target(C_STRUCT)] + res = _mcr(retrievers).search( + targets, SearchRequest(query="bearing", limit=10, mode="vector"), + per_container_limit=2, final_limit=10, + ) + assert all(count <= 2 for count in res.per_container_counts.values()) + + +def test_multi_search_captures_per_container_errors(retrievers) -> None: + from cosmos_retriever.retrieval.models import SearchRequest + + # "content" exists in skf-unstructured but not skf-structured, so the latter + # raises UnknownField; it must be recorded in errors WITHOUT aborting the + # fusion of the collection that succeeded. + targets = [_target(C_UNSTRUCT), _target(C_STRUCT)] + res = _mcr(retrievers).search( + targets, SearchRequest(query="bearing", limit=5, mode="text", text_fields=["content"]) + ) + assert f"{DATABASE}/{C_STRUCT}" in res.errors # field mismatch captured + assert _target(C_UNSTRUCT) in res.searched + assert res.items # good collection still returned fused results + + +# ═══════════════════════ CrossCollectionRetriever ═════════════════════════ + + +def test_cross_schema_is_representative(cross, retrievers) -> None: + assert cross.schema is retrievers[_target(C_UNSTRUCT)].schema + + +def test_cross_search_fuses_multiple_collections(cross) -> None: + from cosmos_retriever.retrieval.models import SearchRequest + + items = cross.search(SearchRequest(query="bearing steel", limit=10, mode="vector")) + assert items + assert [it.rank for it in items] == list(range(len(items))) + assert all("rrf" in it.raw_scores for it in items) + containers = {it.metadata.get("container") for it in items} + assert containers <= {C_UNSTRUCT, C_STRUCT} + assert len(containers) >= 2 # genuinely fused across both collections + + +def test_cross_search_respects_limit(cross) -> None: + from cosmos_retriever.retrieval.models import SearchRequest + + items = cross.search(SearchRequest(query="bearing", limit=4, mode="vector")) + assert len(items) <= 4 + + +def test_cross_search_error_tolerant_on_field_mismatch(cross) -> None: + from cosmos_retriever.retrieval.models import SearchRequest + + # "content" exists in skf-unstructured but not skf-structured -> the latter + # errors out, but the search still returns the former's hits. + items = cross.search( + SearchRequest(query="aerospace bearing", limit=5, mode="text", text_fields=["content"]) + ) + assert items + assert {it.metadata.get("container") for it in items} == {C_UNSTRUCT} + + +def test_cross_grep_fans_out_and_caps(cross) -> None: + from cosmos_retriever.retrieval.models import GrepRequest + + hits = cross.grep_candidates( + GrepRequest(pattern="bearing", text_field=COMMON_TEXT_FIELD, candidate_limit=10) + ) + assert isinstance(hits, list) + assert len(hits) <= 10 + for it in hits: + assert it.item_id.startswith("skf") + + +def test_cross_read_document_round_trip(cross) -> None: + from cosmos_retriever.retrieval.models import ReadDocumentRequest, SearchRequest + + items = cross.search(SearchRequest(query="aerospace bearing", limit=3, mode="vector")) + doc_id = items[0].item_id + doc = cross.read_document(ReadDocumentRequest(document_id=doc_id)) + assert doc.assembled.strip() + + +def test_cross_read_unknown_document_is_empty(cross) -> None: + from cosmos_retriever.retrieval.models import ReadDocumentRequest + + doc = cross.read_document(ReadDocumentRequest(document_id="skf-nonexistent-zzz")) + assert doc.assembled == "" + + +# ═══════════════════════ CosmosRetriever(corpus_name="*") ═════════════════ + + +def test_database_wide_engine_builds_cross_collection(settings) -> None: + from cosmos_retriever.retriever import CosmosRetriever + + engine = CosmosRetriever(settings=settings, corpus_name="*") + assert engine.database_wide is True + inner = engine.toolset.get_tool("search_corpus")._retriever + from cosmos_retriever.retrieval.orchestration import CrossCollectionRetriever + + assert isinstance(inner, CrossCollectionRetriever) + assert len(inner._targets) >= 3 # all searchable skf-database collections diff --git a/cosmos-retriever/tests/end_to_end/test_skf_discovery_live.py b/cosmos-retriever/tests/end_to_end/test_skf_discovery_live.py new file mode 100644 index 0000000..eb5b826 --- /dev/null +++ b/cosmos-retriever/tests/end_to_end/test_skf_discovery_live.py @@ -0,0 +1,315 @@ +"""End-to-end live tests for schema discovery (``binding.py`` _from_live path) +against the real ``skf-rag-test`` containers. + +Exercises every layer of live discovery on real Cosmos DB container metadata: + + container.read() -> parse_container_metadata -> {capabilities, schema} + -> build_capability_retriever_from_live + +and cross-checks the discovered capabilities/schema against each container's +actual vector-embedding / full-text / partition-key policies. Covers all three +skf-database text containers (1536-dim x2 and 3072-dim) so dimension and +field-set adaptation is verified, plus override application and determinism. + +Opt-in: skipped unless ``RUN_SKF_LIVE`` is truthy and ``.env.local`` points at +``skf-rag-test``. Requires ``az login`` with the Cosmos Data Reader role. + +How to run. +1. Run az login for the subscription that owns the skf-rag-test account. +2. Give your identity the Cosmos DB Data Reader role on that account. +3. Fill cosmos-retriever/.env.local with ACCOUNT_URI and the embeddings endpoint and key. +4. Set the environment variable RUN_SKF_LIVE to 1. +5. Run pytest on tests/end_to_end/test_skf_discovery_live.py. + +See tests/README.md for the full setup. +""" +from __future__ import annotations + +import os + +import pytest + + +def _live_enabled() -> bool: + if os.getenv("RUN_SKF_LIVE", "").strip().lower() not in {"1", "true", "yes"}: + return False + try: + from cosmos_retriever.config import get_settings + + s = get_settings() + return bool(s.account_uri and "skf-rag-test" in s.account_uri) + except Exception: + return False + + +pytestmark = pytest.mark.skipif( + not _live_enabled(), + reason="set RUN_SKF_LIVE=1 and configure .env.local for skf-rag-test", +) + +DATABASE = "skf-database" +PRIMARY = "skf-unstructured" + +# Ground truth (from the containers' vector/full-text policies). +EXPECTED = { + "skf-unstructured": { + "dims": 1536, + "fields": {"benefits", "content", "description", "designation", + "long_description", "summary", "taxonomy", "taxonomy_sap", "title"}, + }, + "skf-unstructured-text-large": { + "dims": 3072, + "fields": {"title", "summary", "content"}, + }, + "skf-structured": { + "dims": 1536, + "fields": {"benefits", "description", "designation", + "long_description", "taxonomy", "taxonomy_sap"}, + }, +} + +# ─────────────────────────────── helpers ────────────────────────────────── + + +def _container(settings, name): + corpus = settings.resolve_corpus(name) + client = settings.build_cosmos_client(corpus) + return client.get_database_client(corpus.database).get_container_client(name) + + +def _discover(settings, name): + from cosmos_retriever.retrieval.binding import ( + capabilities_from_metadata, + schema_from_metadata, + ) + from cosmos_retriever.retrieval.discovery.profiler import parse_container_metadata + + container = _container(settings, name) + props = container.read() + md = parse_container_metadata(DATABASE, container.id, props, props.get("_etag")) + return props, md, capabilities_from_metadata(md), schema_from_metadata(md) + + +# ─────────────────────────────── fixtures ───────────────────────────────── + + +@pytest.fixture(scope="session") +def settings(): + from cosmos_retriever.config import get_settings + + return get_settings() + + +@pytest.fixture(scope="session") +def primary(settings): + """(props, metadata, capabilities, schema) for skf-unstructured.""" + try: + return _discover(settings, PRIMARY) + except Exception as exc: + pytest.skip(f"discovery failed: {type(exc).__name__}: {exc}") + + +# ═══════════════════════ raw read + parse_container_metadata ══════════════ + + +def test_container_read_returns_policies(primary) -> None: + props, *_ = primary + assert props["id"] == PRIMARY + assert props.get("partitionKey", {}).get("paths") == ["/id"] + assert props.get("vectorEmbeddingPolicy", {}).get("vectorEmbeddings") + assert props.get("fullTextPolicy", {}).get("fullTextPaths") + + +def test_parse_metadata_partition_and_etag(primary) -> None: + props, md, *_ = primary + assert md.database == DATABASE and md.container == PRIMARY + assert md.partition_key_paths == ["/id"] + assert md.etag == props.get("_etag") + assert md.fetched_at > 0 + + +def test_parse_metadata_vector_field(primary) -> None: + _, md, *_ = primary + assert len(md.vector_fields) == 1 + vf = md.vector_fields[0] + assert vf.path == "/embedding" + assert vf.dimensions == 1536 + assert vf.distance_function == "cosine" + assert vf.indexed is True # a vector index exists for the embedding path + + +def test_parse_metadata_full_text_index_matches_policy(primary) -> None: + _, md, *_ = primary + # The indexed full-text paths should cover the full-text policy paths. + assert set(md.full_text_paths) == set(md.full_text_policy_paths) + assert len(md.full_text_paths) == 9 + + +# ═══════════════════════ capabilities_from_metadata ═══════════════════════ + + +def test_capabilities_flags(primary) -> None: + _, _, caps, _ = primary + assert caps.full_text_supported is True + assert caps.vector_supported is True + assert caps.native_hybrid_supported is True + assert caps.efficient_document_lookup_supported is True + + +def test_capabilities_vector_capability(primary) -> None: + from cosmos_retriever.retrieval.capabilities import SupportLevel + from cosmos_retriever.retrieval.paths import CosmosPath + + _, _, caps, _ = primary + cap = caps.vector_capability_for(CosmosPath.parse("/embedding")) + assert cap is not None + assert cap.dimensions == 1536 + assert cap.distance_function == "cosine" + assert cap.support is SupportLevel.INDEXED + assert caps.vector_capability_for(CosmosPath.parse("/nope")) is None + + +def test_capabilities_full_text_paths(primary) -> None: + from cosmos_retriever.retrieval.paths import CosmosPath + + _, _, caps, _ = primary + for field in ("title", "summary", "content"): + assert caps.has_full_text_path(CosmosPath.parse(f"/{field}")) + assert not caps.has_full_text_path(CosmosPath.parse("/not_a_field")) + assert [str(p) for p in caps.partition_key_paths] == ["/id"] + + +# ═══════════════════════ schema_from_metadata ═════════════════════════════ + + +def test_schema_identity_and_mode(primary) -> None: + _, _, _, schema = primary + assert str(schema.item_id_path) == "/id" + assert schema.is_item_document_mode is True # no document_id_path override + assert schema.partition_key_is_document_id is False + assert [str(p) for p in schema.partition_key_paths] == ["/id"] + + +def test_schema_text_fields(primary) -> None: + _, _, _, schema = primary + assert set(schema.text_field_map()) == EXPECTED[PRIMARY]["fields"] + assert len(schema.text_paths) == 9 + + +def test_schema_vector_field(primary) -> None: + _, _, _, schema = primary + assert set(schema.vector_field_map()) == {"embedding"} + vf = schema.resolve_vector_config(None) + assert vf.dimensions == 1536 + assert str(schema.resolve_vector_field(None)) == "/embedding" + + +def test_schema_resolve_text_fields(primary) -> None: + from cosmos_retriever.retrieval.errors import UnknownField + + _, _, _, schema = primary + assert [str(p) for p in schema.resolve_text_fields(["title"])] == ["/title"] + with pytest.raises(UnknownField): + schema.resolve_text_fields(None) # ambiguous: 9 fields + with pytest.raises(UnknownField): + schema.resolve_text_fields(["not_a_field"]) + + +def test_schema_agent_field_summary(primary) -> None: + _, _, _, schema = primary + summary = schema.agent_field_summary() + assert isinstance(summary, str) and "title" in summary + + +# ═══════════════════════ build_capability_retriever_from_live ═════════════ + + +@pytest.fixture(scope="session") +def live_retriever(settings): + from cosmos_retriever.retrieval import ( + QueryEmbedder, + build_capability_retriever_from_live, + ) + + corpus = settings.resolve_corpus(PRIMARY) + container = _container(settings, PRIMARY) + embedder = QueryEmbedder( + client=settings.build_openai_client(corpus), + model=corpus.embed_model, + query_instruction=corpus.embed_query_instruction, + dimensions=corpus.embed_dimensions, + ) + return build_capability_retriever_from_live( + container=container, database=DATABASE, embedder=embedder + ) + + +def test_from_live_builds_matching_schema(live_retriever, primary) -> None: + _, _, _, schema = primary + assert str(live_retriever.schema.item_id_path) == "/id" + assert set(live_retriever.schema.text_field_map()) == set(schema.text_field_map()) + assert live_retriever.schema.resolve_vector_config(None).dimensions == 1536 + + +def test_from_live_builds_matching_capabilities(live_retriever) -> None: + caps = live_retriever.capabilities + assert caps.native_hybrid_supported and caps.full_text_supported and caps.vector_supported + + +def test_from_live_retriever_actually_searches(live_retriever) -> None: + from cosmos_retriever.retrieval.models import SearchRequest + + items = live_retriever.search( + SearchRequest(query="aerospace bearing", limit=5, mode="vector") + ) + assert items and all(it.item_id.startswith("skf") for it in items) + + +def test_from_live_override_switches_to_chunked_mode(settings) -> None: + from cosmos_retriever.retrieval import ( + QueryEmbedder, + build_capability_retriever_from_live, + ) + from cosmos_retriever.retrieval.schema_override import SchemaOverride + + corpus = settings.resolve_corpus(PRIMARY) + container = _container(settings, PRIMARY) + embedder = QueryEmbedder( + client=settings.build_openai_client(corpus), + model=corpus.embed_model, + dimensions=corpus.embed_dimensions, + ) + override = SchemaOverride(document_id_path="/id", title_path="/title") + retriever = build_capability_retriever_from_live( + container=container, database=DATABASE, embedder=embedder, override=override + ) + schema = retriever.schema + assert str(schema.document_id_path) == "/id" + assert str(schema.title_path) == "/title" + assert schema.is_item_document_mode is False # override provides a document id path + assert schema.partition_key_is_document_id is True # pk /id == document id /id + + +def test_discovery_is_deterministic(settings) -> None: + _, md1, caps1, s1 = _discover(settings, PRIMARY) + _, md2, caps2, s2 = _discover(settings, PRIMARY) + assert set(s1.text_field_map()) == set(s2.text_field_map()) + assert s1.resolve_vector_config(None).dimensions == s2.resolve_vector_config(None).dimensions + assert caps1.native_hybrid_supported == caps2.native_hybrid_supported + assert [str(v.path) for v in md1.vector_fields] == [str(v.path) for v in md2.vector_fields] + + +# ═══════════════════════ cross-container adaptation ═══════════════════════ + + +@pytest.mark.parametrize("name", list(EXPECTED)) +def test_discovery_adapts_per_container(settings, name) -> None: + try: + _, md, caps, schema = _discover(settings, name) + except Exception as exc: + pytest.skip(f"{name}: {type(exc).__name__}: {exc}") + exp = EXPECTED[name] + assert schema.resolve_vector_config(None).dimensions == exp["dims"] + assert md.vector_fields[0].dimensions == exp["dims"] + assert set(schema.text_field_map()) == exp["fields"] + assert caps.native_hybrid_supported is True # all three index text + vector diff --git a/cosmos-retriever/tests/end_to_end/test_skf_live.py b/cosmos-retriever/tests/end_to_end/test_skf_live.py new file mode 100644 index 0000000..5c03e45 --- /dev/null +++ b/cosmos-retriever/tests/end_to_end/test_skf_live.py @@ -0,0 +1,332 @@ +"""End-to-end live tests against the real ``skf-rag-test`` Cosmos DB corpus. + +These exercise the full stack — Azure Cosmos DB hybrid (vector + full-text RRF) +retrieval, Azure OpenAI ``text-embedding-3-small`` (1536-dim) query embeddings, +the local Qwen3-Reranker, and the gpt-5.4 ``/responses`` agent — against the +``skf-database/skf-unstructured`` container (~9k SKF product/industry docs). + +Cost/latency control: the expensive agentic ``CosmosRetriever.search`` runs +ONCE (session fixture) and is asserted many ways; the retrieval layer is covered +cheaply and exhaustively via a directly-built ``CorpusRetriever`` and the built +toolset tools (no LLM). + +Opt-in only: the whole module is skipped unless ``RUN_SKF_LIVE`` is truthy and +``.env.local`` points ``ACCOUNT_URI`` at ``skf-rag-test``. It relies on the real +``.env.local`` (the tests/conftest env-isolation deliberately excludes this +folder). Requires ``az login`` with the Cosmos Data Reader role. + +How to run. +1. Run az login for the subscription that owns the skf-rag-test account. +2. Give your identity the Cosmos DB Data Reader role on that account. +3. Fill cosmos-retriever/.env.local with ACCOUNT_URI, the embeddings endpoint and key, and the chat endpoint, model and key. +4. Start the local Qwen3 reranker so the score endpoint answers on port 8011. +5. Set the environment variable RUN_SKF_LIVE to 1. +6. Run pytest on tests/end_to_end/test_skf_live.py. + +See tests/README.md for the full setup. +""" +from __future__ import annotations + +import os +from dataclasses import asdict + +import pytest + + +def _live_enabled() -> bool: + if os.getenv("RUN_SKF_LIVE", "").strip().lower() not in {"1", "true", "yes"}: + return False + try: + from cosmos_retriever.config import get_settings + + s = get_settings() + return bool(s.account_uri and "skf-rag-test" in s.account_uri) + except Exception: + return False + + +pytestmark = pytest.mark.skipif( + not _live_enabled(), + reason="set RUN_SKF_LIVE=1 and configure .env.local for skf-rag-test", +) + +DATABASE = "skf-database" +CONTAINER = "skf-unstructured" +EXPECTED_DIMS = 1536 +# This corpus exposes 9 full-text fields, so retrieval requires explicit selection. +TEXT_FIELDS = ["title", "summary", "content"] + +# ─────────────────────────────── fixtures ───────────────────────────────── + + +@pytest.fixture(scope="session") +def settings(): + from cosmos_retriever.config import get_settings + + return get_settings() + + +@pytest.fixture(scope="session") +def engine(settings): + """The full CosmosRetriever (toolset + reranker + agent), built once.""" + from cosmos_retriever.retriever import CosmosRetriever + + try: + return CosmosRetriever(settings=settings, corpus_name=CONTAINER) + except Exception as exc: # network/RBAC/embedding misconfig + pytest.skip(f"could not build CosmosRetriever: {type(exc).__name__}: {exc}") + + +@pytest.fixture(scope="session") +def corpus_retriever(settings): + """Low-level CorpusRetriever for direct search/grep/read (no LLM).""" + from cosmos_retriever.retrieval import ( + QueryEmbedder, + build_capability_retriever_from_live, + ) + + corpus = settings.resolve_corpus(CONTAINER) + client = settings.build_cosmos_client(corpus) + container = client.get_database_client(corpus.database).get_container_client(corpus.container) + embedder = QueryEmbedder( + client=settings.build_openai_client(corpus), + model=corpus.embed_model, + query_instruction=corpus.embed_query_instruction, + dimensions=corpus.embed_dimensions, + ) + return build_capability_retriever_from_live( + container=container, + database=corpus.database, + embedder=embedder, + override=corpus.schema_override, + ) + + +@pytest.fixture(scope="session") +def embedder(settings): + from cosmos_retriever.retrieval import QueryEmbedder + + corpus = settings.resolve_corpus(CONTAINER) + return QueryEmbedder( + client=settings.build_openai_client(corpus), + model=corpus.embed_model, + query_instruction=corpus.embed_query_instruction, + dimensions=corpus.embed_dimensions, + ) + + +@pytest.fixture(scope="session") +def agent_result(engine): + """One real agentic search, reused across assertions.""" + return engine.search( + "Which SKF bearings are recommended for high-temperature aerospace " + "engine and gearbox applications?", + max_documents=5, + ) + + +@pytest.fixture(scope="session") +def top_doc_id(corpus_retriever): + from cosmos_retriever.retrieval.models import SearchRequest + + items = corpus_retriever.search( + SearchRequest(query="aerospace bearing", limit=5, text_fields=TEXT_FIELDS) + ) + if not items: + pytest.skip("no documents returned for seed query") + return items[0].item_id + + +def _search(corpus_retriever, **kw): + from cosmos_retriever.retrieval.models import SearchRequest + + return corpus_retriever.search(SearchRequest(**kw)) + + +# ═══════════════════════ config / corpus wiring ═══════════════════════════ + + +def test_settings_point_at_skf(settings) -> None: + assert "skf-rag-test" in settings.account_uri + corpus = settings.resolve_corpus(CONTAINER) + assert corpus.embed_model == "text-embedding-3-small" + assert corpus.embed_dimensions == EXPECTED_DIMS + + +def test_engine_corpus_and_toolset(engine) -> None: + assert engine.database_wide is False + assert engine.corpus.database == DATABASE + assert engine.corpus.container == CONTAINER + names = set(engine.toolset.tools) + assert {"search_corpus", "grep_corpus", "read_document", "prune_chunks"} <= names + assert "execute_query" in names # raw query enabled by default + + +def test_embedder_dimension_matches_corpus(embedder) -> None: + vec = embedder.embed("high temperature aerospace bearing") + assert isinstance(vec, list) and len(vec) == EXPECTED_DIMS + assert all(isinstance(x, float) for x in vec[:8]) + + +# ═══════════════════════ low-level retrieval layer ════════════════════════ + + +def test_hybrid_search_returns_ranked_docs(corpus_retriever) -> None: + items = _search( + corpus_retriever, query="aerospace engine gearbox bearing", limit=10, text_fields=TEXT_FIELDS + ) + assert 1 <= len(items) <= 10 + assert all(it.item_id.startswith("skf") for it in items) + assert all(it.text for it in items) + assert [it.rank for it in items] == list(range(len(items))) + assert items[0].retrieval_strategy in {"native_hybrid", "client_fusion"} + + +def test_hybrid_search_is_relevant(corpus_retriever) -> None: + items = _search( + corpus_retriever, query="high temperature aerospace bearing", limit=5, text_fields=TEXT_FIELDS + ) + blob = " ".join(it.text for it in items).lower() + assert "bearing" in blob + + +def test_vector_search_uses_embedding(corpus_retriever) -> None: + items = _search(corpus_retriever, query="corrosion resistant steel bearing", limit=5, mode="vector") + assert items + assert items[0].retrieval_strategy == "vector" + + +def test_full_text_search(corpus_retriever) -> None: + items = _search( + corpus_retriever, query="aerospace bearing", limit=5, mode="text", text_fields=TEXT_FIELDS + ) + assert items + assert items[0].retrieval_strategy == "full_text" + + +def test_search_respects_limit(corpus_retriever) -> None: + items = _search(corpus_retriever, query="bearing", limit=3, text_fields=TEXT_FIELDS) + assert len(items) <= 3 + + +def test_search_ignored_item_ids_excludes(corpus_retriever, top_doc_id) -> None: + items = _search( + corpus_retriever, + query="aerospace bearing", + limit=10, + ignored_item_ids=[top_doc_id], + text_fields=TEXT_FIELDS, + ) + assert top_doc_id not in {it.item_id for it in items} + + +def test_search_specific_text_field(corpus_retriever) -> None: + items = _search(corpus_retriever, query="bearing", limit=5, mode="text", text_fields=["title"]) + assert isinstance(items, list) # field-scoped search must not error + + +def test_search_without_fields_requires_selection(corpus_retriever) -> None: + from cosmos_retriever.retrieval.errors import UnknownField + from cosmos_retriever.retrieval.models import SearchRequest + + # With 9 full-text fields, a text/hybrid search must name the field(s). + with pytest.raises(UnknownField): + corpus_retriever.search(SearchRequest(query="bearing", mode="text")) + + +def test_grep_finds_literal_term(corpus_retriever) -> None: + from cosmos_retriever.retrieval.models import GrepRequest + + hits = corpus_retriever.grep_candidates( + GrepRequest(pattern="bearing", text_field="title", candidate_limit=50, result_limit=5) + ) + assert isinstance(hits, list) + for it in hits: + assert it.item_id.startswith("skf") + + +def test_grep_all_stopword_pattern_is_empty(corpus_retriever) -> None: + from cosmos_retriever.retrieval.models import GrepRequest + + assert corpus_retriever.grep_candidates(GrepRequest(pattern="the and of")) == [] + + +def test_read_document_round_trip(corpus_retriever, top_doc_id) -> None: + from cosmos_retriever.retrieval.models import ReadDocumentRequest + + doc = corpus_retriever.read_document(ReadDocumentRequest(document_id=top_doc_id)) + assert doc.chunk_texts and doc.assembled.strip() + + +def test_read_unknown_document_is_empty(corpus_retriever) -> None: + from cosmos_retriever.retrieval.models import ReadDocumentRequest + + doc = corpus_retriever.read_document(ReadDocumentRequest(document_id="skf-does-not-exist-xyz")) + assert doc.chunk_texts == [] + + +# ═══════════════════════ toolset integration (with reranker) ══════════════ + + +def test_search_corpus_tool_invokes_reranker(engine) -> None: + tool = engine.toolset.get_tool("search_corpus") + text, meta = tool({"query": "high temperature aerospace bearing", "fields": TEXT_FIELDS}) + assert meta is not None and meta.returned_chunk_ids + assert meta.retrieval_s >= 0.0 + assert meta.rerank_s > 0.0 # the live Qwen3-Reranker actually ran + assert "DOCUMENT ID" in text + + +def test_read_document_tool(engine, top_doc_id) -> None: + tool = engine.toolset.get_tool("read_document") + text, _ = tool({"doc_id": top_doc_id}) + assert isinstance(text, str) and text.strip() + + +def test_execute_query_tool_select_and_write_guard(engine) -> None: + tool = engine.toolset.get_tool("execute_query") + ok, _ = tool({"query": "SELECT TOP 1 c.id FROM c", "database": DATABASE, "container": CONTAINER}) + assert "row(s)" in ok + blocked, _ = tool({"query": "DELETE FROM c", "database": DATABASE, "container": CONTAINER}) + assert "only read-only SELECT queries are allowed" in blocked + + +# ═══════════════════════ full agentic end-to-end ══════════════════════════ + + +def test_agent_result_shape(agent_result) -> None: + from cosmos_retriever.retriever import RetrievalResult, RetrievedDocument + + assert isinstance(agent_result, RetrievalResult) + assert agent_result.num_turns >= 1 + assert agent_result.elapsed_s > 0.0 + assert 1 <= len(agent_result.documents) <= 5 + assert all(isinstance(d, RetrievedDocument) for d in agent_result.documents) + + +def test_agent_documents_are_well_formed(agent_result) -> None: + docs = agent_result.documents + assert [d.rank for d in docs] == list(range(len(docs))) # 0..n-1, ordered + assert len({d.id for d in docs}) == len(docs) # unique ids + for d in docs: + assert d.id and d.id.startswith("skf") + assert d.text and d.text.strip() + assert d.justification and d.justification.strip() + + +def test_agent_result_is_relevant(agent_result) -> None: + blob = " ".join(d.text for d in agent_result.documents).lower() + assert "bearing" in blob + + +def test_agent_result_serializes(agent_result) -> None: + d = asdict(agent_result) + assert d["query"] and isinstance(d["documents"], list) + assert set(d) >= {"query", "documents", "num_turns", "elapsed_s"} + + +def test_agent_pool_covers_returned_docs(agent_result) -> None: + # The candidate pool should be a superset of the finally-returned documents. + returned = {d.id for d in agent_result.documents} + if agent_result.pool_doc_ids: + assert returned <= set(agent_result.pool_doc_ids) diff --git a/cosmos-retriever/tests/end_to_end/test_tokenizer_comparison.py b/cosmos-retriever/tests/end_to_end/test_tokenizer_comparison.py new file mode 100644 index 0000000..9a6bf26 --- /dev/null +++ b/cosmos-retriever/tests/end_to_end/test_tokenizer_comparison.py @@ -0,0 +1,170 @@ +"""Cross-tokenizer drift on ENGLISH text: o200k_harmony (what the service counts +with) vs the real tokenizers of the models that actually run. + +CosmosRetriever budgets every request with ``tiktoken.get_encoding("o200k_harmony")`` +regardless of the configured inference model. These tests quantify how far that +estimate drifts from other model families' real tokenizers (Llama 3.1, Claude, +Gemini/Gemma, Qwen, Mistral, and older OpenAI encodings) on pure English prose, +using the actual ``_text_token_counter`` for the harmony side. + +The corpus, tokenizer loaders, and harmony reference counts live in +``tokenizer_panel`` (run ``python tests/end_to_end/tokenizer_panel.py`` to regenerate the +numbers pinned below). English is the convergence case: modern BPE tokenizers +agree closely here, so the drift is small and the bounds are tight. Tokenizers +come from tiktoken (bundled), the local HF cache (Qwen, gpt-oss), and HF community +ports for the closed models (Claude/Gemini publish no tokenizer file). Anything +unavailable offline is skipped, never faked. + +How to run. +1. Install tiktoken and tokenizers into the virtual environment. +2. Allow network access so the model tokenizers can download from Hugging Face, or prime the local cache first. +3. Run pytest on tests/end_to_end/test_tokenizer_comparison.py. + +See tests/README.md for the full setup. +""" +from __future__ import annotations + +import pytest + +pytest.importorskip("tiktoken") +pytest.importorskip("tokenizers") + +from tokenizer_panel import ( # noqa: E402 + HARMONY_PER_SAMPLE, + HARMONY_TOTAL, +) +from tokenizer_panel import ( + hf_counts as _hf_counts, +) +from tokenizer_panel import ( + local_counts as _local_counts, +) +from tokenizer_panel import ( + tik_counts as _tik_counts, +) + + +def _require(counts): + if counts is None: + pytest.skip("tokenizer not available offline") + return counts + + +def _total(counts) -> int: + return sum(counts) + + +def _ratio(counts) -> float: + return _total(counts) / HARMONY_TOTAL + + +# ═══════════════════════ anchor: harmony itself ═══════════════════════════ + + +def test_harmony_canonical_counts_pinned() -> None: + assert HARMONY_PER_SAMPLE == [20, 18, 23, 22] + assert HARMONY_TOTAL == 83 + + +# ═══════════════════════ exact-match family (o200k) ════════════════════════ + + +def test_o200k_base_identical_to_harmony() -> None: + # GPT-4o's tokenizer: harmony is o200k_base + special tokens -> identical on text. + assert _tik_counts("o200k_base") == HARMONY_PER_SAMPLE + + +def test_gpt_oss_local_identical_to_harmony() -> None: + counts = _require(_local_counts("models--openai--gpt-oss-20b")) + assert counts == HARMONY_PER_SAMPLE + + +# ═══════════════════════ legacy OpenAI encodings ══════════════════════════ + + +def test_cl100k_within_two_percent_on_english() -> None: + # gpt-4 / gpt-3.5 / text-embedding-3: on clean English the gap is tiny. + counts = _tik_counts("cl100k_base") + assert counts is not None + assert 1.0 <= _ratio(counts) <= 1.03 + assert all(c >= h for c, h in zip(counts, HARMONY_PER_SAMPLE, strict=True)) + + +@pytest.mark.parametrize("name", ["p50k_base", "r50k_base"]) +def test_legacy_gpt2_era_close_on_english(name: str) -> None: + # Finding: the large legacy drift comes from CJK/whitespace, NOT English -- + # on clean English prose the gpt-2/codex encodings essentially tie harmony. + counts = _tik_counts(name) + assert counts is not None + assert _ratio(counts) <= 1.03 + + +# ═══════════════════════ real other-model tokenizers ══════════════════════ + + +def test_llama3_drift() -> None: + counts = _require(_hf_counts("NousResearch/Meta-Llama-3.1-8B-Instruct")) + assert 1.02 <= _ratio(counts) <= 1.12 # ~6% more on English + + +def test_claude_drift() -> None: + # Anthropic ships no public tokenizer; Xenova/claude-tokenizer is the port. + counts = _require(_hf_counts("Xenova/claude-tokenizer")) + assert 1.0 <= _ratio(counts) <= 1.10 + + +def test_gemini_gemma_drift() -> None: + # Gemini has no public tokenizer; Gemma is the open proxy. + counts = _require(_hf_counts("Xenova/gemma-tokenizer")) + assert 1.0 <= _ratio(counts) <= 1.10 + + +def test_qwen3_local_drift() -> None: + counts = _require(_local_counts("models--Qwen--Qwen3-8B")) + assert 1.0 <= _ratio(counts) <= 1.06 + + +def test_mistral_is_the_english_outlier() -> None: + # Mistral's 32k-vocab tokenizer is the only common one that drifts materially + # on English (~17% more tokens than harmony). + counts = _require(_hf_counts("Xenova/mistral-tokenizer")) + assert 1.08 <= _ratio(counts) <= 1.30 + + +# ═══════════════════════ the actionable findings ══════════════════════════ + + +def test_harmony_is_the_efficiency_floor_on_english() -> None: + # No common tokenizer counts fewer tokens than harmony on English, so the + # service's estimate is a lower bound: it never over-counts English budgets, + # but under-counts (mildly) for every non-o200k model. + panel = { + "cl100k_base": _tik_counts("cl100k_base"), + "llama3.1": _hf_counts("NousResearch/Meta-Llama-3.1-8B-Instruct"), + "claude": _hf_counts("Xenova/claude-tokenizer"), + "gemma": _hf_counts("Xenova/gemma-tokenizer"), + "qwen3": _local_counts("models--Qwen--Qwen3-8B"), + "mistral": _hf_counts("Xenova/mistral-tokenizer"), + } + available = {k: v for k, v in panel.items() if v is not None} + if not available: + pytest.skip("no comparison tokenizers available offline") + for name, counts in available.items(): + assert _total(counts) >= HARMONY_TOTAL, f"{name} counted fewer than harmony" + + +def test_modern_models_all_within_ten_percent_on_english() -> None: + # General closeness statement: every modern-model tokenizer lands within ~10% + # of harmony on English (Mistral, the legacy-vocab outlier, is excluded). + modern = { + "cl100k_base": _tik_counts("cl100k_base"), + "llama3.1": _hf_counts("NousResearch/Meta-Llama-3.1-8B-Instruct"), + "claude": _hf_counts("Xenova/claude-tokenizer"), + "gemma": _hf_counts("Xenova/gemma-tokenizer"), + "qwen3": _local_counts("models--Qwen--Qwen3-8B"), + } + available = {k: v for k, v in modern.items() if v is not None} + if not available: + pytest.skip("no comparison tokenizers available offline") + for name, counts in available.items(): + assert _ratio(counts) <= 1.10, f"{name} drifted more than 10% on English" diff --git a/cosmos-retriever/tests/end_to_end/tokenizer_panel.py b/cosmos-retriever/tests/end_to_end/tokenizer_panel.py new file mode 100644 index 0000000..0ae0eec --- /dev/null +++ b/cosmos-retriever/tests/end_to_end/tokenizer_panel.py @@ -0,0 +1,125 @@ +"""Reproducible token-count panel for the o200k_harmony drift tests. + +Single source of truth for ``test_tokenizer_comparison.py``: the English corpus, +the tokenizer loaders, and the harmony reference counts live here so the test's +pinned numbers can be regenerated on demand. Run it directly to print the table +used to derive/refresh those assertions:: + + python tests/tokenizer_panel.py + +The harmony side uses the real ``CosmosRetriever._text_token_counter``. Tokenizers +come from tiktoken (bundled), the local HF cache (Qwen, gpt-oss), and HF community +ports for the closed models (Claude/Gemini publish no tokenizer file — the Xenova +Claude port and the open Gemma tokenizer are the closest real artifacts). Any +tokenizer unavailable offline is reported as ``n/a`` and skipped by the tests. +""" +from __future__ import annotations + +import glob +from types import SimpleNamespace + +import tiktoken +from tokenizers import Tokenizer + +from cosmos_retriever.retriever import CosmosRetriever + +# ─────────────────────── canonical English corpus ───────────────────────── + +CORPUS: list[str] = [ + "Retrieval augmented generation grounds a language model in the indexed corpus so its answers cite real source documents.", + "The quarterly report shows revenue increased twelve percent while operating costs remained essentially flat year over year.", + "To reset your password, open the settings page, click security, and follow the emailed verification link within one hour.", + "She argued that the experiment, though elegant, failed to control for several confounding variables in the second cohort.", +] + +_HARMONY = tiktoken.get_encoding("o200k_harmony") + + +def hcount(text: str) -> int: + """Real service counter (CosmosRetriever._text_token_counter).""" + return CosmosRetriever._text_token_counter(SimpleNamespace(_tiktoken=_HARMONY), text) + + +HARMONY_PER_SAMPLE: list[int] = [hcount(s) for s in CORPUS] +HARMONY_TOTAL: int = sum(HARMONY_PER_SAMPLE) + + +# ─────────────────────────── tokenizer loaders ──────────────────────────── + + +def tik_counts(name: str) -> list[int] | None: + try: + enc = tiktoken.get_encoding(name) + except Exception: + return None + return [len(enc.encode(s)) for s in CORPUS] + + +def hf_counts(repo: str) -> list[int] | None: + """Load a tokenizer.json from HF (cache first, then network); None if unavailable.""" + try: + from huggingface_hub import hf_hub_download + except Exception: + return None + path = None + for kwargs in ({"local_files_only": True}, {}): + try: + path = hf_hub_download(repo_id=repo, filename="tokenizer.json", **kwargs) + break + except Exception: + continue + if path is None: + return None + tok = Tokenizer.from_file(path) + return [len(tok.encode(s).ids) for s in CORPUS] + + +def local_counts(model_dir: str) -> list[int] | None: + hits = glob.glob(f"/nvme/hf-cache/hub/{model_dir}/snapshots/*/tokenizer.json") + if not hits: + return None + tok = Tokenizer.from_file(hits[0]) + return [len(tok.encode(s).ids) for s in CORPUS] + + +# Ordered panel: label -> zero-arg loader. Keep in sync with the test assertions. +PANEL: dict[str, object] = { + "o200k_base (gpt-4o)": lambda: tik_counts("o200k_base"), + "gpt-oss (harmony)": lambda: local_counts("models--openai--gpt-oss-20b"), + "cl100k (gpt-4/3.5)": lambda: tik_counts("cl100k_base"), + "p50k (codex)": lambda: tik_counts("p50k_base"), + "r50k (gpt-2)": lambda: tik_counts("r50k_base"), + "qwen3": lambda: local_counts("models--Qwen--Qwen3-8B"), + "llama3.1": lambda: hf_counts("NousResearch/Meta-Llama-3.1-8B-Instruct"), + "claude (xenova)": lambda: hf_counts("Xenova/claude-tokenizer"), + "gemma (gemini proxy)": lambda: hf_counts("Xenova/gemma-tokenizer"), + "mistral": lambda: hf_counts("Xenova/mistral-tokenizer"), +} + + +def build_report() -> list[tuple[str, list[int] | None, int | None, float | None]]: + rows: list[tuple[str, list[int] | None, int | None, float | None]] = [] + for label, loader in PANEL.items(): + counts = loader() # type: ignore[operator] + if counts is None: + rows.append((label, None, None, None)) + else: + total = sum(counts) + rows.append((label, counts, total, total / HARMONY_TOTAL)) + return rows + + +def main() -> None: + print(f"corpus: {len(CORPUS)} English samples") + print(f"harmony per-sample = {HARMONY_PER_SAMPLE} total = {HARMONY_TOTAL}\n") + print(f"{'tokenizer':24} {'total':>6} {'ratio':>7} per-sample") + print(f"{'o200k_harmony (service)':24} {HARMONY_TOTAL:>6} {1.0:>7.3f} {HARMONY_PER_SAMPLE}") + for label, counts, total, ratio in build_report(): + if counts is None: + print(f"{label:24} {'n/a':>6} {'n/a':>7}") + else: + print(f"{label:24} {total:>6} {ratio:>7.3f} {counts}") + + +if __name__ == "__main__": + main() diff --git a/cosmos-retriever/tests/unit/test_anthropic_budget.py b/cosmos-retriever/tests/unit/test_anthropic_budget.py new file mode 100644 index 0000000..a4d258e --- /dev/null +++ b/cosmos-retriever/tests/unit/test_anthropic_budget.py @@ -0,0 +1,145 @@ +from __future__ import annotations + +import copy +from typing import Any + +import cosmos_retriever.inference.agent_loop as agent_loop + + +class _FakeTool: + def __init__(self, name: str, output: str) -> None: + self._name = name + self._output = output + self.received: list[tuple[dict, Any]] = [] + + def get_format(self, provider: Any) -> dict: + return {"name": self._name} + + def __call__(self, args: dict, overrides: Any = None) -> tuple[str, None]: + self.received.append((args, overrides)) + return self._output, None + + +class _FakeToolSet: + def __init__(self, tools: dict[str, _FakeTool]) -> None: + self.tools = tools + + def get_tool(self, name: str) -> _FakeTool | None: + return self.tools.get(name) + + +class _FakeResponse: + def __init__(self, payload: dict) -> None: + self._payload = payload + + def raise_for_status(self) -> None: + return None + + def json(self) -> dict: + return self._payload + + +def _install_fake_http(monkeypatch, scripted: list[dict]) -> list[dict]: + """Patch requests.post to replay `scripted` responses; capture sent payloads.""" + captured: list[dict] = [] + seq = iter(scripted) + + def fake_post(url, json=None, headers=None, timeout=None): # noqa: A002 + captured.append(copy.deepcopy(json)) + return _FakeResponse(next(seq)) + + monkeypatch.setattr(agent_loop.requests, "post", fake_post) + return captured + + +def _tool_use(tid: str, name: str, **inp: Any) -> dict: + return {"content": [{"type": "tool_use", "id": tid, "name": name, "input": inp}]} + + +def _final(text: str) -> dict: + return {"content": [{"type": "text", "text": text}]} + + +def _toolset() -> _FakeToolSet: + return _FakeToolSet( + { + "search_corpus": _FakeTool("search_corpus", "# DOCUMENT ID: doc1\nbody text"), + "prune_chunks": _FakeTool("prune_chunks", "pruned"), + } + ) + + +def test_anthropic_annotates_output_and_passes_budget_overrides(monkeypatch) -> None: + ts = _toolset() + captured = _install_fake_http( + monkeypatch, + [ + _tool_use("t1", "search_corpus", query="a"), + _final("j"), + ], + ) + + result = agent_loop.run_anthropic_search( + toolset=ts, # type: ignore[arg-type] + base_url="https://example/v1", + api_key="k", + model="claude", + query="q", + max_documents=5, + max_turns=5, + text_token_counter=len, + threshold_budget=100_000, + token_budget=200_000, + ) + + # The tool was invoked with a budget overrides dict (not the old bare tool(args)). + args, overrides = ts.tools["search_corpus"].received[0] + assert isinstance(overrides, dict) and "ignore_ids" in overrides + + # The observation returned to the model carries the budget annotation. + turn2_msgs = captured[1]["messages"] + tool_result = turn2_msgs[-1]["content"][0] + assert tool_result["type"] == "tool_result" + assert "[Token usage:" in tool_result["content"] + + assert [d.id for d in result.documents] == ["doc1"] + assert result.timing["llm_s"] >= 0.0 + + +def test_anthropic_rejects_non_prune_tools_when_over_budget(monkeypatch) -> None: + ts = _FakeToolSet( + { + "search_corpus": _FakeTool("search_corpus", "X" * 5000), + "prune_chunks": _FakeTool("prune_chunks", "pruned"), + } + ) + captured = _install_fake_http( + monkeypatch, + [ + _tool_use("t1", "search_corpus", query="a"), # huge output blows the budget + _tool_use("t2", "search_corpus", query="b"), # must be rejected + _final(""), + ], + ) + + agent_loop.run_anthropic_search( + toolset=ts, # type: ignore[arg-type] + base_url="https://example/v1", + api_key="k", + model="claude", + query="q", + max_documents=5, + max_turns=5, + text_token_counter=len, + threshold_budget=1000, + token_budget=2000, + ) + + # Second search_corpus was rejected before execution → tool called only once. + assert len(ts.tools["search_corpus"].received) == 1 + + # The rejection observation is what got sent back on the following turn. + turn3_msgs = captured[2]["messages"] + rejected = turn3_msgs[-1]["content"][0] + assert rejected["type"] == "tool_result" + assert "Token budget exceeded" in rejected["content"] diff --git a/cosmos-retriever/tests/unit/test_binding.py b/cosmos-retriever/tests/unit/test_binding.py new file mode 100644 index 0000000..a20ca86 --- /dev/null +++ b/cosmos-retriever/tests/unit/test_binding.py @@ -0,0 +1,121 @@ +from __future__ import annotations + +from cosmos_retriever.retrieval.binding import ( + capabilities_from_metadata, + schema_from_metadata, +) +from cosmos_retriever.retrieval.discovery.models import ContainerMetadata, VectorIndexInfo +from cosmos_retriever.retrieval.errors import InvalidCorpusSchema +from cosmos_retriever.retrieval.schema import DunderChunkCodec +from cosmos_retriever.retrieval.schema_override import SchemaOverride + +_CHUNKED_OVERRIDE = SchemaOverride( + document_id_path="/docid", + chunk_id_path="/id", + chunk_order_path="/chunk_idx", + use_dunder_codec=True, +) + + +def _meta(*, fts: list[str] | None = None, vec: bool = False) -> ContainerMetadata: + vectors = [] + if vec: + vectors = [ + VectorIndexInfo( + path="/embedding", dimensions=2560, distance_function="cosine", indexed=True + ) + ] + return ContainerMetadata( + database="db", container="c", fetched_at=0.0, + partition_key_paths=["/docid"], full_text_paths=fts or [], vector_fields=vectors, + ) + + +def test_hybrid_metadata_builds_hybrid_schema_and_caps() -> None: + m = _meta(fts=["/text"], vec=True) + caps = capabilities_from_metadata(m) + assert caps.native_hybrid_supported and caps.vector_supported and caps.full_text_supported + s = schema_from_metadata(m) + assert [str(p) for p in s.text_paths] == ["/text"] + assert s.vector_fields[0].dimensions == 2560 + # capability dims must match schema dims so the planner's vector check passes + assert caps.vector_fields[0].dimensions == s.vector_fields[0].dimensions + + +def test_text_only_metadata() -> None: + m = _meta(fts=["/text"], vec=False) + caps = capabilities_from_metadata(m) + assert caps.full_text_supported and not caps.vector_supported + s = schema_from_metadata(m) + assert s.vector_fields == [] + assert [str(p) for p in s.text_paths] == ["/text"] + + +def test_vector_only_metadata_has_no_text_and_is_valid() -> None: + m = _meta(fts=None, vec=True) + caps = capabilities_from_metadata(m) + assert caps.vector_supported and not caps.full_text_supported + s = schema_from_metadata(m) # must not raise despite no text field + assert s.text_paths == [] + assert s.resolve_text_fields(None) == [] + assert s.vector_fields[0].dimensions == 2560 + + +def test_structured_metadata_cannot_build_search_schema() -> None: + m = _meta(fts=None, vec=False) + try: + schema_from_metadata(m) + except InvalidCorpusSchema: + return + raise AssertionError("expected InvalidCorpusSchema for a container with no text or vector") + + +def test_no_override_is_item_document_mode() -> None: + s = schema_from_metadata(_meta(fts=["/text"], vec=True)) + assert s.document_id_path is None + assert s.is_item_document_mode is True # graceful: each item is its own document + + +def test_legacy_override_enables_chunk_reconstruction() -> None: + s = schema_from_metadata(_meta(fts=["/text"], vec=True), _CHUNKED_OVERRIDE) + assert str(s.document_id_path) == "/docid" + assert str(s.chunk_order_path) == "/chunk_idx" + assert s.is_item_document_mode is False + assert isinstance(s.identity_codec, DunderChunkCodec) + + +def test_multiple_text_fields_require_explicit_choice() -> None: + from cosmos_retriever.retrieval.errors import UnknownField + + s = schema_from_metadata(_meta(fts=["/content", "/title"], vec=True)) + # No default: with >1 text field the caller must name the field(s). + try: + s.resolve_text_fields(None) + except UnknownField: + pass + else: + raise AssertionError("expected UnknownField when no text field is specified") + # Explicit choices resolve normally. + assert [str(p) for p in s.resolve_text_fields(["content"])] == ["/content"] + assert [str(p) for p in s.resolve_text_fields(["content", "title"])] == [ + "/content", + "/title", + ] + + +def _run_all() -> int: + fns = [v for k, v in sorted(globals().items()) if k.startswith("test_") and callable(v)] + failed = 0 + for fn in fns: + try: + fn() + print(f"PASS {fn.__name__}") + except BaseException as exc: # noqa: BLE001 + failed += 1 + print(f"FAIL {fn.__name__}: {type(exc).__name__}: {exc}") + print(f"\n{len(fns) - failed}/{len(fns)} passed") + return 1 if failed else 0 + + +if __name__ == "__main__": + raise SystemExit(_run_all()) diff --git a/cosmos-retriever/tests/unit/test_cache.py b/cosmos-retriever/tests/unit/test_cache.py new file mode 100644 index 0000000..f93a7c4 --- /dev/null +++ b/cosmos-retriever/tests/unit/test_cache.py @@ -0,0 +1,104 @@ +from __future__ import annotations + +from cosmos_retriever.cache import BoundedTTLCache + + +class _Clock: + def __init__(self) -> None: + self.t = 0.0 + + def __call__(self) -> float: + return self.t + + def advance(self, dt: float) -> None: + self.t += dt + + +def test_put_get_hit_and_miss() -> None: + c: BoundedTTLCache[str, int] = BoundedTTLCache(max_entries=4, ttl_seconds=10.0) + assert c.get("a") is None + c.put("a", 1) + assert c.get("a") == 1 + s = c.stats() + assert s.hits == 1 and s.misses == 1 and s.entries == 1 + + +def test_ttl_expiry() -> None: + clock = _Clock() + c: BoundedTTLCache[str, int] = BoundedTTLCache( + max_entries=4, ttl_seconds=10.0, time_source=clock + ) + c.put("a", 1) + clock.advance(9.9) + assert c.get("a") == 1 + clock.advance(0.2) # now past ttl + assert c.get("a") is None + assert c.stats().expirations == 1 + + +def test_lru_eviction_order() -> None: + evicted: list[tuple[str, int]] = [] + c: BoundedTTLCache[str, int] = BoundedTTLCache( + max_entries=2, ttl_seconds=100.0, on_evict=lambda k, v: evicted.append((k, v)) + ) + c.put("a", 1) + c.put("b", 2) + c.get("a") # touch a so b is now LRU + c.put("c", 3) # evicts b + assert c.get("b") is None + assert c.get("a") == 1 and c.get("c") == 3 + assert evicted == [("b", 2)] + assert c.stats().evictions == 1 + + +def test_put_replaces_value_and_disposes_old() -> None: + disposed: list[int] = [] + c: BoundedTTLCache[str, int] = BoundedTTLCache( + max_entries=4, ttl_seconds=100.0, on_evict=lambda k, v: disposed.append(v) + ) + c.put("a", 1) + c.put("a", 2) + assert c.get("a") == 2 + assert disposed == [1] + assert len(c) == 1 + + +def test_invalidate_and_clear() -> None: + disposed: list[int] = [] + c: BoundedTTLCache[str, int] = BoundedTTLCache( + max_entries=4, ttl_seconds=100.0, on_evict=lambda k, v: disposed.append(v) + ) + c.put("a", 1) + c.put("b", 2) + assert c.invalidate("a") is True + assert c.invalidate("missing") is False + c.clear() + assert len(c) == 0 + assert sorted(disposed) == [1, 2] + + +def test_construction_validates_bounds() -> None: + for kwargs in ({"max_entries": 0}, {"ttl_seconds": 0.0}, {"ttl_seconds": -1.0}): + try: + BoundedTTLCache(**kwargs) # type: ignore[arg-type] + except ValueError: + continue + raise AssertionError(f"expected ValueError for {kwargs}") + + +def _run_all() -> int: + fns = [v for k, v in sorted(globals().items()) if k.startswith("test_") and callable(v)] + failed = 0 + for fn in fns: + try: + fn() + print(f"PASS {fn.__name__}") + except BaseException as exc: # noqa: BLE001 + failed += 1 + print(f"FAIL {fn.__name__}: {type(exc).__name__}: {exc}") + print(f"\n{len(fns) - failed}/{len(fns)} passed") + return 1 if failed else 0 + + +if __name__ == "__main__": + raise SystemExit(_run_all()) diff --git a/cosmos-retriever/tests/unit/test_compiler.py b/cosmos-retriever/tests/unit/test_compiler.py new file mode 100644 index 0000000..cc7e5a5 --- /dev/null +++ b/cosmos-retriever/tests/unit/test_compiler.py @@ -0,0 +1,295 @@ +from __future__ import annotations + +from typing import Any + +import pytest + +from cosmos_retriever.retrieval.compiler import CosmosQueryCompiler +from cosmos_retriever.retrieval.errors import QueryCompilationError +from cosmos_retriever.retrieval.models import EqualsFilter, InFilter, RangeFilter +from cosmos_retriever.retrieval.paths import CosmosPath +from cosmos_retriever.retrieval.schema import CorpusSchema, VectorFieldConfig + +_VEC = CosmosPath.parse("/embedding") +_TEXT = CosmosPath.parse("/text") +_BODY = CosmosPath.parse("/body") + + +def _schema(*, with_docid: bool = True) -> CorpusSchema: + return CorpusSchema( + item_id_path="/id", + document_id_path="/docid" if with_docid else None, + chunk_id_path="/id", + chunk_order_path="/chunk_idx", + title_path="/title", + source_path="/source_type", + text_paths=["/text"], + vector_fields=[VectorFieldConfig(path="/embedding", dimensions=2560)], + metadata_paths={"year": "/year"}, + ) + + +def _compiler(*, with_docid: bool = True) -> CosmosQueryCompiler: + return CosmosQueryCompiler(_schema(with_docid=with_docid)) + + +def _param(q: Any, name: str) -> dict[str, Any]: + for p in q.parameters: + if p["name"] == name: + return p + raise AssertionError(f"no bound parameter {name!r} in {[p['name'] for p in q.parameters]}") + + +def _param_values(q: Any) -> list[Any]: + return [p["value"] for p in q.parameters] + + +# --- projection ----------------------------------------------------------- + + +def test_projection_emits_logical_columns_and_alias_map() -> None: + select, aliases = _compiler().projection("@k0") + assert select.startswith("SELECT TOP @k0 ") + for col in ( + 'c["id"] AS item_id', + 'c["docid"] AS document_id', + 'c["id"] AS chunk_id', + 'c["chunk_idx"] AS chunk_order', + 'c["title"] AS title', + 'c["source_type"] AS source', + 'c["text"] AS txt_0', + 'c["year"] AS md_year', + ): + assert col in select + # text/metadata aliases resolve back to their logical names + assert aliases["txt_0"] == "text" + assert aliases["md_year"] == "year" + + +# --- structured filters --------------------------------------------------- + + +def test_structured_equals_filter_is_parameterized() -> None: + q = _compiler().compile_structured( + limit=10, + filters=[EqualsFilter(logical_field="year", value=2020)], + ignored_item_ids=[], + partition_key=None, + cross_partition=True, + ) + assert q.strategy == "structured" + assert 'c["year"] = @p1' in q.sql + assert _param(q, "@p1")["value"] == 2020 + assert _param(q, "@k0")["value"] == 10 + + +def test_range_filter_emits_both_bounds() -> None: + q = _compiler().compile_structured( + limit=5, + filters=[RangeFilter(logical_field="year", minimum=2000, maximum=2020)], + ignored_item_ids=[], + partition_key=None, + cross_partition=True, + ) + assert '(c["year"] >= @p1 AND c["year"] <= @p2)' in q.sql + assert _param(q, "@p1")["value"] == 2000 + assert _param(q, "@p2")["value"] == 2020 + + +def test_range_filter_with_only_minimum() -> None: + q = _compiler().compile_structured( + limit=5, + filters=[RangeFilter(logical_field="year", minimum=2000)], + ignored_item_ids=[], + partition_key=None, + cross_partition=True, + ) + assert '(c["year"] >= @p1)' in q.sql + assert "<=" not in q.sql + + +def test_in_filter_uses_array_contains() -> None: + q = _compiler().compile_structured( + limit=5, + filters=[InFilter(logical_field="source", values=["news", "blog"])], + ignored_item_ids=[], + partition_key=None, + cross_partition=True, + ) + assert 'ARRAY_CONTAINS(@p1, c["source_type"])' in q.sql + assert _param(q, "@p1")["value"] == ["news", "blog"] + + +def test_ignored_item_ids_add_not_array_contains() -> None: + q = _compiler().compile_structured( + limit=5, + filters=[], + ignored_item_ids=["a", "b"], + partition_key=None, + cross_partition=True, + ) + assert 'NOT ARRAY_CONTAINS(@p1, c["id"])' in q.sql + assert _param(q, "@p1")["value"] == ["a", "b"] + + +def test_no_filters_emits_no_where_clause() -> None: + q = _compiler().compile_structured( + limit=5, + filters=[], + ignored_item_ids=[], + partition_key=None, + cross_partition=True, + ) + assert "WHERE" not in q.sql + + +# --- vector / full-text / hybrid ----------------------------------------- + + +def test_vector_orders_by_vector_distance() -> None: + q = _compiler().compile_vector( + query_vector=[0.1, 0.2, 0.3], + limit=8, + ignored_item_ids=[], + filters=[], + partition_key=None, + cross_partition=True, + vector_path=_VEC, + ) + assert q.strategy == "vector" + assert 'ORDER BY VectorDistance(c["embedding"], @qVec1)' in q.sql + assert _param(q, "@qVec1")["value"] == [0.1, 0.2, 0.3] + + +def test_full_text_single_path_uses_rank_fulltextscore() -> None: + q = _compiler().compile_full_text( + query="the quick brown fox", + limit=5, + ignored_item_ids=[], + filters=[], + partition_key=None, + cross_partition=True, + text_paths=[_TEXT], + ) + assert q.strategy == "full_text" + # stopword "the" dropped; remaining terms rendered as quoted literals + assert 'ORDER BY RANK FullTextScore(c["text"], "quick", "brown", "fox")' in q.sql + + +def test_full_text_multiple_paths_uses_rank_rrf() -> None: + q = _compiler().compile_full_text( + query="quick brown", + limit=5, + ignored_item_ids=[], + filters=[], + partition_key=None, + cross_partition=True, + text_paths=[_TEXT, _BODY], + ) + assert "ORDER BY RANK RRF(" in q.sql + assert 'FullTextScore(c["text"], "quick", "brown")' in q.sql + assert 'FullTextScore(c["body"], "quick", "brown")' in q.sql + + +def test_hybrid_fuses_vector_and_full_text_in_rrf() -> None: + q = _compiler().compile_hybrid( + query="quick brown", + query_vector=[0.1, 0.2], + limit=7, + ignored_item_ids=[], + filters=[], + partition_key=None, + cross_partition=True, + vector_path=_VEC, + text_paths=[_TEXT], + ) + assert q.strategy == "native_hybrid" + assert ( + 'ORDER BY RANK RRF(VectorDistance(c["embedding"], @qVec1), ' + 'FullTextScore(c["text"], "quick", "brown"))' + ) in q.sql + assert _param(q, "@qVec1")["value"] == [0.1, 0.2] + + +# --- document read -------------------------------------------------------- + + +def test_document_read_filters_by_document_id() -> None: + q = _compiler().compile_document_read( + document_id="doc-1", + max_chunks=50, + partition_key=None, + cross_partition=True, + ) + assert q.strategy == "document_read" + assert 'WHERE c["docid"] = @doc1' in q.sql + assert _param(q, "@doc1")["value"] == "doc-1" + + +def test_document_read_without_document_id_path_raises() -> None: + with pytest.raises(QueryCompilationError): + _compiler(with_docid=False).compile_document_read( + document_id="doc-1", + max_chunks=50, + partition_key=None, + cross_partition=True, + ) + + +# --- errors & injection safety ------------------------------------------- + + +def test_unknown_logical_field_raises() -> None: + with pytest.raises(QueryCompilationError): + _compiler().compile_structured( + limit=5, + filters=[EqualsFilter(logical_field="does_not_exist", value=1)], + ignored_item_ids=[], + partition_key=None, + cross_partition=True, + ) + + +def test_filter_values_are_bound_never_inlined() -> None: + """User-controlled values must go through @params, not the SQL string.""" + malicious = "2020'; DROP TABLE Foo--" + q = _compiler().compile_structured( + limit=5, + filters=[EqualsFilter(logical_field="year", value=malicious)], + ignored_item_ids=[], + partition_key=None, + cross_partition=True, + ) + assert "DROP TABLE" not in q.sql + assert 'c["year"] = @p1' in q.sql + assert malicious in _param_values(q) + + +@pytest.mark.parametrize( + "kind, expected_strategy", + [ + ("hybrid", "native_hybrid"), + ("vector", "vector"), + ("full_text", "full_text"), + ("structured", "structured"), + ], +) +def test_each_query_type_reports_its_strategy(kind: str, expected_strategy: str) -> None: + c = _compiler() + common = dict( + limit=5, + ignored_item_ids=[], + filters=[], + partition_key=None, + cross_partition=True, + ) + if kind == "hybrid": + q = c.compile_hybrid(query="q", query_vector=[0.1], vector_path=_VEC, text_paths=[_TEXT], **common) + elif kind == "vector": + q = c.compile_vector(query_vector=[0.1], vector_path=_VEC, **common) + elif kind == "full_text": + q = c.compile_full_text(query="q", text_paths=[_TEXT], **common) + else: + q = c.compile_structured(**common) + assert q.strategy == expected_strategy + assert q.sql.startswith("SELECT TOP @k0 ") diff --git a/cosmos-retriever/tests/unit/test_config.py b/cosmos-retriever/tests/unit/test_config.py new file mode 100644 index 0000000..8c6c1c3 --- /dev/null +++ b/cosmos-retriever/tests/unit/test_config.py @@ -0,0 +1,610 @@ +"""Exhaustive tests for `cosmos_retriever.config`. + +Covers the internal *resolution* logic of the settings module without touching +Azure / OpenAI / Baseten: all external clients (CosmosClient, OpenAI, +AzureOpenAI, Azure credentials, Baseten PerformanceClient) are patched with +recording fakes, and ``os.environ`` is driven via monkeypatch. + +Not duplicated here (already in test_runtime_config.py): apply_structural_overrides, +RuntimeConfig.structural_key, and RuntimeConfig validators / extra-forbid. + +Deterministic settings are built with ``_env_file=None`` so no .env leaks and +explicit init kwargs win over any ambient environment variables. +""" +from __future__ import annotations + +import dataclasses +import json +import logging +import sys +from types import SimpleNamespace + +import pytest + +from cosmos_retriever import config +from cosmos_retriever.config import ( + CorpusConfig, + RetrieverSettings, + ServerConfigUpdate, + get_config, + get_settings, +) +from cosmos_retriever.retrieval.schema_override import SchemaOverride + +# ────────────────────────────── helpers ─────────────────────────────────── + + +def _settings(**kw) -> RetrieverSettings: + return RetrieverSettings(_env_file=None, **kw) + + +class FakeCosmosClient: + instances: list = [] + + def __init__(self, account_uri, credential=None): + self.account_uri = account_uri + self.credential = credential + self.db_requests: list = [] + FakeCosmosClient.instances.append(self) + + def get_database_client(self, name): + self.db_requests.append(name) + return SimpleNamespace(database=name) + + +class FakeOpenAI: + instances: list = [] + + def __init__(self, **kwargs): + self.kwargs = kwargs + FakeOpenAI.instances.append(self) + + +class FakeAzureOpenAI: + instances: list = [] + + def __init__(self, **kwargs): + self.kwargs = kwargs + FakeAzureOpenAI.instances.append(self) + + +@pytest.fixture(autouse=True) +def _reset_fakes(): + FakeCosmosClient.instances = [] + FakeOpenAI.instances = [] + FakeAzureOpenAI.instances = [] + yield + + +@pytest.fixture +def patch_clients(monkeypatch): + monkeypatch.setattr(config, "CosmosClient", FakeCosmosClient) + monkeypatch.setattr(config, "OpenAI", FakeOpenAI) + monkeypatch.setattr(config, "AzureCliCredential", lambda: "CLI-CRED") + monkeypatch.setattr(config, "DefaultAzureCredential", lambda: "DEFAULT-CRED") + monkeypatch.setattr("openai.AzureOpenAI", FakeAzureOpenAI) + + +# ═══════════════════════════ init_logging ═════════════════════════════════ + + +def test_init_logging_smoke() -> None: + config.init_logging(app_level=logging.DEBUG, lib_level=logging.WARNING, colors=False) + + +# ═══════════════════════════ CorpusConfig ═════════════════════════════════ + + +def test_corpusconfig_is_frozen_with_defaults() -> None: + c = CorpusConfig( + container="c", + account_uri="https://a", + database="db", + embed_base_url=None, + embed_api_key=None, + embed_model="m", + ) + assert c.embed_query_instruction is None + assert c.embed_dimensions is None + assert c.cosmos_key is None and c.schema_override is None + with pytest.raises(dataclasses.FrozenInstanceError): + c.container = "x" # type: ignore[misc] + + +# ═══════════════════════ inference_backend validator ══════════════════════ + + +@pytest.mark.parametrize( + "value,expected", + [("OPENAI_CHAT", "openai_chat"), (" Openai_Responses ", "openai_responses"), + ("anthropic_messages", "anthropic_messages")], +) +def test_inference_backend_normalized(value: str, expected: str) -> None: + assert _settings(inference_backend=value).inference_backend == expected + + +def test_inference_backend_invalid_rejected() -> None: + with pytest.raises(ValueError, match="INFERENCE_BACKEND must be one of"): + _settings(inference_backend="bogus") + + +def test_schema_override_coerced_on_settings() -> None: + s = _settings(cosmos_retriever_schema_override={"item_id_path": "/id"}) + assert isinstance(s.cosmos_retriever_schema_override, SchemaOverride) + assert s.cosmos_retriever_schema_override.item_id_path == "/id" + + +# ═══════════════════════════ _load_registry ═══════════════════════════════ + + +def test_load_registry_none_returns_empty() -> None: + assert _settings()._load_registry() == {} + + +def test_load_registry_inline_json() -> None: + reg = json.dumps({"corp": {"embed_model": "m"}}) + assert _settings(corpus_registry=reg)._load_registry() == {"corp": {"embed_model": "m"}} + + +def test_load_registry_invalid_json_raises() -> None: + with pytest.raises(ValueError, match="not valid JSON"): + _settings(corpus_registry="{not json")._load_registry() + + +def test_load_registry_non_dict_raises() -> None: + with pytest.raises(ValueError, match="must be a JSON object"): + _settings(corpus_registry="[1, 2, 3]")._load_registry() + + +def test_load_registry_missing_file_raises(tmp_path) -> None: + missing = tmp_path / "nope.json" + with pytest.raises(FileNotFoundError, match="missing file"): + _settings(corpus_registry_file=str(missing))._load_registry() + + +def test_load_registry_reads_file(tmp_path) -> None: + path = tmp_path / "reg.json" + path.write_text(json.dumps({"corp": {"embed_model": "fm"}}), encoding="utf-8") + assert _settings(corpus_registry_file=str(path))._load_registry() == { + "corp": {"embed_model": "fm"} + } + + +# ═══════════════════════ _lookup_registry_entry ═══════════════════════════ + + +def test_lookup_prefers_db_container_over_container() -> None: + reg = {"db/corp": {"x": 1}, "corp": {"x": 2}, "db": {"x": 3}} + assert RetrieverSettings._lookup_registry_entry(reg, "db", "corp") == {"x": 1} + + +def test_lookup_falls_back_to_container_then_database() -> None: + assert RetrieverSettings._lookup_registry_entry({"corp": {"x": 2}}, "db", "corp") == {"x": 2} + assert RetrieverSettings._lookup_registry_entry({"db": {"x": 3}}, "db", "corp") == {"x": 3} + + +def test_lookup_without_database_only_tries_container() -> None: + assert RetrieverSettings._lookup_registry_entry({"corp": {"x": 2}}, None, "corp") == {"x": 2} + assert RetrieverSettings._lookup_registry_entry({"db": {"x": 3}}, None, "corp") is None + + +def test_lookup_no_match_returns_none() -> None: + assert RetrieverSettings._lookup_registry_entry({"other": {}}, "db", "corp") is None + + +# ═══════════════════════════ resolve_corpus ═══════════════════════════════ + + +def test_resolve_corpus_no_target_raises() -> None: + with pytest.raises(ValueError, match="No Cosmos container specified"): + _settings().resolve_corpus() + + +def test_resolve_corpus_default_no_database_raises() -> None: + s = _settings(cosmos_corpus_container="corp", account_uri="https://a") + with pytest.raises(ValueError, match="No Cosmos database specified"): + s.resolve_corpus() + + +def test_resolve_corpus_default_no_account_uri_raises() -> None: + s = _settings(cosmos_corpus_container="corp", cosmos_database="db") + with pytest.raises(ValueError, match="no fallback ACCOUNT_URI"): + s.resolve_corpus() + + +def test_resolve_corpus_default_happy() -> None: + s = _settings( + cosmos_corpus_container="corp", + cosmos_database="db", + account_uri="https://acct", + embed_endpoint="https://embed", + openai_api_key="sk-embed", + openai_embedding_model="text-embed", + embed_query_instruction="inst", + openai_embedding_dimensions=256, + cosmos_retriever_schema_override={"item_id_path": "/id"}, + ) + c = s.resolve_corpus() + assert c.container == "corp" and c.database == "db" + assert c.account_uri == "https://acct" + assert c.embed_base_url == "https://embed" + assert c.embed_api_key.get_secret_value() == "sk-embed" + assert c.embed_model == "text-embed" + assert c.embed_query_instruction == "inst" + assert c.embed_dimensions == 256 + assert isinstance(c.schema_override, SchemaOverride) + + +def test_resolve_corpus_entry_env_keys_and_entry_base(monkeypatch) -> None: + monkeypatch.setenv("MY_EMBED_KEY", "embed-secret") + monkeypatch.setenv("MY_COSMOS_KEY", "cosmos-secret") + reg = json.dumps( + { + "corp": { + "embed_base_url": "https://entry-embed", + "embed_model": "entry-model", + "embed_api_key_env": "MY_EMBED_KEY", + "cosmos_key_env": "MY_COSMOS_KEY", + "account_uri": "https://entry-acct", + "database": "entry-db", + } + } + ) + s = _settings(cosmos_corpus_container="corp", corpus_registry=reg) + c = s.resolve_corpus() + assert c.account_uri == "https://entry-acct" + assert c.database == "entry-db" + assert c.embed_base_url == "https://entry-embed" + assert c.embed_model == "entry-model" + assert c.embed_api_key.get_secret_value() == "embed-secret" + assert c.cosmos_key.get_secret_value() == "cosmos-secret" + + +def test_resolve_corpus_entry_without_base_uses_server_endpoint(monkeypatch) -> None: + reg = json.dumps({"corp": {"embed_model": "entry-model", "account_uri": "https://a", "database": "db"}}) + s = _settings( + cosmos_corpus_container="corp", + corpus_registry=reg, + embed_endpoint="https://server-embed", + openai_api_key="server-key", + ) + c = s.resolve_corpus() + assert c.embed_base_url == "https://server-embed" + assert c.embed_api_key.get_secret_value() == "server-key" # inherits server key + + +def test_resolve_corpus_entry_missing_database_raises() -> None: + reg = json.dumps({"corp": {"embed_model": "m", "account_uri": "https://a"}}) + s = _settings(cosmos_corpus_container="corp", corpus_registry=reg) + with pytest.raises(ValueError, match="no database configured"): + s.resolve_corpus() + + +def test_resolve_corpus_entry_missing_account_uri_raises() -> None: + reg = json.dumps({"corp": {"embed_model": "m", "database": "db"}}) + s = _settings(cosmos_corpus_container="corp", corpus_registry=reg) + with pytest.raises(ValueError, match="no account_uri configured"): + s.resolve_corpus() + + +def test_resolve_corpus_entry_dimensions_zero_is_respected() -> None: + reg = json.dumps( + {"corp": {"embed_model": "m", "account_uri": "https://a", "database": "db", "embed_dimensions": 0}} + ) + s = _settings(cosmos_corpus_container="corp", corpus_registry=reg, openai_embedding_dimensions=999) + assert s.resolve_corpus().embed_dimensions == 0 # explicit 0 wins over settings default + + +def test_resolve_corpus_entry_schema_override_and_fallback() -> None: + reg = json.dumps( + {"corp": {"embed_model": "m", "account_uri": "https://a", "database": "db", + "schema_override": {"item_id_path": "/entry"}}} + ) + s = _settings(cosmos_corpus_container="corp", corpus_registry=reg) + assert s.resolve_corpus().schema_override.item_id_path == "/entry" + + reg2 = json.dumps({"corp": {"embed_model": "m", "account_uri": "https://a", "database": "db"}}) + s2 = _settings( + cosmos_corpus_container="corp", + corpus_registry=reg2, + cosmos_retriever_schema_override={"item_id_path": "/default"}, + ) + assert s2.resolve_corpus().schema_override.item_id_path == "/default" + + +def test_resolve_corpus_explicit_container_argument() -> None: + reg = json.dumps({"other": {"embed_model": "m", "account_uri": "https://a", "database": "db"}}) + s = _settings(corpus_registry=reg) + assert s.resolve_corpus(container="other").container == "other" + + +# ═══════════════════════════ _cosmos_credential ═══════════════════════════ + + +def test_cosmos_credential_default_uses_cli(patch_clients, monkeypatch) -> None: + monkeypatch.delenv("COSMOS_USE_DEFAULT_CREDENTIAL", raising=False) + assert _settings()._cosmos_credential() == "CLI-CRED" + + +@pytest.mark.parametrize("flag", ["1", "true", "YES"]) +def test_cosmos_credential_default_azure(patch_clients, monkeypatch, flag: str) -> None: + monkeypatch.setenv("COSMOS_USE_DEFAULT_CREDENTIAL", flag) + assert _settings()._cosmos_credential() == "DEFAULT-CRED" + + +# ═══════════════════════ build_cosmos_client / database ════════════════════ + + +def _corpus(**kw) -> CorpusConfig: + base = dict( + container="c", account_uri="https://acct", database="db", + embed_base_url=None, embed_api_key=None, embed_model="m", + ) + base.update(kw) + return CorpusConfig(**base) + + +def test_build_cosmos_client_with_key(patch_clients) -> None: + from pydantic import SecretStr + + client = _settings().build_cosmos_client(_corpus(cosmos_key=SecretStr("mykey"))) + assert client.account_uri == "https://acct" + assert client.credential == "mykey" # raw secret, not credential object + + +def test_build_cosmos_client_without_key_uses_credential(patch_clients, monkeypatch) -> None: + monkeypatch.delenv("COSMOS_USE_DEFAULT_CREDENTIAL", raising=False) + client = _settings().build_cosmos_client(_corpus()) + assert client.credential == "CLI-CRED" + + +def test_build_cosmos_database_chains(patch_clients) -> None: + db = _settings().build_cosmos_database(_corpus(database="mydb")) + assert db.database == "mydb" + + +# ═══════════════════════════ build_openai_client ══════════════════════════ + + +def test_build_openai_client_with_base_url_and_key(patch_clients) -> None: + from pydantic import SecretStr + + client = _settings().build_openai_client( + _corpus(embed_base_url="https://e", embed_api_key=SecretStr("k")) + ) + assert client.kwargs == {"base_url": "https://e", "api_key": "k"} + + +def test_build_openai_client_no_base_url_empty_key(patch_clients) -> None: + client = _settings().build_openai_client(_corpus(embed_base_url=None, embed_api_key=None)) + assert "base_url" not in client.kwargs + assert client.kwargs["api_key"] == "EMPTY" + + +# ═══════════════════════════ use_*_backend props ══════════════════════════ + + +def test_backend_properties() -> None: + assert _settings(inference_backend="openai_chat").use_chat_backend is True + assert _settings(inference_backend="openai_responses").use_responses_backend is True + assert _settings(inference_backend="anthropic_messages").use_anthropic_backend is True + + +def test_generic_backend_is_chat_or_responses() -> None: + assert _settings(inference_backend="openai_chat").use_generic_llm_backend is True + assert _settings(inference_backend="openai_responses").use_generic_llm_backend is True + assert _settings(inference_backend="anthropic_messages").use_generic_llm_backend is False + + +# ═══════════════════════════ build_chat_client ════════════════════════════ + + +def test_build_chat_client_requires_base_url() -> None: + with pytest.raises(ValueError, match="CHAT_BASE_URL must be set"): + _settings(chat_model="m").build_chat_client() + + +def test_build_chat_client_requires_model() -> None: + with pytest.raises(ValueError, match="CHAT_MODEL"): + _settings(chat_base_url="https://c").build_chat_client() + + +def test_build_chat_client_plain_openai(patch_clients) -> None: + client = _settings(chat_base_url="https://c", chat_model="m").build_chat_client() + assert isinstance(client, FakeOpenAI) + assert client.kwargs == {"base_url": "https://c", "api_key": "EMPTY"} + + +def test_build_chat_client_azure_when_api_version(patch_clients) -> None: + from pydantic import SecretStr + + s = _settings(chat_base_url="https://c", chat_model="m", chat_api_version="2024-01") + s.chat_api_key = SecretStr("chatkey") + client = s.build_chat_client() + assert isinstance(client, FakeAzureOpenAI) + assert client.kwargs == { + "azure_endpoint": "https://c", + "api_key": "chatkey", + "api_version": "2024-01", + } + + +# ═══════════════════════════ apply_server_updates ═════════════════════════ + + +def test_apply_server_updates_maps_and_wraps() -> None: + s = _settings(chat_model="old", cosmos_retriever_cache_max_entries=4) + update = ServerConfigUpdate( + chat_model="new", + cache_max_entries=10, + token_budget=8192, + chat_api_key="secret", + schema_override={"item_id_path": "/id"}, + ) + new = s.apply_server_updates(update) + assert new.chat_model == "new" + assert new.cosmos_retriever_cache_max_entries == 10 + assert new.cosmos_retriever_token_budget == 8192 + assert new.chat_api_key.get_secret_value() == "secret" + assert isinstance(new.cosmos_retriever_schema_override, SchemaOverride) + # original untouched (deep copy) + assert s.chat_model == "old" + assert s.cosmos_retriever_cache_max_entries == 4 + + +def test_apply_server_updates_only_provided_fields() -> None: + s = _settings(chat_model="keep", chat_max_tokens=1000) + new = s.apply_server_updates(ServerConfigUpdate(chat_model="changed")) + assert new.chat_model == "changed" + assert new.chat_max_tokens == 1000 # untouched + + +# ═══════════════════════════ redacted_config ══════════════════════════════ + + +def test_redacted_config_masks_secrets_and_dumps_override() -> None: + from pydantic import SecretStr + + s = _settings( + cosmos_database="db", + cosmos_retriever_schema_override={"item_id_path": "/id"}, + ) + s.chat_api_key = SecretStr("x") + s.cosmos_key = None + red = s.redacted_config() + assert red["chat_api_key"] == "***set***" + assert red["cosmos_key"] is None + assert red["cosmos_database"] == "db" + assert red["cache_max_entries"] == s.cosmos_retriever_cache_max_entries + assert red["schema_override"]["item_id_path"] == "/id" + + +def test_redacted_config_none_schema_override() -> None: + assert _settings().redacted_config()["schema_override"] is None + + +# ═══════════════════════ get_* client delegation ══════════════════════════ + + +def test_get_cosmos_client_delegates(patch_clients, monkeypatch) -> None: + from pydantic import SecretStr + + corpus = _corpus(cosmos_key=SecretStr("k"), account_uri="https://x") + monkeypatch.setattr(RetrieverSettings, "resolve_corpus", lambda self, container=None: corpus) + client = _settings().get_cosmos_client() + assert client.account_uri == "https://x" and client.credential == "k" + + +def test_get_cosmos_database_delegates(patch_clients, monkeypatch) -> None: + corpus = _corpus(database="thedb") + monkeypatch.setattr(RetrieverSettings, "resolve_corpus", lambda self, container=None: corpus) + assert _settings().get_cosmos_database().database == "thedb" + + +def test_get_openai_client_delegates(patch_clients, monkeypatch) -> None: + from pydantic import SecretStr + + corpus = _corpus(embed_base_url="https://e", embed_api_key=SecretStr("k")) + monkeypatch.setattr(RetrieverSettings, "resolve_corpus", lambda self, container=None: corpus) + client = _settings().get_openai_client() + assert client.kwargs == {"base_url": "https://e", "api_key": "k"} + + +# ═══════════════════════════ get_baseten_client ═══════════════════════════ + + +def test_get_baseten_client_requires_credentials() -> None: + with pytest.raises(ValueError, match="BASETEN_API_KEY and BASETEN_MODEL_URL"): + _settings().get_baseten_client() + + +def test_get_baseten_client_happy(monkeypatch) -> None: + from pydantic import SecretStr + + built: list = [] + + class FakePerf: + def __init__(self, base_url, api_key): + self.base_url = base_url + self.api_key = api_key + built.append(self) + + monkeypatch.setitem( + sys.modules, "baseten_performance_client", SimpleNamespace(PerformanceClient=FakePerf) + ) + s = _settings(baseten_model_url="https://bt") + s.baseten_api_key = SecretStr("btkey") + client = s.get_baseten_client() + assert client.base_url == "https://bt" and client.api_key == "btkey" + + +# ═══════════════════ get_settings / get_config / log level ════════════════ + + +def test_get_settings_is_cached_and_calls_init_logging(monkeypatch) -> None: + get_settings.cache_clear() + calls: list = [] + monkeypatch.setattr(config, "init_logging", lambda **kw: calls.append(kw)) + first = get_settings() + second = get_settings() + assert first is second # lru_cache + assert len(calls) == 1 # init_logging only on the cache-miss build + + +def test_get_config_returns_settings() -> None: + get_settings.cache_clear() + assert get_config() is get_settings() + + +@pytest.mark.parametrize( + "level,expected", + [("debug", logging.DEBUG), ("INFO", logging.INFO), ("Warning", logging.WARNING), + ("bogus", logging.INFO)], +) +def test_log_level_to_int(level: str, expected: int) -> None: + assert config._log_level_to_int(level) == expected + + +# ═══════════════════════════ ServerConfigUpdate ═══════════════════════════ + + +def test_server_update_backend_validator() -> None: + assert ServerConfigUpdate(inference_backend="OPENAI_CHAT").inference_backend == "openai_chat" + with pytest.raises(ValueError, match="inference_backend must be one of"): + ServerConfigUpdate(inference_backend="nope") + + +def test_server_update_schema_override_coerced() -> None: + u = ServerConfigUpdate(schema_override={"item_id_path": "/id"}) + assert isinstance(u.schema_override, SchemaOverride) + + +def test_server_update_extra_forbidden() -> None: + with pytest.raises(ValueError): + ServerConfigUpdate(unknown_field="x") + + +@pytest.mark.parametrize( + "kwargs", + [ + {"chat_temperature": 2.5}, + {"chat_max_tokens": 100}, + {"chat_max_turns": 0}, + {"token_budget": 100}, + {"threshold_budget": 100}, + {"cache_max_entries": 2000}, + {"cache_ttl_seconds": 0.0}, + {"search_display_limit": 100}, + ], +) +def test_server_update_numeric_bounds_rejected(kwargs: dict) -> None: + with pytest.raises(ValueError): + ServerConfigUpdate(**kwargs) + + +# ═══════════════════════════════ __all__ ══════════════════════════════════ + + +def test_all_exports_present() -> None: + for name in config.__all__: + assert hasattr(config, name), name diff --git a/cosmos-retriever/tests/unit/test_discovery.py b/cosmos-retriever/tests/unit/test_discovery.py new file mode 100644 index 0000000..d8bf83a --- /dev/null +++ b/cosmos-retriever/tests/unit/test_discovery.py @@ -0,0 +1,206 @@ +from __future__ import annotations + +from typing import Any + +from cosmos_retriever.retrieval.discovery import ( + CapabilityProfiler, + ResourceCatalog, + parse_container_metadata, +) + + +def _props( + *, + pk: list[str] | None = None, + fts_paths: list[str] | None = None, + vectors: list[dict[str, Any]] | None = None, + vector_indexes: list[str] | None = None, +) -> dict[str, Any]: + idx: dict[str, Any] = {"includedPaths": [{"path": "/*"}], "excludedPaths": []} + if fts_paths: + idx["fullTextIndexes"] = [{"path": p} for p in fts_paths] + if vector_indexes: + idx["vectorIndexes"] = [{"path": p, "type": "diskANN"} for p in vector_indexes] + props: dict[str, Any] = { + "id": "c", + "partitionKey": {"paths": pk or ["/id"], "kind": "Hash"}, + "indexingPolicy": idx, + } + if fts_paths: + props["fullTextPolicy"] = {"fullTextPaths": [{"path": p} for p in fts_paths]} + if vectors: + props["vectorEmbeddingPolicy"] = {"vectorEmbeddings": vectors} + return props + + +class _FakeContainer: + def __init__(self, props: dict[str, Any]) -> None: + self._props = props + self.reads = 0 + + class _Conn: + last_response_headers = {"etag": "W/\"1\""} + + self.client_connection = _Conn() + + def read(self) -> dict[str, Any]: + self.reads += 1 + return self._props + + +class _FakeDatabase: + def __init__(self, containers: dict[str, _FakeContainer]) -> None: + self._containers = containers + + def get_container_client(self, name: str) -> _FakeContainer: + return self._containers[name] + + def list_containers(self) -> list[dict[str, str]]: + return [{"id": n} for n in self._containers] + + +class _FakeClient: + def __init__(self, dbs: dict[str, _FakeDatabase]) -> None: + self._dbs = dbs + + def get_database_client(self, name: str) -> _FakeDatabase: + return self._dbs[name] + + def list_databases(self) -> list[dict[str, str]]: + return [{"id": n} for n in self._dbs] + + +class _FakeConnection: + def __init__(self, client: _FakeClient) -> None: + self._client = client + + def client(self) -> _FakeClient: + return self._client + + +def _catalog(scenarios: dict[str, dict[str, Any]], **kw: Any) -> ResourceCatalog: + containers = {name: _FakeContainer(props) for name, props in scenarios.items()} + client = _FakeClient({"db": _FakeDatabase(containers)}) + return ResourceCatalog(_FakeConnection(client), **kw) + + +# ---- capability scenarios (7-10) ------------------------------------------- + +def test_text_only_container() -> None: + p = CapabilityProfiler().profile( + parse_container_metadata("db", "c", _props(fts_paths=["/text"])) + ) + assert p.can_full_text.value is True + assert p.can_vector.value is False + assert p.can_native_hybrid.value is False + assert p.recommended_strategies == ["full_text", "item_lookup"] + + +def test_vector_only_container() -> None: + p = CapabilityProfiler().profile( + parse_container_metadata( + "db", "c", + _props( + vectors=[{"path": "/embedding", "dimensions": 2560, "distanceFunction": "cosine"}], + vector_indexes=["/embedding"], + ), + ) + ) + assert p.can_vector.value is True + assert p.can_full_text.value is False + assert p.can_native_hybrid.value is False + assert p.vector_fields[0].dimensions == 2560 + + +def test_vector_embedding_without_index_is_not_searchable() -> None: + # embedding policy present but NO vector index -> capability must be False + p = CapabilityProfiler().profile( + parse_container_metadata( + "db", "c", + _props(vectors=[{"path": "/embedding", "dimensions": 2560}]), # no vector_indexes + ) + ) + assert p.can_vector.value is False + + +def test_hybrid_container() -> None: + p = CapabilityProfiler().profile( + parse_container_metadata( + "db", "c", + _props( + fts_paths=["/text"], + vectors=[{"path": "/embedding", "dimensions": 2560, "distanceFunction": "cosine"}], + vector_indexes=["/embedding"], + ), + ) + ) + assert p.can_native_hybrid.value is True + assert p.recommended_strategies[0] == "native_hybrid" + + +def test_structured_container_has_only_item_lookup() -> None: + p = CapabilityProfiler().profile(parse_container_metadata("db", "c", _props())) + assert p.can_full_text.value is False + assert p.can_vector.value is False + assert p.can_native_hybrid.value is False + assert p.can_item_lookup.value is True + assert p.recommended_strategies == ["item_lookup"] + + +# ---- discovery + catalog lifecycle ----------------------------------------- + +def test_discovery_lists_databases_and_containers() -> None: + cat = _catalog({"t": _props(fts_paths=["/text"]), "s": _props()}) + assert cat.databases() == ["db"] + assert set(cat.containers("db")) == {"t", "s"} + + +def test_profile_caches_and_refresh_forces_reread() -> None: + scenarios = {"t": _props(fts_paths=["/text"])} + containers = {n: _FakeContainer(p) for n, p in scenarios.items()} + client = _FakeClient({"db": _FakeDatabase(containers)}) + cat = ResourceCatalog(_FakeConnection(client), ttl_seconds=1000.0) + + cat.profile("db", "t") + cat.profile("db", "t") + assert containers["t"].reads == 1 # second call served from cache + + cat.refresh("db", "t") + cat.profile("db", "t") + assert containers["t"].reads == 2 # re-read after refresh + + +def test_invalidate_forces_reread() -> None: + scenarios = {"t": _props(fts_paths=["/text"])} + containers = {n: _FakeContainer(p) for n, p in scenarios.items()} + cat = ResourceCatalog(_FakeConnection(_FakeClient({"db": _FakeDatabase(containers)}))) + cat.profile("db", "t") + cat.invalidate("db", "t") + cat.profile("db", "t") + assert containers["t"].reads == 2 + + +def test_bounded_cache_evicts() -> None: + scenarios = {f"c{i}": _props() for i in range(5)} + cat = _catalog(scenarios, max_entries=2) + for n in scenarios: + cat.profile("db", n) + assert cat.cached_container_count() == 2 + + +def _run_all() -> int: + fns = [v for k, v in sorted(globals().items()) if k.startswith("test_") and callable(v)] + failed = 0 + for fn in fns: + try: + fn() + print(f"PASS {fn.__name__}") + except AssertionError as exc: + failed += 1 + print(f"FAIL {fn.__name__}: {exc}") + print(f"\n{len(fns) - failed}/{len(fns)} passed") + return 1 if failed else 0 + + +if __name__ == "__main__": + raise SystemExit(_run_all()) diff --git a/cosmos-retriever/tests/unit/test_document_resolvers.py b/cosmos-retriever/tests/unit/test_document_resolvers.py new file mode 100644 index 0000000..2084734 --- /dev/null +++ b/cosmos-retriever/tests/unit/test_document_resolvers.py @@ -0,0 +1,267 @@ +"""Exhaustive tests for `cosmos_retriever.retrieval.document_resolvers`. + +Resolvers turn a ReadDocumentRequest into a NormalizedDocument by compiling a +read, executing it, and assembling chunk text. Tests fake the compiler / +executor / schema and patch ``row_text_fields`` / ``assemble_text`` so every +resolver's id derivation, partition/cross-partition decision, chunk ordering, +and warnings are asserted without Cosmos. build_document_resolver's dispatch +matrix is covered too. +""" +from __future__ import annotations + +import pytest + +from cosmos_retriever.retrieval import document_resolvers as dr +from cosmos_retriever.retrieval.document_resolvers import ( + DEFAULT_MAX_CHUNKS, + ChunkedDocumentResolver, + CrossPartitionChunkedDocumentResolver, + DocumentResolver, + ItemIsDocumentResolver, + build_document_resolver, +) +from cosmos_retriever.retrieval.errors import ( + CrossPartitionQueryDisabled, + DocumentResolutionUnsupported, +) +from cosmos_retriever.retrieval.models import ReadDocumentRequest + +# ────────────────────────────── fakes ───────────────────────────────────── + + +class FakeCodec: + def to_document_id(self, raw: str) -> str: + return f"doc::{raw}" + + +class FakeSchema: + def __init__(self, codec=None, item_mode=False, document_id_path="/doc", + pk_is_doc_id=False): + self.identity_codec = codec + self.is_item_document_mode = item_mode + self.document_id_path = document_id_path + self.partition_key_is_document_id = pk_is_doc_id + + +class FakeCompiled: + def __init__(self): + self.projected_aliases = "ALIASES" + self.warnings: list[str] = [] + + +class FakeCompiler: + def __init__(self): + self.calls: list[tuple[str, dict]] = [] + self.compiled: list[FakeCompiled] = [] + + def _mk(self, method, kwargs): + c = FakeCompiled() + self.calls.append((method, kwargs)) + self.compiled.append(c) + return c + + def compile_structured(self, **kw): + return self._mk("structured", kw) + + def compile_document_read(self, **kw): + return self._mk("document_read", kw) + + +class FakeExecutor: + def __init__(self, rows=None): + self.rows = rows if rows is not None else [] + self.ran: list = [] + + def run(self, compiled): + self.ran.append(compiled) + return self.rows + + +class FakePolicy: + def __init__(self, allow_cross_partition_document_read=True): + self.allow_cross_partition_document_read = allow_cross_partition_document_read + + +@pytest.fixture(autouse=True) +def _patch_normalization(monkeypatch): + # row_text_fields returns the row; assemble_text picks its "text". + monkeypatch.setattr(dr, "row_text_fields", lambda r, aliases: r) + monkeypatch.setattr(dr, "assemble_text", lambda fields: fields["text"]) + + +def _row(item_id, text, chunk_order=None): + return {"item_id": item_id, "text": text, "chunk_order": chunk_order} + + +# ═══════════════════════ base helpers ═════════════════════════════════════ + + +def test_document_resolver_is_abstract() -> None: + with pytest.raises(TypeError): + DocumentResolver(FakeSchema(), FakeCompiler(), FakeExecutor(), FakePolicy()) # type: ignore[abstract] + + +def test_derive_document_id_prefers_document_id() -> None: + r = ItemIsDocumentResolver(FakeSchema(), FakeCompiler(), FakeExecutor(), FakePolicy()) + assert r._derive_document_id(ReadDocumentRequest(document_id="D", item_id="I")) == "D" + + +def test_derive_document_id_falls_back_to_item_id() -> None: + r = ItemIsDocumentResolver(FakeSchema(), FakeCompiler(), FakeExecutor(), FakePolicy()) + assert r._derive_document_id(ReadDocumentRequest(item_id="I")) == "I" + + +def test_derive_document_id_empty_when_both_missing() -> None: + r = ItemIsDocumentResolver(FakeSchema(), FakeCompiler(), FakeExecutor(), FakePolicy()) + assert r._derive_document_id(ReadDocumentRequest()) == "" + + +def test_derive_document_id_applies_codec() -> None: + r = ItemIsDocumentResolver(FakeSchema(codec=FakeCodec()), FakeCompiler(), FakeExecutor(), FakePolicy()) + assert r._derive_document_id(ReadDocumentRequest(document_id="D")) == "doc::D" + + +def test_sorted_rows_orders_by_chunk_order_none_as_zero() -> None: + rows = [_row("c", "c", 2), _row("a", "a", None), _row("b", "b", 1)] + ordered = DocumentResolver._sorted_rows(rows) + assert [r["item_id"] for r in ordered] == ["a", "b", "c"] + + +# ═══════════════════════ ItemIsDocumentResolver ═══════════════════════════ + + +def test_item_is_document_resolve_wiring() -> None: + compiler, executor = FakeCompiler(), FakeExecutor(rows=[_row("i1", "t1"), _row("i2", "t2")]) + r = ItemIsDocumentResolver(FakeSchema(item_mode=True), compiler, executor, FakePolicy()) + doc = r.resolve(ReadDocumentRequest(item_id="ITEM")) + + method, kw = compiler.calls[0] + assert method == "structured" + assert kw["limit"] == 1 + assert kw["ignored_item_ids"] == [] + assert kw["cross_partition"] is True # no partition key + only_filter = kw["filters"][0] + assert only_filter.logical_field == "item_id" and only_filter.value == "ITEM" + assert doc.document_id == "ITEM" + assert doc.chunk_texts == ["t1", "t2"] + assert doc.chunk_ids == ["i1", "i2"] + + +def test_item_is_document_prefers_item_id_over_document_id() -> None: + compiler = FakeCompiler() + r = ItemIsDocumentResolver(FakeSchema(), compiler, FakeExecutor(), FakePolicy()) + r.resolve(ReadDocumentRequest(item_id="I", document_id="D")) + assert compiler.calls[0][1]["filters"][0].value == "I" + + +def test_item_is_document_partition_key_sets_cross_false() -> None: + compiler = FakeCompiler() + r = ItemIsDocumentResolver(FakeSchema(), compiler, FakeExecutor(), FakePolicy()) + r.resolve(ReadDocumentRequest(item_id="I", partition_key="pk")) + assert compiler.calls[0][1]["cross_partition"] is False + assert compiler.calls[0][1]["partition_key"] == "pk" + + +# ═══════════════════════ ChunkedDocumentResolver ══════════════════════════ + + +def test_chunked_resolve_wiring_and_sorting() -> None: + compiler = FakeCompiler() + executor = FakeExecutor(rows=[_row("c2", "second", 2), _row("c1", "first", 1)]) + r = ChunkedDocumentResolver(FakeSchema(), compiler, executor, FakePolicy()) + doc = r.resolve(ReadDocumentRequest(document_id="DOC")) + + method, kw = compiler.calls[0] + assert method == "document_read" + assert kw["document_id"] == "DOC" + assert kw["max_chunks"] == DEFAULT_MAX_CHUNKS + assert kw["partition_key"] == "DOC" # falls back to doc id + assert kw["cross_partition"] is False + assert doc.chunk_texts == ["first", "second"] # sorted by chunk_order + assert doc.chunk_ids == ["c1", "c2"] + assert doc.warnings == [] + + +def test_chunked_custom_max_chunks_and_partition_key() -> None: + compiler = FakeCompiler() + r = ChunkedDocumentResolver(FakeSchema(), compiler, FakeExecutor(), FakePolicy()) + r.resolve(ReadDocumentRequest(document_id="DOC", max_chunks=5, partition_key="PK")) + kw = compiler.calls[0][1] + assert kw["max_chunks"] == 5 + assert kw["partition_key"] == "PK" + + +def test_chunked_applies_codec_to_document_id() -> None: + compiler = FakeCompiler() + r = ChunkedDocumentResolver(FakeSchema(codec=FakeCodec()), compiler, FakeExecutor(), FakePolicy()) + r.resolve(ReadDocumentRequest(document_id="raw")) + assert compiler.calls[0][1]["document_id"] == "doc::raw" + + +# ═══════════════════ CrossPartitionChunkedDocumentResolver ════════════════ + + +def test_cross_partition_disabled_raises() -> None: + r = CrossPartitionChunkedDocumentResolver( + FakeSchema(), FakeCompiler(), FakeExecutor(), + FakePolicy(allow_cross_partition_document_read=False), + ) + with pytest.raises(CrossPartitionQueryDisabled): + r.resolve(ReadDocumentRequest(document_id="D")) + + +def test_cross_partition_resolve_wiring_and_warning() -> None: + compiler = FakeCompiler() + executor = FakeExecutor(rows=[_row("c2", "b", 2), _row("c1", "a", 1)]) + r = CrossPartitionChunkedDocumentResolver(FakeSchema(), compiler, executor, FakePolicy()) + doc = r.resolve(ReadDocumentRequest(document_id="DOC")) + + kw = compiler.calls[0][1] + assert kw["document_id"] == "DOC" + assert kw["partition_key"] is None + assert kw["cross_partition"] is True # no partition key + assert doc.chunk_texts == ["a", "b"] # sorted + assert doc.warnings == ["cross-partition document reconstruction"] + + +def test_cross_partition_with_key_sets_cross_false() -> None: + compiler = FakeCompiler() + r = CrossPartitionChunkedDocumentResolver(FakeSchema(), compiler, FakeExecutor(), FakePolicy()) + r.resolve(ReadDocumentRequest(document_id="D", partition_key="PK")) + kw = compiler.calls[0][1] + assert kw["partition_key"] == "PK" + assert kw["cross_partition"] is False + + +# ═══════════════════════ build_document_resolver ══════════════════════════ + + +def test_build_returns_item_resolver_in_item_mode() -> None: + resolver = build_document_resolver( + FakeSchema(item_mode=True), FakeCompiler(), FakeExecutor(), FakePolicy() + ) + assert isinstance(resolver, ItemIsDocumentResolver) + + +def test_build_raises_without_document_id_path() -> None: + with pytest.raises(DocumentResolutionUnsupported): + build_document_resolver( + FakeSchema(item_mode=False, document_id_path=None), + FakeCompiler(), FakeExecutor(), FakePolicy(), + ) + + +def test_build_returns_chunked_when_pk_is_document_id() -> None: + resolver = build_document_resolver( + FakeSchema(item_mode=False, document_id_path="/d", pk_is_doc_id=True), + FakeCompiler(), FakeExecutor(), FakePolicy(), + ) + assert isinstance(resolver, ChunkedDocumentResolver) + + +def test_build_returns_cross_partition_otherwise() -> None: + resolver = build_document_resolver( + FakeSchema(item_mode=False, document_id_path="/d", pk_is_doc_id=False), + FakeCompiler(), FakeExecutor(), FakePolicy(), + ) + assert isinstance(resolver, CrossPartitionChunkedDocumentResolver) diff --git a/cosmos-retriever/tests/unit/test_embeddings.py b/cosmos-retriever/tests/unit/test_embeddings.py new file mode 100644 index 0000000..d23808a --- /dev/null +++ b/cosmos-retriever/tests/unit/test_embeddings.py @@ -0,0 +1,97 @@ +"""Exhaustive tests for `cosmos_retriever.retrieval.embedding`. + +QueryEmbedder wraps an OpenAI embeddings client. Tests use a fake client to +assert the exact create() call (model, input list, encoding_format, optional +dimensions), the instruction-prefixing of the query text, and that the first +embedding vector is returned. +""" +from __future__ import annotations + +from types import SimpleNamespace + +from cosmos_retriever.retrieval.embedding import QueryEmbedder + + +class FakeEmbeddings: + def __init__(self, embedding): + self.embedding = embedding + self.calls: list[dict] = [] + + def create(self, **kwargs): + self.calls.append(kwargs) + return SimpleNamespace(data=[SimpleNamespace(embedding=self.embedding)]) + + +class FakeClient: + def __init__(self, embedding): + self.embeddings = FakeEmbeddings(embedding) + + +def _embedder(embedding=None, **kw) -> tuple[QueryEmbedder, FakeClient]: + client = FakeClient(embedding if embedding is not None else [0.1, 0.2]) + return QueryEmbedder(client=client, model=kw.pop("model", "m"), **kw), client + + +def test_embed_basic_call_and_return() -> None: + emb, client = _embedder(embedding=[0.1, 0.2, 0.3]) + out = emb.embed("hello") + assert out == [0.1, 0.2, 0.3] + assert client.embeddings.calls == [ + {"model": "m", "input": ["hello"], "encoding_format": "float"} + ] + + +def test_embed_no_dimensions_kwarg_when_unset() -> None: + emb, client = _embedder() + emb.embed("q") + assert "dimensions" not in client.embeddings.calls[0] + + +def test_embed_includes_dimensions_when_set() -> None: + emb, client = _embedder(dimensions=256) + emb.embed("q") + assert client.embeddings.calls[0]["dimensions"] == 256 + + +def test_embed_applies_instruction_prefix() -> None: + emb, client = _embedder(query_instruction="Find docs") + emb.embed("cats") + assert client.embeddings.calls[0]["input"] == ["Instruct: Find docs\nQuery: cats"] + + +def test_embed_no_instruction_leaves_text_untouched() -> None: + emb, client = _embedder(query_instruction=None) + emb.embed("plain") + assert client.embeddings.calls[0]["input"] == ["plain"] + + +def test_embed_empty_instruction_is_ignored() -> None: + emb, client = _embedder(query_instruction="") + emb.embed("plain") + assert client.embeddings.calls[0]["input"] == ["plain"] # falsy instruction skipped + + +def test_embed_uses_configured_model() -> None: + emb, client = _embedder(model="text-embedding-3-large") + emb.embed("q") + assert client.embeddings.calls[0]["model"] == "text-embedding-3-large" + + +def test_embed_returns_first_vector_only() -> None: + client = FakeClient(None) + client.embeddings = SimpleNamespace( + calls=[], + create=lambda **kw: SimpleNamespace( + data=[SimpleNamespace(embedding=[1.0]), SimpleNamespace(embedding=[9.0])] + ), + ) + emb = QueryEmbedder(client=client, model="m") + assert emb.embed("q") == [1.0] + + +def test_embed_instruction_and_dimensions_together() -> None: + emb, client = _embedder(query_instruction="Ins", dimensions=64) + emb.embed("q") + call = client.embeddings.calls[0] + assert call["input"] == ["Instruct: Ins\nQuery: q"] + assert call["dimensions"] == 64 diff --git a/cosmos-retriever/tests/unit/test_expressions.py b/cosmos-retriever/tests/unit/test_expressions.py new file mode 100644 index 0000000..e7a33c3 --- /dev/null +++ b/cosmos-retriever/tests/unit/test_expressions.py @@ -0,0 +1,134 @@ +"""Exhaustive tests for `cosmos_retriever.retrieval.expressions`. + +Covers FTS tokenization (Unicode, lowering, dedup, stopwords, term cap, the +all-stopword degenerate case) and, critically for security, the escaping in +``fts_literal_args`` that keeps a hostile term from breaking out of the quoted +full-text literal it is embedded in. +""" +from __future__ import annotations + +import pytest + +from cosmos_retriever.retrieval.expressions import ( + _FTS_MAX_TERMS, + fts_literal_args, + tokenize_for_fts, +) + +# ═══════════════════════════ tokenize_for_fts ═════════════════════════════ + + +@pytest.mark.parametrize("query", ["", " ", "\n\t", "!!! ??? ...", "---"]) +def test_tokenize_empty_or_punctuation_only(query: str) -> None: + assert tokenize_for_fts(query) == [] + + +def test_tokenize_simple() -> None: + assert tokenize_for_fts("hello world") == ["hello", "world"] + + +def test_tokenize_lowercases() -> None: + assert tokenize_for_fts("Hello WORLD FooBar") == ["hello", "world", "foobar"] + + +def test_tokenize_dedupes_preserving_first_order() -> None: + assert tokenize_for_fts("bb aa bb cc aa") == ["bb", "aa", "cc"] + + +def test_tokenize_dedupe_is_case_insensitive() -> None: + assert tokenize_for_fts("Hello hello HELLO") == ["hello"] + + +def test_tokenize_splits_on_punctuation() -> None: + assert tokenize_for_fts("foo, bar. baz! qux?") == ["foo", "bar", "baz", "qux"] + + +def test_tokenize_keeps_digits_and_underscore() -> None: + assert tokenize_for_fts("abc 123 foo_bar") == ["abc", "123", "foo_bar"] + + +def test_tokenize_dedupes_numbers() -> None: + assert tokenize_for_fts("1 1 2 2 3") == ["1", "2", "3"] + + +def test_tokenize_removes_stopwords() -> None: + assert tokenize_for_fts("the cat and the dog") == ["cat", "dog"] + + +def test_tokenize_apostrophe_splits_and_drops_stopword_half() -> None: + # "don" is a stopword, apostrophe is a delimiter, so only "t" survives. + assert tokenize_for_fts("don't") == ["t"] + + +def test_tokenize_unicode_accented_words() -> None: + assert tokenize_for_fts("Café Über") == ["café", "über"] + + +def test_tokenize_unicode_cjk() -> None: + assert tokenize_for_fts("机器 学习 机器") == ["机器", "学习"] + + +def test_tokenize_all_stopwords_reduces_to_empty() -> None: + # SECURITY / degenerate case: an all-English-stopword query yields zero terms. + assert tokenize_for_fts("the and of a to in is it") == [] + assert tokenize_for_fts("THE AND OF") == [] + + +def test_tokenize_caps_at_max_terms() -> None: + query = " ".join(f"w{i}" for i in range(_FTS_MAX_TERMS + 10)) + result = tokenize_for_fts(query) + assert len(result) == _FTS_MAX_TERMS + assert result == [f"w{i}" for i in range(_FTS_MAX_TERMS)] + + +def test_tokenize_cap_counts_distinct_only() -> None: + # Duplicates must not consume the term budget. + distinct = [f"t{i}" for i in range(_FTS_MAX_TERMS)] + query = " ".join(distinct + distinct + ["extra_beyond_cap"]) + result = tokenize_for_fts(query) + assert len(result) == _FTS_MAX_TERMS + assert "extra_beyond_cap" not in result # cap already reached by distinct set + + +def test_tokenize_injection_characters_are_stripped() -> None: + # Quotes / semicolons / brackets are non-word chars -> removed at tokenization. + assert tokenize_for_fts('drop"; SELECT') == ["drop", "select"] + + +# ═══════════════════════════ fts_literal_args ═════════════════════════════ + + +def test_fts_literal_args_empty_is_empty_string() -> None: + assert fts_literal_args([]) == "" + + +def test_fts_literal_args_single_term() -> None: + assert fts_literal_args(["foo"]) == '"foo"' + + +def test_fts_literal_args_multiple_terms_joined() -> None: + assert fts_literal_args(["foo", "bar", "baz"]) == '"foo", "bar", "baz"' + + +def test_fts_literal_args_escapes_embedded_quote() -> None: + # A quote inside a term is backslash-escaped so it can't close the literal. + assert fts_literal_args(['a"b']) == '"a\\"b"' + + +def test_fts_literal_args_escapes_backslash() -> None: + assert fts_literal_args(["a\\b"]) == '"a\\\\b"' + + +def test_fts_literal_args_escapes_backslash_before_quote() -> None: + # Backslash is doubled first so it can't neutralize the quote escaping. + assert fts_literal_args(['\\"']) == '"\\\\\\""' + + +def test_fts_literal_args_quote_breakout_payload_is_neutralized() -> None: + result = fts_literal_args(['" OR "1"="1']) + # every embedded double-quote is preceded by a backslash; no bare '"' survives + # between the outer wrapping quotes. + inner = result[1:-1] + assert '\\"' in inner + assert inner.replace('\\"', "").count('"') == 0 + assert result == '"\\" OR \\"1\\"=\\"1"' diff --git a/cosmos-retriever/tests/unit/test_normalization.py b/cosmos-retriever/tests/unit/test_normalization.py new file mode 100644 index 0000000..b3458ef --- /dev/null +++ b/cosmos-retriever/tests/unit/test_normalization.py @@ -0,0 +1,163 @@ +"""Exhaustive tests for `cosmos_retriever.retrieval.normalization`. + +Covers row_text_fields (txt_ prefix + alias gating, falsy coercion), +assemble_text (empty / single / multi-field joining and name filtering/order), +and normalize_rows (metadata extraction, id str-coercion with None passthrough, +chunk_order int-only, rank offset, channel copying) against real RetrievedItem. +""" +from __future__ import annotations + +from cosmos_retriever.retrieval.normalization import ( + assemble_text, + normalize_rows, + row_text_fields, +) + +# ═══════════════════════════ row_text_fields ══════════════════════════════ + + +def test_row_text_fields_maps_aliased_txt_keys() -> None: + row = {"txt_a": "hello", "txt_b": "world"} + aliases = {"txt_a": "title", "txt_b": "body"} + assert row_text_fields(row, aliases) == {"title": "hello", "body": "world"} + + +def test_row_text_fields_ignores_non_txt_keys() -> None: + row = {"item_id": "x", "md_foo": "m", "title": "t"} + aliases = {"item_id": "item", "title": "T"} + assert row_text_fields(row, aliases) == {} + + +def test_row_text_fields_ignores_txt_key_absent_from_aliases() -> None: + assert row_text_fields({"txt_a": "v"}, {}) == {} + + +def test_row_text_fields_coerces_none_and_falsy_to_empty() -> None: + row = {"txt_a": None, "txt_b": "", "txt_c": 0} + aliases = {"txt_a": "a", "txt_b": "b", "txt_c": "c"} + assert row_text_fields(row, aliases) == {"a": "", "b": "", "c": ""} + + +def test_row_text_fields_empty_row() -> None: + assert row_text_fields({}, {"txt_a": "a"}) == {} + + +# ═══════════════════════════ assemble_text ════════════════════════════════ + + +def test_assemble_text_empty_fields_returns_empty() -> None: + assert assemble_text({}) == "" + + +def test_assemble_text_single_field_no_header() -> None: + assert assemble_text({"body": "hello"}) == "hello" + + +def test_assemble_text_single_field_none_value() -> None: + assert assemble_text({"body": None}) == "" + + +def test_assemble_text_multiple_fields_joined_with_headers() -> None: + result = assemble_text({"title": "T", "body": "B"}) + assert result == "[title]\nT\n\n[body]\nB" + + +def test_assemble_text_names_filter_to_present_only() -> None: + fields = {"title": "T", "body": "B"} + assert assemble_text(fields, ["body"]) == "B" # single selected -> no header + + +def test_assemble_text_names_control_order() -> None: + fields = {"a": "A", "b": "B"} + assert assemble_text(fields, ["b", "a"]) == "[b]\nB\n\n[a]\nA" + + +def test_assemble_text_names_none_present_returns_empty() -> None: + assert assemble_text({"a": "A"}, ["missing"]) == "" + + +def test_assemble_text_names_skip_absent() -> None: + fields = {"a": "A", "b": "B"} + assert assemble_text(fields, ["a", "missing", "b"]) == "[a]\nA\n\n[b]\nB" + + +# ═══════════════════════════ normalize_rows ═══════════════════════════════ + + +def test_normalize_rows_empty() -> None: + assert normalize_rows([], strategy="s") == [] + + +def test_normalize_rows_full_row() -> None: + row = { + "item_id": 42, + "document_id": 7, + "chunk_id": 3, + "chunk_order": 5, + "txt_a": "hello", + "md_score": 0.9, + "md_source_tag": "x", + "title": "T", + "source": "S", + } + aliases = {"txt_a": "body"} + items = normalize_rows( + [row], strategy="vector", channels=["vector"], + projected_aliases=aliases, + ) + it = items[0] + assert it.item_id == "42" # str-coerced + assert it.document_id == "7" and it.chunk_id == "3" + assert it.chunk_order == 5 + assert it.text == "hello" + assert it.text_fields == {"body": "hello"} + assert it.title == "T" and it.source == "S" + assert it.metadata == {"score": 0.9, "source_tag": "x"} # md_ prefix stripped + assert it.retrieval_strategy == "vector" + assert it.retrieval_channels == ["vector"] + assert it.rank == 0 + + +def test_normalize_rows_none_ids_passthrough() -> None: + row = {"item_id": None, "document_id": None, "chunk_id": None} + it = normalize_rows([row], strategy="s")[0] + assert it.item_id == "None" # item_id always str-coerced, even None + assert it.document_id is None # optional ids stay None + assert it.chunk_id is None + + +def test_normalize_rows_non_int_chunk_order_becomes_none() -> None: + for bad in ("5", 1.5, None): + it = normalize_rows([{"item_id": "x", "chunk_order": bad}], strategy="s")[0] + assert it.chunk_order is None + + +def test_normalize_rows_channels_default_empty_and_copied() -> None: + channels = ["vector"] + it = normalize_rows([{"item_id": "x"}], strategy="s", channels=channels)[0] + assert it.retrieval_channels == ["vector"] + assert it.retrieval_channels is not channels # defensive copy + + it2 = normalize_rows([{"item_id": "x"}], strategy="s")[0] + assert it2.retrieval_channels == [] + + +def test_normalize_rows_rank_uses_start_rank_offset() -> None: + rows = [{"item_id": "a"}, {"item_id": "b"}, {"item_id": "c"}] + items = normalize_rows(rows, strategy="s", start_rank=10) + assert [it.rank for it in items] == [10, 11, 12] + + +def test_normalize_rows_queried_text_fields_filter_display() -> None: + row = {"item_id": "x", "txt_a": "A", "txt_b": "B"} + aliases = {"txt_a": "fa", "txt_b": "fb"} + it = normalize_rows( + [row], strategy="s", projected_aliases=aliases, queried_text_fields=["fb"] + )[0] + assert it.text == "B" # only fb selected + + +def test_normalize_rows_metadata_only_md_prefixed() -> None: + row = {"item_id": "x", "md_a": 1, "b": 2, "txt_c": "c"} + it = normalize_rows([row], strategy="s", projected_aliases={"txt_c": "c"})[0] + assert it.metadata == {"a": 1} diff --git a/cosmos-retriever/tests/unit/test_orchestration.py b/cosmos-retriever/tests/unit/test_orchestration.py new file mode 100644 index 0000000..63d2e91 --- /dev/null +++ b/cosmos-retriever/tests/unit/test_orchestration.py @@ -0,0 +1,177 @@ +from __future__ import annotations + +from types import SimpleNamespace + +from cosmos_retriever.retrieval.models import RetrievedItem, SearchRequest +from cosmos_retriever.retrieval.orchestration import ( + ContainerTarget, + MultiContainerRetriever, + fuse_rrf, + select_search_targets, +) + + +def _items(prefix: str, n: int) -> list[RetrievedItem]: + return [RetrievedItem(item_id=f"{prefix}{i}", text=f"{prefix}{i}") for i in range(1, n + 1)] + + +class _FakeRetriever: + def __init__(self, items: list[RetrievedItem]) -> None: + self._items = items + self.last_request: SearchRequest | None = None + + def search(self, request: SearchRequest) -> list[RetrievedItem]: + self.last_request = request + return self._items + + +def _req() -> SearchRequest: + return SearchRequest(query="q", limit=50) + + +def test_fuse_rrf_interleaves_and_tags_metadata() -> None: + ta = ContainerTarget("db", "A") + tb = ContainerTarget("db", "B") + fused = fuse_rrf([(ta, _items("a", 3)), (tb, _items("b", 3))], limit=3) + assert [it.item_id for it in fused] == ["a1", "b1", "a2"] + assert fused[0].metadata["container"] == "A" + assert fused[1].metadata["container"] == "B" + assert "rrf" in fused[0].raw_scores + assert fused[0].rank == 0 and fused[1].rank == 1 + + +def test_multi_search_fans_out_and_fuses() -> None: + ta = ContainerTarget("db", "A") + tb = ContainerTarget("db", "B") + resolvers = {ta: _FakeRetriever(_items("a", 2)), tb: _FakeRetriever(_items("b", 2))} + mcr = MultiContainerRetriever(lambda t: resolvers[t]) + result = mcr.search([ta, tb], _req(), final_limit=3) + assert result.errors == {} + assert set(result.searched) == {ta, tb} + assert result.per_container_counts == {"db/A": 2, "db/B": 2} + assert [it.item_id for it in result.items] == ["a1", "b1", "a2"] + + +def test_partial_failure_is_isolated() -> None: + ta = ContainerTarget("db", "A") + tb = ContainerTarget("db", "B") + + class _Boom: + def search(self, request: SearchRequest) -> list[RetrievedItem]: + raise RuntimeError("container offline") + + resolvers = {ta: _FakeRetriever(_items("a", 2)), tb: _Boom()} + mcr = MultiContainerRetriever(lambda t: resolvers[t]) + result = mcr.search([ta, tb], _req()) + assert result.searched == [ta] + assert "db/B" in result.errors and "container offline" in result.errors["db/B"] + assert [it.item_id for it in result.items] == ["a1", "a2"] + + +def test_per_container_limit_is_applied() -> None: + ta = ContainerTarget("db", "A") + fake = _FakeRetriever(_items("a", 2)) + mcr = MultiContainerRetriever(lambda t: fake) + mcr.search([ta], _req(), per_container_limit=7) + assert fake.last_request is not None and fake.last_request.limit == 7 + + +def test_empty_targets() -> None: + mcr = MultiContainerRetriever(lambda t: _FakeRetriever([])) + result = mcr.search([], _req()) + assert result.items == [] and result.searched == [] + + +def test_duplicate_targets_deduped() -> None: + ta = ContainerTarget("db", "A") + fake = _FakeRetriever(_items("a", 2)) + calls = {"n": 0} + + def resolver(t: ContainerTarget) -> _FakeRetriever: + calls["n"] += 1 + return fake + + mcr = MultiContainerRetriever(resolver) + mcr.search([ta, ta, ta], _req()) + assert calls["n"] == 1 + + +def test_select_search_targets_filters_incapable() -> None: + def cap(fts: bool, vec: bool) -> SimpleNamespace: + return SimpleNamespace( + can_full_text=SimpleNamespace(value=fts), + can_vector=SimpleNamespace(value=vec), + ) + + profiles = {"A": cap(True, True), "B": cap(False, False), "C": cap(False, True)} + + catalog = SimpleNamespace( + containers=lambda db: ["A", "B", "C"], + profile=lambda db, name: profiles[name], + ) + targets = select_search_targets(catalog, "db") + assert [t.container for t in targets] == ["A", "C"] + + +def test_cross_collection_retriever_fuses_probes_and_greps() -> None: + from cosmos_retriever.retrieval.models import ( + GrepRequest, + NormalizedDocument, + ReadDocumentRequest, + ) + from cosmos_retriever.retrieval.orchestration import CrossCollectionRetriever + + ta = ContainerTarget("db", "A") + tb = ContainerTarget("db", "B") + + class _FR: + def __init__(self, items: list[RetrievedItem], doc_text: str) -> None: + self._items = items + self._doc_text = doc_text + self.schema = object() + + def search(self, request: SearchRequest) -> list[RetrievedItem]: + return self._items + + def grep_candidates(self, request: GrepRequest) -> list[RetrievedItem]: + return self._items + + def read_document(self, request: ReadDocumentRequest) -> NormalizedDocument: + return NormalizedDocument( + document_id="d", chunk_texts=[self._doc_text], chunk_ids=["c"] + ) + + ra = _FR(_items("a", 2), doc_text="") # yields an empty document + rb = _FR(_items("b", 2), doc_text="hello") # yields real text + + x = CrossCollectionRetriever([ta, tb], {ta: ra, tb: rb}) + + # search fuses across collections with RRF interleaving + fused = x.search(SearchRequest(query="q", limit=3)) + assert [it.item_id for it in fused] == ["a1", "b1", "a2"] + + # read_document probes past the empty collection to the one with content + doc = x.read_document(ReadDocumentRequest(document_id="d")) + assert doc.assembled == "hello" + + # grep fans out across every collection + g = x.grep_candidates(GrepRequest(pattern="x", candidate_limit=10)) + assert len(g) == 4 + + +def _run_all() -> int: + fns = [v for k, v in sorted(globals().items()) if k.startswith("test_") and callable(v)] + failed = 0 + for fn in fns: + try: + fn() + print(f"PASS {fn.__name__}") + except BaseException as exc: # noqa: BLE001 + failed += 1 + print(f"FAIL {fn.__name__}: {type(exc).__name__}: {exc}") + print(f"\n{len(fns) - failed}/{len(fns)} passed") + return 1 if failed else 0 + + +if __name__ == "__main__": + raise SystemExit(_run_all()) diff --git a/cosmos-retriever/tests/unit/test_paths.py b/cosmos-retriever/tests/unit/test_paths.py new file mode 100644 index 0000000..253c15a --- /dev/null +++ b/cosmos-retriever/tests/unit/test_paths.py @@ -0,0 +1,167 @@ +"""Exhaustive tests for `cosmos_retriever.retrieval.paths`. + +Covers CosmosPath.parse (validation + segment rules), render (alias + +escaping), __str__ / round-trip, frozen-model semantics (immutability, +equality, hashability), and coerce_path. +""" +from __future__ import annotations + +import pytest +from pydantic import ValidationError + +from cosmos_retriever.retrieval.errors import UnsafeCosmosPath +from cosmos_retriever.retrieval.paths import CosmosPath, coerce_path + +# ═══════════════════════════ parse: identity ══════════════════════════════ + + +def test_parse_returns_same_instance_for_cosmospath() -> None: + p = CosmosPath.parse("/id") + assert CosmosPath.parse(p) is p + + +# ═══════════════════════════ parse: type errors ═══════════════════════════ + + +@pytest.mark.parametrize("bad", [123, None, ["/a"], b"/a", 1.5, {"a": 1}]) +def test_parse_non_string_raises(bad) -> None: + with pytest.raises(UnsafeCosmosPath, match="must be a string"): + CosmosPath.parse(bad) + + +def test_parse_error_includes_type_name() -> None: + with pytest.raises(UnsafeCosmosPath, match="got int"): + CosmosPath.parse(5) + + +# ═══════════════════════════ parse: structural rules ══════════════════════ + + +@pytest.mark.parametrize("raw", ["id", "abc", "c/d", ""]) +def test_parse_requires_leading_slash(raw: str) -> None: + with pytest.raises(UnsafeCosmosPath, match="must start with '/'"): + CosmosPath.parse(raw) + + +def test_parse_bare_slash_is_empty() -> None: + with pytest.raises(UnsafeCosmosPath, match="empty or has a trailing"): + CosmosPath.parse("/") + + +@pytest.mark.parametrize("raw", ["/a/", "/id/", "/a/b/"]) +def test_parse_rejects_trailing_slash(raw: str) -> None: + with pytest.raises(UnsafeCosmosPath, match="empty or has a trailing"): + CosmosPath.parse(raw) + + +def test_parse_rejects_empty_middle_segment() -> None: + with pytest.raises(UnsafeCosmosPath, match="unsafe path segment"): + CosmosPath.parse("/a//b") + + +# ═══════════════════════════ parse: segment charset ═══════════════════════ + + +@pytest.mark.parametrize( + "raw,segments", + [ + ("/id", ("id",)), + ("/a/b/c", ("a", "b", "c")), + ("/_id", ("_id",)), + ("/A_b.c-d e", ("A_b.c-d e",)), # underscore, dot, hyphen, space allowed + ("/a1", ("a1",)), + ("/Z", ("Z",)), + ], +) +def test_parse_accepts_valid_paths(raw: str, segments: tuple) -> None: + assert CosmosPath.parse(raw).segments == segments + + +@pytest.mark.parametrize( + "raw", + [ + "/1abc", # leading digit + "/ abc", # leading space + "/.hidden", # leading dot + "/-x", # leading hyphen + "/a@b", # illegal symbol + "/caf\u00e9", # non-ASCII letter + "/a/1b", # bad segment in the middle + "/a#", # illegal symbol + ], +) +def test_parse_rejects_bad_segments(raw: str) -> None: + with pytest.raises(UnsafeCosmosPath, match="unsafe path segment"): + CosmosPath.parse(raw) + + +# ═══════════════════════════════ render ═══════════════════════════════════ + + +def test_render_default_alias() -> None: + assert CosmosPath.parse("/a/b").render() == 'c["a"]["b"]' + + +def test_render_custom_alias() -> None: + assert CosmosPath.parse("/a/b").render("x") == 'x["a"]["b"]' + + +def test_render_single_segment() -> None: + assert CosmosPath.parse("/id").render() == 'c["id"]' + + +def test_render_escapes_quote_and_backslash() -> None: + # Segments with quotes/backslashes can't come from parse, but render must + # still escape them safely when a CosmosPath is built directly. + assert CosmosPath(segments=('a"b',)).render() == 'c["a\\"b"]' + assert CosmosPath(segments=("a\\b",)).render() == 'c["a\\\\b"]' + assert CosmosPath(segments=('\\"',)).render() == 'c["\\\\\\""]' + + +# ═══════════════════════════════ __str__ / round-trip ═════════════════════ + + +def test_str_reconstructs_path() -> None: + assert str(CosmosPath.parse("/a/b/c")) == "/a/b/c" + assert str(CosmosPath.parse("/id")) == "/id" + + +def test_parse_str_round_trip() -> None: + p = CosmosPath.parse("/a/b.c/d e") + assert CosmosPath.parse(str(p)) == p + + +# ═══════════════════════════════ model semantics ══════════════════════════ + + +def test_frozen_cannot_mutate() -> None: + p = CosmosPath.parse("/a") + with pytest.raises(ValidationError): + p.segments = ("b",) # type: ignore[misc] + + +def test_equality_by_value() -> None: + assert CosmosPath.parse("/a/b") == CosmosPath(segments=("a", "b")) + assert CosmosPath.parse("/a") != CosmosPath.parse("/b") + + +def test_hashable_usable_in_set() -> None: + s = {CosmosPath.parse("/a"), CosmosPath(segments=("a",)), CosmosPath.parse("/b")} + assert len(s) == 2 # first two are equal -> collapse + + +# ═══════════════════════════════ coerce_path ══════════════════════════════ + + +def test_coerce_path_returns_same_cosmospath() -> None: + p = CosmosPath.parse("/a") + assert coerce_path(p) is p + + +def test_coerce_path_parses_string() -> None: + assert coerce_path("/a/b").segments == ("a", "b") + + +def test_coerce_path_invalid_raises() -> None: + with pytest.raises(UnsafeCosmosPath): + coerce_path("no-leading-slash") diff --git a/cosmos-retriever/tests/unit/test_planner.py b/cosmos-retriever/tests/unit/test_planner.py new file mode 100644 index 0000000..c4e2739 --- /dev/null +++ b/cosmos-retriever/tests/unit/test_planner.py @@ -0,0 +1,279 @@ +"""Exhaustive tests for `cosmos_retriever.retrieval.planner`. + +RetrievalPlanner chooses which search strategy to use for a request (based on the +schema, capabilities, and policy) but never runs it. Tests fake the schema and +capabilities to drive _vector_ok / _fts_ok through every branch, and override +those two helpers to check the plan_search mode/auto cascade in isolation. No +Cosmos / network. +""" +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +from cosmos_retriever.retrieval.capabilities import SupportLevel +from cosmos_retriever.retrieval.errors import UnsupportedRetrievalCapability +from cosmos_retriever.retrieval.models import GrepRequest, SearchRequest +from cosmos_retriever.retrieval.planner import RetrievalPlanner +from cosmos_retriever.retrieval.strategies import ( + BoundedScanStrategy, + ClientSideFusionStrategy, + FullTextGrepCandidateStrategy, + FullTextSearchStrategy, + NativeHybridStrategy, + VectorSearchStrategy, +) + +# ────────────────────────────── fakes ───────────────────────────────────── + + +class FakeField: + def __init__(self, path="P", dimensions=128): + self.path = path + self.dimensions = dimensions + + +class FakeCap: + def __init__(self, support=SupportLevel.INDEXED, dimensions=128): + self.support = support + self.dimensions = dimensions + + +class FakeSchema: + def __init__(self, vector_fields=("v",), text_paths=("t",), + field=None, resolve_vec_error=False, + resolve_text_error=False, text_paths_resolved=None): + self.vector_fields = list(vector_fields) + self.text_paths = list(text_paths) + self._field = field if field is not None else FakeField() + self._resolve_vec_error = resolve_vec_error + self._resolve_text_error = resolve_text_error + self._text_paths_resolved = text_paths_resolved or ["tp1"] + self.vector_calls: list = [] + self.text_calls: list = [] + + def resolve_vector_config(self, name): + self.vector_calls.append(name) + if self._resolve_vec_error: + raise ValueError("bad vector field") + return self._field + + def resolve_text_fields(self, names): + self.text_calls.append(names) + if self._resolve_text_error: + raise ValueError("bad text field") + return self._text_paths_resolved + + +class FakeCapabilities: + def __init__(self, vector_supported=True, full_text_supported=True, + native_hybrid_supported=False, cap=None, fts_paths=None): + self.vector_supported = vector_supported + self.full_text_supported = full_text_supported + self.native_hybrid_supported = native_hybrid_supported + self._cap = cap if cap is not None else FakeCap() + self._fts_paths = fts_paths # None -> every path has FTS; else a set + + def vector_capability_for(self, path): + return self._cap + + def has_full_text_path(self, path): + if self._fts_paths is None: + return True + return path in self._fts_paths + + +def _planner(schema=None, caps=None, bounded=False): + return RetrievalPlanner( + schema or FakeSchema(), + caps or FakeCapabilities(), + SimpleNamespace(allow_bounded_scan=bounded), + ) + + +def _req(**kw) -> SearchRequest: + base = dict(query="q") + base.update(kw) + return SearchRequest(**base) + + +# ═══════════════════════════════ _vector_ok ═══════════════════════════════ + + +def test_vector_ok_false_when_no_vector_fields() -> None: + p = _planner(FakeSchema(vector_fields=())) + assert p._vector_ok(_req()) is False + + +def test_vector_ok_false_when_capability_disabled() -> None: + p = _planner(caps=FakeCapabilities(vector_supported=False)) + assert p._vector_ok(_req()) is False + + +def test_vector_ok_false_when_resolve_raises() -> None: + p = _planner(FakeSchema(resolve_vec_error=True)) + assert p._vector_ok(_req()) is False + + +def test_vector_ok_false_when_cap_missing() -> None: + p = _planner() + p.capabilities.vector_capability_for = lambda path: None + assert p._vector_ok(_req()) is False + + +@pytest.mark.parametrize("support", [SupportLevel.UNSUPPORTED, SupportLevel.UNKNOWN]) +def test_vector_ok_false_for_weak_support(support) -> None: + p = _planner(caps=FakeCapabilities(cap=FakeCap(support=support))) + assert p._vector_ok(_req()) is False + + +def test_vector_ok_false_on_dimension_mismatch() -> None: + schema = FakeSchema(field=FakeField(dimensions=128)) + caps = FakeCapabilities(cap=FakeCap(dimensions=256)) + assert _planner(schema, caps)._vector_ok(_req()) is False + + +def test_vector_ok_true_when_all_aligned() -> None: + schema = FakeSchema(field=FakeField(path="P", dimensions=128)) + caps = FakeCapabilities(cap=FakeCap(support=SupportLevel.INDEXED, dimensions=128)) + assert _planner(schema, caps)._vector_ok(_req()) is True + + +def test_vector_ok_forwards_requested_field_name() -> None: + schema = FakeSchema() + _planner(schema)._vector_ok(_req(vector_field="myvec")) + assert schema.vector_calls == ["myvec"] + + +def test_vector_ok_none_request_uses_none_name() -> None: + schema = FakeSchema() + _planner(schema)._vector_ok(None) + assert schema.vector_calls == [None] + + +# ═══════════════════════════════ _fts_ok ══════════════════════════════════ + + +def test_fts_ok_false_when_capability_disabled() -> None: + p = _planner(caps=FakeCapabilities(full_text_supported=False)) + assert p._fts_ok(_req()) is False + + +def test_fts_ok_no_names_true_if_any_text_path_has_fts() -> None: + schema = FakeSchema(text_paths=("a", "b")) + caps = FakeCapabilities(fts_paths={"b"}) + assert _planner(schema, caps)._fts_ok(_req()) is True + + +def test_fts_ok_no_names_false_if_no_text_path_has_fts() -> None: + schema = FakeSchema(text_paths=("a", "b")) + caps = FakeCapabilities(fts_paths=set()) + assert _planner(schema, caps)._fts_ok(_req()) is False + + +def test_fts_ok_no_names_false_when_no_text_paths() -> None: + assert _planner(FakeSchema(text_paths=()))._fts_ok(_req()) is False + + +def test_fts_ok_named_fields_resolve_error_false() -> None: + p = _planner(FakeSchema(resolve_text_error=True)) + assert p._fts_ok(_req(text_fields=["x"])) is False + + +def test_fts_ok_named_fields_all_have_fts_true() -> None: + schema = FakeSchema(text_paths_resolved=["p1", "p2"]) + caps = FakeCapabilities(fts_paths={"p1", "p2"}) + assert _planner(schema, caps)._fts_ok(_req(text_fields=["a", "b"])) is True + + +def test_fts_ok_named_fields_missing_one_false() -> None: + schema = FakeSchema(text_paths_resolved=["p1", "p2"]) + caps = FakeCapabilities(fts_paths={"p1"}) # p2 lacks FTS + assert _planner(schema, caps)._fts_ok(_req(text_fields=["a", "b"])) is False + + +# ═══════════════════════════════ plan_search: explicit modes ══════════════ + + +def _p(vector_ok, fts_ok, native=False, bounded=False): + p = _planner(caps=FakeCapabilities(native_hybrid_supported=native), bounded=bounded) + p._vector_ok = lambda req=None: vector_ok # type: ignore[assignment] + p._fts_ok = lambda req=None: fts_ok # type: ignore[assignment] + return p + + +def test_plan_search_vector_mode_ok() -> None: + assert isinstance(_p(True, False).plan_search(_req(mode="vector")), VectorSearchStrategy) + + +def test_plan_search_vector_mode_unavailable_raises() -> None: + with pytest.raises(UnsupportedRetrievalCapability): + _p(False, True).plan_search(_req(mode="vector")) + + +def test_plan_search_text_mode_ok() -> None: + assert isinstance(_p(False, True).plan_search(_req(mode="text")), FullTextSearchStrategy) + + +def test_plan_search_text_mode_unavailable_raises() -> None: + with pytest.raises(UnsupportedRetrievalCapability): + _p(True, False).plan_search(_req(mode="text")) + + +def test_plan_search_hybrid_native() -> None: + assert isinstance( + _p(True, True, native=True).plan_search(_req(mode="hybrid")), NativeHybridStrategy) + + +def test_plan_search_hybrid_client_fusion_when_no_native() -> None: + assert isinstance( + _p(True, True, native=False).plan_search(_req(mode="hybrid")), ClientSideFusionStrategy) + + +@pytest.mark.parametrize("v,f", [(True, False), (False, True), (False, False)]) +def test_plan_search_hybrid_missing_side_raises(v: bool, f: bool) -> None: + with pytest.raises(UnsupportedRetrievalCapability): + _p(v, f).plan_search(_req(mode="hybrid")) + + +# ═══════════════════════════════ plan_search: auto cascade ════════════════ + + +def test_auto_prefers_native_hybrid() -> None: + assert isinstance(_p(True, True, native=True).plan_search(_req()), NativeHybridStrategy) + + +def test_auto_client_fusion_when_no_native() -> None: + assert isinstance(_p(True, True, native=False).plan_search(_req()), ClientSideFusionStrategy) + + +def test_auto_vector_only() -> None: + assert isinstance(_p(True, False).plan_search(_req()), VectorSearchStrategy) + + +def test_auto_full_text_only() -> None: + assert isinstance(_p(False, True).plan_search(_req()), FullTextSearchStrategy) + + +def test_auto_bounded_scan_when_nothing_and_allowed() -> None: + assert isinstance(_p(False, False, bounded=True).plan_search(_req()), BoundedScanStrategy) + + +def test_auto_raises_when_nothing_and_scan_disallowed() -> None: + with pytest.raises(UnsupportedRetrievalCapability): + _p(False, False, bounded=False).plan_search(_req()) + + +# ═══════════════════════════════ plan_grep ════════════════════════════════ + + +def test_plan_grep_returns_full_text_candidate() -> None: + p = _planner(caps=FakeCapabilities(full_text_supported=True)) + assert isinstance(p.plan_grep(GrepRequest(pattern="x")), FullTextGrepCandidateStrategy) + + +def test_plan_grep_raises_without_full_text() -> None: + p = _planner(caps=FakeCapabilities(full_text_supported=False)) + with pytest.raises(UnsupportedRetrievalCapability): + p.plan_grep(GrepRequest(pattern="x")) diff --git a/cosmos-retriever/tests/unit/test_rerank.py b/cosmos-retriever/tests/unit/test_rerank.py new file mode 100644 index 0000000..794e7bf --- /dev/null +++ b/cosmos-retriever/tests/unit/test_rerank.py @@ -0,0 +1,746 @@ +"""Exhaustive tests for the reranker module (`cosmos_retriever.rerank`). + +Covers, with fakes and no network access: + + 1. RerankResult — dataclass fields / equality / mutability + 2. Reranker.__init__ — token_counter / max_tokens invariants + 3. Reranker._truncate_results — token annotation + budget truncation boundaries + 4. Reranker.__call__ — template: _rerank -> slow-warn -> truncate + 5. BasetenReranker — classify() plumbing + yes/no label -> P("yes") score + 6. VLLMQwen3Reranker — /score plumbing, batching, retry/backoff, and + degenerate/malformed model outputs + 7. ContextualReranker — /rerank plumbing, payload/headers, error propagation + 8. Module surface — VLLMReranker alias + a latent-config-bug guard + +All HTTP is intercepted by patching `rerank.requests.post`; the Baseten client +and `get_config` are replaced with fakes. `time.sleep` is patched so retry +backoff does not actually wait. +""" +from __future__ import annotations + +import types + +import pytest +import requests + +from cosmos_retriever import rerank +from cosmos_retriever.rerank import ( + BasetenReranker, + ContextualReranker, + Reranker, + RerankResult, + VLLMQwen3Reranker, + VLLMReranker, +) + +# ───────────────────────── shared fakes / helpers ───────────────────────── + +def x_counter(s: str) -> int: + """Deterministic token counter: each 'x' character counts as one token.""" + return s.count("x") if isinstance(s, str) else 0 + + +class _RecordLogger: + """Captures structlog-style calls so branch coverage can assert on them.""" + + def __init__(self) -> None: + self.warnings: list[tuple] = [] + self.errors: list[tuple] = [] + self.infos: list[tuple] = [] + + def warning(self, *a, **k) -> None: + self.warnings.append((a, k)) + + def error(self, *a, **k) -> None: + self.errors.append((a, k)) + + def info(self, *a, **k) -> None: + self.infos.append((a, k)) + + +class _StubReranker(Reranker): + """Concrete Reranker whose `_rerank` returns a fixed, caller-supplied list. + + Used to exercise the abstract base's template (`__call__`) and truncation in + isolation from any scoring backend. Records what `_rerank` was invoked with. + """ + + def __init__(self, results: list[RerankResult], **kw) -> None: + super().__init__(**kw) + self._results = results + self.seen: list[tuple] = [] + + def _rerank(self, query, documents, instruction=None): + self.seen.append((query, list(documents), instruction)) + # Fresh copies so truncation's in-place `.tokens` writes don't leak back. + return [ + RerankResult(document=r.document, score=r.score, original_index=r.original_index) + for r in self._results + ] + + +def _mk(doc: str, score: float, idx: int) -> RerankResult: + return RerankResult(document=doc, score=score, original_index=idx) + + +# ═════════════════════════ 1. RerankResult ═════════════════════════ + +def test_rerankresult_defaults_tokens_none() -> None: + r = RerankResult(document="d", score=0.5, original_index=3) + assert r.document == "d" and r.score == 0.5 and r.original_index == 3 + assert r.tokens is None + + +def test_rerankresult_equality_and_token_mutation() -> None: + a = RerankResult(document="d", score=0.5, original_index=0) + b = RerankResult(document="d", score=0.5, original_index=0) + assert a == b + a.tokens = 7 # tokens is set later by _truncate_results + assert a != b and a.tokens == 7 + + +# ═════════════════════════ 2. Reranker.__init__ ═════════════════════════ + +def test_init_max_tokens_without_counter_raises() -> None: + with pytest.raises(ValueError, match="token_counter is required"): + _StubReranker([], max_tokens=100) + + +def test_init_defaults_are_none() -> None: + r = _StubReranker([]) + assert r.token_counter is None and r.max_tokens is None + + +def test_init_counter_without_max_tokens_ok() -> None: + r = _StubReranker([], token_counter=x_counter) + assert r.token_counter is x_counter and r.max_tokens is None + + +def test_init_both_set_ok() -> None: + r = _StubReranker([], token_counter=x_counter, max_tokens=50) + assert r.token_counter is x_counter and r.max_tokens == 50 + + +# ═════════════════════════ 3. _truncate_results ═════════════════════════ + +def test_truncate_sets_tokens_on_every_result_even_without_budget() -> None: + results = [_mk("x", 1, 0), _mk("xx", 1, 1)] + r = _StubReranker([], token_counter=x_counter) # no max_tokens -> no truncation + out = r._truncate_results(results) + assert out is results # unchanged list + assert [res.tokens for res in results] == [1, 2] # tokens annotated regardless + + +def test_truncate_no_counter_leaves_tokens_none_and_no_truncation() -> None: + results = [_mk("x", 1, 0), _mk("xxxx", 1, 1)] + r = _StubReranker([]) # no counter + out = r._truncate_results(results, max_tokens=1) + assert out == results + assert all(res.tokens is None for res in results) + + +def test_truncate_boundary_equal_is_kept() -> None: + # tokens: 1, 2, 3 ; budget 3 -> keep [1] then [1+2=3] (== not > 3, kept), drop 3rd + results = [_mk("x", 3, 0), _mk("xx", 2, 1), _mk("xxx", 1, 2)] + r = _StubReranker([], token_counter=x_counter, max_tokens=3) + out = r._truncate_results(results) + assert [res.document for res in out] == ["x", "xx"] + + +def test_truncate_first_doc_over_budget_yields_empty() -> None: + results = [_mk("xxxx", 1, 0)] # 4 tokens + r = _StubReranker([], token_counter=x_counter, max_tokens=3) + assert r._truncate_results(results) == [] + + +def test_truncate_zero_budget_keeps_nothing() -> None: + results = [_mk("x", 1, 0)] + r = _StubReranker([], token_counter=x_counter, max_tokens=0) + assert r._truncate_results(results) == [] + + +def test_truncate_all_fit_keeps_all() -> None: + results = [_mk("x", 1, 0), _mk("x", 1, 1)] + r = _StubReranker([], token_counter=x_counter, max_tokens=100) + out = r._truncate_results(results) + assert len(out) == 2 + + +def test_truncate_call_arg_overrides_instance_max() -> None: + results = [_mk("x", 1, 0), _mk("x", 1, 1), _mk("x", 1, 2)] + r = _StubReranker([], token_counter=x_counter, max_tokens=100) + # call-level budget of 2 wins over the instance's 100 + out = r._truncate_results(results, max_tokens=2) + assert len(out) == 2 + + +def test_truncate_instance_max_used_when_call_arg_none() -> None: + results = [_mk("x", 1, 0), _mk("x", 1, 1), _mk("x", 1, 2)] + r = _StubReranker([], token_counter=x_counter, max_tokens=1) + out = r._truncate_results(results, max_tokens=None) + assert len(out) == 1 + + +def test_truncate_dropped_results_still_annotated(monkeypatch) -> None: + rec = _RecordLogger() + monkeypatch.setattr(rerank, "logger", rec) + results = [_mk("x", 2, 0), _mk("xx", 1, 1), _mk("xxx", 1, 2)] + r = _StubReranker([], token_counter=x_counter, max_tokens=1) + r._truncate_results(results) + # every original result — including the dropped ones — got tokens set + assert [res.tokens for res in results] == [1, 2, 3] + # and a truncation log was emitted + assert rec.infos and rec.infos[0][1]["dropped"] == 2 + + +# ═════════════════════════ 4. Reranker.__call__ template ═════════════════════════ + +def test_call_threads_instruction_and_truncates() -> None: + base = [_mk("x", 2, 0), _mk("x", 1, 1)] + r = _StubReranker(base, token_counter=x_counter, max_tokens=1) + out = r("the query", ["x", "x"], instruction="inst") + # _rerank saw the instruction and documents + assert r.seen == [("the query", ["x", "x"], "inst")] + # truncation applied (each doc is 1 token, budget 1 -> keep 1) + assert len(out) == 1 + + +def test_call_passes_call_level_max_tokens() -> None: + base = [_mk("x", 2, 0), _mk("x", 1, 1), _mk("x", 1, 2)] + r = _StubReranker(base, token_counter=x_counter) # no instance max + out = r("q", ["x", "x", "x"], max_tokens=2) + assert len(out) == 2 + + +def test_call_slow_warning_branch(monkeypatch) -> None: + rec = _RecordLogger() + monkeypatch.setattr(rerank, "logger", rec) + # perf_counter is called twice in __call__: start, then end. 0 -> 2.0s = 2000ms > 1500ms. + seq = iter([0.0, 2.0]) + monkeypatch.setattr(rerank.time, "perf_counter", lambda: next(seq)) + r = _StubReranker([_mk("d", 1.0, 0)]) + out = r("q", ["d"]) + assert len(out) == 1 + assert rec.warnings and "slow" in rec.warnings[0][0][0].lower() + + +def test_call_no_warning_when_fast(monkeypatch) -> None: + rec = _RecordLogger() + monkeypatch.setattr(rerank, "logger", rec) + seq = iter([0.0, 0.1]) # 100ms + monkeypatch.setattr(rerank.time, "perf_counter", lambda: next(seq)) + r = _StubReranker([_mk("d", 1.0, 0)]) + r("q", ["d"]) + assert rec.warnings == [] + + +# ═════════════════════════ 5. BasetenReranker ═════════════════════════ + +def _group(*pairs): + """Build one document's classify group: [(label, score), ...].""" + return [types.SimpleNamespace(label=lbl, score=score) for lbl, score in pairs] + + +class FakeBasetenClient: + def __init__(self, data) -> None: + self._data = data + self.calls: list[dict] = [] + + def classify(self, **kwargs): + self.calls.append(kwargs) + return types.SimpleNamespace(data=self._data) + + +def test_baseten_class_constants() -> None: + assert "yes" in BasetenReranker.PREFIX and "no" in BasetenReranker.PREFIX + assert BasetenReranker.SUFFIX.startswith("<|im_end|>") + assert "assistant" in BasetenReranker.SUFFIX and "" in BasetenReranker.SUFFIX + assert BasetenReranker.DEFAULT_INSTRUCTION + + +def test_baseten_defaults() -> None: + r = BasetenReranker(client=FakeBasetenClient([])) + assert r.batch_size == 16 and r.max_concurrent_requests == 256 and r.timeout_s == 360 + + +def test_baseten_uses_config_client_when_none(monkeypatch) -> None: + sentinel = FakeBasetenClient([]) + fake_cfg = types.SimpleNamespace(get_baseten_client=lambda: sentinel) + monkeypatch.setattr(rerank, "get_config", lambda: fake_cfg) + r = BasetenReranker(client=None) + assert r.client is sentinel + + +def test_baseten_format_input_default_instruction_exact() -> None: + r = BasetenReranker(client=FakeBasetenClient([])) + out = r._format_input(None, "Q?", "DOC") + expected = ( + f"{BasetenReranker.PREFIX}: {BasetenReranker.DEFAULT_INSTRUCTION}\n" + f": Q?\n: DOC{BasetenReranker.SUFFIX}" + ) + assert out == expected + + +def test_baseten_format_input_custom_instruction() -> None: + r = BasetenReranker(client=FakeBasetenClient([])) + out = r._format_input("CUSTOM", "Q", "D") + assert ": CUSTOM" in out and ": Q" in out and ": D" in out + + +def test_baseten_empty_documents_no_client_call() -> None: + client = FakeBasetenClient([]) + r = BasetenReranker(client=client) + assert r._rerank("q", []) == [] + assert client.calls == [] + + +def test_baseten_classify_called_with_expected_kwargs() -> None: + client = FakeBasetenClient([_group(("yes", 0.9), ("no", 0.1))]) + r = BasetenReranker(client=client, batch_size=8, max_concurrent_requests=4, timeout_s=12) + r._rerank("q", ["d0"], instruction="inst") + (kwargs,) = client.calls + assert kwargs["truncate"] is True + assert kwargs["batch_size"] == 8 + assert kwargs["max_concurrent_requests"] == 4 + assert kwargs["timeout_s"] == 12 + assert len(kwargs["inputs"]) == 1 + assert ": inst" in kwargs["inputs"][0] + + +def test_baseten_takes_yes_probability() -> None: + client = FakeBasetenClient([_group(("no", 0.2), ("yes", 0.8))]) + r = BasetenReranker(client=client) + out = r._rerank("q", ["d0"]) + assert out[0].score == 0.8 # picks the "yes" entry regardless of position + + +def test_baseten_missing_yes_scores_zero() -> None: + client = FakeBasetenClient([_group(("no", 0.7))]) + r = BasetenReranker(client=client) + out = r._rerank("q", ["d0"]) + assert out[0].score == 0.0 + + +def test_baseten_unexpected_labels_score_zero() -> None: + # Classifier emits labels the code never expects (never "yes") -> 0.0, no crash. + client = FakeBasetenClient([_group(("maybe", 0.9), ("garbage", 0.8))]) + r = BasetenReranker(client=client) + out = r._rerank("q", ["d0"]) + assert out[0].score == 0.0 + + +def test_baseten_empty_group_scores_zero() -> None: + client = FakeBasetenClient([[]]) + r = BasetenReranker(client=client) + out = r._rerank("q", ["d0"]) + assert out[0].score == 0.0 + + +def test_baseten_breaks_on_first_yes() -> None: + # Two "yes" entries; the loop breaks on the first one it sees. + client = FakeBasetenClient([_group(("yes", 0.55), ("yes", 0.99))]) + r = BasetenReranker(client=client) + out = r._rerank("q", ["d0"]) + assert out[0].score == 0.55 + + +def test_baseten_sorts_descending_and_preserves_original_index() -> None: + client = FakeBasetenClient( + [ + _group(("yes", 0.1)), # d0 + _group(("yes", 0.9)), # d1 + _group(("yes", 0.5)), # d2 + ] + ) + r = BasetenReranker(client=client) + out = r._rerank("q", ["d0", "d1", "d2"]) + assert [x.document for x in out] == ["d1", "d2", "d0"] + assert [x.original_index for x in out] == [1, 2, 0] + assert [x.score for x in out] == [0.9, 0.5, 0.1] + + +def test_baseten_tie_scores_stable_order() -> None: + client = FakeBasetenClient([_group(("yes", 0.5)), _group(("yes", 0.5))]) + r = BasetenReranker(client=client) + out = r._rerank("q", ["d0", "d1"]) + # stable sort keeps original relative order on ties + assert [x.original_index for x in out] == [0, 1] + + +def test_baseten_call_end_to_end_with_truncation() -> None: + client = FakeBasetenClient([_group(("yes", 0.2)), _group(("yes", 0.9))]) + r = BasetenReranker(client=client, token_counter=x_counter, max_tokens=1) + out = r("q", ["x", "x"]) # both 1 token; budget 1 -> keep top-1 after sort + assert len(out) == 1 and out[0].score == 0.9 and out[0].tokens == 1 + + +# ═════════════════════════ 6. VLLMQwen3Reranker ═════════════════════════ + +class FakeHTTPResponse: + def __init__(self, payload=None, raise_exc: Exception | None = None) -> None: + self._payload = payload + self._raise = raise_exc + + def raise_for_status(self) -> None: + if self._raise is not None: + raise self._raise + + def json(self): + return self._payload + + +def test_vllm_alias_is_qwen3() -> None: + assert VLLMReranker is VLLMQwen3Reranker + + +def test_vllm_class_constants_and_defaults() -> None: + r = VLLMQwen3Reranker(base_url="http://host:1/") + assert r.model == "Qwen/Qwen3-Reranker-8B" + assert r.batch_size == 32 and r.timeout_s == 360 + assert "yes" in r.PREFIX and r.SUFFIX.startswith("<|im_end|>") + + +def test_vllm_base_url_trailing_slash_stripped() -> None: + r = VLLMQwen3Reranker(base_url="http://host:8011///") + assert r.base_url == "http://host:8011" + + +def test_vllm_base_url_env_fallback(monkeypatch) -> None: + monkeypatch.setenv("VLLM_RERANKER_URL", "http://env-host:9/") + r = VLLMQwen3Reranker(base_url=None) + assert r.base_url == "http://env-host:9" + + +def test_vllm_base_url_hard_default(monkeypatch) -> None: + monkeypatch.delenv("VLLM_RERANKER_URL", raising=False) + r = VLLMQwen3Reranker(base_url=None) + assert r.base_url == "http://127.0.0.1:8011" + + +def test_vllm_empty_documents_no_http(monkeypatch) -> None: + called = [] + monkeypatch.setattr(rerank.requests, "post", lambda *a, **k: called.append(1)) + r = VLLMQwen3Reranker(base_url="http://h:1") + assert r._rerank("q", []) == [] + assert called == [] + + +def test_vllm_payload_and_endpoint(monkeypatch) -> None: + captured = {} + + def fake_post(url, json=None, timeout=None): # noqa: A002 + captured["url"] = url + captured["json"] = json + captured["timeout"] = timeout + return FakeHTTPResponse({"data": [{"score": 0.5}]}) + + monkeypatch.setattr(rerank.requests, "post", fake_post) + r = VLLMQwen3Reranker(base_url="http://h:8011", timeout_s=42) + r._rerank("Q", ["DOC"], instruction="INST") + assert captured["url"] == "http://h:8011/score" + assert captured["timeout"] == 42 + body = captured["json"] + assert body["model"] == "Qwen/Qwen3-Reranker-8B" + assert body["truncate_prompt_tokens"] == -1 + assert ": INST" in body["text_1"] and ": Q" in body["text_1"] + assert body["text_2"] == [f": DOC{VLLMQwen3Reranker.SUFFIX}"] + + +def test_vllm_default_instruction_used(monkeypatch) -> None: + captured = {} + + def fake_post(url, json=None, timeout=None): # noqa: A002 + captured["json"] = json + return FakeHTTPResponse({"data": [{"score": 0.1}]}) + + monkeypatch.setattr(rerank.requests, "post", fake_post) + r = VLLMQwen3Reranker(base_url="http://h:1") + r._rerank("Q", ["D"]) + assert VLLMQwen3Reranker.DEFAULT_INSTRUCTION in captured["json"]["text_1"] + + +def test_vllm_batches_split_into_multiple_posts(monkeypatch) -> None: + posts: list[list[str]] = [] + + def fake_post(url, json=None, timeout=None): # noqa: A002 + docs = json["text_2"] + posts.append(docs) + return FakeHTTPResponse({"data": [{"score": 0.0} for _ in docs]}) + + monkeypatch.setattr(rerank.requests, "post", fake_post) + r = VLLMQwen3Reranker(base_url="http://h:1", batch_size=2) + r._rerank("q", ["a", "b", "c"]) # 3 docs, batch 2 -> 2 posts (2 + 1) + assert len(posts) == 2 + assert len(posts[0]) == 2 and len(posts[1]) == 1 + + +def test_vllm_scores_accumulate_in_order_and_coerce_float(monkeypatch) -> None: + def fake_post(url, json=None, timeout=None): # noqa: A002 + docs = json["text_2"] + # return integer scores to verify float() coercion + return FakeHTTPResponse({"data": [{"score": i} for i in range(len(docs))]}) + + monkeypatch.setattr(rerank.requests, "post", fake_post) + r = VLLMQwen3Reranker(base_url="http://h:1", batch_size=10) + out = r._rerank("q", ["a", "b", "c"]) + assert all(isinstance(x.score, float) for x in out) + # scores were 0,1,2 -> sorted desc: c(2), b(1), a(0) + assert [x.document for x in out] == ["c", "b", "a"] + assert [x.original_index for x in out] == [2, 1, 0] + + +def test_vllm_all_equal_scores_preserve_original_order(monkeypatch) -> None: + # Degenerate scoring: every document gets the same score. The sort is stable, + # so the reranker keeps the original order instead of shuffling arbitrarily. + def fake_post(url, json=None, timeout=None): # noqa: A002 + docs = json["text_2"] + return FakeHTTPResponse({"data": [{"score": 0.5} for _ in docs]}) + + monkeypatch.setattr(rerank.requests, "post", fake_post) + r = VLLMQwen3Reranker(base_url="http://h:1", batch_size=10) + out = r._rerank("q", ["a", "b", "c", "d"]) + assert [x.original_index for x in out] == [0, 1, 2, 3] + assert [x.document for x in out] == ["a", "b", "c", "d"] + + +def test_vllm_non_numeric_score_raises(monkeypatch) -> None: + # Model returns junk instead of a number -> float() raises. It is not a + # transient RequestException, so it surfaces immediately (no silent 0 score). + def fake_post(url, json=None, timeout=None): # noqa: A002 + return FakeHTTPResponse({"data": [{"score": "not-a-number"}]}) + + monkeypatch.setattr(rerank.requests, "post", fake_post) + r = VLLMQwen3Reranker(base_url="http://h:1") + with pytest.raises(ValueError): + r._rerank("q", ["d"]) + + +def test_vllm_missing_data_key_raises(monkeypatch) -> None: + # A malformed response body with no "data" field surfaces as a KeyError. + def fake_post(url, json=None, timeout=None): # noqa: A002 + return FakeHTTPResponse({"unexpected": []}) + + monkeypatch.setattr(rerank.requests, "post", fake_post) + r = VLLMQwen3Reranker(base_url="http://h:1") + with pytest.raises(KeyError): + r._rerank("q", ["d"]) + + +def test_vllm_fewer_scores_than_documents_drops_unscored(monkeypatch) -> None: + # Model returns fewer scores than documents -> the unscored tail is dropped + # (zip truncation) rather than crashing or inventing scores. + def fake_post(url, json=None, timeout=None): # noqa: A002 + return FakeHTTPResponse({"data": [{"score": 0.9}]}) # one score, three docs + + monkeypatch.setattr(rerank.requests, "post", fake_post) + r = VLLMQwen3Reranker(base_url="http://h:1", batch_size=10) + out = r._rerank("q", ["a", "b", "c"]) + assert [x.document for x in out] == ["a"] + + +def test_vllm_more_scores_than_documents_ignores_extra(monkeypatch) -> None: + def fake_post(url, json=None, timeout=None): # noqa: A002 + return FakeHTTPResponse({"data": [{"score": s} for s in (0.1, 0.2, 0.3)]}) + + monkeypatch.setattr(rerank.requests, "post", fake_post) + r = VLLMQwen3Reranker(base_url="http://h:1", batch_size=10) + out = r._rerank("q", ["only-one"]) + assert len(out) == 1 and out[0].document == "only-one" + + +def test_vllm_retries_then_succeeds(monkeypatch) -> None: + rec = _RecordLogger() + monkeypatch.setattr(rerank, "logger", rec) + sleeps: list[float] = [] + monkeypatch.setattr(rerank.time, "sleep", lambda s: sleeps.append(s)) + + calls = {"n": 0} + + def flaky_post(url, json=None, timeout=None): # noqa: A002 + calls["n"] += 1 + if calls["n"] == 1: + raise requests.exceptions.ConnectionError("boom") + return FakeHTTPResponse({"data": [{"score": 0.7}]}) + + monkeypatch.setattr(rerank.requests, "post", flaky_post) + r = VLLMQwen3Reranker(base_url="http://h:1") + out = r._rerank("q", ["d"]) + assert out[0].score == 0.7 + assert calls["n"] == 2 # one failure, one success + assert sleeps == [1] # backoff 2**0 before the retry + assert rec.warnings # retry warning logged + + +def test_vllm_all_retries_fail_raises(monkeypatch) -> None: + rec = _RecordLogger() + monkeypatch.setattr(rerank, "logger", rec) + monkeypatch.setattr(rerank.time, "sleep", lambda s: None) + + def always_fail(url, json=None, timeout=None): # noqa: A002 + raise requests.exceptions.Timeout("nope") + + monkeypatch.setattr(rerank.requests, "post", always_fail) + r = VLLMQwen3Reranker(base_url="http://h:1") + with pytest.raises(requests.exceptions.Timeout): + r._rerank("q", ["d"]) + assert rec.errors # failure logged after exhausting retries + + +def test_vllm_call_end_to_end_with_truncation(monkeypatch) -> None: + def fake_post(url, json=None, timeout=None): # noqa: A002 + docs = json["text_2"] + return FakeHTTPResponse({"data": [{"score": 0.9} for _ in docs]}) + + monkeypatch.setattr(rerank.requests, "post", fake_post) + r = VLLMQwen3Reranker(base_url="http://h:1", token_counter=x_counter, max_tokens=1) + out = r("q", ["x", "x"]) + assert len(out) == 1 and out[0].tokens == 1 + + +# ═════════════════════════ 7. ContextualReranker ═════════════════════════ + +def test_contextual_class_constants() -> None: + assert ContextualReranker.API_URL.endswith("/v1/rerank") + assert ContextualReranker.DEFAULT_MODEL + assert ContextualReranker.DEFAULT_INSTRUCTION + + +def test_contextual_defaults() -> None: + r = ContextualReranker(api_key="k") + assert r.model == ContextualReranker.DEFAULT_MODEL + assert r.top_n is None and r.timeout_s == 60 + + +def test_contextual_custom_model_and_top_n() -> None: + r = ContextualReranker(api_key="k", model="m", top_n=3) + assert r.model == "m" and r.top_n == 3 + + +def test_contextual_api_key_from_config(monkeypatch) -> None: + fake_cfg = types.SimpleNamespace( + contextual_api_key=types.SimpleNamespace(get_secret_value=lambda: "cfg-key") + ) + monkeypatch.setattr(rerank, "get_config", lambda: fake_cfg) + r = ContextualReranker(api_key=None) + assert r.api_key == "cfg-key" + + +def test_contextual_empty_documents_no_http(monkeypatch) -> None: + called = [] + monkeypatch.setattr(rerank.requests, "post", lambda *a, **k: called.append(1)) + r = ContextualReranker(api_key="k") + assert r._rerank("q", []) == [] + assert called == [] + + +def test_contextual_payload_headers_and_endpoint(monkeypatch) -> None: + captured = {} + + def fake_post(url, json=None, headers=None, timeout=None): # noqa: A002 + captured.update(url=url, json=json, headers=headers, timeout=timeout) + return FakeHTTPResponse({"results": []}) + + monkeypatch.setattr(rerank.requests, "post", fake_post) + r = ContextualReranker(api_key="secret", model="m", top_n=2, timeout_s=15) + r._rerank("Q", ["a", "b"], instruction="INST") + assert captured["url"] == ContextualReranker.API_URL + assert captured["timeout"] == 15 + assert captured["headers"]["Authorization"] == "Bearer secret" + assert captured["headers"]["Content-Type"] == "application/json" + body = captured["json"] + assert body["query"] == "Q" + assert body["documents"] == ["a", "b"] + assert body["model"] == "m" + assert body["top_n"] == 2 + assert body["instruction"] == "INST" + + +def test_contextual_default_instruction_when_none(monkeypatch) -> None: + captured = {} + + def fake_post(url, json=None, headers=None, timeout=None): # noqa: A002 + captured["json"] = json + return FakeHTTPResponse({"results": []}) + + monkeypatch.setattr(rerank.requests, "post", fake_post) + r = ContextualReranker(api_key="k") + r._rerank("q", ["a"]) + assert captured["json"]["instruction"] == ContextualReranker.DEFAULT_INSTRUCTION + + +def test_contextual_top_n_omitted_when_none(monkeypatch) -> None: + captured = {} + + def fake_post(url, json=None, headers=None, timeout=None): # noqa: A002 + captured["json"] = json + return FakeHTTPResponse({"results": []}) + + monkeypatch.setattr(rerank.requests, "post", fake_post) + r = ContextualReranker(api_key="k") # top_n None + r._rerank("q", ["a"]) + assert "top_n" not in captured["json"] + + +def test_contextual_parses_and_sorts_results(monkeypatch) -> None: + payload = { + "results": [ + {"index": 0, "relevance_score": 0.2}, + {"index": 2, "relevance_score": 0.9}, + {"index": 1, "relevance_score": 0.5}, + ] + } + monkeypatch.setattr( + rerank.requests, "post", + lambda *a, **k: FakeHTTPResponse(payload), + ) + r = ContextualReranker(api_key="k") + out = r._rerank("q", ["d0", "d1", "d2"]) + assert [x.document for x in out] == ["d2", "d1", "d0"] + assert [x.original_index for x in out] == [2, 1, 0] + assert [x.score for x in out] == [0.9, 0.5, 0.2] + + +def test_contextual_missing_results_key_returns_empty(monkeypatch) -> None: + monkeypatch.setattr( + rerank.requests, "post", + lambda *a, **k: FakeHTTPResponse({}), # no "results" + ) + r = ContextualReranker(api_key="k") + assert r._rerank("q", ["a"]) == [] + + +def test_contextual_request_exception_propagates(monkeypatch) -> None: + rec = _RecordLogger() + monkeypatch.setattr(rerank, "logger", rec) + + def boom(*a, **k): + raise requests.exceptions.ConnectionError("down") + + monkeypatch.setattr(rerank.requests, "post", boom) + r = ContextualReranker(api_key="k") + with pytest.raises(requests.exceptions.ConnectionError): + r._rerank("q", ["a"]) + assert rec.errors # contextual_rerank_failed logged + + +def test_contextual_raise_for_status_error_propagates(monkeypatch) -> None: + err = requests.exceptions.HTTPError("500") + monkeypatch.setattr( + rerank.requests, "post", + lambda *a, **k: FakeHTTPResponse({"results": []}, raise_exc=err), + ) + r = ContextualReranker(api_key="k") + with pytest.raises(requests.exceptions.HTTPError): + r._rerank("q", ["a"]) + + +# ═════════════════════════ 8. Module surface / latent-bug guard ═════════════════════════ + +def test_contextual_api_key_field_absent_from_settings() -> None: + """Guards a latent bug: ContextualReranker(api_key=None) reads + `config.contextual_api_key`, but RetrieverSettings does not define that field, + so the no-arg path raises AttributeError at runtime. If this assertion starts + failing, the field was added and the config fallback now works.""" + from cosmos_retriever.config import RetrieverSettings + + assert "contextual_api_key" not in RetrieverSettings.model_fields diff --git a/cosmos-retriever/tests/unit/test_retriever.py b/cosmos-retriever/tests/unit/test_retriever.py new file mode 100644 index 0000000..eca504d --- /dev/null +++ b/cosmos-retriever/tests/unit/test_retriever.py @@ -0,0 +1,387 @@ +"""Exhaustive tests for `CorpusRetriever` (`cosmos_retriever.retrieval.retriever`). + +`CorpusRetriever` is a thin orchestrator: it wires collaborators in ``__init__`` +and, per call, resolves schema fields, asks the planner for a strategy, optionally +embeds the query, and delegates execution. These tests isolate it by patching the +five collaborator constructors in the module with recorders returning fakes / +sentinels, so every branch of ``__init__`` / ``search`` / ``grep_candidates`` / +``read_document`` is exercised without Cosmos, embeddings, or real strategies. + +Fakes: FakeSchema, FakePlanner, FakeSearchStrategy, FakeGrepStrategy, +FakeEmbedder, FakeResolver. Real Pydantic models (SearchRequest, GrepRequest, +ReadDocumentRequest, RetrievedItem, NormalizedDocument, PartitionQueryPolicy) +are used directly. +""" +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +from cosmos_retriever.retrieval import retriever as retr_mod +from cosmos_retriever.retrieval.errors import EmbeddingProfileMismatch +from cosmos_retriever.retrieval.models import ( + GrepRequest, + NormalizedDocument, + PartitionQueryPolicy, + ReadDocumentRequest, + RetrievedItem, + SearchRequest, +) +from cosmos_retriever.retrieval.retriever import CorpusRetriever + +# ────────────────────────────── fakes ───────────────────────────────────── + + +class _CallRecorder: + """Records positional/keyword args and returns a fixed value.""" + + def __init__(self, return_value: object) -> None: + self.return_value = return_value + self.calls: list[tuple[tuple, dict]] = [] + + def __call__(self, *args, **kwargs): + self.calls.append((args, kwargs)) + return self.return_value + + +class FakeSchema: + def __init__(self) -> None: + self.vector_calls: list = [] + self.text_calls: list = [] + + def resolve_vector_config(self, name): + self.vector_calls.append(name) + return object() + + def resolve_text_fields(self, names): + self.text_calls.append(names) + return [] + + +class FakeSearchStrategy: + def __init__(self, requires_embedding: bool, result: list[RetrievedItem]) -> None: + self.requires_embedding = requires_embedding + self._result = result + self.execute_calls: list[tuple] = [] + + def execute(self, req, ctx): + self.execute_calls.append((req, ctx)) + return self._result + + +class FakeGrepStrategy: + def __init__(self, result: list[RetrievedItem]) -> None: + self._result = result + self.candidate_calls: list[tuple] = [] + + def candidates(self, req, ctx): + self.candidate_calls.append((req, ctx)) + return self._result + + +class FakePlanner: + def __init__(self, search_strategy=None, grep_strategy=None) -> None: + self._search = search_strategy + self._grep = grep_strategy + self.search_reqs: list = [] + self.grep_reqs: list = [] + + def plan_search(self, req): + self.search_reqs.append(req) + return self._search + + def plan_grep(self, req): + self.grep_reqs.append(req) + return self._grep + + +class FakeEmbedder: + def __init__(self, vector: list[float]) -> None: + self.vector = vector + self.calls: list[str] = [] + + def embed(self, text: str) -> list[float]: + self.calls.append(text) + return self.vector + + +class FakeResolver: + def __init__(self, document: NormalizedDocument) -> None: + self.document = document + self.calls: list = [] + + def resolve(self, req): + self.calls.append(req) + return self.document + + +def _item(item_id: str, text: str = "") -> RetrievedItem: + return RetrievedItem(item_id=item_id, text=text) + + +def _build( + monkeypatch, + *, + search_strategy: FakeSearchStrategy | None = None, + grep_strategy: FakeGrepStrategy | None = None, + embedder: FakeEmbedder | None = None, + policy: PartitionQueryPolicy | None = None, + document: NormalizedDocument | None = None, +) -> SimpleNamespace: + """Patch collaborator constructors and build an isolated CorpusRetriever.""" + schema = FakeSchema() + planner = FakePlanner(search_strategy, grep_strategy) + resolver = FakeResolver(document or NormalizedDocument(chunk_texts=[])) + + compiler_sentinel = object() + executor_sentinel = object() + ctx_sentinel = object() + + comp_rec = _CallRecorder(compiler_sentinel) + exec_rec = _CallRecorder(executor_sentinel) + planner_rec = _CallRecorder(planner) + ctx_rec = _CallRecorder(ctx_sentinel) + resolver_rec = _CallRecorder(resolver) + + monkeypatch.setattr(retr_mod, "CosmosQueryCompiler", comp_rec) + monkeypatch.setattr(retr_mod, "CosmosExecutor", exec_rec) + monkeypatch.setattr(retr_mod, "RetrievalPlanner", planner_rec) + monkeypatch.setattr(retr_mod, "RetrievalContext", ctx_rec) + monkeypatch.setattr(retr_mod, "build_document_resolver", resolver_rec) + + container = object() + capabilities = object() + retriever = CorpusRetriever( + container=container, + schema=schema, # type: ignore[arg-type] + capabilities=capabilities, # type: ignore[arg-type] + query_embedder=embedder, # type: ignore[arg-type] + partition_policy=policy, + ) + return SimpleNamespace( + retriever=retriever, + schema=schema, + planner=planner, + resolver=resolver, + embedder=embedder, + container=container, + capabilities=capabilities, + compiler_sentinel=compiler_sentinel, + executor_sentinel=executor_sentinel, + ctx_sentinel=ctx_sentinel, + comp_rec=comp_rec, + exec_rec=exec_rec, + planner_rec=planner_rec, + ctx_rec=ctx_rec, + resolver_rec=resolver_rec, + ) + + +# ═══════════════════════════ __init__ wiring ══════════════════════════════ + + +def test_init_defaults_policy_to_partition_query_policy(monkeypatch) -> None: + b = _build(monkeypatch) + assert isinstance(b.retriever.policy, PartitionQueryPolicy) + + +def test_init_uses_provided_policy(monkeypatch) -> None: + policy = PartitionQueryPolicy() + b = _build(monkeypatch, policy=policy) + assert b.retriever.policy is policy + + +def test_init_stores_schema_capabilities_embedder(monkeypatch) -> None: + embedder = FakeEmbedder([0.1]) + b = _build(monkeypatch, embedder=embedder) + assert b.retriever.schema is b.schema + assert b.retriever.capabilities is b.capabilities + assert b.retriever._embedder is embedder + + +def test_init_embedder_defaults_to_none(monkeypatch) -> None: + b = _build(monkeypatch) + assert b.retriever._embedder is None + + +def test_init_constructs_compiler_and_executor_with_expected_args(monkeypatch) -> None: + b = _build(monkeypatch) + assert b.comp_rec.calls == [((b.schema,), {})] + assert b.exec_rec.calls == [((b.container,), {})] + assert b.retriever._compiler is b.compiler_sentinel + assert b.retriever._executor is b.executor_sentinel + + +def test_init_constructs_planner_positionally(monkeypatch) -> None: + b = _build(monkeypatch) + (args, kwargs) = b.planner_rec.calls[0] + assert args == (b.schema, b.capabilities, b.retriever.policy) + assert kwargs == {} + assert b.retriever._planner is b.planner + + +def test_init_constructs_context_with_keywords(monkeypatch) -> None: + b = _build(monkeypatch) + (_args, kwargs) = b.ctx_rec.calls[0] + assert kwargs == { + "schema": b.schema, + "compiler": b.compiler_sentinel, + "executor": b.executor_sentinel, + "capabilities": b.capabilities, + "policy": b.retriever.policy, + } + assert b.retriever._ctx is b.ctx_sentinel + + +def test_init_builds_resolver_positionally(monkeypatch) -> None: + b = _build(monkeypatch) + (args, kwargs) = b.resolver_rec.calls[0] + assert args == ( + b.schema, + b.compiler_sentinel, + b.executor_sentinel, + b.retriever.policy, + ) + assert kwargs == {} + assert b.retriever._resolver is b.resolver + + +# ═══════════════════════════════ search ═══════════════════════════════════ + + +def test_search_resolves_vector_field_when_present(monkeypatch) -> None: + strat = FakeSearchStrategy(requires_embedding=False, result=[]) + b = _build(monkeypatch, search_strategy=strat) + b.retriever.search(SearchRequest(query="q", vector_field="vec")) + assert b.schema.vector_calls == ["vec"] + + +def test_search_skips_vector_resolution_when_absent(monkeypatch) -> None: + strat = FakeSearchStrategy(requires_embedding=False, result=[]) + b = _build(monkeypatch, search_strategy=strat) + b.retriever.search(SearchRequest(query="q")) + assert b.schema.vector_calls == [] + + +def test_search_resolves_text_fields_when_present(monkeypatch) -> None: + strat = FakeSearchStrategy(requires_embedding=False, result=[]) + b = _build(monkeypatch, search_strategy=strat) + b.retriever.search(SearchRequest(query="q", text_fields=["a", "b"])) + assert b.schema.text_calls == [["a", "b"]] + + +def test_search_skips_text_resolution_when_none_or_empty(monkeypatch) -> None: + strat = FakeSearchStrategy(requires_embedding=False, result=[]) + b = _build(monkeypatch, search_strategy=strat) + b.retriever.search(SearchRequest(query="q", text_fields=None)) + b.retriever.search(SearchRequest(query="q", text_fields=[])) + assert b.schema.text_calls == [] + + +def test_search_calls_planner_and_returns_execute_result(monkeypatch) -> None: + items = [_item("a"), _item("b")] + strat = FakeSearchStrategy(requires_embedding=False, result=items) + b = _build(monkeypatch, search_strategy=strat) + req = SearchRequest(query="q") + out = b.retriever.search(req) + assert out is items + assert b.planner.search_reqs == [req] + + +def test_search_passes_ctx_identity_to_execute(monkeypatch) -> None: + strat = FakeSearchStrategy(requires_embedding=False, result=[]) + b = _build(monkeypatch, search_strategy=strat) + req = SearchRequest(query="q") + b.retriever.search(req) + (passed_req, passed_ctx) = strat.execute_calls[0] + assert passed_req is req + assert passed_ctx is b.retriever._ctx + + +def test_search_no_embedding_needed_ignores_missing_embedder(monkeypatch) -> None: + strat = FakeSearchStrategy(requires_embedding=False, result=[]) + b = _build(monkeypatch, search_strategy=strat, embedder=None) + req = SearchRequest(query="q") # query_vector None, embedder None + b.retriever.search(req) # must not raise + assert strat.execute_calls[0][0].query_vector is None + + +def test_search_embedding_needed_but_vector_present_skips_embedder(monkeypatch) -> None: + strat = FakeSearchStrategy(requires_embedding=True, result=[]) + embedder = FakeEmbedder([9.9]) + b = _build(monkeypatch, search_strategy=strat, embedder=embedder) + req = SearchRequest(query="q", query_vector=[1.0, 2.0]) + b.retriever.search(req) + assert embedder.calls == [] + assert strat.execute_calls[0][0].query_vector == [1.0, 2.0] + + +def test_search_embeds_query_when_needed_and_updates_request(monkeypatch) -> None: + strat = FakeSearchStrategy(requires_embedding=True, result=[]) + embedder = FakeEmbedder([0.1, 0.2, 0.3]) + b = _build(monkeypatch, search_strategy=strat, embedder=embedder) + req = SearchRequest(query="hello") + b.retriever.search(req) + assert embedder.calls == ["hello"] + executed_req = strat.execute_calls[0][0] + assert executed_req.query_vector == [0.1, 0.2, 0.3] + assert executed_req.query == "hello" + assert req.query_vector is None # original request not mutated + + +def test_search_embedding_needed_without_embedder_raises(monkeypatch) -> None: + strat = FakeSearchStrategy(requires_embedding=True, result=[]) + b = _build(monkeypatch, search_strategy=strat, embedder=None) + with pytest.raises(EmbeddingProfileMismatch): + b.retriever.search(SearchRequest(query="q")) + assert strat.execute_calls == [] # execution never reached + + +# ═══════════════════════════ grep_candidates ══════════════════════════════ + + +def test_grep_resolves_single_text_field_when_present(monkeypatch) -> None: + strat = FakeGrepStrategy(result=[]) + b = _build(monkeypatch, grep_strategy=strat) + b.retriever.grep_candidates(GrepRequest(pattern="p", text_field="body")) + assert b.schema.text_calls == [["body"]] + + +def test_grep_skips_text_resolution_when_absent(monkeypatch) -> None: + strat = FakeGrepStrategy(result=[]) + b = _build(monkeypatch, grep_strategy=strat) + b.retriever.grep_candidates(GrepRequest(pattern="p")) + assert b.schema.text_calls == [] + + +def test_grep_calls_planner_and_returns_candidates(monkeypatch) -> None: + items = [_item("x")] + strat = FakeGrepStrategy(result=items) + b = _build(monkeypatch, grep_strategy=strat) + req = GrepRequest(pattern="p") + out = b.retriever.grep_candidates(req) + assert out is items + assert b.planner.grep_reqs == [req] + + +def test_grep_passes_ctx_identity(monkeypatch) -> None: + strat = FakeGrepStrategy(result=[]) + b = _build(monkeypatch, grep_strategy=strat) + req = GrepRequest(pattern="p") + b.retriever.grep_candidates(req) + (passed_req, passed_ctx) = strat.candidate_calls[0] + assert passed_req is req + assert passed_ctx is b.retriever._ctx + + +# ═══════════════════════════ read_document ════════════════════════════════ + + +def test_read_document_delegates_to_resolver(monkeypatch) -> None: + doc = NormalizedDocument(chunk_texts=["hello"]) + b = _build(monkeypatch, document=doc) + req = ReadDocumentRequest(document_id="d1") + out = b.retriever.read_document(req) + assert out is doc + assert b.resolver.calls == [req] diff --git a/cosmos-retriever/tests/unit/test_runtime_config.py b/cosmos-retriever/tests/unit/test_runtime_config.py new file mode 100644 index 0000000..a78da47 --- /dev/null +++ b/cosmos-retriever/tests/unit/test_runtime_config.py @@ -0,0 +1,138 @@ +from __future__ import annotations + +import asyncio + +from cosmos_retriever.config import RetrieverSettings, RuntimeConfig +from cosmos_retriever.server import _RetrieverPool + + +def _settings() -> RetrieverSettings: + return RetrieverSettings( + account_uri="https://x.documents.azure.com:443/", + cosmos_database="db", + cosmos_corpus_container="corpus", + ) + + +def test_apply_structural_overrides_produces_new_settings() -> None: + s = _settings() + rc = RuntimeConfig( + inference_backend="openai_chat", + chat_base_url="http://chat/v1", + chat_api_key="secret123", + chat_model="my-model", + openai_api_key="embkey", + openai_embedding_model="emb-model", + embed_endpoint="http://embed/v1", + embed_query_instruction="do the thing", + account_uri="https://acct.documents.azure.com:443/", + schema_override={"document_id_path": "/docid", "use_dunder_codec": True}, + search_display_limit=7, + ) + eff = s.apply_structural_overrides(rc) + assert eff.inference_backend == "openai_chat" + assert eff.chat_base_url == "http://chat/v1" + assert eff.chat_api_key.get_secret_value() == "secret123" + assert eff.chat_model == "my-model" + assert eff.openai_api_key.get_secret_value() == "embkey" + assert eff.openai_embedding_model == "emb-model" + assert eff.embed_endpoint == "http://embed/v1" + assert eff.account_uri == "https://acct.documents.azure.com:443/" + assert eff.cosmos_retriever_schema_override is not None + assert str(eff.cosmos_retriever_schema_override.document_id_path) == "/docid" + assert eff.cosmos_retriever_schema_override.use_dunder_codec is True + assert eff.cosmos_retriever_search_display_limit == 7 + # base is untouched + assert s.chat_model is None and s.inference_backend == "openai_responses" + + +def test_apply_none_returns_same_object() -> None: + s = _settings() + assert s.apply_structural_overrides(None) is s + + +def test_structural_key_ignores_execution_fields() -> None: + a = RuntimeConfig(chat_model="m", chat_max_turns=1, chat_temperature=0.1, max_documents=5) + b = RuntimeConfig(chat_model="m", chat_max_turns=999, chat_temperature=1.9, max_documents=30) + assert a.structural_key() == b.structural_key() + + +def test_structural_key_distinguishes_structural_fields() -> None: + a = RuntimeConfig(chat_model="m") + b = RuntimeConfig(chat_model="other") + assert a.structural_key() != b.structural_key() + + +def test_structural_key_hashes_secrets() -> None: + key = RuntimeConfig(chat_api_key="topsecret", openai_api_key="alsosecret").structural_key() + flat = str(key) + assert "topsecret" not in flat and "alsosecret" not in flat + + +def test_validators_and_extra_forbid() -> None: + for kwargs in ( + {"inference_backend": "bogus"}, + {"schema_override": "weird"}, + {"unknown_field": 1}, + ): + try: + RuntimeConfig(**kwargs) + except Exception: + continue + raise AssertionError(f"expected validation error for {kwargs}") + + +def test_pool_shares_retriever_for_execution_only_overrides() -> None: + pool = _RetrieverPool(_settings()) + builds = {"n": 0} + + def fake_build(scope, overrides): + builds["n"] += 1 + return object() + + pool._build = fake_build # type: ignore[method-assign] + + async def run() -> None: + r1, _ = await pool.get(None, None, RuntimeConfig(chat_max_turns=5)) + r2, _ = await pool.get(None, None, RuntimeConfig(chat_max_turns=99)) + assert r1 is r2 + + asyncio.run(run()) + assert builds["n"] == 1 + + +def test_pool_rebuilds_for_structural_overrides() -> None: + pool = _RetrieverPool(_settings()) + builds = {"n": 0} + + def fake_build(scope, overrides): + builds["n"] += 1 + return object() + + pool._build = fake_build # type: ignore[method-assign] + + async def run() -> None: + await pool.get(None, None, RuntimeConfig(chat_model="a")) + await pool.get(None, None, RuntimeConfig(chat_model="b")) + await pool.get(None, None, None) + + asyncio.run(run()) + assert builds["n"] == 3 + + +def _run_all() -> int: + fns = [v for k, v in sorted(globals().items()) if k.startswith("test_") and callable(v)] + failed = 0 + for fn in fns: + try: + fn() + print(f"PASS {fn.__name__}") + except BaseException as exc: # noqa: BLE001 + failed += 1 + print(f"FAIL {fn.__name__}: {type(exc).__name__}: {exc}") + print(f"\n{len(fns) - failed}/{len(fns)} passed") + return 1 if failed else 0 + + +if __name__ == "__main__": + raise SystemExit(_run_all()) diff --git a/cosmos-retriever/tests/unit/test_security.py b/cosmos-retriever/tests/unit/test_security.py new file mode 100644 index 0000000..a796769 --- /dev/null +++ b/cosmos-retriever/tests/unit/test_security.py @@ -0,0 +1,246 @@ +"""Security-focused tests for cosmos-retriever. + +Consolidates the safety-critical behaviours that are easy to regress: + + 1. Secret masking / non-leakage — RetrieverSettings.redacted_config, + CorpusConfig / settings repr, and RuntimeConfig.structural_key never + expose raw secret values. + 2. Path-injection defense — CosmosPath.parse rejects SQL/traversal/control + payloads, and render() escapes so a segment cannot break out of the + ["..."] bracketing. + 3. Read-only SQL enforcement — _SELECT_RE and RunQueryTool reject write / + DDL statements (case- and whitespace-insensitive) and never reach the + Cosmos client for them. + +Distinct sentinel secret values are used so a leak is detectable by substring +search of the serialized output. +""" +from __future__ import annotations + +import json +from types import SimpleNamespace + +import pytest +from pydantic import SecretStr + +from cosmos_retriever.config import CorpusConfig, RetrieverSettings, RuntimeConfig +from cosmos_retriever.retrieval.errors import UnsafeCosmosPath +from cosmos_retriever.retrieval.paths import CosmosPath +from cosmos_retriever.tools import _SELECT_RE, RunQueryTool + +# ════════════════════════ 1. secret masking / leakage ═════════════════════ + +# Unique per-field sentinels so any leak is greppable. +_CHAT_SECRET = "CHAT-SENTINEL-a1" +_OPENAI_SECRET = "OPENAI-SENTINEL-b2" +_COSMOS_SECRET = "COSMOS-SENTINEL-c3" +_BASETEN_SECRET = "BASETEN-SENTINEL-d4" + + +def _settings_with_secrets() -> RetrieverSettings: + return RetrieverSettings( + _env_file=None, + chat_api_key=_CHAT_SECRET, + openai_api_key=_OPENAI_SECRET, + cosmos_key=_COSMOS_SECRET, + baseten_api_key=_BASETEN_SECRET, + cosmos_database="db", + ) + + +def test_redacted_config_masks_all_secret_fields() -> None: + red = _settings_with_secrets().redacted_config() + assert red["chat_api_key"] == "***set***" + assert red["openai_api_key"] == "***set***" + assert red["cosmos_key"] == "***set***" + assert red["baseten_api_key"] == "***set***" + + +def test_redacted_config_never_contains_raw_secret_values() -> None: + blob = json.dumps(_settings_with_secrets().redacted_config()) + for secret in (_CHAT_SECRET, _OPENAI_SECRET, _COSMOS_SECRET, _BASETEN_SECRET): + assert secret not in blob + + +def test_redacted_config_unset_secrets_are_none() -> None: + red = RetrieverSettings(_env_file=None).redacted_config() + for key in ("chat_api_key", "openai_api_key", "cosmos_key", "baseten_api_key"): + assert red[key] is None + + +def test_settings_repr_and_str_do_not_leak_secrets() -> None: + s = _settings_with_secrets() + for text in (repr(s), str(s)): + for secret in (_CHAT_SECRET, _OPENAI_SECRET, _COSMOS_SECRET, _BASETEN_SECRET): + assert secret not in text + + +def test_settings_model_dump_json_does_not_leak_secrets() -> None: + dumped = _settings_with_secrets().model_dump_json() + for secret in (_CHAT_SECRET, _OPENAI_SECRET, _COSMOS_SECRET, _BASETEN_SECRET): + assert secret not in dumped + + +def test_corpusconfig_repr_does_not_leak_secrets() -> None: + c = CorpusConfig( + container="c", + account_uri="https://a", + database="db", + embed_base_url=None, + embed_api_key=SecretStr("EMBED-SENTINEL-e5"), + embed_model="m", + cosmos_key=SecretStr("CK-SENTINEL-f6"), + ) + assert "EMBED-SENTINEL-e5" not in repr(c) + assert "CK-SENTINEL-f6" not in repr(c) + + +def test_structural_key_hashes_and_never_exposes_raw_secrets() -> None: + rc = RuntimeConfig(chat_api_key="CK-RAW-xyz", openai_api_key="OK-RAW-xyz") + key = rc.structural_key() + blob = str(key) + assert "CK-RAW-xyz" not in blob + assert "OK-RAW-xyz" not in blob + # secrets are represented by 16-char sha256 prefixes, not the plaintext + assert any(isinstance(v, str) and len(v) == 16 for v in key) + + +def test_structural_key_none_secrets_stay_none() -> None: + key = RuntimeConfig().structural_key() + # positions 2 and 6 are the hashed chat/openai keys + assert key[2] is None and key[6] is None + + +def test_get_baseten_client_error_does_not_echo_secret() -> None: + # Only model_url set: the error path must not include any key material. + s = RetrieverSettings(_env_file=None, baseten_model_url="https://bt") + with pytest.raises(ValueError) as exc: + s.get_baseten_client() + assert "BASETEN_MODEL_URL" in str(exc.value) + + +# ════════════════════════ 2. path-injection defense ═══════════════════════ + + +@pytest.mark.parametrize( + "payload", + [ + '/a"] OR 1=1', # quote + bracket breakout attempt + "/a'; DROP TABLE x", # single-quote SQL injection + "/a`b", # backtick + "/a[0]", # bracket indexing + "/a*", # wildcard + "/a;b", # statement separator + "/a=b", # operator + "/a(b)", # parens + "/a|b", # pipe + "/a$b", # dollar + "/a%b", # percent + "/a\nb", # newline (control char) + "/a\tb", # tab + "/a\\b", # backslash + "/..", # path traversal segment + "/../etc/passwd", # traversal chain + "/a/../b", # embedded traversal + ], +) +def test_cosmospath_rejects_injection_payloads(payload: str) -> None: + with pytest.raises(UnsafeCosmosPath): + CosmosPath.parse(payload) + + +def test_render_escapes_bracket_breakout_segment() -> None: + # A directly-constructed hostile segment must not break out of ["..."]: + # every embedded quote is backslash-escaped. + rendered = CosmosPath(segments=('x"][\"y',)).render() + assert rendered == 'c["x\\"][\\"y"]' + # no unescaped closing-then-opening bracket sequence survives + assert '"]["' not in rendered.replace('\\"', "") + + +def test_render_escapes_backslash_before_quote() -> None: + # Backslash is escaped first so it can't neutralize the quote escaping. + assert CosmosPath(segments=('\\"',)).render() == 'c["\\\\\\""]' + + +# ════════════════════════ 3. read-only SQL enforcement ════════════════════ + + +class _FakeContainer: + def __init__(self, rows=None): + self.rows = rows or [] + self.received = None # set only if query_items is actually called + + def query_items(self, query, enable_cross_partition_query, max_item_count): + self.received = query + yield from self.rows + + +class _FakeCosmosClient: + def __init__(self, container): + self._container = container + + def get_database_client(self, name): + return SimpleNamespace(get_container_client=lambda _n: self._container) + + +def _run_tool(container) -> RunQueryTool: + return RunQueryTool( + client=_FakeCosmosClient(container), + default_database="db", + default_container="cont", + ) + + +@pytest.mark.parametrize( + "query", + [ + "INSERT INTO c VALUES (1)", + " \n\t update c set x = 1", + "(( DELETE FROM c ))", + "DrOp TaBlE c", + "MERGE INTO c", + "UPSERT c", + "EXEC sp_bad", + "CALL something()", + "ALTER CONTAINER c", + "CREATE INDEX i", + "TRUNCATE c", + "GRANT ALL", + "REPLACE INTO c", + "WITH t AS (SELECT 1) DELETE FROM t", # CTE prefix is not a SELECT start + ], +) +def test_select_regex_rejects_non_select(query: str) -> None: + assert _SELECT_RE.match(query) is None + + +@pytest.mark.parametrize( + "query", + ["SELECT * FROM c", " \n select c.id from c", "(select 1)", "SeLeCt x"], +) +def test_select_regex_accepts_read_queries(query: str) -> None: + assert _SELECT_RE.match(query) is not None + + +def test_run_query_blocks_write_and_never_touches_client() -> None: + container = _FakeContainer() + text, _ = _run_tool(container)({"query": "DELETE FROM c"}) + assert "only read-only SELECT queries are allowed" in text + assert container.received is None # client never queried + + +def test_run_query_allows_select_positive_control() -> None: + container = _FakeContainer(rows=[{"id": 1}]) + text, _ = _run_tool(container)({"query": "SELECT * FROM c"}) + assert container.received == "SELECT * FROM c" # guard let the read through + assert "row(s)" in text + + +def test_select_regex_multistatement_is_a_documented_limitation() -> None: + # KNOWN LIMITATION: _SELECT_RE only anchors the *start* of the query, so a + # chained statement after a ';' still matches. This is not exploitable via + # Azure Cosmos, whose query_items executes read-only SQL and rejects DDL/DML; + # the guard is defense-in-depth, not the sole control. Pinning the current + # behavior so any future hardening (rejecting ';') updates this test. + assert _SELECT_RE.match("SELECT * FROM c; DROP TABLE x") is not None diff --git a/cosmos-retriever/tests/unit/test_server.py b/cosmos-retriever/tests/unit/test_server.py new file mode 100644 index 0000000..2f3682c --- /dev/null +++ b/cosmos-retriever/tests/unit/test_server.py @@ -0,0 +1,346 @@ +"""Exhaustive tests for the HTTP service (`cosmos_retriever.server`). + +Covers, without Cosmos / OpenAI / real engines: + + 1. RetrievalScope.resolve — db/container fallbacks + "*" whole-database mode + 2. SearchRequest — defaults, alias, populate_by_name, bounds + 3. _RetrieverPool — stats mapping, build-once caching, per-key locks, + _build deep-copy + scope wiring, update_settings + (clear vs rebuild), settings property + 4. create_app routes — /health (with/without pool), GET/PATCH /config, + POST /search (happy, missing db, engine error), + get_settings fallback + +`server.CosmosRetriever` is patched with a fake recording engine; real +`RetrieverSettings` (with ``_env_file=None`` so no .env leaks) drives config +logic. Async pool methods are driven with ``anyio.run``; routes via +``fastapi.testclient.TestClient``. +""" +from __future__ import annotations + +from dataclasses import dataclass + +import anyio +import pytest +from fastapi.testclient import TestClient + +from cosmos_retriever import config, server +from cosmos_retriever.config import RetrieverSettings, RuntimeConfig +from cosmos_retriever.server import RetrievalScope, SearchRequest, _RetrieverPool, create_app + +# ────────────────────────────── helpers ─────────────────────────────────── + + +def _settings(**kw) -> RetrieverSettings: + """Deterministic settings: init kwargs win over env, .env disabled.""" + base: dict = { + "cosmos_database": None, + "cosmos_corpus_container": None, + "cosmos_retriever_cache_max_entries": 4, + "cosmos_retriever_cache_ttl_seconds": 100.0, + } + base.update(kw) + return RetrieverSettings(_env_file=None, **base) + + +@dataclass +class FakeResult: + answer: str + documents: list + + +def _install_fake_retriever(monkeypatch, search_fn=None) -> list: + """Patch server.CosmosRetriever with a recorder; return list of built engines.""" + built: list = [] + + class FakeCosmosRetriever: + def __init__(self, *, settings, corpus_name): + self.settings = settings + self.corpus_name = corpus_name + self.search_calls: list = [] + built.append(self) + + def search(self, query, max_documents=20, overrides=None): + self.search_calls.append((query, max_documents, overrides)) + if search_fn is not None: + return search_fn(query, max_documents, overrides) + return FakeResult(answer=f"ans:{query}", documents=[1, 2]) + + monkeypatch.setattr(server, "CosmosRetriever", FakeCosmosRetriever) + return built + + +def _run(func, *args): + return anyio.run(func, *args) + + +# ═══════════════════════ 1. RetrievalScope.resolve ════════════════════════ + + +def test_scope_is_namedtuple_fields() -> None: + s = RetrievalScope(database="d", container="c") + assert (s.database, s.container) == ("d", "c") + + +def test_scope_uses_settings_defaults() -> None: + s = RetrievalScope.resolve(_settings(cosmos_database="D", cosmos_corpus_container="C"), None, None) + assert s == ("D", "C") + + +def test_scope_container_defaults_to_star() -> None: + s = RetrievalScope.resolve(_settings(cosmos_database="D"), None, None) + assert s == ("D", "*") + + +def test_scope_database_none_when_unset() -> None: + s = RetrievalScope.resolve(_settings(), None, None) + assert s == (None, "*") + + +def test_scope_explicit_args_override_settings() -> None: + s = RetrievalScope.resolve( + _settings(cosmos_database="D", cosmos_corpus_container="C"), "D2", "C2" + ) + assert s == ("D2", "C2") + + +def test_scope_explicit_db_container_omitted_uses_star() -> None: + s = RetrievalScope.resolve(_settings(cosmos_database="D"), "D2", None) + assert s == ("D2", "*") + + +# ═══════════════════════════ 2. SearchRequest ═════════════════════════════ + + +def test_search_request_defaults() -> None: + r = SearchRequest(query="q") + assert r.max_documents == 20 + assert r.database is None and r.container is None and r.overrides is None + + +def test_search_request_alias_and_field_name() -> None: + assert SearchRequest(**{"query": "q", "maxDocuments": 5}).max_documents == 5 + assert SearchRequest(query="q", max_documents=7).max_documents == 7 # populate_by_name + + +def test_search_request_query_min_length() -> None: + with pytest.raises(ValueError): + SearchRequest(query="") + + +@pytest.mark.parametrize("n", [1, 30]) +def test_search_request_max_documents_bounds_ok(n: int) -> None: + assert SearchRequest(query="q", max_documents=n).max_documents == n + + +@pytest.mark.parametrize("n", [0, 31]) +def test_search_request_max_documents_out_of_range(n: int) -> None: + with pytest.raises(ValueError): + SearchRequest(query="q", max_documents=n) + + +def test_search_request_overrides_parsed_to_runtime_config() -> None: + r = SearchRequest(query="q", overrides={"chat_model": "m"}) + assert isinstance(r.overrides, RuntimeConfig) + assert r.overrides.chat_model == "m" + + +# ═══════════════════════════ 3. _RetrieverPool ════════════════════════════ + + +async def _get(pool, a, b, o): + return await pool.get(a, b, o) + + +async def _get_twice(pool, a, b, o1, o2): + r1 = await pool.get(a, b, o1) + r2 = await pool.get(a, b, o2) + return r1, r2 + + +def test_pool_stats_shape_initial() -> None: + pool = _RetrieverPool(_settings(cosmos_retriever_cache_max_entries=9, cosmos_retriever_cache_ttl_seconds=42.0)) + s = pool.stats() + assert set(s) == { + "entries", "max_entries", "ttl_seconds", "hits", "misses", "evictions", "expirations" + } + assert s["entries"] == 0 + assert s["max_entries"] == 9 + assert s["ttl_seconds"] == 42.0 + + +def test_pool_settings_property() -> None: + st = _settings(cosmos_database="D") + pool = _RetrieverPool(st) + assert pool.settings is st + + +def test_pool_get_builds_once_and_caches(monkeypatch) -> None: + built = _install_fake_retriever(monkeypatch) + pool = _RetrieverPool(_settings(cosmos_database="D")) + (r1, lock1), (r2, lock2) = _run(_get_twice, pool, None, None, None, None) + assert r1 is r2 # cached engine reused + assert lock1 is lock2 # stable per-key lock + assert len(built) == 1 + stats = pool.stats() + # First get misses twice (double-checked lock re-reads the cache); second hits once. + assert stats["entries"] == 1 and stats["hits"] == 1 and stats["misses"] == 2 + + +def test_pool_get_scope_wires_database_and_container(monkeypatch) -> None: + built = _install_fake_retriever(monkeypatch) + pool = _RetrieverPool(_settings(cosmos_database="D")) + _run(_get, pool, None, None, None) + engine = built[0] + assert engine.settings.cosmos_database == "D" + assert engine.corpus_name == "*" # container omitted -> whole database + + +def test_pool_build_deep_copies_and_does_not_mutate_pool_settings(monkeypatch) -> None: + built = _install_fake_retriever(monkeypatch) + st = _settings(cosmos_database="D") + pool = _RetrieverPool(st) + _run(_get, pool, "OTHER", None, None) + engine = built[0] + assert engine.settings is not st # deep copy, not the pool's own settings + assert engine.settings.cosmos_database == "OTHER" + assert st.cosmos_database == "D" # pool settings untouched + + +def test_pool_distinct_overrides_build_separate_engines(monkeypatch) -> None: + built = _install_fake_retriever(monkeypatch) + pool = _RetrieverPool(_settings(cosmos_database="D")) + o1 = RuntimeConfig(chat_model="a") + o2 = RuntimeConfig(chat_model="b") + (r1, lock1), (r2, lock2) = _run(_get_twice, pool, None, None, o1, o2) + assert r1 is not r2 + assert lock1 is not lock2 + assert len(built) == 2 + + +def test_pool_update_settings_clears_cache_when_size_unchanged(monkeypatch) -> None: + _install_fake_retriever(monkeypatch) + pool = _RetrieverPool(_settings(cosmos_database="D", cosmos_retriever_cache_max_entries=4)) + _run(_get, pool, None, None, None) + assert pool.stats()["entries"] == 1 + cache_before = pool._cache + new = _settings(cosmos_database="D2", cosmos_retriever_cache_max_entries=4, cosmos_retriever_cache_ttl_seconds=100.0) + _run(pool.update_settings, new) + assert pool._cache is cache_before # same cache object, just cleared + assert pool.stats()["entries"] == 0 + assert pool.settings is new + assert pool._locks == {} + + +def test_pool_update_settings_rebuilds_cache_on_size_change(monkeypatch) -> None: + _install_fake_retriever(monkeypatch) + pool = _RetrieverPool(_settings(cosmos_database="D", cosmos_retriever_cache_max_entries=4)) + _run(_get, pool, None, None, None) + cache_before = pool._cache + new = _settings(cosmos_database="D", cosmos_retriever_cache_max_entries=8) + _run(pool.update_settings, new) + assert pool._cache is not cache_before # rebuilt due to size change + assert pool.stats()["max_entries"] == 8 + assert pool.stats()["entries"] == 0 + + +# ═══════════════════════════ 4. create_app routes ═════════════════════════ + + +def test_health_without_pool_returns_empty_cache() -> None: + app = create_app(_settings()) + client = TestClient(app) # no context manager -> lifespan not run + resp = client.get("/health") + assert resp.status_code == 200 + assert resp.json() == {"status": "ok", "retriever_cache": {}} + + +def test_health_with_pool_returns_stats() -> None: + app = create_app(_settings()) + with TestClient(app) as client: + body = client.get("/health").json() + assert body["status"] == "ok" + assert set(body["retriever_cache"]) >= {"entries", "max_entries", "hits", "misses"} + + +def test_get_config_returns_redacted_and_pool() -> None: + app = create_app(_settings(cosmos_database="D", cosmos_retriever_cache_max_entries=6)) + with TestClient(app) as client: + body = client.get("/config").json() + assert body["config"]["cosmos_database"] == "D" + assert body["config"]["cache_max_entries"] == 6 + assert "pool" in body and body["pool"]["entries"] == 0 + + +def test_patch_config_applies_and_propagates(monkeypatch) -> None: + built = _install_fake_retriever(monkeypatch) + app = create_app(_settings(cosmos_database="D")) + with TestClient(app) as client: + resp = client.patch("/config", json={"chat_model": "newmodel"}) + body = resp.json() + assert resp.status_code == 200 + assert body["status"] == "ok" + assert body["changed"] == ["chat_model"] + assert body["config"]["chat_model"] == "newmodel" + # New engine built after update carries the new setting. + client.post("/search", json={"query": "q", "database": "D"}) + assert built[-1].settings.chat_model == "newmodel" + + +def test_patch_config_error_returns_400(monkeypatch) -> None: + def _boom(self, update): + raise ValueError("bad update") + + monkeypatch.setattr(config.RetrieverSettings, "apply_server_updates", _boom) + app = create_app(_settings(cosmos_database="D")) + with TestClient(app) as client: + resp = client.patch("/config", json={"chat_model": "x"}) + assert resp.status_code == 400 + assert resp.json() == {"error": "bad update", "type": "ValueError"} + + +def test_search_happy_path_returns_result_dict(monkeypatch) -> None: + built = _install_fake_retriever(monkeypatch) + app = create_app(_settings(cosmos_database="D")) + with TestClient(app) as client: + resp = client.post( + "/search", + json={"query": "hi", "database": "D", "maxDocuments": 5, + "overrides": {"chat_model": "m"}}, + ) + assert resp.status_code == 200 + assert resp.json() == {"answer": "ans:hi", "documents": [1, 2]} + query, max_docs, overrides = built[0].search_calls[0] + assert query == "hi" and max_docs == 5 + assert isinstance(overrides, RuntimeConfig) and overrides.chat_model == "m" + + +def test_search_missing_database_returns_400(monkeypatch) -> None: + _install_fake_retriever(monkeypatch) + app = create_app(_settings()) # no default database + with TestClient(app) as client: + resp = client.post("/search", json={"query": "hi"}) + assert resp.status_code == 400 + assert resp.json()["type"] == "ValueError" + assert "Missing required field: database" in resp.json()["error"] + + +def test_search_engine_exception_returns_500(monkeypatch) -> None: + def _raise(query, max_documents, overrides): + raise RuntimeError("kaboom") + + _install_fake_retriever(monkeypatch, search_fn=_raise) + app = create_app(_settings(cosmos_database="D")) + with TestClient(app) as client: + resp = client.post("/search", json={"query": "hi", "database": "D"}) + assert resp.status_code == 500 + assert resp.json() == {"error": "kaboom", "type": "RuntimeError"} + + +def test_create_app_falls_back_to_get_settings(monkeypatch) -> None: + sentinel = _settings(cosmos_database="FROM_GET_SETTINGS") + monkeypatch.setattr(server, "get_settings", lambda: sentinel) + app = create_app() # settings=None -> get_settings() + with TestClient(app) as client: + body = client.get("/config").json() + assert body["config"]["cosmos_database"] == "FROM_GET_SETTINGS" diff --git a/cosmos-retriever/tests/unit/test_strategies.py b/cosmos-retriever/tests/unit/test_strategies.py new file mode 100644 index 0000000..89552b6 --- /dev/null +++ b/cosmos-retriever/tests/unit/test_strategies.py @@ -0,0 +1,376 @@ +"""Exhaustive tests for `cosmos_retriever.retrieval.strategies`. + +Each strategy is a thin orchestrator: resolve fields -> decide cross-partition +-> compile -> execute -> normalize. Tests isolate that wiring by faking the +schema / compiler / executor / policy and patching ``strategies.normalize_rows`` +with a recorder, so every ``compile_*`` argument and every ``normalize_rows`` +argument (strategy name, channels, projected_aliases, queried_text_fields) is +asserted without Cosmos. ClientSideFusion's RRF math is checked against real +RetrievedItem instances; grep's all-stopword short-circuit is verified to touch +no compiler/executor. +""" +from __future__ import annotations + +import pytest + +from cosmos_retriever.retrieval import strategies as strat_mod +from cosmos_retriever.retrieval.errors import ( + CrossPartitionQueryDisabled, + UnboundedScanRejected, +) +from cosmos_retriever.retrieval.models import GrepRequest, RetrievedItem, SearchRequest +from cosmos_retriever.retrieval.strategies import ( + BoundedScanStrategy, + ClientSideFusionStrategy, + FullTextGrepCandidateStrategy, + FullTextSearchStrategy, + GrepCandidateStrategy, + NativeHybridStrategy, + RetrievalContext, + SearchStrategy, + VectorSearchStrategy, + _resolve_cross_partition, +) + +# ────────────────────────────── fakes ───────────────────────────────────── + + +class FakeSchema: + def __init__(self, vec="VEC_PATH", text=("T1", "T2")): + self.vec = vec + self.text = list(text) + self.vector_calls: list = [] + self.text_calls: list = [] + + def resolve_vector_field(self, name): + self.vector_calls.append(name) + return self.vec + + def resolve_text_fields(self, names): + self.text_calls.append(names) + return self.text + + +class FakeCompiled: + def __init__(self, aliases="ALIASES"): + self.projected_aliases = aliases + self.warnings: list[str] = [] + + +class FakeCompiler: + def __init__(self): + self.calls: list[tuple[str, dict]] = [] + self.compiled: list[FakeCompiled] = [] + + def _mk(self, method: str, kwargs: dict) -> FakeCompiled: + c = FakeCompiled() + self.calls.append((method, kwargs)) + self.compiled.append(c) + return c + + def compile_hybrid(self, **kw): + return self._mk("hybrid", kw) + + def compile_vector(self, **kw): + return self._mk("vector", kw) + + def compile_full_text(self, **kw): + return self._mk("full_text", kw) + + def compile_structured(self, **kw): + return self._mk("structured", kw) + + +class FakeExecutor: + def __init__(self, rows=None): + self.rows = rows if rows is not None else [{"row": 1}] + self.ran: list = [] + + def run(self, compiled): + self.ran.append(compiled) + return self.rows + + +class FakePolicy: + def __init__(self, cross=True, bounded=True): + self.allow_cross_partition_search = cross + self.allow_bounded_scan = bounded + + +class FakeNormalize: + def __init__(self): + self.calls: list[tuple] = [] + self.by_strategy: dict[str, list] = {} + self.default: list = [] + + def __call__(self, rows, **kwargs): + self.calls.append((rows, kwargs)) + return self.by_strategy.get(kwargs.get("strategy"), self.default) + + +@pytest.fixture +def norm(monkeypatch) -> FakeNormalize: + fake = FakeNormalize() + monkeypatch.setattr(strat_mod, "normalize_rows", fake) + return fake + + +def _ctx(schema=None, compiler=None, executor=None, policy=None) -> RetrievalContext: + return RetrievalContext( + schema=schema or FakeSchema(), + compiler=compiler or FakeCompiler(), + executor=executor or FakeExecutor(), + capabilities=None, # unused by execute + policy=policy or FakePolicy(), + ) + + +def _req(**kw) -> SearchRequest: + base = dict(query="q", limit=10) + base.update(kw) + return SearchRequest(**base) + + +# ═══════════════════════ _resolve_cross_partition ═════════════════════════ + + +def test_cross_partition_with_key_is_false() -> None: + assert _resolve_cross_partition("pk", FakePolicy(cross=False)) is False + + +def test_cross_partition_none_allowed_true() -> None: + assert _resolve_cross_partition(None, FakePolicy(cross=True)) is True + + +def test_cross_partition_none_disallowed_raises() -> None: + with pytest.raises(CrossPartitionQueryDisabled): + _resolve_cross_partition(None, FakePolicy(cross=False)) + + +# ═══════════════════════ base / class attributes ══════════════════════════ + + +def test_search_strategy_is_abstract() -> None: + with pytest.raises(TypeError): + SearchStrategy() # type: ignore[abstract] + + +def test_grep_strategy_is_abstract() -> None: + with pytest.raises(TypeError): + GrepCandidateStrategy() # type: ignore[abstract] + + +def test_strategy_names_and_embedding_flags() -> None: + assert (NativeHybridStrategy.name, NativeHybridStrategy.requires_embedding) == ( + "native_hybrid", True) + assert (VectorSearchStrategy.name, VectorSearchStrategy.requires_embedding) == ( + "vector", True) + assert (FullTextSearchStrategy.name, FullTextSearchStrategy.requires_embedding) == ( + "full_text", False) + assert (ClientSideFusionStrategy.name, ClientSideFusionStrategy.requires_embedding) == ( + "client_fusion", True) + assert ClientSideFusionStrategy._RRF_K == 60 + assert (BoundedScanStrategy.name, BoundedScanStrategy.requires_embedding) == ( + "bounded_scan", False) + + +# ═══════════════════════ NativeHybridStrategy ═════════════════════════════ + + +def test_native_hybrid_execute_wiring(norm) -> None: + schema, compiler, executor = FakeSchema(), FakeCompiler(), FakeExecutor() + ctx = _ctx(schema, compiler, executor) + norm.default = ["RESULT"] + req = _req(query="find", query_vector=[0.1], vector_field="vf", + text_fields=["a"], ignored_item_ids=["x"], limit=5) + out = NativeHybridStrategy().execute(req, ctx) + + assert out == ["RESULT"] + assert schema.vector_calls == ["vf"] + assert schema.text_calls == [["a"]] + method, kw = compiler.calls[0] + assert method == "hybrid" + assert kw == { + "query": "find", "query_vector": [0.1], "limit": 5, + "ignored_item_ids": ["x"], "filters": [], "partition_key": None, + "cross_partition": True, "vector_path": "VEC_PATH", "text_paths": ["T1", "T2"], + } + _rows, nkw = norm.calls[0] + assert nkw["strategy"] == "native_hybrid" + assert nkw["channels"] == ["vector", "full_text"] + assert nkw["projected_aliases"] == "ALIASES" + assert nkw["queried_text_fields"] == ["a"] + + +def test_native_hybrid_none_query_vector_becomes_empty(norm) -> None: + compiler = FakeCompiler() + NativeHybridStrategy().execute(_req(query_vector=None), _ctx(compiler=compiler)) + assert compiler.calls[0][1]["query_vector"] == [] + + +# ═══════════════════════ VectorSearchStrategy ═════════════════════════════ + + +def test_vector_execute_wiring(norm) -> None: + schema, compiler = FakeSchema(), FakeCompiler() + ctx = _ctx(schema, compiler) + VectorSearchStrategy().execute(_req(query_vector=[1.0], vector_field="vf"), ctx) + assert schema.vector_calls == ["vf"] + assert schema.text_calls == [] # vector never resolves text fields + method, kw = compiler.calls[0] + assert method == "vector" + assert kw["vector_path"] == "VEC_PATH" and kw["query_vector"] == [1.0] + nkw = norm.calls[0][1] + assert nkw["strategy"] == "vector" + assert nkw["channels"] == ["vector"] + assert "queried_text_fields" not in nkw # vector omits it + + +def test_vector_none_query_vector_becomes_empty(norm) -> None: + compiler = FakeCompiler() + VectorSearchStrategy().execute(_req(query_vector=None), _ctx(compiler=compiler)) + assert compiler.calls[0][1]["query_vector"] == [] + + +# ═══════════════════════ FullTextSearchStrategy ═══════════════════════════ + + +def test_full_text_execute_wiring(norm) -> None: + schema, compiler = FakeSchema(), FakeCompiler() + ctx = _ctx(schema, compiler) + FullTextSearchStrategy().execute(_req(query="hello", text_fields=["body"]), ctx) + assert schema.text_calls == [["body"]] + method, kw = compiler.calls[0] + assert method == "full_text" + assert kw["query"] == "hello" and kw["text_paths"] == ["T1", "T2"] + nkw = norm.calls[0][1] + assert nkw["strategy"] == "full_text" + assert nkw["channels"] == ["full_text"] + assert nkw["queried_text_fields"] == ["body"] + + +def test_full_text_cross_partition_false_with_partition_key(norm) -> None: + compiler = FakeCompiler() + FullTextSearchStrategy().execute(_req(partition_key="pk"), _ctx(compiler=compiler)) + assert compiler.calls[0][1]["cross_partition"] is False + + +def test_full_text_cross_partition_disabled_raises() -> None: + with pytest.raises(CrossPartitionQueryDisabled): + FullTextSearchStrategy().execute(_req(), _ctx(policy=FakePolicy(cross=False))) + + +# ═══════════════════════ ClientSideFusionStrategy ═════════════════════════ + + +def _fusion_ctx(norm, vector_hits, fts_hits): + norm.by_strategy = {"vector": vector_hits, "full_text": fts_hits} + return _ctx() + + +def test_client_fusion_rrf_ranking_and_metadata(norm) -> None: + vhits = [RetrievedItem(item_id="a", text="va"), + RetrievedItem(item_id="b", text="vb_vector")] + fhits = [RetrievedItem(item_id="b", text="vb_fts"), + RetrievedItem(item_id="c", text="vc")] + out = ClientSideFusionStrategy().execute(_req(limit=10), _fusion_ctx(norm, vhits, fhits)) + + assert [r.item_id for r in out] == ["b", "a", "c"] # b highest (in both) + assert [r.rank for r in out] == [0, 1, 2] + assert all(r.retrieval_strategy == "client_fusion" for r in out) + assert out[0].retrieval_channels == ["vector", "full_text"] + assert out[1].retrieval_channels == ["vector"] + assert out[2].retrieval_channels == ["full_text"] + # setdefault keeps the first (vector) instance for the shared id + assert out[0].text == "vb_vector" + # RRF scores: b = 1/61 + 1/60, a = 1/60, c = 1/61 + assert out[0].raw_scores["rrf"] == pytest.approx(1 / 61 + 1 / 60) + assert out[1].raw_scores["rrf"] == pytest.approx(1 / 60) + assert out[2].raw_scores["rrf"] == pytest.approx(1 / 61) + + +def test_client_fusion_truncates_to_limit(norm) -> None: + vhits = [RetrievedItem(item_id="a"), RetrievedItem(item_id="b")] + fhits = [RetrievedItem(item_id="b"), RetrievedItem(item_id="c")] + out = ClientSideFusionStrategy().execute(_req(limit=2), _fusion_ctx(norm, vhits, fhits)) + assert [r.item_id for r in out] == ["b", "a"] + + +def test_client_fusion_empty_inputs(norm) -> None: + assert ClientSideFusionStrategy().execute(_req(), _fusion_ctx(norm, [], [])) == [] + + +# ═══════════════════════ BoundedScanStrategy ══════════════════════════════ + + +def test_bounded_scan_disabled_raises() -> None: + with pytest.raises(UnboundedScanRejected): + BoundedScanStrategy().execute(_req(), _ctx(policy=FakePolicy(bounded=False))) + + +def test_bounded_scan_execute_wiring(norm) -> None: + compiler, executor = FakeCompiler(), FakeExecutor() + ctx = _ctx(compiler=compiler, executor=executor) + BoundedScanStrategy().execute(_req(limit=7, ignored_item_ids=["z"]), ctx) + method, kw = compiler.calls[0] + assert method == "structured" + assert kw == { + "limit": 7, "filters": [], "ignored_item_ids": ["z"], + "partition_key": None, "cross_partition": True, + } + assert compiler.compiled[0].warnings == ["bounded scan active"] + nkw = norm.calls[0][1] + assert nkw["strategy"] == "bounded_scan" + assert "channels" not in nkw + assert "queried_text_fields" not in nkw + + +# ═══════════════════════ FullTextGrepCandidateStrategy ════════════════════ + + +def test_grep_all_stopword_pattern_short_circuits(norm) -> None: + compiler, executor = FakeCompiler(), FakeExecutor() + ctx = _ctx(compiler=compiler, executor=executor) + out = FullTextGrepCandidateStrategy().candidates(GrepRequest(pattern="the and of"), ctx) + assert out == [] + assert compiler.calls == [] # never compiled + assert executor.ran == [] # never executed + + +def test_grep_empty_pattern_short_circuits(norm) -> None: + compiler = FakeCompiler() + out = FullTextGrepCandidateStrategy().candidates( + GrepRequest(pattern="!!! ???"), _ctx(compiler=compiler) + ) + assert out == [] + assert compiler.calls == [] + + +def test_grep_execute_wiring_with_field(norm) -> None: + schema, compiler = FakeSchema(), FakeCompiler() + ctx = _ctx(schema, compiler) + norm.default = ["G"] + req = GrepRequest(pattern="machine learning", text_field="body", candidate_limit=25) + out = FullTextGrepCandidateStrategy().candidates(req, ctx) + + assert out == ["G"] + assert schema.text_calls == [["body"]] + method, kw = compiler.calls[0] + assert method == "full_text" + assert kw["query"] == "machine learning" + assert kw["limit"] == 25 + assert kw["ignored_item_ids"] == [] # grep never carries ignored ids + assert kw["strategy"] == "grep_full_text" + nkw = norm.calls[0][1] + assert nkw["strategy"] == "grep_full_text" + assert nkw["channels"] == ["full_text"] + assert nkw["queried_text_fields"] == ["body"] + + +def test_grep_without_field_resolves_none(norm) -> None: + schema, compiler = FakeSchema(), FakeCompiler() + FullTextGrepCandidateStrategy().candidates( + GrepRequest(pattern="hello world"), _ctx(schema, compiler) + ) + assert schema.text_calls == [None] + assert norm.calls[0][1]["queried_text_fields"] is None diff --git a/cosmos-retriever/tests/unit/test_token_budget.py b/cosmos-retriever/tests/unit/test_token_budget.py new file mode 100644 index 0000000..1ef41be --- /dev/null +++ b/cosmos-retriever/tests/unit/test_token_budget.py @@ -0,0 +1,381 @@ +"""Comprehensive tests for the token-budget system. + +Covers, against the REAL agent loops (with a scripted fake LLM client + fake +tools), every budgeting feature: + + 1. Real pruning (prune_chunks_from_trajectory) + 2. Cross-turn dedup (DeduplicatingPruningSearchAgent) + 3. Spillage/rejection (rejection_budget cutoff) + 4. Tool-output clamping (max_tokens override when budget tight) + 5. Token budgeting (threshold -> prune/conclude + tool restriction, marker) + +The token counter used throughout counts occurrences of the sentinel "TOK", +so budgets are fully deterministic and independent of the (large) system prompt. +""" +from __future__ import annotations + +import json +import types + +import pytest + +from cosmos_retriever.inference.agent_loop import ( + _BudgetController, + _remove_chunks_from_text, + run_chat_search, + run_responses_search, +) + +TOK = "TOK" # sentinel word the fake counter counts + + +def counter(s: str) -> int: + return s.count(TOK) if isinstance(s, str) else 0 + + +# ───────────────────────── fakes ───────────────────────── + +class _Schema: + def __init__(self, name): + self.name = name + + +class FakeMeta: + def __init__(self, returned_chunk_ids=None): + self.returned_chunk_ids = list(returned_chunk_ids or []) + self.retrieval_s = 0.0 + self.rerank_s = 0.0 + + +class FakeTool: + """Records (params, overrides) per call; returns scripted outputs.""" + + def __init__(self, name, outputs=None, returned_chunk_ids=None): + self.tool_schema = _Schema(name) + self._outputs = list(outputs or []) + self._rcids = returned_chunk_ids + self.calls = [] # list of (params, overrides) + + def get_format(self, provider): + return {"type": "function", "name": self.tool_schema.name} + + def __call__(self, params, overrides=None): + i = len(self.calls) + self.calls.append((params, overrides)) + out = self._outputs[min(i, len(self._outputs) - 1)] if self._outputs else "ok" + meta = FakeMeta(self._rcids) if self.tool_schema.name == "search_corpus" else None + return out, meta + + +class FakeToolSet: + def __init__(self, tools): + self.tools = {t.tool_schema.name: t for t in tools} + + def get_tool(self, name): + return self.tools.get(name) + + +class FC: + """Fake /responses function_call output item.""" + + type = "function_call" + + def __init__(self, name, args, call_id): + self.name = name + self.arguments = json.dumps(args) + self.call_id = call_id + + +def _usage(): + return types.SimpleNamespace( + input_tokens=1, output_tokens=1, total_tokens=2, output_tokens_details=None + ) + + +class Resp: + def __init__(self, output, output_text="", rid="r"): + self.output = output + self.output_text = output_text + self.id = rid + self.usage = _usage() + + +class _RespCreate: + def __init__(self, outer): + self.outer = outer + + def create(self, **kwargs): + self.outer.calls.append(kwargs) + return self.outer.queue.pop(0) + + +class FakeResponsesClient: + def __init__(self, queue): + self.queue = list(queue) + self.calls = [] # kwargs of each create() + self.responses = _RespCreate(self) + + +class _Msg: + def __init__(self, content, tool_calls=None): + self.content = content + self.tool_calls = tool_calls + + +class _ChatTC: + def __init__(self, name, args, tid): + self.id = tid + self.function = types.SimpleNamespace(name=name, arguments=json.dumps(args)) + + +class _ChatResp: + def __init__(self, msg): + self.choices = [types.SimpleNamespace(message=msg)] + self.usage = types.SimpleNamespace(prompt_tokens=1, completion_tokens=1, total_tokens=2) + + +class _ChatCreate: + def __init__(self, outer): + self.outer = outer + + def create(self, **kwargs): + self.outer.calls.append(kwargs) + return self.outer.queue.pop(0) + + +class FakeChatClient: + def __init__(self, queue): + self.queue = list(queue) + self.calls = [] + self.completions = _ChatCreate(self) + + @property + def chat(self): + return types.SimpleNamespace(completions=self.completions) + + +def _base_tools(): + return [ + FakeTool("search_corpus", outputs=["SEARCH-OUT"], returned_chunk_ids=["c1", "c2"]), + FakeTool("grep_corpus"), + FakeTool("read_document"), + FakeTool("prune_chunks", outputs=["Pruned"]), + ] + + +# ═════════════════════════ 1. prune helper ═════════════════════════ + +def test_prune_removes_matching_block_keeps_others(): + text = "head\n# DOCUMENT ID: A \nbodyA\n# DOCUMENT ID: B \nbodyB\n\n[Token usage: 3/16]" + out = _remove_chunks_from_text(text, {"A"}) + assert "bodyA" not in out + assert "bodyB" in out + assert "[Token usage:" in out # marker preserved + + +def test_prune_multiple_blocks(): + text = "# DOCUMENT ID: A \nbodyA\n# DOCUMENT ID: B \nbodyB\n# DOCUMENT ID: C \nbodyC" + out = _remove_chunks_from_text(text, {"A", "C"}) + assert "bodyA" not in out and "bodyC" not in out and "bodyB" in out + + +def test_prune_noop_when_no_ids_or_no_matches(): + text = "# DOCUMENT ID: A \nbodyA" + assert _remove_chunks_from_text(text, set()) == text + assert _remove_chunks_from_text("plain text", {"A"}) == "plain text" + + +def test_prune_collapses_blank_lines(): + text = "# DOCUMENT ID: A \nbodyA\n# DOCUMENT ID: B \nbodyB" + out = _remove_chunks_from_text(text, {"A"}) + assert "\n\n\n" not in out + + +# ═════════════════════════ 2. controller units ═════════════════════════ + +def test_rejection_budget_formula(): + c = _BudgetController(text_token_counter=counter, threshold_budget=16384, token_budget=32268) + assert c.rejection_budget == 16384 + int((32268 - 16384) * 0.5) == 24326 + + +def test_dedup_records_and_exposes_ignore_ids(): + c = _BudgetController(text_token_counter=counter, threshold_budget=100, token_budget=200) + assert c.search_overrides() == {"ignore_ids": []} + c.record_search(["a_1", "b_2"], "q1") + assert set(c.search_overrides()["ignore_ids"]) == {"a_1", "b_2"} + # first query wins for read reranking + c.record_search(["a_1"], "q2") + assert c.read_overrides({"doc_id": "a_1"}) == {"query": "q1"} + assert c.read_overrides({"doc_id": "unknown"}) == {} + + +def test_prune_state_removes_recorded_chunks(): + c = _BudgetController(text_token_counter=counter, threshold_budget=100, token_budget=200) + c.record_prune(["X"]) + assert c.prune_text("# DOCUMENT ID: X \nbody") == "# DOCUMENT ID: X [pruned]" + + +def test_reject_and_clamp_and_marker(): + c = _BudgetController(text_token_counter=counter, threshold_budget=3, token_budget=100) + # reject non-prune past rejection budget (=51); allow prune + assert c.should_reject("search_corpus", 60) is True + assert c.should_reject("prune_chunks", 60) is False + assert c.should_reject("search_corpus", 10) is False # under rejection + # clamp when remaining < tool_output_budget(4096) + assert c.tool_max_tokens("search_corpus", 99) == max(512, (100 - 99) // 2) + # ample budget -> no clamp (needs remaining >= 4096, so use a real-sized budget) + big = _BudgetController(text_token_counter=counter, threshold_budget=16384, token_budget=32268) + assert big.tool_max_tokens("search_corpus", 10) is None # ample + assert big.tool_max_tokens("grep_corpus", 32000) is None # not a clamped tool + # marker + assert c.annotate("x", 20) == "x\n\n[Token usage: 20/3]" + assert c.over_threshold(4) and not c.over_threshold(2) + assert c.over_token_budget(101) and not c.over_token_budget(100) + + +# ═════════════════════════ 3. responses loop integration ═════════════════════════ + +def test_responses_loop_cross_turn_dedup(): + """2nd search must receive ignore_ids from the 1st search's returned ids.""" + tools = _base_tools() + ts = FakeToolSet(tools) + client = FakeResponsesClient([ + Resp([FC("search_corpus", {"query": "q1"}, "1")]), # turn 1 + Resp([FC("search_corpus", {"query": "q2"}, "2")]), # turn 2 + Resp([], output_text=""), # turn 3: conclude + ]) + run_responses_search( + toolset=ts, client=client, model="m", query="Q", + max_turns=10, text_token_counter=counter, + threshold_budget=10_000, token_budget=20_000, + ) + search = ts.get_tool("search_corpus") + assert len(search.calls) == 2 + # 2nd call overrides carry ignore_ids from 1st (c1,c2) + _, ov2 = search.calls[1] + assert set(ov2["ignore_ids"]) == {"c1", "c2"} + + +def test_responses_loop_real_prune_shrinks_transcript(): + """After prune_chunks, the pruned block is gone from the resent transcript.""" + tools = [ + FakeTool("search_corpus", + outputs=[f"\n# DOCUMENT ID: c1 \n{TOK} {TOK} bodyone"], + returned_chunk_ids=["c1"]), + FakeTool("prune_chunks", outputs=["Pruned"]), + ] + ts = FakeToolSet(tools) + client = FakeResponsesClient([ + Resp([FC("search_corpus", {"query": "q1"}, "1")]), # turn 1 + Resp([FC("prune_chunks", {"chunk_ids": ["c1"]}, "2")]), # turn 2 prune c1 + Resp([FC("search_corpus", {"query": "q3"}, "3")]), # turn 3 (triggers re-render) + Resp([], output_text="done"), # turn 4 conclude + ]) + run_responses_search( + toolset=ts, client=client, model="m", query="Q", + max_turns=10, text_token_counter=counter, + threshold_budget=10_000, token_budget=20_000, + ) + # turn 3's create() input must NOT contain the pruned body + turn3_input = client.calls[2]["input"] + blob = json.dumps(turn3_input) + assert "bodyone" not in blob # real pruning happened + + +def test_responses_loop_threshold_restricts_to_prune_and_injects_message(): + tools = [ + FakeTool("search_corpus", + outputs=[f"# DOCUMENT ID: c1 \n{TOK} {TOK} {TOK} {TOK} {TOK}"], # 5 TOK + returned_chunk_ids=["c1"]), + FakeTool("prune_chunks", outputs=["Pruned"]), + ] + ts = FakeToolSet(tools) + client = FakeResponsesClient([ + Resp([FC("search_corpus", {"query": "q1"}, "1")]), # turn 1 -> 5 TOK in transcript + Resp([FC("prune_chunks", {"chunk_ids": ["c1"]}, "2")]), # turn 2 (restricted) + Resp([], output_text="done"), + ]) + run_responses_search( + toolset=ts, client=client, model="m", query="Q", + max_turns=10, text_token_counter=counter, + threshold_budget=3, token_budget=100, # 5 TOK > threshold 3 + ) + # turn 2's create(): tools restricted to prune only + budget message injected + turn2 = client.calls[1] + tool_names = {t["name"] for t in turn2["tools"]} + assert tool_names == {"prune_chunks"} + assert "OVER BUDGET" in json.dumps(turn2["input"]) + + +def test_responses_loop_rejection_blocks_non_prune(): + """Past rejection budget, a non-prune tool call is not executed; model gets the error.""" + tools = [ + FakeTool("search_corpus", + outputs=[" ".join([TOK] * 60)], # 60 TOK -> over rejection (51) + returned_chunk_ids=["c1"]), + FakeTool("prune_chunks", outputs=["Pruned"]), + ] + ts = FakeToolSet(tools) + client = FakeResponsesClient([ + Resp([FC("search_corpus", {"query": "q1"}, "1")]), # turn 1 -> 60 TOK + Resp([FC("search_corpus", {"query": "q2"}, "2")]), # turn 2 non-prune -> rejected + Resp([], output_text="done"), + ]) + run_responses_search( + toolset=ts, client=client, model="m", query="Q", + max_turns=10, text_token_counter=counter, + threshold_budget=3, token_budget=100, # rejection = 3 + 0.5*97 = 51 + ) + search = ts.get_tool("search_corpus") + assert len(search.calls) == 1 # 2nd search NOT executed (rejected) + # rejection message surfaced to the model on turn 3 + assert "Token budget exceeded" in json.dumps(client.calls[2]["input"]) + + +def test_responses_loop_tool_output_clamp(): + tools = [ + FakeTool("search_corpus", outputs=[f"# DOCUMENT ID: c1 \n{TOK} {TOK}"], + returned_chunk_ids=["c1"]), + ] + ts = FakeToolSet(tools) + client = FakeResponsesClient([ + Resp([FC("search_corpus", {"query": "q1"}, "1")]), # turn1 -> 2 TOK + Resp([FC("search_corpus", {"query": "q2"}, "2")]), # turn2 -> clamp (remaining tight) + Resp([], output_text="done"), + ]) + run_responses_search( + toolset=ts, client=client, model="m", query="Q", + max_turns=10, text_token_counter=counter, + threshold_budget=1000, token_budget=4, # remaining = 4-2 = 2 < 4096 -> clamp + ) + search = ts.get_tool("search_corpus") + _, ov2 = search.calls[1] + assert ov2 is not None and "max_tokens" in ov2 and ov2["max_tokens"] >= 512 + + +# ═════════════════════════ 4. chat loop integration ═════════════════════════ + +def test_chat_loop_dedup_and_prune(): + tools = [ + FakeTool("search_corpus", + outputs=[f"# DOCUMENT ID: c1 \n{TOK} bodyone"], + returned_chunk_ids=["c1"]), + FakeTool("prune_chunks", outputs=["Pruned"]), + ] + ts = FakeToolSet(tools) + client = FakeChatClient([ + _ChatResp(_Msg("", [_ChatTC("search_corpus", {"query": "q1"}, "t1")])), + _ChatResp(_Msg("", [_ChatTC("prune_chunks", {"chunk_ids": ["c1"]}, "t2")])), + _ChatResp(_Msg("", [_ChatTC("search_corpus", {"query": "q2"}, "t3")])), + _ChatResp(_Msg("", None)), + ]) + run_chat_search( + toolset=ts, client=client, model="m", query="Q", + max_turns=10, text_token_counter=counter, + threshold_budget=10_000, token_budget=20_000, + ) + search = ts.get_tool("search_corpus") + # dedup: 2nd search got ignore_ids from 1st + assert set(search.calls[1][1]["ignore_ids"]) == {"c1"} + # real prune: 3rd chat call's messages no longer contain the pruned body + assert "bodyone" not in json.dumps(client.calls[2]["messages"]) diff --git a/cosmos-retriever/tests/unit/test_token_count.py b/cosmos-retriever/tests/unit/test_token_count.py new file mode 100644 index 0000000..9752d96 --- /dev/null +++ b/cosmos-retriever/tests/unit/test_token_count.py @@ -0,0 +1,221 @@ +"""Token-count accuracy tests for the o200k_harmony counter. + +The retrieval service budgets/truncates using ``CosmosRetriever._text_token_counter``, +a one-line wrapper over ``tiktoken.get_encoding("o200k_harmony").encode``. These +tests treat tiktoken's o200k_harmony as the ground-truth tokenizer for +gpt-oss/Harmony and o200k-family models and verify: + + * exactness of the real counter against the encoder, + * cross-model agreement (o200k_harmony == o200k_base on real text) and the + divergence from the older cl100k_base tokenizer, + * accuracy of the crude ``len//4`` fallback estimate vs the real count, + * edge behaviour (empty, whitespace, Unicode/CJK/emoji, determinism, + monotonicity, sub-additivity, and special-token rejection). + +Ground-truth counts are pinned for tiktoken 0.13.0. + +Sample results (o200k_harmony token counts): + + "" 0 + "hello" 1 + "Hello, world!" 4 + 9-word English pangram 9 + "机器学习很有趣" (CJK) 5 + "I love 😀 pizza 🍕" 6 + +Cross-model drift on English prose, tokens relative to o200k_harmony (= 1.00), +from the comparison experiment in ``test_tokenizer_comparison.py``: + + o200k_base (GPT-4o) 1.00 + gpt-oss (Harmony) 1.00 + cl100k (GPT-4/3.5) 1.01 + Qwen3 1.01 + Claude 1.02 + Gemma (Gemini) 1.02 + Llama 3.1 1.06 + Mistral 1.17 + +Full panel and regeneration live in ``tests/end_to_end/test_tokenizer_comparison.py`` +(run ``python tests/end_to_end/tokenizer_panel.py``); see also ``tests/README.md``. +""" +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import tiktoken + +from cosmos_retriever.retriever import CosmosRetriever + +# ────────────────────────── encoders / real counter ─────────────────────── + +HARMONY = tiktoken.get_encoding("o200k_harmony") +O200K_BASE = tiktoken.get_encoding("o200k_base") +CL100K = tiktoken.get_encoding("cl100k_base") + + +def count(text: str) -> int: + """Invoke the real CosmosRetriever._text_token_counter with a fake self.""" + fake = SimpleNamespace(_tiktoken=HARMONY) + return CosmosRetriever._text_token_counter(fake, text) + + +# Pinned ground truth (tiktoken 0.13.0, o200k_harmony). +_ANCHORS = { + "": 0, + "hello": 1, + "Hello, world!": 4, + "The quick brown fox jumps over the lazy dog": 9, + "café résumé naïve": 5, + "I love 😀 pizza 🍕": 6, + "机器学习很有趣": 5, + " \n\n\t ": 2, + "ha" * 100: 51, +} + + +# ═══════════════════════ encoding identity ════════════════════════════════ + + +def test_harmony_encoding_name_and_superset_of_base() -> None: + assert HARMONY.name == "o200k_harmony" + # harmony is o200k_base plus Harmony special tokens -> strictly larger vocab + assert HARMONY.n_vocab > O200K_BASE.n_vocab + + +# ═══════════════════════ counter exactness ════════════════════════════════ + + +@pytest.mark.parametrize("text,expected", list(_ANCHORS.items())) +def test_counter_matches_pinned_ground_truth(text: str, expected: int) -> None: + assert count(text) == expected + + +@pytest.mark.parametrize("text", list(_ANCHORS)) +def test_counter_equals_encoder_length(text: str) -> None: + assert count(text) == len(HARMONY.encode(text)) + + +def test_counter_returns_nonnegative_int() -> None: + for text in _ANCHORS: + n = count(text) + assert isinstance(n, int) and n >= 0 + + +def test_empty_string_is_zero() -> None: + assert count("") == 0 + + +def test_counter_is_deterministic() -> None: + text = "The quick brown fox jumps over the lazy dog" + assert count(text) == count(text) == 9 + + +def test_counter_is_monotonic_under_append() -> None: + base = "the quick brown fox" + assert count(base) <= count(base + " jumps over the lazy dog") + + +@pytest.mark.parametrize( + "a,b", + [ + ("hello", "world"), + ("machine", "learning"), + ("The quick brown", " fox jumps"), + ("café", "résumé"), + ], +) +def test_concatenation_is_subadditive(a: str, b: str) -> None: + # BPE only re-merges at the boundary, never splits existing tokens. + assert count(a + b) <= count(a) + count(b) + + +# ═══════════════════ cross-model agreement / divergence ═══════════════════ + + +@pytest.mark.parametrize("text", list(_ANCHORS)) +def test_harmony_equals_o200k_base_on_real_text(text: str) -> None: + # For non-special text, harmony reduces to o200k_base -> identical counts. + assert len(HARMONY.encode(text)) == len(O200K_BASE.encode(text)) + + +@pytest.mark.parametrize("text", list(_ANCHORS)) +def test_harmony_never_exceeds_cl100k(text: str) -> None: + # The newer o200k tokenizer is at least as efficient as legacy cl100k. + assert len(HARMONY.encode(text)) <= len(CL100K.encode(text)) + + +@pytest.mark.parametrize("text", ["café résumé naïve", "I love 😀 pizza 🍕", "机器学习很有趣"]) +def test_harmony_strictly_more_efficient_than_cl100k_on_non_ascii(text: str) -> None: + assert len(HARMONY.encode(text)) < len(CL100K.encode(text)) + + +def test_cjk_divergence_is_large() -> None: + # 机器学习很有趣: harmony 5 vs cl100k 10 -> harmony halves the CJK cost. + assert len(HARMONY.encode("机器学习很有趣")) == 5 + assert len(CL100K.encode("机器学习很有趣")) == 10 + + +# ═══════════════════ crude len//4 estimate vs real count ══════════════════ + + +def _crude(s: str) -> int: + # Mirrors the agent_loop ContextTracker default fallback counter. + return len(s) // 4 + + +def test_crude_estimate_close_for_english_prose() -> None: + prose = ( + "Retrieval augmented generation combines a search index with a language " + "model so answers stay grounded in the underlying corpus documents." + ) + real = count(prose) + est = _crude(prose) + # For English prose the char/4 heuristic stays within ~40% of the real count. + assert 0.6 <= est / real <= 1.6 + + +def test_crude_estimate_underestimates_dense_scripts() -> None: + cjk = "机器学习很有趣" + # len//4 = 1 but the real cost is 5 -> the char heuristic massively underestimates. + assert _crude(cjk) < count(cjk) + assert count(cjk) >= 3 * max(_crude(cjk), 1) + + +def test_real_count_beats_crude_as_accuracy_reference() -> None: + # The tiktoken count equals o200k_base (exact for the target models); the crude + # estimate does not — demonstrating why the service uses tiktoken, not len//4. + for text in ("Hello, world!", "机器学习很有趣", "I love 😀 pizza 🍕"): + assert count(text) == len(O200K_BASE.encode(text)) + # crude only coincidentally matches; assert it diverges on at least the CJK case + assert _crude("机器学习很有趣") != count("机器学习很有趣") + + +# ═══════════════════════ special-token handling ═══════════════════════════ + + +@pytest.mark.parametrize("marker", ["<|end|>", "<|start|>", "<|message|>", "<|return|>"]) +def test_special_token_text_is_rejected(marker: str) -> None: + # Default encode disallows special tokens; document text containing Harmony + # markers would raise rather than silently mis-count. + with pytest.raises(ValueError): + count(marker) + + +def test_special_token_counts_as_single_when_explicitly_allowed() -> None: + # Sanity check on the tokenizer itself: the marker is one special token. + assert HARMONY.encode("<|end|>", allowed_special="all") == [200007] + + +# ═══════════════════════ scaling / whitespace ═════════════════════════════ + + +def test_whitespace_and_newlines_are_counted() -> None: + assert count(" \n\n\t ") == 2 + + +def test_repeated_token_scales_sublinearly_in_chars() -> None: + text = "ha" * 100 # 200 chars + n = count(text) + assert n == 51 # far below 200 chars thanks to BPE merges + assert n < len(text) // 2 diff --git a/cosmos-retriever/tests/unit/test_tools.py b/cosmos-retriever/tests/unit/test_tools.py new file mode 100644 index 0000000..52779d9 --- /dev/null +++ b/cosmos-retriever/tests/unit/test_tools.py @@ -0,0 +1,815 @@ +"""Exhaustive tests for the agent tools module (`cosmos_retriever.tools`). + +Covers, with fakes and no network / Cosmos / OpenAI access: + + 1. ToolSchema — provider wire formats + dispatch + required default + 2. Static schema consts — search / read / grep / multi / prune shapes + 3. Tool / SerializedTool — get_format, __repr__, abstractness, placeholder raise + 4. Dynamic schema builders— _search_schema_for / _grep_schema_for field matrices + 5. SearchCorpusTool — validation, coercion, errors, rerank reorder, limits + 6. GrepCorpusTool — validation, field errors, empty, regex fallback, tokens + 7. ReadDocumentTool — ctor invariant, id aliases, item-doc / rerank / truncate + 8. PruneChunksTool — validation + fixed output + 9. _sanitize_query_value — vector collapse, recursion, long-string truncation + 10. _SELECT_RE — accepts SELECT variants, rejects writes + 11. RunQueryTool — validation, guards, happy path, cap, truncate, error + 12. MultiToolUseTool — dispatch, unknown-tool raise, json encoding + 13. UserTextTool — always raises + 14. ToolSet / build — add/remove/get/formats/repr + build wiring + +Fakes: FakeSchema, FakeRetriever, FakeReranker, FakeCosmosClient. Real Pydantic +models (RetrievedItem, NormalizedDocument, RerankResult) are used directly. +`tools.logger` is replaced with a recording stub; `tools.time.perf_counter` is +patched where deterministic timing is asserted. +""" +from __future__ import annotations + +import json + +import pytest + +from cosmos_retriever import tools +from cosmos_retriever.rerank import RerankResult +from cosmos_retriever.retrieval.errors import ( + UnknownField, + UnsupportedRetrievalCapability, +) +from cosmos_retriever.retrieval.models import NormalizedDocument, RetrievedItem +from cosmos_retriever.tools import ( + _SELECT_RE, + GREP_CORPUS_SCHEMA, + MULTI_TOOL_USE_SCHEMA, + PRUNE_CHUNKS_SCHEMA, + READ_DOCUMENT_SCHEMA, + RUN_QUERY_SCHEMA, + SEARCH_CORPUS_SCHEMA, + GrepCorpusTool, + MultiToolUseTool, + PruneChunksTool, + ReadDocumentTool, + RunQueryTool, + SearchCorpusTool, + SearchCorpusToolCallMetadata, + SerializedTool, + Tool, + ToolSchema, + ToolSet, + UserTextTool, + _grep_schema_for, + _sanitize_query_value, + _search_schema_for, +) +from cosmos_retriever.utils import ProviderFormat + +# ────────────────────────────── shared fakes ────────────────────────────── + + +class _RecordLogger: + """structlog-like stub that records (event, kwargs) and binds to itself.""" + + def __init__(self) -> None: + self.events: list[tuple[str, str, dict]] = [] + + def bind(self, **kwargs): + return self + + def _log(self, level, event, **kwargs): + self.events.append((level, event, kwargs)) + + def info(self, event, **kwargs): + self._log("info", event, **kwargs) + + def warning(self, event, **kwargs): + self._log("warning", event, **kwargs) + + def error(self, event, **kwargs): + self._log("error", event, **kwargs) + + +class FakeSchema: + """Stands in for CorpusSchema; only the surface tools.py touches.""" + + def __init__( + self, + text_fields: list[str] | None = None, + vector_fields: list[str] | None = None, + item_document_mode: bool = False, + summary: str = "SCHEMA-SUMMARY", + ) -> None: + self._text = text_fields if text_fields is not None else ["body"] + self._vector = vector_fields if vector_fields is not None else [] + self.is_item_document_mode = item_document_mode + self._summary = summary + + def text_field_map(self) -> dict[str, object]: + return {name: object() for name in self._text} + + def vector_field_map(self) -> dict[str, object]: + return {name: object() for name in self._vector} + + def agent_field_summary(self) -> str: + return self._summary + + +class FakeRetriever: + """Stands in for CorpusRetriever; records requests, returns canned data.""" + + def __init__( + self, + schema: FakeSchema | None = None, + items: list[RetrievedItem] | None = None, + candidates: list[RetrievedItem] | None = None, + document: NormalizedDocument | None = None, + search_exc: Exception | None = None, + grep_exc: Exception | None = None, + ) -> None: + self.schema = schema or FakeSchema() + self._items = items or [] + self._candidates = candidates if candidates is not None else [] + self._document = document or NormalizedDocument(chunk_texts=[]) + self._search_exc = search_exc + self._grep_exc = grep_exc + self.search_requests: list = [] + self.grep_requests: list = [] + self.read_requests: list = [] + + def search(self, request): + self.search_requests.append(request) + if self._search_exc is not None: + raise self._search_exc + return self._items + + def grep_candidates(self, request): + self.grep_requests.append(request) + if self._grep_exc is not None: + raise self._grep_exc + return self._candidates + + def read_document(self, request): + self.read_requests.append(request) + return self._document + + +class FakeReranker: + """Callable stand-in for Reranker; emits RerankResults in a fixed order.""" + + def __init__(self, order: list[int], tokens: list[int | None] | None = None) -> None: + self.order = order + self.tokens = tokens + self.calls: list[tuple] = [] + + def __call__(self, query, documents, max_tokens=None): + self.calls.append((query, list(documents), max_tokens)) + out: list[RerankResult] = [] + for i, idx in enumerate(self.order): + tok = self.tokens[i] if self.tokens is not None else None + out.append( + RerankResult( + document=documents[idx], + score=1.0 - 0.1 * i, + original_index=idx, + tokens=tok, + ) + ) + return out + + +class FakeContainer: + def __init__(self, rows=None, exc: Exception | None = None) -> None: + self.rows = rows or [] + self.exc = exc + self.received: tuple | None = None + + def query_items(self, query, enable_cross_partition_query, max_item_count): + self.received = (query, enable_cross_partition_query, max_item_count) + if self.exc is not None: + raise self.exc + yield from self.rows + + +class FakeDatabase: + def __init__(self, container: FakeContainer) -> None: + self._container = container + self.requested_container: str | None = None + + def get_container_client(self, name): + self.requested_container = name + return self._container + + +class FakeCosmosClient: + def __init__(self, container: FakeContainer) -> None: + self._database = FakeDatabase(container) + self.requested_database: str | None = None + + def get_database_client(self, name): + self.requested_database = name + return self._database + + +def _item(item_id: str, text: str) -> RetrievedItem: + return RetrievedItem(item_id=item_id, text=text) + + +@pytest.fixture(autouse=True) +def _silence_logger(monkeypatch): + """Replace module logger with a recorder for every test.""" + rec = _RecordLogger() + monkeypatch.setattr(tools, "logger", rec) + return rec + + +# ═══════════════════════ 1. ToolSchema wire formats ═══════════════════════ + + +def _schema() -> ToolSchema: + return ToolSchema( + name="t", + description="d", + parameters={"q": {"type": "string"}}, + required=["q"], + ) + + +def test_toolschema_required_defaults_to_empty_list() -> None: + s = ToolSchema(name="t", description="d", parameters={}) + assert s.required == [] + + +def test_openai_format_is_flat_function() -> None: + out = _schema()._to_openai_format() + assert out == { + "type": "function", + "name": "t", + "description": "d", + "parameters": { + "type": "object", + "properties": {"q": {"type": "string"}}, + "required": ["q"], + }, + } + + +def test_openai_harmony_format_nests_under_function() -> None: + out = _schema()._to_openai_harmony_format() + assert out["type"] == "function" + assert out["function"]["name"] == "t" + assert out["function"]["parameters"]["required"] == ["q"] + + +def test_anthropic_format_uses_input_schema() -> None: + out = _schema()._to_anthropic_format() + assert set(out) == {"name", "description", "input_schema"} + assert "parameters" not in out + assert out["input_schema"]["properties"] == {"q": {"type": "string"}} + + +def test_to_provider_format_dispatch() -> None: + s = _schema() + assert s.to_provider_format(ProviderFormat.OPENAI) == s._to_openai_format() + assert s.to_provider_format(ProviderFormat.OPENAI_HARMONY) == s._to_openai_harmony_format() + assert s.to_provider_format(ProviderFormat.ANTHROPIC) == s._to_anthropic_format() + + +def test_to_provider_format_unsupported_raises() -> None: + with pytest.raises(ValueError, match="Unsupported provider format"): + _schema().to_provider_format(object()) # type: ignore[arg-type] + + +# ═══════════════════════ 2. static schema constants ═══════════════════════ + + +def test_static_schema_names_and_required() -> None: + assert SEARCH_CORPUS_SCHEMA.name == "search_corpus" + assert SEARCH_CORPUS_SCHEMA.required == ["query"] + assert READ_DOCUMENT_SCHEMA.name == "read_document" + assert READ_DOCUMENT_SCHEMA.required == ["doc_id"] + assert GREP_CORPUS_SCHEMA.name == "grep_corpus" + assert GREP_CORPUS_SCHEMA.required == ["pattern"] + assert PRUNE_CHUNKS_SCHEMA.name == "prune_chunks" + assert PRUNE_CHUNKS_SCHEMA.required == ["chunk_ids"] + assert RUN_QUERY_SCHEMA.name == "execute_query" + assert RUN_QUERY_SCHEMA.required == ["query"] + + +def test_multi_tool_use_schema_item_shape() -> None: + items = MULTI_TOOL_USE_SCHEMA.parameters["tool_calls"]["items"] + assert items["required"] == ["tool_name", "parameters"] + assert set(items["properties"]) == {"tool_name", "parameters"} + + +# ═══════════════════════ 3. Tool / SerializedTool ═════════════════════════ + + +def test_tool_is_abstract() -> None: + with pytest.raises(TypeError): + Tool(tool_schema=_schema()) # type: ignore[abstract] + + +def test_tool_get_format_and_repr() -> None: + tool = PruneChunksTool() + assert tool.__repr__() == "Tool(name='prune_chunks')" + fmt = tool.get_format(ProviderFormat.OPENAI) + assert fmt["name"] == "prune_chunks" + + +def test_serialized_tool_cannot_execute() -> None: + tool = SerializedTool(tool_schema=_schema()) + with pytest.raises(NotImplementedError): + tool({"q": "x"}) + + +# ═══════════════════ 4. dynamic schema builders ═══════════════════════════ + + +def test_search_schema_single_text_no_vector() -> None: + s = _search_schema_for(FakeSchema(text_fields=["body"], vector_fields=[])) + assert s.required == ["query"] + assert set(s.parameters) == {"query"} + assert "SCHEMA-SUMMARY" in s.description + + +def test_search_schema_multi_text_adds_fields_and_requires_it() -> None: + s = _search_schema_for(FakeSchema(text_fields=["a", "b"], vector_fields=[])) + assert "fields" in s.parameters + assert s.required == ["query", "fields"] + assert "mode" not in s.parameters # no vector -> no mode + + +def test_search_schema_multi_vector_adds_vector_field_and_mode() -> None: + s = _search_schema_for(FakeSchema(text_fields=["a"], vector_fields=["v1", "v2"])) + assert "vector_field" in s.parameters + assert "mode" in s.parameters # text and vector both present + assert "fields" not in s.parameters # single text + assert s.required == ["query"] + + +def test_search_schema_multi_text_and_vector_full_params() -> None: + s = _search_schema_for(FakeSchema(text_fields=["a", "b"], vector_fields=["v1", "v2"])) + assert {"query", "fields", "vector_field", "mode"} <= set(s.parameters) + assert s.required == ["query", "fields"] + + +def test_search_schema_single_vector_no_vector_field_param() -> None: + s = _search_schema_for(FakeSchema(text_fields=["a"], vector_fields=["only"])) + assert "vector_field" not in s.parameters # needs >1 vector + assert "mode" in s.parameters # still text+vector present + + +def test_grep_schema_single_text() -> None: + s = _grep_schema_for(FakeSchema(text_fields=["body"])) + assert s.required == ["pattern"] + assert "field" not in s.parameters + + +def test_grep_schema_multi_text_requires_field() -> None: + s = _grep_schema_for(FakeSchema(text_fields=["a", "b"])) + assert "field" in s.parameters + assert s.required == ["pattern", "field"] + + +# ═══════════════════════ 5. SearchCorpusTool ══════════════════════════════ + + +def _search_tool(retriever, reranker=None, search_limit=50, display_limit=10): + return SearchCorpusTool( + retriever=retriever, + reranker=reranker, + search_limit=search_limit, + display_limit=display_limit, + ) + + +def test_search_invalid_params_raises() -> None: + tool = _search_tool(FakeRetriever()) + with pytest.raises(ValueError, match="Invalid params type"): + tool({"not_query": 1}) + with pytest.raises(ValueError): + tool(["query"]) # type: ignore[arg-type] + + +def test_search_passes_request_fields_and_coerces() -> None: + retr = FakeRetriever(items=[_item("a", "ta")]) + tool = _search_tool(retr, search_limit=7) + tool( + {"query": "Q", "fields": "title", "vector_field": "vec", "mode": "hybrid"}, + overrides={"ignore_ids": ["x", "y"]}, + ) + req = retr.search_requests[0] + assert req.query == "Q" + assert req.limit == 7 + assert req.text_fields == ["title"] # str coerced to list + assert req.vector_field == "vec" + assert req.mode == "hybrid" + assert req.ignored_item_ids == ["x", "y"] + + +def test_search_invalid_mode_falls_back_to_auto() -> None: + retr = FakeRetriever(items=[]) + tool = _search_tool(retr) + tool({"query": "Q", "mode": "nonsense"}) + assert retr.search_requests[0].mode == "auto" + + +def test_search_field_error_returns_message_not_raise() -> None: + retr = FakeRetriever(search_exc=UnknownField("bad field")) + tool = _search_tool(retr) + text, meta = tool({"query": "Q"}) + assert "Search field/mode error: bad field" in text + assert isinstance(meta, SearchCorpusToolCallMetadata) + assert meta.returned_chunk_ids == [] + + +def test_search_capability_error_returns_message() -> None: + retr = FakeRetriever(search_exc=UnsupportedRetrievalCapability("nope")) + text, meta = _search_tool(retr)({"query": "Q"}) + assert "Search field/mode error: nope" in text + assert meta.returned_chunk_ids == [] + + +def test_search_no_reranker_passthrough_and_metadata() -> None: + retr = FakeRetriever(items=[_item("a", "ta"), _item("b", "tb")]) + text, meta = _search_tool(retr)({"query": "Q"}) + assert meta.returned_chunk_ids == ["a", "b"] + assert meta.rerank_s == 0.0 + assert "ta" in text and "tb" in text + + +def test_search_empty_results() -> None: + retr = FakeRetriever(items=[]) + reranker = FakeReranker(order=[0]) + text, meta = _search_tool(retr, reranker=reranker)({"query": "Q"}) + assert meta.returned_chunk_ids == [] + assert reranker.calls == [] # reranker skipped when no ids + assert text == "No results found" + + +def test_search_with_reranker_reorders_by_original_index() -> None: + retr = FakeRetriever(items=[_item("a", "ta"), _item("b", "tb"), _item("c", "tc")]) + reranker = FakeReranker(order=[2, 0, 1], tokens=[5, 6, 7]) + text, meta = _search_tool(retr, reranker=reranker)({"query": "Q"}) + assert meta.returned_chunk_ids == ["c", "a", "b"] + assert reranker.calls[0][0] == "Q" + assert "(5 tokens)" in text # token counts propagated into formatting + + +def test_search_passes_max_tokens_override_to_reranker() -> None: + retr = FakeRetriever(items=[_item("a", "ta")]) + reranker = FakeReranker(order=[0]) + _search_tool(retr, reranker=reranker)({"query": "Q"}, overrides={"max_tokens": 123}) + assert reranker.calls[0][2] == 123 + + +def test_search_display_limit_truncates() -> None: + retr = FakeRetriever(items=[_item(str(i), f"t{i}") for i in range(5)]) + _, meta = _search_tool(retr, display_limit=2)({"query": "Q"}) + assert meta.returned_chunk_ids == ["0", "1"] + + +def test_search_timing_rounded_to_three_decimals(monkeypatch) -> None: + seq = iter([0.0, 0.512812, 1.0, 2.517001]) + monkeypatch.setattr(tools.time, "perf_counter", lambda: next(seq)) + retr = FakeRetriever(items=[_item("a", "ta")]) + reranker = FakeReranker(order=[0]) + _, meta = _search_tool(retr, reranker=reranker)({"query": "Q"}) + assert meta.retrieval_s == 0.513 + assert meta.rerank_s == 1.517 + + +# ═══════════════════════ 6. GrepCorpusTool ════════════════════════════════ + + +def test_grep_invalid_params_raises() -> None: + with pytest.raises(ValueError, match="Invalid params type"): + GrepCorpusTool(FakeRetriever())({"nope": 1}) + + +def test_grep_field_error_returns_message() -> None: + retr = FakeRetriever(grep_exc=UnknownField("bad")) + text, meta = GrepCorpusTool(retr)({"pattern": "x"}) + assert "Grep field error: bad" in text + assert meta.returned_chunk_ids == [] + + +def test_grep_no_candidates_returns_no_results() -> None: + retr = FakeRetriever(candidates=[]) + text, meta = GrepCorpusTool(retr)({"pattern": "x"}) + assert text == "No results found" + assert meta.returned_chunk_ids == [] + + +def test_grep_filters_by_case_insensitive_regex() -> None: + retr = FakeRetriever( + candidates=[_item("a", "has FOO here"), _item("b", "no match"), _item("c", "foobar")] + ) + _, meta = GrepCorpusTool(retr)({"pattern": "foo"}) + assert meta.returned_chunk_ids == ["a", "c"] + + +def test_grep_invalid_regex_falls_back_to_candidates() -> None: + cands = [_item(str(i), f"t{i}") for i in range(7)] + retr = FakeRetriever(candidates=cands) + _, meta = GrepCorpusTool(retr)({"pattern": "["}) # invalid regex + assert meta.returned_chunk_ids == ["0", "1", "2", "3", "4"] # first 5 + + +def test_grep_token_counter_annotates() -> None: + retr = FakeRetriever(candidates=[_item("a", "abc")]) + text, _ = GrepCorpusTool(retr, token_counter=len)({"pattern": "abc"}) + assert "(3 tokens)" in text + + +def test_grep_field_forwarded_to_request() -> None: + retr = FakeRetriever(candidates=[_item("a", "abc")]) + GrepCorpusTool(retr)({"pattern": "abc", "field": "body"}) + assert retr.grep_requests[0].text_field == "body" + + +# ═══════════════════════ 7. ReadDocumentTool ══════════════════════════════ + + +def test_read_ctor_max_tokens_requires_counter() -> None: + with pytest.raises(ValueError, match="token_counter is required"): + ReadDocumentTool(FakeRetriever(), max_tokens=10) + + +def test_read_invalid_params_raises() -> None: + with pytest.raises(ValueError, match="Invalid params type"): + ReadDocumentTool(FakeRetriever())({"nope": 1}) + + +def test_read_accepts_doc_id_and_id_aliases() -> None: + retr = FakeRetriever(document=NormalizedDocument(chunk_texts=["body"])) + ReadDocumentTool(retr)({"doc_id": "d1"}) + ReadDocumentTool(retr)({"id": "d2"}) + assert retr.read_requests[0].document_id == "d1" + assert retr.read_requests[1].document_id == "d2" + + +def test_read_item_document_mode_verbatim_no_counter() -> None: + retr = FakeRetriever( + schema=FakeSchema(item_document_mode=True), + document=NormalizedDocument(chunk_texts=["whole doc"]), + ) + text, meta = ReadDocumentTool(retr)({"doc_id": "d"}) + assert text == "whole doc" + assert meta is None + + +def test_read_item_document_mode_with_counter_header() -> None: + retr = FakeRetriever( + schema=FakeSchema(item_document_mode=True), + document=NormalizedDocument(chunk_texts=["abcde"]), + ) + text, _ = ReadDocumentTool(retr, token_counter=len)({"doc_id": "d"}) + assert text == "# Document (5 tokens)\nabcde" + + +def test_read_chunk_mode_rerank_filters_kept_chunks() -> None: + retr = FakeRetriever(document=NormalizedDocument(chunk_texts=["A", "B", "C"])) + reranker = FakeReranker(order=[0, 2]) # keep indices 0 and 2 + tool = ReadDocumentTool(retr, reranker=reranker, token_counter=len, max_tokens=100) + text, _ = tool({"doc_id": "d"}, overrides={"query": "Q", "max_tokens": 100}) + assert text.endswith("AC") # original order preserved, B dropped + + +def test_read_chunk_mode_token_truncation_without_reranker() -> None: + retr = FakeRetriever(document=NormalizedDocument(chunk_texts=["aa", "bb", "cc"])) + tool = ReadDocumentTool(retr, token_counter=len, max_tokens=4) + text, _ = tool({"doc_id": "d"}) + assert text == "# Document (4 tokens)\naabb" # cc dropped by budget + + +def test_read_chunk_mode_plain_no_counter() -> None: + retr = FakeRetriever(document=NormalizedDocument(chunk_texts=["x", "y"])) + text, meta = ReadDocumentTool(retr)({"doc_id": "d"}) + assert text == "xy" + assert meta is None + + +# ═══════════════════════ 8. PruneChunksTool ═══════════════════════════════ + + +def test_prune_invalid_params_raises() -> None: + with pytest.raises(ValueError, match="Invalid params type"): + PruneChunksTool()({"nope": 1}) + + +def test_prune_returns_fixed_output() -> None: + text, meta = PruneChunksTool()({"chunk_ids": ["a", "b"]}) + assert text == "Pruned" + assert meta is None + + +# ═══════════════════ 9. _sanitize_query_value ═════════════════════════════ + + +def test_sanitize_collapses_numeric_vector() -> None: + assert _sanitize_query_value(list(range(40))) == "" + + +def test_sanitize_short_list_recursed_not_collapsed() -> None: + assert _sanitize_query_value([1, 2, 3]) == [1, 2, 3] + + +def test_sanitize_non_numeric_long_list_recursed() -> None: + val = ["s"] * 40 + assert _sanitize_query_value(val) == ["s"] * 40 # not all-numeric -> not collapsed + + +def test_sanitize_nested_dict_and_long_string() -> None: + long = "z" * 5000 + out = _sanitize_query_value({"k": long, "n": {"vec": list(range(50))}}) + assert out["k"].endswith("\u2026") and len(out["k"]) == 4001 + assert out["n"]["vec"] == "" + + +def test_sanitize_scalar_passthrough() -> None: + assert _sanitize_query_value(7) == 7 + assert _sanitize_query_value("short") == "short" + + +# ═══════════════════════ 10. _SELECT_RE ═══════════════════════════════════ + + +@pytest.mark.parametrize( + "q", + ["select * from c", " SELECT c.id", "(select 1)", "(( select x", "\n\tSELECT a"], +) +def test_select_re_accepts(q: str) -> None: + assert _SELECT_RE.match(q) + + +@pytest.mark.parametrize( + "q", + ["insert into c", "update c set x=1", "delete from c", "drop table c", "with t as ()"], +) +def test_select_re_rejects(q: str) -> None: + assert _SELECT_RE.match(q) is None + + +# ═══════════════════════ 11. RunQueryTool ═════════════════════════════════ + + +def _run_tool(container, **kw): + return RunQueryTool( + client=FakeCosmosClient(container), + default_database=kw.pop("db", "corpusdb"), + default_container=kw.pop("cont", "corpuscont"), + **kw, + ) + + +def test_run_invalid_params_raises() -> None: + with pytest.raises(ValueError, match="Invalid params type"): + _run_tool(FakeContainer())({"nope": 1}) + + +def test_run_empty_query_error() -> None: + text, _ = _run_tool(FakeContainer())({"query": " "}) + assert "must be a non-empty SELECT query" in text + + +def test_run_non_select_rejected() -> None: + text, _ = _run_tool(FakeContainer())({"query": "DELETE FROM c"}) + assert "only read-only SELECT queries are allowed" in text + + +def test_run_missing_db_or_container_error() -> None: + tool = RunQueryTool(client=FakeCosmosClient(FakeContainer())) + text, _ = tool({"query": "SELECT * FROM c"}) + assert "specify both 'database' and 'container'" in text + + +def test_run_happy_path_header_and_body() -> None: + container = FakeContainer(rows=[{"id": 1}, {"id": 2}]) + text, _ = _run_tool(container)({"query": "SELECT * FROM c"}) + header, body = text.split("\n", 1) + assert header == "# execute_query: 2 row(s) from corpusdb/corpuscont" + assert json.loads(body) == [{"id": 1}, {"id": 2}] + + +def test_run_caps_rows_and_marks_capped() -> None: + container = FakeContainer(rows=[{"id": i} for i in range(25)]) + text, _ = _run_tool(container, max_rows=3)({"query": "SELECT * FROM c"}) + header = text.split("\n", 1)[0] + assert "3 row(s)" in header and "capped at 3" in header + + +def test_run_truncates_large_body() -> None: + container = FakeContainer(rows=[{"blob": "x" * 500}]) + text, _ = _run_tool(container, max_chars=50)({"query": "SELECT * FROM c"}) + header, body = text.split("\n", 1) + assert "[output truncated]" in header + assert body.endswith("\u2026") + + +def test_run_client_exception_returned_as_text() -> None: + container = FakeContainer(exc=ValueError("boom")) + text, _ = _run_tool(container)({"query": "SELECT * FROM c"}) + assert text == "execute_query error: ValueError: boom" + + +def test_run_explicit_db_container_override_defaults() -> None: + container = FakeContainer(rows=[]) + tool = _run_tool(container) + tool({"query": "SELECT * FROM c", "database": "d2", "container": "c2"}) + assert tool._client.requested_database == "d2" + + +# ═══════════════════════ 12. MultiToolUseTool ═════════════════════════════ + + +class _EchoTool(Tool): + def __call__(self, params, overrides=None): + return f"echo:{params.get('v')}", None + + +def _echo_toolset() -> ToolSet: + ts = ToolSet() + ts.add_tool(_EchoTool(tool_schema=ToolSchema(name="echo", description="d", parameters={}))) + return ts + + +def test_multi_tool_dispatches_and_json_encodes() -> None: + tool = MultiToolUseTool(_echo_toolset()) + text, meta = tool( + {"tool_calls": [{"tool_name": "echo", "parameters": {"v": 1}}, + {"tool_name": "echo", "parameters": {"v": 2}}]} + ) + assert json.loads(text) == ["echo:1", "echo:2"] + assert meta is None + + +def test_multi_tool_unknown_tool_raises() -> None: + tool = MultiToolUseTool(_echo_toolset()) + with pytest.raises(ValueError, match="not found in toolset"): + tool({"tool_calls": [{"tool_name": "missing", "parameters": {}}]}) + + +# ═══════════════════════ 13. UserTextTool ═════════════════════════════════ + + +def test_user_text_tool_always_raises() -> None: + with pytest.raises(ValueError, match="should not be called directly"): + UserTextTool()({}) + + +# ═══════════════════════ 14. ToolSet / build ══════════════════════════════ + + +def test_toolset_add_duplicate_raises() -> None: + ts = ToolSet() + ts.add_tool(PruneChunksTool()) + with pytest.raises(ValueError, match="already exists"): + ts.add_tool(PruneChunksTool()) + + +def test_toolset_remove_missing_is_noop_and_get() -> None: + ts = ToolSet() + tool = PruneChunksTool() + ts.add_tool(tool) + ts.remove_tool("does_not_exist") # no raise + assert ts.get_tool("prune_chunks") is tool + ts.remove_tool("prune_chunks") + assert ts.get_tool("prune_chunks") is None + + +def test_toolset_get_formats_one_per_tool() -> None: + ts = ToolSet() + ts.add_tool(PruneChunksTool()) + ts.add_tool(UserTextTool()) + fmts = ts.get_formats(ProviderFormat.OPENAI) + assert {f["name"] for f in fmts} == {"prune_chunks", "user_text"} + + +def test_toolset_repr_sorted_with_name() -> None: + ts = ToolSet(name="mine") + ts.add_tool(UserTextTool()) + ts.add_tool(PruneChunksTool()) + assert repr(ts) == "ToolSet (mine)[2 tools: prune_chunks, user_text]" + + +def test_build_requires_retriever_or_deps() -> None: + with pytest.raises(ValueError, match="requires either 'retriever'"): + ToolSet.build() + + +def test_build_default_toolset_has_four_tools() -> None: + ts = ToolSet.build(retriever=FakeRetriever()) + assert set(ts.tools) == {"search_corpus", "grep_corpus", "read_document", "prune_chunks"} + + +def test_build_enable_raw_query_adds_execute_query() -> None: + ts = ToolSet.build( + retriever=FakeRetriever(), + enable_raw_query=True, + cosmos_client=FakeCosmosClient(FakeContainer()), + ) + assert "execute_query" in ts.tools + assert len(ts.tools) == 5 + + +def test_build_raw_query_not_added_without_client() -> None: + ts = ToolSet.build(retriever=FakeRetriever(), enable_raw_query=True) + assert "execute_query" not in ts.tools 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/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..f394bd1 100644 --- a/src/AzureCosmosDB.MCP.Toolkit/Controllers/MCPProtocolController.cs +++ b/src/AzureCosmosDB.MCP.Toolkit/Controllers/MCPProtocolController.cs @@ -259,6 +259,29 @@ 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 container to narrow to. Omit to search the whole database (all searchable collections).", maxLength = 256 }, + temperature = new { type = "number", description = "Optional LLM sampling temperature for this call (0.0-2.0). Lower is more deterministic.", minimum = 0.0, maximum = 2.0 }, + maxTurns = new { type = "integer", description = "Optional cap on the agent's reasoning/search turns for this call (1-200)", minimum = 1, maximum = 200 }, + reasoningEffort = new { type = "string", description = "Optional reasoning effort for reasoning models: 'low', 'medium', or 'high'", @enum = new[] { "low", "medium", "high" } }, + schemaOverride = new { type = "object", description = "Optional schema override as a JSON object (keys: document_id_path, chunk_id_path, chunk_order_path, title_path, source_path, item_id_path, use_dunder_codec). Omit for pure discovery." }, + searchDisplayLimit = new { type = "integer", description = "Optional cap on how many hits each internal search surfaces (1-50)", minimum = 1, maximum = 50 }, + accountUri = new { type = "string", description = "Optional Cosmos account endpoint URL to target a different account for this call (else the server's configured account), e.g. https://.documents.azure.com:443/", maxLength = 512 }, + embeddingModel = new { type = "string", description = "Optional embedding model/deployment name for this call (must match how the target container was embedded)", maxLength = 256 }, + embeddingEndpoint = new { type = "string", description = "Optional embedding endpoint base URL for this call, e.g. https://.services.ai.azure.com/openai/v1", maxLength = 512 } + }, + required = new string[] { "query" }, + additionalProperties = false + } } } } @@ -452,6 +475,20 @@ private async Task ExecuteTool(string toolName, Dictionary await _cosmosDbTools.AgenticSearch( + GetStringArg(args, "query"), + GetOptionalIntArg(args, "maxDocuments", 20), + GetOptionalStringArg(args, "database"), + GetOptionalStringArg(args, "container"), + GetNullableDoubleArg(args, "temperature"), + GetNullableIntArg(args, "maxTurns"), + GetOptionalStringArg(args, "reasoningEffort"), + GetOptionalSchemaOverrideArg(args, "schemaOverride"), + GetNullableIntArg(args, "searchDisplayLimit"), + GetOptionalStringArg(args, "accountUri"), + GetOptionalStringArg(args, "embeddingModel"), + GetOptionalStringArg(args, "embeddingEndpoint"), + cancellationToken), _ => throw new ArgumentException($"Unknown tool: {toolName}") }; } @@ -461,6 +498,57 @@ 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; + } + + // schemaOverride may arrive as a JSON object (JsonElement) or a string. Return + // its JSON text form so it can be forwarded to the retriever, which parses it. + private static string? GetOptionalSchemaOverrideArg(Dictionary args, string key) + { + if (!args.TryGetValue(key, out var value) || value is null) return null; + if (value is System.Text.Json.JsonElement el) + { + return el.ValueKind switch + { + System.Text.Json.JsonValueKind.Null => null, + System.Text.Json.JsonValueKind.String => el.GetString(), + System.Text.Json.JsonValueKind.Object => el.GetRawText(), + _ => el.GetRawText(), + }; + } + var s = value.ToString(); + return string.IsNullOrWhiteSpace(s) ? null : s; + } + + private static int? GetNullableIntArg(Dictionary args, string key) + { + if (!args.TryGetValue(key, out var value)) return null; + if (value is int intValue) return intValue; + if (int.TryParse(value?.ToString(), out var parsed)) return parsed; + return null; + } + + private static double? GetNullableDoubleArg(Dictionary args, string key) + { + if (!args.TryGetValue(key, out var value)) return null; + if (value is double dblValue) return dblValue; + if (value is int intValue) return intValue; + if (double.TryParse(value?.ToString(), System.Globalization.NumberStyles.Float, System.Globalization.CultureInfo.InvariantCulture, out var parsed)) return parsed; + return null; + } + + private static bool? GetNullableBoolArg(Dictionary args, string key) + { + if (!args.TryGetValue(key, out var value)) return null; + if (value is bool boolValue) return boolValue; + if (bool.TryParse(value?.ToString(), out var parsed)) return parsed; + return null; + } + 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..92c3d54 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,28 @@ 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; + double? temperature = parameters.ContainsKey("temperature") ? GetRequiredParameter(parameters, "temperature") : null; + int? maxTurns = parameters.ContainsKey("maxTurns") ? GetRequiredParameter(parameters, "maxTurns") : null; + string? reasoningEffort = parameters.ContainsKey("reasoningEffort") ? GetRequiredParameter(parameters, "reasoningEffort") : null; + string? schemaOverride = parameters.ContainsKey("schemaOverride") ? GetRequiredParameter(parameters, "schemaOverride") : null; + int? searchDisplayLimit = parameters.ContainsKey("searchDisplayLimit") ? GetRequiredParameter(parameters, "searchDisplayLimit") : null; + string? accountUri = parameters.ContainsKey("accountUri") ? GetRequiredParameter(parameters, "accountUri") : null; + string? embeddingModel = parameters.ContainsKey("embeddingModel") ? GetRequiredParameter(parameters, "embeddingModel") : null; + string? embeddingEndpoint = parameters.ContainsKey("embeddingEndpoint") ? GetRequiredParameter(parameters, "embeddingEndpoint") : null; + return await _cosmosDbTools.AgenticSearch( + query, maxDocuments, database, container, + temperature, maxTurns, reasoningEffort, schemaOverride, searchDisplayLimit, + accountUri, embeddingModel, embeddingEndpoint); + } + 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..775ea1e 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 @@ -402,6 +404,13 @@ public static class CosmosDbTools // OPENAI_EMBEDDING_DEPLOYMENT - Embedding model deployment name (e.g. text-embedding-3-small) // Auth uses Entra ID via DefaultAzureCredential (supports Managed Identity and service principals). + // Shared credential for the native Cosmos tools. Excludes Managed Identity so + // that on Azure VMs (where IMDS is present but the VM identity lacks Cosmos + // RBAC → "SSO failure") the chain falls through to the developer's Azure CLI + // login, matching CosmosClientFactory and the Python retriever. + private static DefaultAzureCredential CreateCosmosCredential() => + new(new DefaultAzureCredentialOptions { ExcludeManagedIdentityCredential = true }); + [McpServerTool, Description("Lists databases available in the Cosmos DB account.")] public static async Task ListDatabases() { @@ -413,7 +422,7 @@ public static async Task ListDatabases() return JsonSerializer.Serialize(new { error = "Missing required environment variable COSMOS_ENDPOINT." }); } - var credential = new DefaultAzureCredential(); + var credential = CreateCosmosCredential(); using var client = new CosmosClient(endpoint, credential, new CosmosClientOptions { ApplicationName = "AzureCosmosDBMCP" @@ -459,7 +468,7 @@ public static async Task ListCollections( return JsonSerializer.Serialize(new { error = "Parameter 'databaseId' is required." }); } - var credential = new DefaultAzureCredential(); + var credential = CreateCosmosCredential(); using var client = new CosmosClient(endpoint, credential, new CosmosClientOptions { ApplicationName = "AzureCosmosDBMCP" @@ -511,7 +520,7 @@ public static async Task GetRecentDocuments( return JsonSerializer.Serialize(new { error = "Parameter 'n' must be a whole number between 1 and 20." }); } - var credential = new DefaultAzureCredential(); + var credential = CreateCosmosCredential(); using var client = new CosmosClient(endpoint, credential, new CosmosClientOptions { ApplicationName = "AzureCosmosDBMCP" @@ -583,7 +592,7 @@ public static async Task TextSearch( return JsonSerializer.Serialize(new { error = "Invalid property name. Use dot notation with letters, digits, and underscores only (e.g., name or profile.name)." }); } - var credential = new DefaultAzureCredential(); + var credential = CreateCosmosCredential(); using var client = new CosmosClient(endpoint, credential, new CosmosClientOptions { ApplicationName = "AzureCosmosDBMCP" @@ -641,7 +650,7 @@ public static async Task FindDocumentByID( return JsonSerializer.Serialize(new { error = "Parameter 'id' is required." }); } - var credential = new DefaultAzureCredential(); + var credential = CreateCosmosCredential(); using var client = new CosmosClient(endpoint, credential, new CosmosClientOptions { ApplicationName = "AzureCosmosDBMCP" @@ -691,7 +700,7 @@ public static async Task GetApproximateSchema( return JsonSerializer.Serialize(new { error = "Parameters 'databaseId' and 'containerId' are required." }); } - var credential = new DefaultAzureCredential(); + var credential = CreateCosmosCredential(); using var client = new CosmosClient(endpoint, credential, new CosmosClientOptions { ApplicationName = "AzureCosmosDBMCP" @@ -868,7 +877,7 @@ public static async Task VectorSearch( } } - var credential = new DefaultAzureCredential(); + var credential = CreateCosmosCredential(); // Generate embedding using the appropriate embedding service // (Azure AI Services, OpenAI native, or Azure AI Foundry) @@ -1025,7 +1034,7 @@ public static async Task HybridSearch( } } - var credential = new DefaultAzureCredential(); + var credential = CreateCosmosCredential(); // Generate embedding using the configured embedding service float[] embedding; @@ -1093,4 +1102,76 @@ 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. Optional tuning knobs (temperature, maxTurns, reasoningEffort, schemaOverride, searchDisplayLimit) override the retriever's defaults for this call only.")] + 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 container to narrow the search to. Omit to search the WHOLE database (all searchable collections) — recommended default.")] string? container = null, + [Description("Optional LLM sampling temperature for this call (0.0-2.0). Lower is more deterministic.")] double? temperature = null, + [Description("Optional cap on the agent's reasoning/search turns for this call (1-200).")] int? maxTurns = null, + [Description("Optional reasoning effort for reasoning models: 'low', 'medium', or 'high'.")] string? reasoningEffort = null, + [Description("Optional schema override as a JSON object (keys: document_id_path, chunk_id_path, chunk_order_path, title_path, source_path, item_id_path, use_dunder_codec), or 'none' for pure discovery. Example: {\"document_id_path\":\"/docid\",\"chunk_order_path\":\"/chunk_idx\",\"use_dunder_codec\":true}")] string? schemaOverride = null, + [Description("Optional cap on how many hits each internal search surfaces (1-50).")] int? searchDisplayLimit = null, + [Description("Optional Cosmos account endpoint URL to target a different account for this call (else the server's configured account). Example: https://.documents.azure.com:443/")] string? accountUri = null, + [Description("Optional embedding model/deployment name to use for this call (must match how the target container was embedded).")] string? embeddingModel = null, + [Description("Optional embedding endpoint base URL to use for this call, e.g. https://.services.ai.azure.com/openai/v1 or http://host:port/v1.")] string? embeddingEndpoint = 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." }); + } + if (temperature is < 0.0 or > 2.0) + { + return JsonSerializer.Serialize(new { error = "Parameter 'temperature' must be between 0.0 and 2.0." }); + } + if (maxTurns is < 1 or > 200) + { + return JsonSerializer.Serialize(new { error = "Parameter 'maxTurns' must be between 1 and 200." }); + } + if (searchDisplayLimit is < 1 or > 50) + { + return JsonSerializer.Serialize(new { error = "Parameter 'searchDisplayLimit' must be between 1 and 50." }); + } + if (reasoningEffort is not null && reasoningEffort is not ("low" or "medium" or "high")) + { + return JsonSerializer.Serialize(new { error = "Parameter 'reasoningEffort' must be 'low', 'medium', or 'high'." }); + } + if (schemaOverride is not null && !string.Equals(schemaOverride, "none", StringComparison.OrdinalIgnoreCase)) + { + try + { + using var doc = JsonDocument.Parse(schemaOverride); + if (doc.RootElement.ValueKind != JsonValueKind.Object) + { + return JsonSerializer.Serialize(new { error = "Parameter 'schemaOverride' must be a JSON object or 'none'." }); + } + } + catch (JsonException) + { + return JsonSerializer.Serialize(new { error = "Parameter 'schemaOverride' must be valid JSON (an object) or 'none'." }); + } + } + if (accountUri is not null && !Uri.TryCreate(accountUri, UriKind.Absolute, out _)) + { + return JsonSerializer.Serialize(new { error = "Parameter 'accountUri' must be an absolute URL, e.g. https://.documents.azure.com:443/." }); + } + if (embeddingEndpoint is not null && !Uri.TryCreate(embeddingEndpoint, UriKind.Absolute, out _)) + { + return JsonSerializer.Serialize(new { error = "Parameter 'embeddingEndpoint' must be an absolute URL." }); + } + + return await AgenticSearchExecutor.RunAsync( + query, maxDocuments, logger, database, container, + temperature, maxTurns, reasoningEffort, schemaOverride, searchDisplayLimit, + accountUri, embeddingModel, embeddingEndpoint); + } } diff --git a/src/AzureCosmosDB.MCP.Toolkit/Services/AgenticSearchExecutor.cs b/src/AzureCosmosDB.MCP.Toolkit/Services/AgenticSearchExecutor.cs new file mode 100644 index 0000000..342464c --- /dev/null +++ b/src/AzureCosmosDB.MCP.Toolkit/Services/AgenticSearchExecutor.cs @@ -0,0 +1,285 @@ +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"; + + /// + /// Optional JSON map of { "<database>": "<retriever base url>" }. + /// When a request targets a database present in this map, that retriever + /// endpoint is used instead of . Lets different + /// databases be served by different retriever deployments. + /// + public const string BaseUrlMapEnvVar = "COSMOS_RETRIEVER_URL_MAP"; + + 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. + /// Optional LLM sampling temperature (0.0–2.0). + /// Optional cap on agent reasoning turns (1–200). + /// Optional reasoning effort ("low"/"medium"/"high"). + /// Optional schema override as a JSON object (keys: + /// document_id_path, chunk_id_path, chunk_order_path, title_path, source_path, + /// item_id_path, use_dunder_codec), or "none" for pure discovery. + /// Optional cap on hits surfaced per search (1–50). + /// Optional Cosmos account endpoint override for this call. + /// Optional embedding model/deployment override for this call. + /// Optional embedding endpoint base URL override for this call. + /// 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. + /// + /// + /// The optional tuning knobs are forwarded to the retriever service as a + /// per-request overrides object (mapping to the Python + /// RuntimeConfig). Only non-null values are sent; anything omitted + /// falls back to the retriever's own configured defaults. + /// + public static async Task RunAsync( + string query, + int maxDocuments, + ILogger logger, + string? database = null, + string? container = null, + double? temperature = null, + int? maxTurns = null, + string? reasoningEffort = null, + string? schemaOverride = null, + int? searchDisplayLimit = null, + string? accountUri = null, + string? embeddingModel = null, + string? embeddingEndpoint = null, + CancellationToken cancellationToken = default) + { + var baseUrl = ResolveBaseUrl(database).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; + + // Per-request tuning knobs -> the retriever's RuntimeConfig overrides. + // Only include knobs the caller actually set; omit the rest so the + // service applies its own defaults. + var overrides = new Dictionary(); + if (temperature is not null) overrides["chat_temperature"] = temperature; + if (maxTurns is not null) overrides["chat_max_turns"] = maxTurns; + if (!string.IsNullOrWhiteSpace(reasoningEffort)) overrides["chat_reasoning_effort"] = reasoningEffort; + if (!string.IsNullOrWhiteSpace(schemaOverride) && !string.Equals(schemaOverride, "none", StringComparison.OrdinalIgnoreCase)) + { + // Forward the schema override as a nested JSON object so the retriever + // receives a structured override (its RuntimeConfig coerces it). + try + { + using var doc = JsonDocument.Parse(schemaOverride); + overrides["schema_override"] = doc.RootElement.Clone(); + } + catch (JsonException) + { + overrides["schema_override"] = schemaOverride; + } + } + if (searchDisplayLimit is not null) overrides["search_display_limit"] = searchDisplayLimit; + if (!string.IsNullOrWhiteSpace(accountUri)) overrides["account_uri"] = accountUri; + if (!string.IsNullOrWhiteSpace(embeddingModel)) overrides["openai_embedding_model"] = embeddingModel; + if (!string.IsNullOrWhiteSpace(embeddingEndpoint)) overrides["embed_endpoint"] = embeddingEndpoint; + if (overrides.Count > 0) payload["overrides"] = overrides; + + logger.LogInformation( + "agentic_search: POST {RequestUri} (database={Database} container={Container} timeout={Timeout}s overrides={OverrideCount})", + requestUri, database ?? "", container ?? "", timeoutSeconds, overrides.Count); + + 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; + } + + // Resolve the retriever base URL for a request, preferring a per-database + // override from COSMOS_RETRIEVER_URL_MAP (JSON {"":""}) and falling + // back to COSMOS_RETRIEVER_URL / the built-in default. + private static string ResolveBaseUrl(string? database) + { + if (!string.IsNullOrWhiteSpace(database)) + { + var raw = Environment.GetEnvironmentVariable(BaseUrlMapEnvVar); + if (!string.IsNullOrWhiteSpace(raw) && LooksLikeJson(raw)) + { + try + { + var map = JsonSerializer.Deserialize>( + raw, new JsonSerializerOptions { PropertyNameCaseInsensitive = true }); + if (map is not null) + { + foreach (var kv in map) + { + if (string.Equals(kv.Key, database, StringComparison.OrdinalIgnoreCase) + && !string.IsNullOrWhiteSpace(kv.Value)) + { + return kv.Value; + } + } + } + } + catch (JsonException) + { + // Malformed map -> fall through to the default endpoint. + } + } + } + return ResolveString(BaseUrlEnvVar, defaultValue: DefaultBaseUrl); + } + + 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..79122a6 100644 --- a/src/AzureCosmosDB.MCP.Toolkit/Services/CosmosDbToolsService.cs +++ b/src/AzureCosmosDB.MCP.Toolkit/Services/CosmosDbToolsService.cs @@ -700,4 +700,34 @@ 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, + double? temperature = null, + int? maxTurns = null, + string? reasoningEffort = null, + string? schemaOverride = null, + int? searchDisplayLimit = null, + string? accountUri = null, + string? embeddingModel = null, + string? embeddingEndpoint = null, + CancellationToken cancellationToken = default) + { + var raw = await AgenticSearchExecutor.RunAsync( + query, maxDocuments, _logger, database, container, + temperature, maxTurns, reasoningEffort, schemaOverride, searchDisplayLimit, + accountUri, embeddingModel, embeddingEndpoint, + 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..93c07b1 100644 --- a/src/AzureCosmosDB.MCP.Toolkit/Services/McpToolRequestValidator.cs +++ b/src/AzureCosmosDB.MCP.Toolkit/Services/McpToolRequestValidator.cs @@ -58,6 +58,21 @@ 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), + ["temperature"] = ToolArgumentSchema.Number(required: false, minValue: 0.0, maxValue: 2.0), + ["maxTurns"] = ToolArgumentSchema.Integer(required: false, minValue: 1, maxValue: 200), + ["reasoningEffort"] = ToolArgumentSchema.String(required: false, maxLength: 16), + ["schemaOverride"] = ToolArgumentSchema.Object(required: false, maxLength: 2048), + ["searchDisplayLimit"] = ToolArgumentSchema.Integer(required: false, minValue: 1, maxValue: 50), + ["accountUri"] = ToolArgumentSchema.String(required: false, maxLength: 512), + ["embeddingModel"] = ToolArgumentSchema.String(required: false, maxLength: 256), + ["embeddingEndpoint"] = ToolArgumentSchema.String(required: false, maxLength: 512) }) }; @@ -68,7 +83,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) { @@ -116,8 +133,11 @@ private static Dictionary ValidateArguments(string toolName, Too validated[argument.Key] = argument.Value.Kind switch { - JsonValueKind.String => ValidateString(valueElement.GetString(), argument.Key, argument.Value.MaxLength), - JsonValueKind.Number => ValidateInteger(valueElement, argument.Key, argument.Value.MinValue, argument.Value.MaxValue), + ArgKind.String => ValidateString(valueElement.GetString(), argument.Key, argument.Value.MaxLength), + ArgKind.Integer => ValidateInteger(valueElement, argument.Key, (int)argument.Value.MinValue, (int)argument.Value.MaxValue), + ArgKind.Number => ValidateNumber(valueElement, argument.Key, argument.Value.MinValue, argument.Value.MaxValue), + ArgKind.Boolean => ValidateBoolean(valueElement, argument.Key), + ArgKind.Object => ValidateObject(valueElement, argument.Key, argument.Value.MaxLength), _ => throw new ToolInputValidationException($"Unsupported schema for argument '{argument.Key}'.") }; } @@ -166,6 +186,49 @@ private static int ValidateInteger(JsonElement valueElement, string fieldName, i return value; } + private static double ValidateNumber(JsonElement valueElement, string fieldName, double minValue, double maxValue) + { + if (valueElement.ValueKind != JsonValueKind.Number || !valueElement.TryGetDouble(out var value)) + { + throw new ToolInputValidationException($"'{fieldName}' must be a number."); + } + + if (value < minValue || value > maxValue) + { + throw new ToolInputValidationException($"'{fieldName}' must be between {minValue} and {maxValue}."); + } + + return value; + } + + private static bool ValidateBoolean(JsonElement valueElement, string fieldName) + { + if (valueElement.ValueKind != JsonValueKind.True && valueElement.ValueKind != JsonValueKind.False) + { + throw new ToolInputValidationException($"'{fieldName}' must be a boolean."); + } + + return valueElement.GetBoolean(); + } + + // Accepts a JSON object and returns its raw JSON text (length-capped). Used + // for structured arguments like schemaOverride that are forwarded verbatim. + private static string ValidateObject(JsonElement valueElement, string fieldName, int maxLength) + { + if (valueElement.ValueKind != JsonValueKind.Object) + { + throw new ToolInputValidationException($"'{fieldName}' must be a JSON object."); + } + + var raw = valueElement.GetRawText(); + if (raw.Length > maxLength) + { + throw new ToolInputValidationException($"'{fieldName}' exceeds the maximum length of {maxLength} characters."); + } + + return raw; + } + private static bool TryGetProperty(JsonElement element, string propertyName, out JsonElement value) { if (element.ValueKind == JsonValueKind.Object && element.TryGetProperty(propertyName, out value)) @@ -196,13 +259,24 @@ private static void RejectUnknownProperties(JsonElement element, IEnumerable Arguments); - private sealed record ToolArgumentSchema(JsonValueKind Kind, bool Required, int MaxLength = 0, int MinValue = 0, int MaxValue = 0) + private enum ArgKind { String, Integer, Number, Boolean, Object } + + private sealed record ToolArgumentSchema(ArgKind Kind, bool Required, int MaxLength = 0, double MinValue = 0, double MaxValue = 0) { public static ToolArgumentSchema String(bool required, int maxLength) - => new(JsonValueKind.String, required, MaxLength: maxLength); + => new(ArgKind.String, required, MaxLength: maxLength); public static ToolArgumentSchema Integer(bool required, int minValue, int maxValue) - => new(JsonValueKind.Number, required, MinValue: minValue, MaxValue: maxValue); + => new(ArgKind.Integer, required, MinValue: minValue, MaxValue: maxValue); + + public static ToolArgumentSchema Number(bool required, double minValue, double maxValue) + => new(ArgKind.Number, required, MinValue: minValue, MaxValue: maxValue); + + public static ToolArgumentSchema Boolean(bool required) + => new(ArgKind.Boolean, required); + + public static ToolArgumentSchema Object(bool required, int maxLength) + => new(ArgKind.Object, required, MaxLength: maxLength); } } diff --git a/tests/AzureCosmosDB.MCP.Toolkit.Tests/AgenticSearchExecutorTests.cs b/tests/AzureCosmosDB.MCP.Toolkit.Tests/AgenticSearchExecutorTests.cs new file mode 100644 index 0000000..efd0505 --- /dev/null +++ b/tests/AzureCosmosDB.MCP.Toolkit.Tests/AgenticSearchExecutorTests.cs @@ -0,0 +1,226 @@ +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 the AgenticSearchExecutor. Stands in for the cosmos-retriever +/// FastAPI service with a tiny in-process HttpListener 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 HttpListener. + /// 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(); + } + } +}