Cross-collection agentic retrieval + live schema discovery#150
Open
aryan-410 wants to merge 17 commits into
Open
Cross-collection agentic retrieval + live schema discovery#150aryan-410 wants to merge 17 commits into
aryan-410 wants to merge 17 commits into
Conversation
FastAPI service (POST /search, GET /health) wrapping CosmosRetriever, which runs a multi-turn retrieval agent over a Cosmos DB corpus. Pluggable inference backend: harmony_vllm (fine-tuned pat-jj/harness-1), openai_chat (any OpenAI-compatible chat model), or openai_responses (reasoning models such as gpt-5.4 on Azure AI Foundry). Includes tests for the server and agent loops. The .NET agentic_search tool calls this service over HTTP.
Add a 9th MCP tool, agentic_search, that runs the Cosmos Retriever agent over a Cosmos DB corpus and returns ranked, curated documents. - AgenticSearchExecutor: calls the cosmos-retriever service over HTTP (COSMOS_RETRIEVER_URL, COSMOS_RETRIEVER_TIMEOUT_S); always returns parseable JSON (error envelope on failure). - Wire into Program.cs, MCPProtocolController (tools/list + tools/call), MCPTestController, and McpToolRequestValidator. - CosmosClientFactory: exclude ManagedIdentityCredential (fall through to az login); accept the standard MCP _meta params field. - Docs: docs/AGENTIC_SEARCH.md, README + CHANGELOG + .env.example.
…rness - Move VllmTokenCompleter + run_single_episode into inference/vllm_policy.py - Delete inference/evaluate_harness1_vllm.py (eval/benchmark code) - Repoint retriever.py to vllm_policy; update env_rl docstring - Include pre-existing: pool_doc_ids trajectory pooling (openai_chat), optional baseten import (rerank)
…s to 1-50, de-brand harness references
The datagen/ package deletion was previously only staged, never committed, so it still appeared in the PR. Actually remove it (search_dataset.py, generate_sft_rl_splits.py, BrowseComp-Plus, README, __init__) along with the unit tests folder, the datagen TYPE_CHECKING import in tasks.py, the stale datagen comment in config.py, and the now-dangling pytest/respx dev deps and pytest/ruff test config in pyproject.toml.
…lResult Add a trajectory field to RetrievalResult populated by the harmony_vllm backend: the search queries issued (search_history), per-turn tool calls (turn_tools), programmatic per-turn status summaries (turn_summaries), and the final per-doc importance tags from curation (curated_importance).
…ed retrieval layer - Add schema-driven retrieval package (planner/compiler/executor/strategies/ document-resolvers) decoupled from physical Cosmos paths - Endpoint-only inference (OpenAI chat/responses); drop vLLM/RL runtime - Add foundry-harness and retrieval/agent benchmark scripts
…lt corpus profile); rename LegacyDunderCodec -> DunderChunkCodec
…ratch, endpoint-only backends
…loop.py and ChatSearchResult -> AgentSearchResult (both backends live there); keep backend-id strings
…undry Messages API); refresh backend docs/.env
- Live discovery of container schema/capabilities (retrieval/discovery/, binding.py), replacing the hardcoded default_chunked_schema (defaults.py removed) - Cross-collection retrieval: container="*" fans out across all searchable collections in a database and fuses with RRF (orchestration.py); per-collection depth tracks the fused limit so gold docs ranked deep in a collection aren't dropped - Structured, user-supplied schema overrides (schema_override.py) via the corpus registry, replacing the none/legacy_chunked enum - Database-aware corpus registry: per-database embedding endpoint/model with default-endpoint fallback - container is now optional (defaults to whole-database); database required, no silent server defaults - /config admin endpoints, per-request RuntimeConfig overrides, timing instrumentation, bounded-TTL retriever cache - .NET agentic_search: optional container, JSON schemaOverride, Managed-Identity- excluded Cosmos credential for the native tools - Tests: binding, discovery, orchestration, cache, runtime config
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Upgrades the
cosmos-retrieveragentic search service to (1) search across all collections in a database and (2) discover each container's schema/capabilities live instead of assuming a hardcoded chunked schema.Key changes
retrieval/discovery/,retrieval/binding.py): infers vector/full-text indexes, partition keys, and dimensions from container metadata. Removes the hardcodeddefault_chunked_schema(defaults.pydeleted).retrieval/orchestration.py):container="*"fans out across every searchable collection in a database and fuses hits with RRF. Per-collection depth tracks the fused limit so a gold doc ranked deep within its own collection isn't truncated before fusion.retrieval/schema_override.py): user-supplied JSON override (document_id_path, chunk_order_path, dunder codec, …) via the corpus registry, replacing thenone/legacy_chunkedenum.containeris now optional — omit to search the whole database;databaseis required (no silent server defaults).GET/PATCH /config, per-requestRuntimeConfigoverrides, per-turn timing instrumentation, bounded-TTL retriever cache.agentic_search: optionalcontainer, JSONschemaOverride, and a Managed-Identity-excluded Cosmos credential so the native tools fall through to developer CLI auth locally.Tests
Adds unit tests for binding, discovery, orchestration, cache, and runtime config. All pass (
pytest, 38).Notes