diff --git a/.secrets.baseline b/.secrets.baseline index ae2853c00..08dc3e961 100644 --- a/.secrets.baseline +++ b/.secrets.baseline @@ -306,7 +306,7 @@ "filename": "sources/knowledge_layer/KNOWLEDGE-LAYER-SETUP.md", "hashed_secret": "9d4e1e23bd5b727046a9e3b4b7db57bd8d6ee684", "is_verified": false, - "line_number": 571 + "line_number": 577 } ], "sources/knowledge_layer/src/foundational_rag/README.md": [ diff --git a/configs/config_web_azure_ai_search.yml b/configs/config_web_azure_ai_search.yml new file mode 100644 index 000000000..32cb949f6 --- /dev/null +++ b/configs/config_web_azure_ai_search.yml @@ -0,0 +1,211 @@ +# This is the Web mode configuration for Azure AI Search. +# It has the following features: +# - Knowledge retrieval using Azure AI Search +# - Web search tools by default; Paper search is optional + +general: + use_uvloop: true + telemetry: + logging: + console: + _type: console + level: INFO + # tracing: + # langsmith: # Optional: LangSmith tracing - requires langsmith API key. Set using `export LANGSMITH_API_KEY=` + # _type: langsmith + # project: nvidia-aiq + + front_end: + _type: aiq_api + runner_class: aiq_api.plugin.AIQAPIWorker + # ========================================================================= + # Knowledge API is automatically enabled when knowledge_retrieval function + # is configured + # ========================================================================= + # Async Job API Settings + # ========================================================================= + # Async job infrastructure database (NAT JobStore + EventStore) + # Used by: /v1/jobs/async routes, SSE streaming, job status persistence + # Requires async driver for SQLite (aiosqlite) or PostgreSQL (asyncpg) + # Environment overrides: + # - NAT_JOB_STORE_DB_URL (direct override) + # - NAT_JOB_STORE_DB_URL_DEV / NAT_JOB_STORE_DB_URL_PROD (via NAT_ENV) + db_url: ${NAT_JOB_STORE_DB_URL:-sqlite+aiosqlite:///./jobs.db} + # Job expiry - how long completed jobs stay in database before cleanup + expiry_seconds: 86400 # 24 hours (min: 600, max: 604800/7 days) + cors: + allow_origin_regex: 'http://localhost(:\d+)?|http://127.0.0.1(:\d+)?' + allow_methods: + - GET + - POST + - DELETE + - OPTIONS + allow_headers: + - "*" + allow_credentials: true + expose_headers: + - "*" + +llms: + nemotron_llm_intent: + _type: nim + model_name: nvidia/nemotron-3-super-120b-a12b + base_url: "https://integrate.api.nvidia.com/v1" + temperature: 0.5 + top_p: 0.9 + max_tokens: 4096 + num_retries: 5 + chat_template_kwargs: + enable_thinking: true + + nemotron_super_llm: + _type: nim + model_name: nvidia/nemotron-3-super-120b-a12b + base_url: "https://integrate.api.nvidia.com/v1" + temperature: 0.7 + top_p: 0.7 + max_tokens: 65536 + num_retries: 5 + chat_template_kwargs: + enable_thinking: true + + # LLM for document summaries (required when generate_summary: true) + summary_llm: + _type: nim + model_name: nvidia/nemotron-mini-4b-instruct + base_url: "https://integrate.api.nvidia.com/v1" + api_key: ${NVIDIA_API_KEY} + temperature: 0.3 + max_tokens: 100 + +functions: + # ========================================================================= + # Data Source Registry + # ========================================================================= + # Central registry that controls: + # 1. UI toggles — each source appears as an on/off switch in the frontend + # 2. Per-message filtering — users can select active sources per request + # 3. Tool auto-inheritance — agents with no explicit `tools` list receive + # every tool listed here (use `exclude_tools` on agents to specialize) + # + # Adding a new tool or MCP function group? Just add it here — all agents + # pick it up automatically. No per-agent config changes needed. + # + # Source entry fields: + # id — Unique key used in API payloads and filtering + # name — Display name shown in the UI + # description — Human-readable description shown in the UI + # tools — List of NAT function names or function group names + # requires_auth — (default: false) If true, the UI greys out this source + # until the user signs in. Use for sources that need + # user-level OAuth tokens (e.g., enterprise SSO). + # Sources using backend API keys (Tavily, Serper) should + # leave this false. + # default_enabled — (default: true) Whether enabled by default + # + # See docs/source/customization/tools-and-sources.md for full details. + # ========================================================================= + data_sources: + _type: data_source_registry + sources: + - id: web_search + name: "Web Search" + description: "Search the web for real-time information." + tools: + - web_search_tool + - advanced_web_search_tool + - id: knowledge_layer + name: "Knowledge Base" + description: "Search uploaded documents and files." + tools: + - knowledge_search + # Uncomment when paper_search_tool is enabled (requires SERPER_API_KEY) + # - id: paper_search + # name: "Academic Papers" + # description: "Search academic papers and scientific publications." + # tools: + # - paper_search_tool + + web_search_tool: + _type: tavily_web_search + max_results: 5 + max_content_length: 1000 + + advanced_web_search_tool: + _type: tavily_web_search + max_results: 2 + advanced_search: true + + # Knowledge Retrieval (see sources/knowledge_layer/KNOWLEDGE-LAYER-SETUP.md) + knowledge_search: + _type: knowledge_retrieval + backend: azure_ai_search + collection_name: ${COLLECTION_NAME:-aiq_default} + generate_summary: true + summary_model: summary_llm # Required when generate_summary: true + summary_db: ${AIQ_SUMMARY_DB:-sqlite+aiosqlite:///./summaries.db} + top_k: 5 + + # Paper Search (optional - requires SERPER_API_KEY) + # Uncomment the block below and set SERPER_API_KEY to enable academic paper search. + # paper_search_tool: + # _type: paper_search + # max_results: 5 + # serper_api_key: ${SERPER_API_KEY} + + # ========================================================================= + # Agents + # ========================================================================= + # Tool inheritance: agents with no `tools` list inherit ALL tools from the + # data_source_registry above. Use `exclude_tools` to remove specific tools + # from an agent (e.g., give shallow the basic search, deep the advanced). + # To bypass auto-inherit entirely, set an explicit `tools` list. + # ========================================================================= + intent_classifier: + _type: intent_classifier + llm: nemotron_llm_intent + # tools: omitted -> inherits all from data_source_registry + # exclude_tools: [] + verbose: true + + clarifier_agent: + _type: clarifier_agent + llm: nemotron_super_llm + planner_llm: nemotron_super_llm + # tools: omitted -> inherits all from data_source_registry + # exclude_tools: [] + max_turns: 3 + enable_plan_approval: true + log_response_max_chars: 2000 + verbose: true + + shallow_research_agent: + _type: shallow_research_agent + llm: nemotron_super_llm + # tools: omitted -> inherits all from data_source_registry + exclude_tools: # Remove advanced variant; shallow uses web_search_tool + - advanced_web_search_tool + verbose: true + max_llm_turns: 10 + max_tool_iterations: 5 + + deep_research_agent: + _type: deep_research_agent + enable_citation_verification: true + orchestrator_llm: nemotron_super_llm + source_router_llm: nemotron_super_llm + researcher_llm: nemotron_super_llm + planner_llm: nemotron_super_llm + writer_llm: nemotron_super_llm + # tools: omitted -> inherits all from data_source_registry + exclude_tools: # Remove basic variant; deep uses advanced_web_search_tool + - web_search_tool + verbose: true + +workflow: + _type: chat_deepresearcher_agent + verbose: true + enable_escalation: true + enable_clarifier: true + use_async_deep_research: true + checkpoint_db: ${AIQ_CHECKPOINT_DB:-./checkpoints.db} diff --git a/deploy/.env.example b/deploy/.env.example index b5b9d3b90..0f7b367b1 100644 --- a/deploy/.env.example +++ b/deploy/.env.example @@ -170,6 +170,16 @@ FILE_EXPIRATION_CHECK_INTERVAL_HOURS=0 # RAG_SERVER_URL=http://localhost:8081/v1 # RAG_INGEST_URL=http://localhost:8082/v1 +# ----------------------------------------------------------------------------- +# Azure AI Search knowledge layer (required only with the Azure backend) +# ----------------------------------------------------------------------------- +# AZURE_SEARCH_ENDPOINT=https://.search.windows.net +# API-key authentication: +# AZURE_SEARCH_API_KEY= +# User-assigned managed identity: +# AZURE_CLIENT_ID= +# Optional deployment-unique prefix (default: aiq): +# AIQ_AZURE_SEARCH_INDEX_PREFIX=aiq # ----------------------------------------------------------------------------- # Frontend (UI) runtime configuration (Docker Compose) diff --git a/docs/source/customization/configuration-reference.md b/docs/source/customization/configuration-reference.md index 91a4fd8d2..f996fa920 100644 --- a/docs/source/customization/configuration-reference.md +++ b/docs/source/customization/configuration-reference.md @@ -228,8 +228,8 @@ functions: ### `knowledge_retrieval` -Semantic search over ingested documents. AI-Q supports three backends: LlamaIndex (local ChromaDB), Foundational RAG -(hosted NVIDIA RAG Blueprint), and OpenSearch (self-hosted OpenSearch or Amazon OpenSearch Serverless). +Semantic search over ingested documents. Supports LlamaIndex (local ChromaDB), Foundational RAG +(hosted NVIDIA RAG Blueprint), OpenSearch (self-hosted OpenSearch or Amazon OpenSearch Serverless), and Azure AI Search. ```yaml functions: @@ -256,9 +256,22 @@ functions: rag_url: ${RAG_SERVER_URL:-http://localhost:8081/v1} ingest_url: ${RAG_INGEST_URL:-http://localhost:8082/v1} timeout: 300 - # verify_ssl: false # Only set to false for self-signed certs + # verify_ssl: false # Only set to false for self-signed certs ``` +```yaml +functions: + # Azure AI Search backend + knowledge_search: + _type: knowledge_retrieval + backend: azure_ai_search + collection_name: ${COLLECTION_NAME:-test_collection} +``` + +This example reads `AZURE_SEARCH_ENDPOINT` and `NVIDIA_API_KEY` from the +environment. `AZURE_SEARCH_API_KEY` is optional; when absent, the adapter uses +`DefaultAzureCredential`. + ```yaml functions: # OpenSearch backend @@ -278,7 +291,7 @@ functions: | Parameter | Type | Default | Description | |-----------|------|---------|-------------| -| `backend` | `str` | `llamaindex` | Backend type: `llamaindex`, `foundational_rag`, or `opensearch`. | +| `backend` | `str` | `llamaindex` | Backend type: `llamaindex`, `opensearch`, `foundational_rag`, or `azure_ai_search`. | | `collection_name` | `str` | `default` | Name of the document collection/index. | | `top_k` | `int` | `5` | Number of results to return per query. | | `generate_summary` | `bool` | `false` | Generate one-sentence summaries for ingested documents. | @@ -289,6 +302,10 @@ functions: | `ingest_url` | `str` | `http://localhost:8082/v1` | RAG ingestion server URL. Foundational RAG backend only. | | `timeout` | `int` | `120` | Request timeout in seconds. Foundational RAG backend only. | | `verify_ssl` | `bool` | `true` | Verify SSL certificates. Set `false` for self-signed certs. Foundational RAG backend only. | +| `azure_search_endpoint` | `URL` | `AZURE_SEARCH_ENDPOINT` | Azure AI Search service endpoint. Required for Azure AI Search. | +| `azure_search_api_key` | `SecretStr` | `AZURE_SEARCH_API_KEY` | Optional admin API key. | +| `azure_search_index_prefix` | `str` | `AIQ_AZURE_SEARCH_INDEX_PREFIX` or `aiq` | Deployment-unique namespace for the shared AI-Q index. | +| `embed_dim` | `int` | `AIQ_EMBED_DIM` or `2048` | Embedding dimensions; must match the model and existing index schema. | | `opensearch_url` | `str` | `http://localhost:9200` | OpenSearch endpoint. OpenSearch backend only. | | `opensearch_auth_type` | `str` | `none` | Authentication mode: `none`, `basic`, or `sigv4`. | | `opensearch_username` | `str` | `None` | Username for basic authentication. Also read from `OPENSEARCH_USERNAME`. | @@ -585,6 +602,7 @@ only the additional sections you need. |------|------|------------------------------| | `configs/config_cli_default.yml` | CLI | Chat pipeline with Tavily web search and clarification. No knowledge backend. Paper search is present only as a commented opt-in. | | `configs/config_web_default_llamaindex.yml` | Web API | Default chat pipeline with LlamaIndex/ChromaDB knowledge retrieval and Tavily. Paper search is commented out. | +| `configs/config_web_azure_ai_search.yml` | Web API | Azure AI Search knowledge retrieval and web search | | `configs/config_web_frag.yml` | Web API / Helm base | Foundational RAG plus Tavily. Requires separately deployed RAG query and ingestion services. Paper search is commented out. | | `configs/config_web_opensearch.yml` | Web API | Built-in OpenSearch knowledge backend plus Tavily. Supports unauthenticated or basic self-hosted OpenSearch and SigV4 (`es` or `aoss`); infrastructure and credentials are deployment opt-ins. | | `configs/config_frontier_models.yml` | Web API | LlamaIndex plus explicit per-agent tools, Nemotron researcher roles, and an OpenAI frontier model for orchestration/planning/writing. Requires `OPENAI_API_KEY`; paper search is commented out. | diff --git a/docs/source/customization/knowledge-layer.md b/docs/source/customization/knowledge-layer.md index e964d70c1..df7e6b420 100644 --- a/docs/source/customization/knowledge-layer.md +++ b/docs/source/customization/knowledge-layer.md @@ -43,6 +43,7 @@ A pluggable abstraction for document ingestion and retrieval. Swap backends with |---------|-------------|------|--------------|----------| | `llamaindex` | `"llamaindex"` | Local Library | ChromaDB | Dev, prototyping, macOS/Linux | | `foundational_rag` | `"foundational_rag"` | Hosted Service | Remote Milvus | Production, multi-user | +| `azure_ai_search` | `"azure_ai_search"` | Managed Service | Azure AI Search | Managed hybrid retrieval | | `opensearch` | `"opensearch"` | External Service | OpenSearch k-NN index | Self-hosted OpenSearch, Amazon OpenSearch Service, or Serverless | **Local Library Mode** - Everything runs in your Python process. No external services needed. @@ -53,6 +54,8 @@ A pluggable abstraction for document ingestion and retrieval. Swap backends with - Tested with: **NVIDIA RAG Blueprint `v2.4.0`** (Helm chart `nvidia-blueprint-rag`) - [Deployment Guide](https://github.com/NVIDIA-AI-Blueprints/rag/blob/main/docs/deploy-docker-self-hosted.md) - Backend-specific documentation: `sources/knowledge_layer/src/foundational_rag/README.md` +- **`azure_ai_search`** - Stores client-generated embeddings in namespaced Azure AI Search indexes and supports + vector, hybrid, and semantic-ranked retrieval. - **`opensearch`** - Uses one vector index per AI-Q collection with `none`, `basic`, or SigV4 authentication. - Supports self-hosted OpenSearch, Amazon OpenSearch Service (`es`), and Amazon OpenSearch Serverless (`aoss`). - Can ingest in the local process or dispatch ingestion to Dask workers. @@ -75,6 +78,7 @@ export NVIDIA_API_KEY=nvapi-your-key-here # 2. Install backend (choose one) uv pip install -e "sources/knowledge_layer[llamaindex]" # Recommended for local dev - works on macOS/Linux uv pip install -e "sources/knowledge_layer[foundational_rag]" # Requires deployed server +uv pip install -e "sources/knowledge_layer[azure_ai_search]" # Requires an Azure AI Search service uv pip install -e "sources/knowledge_layer[opensearch]" # Requires an OpenSearch endpoint ``` @@ -164,6 +168,36 @@ functions: timeout: 120 ``` +**Azure AI Search (Managed Service)** + +```yaml +functions: + knowledge_search: + _type: knowledge_retrieval + backend: azure_ai_search + collection_name: my_docs +``` + +Set `AZURE_SEARCH_ENDPOINT` and `NVIDIA_API_KEY` in the environment. Setting +`AZURE_SEARCH_API_KEY` selects key authentication; otherwise Azure +`DefaultAzureCredential` is used. The workload identity needs `Search Service +Contributor` for index management and `Search Index Data Contributor` for +document ingestion and retrieval. Embedding defaults can be shared with the +LlamaIndex backend through `AIQ_EMBED_BASE_URL` and `AIQ_EMBED_MODEL`; set +`AIQ_EMBED_DIM` when changing the model dimensions. Set a deployment-unique +`AIQ_AZURE_SEARCH_INDEX_PREFIX` when multiple AI-Q deployments share a search +service. + +Azure stores all logical collections in one physical index selected by the +prefix, schema version, embedding model, and dimension. Collection, file, and +chunk manifests enforce logical isolation. Retrieval is always hybrid, and +chunking is fixed at 1024 tokens with 128-token overlap. + +Upload responses return canonical UUID file IDs. Same-name uploads coexist as +independent files. Collection cleanup uses `AIQ_COLLECTION_TTL_HOURS` (24 hours +by default) and `AIQ_TTL_CLEANUP_INTERVAL_SECONDS` (one hour by default), +matching the other knowledge backends. + **OpenSearch (Self-Hosted or AWS)** ```yaml @@ -264,10 +298,13 @@ File type support depends on the configured backend: | **LlamaIndex** | PDF, DOCX, TXT, MD, HTML, JSON, CSV | | **Foundational RAG** | PDF, DOCX, PPTX, TXT, MD, HTML, images (PNG, JPG) | | **OpenSearch** | PDF, DOCX, PPTX, TXT, MD, CSV, JSON, YAML, YML, LOG | +| **Azure AI Search** | PDF, DOCX, TXT, MD | For custom backends, supported types are determined by the backend implementation. -> **Note:** The backends support more types than the frontend currently allows. The frontend only supports uploading `.pdf,.docx,.txt,.md` (the common subset across the shipped backends). Types such as PPTX, HTML, JSON, CSV, YAML, logs, and images depend on the backend and are not accepted by the default frontend upload flow. +> **Note:** The backends support more types than the frontend currently allows. The frontend only supports uploading +> `.pdf,.docx,.txt,.md` (the common subset across all backends). Types like HTML, JSON, CSV, and images are supported by +> some backends but the frontend upload flow does not handle them yet -- this is a separate task. To change the accepted types in the frontend, set `FILE_UPLOAD_ACCEPTED_TYPES` for your deployment method: @@ -448,6 +485,8 @@ Configuration values are resolved in the following order (highest to lowest prio | `KNOWLEDGE_RETRIEVER_BACKEND` | All | Default retriever backend (fallback if not in YAML) | | `KNOWLEDGE_INGESTOR_BACKEND` | All | Default ingestor backend (fallback if not in YAML) | | `AIQ_CHROMA_DIR` | llamaindex | ChromaDB persistence path | +| `AIQ_COLLECTION_TTL_HOURS` | all local/managed backends | Hours before stale collections are deleted (default: 24) | +| `AIQ_TTL_CLEANUP_INTERVAL_SECONDS` | all local/managed backends | Collection cleanup interval (default: 3600) | | `RAG_SERVER_URL` | foundational_rag | Query server URL (port 8081) | | `RAG_INGEST_URL` | foundational_rag | Ingestion server URL (port 8082) | | `OPENSEARCH_URL` | opensearch | OpenSearch endpoint URL | diff --git a/docs/source/examples/azure-ai-search.md b/docs/source/examples/azure-ai-search.md new file mode 100644 index 000000000..e05cb7a13 --- /dev/null +++ b/docs/source/examples/azure-ai-search.md @@ -0,0 +1,88 @@ + + +# Example: Azure AI Search knowledge layer + +Use Azure AI Search as the document store while retaining AI-Q's existing +upload API, per-conversation collection routing, document summaries, and +citations. This example assumes the Azure AI Search service and embedding +endpoint already exist; it does not deploy Azure infrastructure. + +## Prerequisites + +- Create or select an [Azure AI Search service](https://learn.microsoft.com/azure/search/search-create-service-portal). +- Choose a tier whose [document and storage limits](https://learn.microsoft.com/azure/search/search-limits-quotas-capacity) + cover the expected data volume. AI-Q stores logical collections in one physical index. +- Copy the service endpoint from the Azure portal. For key authentication, also copy an admin key. Otherwise, + [enable role-based access](https://learn.microsoft.com/azure/search/keyless-connections) and assign the roles below. + +Install the backend dependency: + +```bash +uv pip install -e "sources/knowledge_layer[azure_ai_search]" +``` + +Create `deploy/.env` when needed. The Azure entries are commented in the +template so non-Azure runs do not change; uncomment the endpoint and the +settings required for your authentication mode: + +```bash +cp -n deploy/.env.example deploy/.env +``` + +```text +NVIDIA_API_KEY= +AZURE_SEARCH_ENDPOINT=https://.search.windows.net +# API-key authentication: +# AZURE_SEARCH_API_KEY= +# User-assigned managed identity: +# AZURE_CLIENT_ID= +# Optional deployment-unique prefix (default: aiq): +# AIQ_AZURE_SEARCH_INDEX_PREFIX=aiq +``` + +## Grant managed identity access + +When `AZURE_SEARCH_API_KEY` is absent, enable role-based access on the Azure AI +Search service and grant the workload identity both of these built-in roles: + +| Role | Used for | +|------|----------| +| `Search Service Contributor` | Create and inspect the shared AI-Q index. | +| `Search Index Data Contributor` | Upload, query, and delete index documents. | + +Assign the roles at the search-service scope. The principal ID is the object ID +of the system-assigned or user-assigned managed identity running AI-Q. + +Start the backend and UI with the checked-in Azure configuration: + +```bash +./scripts/start_e2e.sh --config_file configs/config_web_azure_ai_search.yml +``` + +`start_e2e.sh` sources `deploy/.env` before starting the backend, so +uncommented values in that file replace same-named shell exports. Keep the +Azure settings there, or leave them commented before relying on exported +values. + +`AZURE_SEARCH_API_KEY` selects API-key authentication when present; otherwise +the adapter uses `DefaultAzureCredential`. Set `AZURE_CLIENT_ID` to select a +user-assigned identity. Embeddings share `AIQ_EMBED_BASE_URL`, +`AIQ_EMBED_MODEL`, and `NVIDIA_API_KEY` with the LlamaIndex backend. +Azure-specific optional settings are `AIQ_EMBED_DIM` and +`AIQ_AZURE_SEARCH_INDEX_PREFIX`. The prefix must uniquely identify one AI-Q +deployment when a search service is shared. + +Changing the embedding model or dimension selects a different physical index +and requires re-ingestion. Frontend WebSocket queries use the conversation ID +as the collection; direct API tests must supply equivalent context or query the +configured fallback collection. + +The backend stores collection, file, and chunk records in one physical index +selected by `azure_search_index_prefix`, schema version, embedding model, and +dimension. Every operation applies internal collection filters. Retrieval is +always hybrid, and ingestion uses fixed 1024-token chunks with 128-token +overlap. File IDs returned by upload are authoritative for status and +delete operations; same-name uploads coexist independently. diff --git a/docs/source/examples/index.md b/docs/source/examples/index.md index 1fd9a263e..a1a916d6a 100644 --- a/docs/source/examples/index.md +++ b/docs/source/examples/index.md @@ -12,6 +12,7 @@ Complete, annotated configuration examples for common use cases. | [Minimal Shallow Only](./minimal-shallow-only.md) | Simplest setup — shallow research with web search | Custom minimal | | [Full Pipeline -- LlamaIndex](./full-pipeline-llamaindex.md) | Complete local setup with LlamaIndex + ChromaDB | `config_web_default_llamaindex.yml` | | [Full Pipeline -- Foundational RAG](./full-pipeline-web.md) | Complete production setup with hosted RAG | `config_web_frag.yml` | +| [Azure AI Search Knowledge Layer](./azure-ai-search.md) | Managed hybrid document retrieval | `config_web_azure_ai_search.yml` | | [CLI with Local NIMs](./cli-with-local-nims.md) | Interactive CLI mode with self-hosted NIM models | `config_cli_default.yml` | | [Hybrid Frontier Model](./hybrid-frontier-model.md) | NIM for shallow + frontier model for deep research | Custom hybrid | | [Deep Research Skills and Sandbox](./skills-sandbox/index.md) | DeepAgents skills with provider-backed sandbox execution for quantitative research workflows | `config_domain_routing_and_skills.yml` | diff --git a/docs/source/index.md b/docs/source/index.md index ec6498969..26ef008b0 100644 --- a/docs/source/index.md +++ b/docs/source/index.md @@ -142,6 +142,7 @@ FAQ <./resources/faq.md> ./examples/minimal-shallow-only.md ./examples/full-pipeline-llamaindex.md ./examples/full-pipeline-web.md +./examples/azure-ai-search.md ./examples/cli-with-local-nims.md ./examples/hybrid-frontier-model.md ./examples/skills-sandbox/index.md diff --git a/sources/knowledge_layer/KNOWLEDGE-LAYER-SETUP.md b/sources/knowledge_layer/KNOWLEDGE-LAYER-SETUP.md index 6adb54aff..af6d9fe69 100644 --- a/sources/knowledge_layer/KNOWLEDGE-LAYER-SETUP.md +++ b/sources/knowledge_layer/KNOWLEDGE-LAYER-SETUP.md @@ -9,7 +9,7 @@ A pluggable abstraction for document ingestion and retrieval. Swap backends with - **Collection Management** - create/delete/list collections per session or use case - **File Management** - upload/delete/list files with status tracking (UPLOADING → INGESTING → SUCCESS/FAILED) - **Content Typing** - TEXT, TABLE, CHART, IMAGE enums for frontend rendering -- **Backend Agnostic** - Swap between local (LlamaIndex), OpenSearch, and hosted RAG Blueprint without core agent code changes +- **Backend Agnostic** - Swap between local (LlamaIndex), OpenSearch, Azure AI Search, and hosted RAG Blueprint without core agent code changes --- @@ -36,6 +36,7 @@ A pluggable abstraction for document ingestion and retrieval. Swap backends with | `llamaindex` | `"llamaindex"` | Local Library | ChromaDB | Dev, prototyping, macOS/Linux | | `opensearch` | `"opensearch"` | Direct Client | OpenSearch k-NN | Self-hosted OpenSearch, Amazon OpenSearch Serverless | | `foundational_rag` | `"foundational_rag"` | Hosted Service | Remote Milvus | Production, multi-user | +| `azure_ai_search` | `"azure_ai_search"` | Managed Service | Azure AI Search | Managed hybrid retrieval | **Local Library Mode** - Everything runs in your Python process. No external services needed. - **`llamaindex`** - LlamaIndex + ChromaDB. Lightweight, great for development. Works on macOS and Linux. @@ -43,6 +44,9 @@ A pluggable abstraction for document ingestion and retrieval. Swap backends with **Hosted Service Mode** - Connects to deployed services via HTTP. Requires infrastructure but scales better. - **`foundational_rag`** - Connects to [NVIDIA RAG Blueprint](https://github.com/NVIDIA-AI-Blueprints/rag) via HTTP. - [Deployment Guide](https://github.com/NVIDIA-AI-Blueprints/rag/blob/main/docs/deploy-docker-self-hosted.md) +- **`azure_ai_search`** - Uses one AI-Q-owned shared index in a managed Azure AI Search service. Collection and file + manifests isolate logical collections. Canonical UUID file IDs support status and deletion, while same-name uploads + coexist independently. See [`src/azure_ai_search/README.md`](src/azure_ai_search/README.md). **OpenSearch Mode** - Stores AIQ collections directly in OpenSearch vector indexes. - **`opensearch`** - Uses one OpenSearch index per AIQ collection/session. Supports unauthenticated local clusters, @@ -64,6 +68,7 @@ export NVIDIA_API_KEY=nvapi-your-key-here uv pip install -e "sources/knowledge_layer[llamaindex]" # Recommended for local dev - works on macOS/Linux uv pip install -e "sources/knowledge_layer[foundational_rag]" # Requires deployed server uv pip install -e "sources/knowledge_layer[opensearch]" # Requires OpenSearch/OpenSearch Serverless +uv pip install -e "sources/knowledge_layer[azure_ai_search]" # Requires an Azure AI Search service ``` > **New to Knowledge Layer?** Start with `llamaindex` - it requires no external services and works on macOS and Linux. @@ -183,6 +188,7 @@ functions: > **Separate Docker stacks:** When AI-Q and RAG run as separate Docker Compose stacks, connect the AI-Q backend to the RAG network: `docker network connect nvidia-rag aiq-agent`. See the [Docker Compose README](../../deploy/compose/README.md#networking-when-aiq-and-rag-run-as-separate-compose-stacks) for details. **OpenSearch (Self-hosted)** + ```yaml functions: knowledge_search: @@ -1038,6 +1044,13 @@ Configuration values are resolved in the following order (highest to lowest prio | `KNOWLEDGE_RETRIEVER_BACKEND` | All | Default retriever backend (fallback if not in YAML) | | `KNOWLEDGE_INGESTOR_BACKEND` | All | Default ingestor backend (fallback if not in YAML) | | `AIQ_CHROMA_DIR` | llamaindex | ChromaDB persistence path | +| `AZURE_SEARCH_ENDPOINT` | azure_ai_search | Azure AI Search service endpoint | +| `AZURE_SEARCH_API_KEY` | azure_ai_search | Optional admin key; omit to use `DefaultAzureCredential` | +| `AZURE_CLIENT_ID` | azure_ai_search | Client ID for the user-assigned managed identity used by `DefaultAzureCredential` | +| `AIQ_AZURE_SEARCH_INDEX_PREFIX` | azure_ai_search | Deployment-unique prefix for the shared AI-Q index (default: `aiq`) | +| `AIQ_EMBED_MODEL` | llamaindex, azure_ai_search | Embedding model name | +| `AIQ_EMBED_BASE_URL` | llamaindex, azure_ai_search | Embedding API base URL | +| `AIQ_EMBED_DIM` | azure_ai_search | Embedding dimensions (default: `2048`) | | `AIQ_SUMMARY_DB` | All | Summary database URL (SQLite or PostgreSQL) | | `RAG_SERVER_URL` | foundational_rag | Query server URL (port 8081) | | `RAG_INGEST_URL` | foundational_rag | Ingestion server URL (port 8082) | diff --git a/sources/knowledge_layer/README.md b/sources/knowledge_layer/README.md index 53dd9b250..cc2dd1542 100644 --- a/sources/knowledge_layer/README.md +++ b/sources/knowledge_layer/README.md @@ -15,6 +15,9 @@ uv pip install -e "sources/knowledge_layer[foundational_rag]" # With OpenSearch (self-hosted or Amazon OpenSearch) uv pip install -e "sources/knowledge_layer[opensearch]" + +# With Azure AI Search +uv pip install -e "sources/knowledge_layer[azure_ai_search]" ``` ## Available Backends @@ -24,7 +27,11 @@ uv pip install -e "sources/knowledge_layer[opensearch]" | `llamaindex` | ChromaDB | Development, prototyping | | `opensearch` | OpenSearch k-NN | Self-hosted OpenSearch, Amazon OpenSearch Serverless | | `foundational_rag` | Remote Milvus | Production, multi-user | +| `azure_ai_search` | Azure AI Search | Managed hybrid search | ## Usage See [Web UI Mode](./KNOWLEDGE-LAYER-SETUP.md#web-ui-mode) for document upload and chat interfaces. + +Azure AI Search configuration and operational notes are in the +[backend README](./src/azure_ai_search/README.md). diff --git a/sources/knowledge_layer/pyproject.toml b/sources/knowledge_layer/pyproject.toml index 6a8767179..4823d87a5 100644 --- a/sources/knowledge_layer/pyproject.toml +++ b/sources/knowledge_layer/pyproject.toml @@ -18,7 +18,7 @@ build-backend = "setuptools.build_meta" requires = ["setuptools >= 64", "setuptools-scm>=8"] [tool.setuptools] -packages = ["knowledge_layer", "knowledge_layer.llamaindex", "knowledge_layer.foundational_rag", "knowledge_layer.opensearch"] +packages = ["knowledge_layer", "knowledge_layer.llamaindex", "knowledge_layer.foundational_rag", "knowledge_layer.opensearch", "knowledge_layer.azure_ai_search"] package-dir = {"knowledge_layer" = "src"} [project] @@ -61,8 +61,17 @@ opensearch = [ "python-pptx>=0.6.21", "distributed>=2024.1.0", ] +azure_ai_search = [ + "azure-identity>=1.15.0,<1.26", + "azure-search-documents>=11.6.0,<11.7", + "docx2txt>=0.8", + "llama-index-core>=0.11.0", + "llama-index-embeddings-nvidia>=0.2.0", + "llama-index-readers-file>=0.2.0", + "pypdf>=4.0.0", +] all = [ - "knowledge-layer[llamaindex,foundational_rag,opensearch]", + "knowledge-layer[llamaindex,foundational_rag,opensearch,azure_ai_search]", ] [project.entry-points."nat.plugins"] diff --git a/sources/knowledge_layer/src/__init__.py b/sources/knowledge_layer/src/__init__.py index 1ea7ed451..425481b4b 100644 --- a/sources/knowledge_layer/src/__init__.py +++ b/sources/knowledge_layer/src/__init__.py @@ -23,6 +23,7 @@ - llamaindex: LlamaIndex + ChromaDB (lightweight, local) - foundational_rag: Hosted NVIDIA RAG Blueprint (production, multi-user) - opensearch: OpenSearch vector search for self-hosted clusters and Amazon OpenSearch Serverless +- azure_ai_search: Azure AI Search with client-side embeddings Note: NAT tool registrations require NAT to be installed. The adapter modules can be used standalone without NAT. diff --git a/sources/knowledge_layer/src/azure_ai_search/README.md b/sources/knowledge_layer/src/azure_ai_search/README.md new file mode 100644 index 000000000..766c15d31 --- /dev/null +++ b/sources/knowledge_layer/src/azure_ai_search/README.md @@ -0,0 +1,81 @@ + + +# Azure AI Search backend + +This backend stores AI-Q document chunks and vectors in Azure AI Search. It +uses the shared `knowledge_retrieval` NAT function, Knowledge API, session +collection routing, summary store, and citation formatter. + +## Install + +```bash +uv pip install -e "sources/knowledge_layer[azure_ai_search]" +``` + +## Configure + +Set service and model credentials in the environment: + +```bash +export AZURE_SEARCH_ENDPOINT=https://.search.windows.net +export NVIDIA_API_KEY= +# Optional: setting this selects API-key auth instead of DefaultAzureCredential. +export AZURE_SEARCH_API_KEY= +``` + +```yaml +functions: + knowledge_search: + _type: knowledge_retrieval + backend: azure_ai_search + collection_name: ${COLLECTION_NAME:-aiq_default} + top_k: 5 + + generate_summary: true + summary_model: summary_llm + summary_db: ${AIQ_SUMMARY_DB:-sqlite+aiosqlite:///./summaries.db} +``` + +Explicit YAML values still override the environment-backed defaults. Azure +Search uses `AZURE_SEARCH_ENDPOINT` and optional `AZURE_SEARCH_API_KEY`. +Embedding configuration shares `AIQ_EMBED_BASE_URL`, `AIQ_EMBED_MODEL`, and +`NVIDIA_API_KEY` with the LlamaIndex backend; Azure additionally accepts +`AIQ_EMBED_DIM` and `AIQ_AZURE_SEARCH_INDEX_PREFIX`. The index prefix must be +unique to one AI-Q deployment sharing a search service. + +When `AZURE_SEARCH_API_KEY` is absent, the adapter uses +`DefaultAzureCredential`. Set `AZURE_CLIENT_ID` when a user-assigned identity +should be selected. Enable role-based access on the search service and grant +the identity `Search Service Contributor` for index management plus `Search +Index Data Contributor` for document ingestion and retrieval. Assign both roles +at the search-service scope because logical collections share one physical index. + +The adapter parses PDF, DOCX, TXT, and Markdown uploads with LlamaIndex, +creates one namespaced Azure AI Search index per deployment prefix, schema +version, embedding model, and dimension, and always performs balanced hybrid +retrieval. Collection and file manifests isolate logical collections in that +index. Documents use fixed 1024-token chunks with 128-token overlap. Only the +index carrying the matching AI-Q ownership/schema marker is visible or mutable. + +Upload responses return canonical UUID file IDs used by job progress, list, +status, and delete operations. Same-name uploads coexist independently under +different file IDs. Upload and delete requests stay below Azure's 1,000-action +and 16 MiB limits, and every per-document result is checked. + +Collections use the shared Knowledge Layer TTL settings: +`AIQ_COLLECTION_TTL_HOURS` defaults to 24 hours and +`AIQ_TTL_CLEANUP_INTERVAL_SECONDS` defaults to 3600 seconds. Successful file +and collection deletion also clears corresponding summary records. + +`embed_dim` must match both the embedding model output and the selected index. +Changing from a 2048-dimensional model to `nvidia/nv-embed-v1` at 4096 +dimensions selects a different physical index and requires re-ingestion. The +adapter validates ownership, fields, vector profile, and dimensions before use; +it does not alter an incompatible schema. + +For direct API tests, use the same collection or conversation context used for +upload. A standalone chat request without that context falls back to the +configured `collection_name`. diff --git a/sources/knowledge_layer/src/azure_ai_search/__init__.py b/sources/knowledge_layer/src/azure_ai_search/__init__.py new file mode 100644 index 000000000..6ca6bdb22 --- /dev/null +++ b/sources/knowledge_layer/src/azure_ai_search/__init__.py @@ -0,0 +1,9 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Azure AI Search Knowledge Layer backend.""" + +from .adapter import AzureAISearchIngestor +from .adapter import AzureAISearchRetriever + +__all__ = ["AzureAISearchIngestor", "AzureAISearchRetriever"] diff --git a/sources/knowledge_layer/src/azure_ai_search/adapter.py b/sources/knowledge_layer/src/azure_ai_search/adapter.py new file mode 100644 index 000000000..a9c2db159 --- /dev/null +++ b/sources/knowledge_layer/src/azure_ai_search/adapter.py @@ -0,0 +1,1344 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import asyncio +import json +import logging +import os +import re +import threading +import time +import uuid +from collections.abc import Iterator +from datetime import UTC +from datetime import datetime +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +from azure.core.credentials import AzureKeyCredential +from azure.core.exceptions import ClientAuthenticationError +from azure.core.exceptions import HttpResponseError +from azure.core.exceptions import ResourceNotFoundError +from azure.core.exceptions import ServiceRequestError +from azure.identity import DefaultAzureCredential +from azure.search.documents import SearchClient +from azure.search.documents.indexes import SearchIndexClient +from azure.search.documents.indexes.models import HnswAlgorithmConfiguration +from azure.search.documents.indexes.models import HnswParameters +from azure.search.documents.indexes.models import SearchableField +from azure.search.documents.indexes.models import SearchField +from azure.search.documents.indexes.models import SearchFieldDataType +from azure.search.documents.indexes.models import SearchIndex +from azure.search.documents.indexes.models import SimpleField +from azure.search.documents.indexes.models import VectorSearch +from azure.search.documents.indexes.models import VectorSearchAlgorithmMetric +from azure.search.documents.indexes.models import VectorSearchProfile +from azure.search.documents.models import VectorizedQuery + +from aiq_agent.knowledge import BaseIngestor +from aiq_agent.knowledge import BaseRetriever +from aiq_agent.knowledge import Chunk +from aiq_agent.knowledge import ContentType +from aiq_agent.knowledge import FileProgress +from aiq_agent.knowledge import IngestionJobStatus +from aiq_agent.knowledge import JobState +from aiq_agent.knowledge import RetrievalResult +from aiq_agent.knowledge import clear_collection_summaries +from aiq_agent.knowledge import register_ingestor +from aiq_agent.knowledge import register_retriever +from aiq_agent.knowledge import register_summary +from aiq_agent.knowledge import unregister_summary +from aiq_agent.knowledge.base import CollectionInfo +from aiq_agent.knowledge.base import FileInfo +from aiq_agent.knowledge.base import TTLCleanupMixin +from aiq_agent.knowledge.schema import FileStatus + +logger = logging.getLogger(__name__) + +_BACKEND_NAME = "azure_ai_search" +_SCHEMA_VERSION = 1 +_MARKER_PREFIX = "aiq.azure_ai_search:" +_MARKER_MAX_CHARS = 4000 +_MAX_INDEX_NAME_LENGTH = 128 +_MAX_BATCH_ACTIONS = 1000 +_MAX_BATCH_BYTES = 16 * 1024 * 1024 +_PAGE_SIZE = 1000 +_DELETE_ATTEMPTS = 3 +_CONSISTENCY_ATTEMPTS = 20 +_CONSISTENCY_DELAY_SECONDS = 0.25 +_CHUNK_SIZE = 1024 +_CHUNK_OVERLAP = 128 +_SUMMARY_MAX_CHARS = 1000 +_RECORD_COLLECTION = "collection" +_RECORD_FILE = "file" +_RECORD_CHUNK = "chunk" +_COLLECTION_ACTIVE = "active" +_COLLECTION_DELETING = "deleting" + +COLLECTION_TTL_HOURS = float(os.environ.get("AIQ_COLLECTION_TTL_HOURS", "24")) +TTL_CLEANUP_INTERVAL_SECONDS = int(os.environ.get("AIQ_TTL_CLEANUP_INTERVAL_SECONDS", "3600")) + + +def _coerce_config(config: dict[str, Any] | None) -> SimpleNamespace: + """Apply adapter defaults so direct factory usage matches YAML usage.""" + values: dict[str, Any] = { + "endpoint": os.environ.get("AZURE_SEARCH_ENDPOINT"), + "api_key": os.environ.get("AZURE_SEARCH_API_KEY"), + "embed_base_url": os.environ.get("AIQ_EMBED_BASE_URL") or "https://integrate.api.nvidia.com/v1", + "embed_model": os.environ.get("AIQ_EMBED_MODEL", "nvidia/llama-nemotron-embed-vl-1b-v2"), + "embed_dim": int(os.environ.get("AIQ_EMBED_DIM", "2048")), + "collection_name": "default", + "cleanup_files": False, + "generate_summary": False, + "summary_llm": None, + "index_prefix": os.environ.get("AIQ_AZURE_SEARCH_INDEX_PREFIX", "aiq"), + "start_ttl_cleanup": True, + } + values.update({key: value for key, value in (config or {}).items() if key in values}) + if not values["endpoint"]: + raise ValueError("Azure AI Search configuration requires `endpoint`") + return SimpleNamespace(**values) + + +def _build_search_credential(cfg: SimpleNamespace): + """Pick the Azure SDK credential without exposing secret values.""" + api_key = _secret_value(cfg.api_key) + return AzureKeyCredential(api_key) if api_key else DefaultAzureCredential() + + +def _secret_value(value: Any) -> str | None: + if value is None: + return None + if hasattr(value, "get_secret_value"): + return value.get_secret_value() + return str(value) + + +def _utc_now() -> datetime: + return datetime.now(UTC) + + +def _parse_timestamp(value: Any) -> datetime | None: + if isinstance(value, datetime): + return value if value.tzinfo else value.replace(tzinfo=UTC) + if isinstance(value, str): + try: + parsed = datetime.fromisoformat(value.replace("Z", "+00:00")) + return parsed if parsed.tzinfo else parsed.replace(tzinfo=UTC) + except ValueError: + return None + return None + + +def _sanitize_index_part(value: str, fallback: str = "default") -> str: + normalized = re.sub(r"[^a-z0-9]+", "-", value.lower()).strip("-") + return normalized or fallback + + +def _index_name_for_config(prefix: str, embed_model: str, embed_dim: int) -> str: + """Map one deployment and embedding schema to one physical Azure index.""" + suffix = uuid.uuid5( + uuid.NAMESPACE_URL, + f"{_SCHEMA_VERSION}\0{prefix}\0{embed_model}\0{embed_dim}", + ).hex[:12] + tail = f"knowledge-v{_SCHEMA_VERSION}-{suffix}" + available = _MAX_INDEX_NAME_LENGTH - len(tail) - 1 + prefix_part = _sanitize_index_part(prefix, "aiq")[:available].rstrip("-") or "aiq" + return f"{prefix_part}-{tail}" + + +def _validate_index_name(name: str) -> None: + if not re.fullmatch(r"[a-z0-9](?:[a-z0-9-]{0,126}[a-z0-9])?", name) or "--" in name: + raise ValueError( + f"Invalid Azure AI Search index name {name!r}. Names must use lowercase letters, digits, and single " + "hyphens; be 2-128 characters; and start and end with a letter or digit." + ) + + +def _encode_marker(marker: dict[str, Any]) -> str: + encoded = _MARKER_PREFIX + json.dumps(marker, separators=(",", ":"), sort_keys=True) + if len(encoded) > _MARKER_MAX_CHARS: + raise ValueError(f"Azure AI Search ownership marker exceeds {_MARKER_MAX_CHARS} characters") + return encoded + + +def _decode_marker(description: str | None) -> dict[str, Any] | None: + if not description or not description.startswith(_MARKER_PREFIX): + return None + try: + marker = json.loads(description[len(_MARKER_PREFIX) :]) + except (TypeError, ValueError): + return None + return marker if isinstance(marker, dict) else None + + +def _new_marker(cfg: SimpleNamespace) -> dict[str, Any]: + marker = { + "backend": _BACKEND_NAME, + "schema_version": _SCHEMA_VERSION, + "index_prefix": cfg.index_prefix, + "embedding_model": cfg.embed_model, + "embedding_dim": cfg.embed_dim, + "created_at": _utc_now().isoformat(), + } + _encode_marker(marker) + return marker + + +def _record_id(record_type: str, collection_name: str, file_id: str | None = None) -> str: + value = f"{record_type}\0{collection_name}\0{file_id or ''}" + return f"{record_type}-{uuid.uuid5(uuid.NAMESPACE_URL, value).hex}" + + +def _odata_literal(value: str) -> str: + return "'" + value.replace("'", "''") + "'" + + +def _and_filter(*filters: str | None) -> str | None: + values = [f"({value})" for value in filters if value] + return " and ".join(values) or None + + +def _record_filter( + record_type: str, + collection_name: str | None = None, + *, + file_id: str | None = None, + file_name: str | None = None, + status: str | None = None, +) -> str: + return ( + _and_filter( + f"record_type eq {_odata_literal(record_type)}", + f"collection_id eq {_odata_literal(collection_name)}" if collection_name is not None else None, + f"file_id eq {_odata_literal(file_id)}" if file_id is not None else None, + f"file_name eq {_odata_literal(file_name)}" if file_name is not None else None, + f"status eq {_odata_literal(status)}" if status is not None else None, + ) + or "" + ) + + +def _parse_metadata(value: Any) -> dict[str, Any]: + if isinstance(value, dict): + return dict(value) + if not value: + return {} + try: + parsed = json.loads(str(value)) + except (TypeError, ValueError): + logger.warning("Ignoring malformed Azure AI Search metadata JSON") + return {} + return parsed if isinstance(parsed, dict) else {} + + +def _json_default(value: Any) -> str: + if isinstance(value, datetime): + return value.isoformat() + raise TypeError(f"Object of type {type(value).__name__} is not JSON serializable") + + +def _iter_index_batches( + documents: list[dict[str, Any]], + action: str, + max_actions: int = _MAX_BATCH_ACTIONS, + max_bytes: int = _MAX_BATCH_BYTES, +) -> Iterator[list[dict[str, Any]]]: + """Batch actions below both Azure's count and serialized payload limits.""" + batch: list[dict[str, Any]] = [] + batch_bytes = len(b'{"value":[]}') + for document in documents: + payload = {"@search.action": action, **document} + action_bytes = ( + len(json.dumps(payload, default=_json_default, ensure_ascii=False, separators=(",", ":")).encode("utf-8")) + + 1 + ) + if action_bytes + len(b'{"value":[]}') > max_bytes: + raise ValueError(f"Azure AI Search {action} action exceeds the 16 MiB request limit") + if batch and (len(batch) >= max_actions or batch_bytes + action_bytes > max_bytes): + yield batch + batch = [] + batch_bytes = len(b'{"value":[]}') + batch.append(document) + batch_bytes += action_bytes + if batch: + yield batch + + +def _indexing_outcome(results: list[Any], expected: list[dict[str, Any]]) -> tuple[list[str], list[str]]: + expected_keys = [str(document["id"]) for document in expected] + by_key = {str(getattr(result, "key", "")): result for result in results} + succeeded: list[str] = [] + failures: list[str] = [] + for key in expected_keys: + result = by_key.get(key) + if result is not None and getattr(result, "succeeded", False): + succeeded.append(key) + else: + message = getattr(result, "error_message", None) if result is not None else "missing result" + failures.append(f"{key}: {message or 'rejected'}") + return succeeded, failures + + +def _build_index_schema(name: str, embed_dim: int, description: str | None = None) -> SearchIndex: + return SearchIndex( + name=name, + description=description, + fields=[ + SimpleField( + name="id", + type=SearchFieldDataType.String, + key=True, + filterable=True, + sortable=True, + ), + SimpleField( + name="record_type", + type=SearchFieldDataType.String, + filterable=True, + facetable=True, + sortable=True, + ), + SimpleField(name="collection_id", type=SearchFieldDataType.String, filterable=True, sortable=True), + SearchableField(name="chunk", analyzer_name="standard.lucene"), + SearchField( + name="embedding", + type=SearchFieldDataType.Collection(SearchFieldDataType.Single), + searchable=True, + vector_search_dimensions=embed_dim, + vector_search_profile_name="hnsw-profile", + ), + SimpleField(name="file_id", type=SearchFieldDataType.String, filterable=True, sortable=True), + SearchableField(name="file_name", filterable=True, sortable=True), + SimpleField(name="status", type=SearchFieldDataType.String, filterable=True, sortable=True), + SimpleField(name="error_message", type=SearchFieldDataType.String), + SimpleField(name="description", type=SearchFieldDataType.String), + SimpleField(name="summary", type=SearchFieldDataType.String), + SimpleField(name="page_number", type=SearchFieldDataType.Int32, filterable=True, sortable=True), + SimpleField(name="chunk_index", type=SearchFieldDataType.Int32, filterable=True, sortable=True), + SimpleField(name="chunk_count", type=SearchFieldDataType.Int32), + SimpleField(name="file_size", type=SearchFieldDataType.Int64, filterable=True), + SimpleField(name="created_at", type=SearchFieldDataType.DateTimeOffset, filterable=True, sortable=True), + SimpleField(name="updated_at", type=SearchFieldDataType.DateTimeOffset, filterable=True, sortable=True), + SimpleField(name="uploaded_at", type=SearchFieldDataType.DateTimeOffset, filterable=True, sortable=True), + SimpleField(name="ingested_at", type=SearchFieldDataType.DateTimeOffset, filterable=True, sortable=True), + SimpleField(name="metadata", type=SearchFieldDataType.String), + ], + vector_search=VectorSearch( + algorithms=[ + HnswAlgorithmConfiguration( + name="hnsw-default", + parameters=HnswParameters( + m=4, + ef_construction=400, + ef_search=500, + metric=VectorSearchAlgorithmMetric.COSINE, + ), + ), + ], + profiles=[VectorSearchProfile(name="hnsw-profile", algorithm_configuration_name="hnsw-default")], + ), + ) + + +def _validate_index_schema(index: SearchIndex, cfg: SimpleNamespace) -> dict[str, Any]: + marker = _decode_marker(index.description) + if marker is None: + raise RuntimeError(f"Azure AI Search index {index.name!r} is not owned by AI-Q") + if marker.get("backend") != _BACKEND_NAME or marker.get("schema_version") != _SCHEMA_VERSION: + raise RuntimeError(f"Azure AI Search index {index.name!r} has an incompatible AI-Q ownership marker") + if ( + marker.get("index_prefix") != cfg.index_prefix + or marker.get("embedding_dim") != cfg.embed_dim + or marker.get("embedding_model") != cfg.embed_model + ): + raise RuntimeError(f"Azure AI Search index {index.name!r} ownership or embedding configuration does not match") + + fields = {field.name: field for field in index.fields} + required_types = { + "id": SearchFieldDataType.String, + "record_type": SearchFieldDataType.String, + "collection_id": SearchFieldDataType.String, + "chunk": SearchFieldDataType.String, + "embedding": SearchFieldDataType.Collection(SearchFieldDataType.Single), + "file_id": SearchFieldDataType.String, + "file_name": SearchFieldDataType.String, + "status": SearchFieldDataType.String, + "error_message": SearchFieldDataType.String, + "description": SearchFieldDataType.String, + "summary": SearchFieldDataType.String, + "page_number": SearchFieldDataType.Int32, + "chunk_index": SearchFieldDataType.Int32, + "chunk_count": SearchFieldDataType.Int32, + "file_size": SearchFieldDataType.Int64, + "created_at": SearchFieldDataType.DateTimeOffset, + "updated_at": SearchFieldDataType.DateTimeOffset, + "uploaded_at": SearchFieldDataType.DateTimeOffset, + "ingested_at": SearchFieldDataType.DateTimeOffset, + "metadata": SearchFieldDataType.String, + } + missing = [name for name in required_types if name not in fields] + mismatched = [ + name for name, field_type in required_types.items() if name in fields and fields[name].type != field_type + ] + if missing or mismatched: + raise RuntimeError( + f"Azure AI Search index {index.name!r} schema mismatch: missing={missing}, wrong_type={mismatched}" + ) + if not fields["id"].key or not fields["id"].filterable or not fields["id"].sortable: + raise RuntimeError(f"Azure AI Search index {index.name!r} requires id to be key/filterable/sortable") + if not fields["record_type"].filterable or not fields["record_type"].facetable: + raise RuntimeError(f"Azure AI Search index {index.name!r} requires facetable record_type") + if not fields["collection_id"].filterable or not fields["file_id"].filterable: + raise RuntimeError(f"Azure AI Search index {index.name!r} requires filterable identity fields") + embedding = fields["embedding"] + if embedding.vector_search_dimensions != cfg.embed_dim or embedding.vector_search_profile_name != "hnsw-profile": + raise RuntimeError(f"Azure AI Search index {index.name!r} vector profile or dimensions do not match") + profile_names = {profile.name for profile in (index.vector_search.profiles if index.vector_search else [])} + if "hnsw-profile" not in profile_names: + raise RuntimeError(f"Azure AI Search index {index.name!r} is missing hnsw-profile") + return marker + + +class _AzureIndexMixin: + cfg: SimpleNamespace + _credential: Any + _embedding: Any + _index_client: SearchIndexClient + _search_client: SearchClient | None + _index_validated: bool + + def _initialize_azure(self, config: dict[str, Any]) -> None: + self.cfg = _coerce_config(config) + self._credential = _build_search_credential(self.cfg) + self._index_client = SearchIndexClient(endpoint=str(self.cfg.endpoint), credential=self._credential) + self._search_client = None + self._index_validated = False + self._embedding = None + + @property + def embedding(self): + if self._embedding is None: + from llama_index.embeddings.nvidia import NVIDIAEmbedding + + self._embedding = NVIDIAEmbedding( + model=self.cfg.embed_model, + base_url=str(self.cfg.embed_base_url), + ) + return self._embedding + + def _physical_index_name(self) -> str: + name = _index_name_for_config(self.cfg.index_prefix, self.cfg.embed_model, self.cfg.embed_dim) + _validate_index_name(name) + return name + + def _get_search_client(self) -> SearchClient: + if self._search_client is None: + self._search_client = self._index_client.get_search_client(self._physical_index_name()) + return self._search_client + + def _get_owned_index(self) -> tuple[SearchIndex, dict[str, Any]]: + index = self._index_client.get_index(self._physical_index_name()) + return index, _validate_index_schema(index, self.cfg) + + def _ensure_index(self) -> tuple[SearchIndex, dict[str, Any]]: + try: + index, marker = self._get_owned_index() + except ResourceNotFoundError: + marker = _new_marker(self.cfg) + schema = _build_index_schema(self._physical_index_name(), self.cfg.embed_dim, _encode_marker(marker)) + try: + index = self._index_client.create_index(schema) + except Exception as create_error: # noqa: BLE001 + try: + index = self._index_client.get_index(self._physical_index_name()) + except ResourceNotFoundError: + raise create_error + marker = _validate_index_schema(index, self.cfg) + self._index_validated = True + return index, marker + + def _get_validated_search_client(self) -> SearchClient: + if not self._index_validated: + self._get_owned_index() + self._index_validated = True + return self._get_search_client() + + def _get_document(self, document_id: str) -> dict[str, Any] | None: + try: + return dict(self._get_validated_search_client().get_document(key=document_id)) + except ResourceNotFoundError: + return None + + def _get_collection_manifest(self, collection_name: str) -> dict[str, Any] | None: + document = self._get_document(_record_id(_RECORD_COLLECTION, collection_name)) + if ( + document + and document.get("record_type") == _RECORD_COLLECTION + and document.get("collection_id") == collection_name + ): + return document + return None + + +@register_retriever(_BACKEND_NAME) +class AzureAISearchRetriever(_AzureIndexMixin, BaseRetriever): + """Hybrid retriever backed by one owned Azure index.""" + + backend_name = _BACKEND_NAME + + def __init__(self, config: dict[str, Any] | None = None): + super().__init__(config) + self._initialize_azure(self.config) + + def _get_client(self, collection_name: str) -> SearchClient: + if self._get_collection_manifest(collection_name) is None: + raise ResourceNotFoundError(f"Collection {collection_name!r} not found") + return self._get_validated_search_client() + + async def retrieve( + self, + query: str, + collection_name: str, + top_k: int = 10, + filters: dict[str, Any] | None = None, + ) -> RetrievalResult: + return await asyncio.to_thread(self._retrieve_sync, query, collection_name, top_k, filters) + + def _retrieve_sync( + self, + query: str, + collection_name: str, + top_k: int, + filters: dict[str, Any] | None, + ) -> RetrievalResult: + if filters: + return RetrievalResult( + query=query, + backend=_BACKEND_NAME, + chunks=[], + success=False, + error_message="Azure AI Search metadata filters are not supported", + ) + try: + client = self._get_client(collection_name) + query_vector = self.embedding.get_query_embedding(query) + vector_query = VectorizedQuery( + vector=query_vector, + k_nearest_neighbors=max(top_k * 3, 20), + fields="embedding", + ) + search_params: dict[str, Any] = { + "search_text": query, + "vector_queries": [vector_query], + "filter": _record_filter(_RECORD_CHUNK, collection_name), + "top": top_k, + "select": ["id", "chunk", "file_id", "file_name", "page_number", "metadata"], + } + chunks = [self.normalize(hit) for hit in client.search(**search_params)] + return RetrievalResult(query=query, backend=_BACKEND_NAME, chunks=chunks, success=True) + except ResourceNotFoundError: + message = f"AI-Q Azure AI Search collection {collection_name!r} not found" + except ClientAuthenticationError as error: + message = f"AI Search authentication failed: {error!s}" + except ServiceRequestError as error: + message = f"AI Search service unavailable: {error!s}" + except HttpResponseError as error: + message = f"AI Search request failed: {error.status_code} {error.reason or error!s}" + except Exception as error: # noqa: BLE001 + message = f"Unexpected error during retrieval: {error!s}" + logger.exception("Azure AI Search retrieval failed") + return RetrievalResult(query=query, backend=_BACKEND_NAME, chunks=[], success=False, error_message=message) + + def normalize(self, raw_result: Any) -> Chunk: + chunk_id = raw_result.get("id") or str(uuid.uuid4()) + content = raw_result.get("chunk") or "" + file_name = raw_result.get("file_name") or "unknown" + page_number = _coerce_page_number(raw_result.get("page_number")) + search_score = raw_result.get("@search.score") or 0.0 + score = float(search_score) + score = min(max(score, 0.0), 1.0) + metadata = _parse_metadata(raw_result.get("metadata")) + if file_id := raw_result.get("file_id"): + metadata["file_id"] = file_id + return Chunk( + chunk_id=str(chunk_id), + content=str(content), + score=score, + file_name=str(file_name), + page_number=page_number, + display_citation=f"{file_name}, p.{page_number}" if page_number else str(file_name), + content_type=ContentType.TEXT, + metadata=metadata, + ) + + async def health_check(self) -> bool: + def _check() -> bool: + list(self._index_client.list_index_names()) + return True + + try: + return await asyncio.to_thread(_check) + except Exception: # noqa: BLE001 + logger.exception("Azure AI Search retriever health_check failed") + return False + + +@register_ingestor(_BACKEND_NAME) +class AzureAISearchIngestor(TTLCleanupMixin, _AzureIndexMixin, BaseIngestor): + """Parse, embed, and persist documents in owned Azure AI Search indexes.""" + + backend_name = _BACKEND_NAME + + def __init__(self, config: dict[str, Any] | None = None): + super().__init__(config) + self._initialize_azure(self.config) + self._splitter = None + self._summary_llm = self.cfg.summary_llm + self._jobs_lock = threading.RLock() + self._jobs: dict[str, IngestionJobStatus] = {} + self._files: dict[str, FileInfo] = {} + self._deleted_collections: set[str] = set() + self._deleted_files: set[tuple[str, str]] = set() + if self.cfg.start_ttl_cleanup: + self._start_ttl_cleanup_task(COLLECTION_TTL_HOURS, TTL_CLEANUP_INTERVAL_SECONDS) + + @property + def splitter(self): + if self._splitter is None: + from llama_index.core.node_parser import SentenceSplitter + + self._splitter = SentenceSplitter(chunk_size=_CHUNK_SIZE, chunk_overlap=_CHUNK_OVERLAP) + return self._splitter + + def _write_document(self, document: dict[str, Any]) -> None: + self._upload_documents(self._get_validated_search_client(), [document]) + + def _wait_for_search_state(self, filter_text: str, *, present: bool, label: str) -> None: + client = self._get_validated_search_client() + for attempt in range(_CONSISTENCY_ATTEMPTS): + found = bool(list(client.search(search_text="*", filter=filter_text, select=["id"], top=1))) + if found == present: + return + if attempt + 1 < _CONSISTENCY_ATTEMPTS: + time.sleep(_CONSISTENCY_DELAY_SECONDS) + state = "visible" if present else "absent" + raise RuntimeError(f"Timed out waiting for {label} to become {state} in Azure AI Search") + + def _write_collection_manifest( + self, + collection_name: str, + description: str | None = None, + metadata: dict[str, Any] | None = None, + *, + status: str = _COLLECTION_ACTIVE, + ) -> dict[str, Any]: + self._ensure_index() + existing = self._get_collection_manifest(collection_name) or {} + if status == _COLLECTION_ACTIVE and existing.get("status") == _COLLECTION_DELETING: + raise ValueError(f"Collection {collection_name!r} is being deleted") + now = _utc_now() + existing_metadata = _parse_metadata(existing.get("metadata")) + merged_metadata = {**existing_metadata, **(metadata or {})} + document = { + "id": _record_id(_RECORD_COLLECTION, collection_name), + "record_type": _RECORD_COLLECTION, + "collection_id": collection_name, + "status": status, + "description": description if description is not None else existing.get("description"), + "metadata": json.dumps(merged_metadata, separators=(",", ":"), sort_keys=True), + "created_at": _parse_timestamp(existing.get("created_at")) or now, + "updated_at": now, + } + self._write_document(document) + if status == _COLLECTION_ACTIVE: + self._deleted_collections.discard(collection_name) + return document + + def _update_collection_timestamp(self, collection_name: str) -> None: + manifest = self._get_collection_manifest(collection_name) + if manifest is None or manifest.get("status") != _COLLECTION_ACTIVE: + raise ResourceNotFoundError(f"Collection {collection_name!r} not found") + manifest["updated_at"] = _utc_now() + self._write_document(manifest) + + def _get_file_manifest(self, file_id: str, collection_name: str) -> dict[str, Any] | None: + documents = list( + self._get_validated_search_client().search( + search_text="*", + filter=_record_filter(_RECORD_FILE, collection_name, file_id=file_id), + top=1, + ) + ) + return dict(documents[0]) if documents else None + + def _write_file_manifest(self, info: FileInfo, *, summary: str | None = None) -> dict[str, Any]: + metadata = {key: value for key, value in info.metadata.items() if key not in {"job_id", "summary"}} + document = { + "id": _record_id(_RECORD_FILE, info.collection_name, info.file_id), + "record_type": _RECORD_FILE, + "collection_id": info.collection_name, + "file_id": info.file_id, + "file_name": info.file_name, + "status": info.status.value, + "error_message": info.error_message, + "summary": summary or info.metadata.get("summary"), + "chunk_count": info.chunk_count, + "file_size": info.file_size, + "uploaded_at": info.uploaded_at, + "ingested_at": info.ingested_at, + "metadata": json.dumps(metadata, separators=(",", ":"), sort_keys=True), + } + self._write_document(document) + self._deleted_files.discard((info.collection_name, info.file_id)) + return document + + @staticmethod + def _file_info_from_manifest(document: dict[str, Any]) -> FileInfo: + metadata = _parse_metadata(document.get("metadata")) + if summary := document.get("summary"): + metadata["summary"] = summary + try: + status = FileStatus(document.get("status")) + except ValueError: + status = FileStatus.FAILED + return FileInfo( + file_id=str(document.get("file_id") or ""), + file_name=str(document.get("file_name") or "unknown"), + collection_name=str(document.get("collection_id") or ""), + status=status, + error_message=document.get("error_message"), + file_size=document.get("file_size"), + chunk_count=int(document.get("chunk_count") or 0), + uploaded_at=_parse_timestamp(document.get("uploaded_at")), + ingested_at=_parse_timestamp(document.get("ingested_at")), + metadata=metadata, + ) + + def _collection_counts(self, collection_name: str) -> tuple[int, int]: + results = self._get_validated_search_client().search( + search_text="*", + filter=f"collection_id eq {_odata_literal(collection_name)}", + facets=["record_type,count:0"], + top=0, + ) + facets = results.get_facets() or {} + counts = {str(item.get("value")): int(item.get("count") or 0) for item in facets.get("record_type", [])} + return counts.get(_RECORD_FILE, 0), counts.get(_RECORD_CHUNK, 0) + + def _collection_info(self, manifest: dict[str, Any]) -> CollectionInfo: + name = str(manifest["collection_id"]) + file_count, chunk_count = self._collection_counts(name) + return CollectionInfo( + name=name, + description=manifest.get("description"), + file_count=file_count, + chunk_count=chunk_count, + backend=_BACKEND_NAME, + metadata={ + **_parse_metadata(manifest.get("metadata")), + "index_name": self._physical_index_name(), + "embedding_model": self.cfg.embed_model, + "embedding_dim": self.cfg.embed_dim, + }, + created_at=_parse_timestamp(manifest.get("created_at")), + updated_at=_parse_timestamp(manifest.get("updated_at")), + ) + + def submit_job( + self, + file_paths: list[str], + collection_name: str, + config: dict[str, Any] | None = None, + ) -> str: + job_id = str(uuid.uuid4()) + job_config = {**self.config, **(config or {})} + original_filenames = _resolve_filenames(file_paths, job_config.get("original_filenames")) + validated = [(path, original_filenames[index]) for index, path in enumerate(file_paths) if Path(path).is_file()] + file_metadata = job_config.get("metadata") or {} + json.dumps(file_metadata, separators=(",", ":"), sort_keys=True) + + if not validated: + with self._jobs_lock: + self._jobs[job_id] = IngestionJobStatus( + job_id=job_id, + status=JobState.FAILED, + collection_name=collection_name, + backend=_BACKEND_NAME, + submitted_at=_utc_now(), + completed_at=_utc_now(), + total_files=len(file_paths), + error_message="No valid file paths provided", + ) + return job_id + + submitted_at = _utc_now() + file_progress: list[FileProgress] = [] + files: dict[str, FileInfo] = {} + for path, file_name in validated: + file_id = str(uuid.uuid4()) + file_progress.append(FileProgress(file_id=file_id, file_name=file_name, status=FileStatus.UPLOADING)) + files[file_id] = FileInfo( + file_id=file_id, + file_name=file_name, + collection_name=collection_name, + status=FileStatus.UPLOADING, + file_size=Path(path).stat().st_size, + uploaded_at=submitted_at, + metadata={**file_metadata, "job_id": job_id}, + ) + + with self._jobs_lock: + self._files.update(files) + self._jobs[job_id] = IngestionJobStatus( + job_id=job_id, + status=JobState.PENDING, + collection_name=collection_name, + backend=_BACKEND_NAME, + submitted_at=submitted_at, + total_files=len(validated), + file_details=file_progress, + ) + + try: + threading.Thread( + target=self._process_job, + args=(job_id, [path for path, _ in validated], collection_name, job_config), + daemon=True, + name=f"aiq-azure-search-ingest-{job_id[:8]}", + ).start() + except Exception as error: # noqa: BLE001 + message = f"Failed to start ingestion worker: {self._translate_error(error)}" + self._fail_job_setup(job_id, message) + if job_config.get("cleanup_files", self.cfg.cleanup_files): + self._cleanup_paths([path for path, _ in validated]) + logger.exception("Failed to start Azure AI Search ingestion job %s", job_id) + return job_id + + def get_job_status(self, job_id: str) -> IngestionJobStatus: + with self._jobs_lock: + job = self._jobs.get(job_id) + if job is not None: + return job.model_copy(deep=True) + return IngestionJobStatus( + job_id=job_id, + status=JobState.FAILED, + collection_name="", + backend=_BACKEND_NAME, + submitted_at=_utc_now(), + completed_at=_utc_now(), + error_message="Job not found", + ) + + def _process_job( + self, + job_id: str, + file_paths: list[str], + collection_name: str, + config: dict[str, Any], + ) -> None: + cleanup = bool(config.get("cleanup_files", self.cfg.cleanup_files)) + self._update_job(job_id, status=JobState.PROCESSING, started_at=_utc_now()) + client: SearchClient | None = None + try: + self._ensure_index() + client = self._get_validated_search_client() + collection = self._get_collection_manifest(collection_name) + if collection is None: + self._write_collection_manifest(collection_name) + elif collection.get("status") != _COLLECTION_ACTIVE: + raise ValueError(f"Collection {collection_name!r} is being deleted") + + job = self.get_job_status(job_id) + for detail in job.file_details: + self._write_file_manifest(self._files[detail.file_id]) + except Exception as error: # noqa: BLE001 + message = self._translate_error(error) + if client is not None: + try: + job = self.get_job_status(job_id) + self._delete_document_ids( + client, + [_record_id(_RECORD_FILE, collection_name, detail.file_id) for detail in job.file_details], + ) + except Exception as rollback_error: # noqa: BLE001 + message = f"{message}; manifest rollback failed: {self._translate_error(rollback_error)}" + self._fail_job_setup(job_id, message) + if cleanup: + self._cleanup_paths(file_paths) + logger.exception("Failed to initialize Azure AI Search ingestion job %s", job_id) + return + + failed = 0 + for index, path in enumerate(file_paths): + job = self.get_job_status(job_id) + detail = job.file_details[index] + tracked = self._files[detail.file_id] + self._update_file_progress(job_id, index, status=FileStatus.INGESTING) + try: + chunk_count, summary, ingested_at = self._process_file( + path=path, + collection_name=collection_name, + file_id=detail.file_id, + file_name=detail.file_name, + file_size=tracked.file_size or 0, + uploaded_at=tracked.uploaded_at or _utc_now(), + metadata={key: value for key, value in tracked.metadata.items() if key != "job_id"}, + ) + self._update_file_progress( + job_id, + index, + status=FileStatus.SUCCESS, + progress_percent=100.0, + chunks_created=chunk_count, + summary=summary, + ingested_at=ingested_at, + ) + except Exception as error: # noqa: BLE001 + failed += 1 + message = self._translate_error(error) + try: + self._delete_file_documents(detail.file_id, collection_name) + except Exception as rollback_error: # noqa: BLE001 + message = f"{message}; chunk rollback failed: {self._translate_error(rollback_error)}" + self._update_file_progress( + job_id, + index, + status=FileStatus.FAILED, + progress_percent=0.0, + chunks_created=0, + error_message=message, + ) + logger.exception("Failed to ingest %s", detail.file_name) + finally: + if cleanup: + self._cleanup_paths([path]) + self._update_job(job_id, processed_files=index + 1) + + if failed == len(file_paths): + self._fail_job(job_id, f"All {failed} file(s) failed to ingest") + else: + self._update_job(job_id, status=JobState.COMPLETED, completed_at=_utc_now()) + + def _process_file( + self, + *, + path: str, + collection_name: str, + file_id: str, + file_name: str, + file_size: int, + uploaded_at: datetime, + metadata: dict[str, Any], + ) -> tuple[int, str | None, datetime]: + from llama_index.core import SimpleDirectoryReader + + documents = SimpleDirectoryReader(input_files=[path]).load_data() + if not documents: + raise ValueError(f"No content extracted from {file_name}") + nodes = self.splitter.get_nodes_from_documents(documents) + if not nodes: + raise ValueError(f"Chunking produced 0 chunks for {file_name}") + texts = [node.get_content() for node in nodes] + embeddings = self.embedding.get_text_embedding_batch(texts) + if any(len(vector) != self.cfg.embed_dim for vector in embeddings): + raise ValueError(f"Embedding dimensions do not match configured embed_dim={self.cfg.embed_dim}") + + ingested_at = _utc_now() + encoded_metadata = json.dumps(metadata, separators=(",", ":"), sort_keys=True) + search_documents: list[dict[str, Any]] = [] + for chunk_index, (node, vector) in enumerate(zip(nodes, embeddings, strict=True)): + search_documents.append( + { + "id": f"chunk-{file_id}-{chunk_index:08d}", + "record_type": _RECORD_CHUNK, + "collection_id": collection_name, + "chunk": node.get_content(), + "embedding": list(vector), + "file_id": file_id, + "file_name": file_name, + "page_number": _coerce_page_number(node.metadata.get("page_label")), + "chunk_index": chunk_index, + "file_size": file_size, + "uploaded_at": uploaded_at, + "ingested_at": ingested_at, + "metadata": encoded_metadata, + } + ) + + client = self._get_validated_search_client() + self._upload_documents(client, search_documents) + collection = self._get_collection_manifest(collection_name) + if collection is None or collection.get("status") != _COLLECTION_ACTIVE: + self._delete_document_ids(client, [str(document["id"]) for document in search_documents]) + raise RuntimeError(f"Collection {collection_name!r} became unavailable during ingestion") + + summary = self._generate_summary("\n".join(texts), file_name) if self.cfg.generate_summary else None + self._update_collection_timestamp(collection_name) + return len(search_documents), summary, ingested_at + + def _upload_documents(self, client: SearchClient, documents: list[dict[str, Any]]) -> None: + uploaded_ids: list[str] = [] + try: + for batch in _iter_index_batches(documents, "upload"): + results = client.upload_documents(documents=batch) + succeeded, failures = _indexing_outcome(results, batch) + uploaded_ids.extend(succeeded) + if failures: + raise RuntimeError(f"Azure AI Search rejected upload actions: {'; '.join(failures)}") + except Exception as upload_error: # noqa: BLE001 + if uploaded_ids: + try: + self._delete_document_ids(client, uploaded_ids) + except Exception as rollback_error: # noqa: BLE001 + raise RuntimeError( + f"Upload failed ({upload_error}); rollback failed ({rollback_error})" + ) from upload_error + raise + + def _delete_document_ids(self, client: SearchClient, document_ids: list[str]) -> None: + for batch in _iter_index_batches([{"id": item} for item in document_ids], "delete"): + pending = batch + failures: list[str] = [] + for _attempt in range(_DELETE_ATTEMPTS): + results = client.delete_documents(documents=pending) + _succeeded, failures = _indexing_outcome(results, pending) + if not failures: + break + failed_keys = {failure.split(":", 1)[0] for failure in failures} + pending = [document for document in pending if str(document["id"]) in failed_keys] + if failures: + raise RuntimeError(f"Azure AI Search rejected delete actions: {'; '.join(failures)}") + + def _iter_documents( + self, + client: SearchClient, + *, + filter_text: str | None = None, + select: list[str] | None = None, + ) -> Iterator[dict[str, Any]]: + last_id: str | None = None + while True: + cursor_filter = f"id gt {_odata_literal(last_id)}" if last_id is not None else None + page = list( + client.search( + search_text="*", + filter=_and_filter(filter_text, cursor_filter), + select=select, + order_by=["id asc"], + top=_PAGE_SIZE, + ) + ) + if not page: + return + yield from page + next_id = str(page[-1].get("id") or "") + if len(page) < _PAGE_SIZE: + return + if not next_id or next_id == last_id: + raise RuntimeError("Azure AI Search pagination did not advance") + last_id = next_id + + def _delete_file_documents(self, file_id: str, collection_name: str) -> int: + client = self._get_validated_search_client() + filter_text = _record_filter(_RECORD_CHUNK, collection_name, file_id=file_id) + ids = [ + str(hit["id"]) + for hit in self._iter_documents(client, filter_text=filter_text, select=["id"]) + if hit.get("id") + ] + if ids: + self._delete_document_ids(client, ids) + return len(ids) + + def _generate_summary(self, text: str, file_name: str) -> str | None: + if self._summary_llm is None: + return None + snippet = text[:_SUMMARY_MAX_CHARS] + prompt = f"Summarise the following document ({file_name}) in one sentence (max 30 words):\n\n{snippet}" + try: + response = self._summary_llm.invoke(prompt) + summary = getattr(response, "content", None) or str(response) + return summary.strip() or None + except Exception: # noqa: BLE001 + logger.exception("Summary generation failed for %s", file_name) + return None + + def _update_job(self, job_id: str, **fields: Any) -> None: + with self._jobs_lock: + job = self._jobs.get(job_id) + if job: + self._jobs[job_id] = job.model_copy(update=fields) + + def _fail_job(self, job_id: str, error_message: str) -> None: + self._update_job(job_id, status=JobState.FAILED, error_message=error_message, completed_at=_utc_now()) + + def _fail_job_setup(self, job_id: str, error_message: str) -> None: + with self._jobs_lock: + job = self._jobs.get(job_id) + if job is None: + return + details = [ + detail.model_copy(update={"status": FileStatus.FAILED, "error_message": error_message}) + for detail in job.file_details + ] + for detail in details: + if tracked := self._files.get(detail.file_id): + tracked.status = FileStatus.FAILED + tracked.error_message = error_message + self._deleted_files.add((job.collection_name, detail.file_id)) + self._jobs[job_id] = job.model_copy( + update={ + "status": JobState.FAILED, + "processed_files": job.total_files, + "file_details": details, + "error_message": error_message, + "completed_at": _utc_now(), + } + ) + + def _update_file_progress(self, job_id: str, index: int, **fields: Any) -> None: + summary = fields.pop("summary", None) + ingested_at = fields.pop("ingested_at", None) + snapshot: FileInfo | None = None + with self._jobs_lock: + job = self._jobs.get(job_id) + if job is None or index >= len(job.file_details): + return + details = list(job.file_details) + details[index] = details[index].model_copy(update=fields) + tracked = self._files.get(details[index].file_id) + if tracked: + tracked.status = details[index].status + tracked.error_message = details[index].error_message + tracked.chunk_count = details[index].chunks_created + if tracked.status == FileStatus.SUCCESS: + tracked.ingested_at = ingested_at or _utc_now() + if summary: + tracked.metadata["summary"] = summary + else: + tracked.ingested_at = None + tracked.metadata.pop("summary", None) + snapshot = tracked.model_copy(deep=True) + self._jobs[job_id] = job.model_copy(update={"file_details": details}) + if snapshot and (collection := self._get_collection_manifest(snapshot.collection_name)): + if collection.get("status") == _COLLECTION_ACTIVE: + self._write_file_manifest(snapshot, summary=summary) + if snapshot.status == FileStatus.SUCCESS and summary: + register_summary(snapshot.collection_name, snapshot.file_name, summary) + + @staticmethod + def _cleanup_paths(paths: list[str]) -> None: + for path in paths: + try: + os.unlink(path) + except OSError: + pass + + @staticmethod + def _translate_error(error: Exception) -> str: + if isinstance(error, ResourceNotFoundError): + return f"AI Search index not found: {error!s}" + if isinstance(error, ClientAuthenticationError): + return f"AI Search authentication failed: {error!s}" + if isinstance(error, ServiceRequestError): + return f"AI Search service unavailable: {error!s}" + if isinstance(error, HttpResponseError): + return f"AI Search request failed ({error.status_code}): {error.reason or error!s}" + return str(error) + + def create_collection( + self, + name: str, + description: str | None = None, + metadata: dict[str, Any] | None = None, + ) -> CollectionInfo: + manifest = self._write_collection_manifest(name, description, metadata) + self._wait_for_search_state( + _record_filter(_RECORD_COLLECTION, name, status=_COLLECTION_ACTIVE), + present=True, + label=f"collection {name!r}", + ) + return self._collection_info(manifest) + + def delete_collection(self, name: str) -> bool: + manifest = self._get_collection_manifest(name) + if manifest is None: + return False + file_manifests = self._iter_documents( + self._get_validated_search_client(), + filter_text=_record_filter(_RECORD_FILE, name), + ) + if any( + self._file_info_from_manifest(document).status in {FileStatus.UPLOADING, FileStatus.INGESTING} + for document in file_manifests + if (name, str(document.get("file_id") or "")) not in self._deleted_files + ): + raise ValueError(f"Cannot delete collection {name!r} while files are ingesting") + + self._write_collection_manifest(name, status=_COLLECTION_DELETING) + client = self._get_validated_search_client() + content_filter = _and_filter( + f"collection_id eq {_odata_literal(name)}", + f"record_type ne {_odata_literal(_RECORD_COLLECTION)}", + ) + document_ids = [ + str(document["id"]) + for document in self._iter_documents(client, filter_text=content_filter, select=["id"]) + if document.get("id") + ] + if document_ids: + self._delete_document_ids(client, document_ids) + self._delete_document_ids(client, [str(manifest["id"])]) + with self._jobs_lock: + self._deleted_collections.add(name) + self._files = {file_id: info for file_id, info in self._files.items() if info.collection_name != name} + if self.cfg.generate_summary: + clear_collection_summaries(name) + return True + + def list_collections(self) -> list[CollectionInfo]: + try: + manifests = self._iter_documents( + self._get_validated_search_client(), + filter_text=_record_filter(_RECORD_COLLECTION, status=_COLLECTION_ACTIVE), + ) + return [ + self._collection_info(manifest) + for manifest in manifests + if manifest.get("collection_id") not in self._deleted_collections + ] + except ResourceNotFoundError: + return [] + + def get_collection(self, name: str) -> CollectionInfo | None: + if name in self._deleted_collections: + return None + manifest = self._get_collection_manifest(name) + if manifest is None or manifest.get("status") != _COLLECTION_ACTIVE: + return None + return self._collection_info(manifest) + + def upload_file( + self, + file_path: str, + collection_name: str, + metadata: dict[str, Any] | None = None, + ) -> FileInfo: + path = Path(file_path) + if not path.is_file(): + return FileInfo( + file_id=str(uuid.uuid4()), + file_name=path.name, + collection_name=collection_name, + status=FileStatus.FAILED, + error_message=f"File not found: {file_path}", + ) + job_id = self.submit_job( + [file_path], + collection_name, + {"original_filenames": [path.name], "metadata": metadata or {}}, + ) + with self._jobs_lock: + file_id = self._jobs[job_id].file_details[0].file_id + info = self._files[file_id].model_copy(deep=True) + info.metadata["job_id"] = job_id + return info + + def delete_file(self, file_id: str, collection_name: str) -> bool: + info = self.get_file_status(file_id, collection_name) + if info is None: + return False + if info.status in {FileStatus.UPLOADING, FileStatus.INGESTING}: + raise ValueError(f"Cannot delete file {file_id!r} while it is ingesting") + + self._delete_file_documents(file_id, collection_name) + self._delete_document_ids( + self._get_validated_search_client(), + [_record_id(_RECORD_FILE, collection_name, file_id)], + ) + self._deleted_files.add((collection_name, file_id)) + with self._jobs_lock: + self._files.pop(file_id, None) + + if self.cfg.generate_summary: + remaining = [ + item + for item in self.list_files(collection_name) + if item.file_name == info.file_name and item.status == FileStatus.SUCCESS + ] + newest = max(remaining, key=lambda item: item.ingested_at or datetime.min.replace(tzinfo=UTC), default=None) + if newest and (summary := newest.metadata.get("summary")): + register_summary(collection_name, info.file_name, str(summary)) + else: + unregister_summary(collection_name, info.file_name) + self._update_collection_timestamp(collection_name) + self._wait_for_search_state( + _record_filter(_RECORD_FILE, collection_name, file_id=file_id), + present=False, + label=f"file {file_id!r}", + ) + return True + + def list_files(self, collection_name: str) -> list[FileInfo]: + if collection_name in self._deleted_collections: + return [] + try: + collection = self._get_collection_manifest(collection_name) + if collection is None: + return [] + manifests = self._iter_documents( + self._get_validated_search_client(), + filter_text=_record_filter(_RECORD_FILE, collection_name), + ) + files = [ + self._file_info_from_manifest(manifest) + for manifest in manifests + if (collection_name, str(manifest.get("file_id") or "")) not in self._deleted_files + ] + except ResourceNotFoundError: + return [] + except Exception: # noqa: BLE001 + logger.exception("Failed to list files for %r", collection_name) + return [] + return sorted(files, key=lambda item: (item.file_name, item.file_id)) + + def get_file_status(self, file_id: str, collection_name: str) -> FileInfo | None: + if (collection_name, file_id) in self._deleted_files: + return None + manifest = self._get_file_manifest(file_id, collection_name) + return self._file_info_from_manifest(manifest) if manifest else None + + async def health_check(self) -> bool: + def _check() -> bool: + list(self._index_client.list_index_names()) + return True + + try: + return await asyncio.to_thread(_check) + except Exception: # noqa: BLE001 + logger.exception("Azure AI Search ingestor health_check failed") + return False + + +def _resolve_filenames(file_paths: list[str], raw: Any) -> list[str]: + if isinstance(raw, dict): + return [raw.get(path) or Path(path).name for path in file_paths] + if isinstance(raw, list): + return [ + raw[index] if index < len(raw) and raw[index] else Path(path).name for index, path in enumerate(file_paths) + ] + return [Path(path).name for path in file_paths] + + +def _coerce_page_number(page_label: Any) -> int | None: + if page_label is None: + return None + try: + page_number = int(str(page_label)) + return page_number if page_number > 0 else None + except (TypeError, ValueError): + return None diff --git a/sources/knowledge_layer/src/register.py b/sources/knowledge_layer/src/register.py index 292873c1f..02948bc24 100644 --- a/sources/knowledge_layer/src/register.py +++ b/sources/knowledge_layer/src/register.py @@ -26,6 +26,7 @@ from typing import Literal from pydantic import Field +from pydantic import HttpUrl from pydantic import SecretStr from pydantic import model_validator @@ -33,13 +34,24 @@ from nat.builder.context import Context from nat.builder.function_info import FunctionInfo from nat.cli.register_workflow import register_function +from nat.data_models.common import OptionalSecretStr from nat.data_models.function import FunctionBaseConfig logger = logging.getLogger(__name__) +def _url_from_env(name: str, default: str | None = None) -> HttpUrl | None: + value = os.environ.get(name) or default + return HttpUrl(value) if value else None + + +def _secret_from_env(name: str) -> SecretStr | None: + value = os.environ.get(name) + return SecretStr(value) if value else None + + # Type-safe backend selection - Pydantic validates at config load time -BackendType = Literal["llamaindex", "foundational_rag", "opensearch"] +BackendType = Literal["llamaindex", "foundational_rag", "opensearch", "azure_ai_search"] OpenSearchAuthType = Literal["none", "basic", "sigv4"] OpenSearchAwsService = Literal["aoss", "es"] OpenSearchIngestionMode = Literal["local", "dask", "auto"] @@ -244,12 +256,31 @@ class KnowledgeRetrievalConfig(FunctionBaseConfig, name="knowledge_retrieval"): ) embed_model: str = Field( default_factory=lambda: _env_value("AIQ_EMBED_MODEL", default="nvidia/llama-nemotron-embed-vl-1b-v2"), - description="Embedding model for OpenSearch vector ingestion and retrieval.", + description="Embedding model for OpenSearch and Azure AI Search ingestion and retrieval.", ) embed_base_url: str = Field( default_factory=lambda: _env_value("AIQ_EMBED_BASE_URL", default="https://integrate.api.nvidia.com/v1"), description="OpenAI-compatible embeddings endpoint base URL.", ) + # Azure AI Search options + azure_search_endpoint: HttpUrl | None = Field( + default_factory=lambda: _url_from_env("AZURE_SEARCH_ENDPOINT"), + description="Azure AI Search service URL; defaults to AZURE_SEARCH_ENDPOINT", + ) + azure_search_api_key: OptionalSecretStr = Field( + default_factory=lambda: _secret_from_env("AZURE_SEARCH_API_KEY"), + description="Optional Azure AI Search admin key; defaults to AZURE_SEARCH_API_KEY", + ) + azure_search_index_prefix: str = Field( + default_factory=lambda: _env_value("AIQ_AZURE_SEARCH_INDEX_PREFIX", default="aiq"), + min_length=1, + description="Unique deployment namespace for the shared AI-Q index", + ) + embed_dim: int = Field( + default_factory=lambda: _env_int("AIQ_EMBED_DIM", 2048), + gt=0, + description="Embedding dimensions; defaults to AIQ_EMBED_DIM and must match existing indexes", + ) @model_validator(mode="after") def validate_backend_config(self): @@ -298,6 +329,9 @@ def validate_backend_config(self): ) if not self.opensearch_verify_certs: logger.warning("TLS verification disabled for opensearch. Use only in trusted environments.") + elif backend == "azure_ai_search": + if self.azure_search_endpoint is None: + raise ValueError("azure_ai_search requires azure_search_endpoint") return self @@ -384,8 +418,25 @@ def _setup_backend(config: KnowledgeRetrievalConfig, summary_llm_obj=None) -> tu **summary_config, } + elif backend == "azure_ai_search": + import knowledge_layer.azure_ai_search.adapter # noqa: F401 + + backend_config = { + "endpoint": str(config.azure_search_endpoint), + "api_key": config.azure_search_api_key, + "index_prefix": config.azure_search_index_prefix, + "embed_base_url": str(config.embed_base_url), + "embed_model": config.embed_model, + "embed_dim": config.embed_dim, + "collection_name": config.collection_name, + "cleanup_files": False, + **summary_config, + } + else: - raise ValueError(f"Unknown backend: {backend}. Use 'llamaindex', 'foundational_rag', or 'opensearch'.") + raise ValueError( + f"Unknown backend: {backend}. Use 'llamaindex', 'foundational_rag', 'opensearch', or 'azure_ai_search'." + ) os.environ["KNOWLEDGE_RETRIEVER_BACKEND"] = backend os.environ["KNOWLEDGE_INGESTOR_BACKEND"] = backend @@ -477,7 +528,7 @@ async def knowledge_retrieval(config: KnowledgeRetrievalConfig, _builder: Builde This function provides semantic search over documents that have been previously ingested into the knowledge layer. It supports multiple - backends (LlamaIndex, Foundational RAG, OpenSearch) and returns formatted results + backends (LlamaIndex, Foundational RAG, OpenSearch, Azure AI Search) and returns formatted results suitable for LLM consumption. The retriever and ingestor are initialized once when the function is diff --git a/tests/knowledge_layer_tests/run_adapter_compliance.py b/tests/knowledge_layer_tests/run_adapter_compliance.py index 99ef916e2..52b051a69 100755 --- a/tests/knowledge_layer_tests/run_adapter_compliance.py +++ b/tests/knowledge_layer_tests/run_adapter_compliance.py @@ -19,12 +19,23 @@ python tests/knowledge_layer_tests/run_adapter_compliance.py --backend llamaindex --quick python tests/knowledge_layer_tests/run_adapter_compliance.py --backend foundational_rag --quick python tests/knowledge_layer_tests/run_adapter_compliance.py --backend opensearch --quick + python tests/knowledge_layer_tests/run_adapter_compliance.py --backend azure_ai_search --quick \ + --config '{"endpoint":"https://example.search.windows.net","start_ttl_cleanup":false}' # Full mode - complete ingestion + retrieval test python tests/knowledge_layer_tests/run_adapter_compliance.py --backend llamaindex python tests/knowledge_layer_tests/run_adapter_compliance.py --backend foundational_rag python tests/knowledge_layer_tests/run_adapter_compliance.py --backend opensearch + export AZURE_SEARCH_ENDPOINT=https://your-service.search.windows.net + export AZURE_SEARCH_API_KEY=your-search-admin-key + export AIQ_AZURE_SEARCH_INDEX_PREFIX="aiq-${USER}" + python tests/knowledge_layer_tests/run_adapter_compliance.py --backend azure_ai_search \ + --config '{"start_ttl_cleanup":false}' + + # For managed identity, omit AZURE_SEARCH_API_KEY. Set AZURE_CLIENT_ID for a + # user-assigned identity. NVIDIA_API_KEY remains required for embeddings. + Exit codes: 0 - All tests passed 1 - Some tests failed @@ -95,6 +106,7 @@ def _import_backend(self): "llamaindex": "knowledge_layer.llamaindex", "foundational_rag": "knowledge_layer.foundational_rag", "opensearch": "knowledge_layer.opensearch", + "azure_ai_search": "knowledge_layer.azure_ai_search", } module_name = backend_imports.get(self.backend.lower()) @@ -390,8 +402,9 @@ def _test_delete_file(self): if not files: return True, "No files to delete (already empty)" + file_id = files[0].file_id filename = files[0].file_name - result = self.ingestor.delete_file(filename, self.collection_name) + result = self.ingestor.delete_file(file_id, self.collection_name) if not result: return False, f"delete_file returned False for '{filename}'" @@ -455,7 +468,10 @@ def main(): ) parser.add_argument( - "--backend", "-b", required=True, help="Backend name (e.g., llamaindex, foundational_rag, opensearch)" + "--backend", + "-b", + required=True, + help="Backend name (e.g., llamaindex, foundational_rag, opensearch, azure_ai_search)", ) parser.add_argument("--config", "-c", default="{}", help="Backend config as JSON string (default: {})") diff --git a/tests/knowledge_layer_tests/test_azure_ai_search.py b/tests/knowledge_layer_tests/test_azure_ai_search.py new file mode 100644 index 000000000..576910af0 --- /dev/null +++ b/tests/knowledge_layer_tests/test_azure_ai_search.py @@ -0,0 +1,1140 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import asyncio +import re +from datetime import UTC +from datetime import datetime +from datetime import timedelta +from types import SimpleNamespace +from uuid import UUID + +import pytest +from azure.core.credentials import AzureKeyCredential +from azure.core.exceptions import ResourceNotFoundError +from azure.core.exceptions import ServiceRequestError +from azure.identity import DefaultAzureCredential +from azure.search.documents.indexes.models import SearchIndex +from knowledge_layer.azure_ai_search import adapter as azure_adapter +from knowledge_layer.azure_ai_search.adapter import AzureAISearchIngestor +from knowledge_layer.azure_ai_search.adapter import AzureAISearchRetriever +from knowledge_layer.azure_ai_search.adapter import _build_index_schema +from knowledge_layer.azure_ai_search.adapter import _coerce_page_number +from knowledge_layer.azure_ai_search.adapter import _encode_marker +from knowledge_layer.azure_ai_search.adapter import _index_name_for_config +from knowledge_layer.azure_ai_search.adapter import _iter_index_batches +from knowledge_layer.azure_ai_search.adapter import _new_marker +from knowledge_layer.azure_ai_search.adapter import _resolve_filenames +from knowledge_layer.azure_ai_search.adapter import _validate_index_name +from knowledge_layer.azure_ai_search.adapter import _validate_index_schema +from knowledge_layer.register import KnowledgeRetrievalConfig +from knowledge_layer.register import _format_results +from knowledge_layer.register import _setup_backend +from pydantic import SecretStr + +from aiq_agent.knowledge import BaseIngestor +from aiq_agent.knowledge import BaseRetriever +from aiq_agent.knowledge import ContentType +from aiq_agent.knowledge import RetrievalResult +from aiq_agent.knowledge.factory import is_ingestor_registered +from aiq_agent.knowledge.factory import is_retriever_registered +from aiq_agent.knowledge.schema import FileStatus +from aiq_agent.knowledge.schema import JobState + + +class FakeIndexingResult: + def __init__(self, key: str, succeeded: bool = True, error_message: str | None = None): + self.key = key + self.succeeded = succeeded + self.error_message = error_message + + +class FakeSearchResults(list): + def __init__(self, documents: list[dict], facets: dict | None = None): + super().__init__(documents) + self._facets = facets or {} + + def get_facets(self): + return self._facets + + +def _literal(filter_text: str | None, field: str, operator: str = "eq") -> str | None: + if not filter_text: + return None + match = re.search(rf"{field} {operator} '((?:''|[^'])*)'", filter_text) + return match.group(1).replace("''", "'") if match else None + + +class FakeSearchClient: + def __init__(self): + self.documents: dict[str, dict] = {} + self.upload_batches: list[list[dict]] = [] + self.delete_batches: list[list[dict]] = [] + self.search_filters: list[str | None] = [] + self.search_requests: list[dict] = [] + self.fail_upload_ids: set[str] = set() + self.delete_failures_remaining: dict[str, int] = {} + self.hidden_searches_remaining = 0 + self.stale_deleted_searches_remaining = 0 + self.deleted_documents: dict[str, dict] = {} + + def upload_documents(self, documents: list[dict]): + self.upload_batches.append(documents) + results = [] + for document in documents: + key = document["id"] + succeeded = key not in self.fail_upload_ids + if succeeded: + self.documents[key] = dict(document) + results.append(FakeIndexingResult(key, succeeded, None if succeeded else "rejected")) + return results + + def delete_documents(self, documents: list[dict]): + self.delete_batches.append(documents) + results = [] + for document in documents: + key = document["id"] + remaining = self.delete_failures_remaining.get(key, 0) + succeeded = remaining <= 0 + if remaining > 0: + self.delete_failures_remaining[key] = remaining - 1 + if succeeded: + deleted = self.documents.pop(key, None) + if deleted: + self.deleted_documents[key] = deleted + results.append(FakeIndexingResult(key, succeeded, None if succeeded else "retry")) + return results + + def get_document(self, key: str): + if key not in self.documents: + raise ResourceNotFoundError("missing") + return dict(self.documents[key]) + + def search(self, search_text="*", filter=None, select=None, order_by=None, top=None, **kwargs): + del order_by + self.search_filters.append(filter) + self.search_requests.append( + {"search_text": search_text, "filter": filter, "select": select, "top": top, **kwargs} + ) + equals = { + field: _literal(filter, field) + for field in ("record_type", "collection_id", "file_id", "file_name", "status") + } + excluded_record_type = _literal(filter, "record_type", "ne") + after_id = _literal(filter, "id", "gt") + documents = list(self.documents.values()) + if self.stale_deleted_searches_remaining > 0 and self.deleted_documents: + documents.extend(self.deleted_documents.values()) + self.stale_deleted_searches_remaining -= 1 + if self.hidden_searches_remaining > 0: + documents = [] + self.hidden_searches_remaining -= 1 + documents = sorted(documents, key=lambda item: item["id"]) + for field, value in equals.items(): + if value is not None: + documents = [document for document in documents if document.get(field) == value] + if excluded_record_type is not None: + documents = [document for document in documents if document.get("record_type") != excluded_record_type] + if after_id is not None: + documents = [document for document in documents if document["id"] > after_id] + facet_documents = documents + if top is not None: + documents = documents[:top] + facets = {} + if "facets" in kwargs and any(str(value).startswith("record_type") for value in kwargs["facets"]): + counts: dict[str, int] = {} + for document in facet_documents: + record_type = document.get("record_type") + if record_type: + counts[record_type] = counts.get(record_type, 0) + 1 + facets["record_type"] = [{"value": value, "count": count} for value, count in counts.items()] + if select: + documents = [{key: document.get(key) for key in select} for document in documents] + else: + documents = [dict(document) for document in documents] + return FakeSearchResults(documents, facets) + + def get_document_count(self): + return len(self.documents) + + +class FakeIndexClient: + def __init__(self): + self.indexes: dict[str, SearchIndex] = {} + self.race_on_create = False + self.fail_create = False + + def get_index(self, name: str): + if name not in self.indexes: + raise ResourceNotFoundError("missing") + return self.indexes[name] + + def create_index(self, index: SearchIndex): + if self.race_on_create: + self.indexes[index.name] = index + raise RuntimeError("concurrent create") + if self.fail_create: + raise RuntimeError("create failed") + self.indexes[index.name] = index + return index + + def create_or_update_index(self, index: SearchIndex, **kwargs): + del kwargs + self.indexes[index.name] = index + return index + + def list_indexes(self): + return list(self.indexes.values()) + + def list_index_names(self): + return list(self.indexes) + + def delete_index(self, name: str): + if name not in self.indexes: + raise ResourceNotFoundError("missing") + del self.indexes[name] + + +class FakeEmbedding: + def __init__(self, dimensions: int = 4): + self.dimensions = dimensions + + def get_query_embedding(self, query): + del query + return [0.1] * self.dimensions + + def get_text_embedding_batch(self, texts): + return [[0.1] * self.dimensions for _ in texts] + + +class FakeNode: + def __init__(self, text: str, page: str = "1"): + self.text = text + self.metadata = {"page_label": page} + + def get_content(self): + return self.text + + +class FakeSplitter: + def __init__(self, nodes: list[FakeNode]): + self.nodes = nodes + + def get_nodes_from_documents(self, documents): + del documents + return self.nodes + + +def _config(**overrides): + config = { + "endpoint": "https://example.search.windows.net", + "api_key": SecretStr("test-key"), + "embed_model": "test-embed", + "embed_dim": 4, + "embed_base_url": "https://integrate.api.nvidia.com/v1", + "start_ttl_cleanup": False, + "index_prefix": "aiq-test", + } + config.update(overrides) + return config + + +def _ingestor(**overrides): + ingestor = AzureAISearchIngestor(_config(**overrides)) + ingestor._index_client = FakeIndexClient() + search_client = FakeSearchClient() + ingestor._search_client = search_client + ingestor._index_validated = False + return ingestor, search_client + + +def _write_file_manifest( + ingestor, + file_id: str, + collection_name: str = "docs", + file_name: str = "report.pdf", + *, + status: FileStatus = FileStatus.SUCCESS, + chunk_count: int = 0, + summary: str | None = None, + ingested_at: datetime | None = None, +): + info = azure_adapter.FileInfo( + file_id=file_id, + file_name=file_name, + collection_name=collection_name, + status=status, + error_message="bad file" if status == FileStatus.FAILED else None, + file_size=42, + chunk_count=chunk_count, + uploaded_at=datetime.now(UTC), + ingested_at=ingested_at or datetime.now(UTC), + metadata={"section": "finance"}, + ) + ingestor._write_file_manifest(info, summary=summary) + return info + + +def _install_reader(monkeypatch): + class FakeReader: + def __init__(self, input_files): + self.input_files = input_files + + def load_data(self): + return [object()] + + monkeypatch.setattr("llama_index.core.SimpleDirectoryReader", FakeReader) + + +def test_backend_registered_and_implements_sdk_contracts(): + assert is_ingestor_registered("azure_ai_search") + assert is_retriever_registered("azure_ai_search") + assert issubclass(AzureAISearchIngestor, BaseIngestor) + assert issubclass(AzureAISearchRetriever, BaseRetriever) + + +def test_config_requires_endpoint_and_omits_removed_azure_options(monkeypatch): + monkeypatch.delenv("AZURE_SEARCH_ENDPOINT", raising=False) + monkeypatch.delenv("AZURE_SEARCH_API_KEY", raising=False) + + with pytest.raises(ValueError, match="azure_search_endpoint"): + KnowledgeRetrievalConfig(backend="azure_ai_search") + config = KnowledgeRetrievalConfig( + backend="azure_ai_search", + azure_search_endpoint="https://example.search.windows.net", + ) + assert config.azure_search_api_key is None + for field in ("chunk_size", "chunk_overlap", "use_hybrid", "use_semantic_ranker", "summary_max_chars"): + assert field not in KnowledgeRetrievalConfig.model_fields + + +def test_config_uses_shared_environment_defaults(monkeypatch): + monkeypatch.setenv("AZURE_SEARCH_ENDPOINT", "https://env.search.windows.net") + monkeypatch.setenv("AZURE_SEARCH_API_KEY", "env-search-key") + monkeypatch.setenv("AIQ_AZURE_SEARCH_INDEX_PREFIX", "env-aiq") + monkeypatch.setenv("AIQ_EMBED_BASE_URL", "https://embed.example.com/v1") + monkeypatch.setenv("AIQ_EMBED_MODEL", "env-embed") + monkeypatch.setenv("AIQ_EMBED_DIM", "8") + + config = KnowledgeRetrievalConfig(backend="azure_ai_search") + backend, backend_config = _setup_backend(config) + adapter_config = azure_adapter._coerce_config(None) + + assert backend == "azure_ai_search" + assert backend_config["endpoint"] == "https://env.search.windows.net/" + assert backend_config["api_key"].get_secret_value() == "env-search-key" + assert backend_config["index_prefix"] == "env-aiq" + assert backend_config["embed_base_url"] == "https://embed.example.com/v1" + assert backend_config["embed_model"] == "env-embed" + assert backend_config["embed_dim"] == 8 + assert adapter_config.endpoint == "https://env.search.windows.net" + + +def test_config_uses_defaults_for_empty_azure_environment_values(monkeypatch): + monkeypatch.setenv("AIQ_AZURE_SEARCH_INDEX_PREFIX", "") + monkeypatch.setenv("AIQ_EMBED_DIM", "") + + config = KnowledgeRetrievalConfig( + backend="azure_ai_search", + azure_search_endpoint="https://example.search.windows.net", + ) + + assert config.azure_search_index_prefix == "aiq" + assert config.embed_dim == 2048 + + +def test_shared_embedding_defaults_match_adapter(monkeypatch): + monkeypatch.setenv("AIQ_EMBED_BASE_URL", "") + monkeypatch.delenv("AIQ_EMBED_MODEL", raising=False) + monkeypatch.delenv("AIQ_EMBED_DIM", raising=False) + + config = KnowledgeRetrievalConfig( + backend="azure_ai_search", + azure_search_endpoint="https://example.search.windows.net", + ) + adapter_config = azure_adapter._coerce_config({"endpoint": "https://example.search.windows.net"}) + + assert str(config.embed_base_url) == "https://integrate.api.nvidia.com/v1" + assert adapter_config.embed_base_url == "https://integrate.api.nvidia.com/v1" + assert config.embed_model == "nvidia/llama-nemotron-embed-vl-1b-v2" + assert config.embed_dim == 2048 + assert adapter_config.embed_model == config.embed_model + assert adapter_config.embed_dim == config.embed_dim + + +def test_search_credential_uses_api_key_or_default_credential(): + with_key = azure_adapter._coerce_config( + {"endpoint": "https://example.search.windows.net", "api_key": SecretStr("test-key")} + ) + without_key = azure_adapter._coerce_config({"endpoint": "https://example.search.windows.net", "api_key": None}) + + assert isinstance(azure_adapter._build_search_credential(with_key), AzureKeyCredential) + assert isinstance(azure_adapter._build_search_credential(without_key), DefaultAzureCredential) + + +def test_setup_backend_preserves_secrets_and_prefix(monkeypatch): + monkeypatch.setenv("KNOWLEDGE_RETRIEVER_BACKEND", "llamaindex") + monkeypatch.setenv("KNOWLEDGE_INGESTOR_BACKEND", "llamaindex") + config = KnowledgeRetrievalConfig( + backend="azure_ai_search", + azure_search_endpoint="https://example.search.windows.net", + azure_search_api_key="test-search-key", # pragma: allowlist secret + azure_search_index_prefix="tenant-aiq", + ) + + backend, backend_config = _setup_backend(config) + + assert backend == "azure_ai_search" + assert isinstance(backend_config["api_key"], SecretStr) + assert backend_config["index_prefix"] == "tenant-aiq" + + +def test_search_api_key_survives_nat_json_serialization(): + api_key = "test-search-key" # pragma: allowlist secret + config = KnowledgeRetrievalConfig( + backend="azure_ai_search", + azure_search_endpoint="https://example.search.windows.net", + azure_search_api_key=api_key, + ) + + serialized = config.model_dump(mode="json", by_alias=True, round_trip=True) + + assert serialized["azure_search_api_key"] == api_key + + +def test_index_name_is_stable_and_changes_with_embedding_configuration(): + index = _index_name_for_config("AIQ Prod", "test/embed", 2048) + + assert index == _index_name_for_config("AIQ Prod", "test/embed", 2048) + assert index != _index_name_for_config("AIQ Prod", "other/embed", 2048) + assert index != _index_name_for_config("AIQ Prod", "test/embed", 1024) + assert len(index) <= 128 + assert re.fullmatch(r"[a-z0-9-]+", index) + _validate_index_name(index) + with pytest.raises(ValueError, match="Invalid Azure AI Search index name"): + _validate_index_name("invalid_name") + + +def test_marker_and_schema_validation_reject_unowned_and_mismatched_indexes(): + cfg = SimpleNamespace(embed_model="test-embed", embed_dim=4, index_prefix="aiq-test") + marker = _new_marker(cfg) + index = _build_index_schema("aiq-docs-123456789abc", 4, _encode_marker(marker)) + + assert _validate_index_schema(index, cfg)["schema_version"] == 1 + assert "metadata" not in marker + assert "collection" not in marker + index.description = "unmanaged" + with pytest.raises(RuntimeError, match="not owned"): + _validate_index_schema(index, cfg) + index.description = _encode_marker(marker) + embedding = next(field for field in index.fields if field.name == "embedding") + embedding.vector_search_dimensions = 8 + with pytest.raises(RuntimeError, match="vector profile or dimensions"): + _validate_index_schema(index, cfg) + + +def test_create_collection_handles_race_and_ignores_unmanaged_indexes(): + ingestor, _client = _ingestor() + ingestor._index_client.race_on_create = True + created = ingestor.create_collection("docs", description="Documents", metadata={"tenant": "alpha"}) + ingestor._index_client.indexes["unrelated"] = SearchIndex(name="unrelated", fields=[], description="other") + + assert created.name == "docs" + assert created.metadata["tenant"] == "alpha" + assert [collection.name for collection in ingestor.list_collections()] == ["docs"] + + +def test_multiple_collections_share_one_physical_index_and_keep_marker_bounded(): + ingestor, _client = _ingestor() + large_value = "x" * 10_000 + + ingestor.create_collection("docs", description=large_value, metadata={"large": large_value}) + ingestor.create_collection("private") + + assert len(ingestor._index_client.indexes) == 1 + index = ingestor._index_client.indexes[ingestor._physical_index_name()] + assert len(index.description) < 4_000 + assert large_value not in index.description + assert {item.name for item in ingestor.list_collections()} == {"docs", "private"} + + +def test_create_waits_for_collection_manifest_search_visibility(monkeypatch): + ingestor, client = _ingestor() + client.hidden_searches_remaining = 2 + monkeypatch.setattr(azure_adapter, "_CONSISTENCY_DELAY_SECONDS", 0) + + ingestor.create_collection("docs") + + assert client.hidden_searches_remaining == 0 + assert [item.name for item in ingestor.list_collections()] == ["docs"] + + +def test_mismatched_and_foreign_indexes_are_ignored(): + ingestor, _client = _ingestor() + mismatched_marker = _new_marker(ingestor.cfg) + mismatched_marker["schema_version"] = azure_adapter._SCHEMA_VERSION + 1 + ingestor._index_client.indexes["aiq-test-docs-mismatched"] = _build_index_schema( + "aiq-test-docs-mismatched", + ingestor.cfg.embed_dim, + _encode_marker(mismatched_marker), + ) + ingestor._index_client.indexes["foreign"] = SearchIndex(name="foreign", fields=[], description="foreign") + + assert ingestor.list_collections() == [] + + +def test_create_collection_propagates_real_create_failure(): + ingestor, _client = _ingestor() + ingestor._index_client.fail_create = True + + with pytest.raises(RuntimeError, match="create failed"): + ingestor.create_collection("docs") + + +@pytest.mark.parametrize( + ("hit", "expected"), + [ + ({"@search.reranker_score": 2.0, "@search.score": 0.9}, 0.9), + ({"@search.score": 8.0}, 1.0), + ({"@search.score": -1.0}, 0.0), + ({"@search.score": 0.75}, 0.75), + ], +) +def test_normalize_clamps_search_score(hit, expected): + chunk = AzureAISearchRetriever.__new__(AzureAISearchRetriever).normalize( + {"id": "chunk-1", "chunk": "text", "file_name": "report.pdf", **hit} + ) + assert chunk.score == expected + + +def test_normalize_parses_metadata_and_populates_citation(): + retriever = AzureAISearchRetriever.__new__(AzureAISearchRetriever) + cited = retriever.normalize( + { + "id": "chunk-1", + "chunk": "text", + "file_id": "file-1", + "file_name": "report.pdf", + "page_number": "3", + "metadata": '{"section":"intro"}', + } + ) + fallback = retriever.normalize({"page_number": 0}) + + assert cited.display_citation == "report.pdf, p.3" + assert cited.content_type is ContentType.TEXT + assert cited.metadata == {"section": "intro", "file_id": "file-1"} + assert fallback.content == "" + assert fallback.file_name == "unknown" + UUID(fallback.chunk_id) + + +def test_shared_formatter_retains_source_and_citation_lines(): + chunk = AzureAISearchRetriever.__new__(AzureAISearchRetriever).normalize( + {"id": "chunk-1", "chunk": "text", "file_name": "report.pdf", "page_number": 3} + ) + formatted = _format_results(RetrievalResult(query="query", backend="azure_ai_search", chunks=[chunk]), "query") + + assert "Source: report.pdf" in formatted + assert "Citation: report.pdf, p.3" in formatted + + +@pytest.mark.asyncio +async def test_retrieve_builds_scoped_hybrid_search_request(): + class FakeClient: + def search(self, **kwargs): + self.kwargs = kwargs + return [{"id": "chunk-1", "chunk": "answer", "file_name": "report.pdf"}] + + client = FakeClient() + retriever = AzureAISearchRetriever.__new__(AzureAISearchRetriever) + retriever._embedding = FakeEmbedding(2) + retriever._get_client = lambda collection_name: client + + result = await retriever.retrieve("hello", "session-1", top_k=5, filters={}) + + assert result.success + assert client.kwargs["search_text"] == "hello" + assert len(client.kwargs["vector_queries"]) == 1 + assert "record_type eq 'chunk'" in client.kwargs["filter"] + assert "collection_id eq 'session-1'" in client.kwargs["filter"] + assert "query_type" not in client.kwargs + assert client.kwargs["select"] == ["id", "chunk", "file_id", "file_name", "page_number", "metadata"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("filters", [{"$filter": "file_id ne ''"}, {"file_name": "report.pdf"}]) +async def test_retrieve_rejects_filters(filters): + retriever = AzureAISearchRetriever.__new__(AzureAISearchRetriever) + retriever._get_client = lambda collection_name: pytest.fail(f"unexpected Azure request for {collection_name}") + + result = await retriever.retrieve("hello", "session-1", filters=filters) + + assert not result.success + assert result.error_message == "Azure AI Search metadata filters are not supported" + + +def test_submit_job_uses_one_canonical_file_id(monkeypatch, tmp_path): + class NoStartThread: + def __init__(self, *args, **kwargs): + pass + + def start(self): + pass + + monkeypatch.setattr(azure_adapter.threading, "Thread", NoStartThread) + path = tmp_path / "document.txt" + path.write_text("content", encoding="utf-8") + ingestor, _client = _ingestor() + + job_id = ingestor.submit_job( + [str(path)], + "docs", + {"original_filenames": ["original.txt"], "metadata": {"large": "x" * 10_000}}, + ) + job = ingestor.get_job_status(job_id) + file_id = job.file_details[0].file_id + + assert job.status == JobState.PENDING + assert file_id in ingestor._files + assert ingestor._files[file_id].file_name == "original.txt" + assert ingestor._files[file_id].metadata["large"] == "x" * 10_000 + assert _client.documents == {} + + +@pytest.mark.parametrize("rollback_fails", [False, True]) +def test_submit_job_rolls_back_manifests_when_second_write_fails(monkeypatch, tmp_path, rollback_fails): + class SynchronousThread: + def __init__(self, *, target, args, **kwargs): + del kwargs + self.target = target + self.args = args + + def start(self): + self.target(*self.args) + + monkeypatch.setattr(azure_adapter.threading, "Thread", SynchronousThread) + paths = [tmp_path / "first.txt", tmp_path / "second.txt"] + for path in paths: + path.write_text("content", encoding="utf-8") + ingestor, client = _ingestor() + write_manifest = ingestor._write_file_manifest + writes = [] + + def fail_second_manifest(info, **kwargs): + assert info.metadata["job_id"] in ingestor._jobs + writes.append(info.file_id) + if len(writes) == 2: + raise RuntimeError("second manifest failed") + return write_manifest(info, **kwargs) + + monkeypatch.setattr(ingestor, "_write_file_manifest", fail_second_manifest) + if rollback_fails: + + def fail_rollback(*args, **kwargs): + raise RuntimeError("rollback unavailable") + + monkeypatch.setattr(ingestor, "_delete_document_ids", fail_rollback) + + job_id = ingestor.submit_job([str(path) for path in paths], "docs") + job = ingestor.get_job_status(job_id) + + assert job.status == JobState.FAILED + assert job.processed_files == 2 + assert all(detail.status == FileStatus.FAILED for detail in job.file_details) + assert all(ingestor._files[detail.file_id].status == FileStatus.FAILED for detail in job.file_details) + assert "second manifest failed" in job.error_message + if rollback_fails: + assert "manifest rollback failed: rollback unavailable" in job.error_message + else: + assert not [document for document in client.documents.values() if document.get("record_type") == "file"] + assert client.delete_batches + assert all(path.exists() for path in paths) + + +def test_submit_job_records_thread_start_failure(monkeypatch, tmp_path): + class FailingThread: + def __init__(self, *args, **kwargs): + pass + + def start(self): + raise RuntimeError("thread unavailable") + + monkeypatch.setattr(azure_adapter.threading, "Thread", FailingThread) + path = tmp_path / "document.txt" + path.write_text("content", encoding="utf-8") + ingestor, client = _ingestor() + + job_id = ingestor.submit_job([str(path)], "docs", {"cleanup_files": True}) + job = ingestor.get_job_status(job_id) + + assert job.status == JobState.FAILED + assert job.processed_files == 1 + assert job.file_details[0].status == FileStatus.FAILED + assert "thread unavailable" in job.error_message + assert client.documents == {} + assert not path.exists() + + +def test_upload_file_preserves_caller_owned_source_by_default(monkeypatch, tmp_path): + class SynchronousThread: + def __init__(self, *, target, args, **kwargs): + del kwargs + self.target = target + self.args = args + + def start(self): + self.target(*self.args) + + _install_reader(monkeypatch) + monkeypatch.setattr(azure_adapter.threading, "Thread", SynchronousThread) + path = tmp_path / "document.txt" + path.write_text("content", encoding="utf-8") + ingestor, _client = _ingestor() + ingestor._embedding = FakeEmbedding() + ingestor._splitter = FakeSplitter([FakeNode("content")]) + + ingestor.upload_file(str(path), "docs") + + assert path.exists() + + +def test_post_upload_failure_rolls_back_chunks_and_retains_failed_manifest(monkeypatch, tmp_path): + class SynchronousThread: + def __init__(self, *, target, args, **kwargs): + del kwargs + self.target = target + self.args = args + + def start(self): + self.target(*self.args) + + _install_reader(monkeypatch) + monkeypatch.setattr(azure_adapter.threading, "Thread", SynchronousThread) + path = tmp_path / "document.txt" + path.write_text("content", encoding="utf-8") + ingestor, client = _ingestor() + ingestor._embedding = FakeEmbedding() + ingestor._splitter = FakeSplitter([FakeNode("content")]) + + def fail_timestamp(collection): + del collection + raise RuntimeError("fail") + + monkeypatch.setattr(ingestor, "_update_collection_timestamp", fail_timestamp) + + uploaded = ingestor.upload_file(str(path), "docs") + status = ingestor.get_file_status(uploaded.file_id, "docs") + + assert status.status == FileStatus.FAILED + assert not [document for document in client.documents.values() if document.get("record_type") == "chunk"] + + +def test_batches_respect_action_count_and_payload_size(): + documents = [{"id": str(index), "chunk": "x" * 10} for index in range(5)] + count_batches = list(_iter_index_batches(documents, "upload", max_actions=2, max_bytes=10_000)) + byte_batches = list(_iter_index_batches(documents, "upload", max_actions=100, max_bytes=120)) + + assert [len(batch) for batch in count_batches] == [2, 2, 1] + assert len(byte_batches) > 1 + with pytest.raises(ValueError, match="16 MiB"): + list(_iter_index_batches([{"id": "large", "chunk": "x" * 500}], "upload", max_bytes=100)) + + +def test_partial_upload_rolls_back_successful_actions(): + ingestor, client = _ingestor() + documents = [{"id": f"file-{index}", "chunk": "text"} for index in range(3)] + client.fail_upload_ids.add("file-1") + + with pytest.raises(RuntimeError, match="rejected upload"): + ingestor._upload_documents(client, documents) + + assert client.documents == {} + assert client.delete_batches + + +def test_delete_retries_per_document_failures(): + ingestor, client = _ingestor() + client.documents["chunk-1"] = {"id": "chunk-1"} + client.delete_failures_remaining["chunk-1"] = 1 + + ingestor._delete_document_ids(client, ["chunk-1"]) + + assert "chunk-1" not in client.documents + assert len(client.delete_batches) == 2 + + +def test_persistent_delete_failure_is_reported(): + ingestor, client = _ingestor() + client.documents["chunk-1"] = {"id": "chunk-1"} + client.delete_failures_remaining["chunk-1"] = 10 + + with pytest.raises(RuntimeError, match="rejected delete"): + ingestor._delete_document_ids(client, ["chunk-1"]) + + assert "chunk-1" in client.documents + + +def test_pagination_is_exhaustive_and_file_ids_are_odata_escaped(monkeypatch): + monkeypatch.setattr(azure_adapter, "_PAGE_SIZE", 2) + ingestor, client = _ingestor() + ingestor.create_collection("docs") + for index in range(5): + _write_file_manifest(ingestor, f"file-{index}", file_name=f"file-{index}.txt") + + assert len(ingestor.list_files("docs")) == 5 + crafted = "x' or file_id ne ''" + assert ingestor._delete_file_documents(crafted, "docs") == 0 + assert any("file_id eq 'x'' or file_id ne '''''" in (item or "") for item in client.search_filters) + + +def test_same_name_upload_keeps_existing_file_chunks(monkeypatch, tmp_path): + _install_reader(monkeypatch) + path = tmp_path / "document.txt" + path.write_text("new content", encoding="utf-8") + ingestor, client = _ingestor() + ingestor.create_collection("docs") + ingestor._embedding = FakeEmbedding() + ingestor._splitter = FakeSplitter([FakeNode("new")]) + client.documents["old-00000000"] = { + "id": "old-00000000", + "record_type": "chunk", + "collection_id": "docs", + "chunk": "old", + "embedding": [0.1] * 4, + "file_id": "old", + "file_name": "document.txt", + "page_number": 1, + "chunk_index": 0, + "file_size": 100, + "uploaded_at": datetime.now(UTC), + "ingested_at": datetime.now(UTC), + "metadata": "{}", + } + + count, summary, _ingested_at = ingestor._process_file( + path=str(path), + collection_name="docs", + file_id="new", + file_name="document.txt", + file_size=11, + uploaded_at=datetime.now(UTC), + metadata={"tenant": "alpha"}, + ) + + assert count == 1 + assert summary is None + assert {document["id"] for document in client.documents.values() if document.get("record_type") == "chunk"} == { + "old-00000000", + "chunk-new-00000000", + } + assert client.documents["chunk-new-00000000"]["metadata"] == '{"tenant":"alpha"}' + + +def test_failed_upload_rolls_back_new_chunks_and_preserves_existing_duplicate(monkeypatch, tmp_path): + _install_reader(monkeypatch) + path = tmp_path / "document.txt" + path.write_text("new content", encoding="utf-8") + ingestor, client = _ingestor() + ingestor.create_collection("docs") + ingestor._embedding = FakeEmbedding() + ingestor._splitter = FakeSplitter([FakeNode("first"), FakeNode("second")]) + client.documents["old-00000000"] = { + "id": "old-00000000", + "record_type": "chunk", + "collection_id": "docs", + "file_id": "old", + "file_name": "document.txt", + } + client.fail_upload_ids.add("chunk-new-00000001") + + with pytest.raises(RuntimeError, match="rejected upload"): + ingestor._process_file( + path=str(path), + collection_name="docs", + file_id="new", + file_name="document.txt", + file_size=11, + uploaded_at=datetime.now(UTC), + metadata={}, + ) + + chunk_ids = {document["id"] for document in client.documents.values() if document.get("record_type") == "chunk"} + assert chunk_ids == {"old-00000000"} + + +def test_duplicate_file_manifests_are_independent(): + ingestor, _client = _ingestor() + ingestor.create_collection("docs") + _write_file_manifest(ingestor, "old", file_name="document.txt") + _write_file_manifest(ingestor, "new", file_name="document.txt") + + assert [(item.file_id, item.file_name) for item in ingestor.list_files("docs")] == [ + ("new", "document.txt"), + ("old", "document.txt"), + ] + + +def test_metadata_counts_and_status_round_trip(): + ingestor, client = _ingestor() + ingestor.create_collection("docs", metadata={"tenant": "alpha"}) + _write_file_manifest(ingestor, "file-1", chunk_count=3) + for index in range(3): + client.documents[f"chunk-file-1-{index:08d}"] = { + "id": f"chunk-file-1-{index:08d}", + "record_type": "chunk", + "collection_id": "docs", + "file_id": "file-1", + "file_name": "report.pdf", + } + + files = ingestor.list_files("docs") + collection = ingestor.get_collection("docs") + + assert files[0].file_id == "file-1" + assert files[0].chunk_count == 3 + assert files[0].metadata == {"section": "finance"} + assert ingestor.get_file_status("file-1", "docs") == files[0] + assert collection.file_count == 1 + assert collection.chunk_count == 3 + assert collection.metadata["tenant"] == "alpha" + facet_request = next(request for request in client.search_requests if request.get("facets")) + assert facet_request["top"] == 0 + assert facet_request["facets"] == ["record_type,count:0"] + + +def test_file_status_and_deletion_do_not_cross_collection_boundaries(): + ingestor, client = _ingestor() + ingestor.create_collection("docs") + ingestor.create_collection("private") + _write_file_manifest(ingestor, "shared-id", "docs") + _write_file_manifest(ingestor, "shared-id", "private") + for collection_name in ("docs", "private"): + key = f"chunk-{collection_name}" + client.documents[key] = { + "id": key, + "record_type": "chunk", + "collection_id": collection_name, + "file_id": "shared-id", + } + + assert ingestor.delete_file("shared-id", "docs") + assert ingestor.get_file_status("shared-id", "docs") is None + assert ingestor.get_file_status("shared-id", "private") is not None + assert "chunk-docs" not in client.documents + assert "chunk-private" in client.documents + + +def test_delete_waits_for_stale_file_manifest_to_disappear(monkeypatch): + ingestor, client = _ingestor() + ingestor.create_collection("docs") + _write_file_manifest(ingestor, "file-1") + client.stale_deleted_searches_remaining = 2 + monkeypatch.setattr(azure_adapter, "_CONSISTENCY_DELAY_SECONDS", 0) + + assert ingestor.delete_file("file-1", "docs") + assert client.stale_deleted_searches_remaining == 0 + client.stale_deleted_searches_remaining = 1 + assert ingestor.list_files("docs") == [] + + +def test_delete_timeout_still_commits_bookkeeping(monkeypatch): + ingestor, client = _ingestor(generate_summary=True) + ingestor.create_collection("docs") + now = datetime.now(UTC) + _write_file_manifest( + ingestor, + "older", + summary="Older summary", + ingested_at=now - timedelta(minutes=1), + ) + newer = _write_file_manifest(ingestor, "newer", summary="Newer summary", ingested_at=now) + ingestor._files["newer"] = newer + client.stale_deleted_searches_remaining = 2 + registered = [] + timestamp_updates = [] + update_collection_timestamp = ingestor._update_collection_timestamp + + def track_collection_timestamp(collection_name): + timestamp_updates.append(collection_name) + update_collection_timestamp(collection_name) + + monkeypatch.setattr(azure_adapter, "_CONSISTENCY_ATTEMPTS", 1) + monkeypatch.setattr( + azure_adapter, + "register_summary", + lambda collection, file_name, summary: registered.append((collection, file_name, summary)), + ) + monkeypatch.setattr(ingestor, "_update_collection_timestamp", track_collection_timestamp) + + with pytest.raises(RuntimeError, match="Timed out waiting for file 'newer' to become absent"): + ingestor.delete_file("newer", "docs") + + assert "newer" not in ingestor._files + assert ingestor.get_file_status("newer", "docs") is None + assert registered == [("docs", "report.pdf", "Older summary")] + assert timestamp_updates == ["docs"] + assert not ingestor.delete_file("newer", "docs") + assert registered == [("docs", "report.pdf", "Older summary")] + assert timestamp_updates == ["docs"] + + +def test_failed_uploads_remain_visible(): + ingestor, _client = _ingestor() + ingestor.create_collection("docs") + _write_file_manifest(ingestor, "failed", file_name="failed.txt", status=FileStatus.FAILED) + + files = ingestor.list_files("docs") + + assert [(item.file_id, item.status) for item in files] == [("failed", FileStatus.FAILED)] + + +def test_ttl_deletes_only_expired_owned_collection_and_clears_summary(monkeypatch): + ingestor, client = _ingestor(generate_summary=True) + ingestor.create_collection("old") + ingestor.create_collection("new") + old_manifest = next(document for document in client.documents.values() if document.get("collection_id") == "old") + old_manifest["updated_at"] = datetime.now(UTC) - timedelta(hours=25) + ingestor._index_client.indexes["unrelated"] = SearchIndex(name="unrelated", fields=[], description="other") + cleared = [] + monkeypatch.setattr(azure_adapter, "clear_collection_summaries", cleared.append) + ingestor._ttl_hours = 24 + + ingestor._cleanup_expired_collections() + + assert ingestor.get_collection("old") is None + assert ingestor.get_collection("new") is not None + assert ingestor._physical_index_name() in ingestor._index_client.indexes + assert "unrelated" in ingestor._index_client.indexes + assert cleared == ["old"] + + +def test_collection_summary_clears_only_after_confirmed_delete(monkeypatch): + ingestor, _client = _ingestor(generate_summary=True) + ingestor.create_collection("docs") + cleared = [] + monkeypatch.setattr(azure_adapter, "clear_collection_summaries", cleared.append) + + assert ingestor.delete_collection("docs") + assert cleared == ["docs"] + assert not ingestor.delete_collection("docs") + assert cleared == ["docs"] + + +def test_file_summary_clears_only_after_confirmed_delete(monkeypatch): + ingestor, client = _ingestor(generate_summary=True) + ingestor.create_collection("docs") + _write_file_manifest(ingestor, "file-1", summary="Summary") + client.documents["chunk-file-1-00000000"] = { + "id": "chunk-file-1-00000000", + "record_type": "chunk", + "collection_id": "docs", + "file_id": "file-1", + "file_name": "report.pdf", + } + client.delete_failures_remaining["chunk-file-1-00000000"] = 10 + cleared = [] + monkeypatch.setattr(azure_adapter, "unregister_summary", lambda collection, file_name: cleared.append(file_name)) + + with pytest.raises(RuntimeError, match="rejected delete"): + ingestor.delete_file("file-1", "docs") + + assert cleared == [] + + +def test_summary_cleanup_is_skipped_when_summaries_are_disabled(monkeypatch): + ingestor, _client = _ingestor() + ingestor.create_collection("docs") + _write_file_manifest(ingestor, "file-1") + cleared = [] + monkeypatch.setattr( + azure_adapter, + "unregister_summary", + lambda collection, file_name: cleared.append((collection, file_name)), + ) + monkeypatch.setattr(azure_adapter, "clear_collection_summaries", lambda collection: cleared.append((collection,))) + + assert ingestor.delete_file("file-1", "docs") + assert ingestor.delete_collection("docs") + assert cleared == [] + + +def test_deleting_newest_duplicate_restores_previous_summary(monkeypatch): + ingestor, _client = _ingestor(generate_summary=True) + ingestor.create_collection("docs") + now = datetime.now(UTC) + _write_file_manifest( + ingestor, + "older", + file_name="report.pdf", + summary="Older summary", + ingested_at=now - timedelta(minutes=1), + ) + _write_file_manifest(ingestor, "newer", file_name="report.pdf", summary="Newer summary", ingested_at=now) + registered = [] + monkeypatch.setattr( + azure_adapter, + "register_summary", + lambda collection, file_name, summary: registered.append((collection, file_name, summary)), + ) + + assert ingestor.delete_file("newer", "docs") + assert registered == [("docs", "report.pdf", "Older summary")] + + +def test_active_file_and_collection_deletion_are_rejected(): + ingestor, _client = _ingestor() + ingestor.create_collection("docs") + _write_file_manifest(ingestor, "active", status=FileStatus.INGESTING) + + with pytest.raises(ValueError, match="while it is ingesting"): + ingestor.delete_file("active", "docs") + with pytest.raises(ValueError, match="while files are ingesting"): + ingestor.delete_collection("docs") + + +@pytest.mark.asyncio +async def test_health_checks_run_sync_sdk_off_event_loop(monkeypatch): + ingestor, _client = _ingestor() + calls = [] + + async def fake_to_thread(function, *args): + calls.append(function) + return function(*args) + + monkeypatch.setattr(asyncio, "to_thread", fake_to_thread) + + assert await ingestor.health_check() + assert len(calls) == 1 + + +def test_index_schema_uses_requested_dimensions_and_fields(): + schema = _build_index_schema("aiq-session-123456789abc", 4096) + fields = {field.name: field for field in schema.fields} + + assert fields["embedding"].vector_search_dimensions == 4096 + assert fields["file_id"].filterable + assert fields["id"].sortable + assert fields["record_type"].facetable + assert fields["collection_id"].filterable + assert schema.semantic_search is None + + +def test_splitter_uses_fixed_chunking_configuration(): + ingestor, _client = _ingestor() + + assert ingestor.splitter.chunk_size == 1024 + assert ingestor.splitter.chunk_overlap == 128 + + +def test_filename_page_normalization_and_error_translation(tmp_path): + paths = [str(tmp_path / "tmp-one"), str(tmp_path / "tmp-two")] + + assert _resolve_filenames(paths, ["one.pdf", "two.docx"]) == ["one.pdf", "two.docx"] + assert _resolve_filenames(paths, {paths[0]: "mapped.txt"}) == ["mapped.txt", "tmp-two"] + assert _coerce_page_number("4") == 4 + assert _coerce_page_number("cover") is None + assert ( + AzureAISearchIngestor._translate_error(ServiceRequestError("connection refused")) + == "AI Search service unavailable: connection refused" + ) diff --git a/uv.lock b/uv.lock index cecfd63ca..5a0c9c03a 100644 --- a/uv.lock +++ b/uv.lock @@ -309,7 +309,7 @@ requires-dist = [ { name = "sphinxcontrib-mermaid", marker = "extra == 'docs'", specifier = ">=2.0.1" }, { name = "yapf", marker = "extra == 'dev'", specifier = ">=0.40.0" }, ] -provides-extras = ["dev", "docs", "s3", "viz"] +provides-extras = ["s3", "dev", "docs", "viz"] [package.metadata.requires-dev] dev = [ @@ -593,6 +593,59 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/fb/95/adcb68e20c34162e9135f370d6e31737719c2b6f94bc953fe7ed1f10fe21/authlib-1.7.2-py2.py3-none-any.whl", hash = "sha256:3e1faedc9d87e7d56a164eca3ccb6ace0d61b94abe83e92242f8dc8bba9b4a9f", size = 259548, upload-time = "2026-05-06T08:10:21.436Z" }, ] +[[package]] +name = "azure-common" +version = "1.1.28" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/3e/71/f6f71a276e2e69264a97ad39ef850dca0a04fce67b12570730cb38d0ccac/azure-common-1.1.28.zip", hash = "sha256:4ac0cd3214e36b6a1b6a442686722a5d8cc449603aa833f3f0f40bda836704a3", size = 20914, upload-time = "2022-02-03T19:39:44.373Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/62/55/7f118b9c1b23ec15ca05d15a578d8207aa1706bc6f7c87218efffbbf875d/azure_common-1.1.28-py2.py3-none-any.whl", hash = "sha256:5c12d3dcf4ec20599ca6b0d3e09e86e146353d443e7fcc050c9a19c1f9df20ad", size = 14462, upload-time = "2022-02-03T19:39:42.417Z" }, +] + +[[package]] +name = "azure-core" +version = "1.41.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "requests" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/a6/f3/b416179e408990df5db0d516283022dde0f5d0111d98c1a848e41853e81c/azure_core-1.41.0.tar.gz", hash = "sha256:f46ff5dfcd230f25cf1c19e8a34b8dc08a337b2503e268bb600a16c00db8ad5a", size = 381042, upload-time = "2026-05-07T23:30:54.302Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5b/db/325c6d7312d2200251c52323878281045aaffcb5586612296484e4280eaa/azure_core-1.41.0-py3-none-any.whl", hash = "sha256:522b4011e8180b1a3dcd2024396a4e7fe9ac37fb8597db47163d230b5efe892d", size = 220920, upload-time = "2026-05-07T23:30:56.357Z" }, +] + +[[package]] +name = "azure-identity" +version = "1.25.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "azure-core" }, + { name = "cryptography" }, + { name = "msal" }, + { name = "msal-extensions" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/c5/0e/3a63efb48aa4a5ae2cfca61ee152fbcb668092134d3eb8bfda472dd5c617/azure_identity-1.25.3.tar.gz", hash = "sha256:ab23c0d63015f50b630ef6c6cf395e7262f439ce06e5d07a64e874c724f8d9e6", size = 286304, upload-time = "2026-03-13T01:12:20.892Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/49/9a/417b3a533e01953a7c618884df2cb05a71e7b68bdbce4fbdb62349d2a2e8/azure_identity-1.25.3-py3-none-any.whl", hash = "sha256:f4d0b956a8146f30333e071374171f3cfa7bdb8073adb8c3814b65567aa7447c", size = 192138, upload-time = "2026-03-13T01:12:22.951Z" }, +] + +[[package]] +name = "azure-search-documents" +version = "11.6.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "azure-common" }, + { name = "azure-core" }, + { name = "isodate" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/cf/68/9d59a0bed5fd9581b45444e8abc3ecda97e0466ae0f03affc7cddfb9fa74/azure_search_documents-11.6.0.tar.gz", hash = "sha256:fcc807076ff82024be576ffccb0d0f3261e5c2a112a6666b86ec70bbdb2e1d64", size = 311194, upload-time = "2025-10-09T22:04:03.655Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c5/4c/d74e5c3ccc0b9ead0e400a2d70ded67554b56a5d799aaa8bf5baaacf4aea/azure_search_documents-11.6.0-py3-none-any.whl", hash = "sha256:c3eb2deaf7926844e99a881830861225ef68e8b3bc067a76019e87fc7f5586dc", size = 307935, upload-time = "2025-10-09T22:04:05.008Z" }, +] + [[package]] name = "babel" version = "2.18.0" @@ -1370,6 +1423,15 @@ requires-dist = [ { name = "pydantic", specifier = ">=2.0.0" }, ] +[[package]] +name = "defusedxml" +version = "0.8.0rc2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/5e/3b/b8849dcc3f96913924137dc4ea041d74aa513a3c5dda83d8366491290c74/defusedxml-0.8.0rc2.tar.gz", hash = "sha256:138c7d540a78775182206c7c97fe65b246a2f40b29471e1a2f1b0da76e7a3942", size = 52575, upload-time = "2023-09-29T08:01:27.517Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5d/c7/6b4ad89ca6f7732ff97ce5e9caa6fe739600d26c5d53c20d0bf9abb79ec5/defusedxml-0.8.0rc2-py2.py3-none-any.whl", hash = "sha256:1c812964311154c3bf4aaf3bc1443b31ee13530b7f255eaaa062c0553c76103d", size = 25756, upload-time = "2023-09-29T08:01:25.515Z" }, +] + [[package]] name = "deprecated" version = "1.3.1" @@ -2221,6 +2283,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12", size = 7484, upload-time = "2025-10-18T21:55:41.639Z" }, ] +[[package]] +name = "isodate" +version = "0.7.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/54/4d/e940025e2ce31a8ce1202635910747e5a87cc3a6a6bb2d00973375014749/isodate-0.7.2.tar.gz", hash = "sha256:4cd1aa0f43ca76f4a6c6c0292a85f40b35ec2e43e315b59f06e6d32171a953e6", size = 29705, upload-time = "2024-10-08T23:04:11.5Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/15/aa/0aca39a37d3c7eb941ba736ede56d689e7be91cab5d9ca846bde3999eba6/isodate-0.7.2-py3-none-any.whl", hash = "sha256:28009937d8031054830160fce6d409ed342816b543597cece116d966c6d99e15", size = 22320, upload-time = "2024-10-08T23:04:09.501Z" }, +] + [[package]] name = "jinja2" version = "3.1.6" @@ -2471,12 +2542,16 @@ dependencies = [ [package.optional-dependencies] all = [ + { name = "azure-identity" }, + { name = "azure-search-documents" }, { name = "boto3" }, { name = "chromadb" }, { name = "distributed" }, { name = "docx2txt" }, { name = "llama-index" }, + { name = "llama-index-core" }, { name = "llama-index-embeddings-nvidia" }, + { name = "llama-index-readers-file" }, { name = "llama-index-vector-stores-chroma" }, { name = "openai" }, { name = "opensearch-py" }, @@ -2488,6 +2563,15 @@ all = [ { name = "requests" }, { name = "urllib3" }, ] +azure-ai-search = [ + { name = "azure-identity" }, + { name = "azure-search-documents" }, + { name = "docx2txt" }, + { name = "llama-index-core" }, + { name = "llama-index-embeddings-nvidia" }, + { name = "llama-index-readers-file" }, + { name = "pypdf" }, +] foundational-rag = [ { name = "docx2txt" }, { name = "python-pptx" }, @@ -2517,16 +2601,22 @@ opensearch = [ [package.metadata] requires-dist = [ + { name = "azure-identity", marker = "extra == 'azure-ai-search'", specifier = ">=1.15.0,<1.26" }, + { name = "azure-search-documents", marker = "extra == 'azure-ai-search'", specifier = ">=11.6.0,<11.7" }, { name = "boto3", marker = "extra == 'opensearch'", specifier = ">=1.28.0" }, { name = "chromadb", marker = "extra == 'llamaindex'", specifier = ">=0.4.0" }, { name = "distributed", marker = "extra == 'opensearch'", specifier = ">=2024.1.0" }, + { name = "docx2txt", marker = "extra == 'azure-ai-search'", specifier = ">=0.8" }, { name = "docx2txt", marker = "extra == 'foundational-rag'", specifier = ">=0.8" }, { name = "docx2txt", marker = "extra == 'llamaindex'", specifier = ">=0.8" }, { name = "docx2txt", marker = "extra == 'opensearch'", specifier = ">=0.8" }, { name = "httpx", specifier = ">=0.24.0" }, - { name = "knowledge-layer", extras = ["llamaindex", "foundational-rag", "opensearch"], marker = "extra == 'all'", editable = "sources/knowledge_layer" }, + { name = "knowledge-layer", extras = ["llamaindex", "foundational-rag", "opensearch", "azure-ai-search"], marker = "extra == 'all'", editable = "sources/knowledge_layer" }, { name = "llama-index", marker = "extra == 'llamaindex'", specifier = ">=0.10.0" }, + { name = "llama-index-core", marker = "extra == 'azure-ai-search'", specifier = ">=0.11.0" }, + { name = "llama-index-embeddings-nvidia", marker = "extra == 'azure-ai-search'", specifier = ">=0.2.0" }, { name = "llama-index-embeddings-nvidia", marker = "extra == 'llamaindex'", specifier = ">=0.1.0" }, + { name = "llama-index-readers-file", marker = "extra == 'azure-ai-search'", specifier = ">=0.2.0" }, { name = "llama-index-vector-stores-chroma", marker = "extra == 'llamaindex'", specifier = ">=0.1.0" }, { name = "openai", marker = "extra == 'llamaindex'", specifier = ">=1.0.0" }, { name = "openai", marker = "extra == 'opensearch'", specifier = ">=1.0.0" }, @@ -2534,6 +2624,7 @@ requires-dist = [ { name = "pdfplumber", marker = "extra == 'llamaindex'", specifier = ">=0.10.0" }, { name = "pillow", marker = "extra == 'llamaindex'", specifier = ">=9.0.0" }, { name = "pydantic", specifier = ">=2.0.0" }, + { name = "pypdf", marker = "extra == 'azure-ai-search'", specifier = ">=4.0.0" }, { name = "pypdf", marker = "extra == 'opensearch'", specifier = ">=4.0.0" }, { name = "pypdfium2", marker = "extra == 'llamaindex'", specifier = ">=5.0.0" }, { name = "python-dotenv", specifier = ">=1.0.0" }, @@ -2542,7 +2633,7 @@ requires-dist = [ { name = "requests", marker = "extra == 'foundational-rag'", specifier = ">=2.28.0" }, { name = "urllib3", marker = "extra == 'foundational-rag'", specifier = ">=2.7.0,<3" }, ] -provides-extras = ["llamaindex", "foundational-rag", "opensearch", "all"] +provides-extras = ["llamaindex", "foundational-rag", "opensearch", "azure-ai-search", "all"] [[package]] name = "kubernetes" @@ -3134,6 +3225,23 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/62/c2/cfa341cff00154b58abed3fc559fb43c2531d5dd70eefb75f230704d4c3c/llama_index_llms_openai-0.7.9-py3-none-any.whl", hash = "sha256:0bc8f59faddf8dbc9f90c5576127ab2fa6d368be5078885dc7e611af129b5a79", size = 28648, upload-time = "2026-05-29T15:32:35.997Z" }, ] +[[package]] +name = "llama-index-readers-file" +version = "0.6.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "beautifulsoup4" }, + { name = "defusedxml" }, + { name = "llama-index-core" }, + { name = "pandas" }, + { name = "pypdf" }, + { name = "striprtf" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/29/b9/5d5ecb81febd544a04b38f4ef7d3a69ec0625204f3c0f3d8200b70f16c2d/llama_index_readers_file-0.6.0.tar.gz", hash = "sha256:ff366d6ff5ecb7119275ac859310d8b672d8b6b3261afae02f4084fce9076bd0", size = 32537, upload-time = "2026-03-12T20:34:32.309Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7e/ab/ff4b2e14a63708abc0115ccb13e41b324972c0346d9b82d1fb8dbb77873a/llama_index_readers_file-0.6.0-py3-none-any.whl", hash = "sha256:1026d94f2d5902152373bc2c3b7caa7e216d956620b22d510e516850b6a7440d", size = 51830, upload-time = "2026-03-12T20:34:33.315Z" }, +] + [[package]] name = "llama-index-vector-stores-chroma" version = "0.5.5" @@ -3506,6 +3614,32 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/56/54/400262056c144ceee5edab40efa2541ae8928ae5f244fd9025f3ad26c909/modal-1.4.3-py3-none-any.whl", hash = "sha256:802917181f576458a0cb833322157dab09c4f367326426c5a732661a0c519577", size = 826232, upload-time = "2026-05-18T22:34:43.335Z" }, ] +[[package]] +name = "msal" +version = "1.38.0rc1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cryptography" }, + { name = "pyjwt", extra = ["crypto"] }, + { name = "requests" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/b1/ba/afc0474f72674e1a19155afe7b6b3a8b12359d2aee458e216f4dd8030f5b/msal-1.38.0rc1.tar.gz", hash = "sha256:c0160b98217f84705339189d0fa5099cdec0ffa5986e3d2053f450007a2f1a89", size = 182430, upload-time = "2026-06-05T15:20:47.287Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d9/9f/873565789342901574aa85d228c6ed6761addf9f5a4829217e7f94671700/msal-1.38.0rc1-py3-none-any.whl", hash = "sha256:4c3528a473d856f725c8f9b302c364c576e21b5cf43a581273c2682d973c4da3", size = 123766, upload-time = "2026-06-05T15:20:48.981Z" }, +] + +[[package]] +name = "msal-extensions" +version = "1.3.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "msal" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/01/99/5d239b6156eddf761a636bded1118414d161bd6b7b37a9335549ed159396/msal_extensions-1.3.1.tar.gz", hash = "sha256:c5b0fd10f65ef62b5f1d62f4251d51cbcaf003fcedae8c91b040a488614be1a4", size = 23315, upload-time = "2025-03-14T23:51:03.902Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5e/75/bd9b7bb966668920f06b200e84454c8f3566b102183bc55c5473d96cb2b9/msal_extensions-1.3.1-py3-none-any.whl", hash = "sha256:96d3de4d034504e969ac5e85bae8106c8373b5c6568e4c8fa7af2eca9dbe6bca", size = 20583, upload-time = "2025-03-14T23:51:03.016Z" }, +] + [[package]] name = "msgpack" version = "1.1.2" @@ -5962,23 +6096,23 @@ resolution-markers = [ "python_full_version < '3.12'", ] dependencies = [ - { name = "alabaster", marker = "python_full_version < '3.12'" }, - { name = "babel", marker = "python_full_version < '3.12'" }, - { name = "colorama", marker = "python_full_version < '3.12' and sys_platform == 'win32'" }, - { name = "docutils", marker = "python_full_version < '3.12'" }, - { name = "imagesize", marker = "python_full_version < '3.12'" }, - { name = "jinja2", marker = "python_full_version < '3.12'" }, - { name = "packaging", marker = "python_full_version < '3.12'" }, - { name = "pygments", marker = "python_full_version < '3.12'" }, - { name = "requests", marker = "python_full_version < '3.12'" }, - { name = "roman-numerals", marker = "python_full_version < '3.12'" }, - { name = "snowballstemmer", marker = "python_full_version < '3.12'" }, - { name = "sphinxcontrib-applehelp", marker = "python_full_version < '3.12'" }, - { name = "sphinxcontrib-devhelp", marker = "python_full_version < '3.12'" }, - { name = "sphinxcontrib-htmlhelp", marker = "python_full_version < '3.12'" }, - { name = "sphinxcontrib-jsmath", marker = "python_full_version < '3.12'" }, - { name = "sphinxcontrib-qthelp", marker = "python_full_version < '3.12'" }, - { name = "sphinxcontrib-serializinghtml", marker = "python_full_version < '3.12'" }, + { name = "alabaster" }, + { name = "babel" }, + { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "docutils" }, + { name = "imagesize" }, + { name = "jinja2" }, + { name = "packaging" }, + { name = "pygments" }, + { name = "requests" }, + { name = "roman-numerals" }, + { name = "snowballstemmer" }, + { name = "sphinxcontrib-applehelp" }, + { name = "sphinxcontrib-devhelp" }, + { name = "sphinxcontrib-htmlhelp" }, + { name = "sphinxcontrib-jsmath" }, + { name = "sphinxcontrib-qthelp" }, + { name = "sphinxcontrib-serializinghtml" }, ] sdist = { url = "https://files.pythonhosted.org/packages/42/50/a8c6ccc36d5eacdfd7913ddccd15a9cee03ecafc5ee2bc40e1f168d85022/sphinx-9.0.4.tar.gz", hash = "sha256:594ef59d042972abbc581d8baa577404abe4e6c3b04ef61bd7fc2acbd51f3fa3", size = 8710502, upload-time = "2025-12-04T07:45:27.343Z" } wheels = [ @@ -5994,23 +6128,23 @@ resolution-markers = [ "python_full_version == '3.12.*'", ] dependencies = [ - { name = "alabaster", marker = "python_full_version >= '3.12'" }, - { name = "babel", marker = "python_full_version >= '3.12'" }, - { name = "colorama", marker = "python_full_version >= '3.12' and sys_platform == 'win32'" }, - { name = "docutils", marker = "python_full_version >= '3.12'" }, - { name = "imagesize", marker = "python_full_version >= '3.12'" }, - { name = "jinja2", marker = "python_full_version >= '3.12'" }, - { name = "packaging", marker = "python_full_version >= '3.12'" }, - { name = "pygments", marker = "python_full_version >= '3.12'" }, - { name = "requests", marker = "python_full_version >= '3.12'" }, - { name = "roman-numerals", marker = "python_full_version >= '3.12'" }, - { name = "snowballstemmer", marker = "python_full_version >= '3.12'" }, - { name = "sphinxcontrib-applehelp", marker = "python_full_version >= '3.12'" }, - { name = "sphinxcontrib-devhelp", marker = "python_full_version >= '3.12'" }, - { name = "sphinxcontrib-htmlhelp", marker = "python_full_version >= '3.12'" }, - { name = "sphinxcontrib-jsmath", marker = "python_full_version >= '3.12'" }, - { name = "sphinxcontrib-qthelp", marker = "python_full_version >= '3.12'" }, - { name = "sphinxcontrib-serializinghtml", marker = "python_full_version >= '3.12'" }, + { name = "alabaster" }, + { name = "babel" }, + { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "docutils" }, + { name = "imagesize" }, + { name = "jinja2" }, + { name = "packaging" }, + { name = "pygments" }, + { name = "requests" }, + { name = "roman-numerals" }, + { name = "snowballstemmer" }, + { name = "sphinxcontrib-applehelp" }, + { name = "sphinxcontrib-devhelp" }, + { name = "sphinxcontrib-htmlhelp" }, + { name = "sphinxcontrib-jsmath" }, + { name = "sphinxcontrib-qthelp" }, + { name = "sphinxcontrib-serializinghtml" }, ] sdist = { url = "https://files.pythonhosted.org/packages/cd/bd/f08eb0f4eed5c83f1ba2a3bd18f7745a2b1525fad70660a1c00224ec468a/sphinx-9.1.0.tar.gz", hash = "sha256:7741722357dd75f8190766926071fed3bdc211c74dd2d7d4df5404da95930ddb", size = 8718324, upload-time = "2025-12-31T15:09:27.646Z" } wheels = [ @@ -6247,6 +6381,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/1c/54/196d0c1db10af76baa4f64894448505d60d3cdf70ef92cbb35f46a4e4c71/starlette-1.2.1-py3-none-any.whl", hash = "sha256:4de0082d08c8f6764a85a54cf1120d6939507a19905c7768acad2a9f875d2b89", size = 73350, upload-time = "2026-05-31T01:07:50.09Z" }, ] +[[package]] +name = "striprtf" +version = "0.0.26" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/25/20/3d419008265346452d09e5dadfd5d045b64b40d8fc31af40588e6c76997a/striprtf-0.0.26.tar.gz", hash = "sha256:fdb2bba7ac440072d1c41eab50d8d74ae88f60a8b6575c6e2c7805dc462093aa", size = 6258, upload-time = "2023-07-20T14:30:36.29Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a3/cf/0fea4f4ba3fc2772ac2419278aa9f6964124d4302117d61bc055758e000c/striprtf-0.0.26-py3-none-any.whl", hash = "sha256:8c8f9d32083cdc2e8bfb149455aa1cc5a4e0a035893bedc75db8b73becb3a1bb", size = 6914, upload-time = "2023-07-20T14:30:35.338Z" }, +] + [[package]] name = "synchronicity" version = "0.12.3"