From 79244d3956ddc8742bf3eb7534fc64295726f022 Mon Sep 17 00:00:00 2001 From: Harmke Alkemade Date: Fri, 3 Jul 2026 11:58:51 +0200 Subject: [PATCH 01/27] feat(knowledge): add Azure AI Search backend Signed-off-by: Harmke Alkemade --- docs/source/examples/azure-ai-search.md | 53 + docs/source/examples/index.md | 1 + .../knowledge_layer/KNOWLEDGE-LAYER-SETUP.md | 2 + sources/knowledge_layer/README.md | 7 + sources/knowledge_layer/pyproject.toml | 13 +- sources/knowledge_layer/src/__init__.py | 1 + .../src/azure_ai_search/README.md | 67 ++ .../src/azure_ai_search/__init__.py | 9 + .../src/azure_ai_search/adapter.py | 971 ++++++++++++++++++ sources/knowledge_layer/src/register.py | 66 +- .../test_azure_ai_search.py | 182 ++++ uv.lock | 217 +++- 12 files changed, 1548 insertions(+), 41 deletions(-) create mode 100644 docs/source/examples/azure-ai-search.md create mode 100644 sources/knowledge_layer/src/azure_ai_search/README.md create mode 100644 sources/knowledge_layer/src/azure_ai_search/__init__.py create mode 100644 sources/knowledge_layer/src/azure_ai_search/adapter.py create mode 100644 tests/knowledge_layer_tests/test_azure_ai_search.py diff --git a/docs/source/examples/azure-ai-search.md b/docs/source/examples/azure-ai-search.md new file mode 100644 index 000000000..c4bc37d4b --- /dev/null +++ b/docs/source/examples/azure-ai-search.md @@ -0,0 +1,53 @@ + + +# 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. + +Install the backend dependency: + +```bash +uv pip install -e "sources/knowledge_layer[azure_ai_search]" +``` + +Replace the `knowledge_search` block in a web configuration such as +`configs/config_web_default_llamaindex.yml`: + +```yaml +functions: + knowledge_search: + _type: knowledge_retrieval + backend: azure_ai_search + collection_name: ${COLLECTION_NAME:-aiq_default} + top_k: 5 + + azure_search_endpoint: https://.search.windows.net + azure_search_auth_mode: managed_identity + + embed_endpoint: https://integrate.api.nvidia.com/v1 + embed_model: nvidia/nv-embed-v1 + embed_dim: 4096 + use_hybrid: true + use_semantic_ranker: true + + generate_summary: true + summary_model: summary_llm + summary_db: ${AIQ_SUMMARY_DB:-sqlite+aiosqlite:///./summaries.db} +``` + +For API-key authentication, set `azure_search_auth_mode: api_key` and add +`azure_search_api_key: ${AZURE_SEARCH_API_KEY}`. Managed identity uses +`DefaultAzureCredential`; set `AZURE_CLIENT_ID` to select a user-assigned +identity. If `embed_api_key` is omitted, the NVIDIA embedding client reads +`NVIDIA_API_KEY`. + +Existing indexes must use the configured `embed_dim`. Delete and re-ingest a +collection when changing embedding dimensions. Frontend WebSocket queries use +the conversation ID as the collection; direct API tests must supply equivalent +context or query the configured fallback collection. diff --git a/docs/source/examples/index.md b/docs/source/examples/index.md index 8c0854156..1360fe58d 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 and semantic document retrieval | `config_web_default_llamaindex.yml` base | | [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 Modal sandbox execution for quantitative research workflows | `config_domain_routing_and_skills.yml` | diff --git a/sources/knowledge_layer/KNOWLEDGE-LAYER-SETUP.md b/sources/knowledge_layer/KNOWLEDGE-LAYER-SETUP.md index 6adb54aff..350d169d9 100644 --- a/sources/knowledge_layer/KNOWLEDGE-LAYER-SETUP.md +++ b/sources/knowledge_layer/KNOWLEDGE-LAYER-SETUP.md @@ -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 | Hybrid and semantic 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. @@ -64,6 +65,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. diff --git a/sources/knowledge_layer/README.md b/sources/knowledge_layer/README.md index 53dd9b250..e1812a239 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 and semantic 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..435a84779 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.5.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..dda36060c --- /dev/null +++ b/sources/knowledge_layer/src/azure_ai_search/README.md @@ -0,0 +1,67 @@ + + +# 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 + +```yaml +functions: + knowledge_search: + _type: knowledge_retrieval + backend: azure_ai_search + collection_name: ${COLLECTION_NAME:-aiq_default} + + azure_search_endpoint: https://.search.windows.net + azure_search_auth_mode: managed_identity + # For key authentication instead: + # azure_search_auth_mode: api_key + # azure_search_api_key: ${AZURE_SEARCH_API_KEY} + + embed_endpoint: https://integrate.api.nvidia.com/v1 + embed_model: nvidia/nv-embed-v1 + embed_dim: 4096 + # NVIDIAEmbedding reads NVIDIA_API_KEY when this field is omitted. + # embed_api_key: ${NVIDIA_API_KEY} + + use_hybrid: true + use_semantic_ranker: true + top_k: 5 + chunk_size: 512 + chunk_overlap: 64 + + generate_summary: true + summary_model: summary_llm + summary_db: ${AIQ_SUMMARY_DB:-sqlite+aiosqlite:///./summaries.db} +``` + +Managed identity uses `DefaultAzureCredential`. Set `AZURE_CLIENT_ID` when a +user-assigned identity should be selected. The identity needs permission to +create and delete indexes and to read, write, and delete index documents. + +The adapter parses PDF, DOCX, TXT, and Markdown uploads with LlamaIndex, +creates one Azure AI Search index per AI-Q collection, and performs vector or +hybrid retrieval with optional semantic ranking. AI-Q frontend conversations +use their conversation ID as the collection name, keeping uploads and +WebSocket retrieval in the same index. + +`embed_dim` must match both the embedding model output and any existing index. +Changing from a 2048-dimensional model to `nvidia/nv-embed-v1` at 4096 +dimensions requires deleting and re-ingesting the old collection. The adapter +does not alter an existing index 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..0fba51c04 --- /dev/null +++ b/sources/knowledge_layer/src/azure_ai_search/adapter.py @@ -0,0 +1,971 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import asyncio +import hashlib +import logging +import os +import re +import threading +import uuid +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 ResourceExistsError +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 SemanticConfiguration +from azure.search.documents.indexes.models import SemanticField +from azure.search.documents.indexes.models import SemanticPrioritizedFields +from azure.search.documents.indexes.models import SemanticSearch +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 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.schema import FileStatus + +logger = logging.getLogger(__name__) + +_BACKEND_NAME = "azure_ai_search" +_SEMANTIC_CONFIG = "default-semantic" + + +def _coerce_config( + config: dict[str, Any] | None, +) -> SimpleNamespace: + """Expose the validated KnowledgeRetrievalConfig values as attributes.""" + if not config: + raise ValueError("Azure AI Search configuration is required") + return SimpleNamespace(**config) + + +def _build_search_credential(cfg: SimpleNamespace): + """Pick the right Azure SDK credential based on auth_mode.""" + if cfg.auth_mode == "api_key": + api_key = _secret_value(cfg.api_key) + if api_key is None: + raise ValueError("auth_mode=api_key requires the `api_key` field to be set") + return AzureKeyCredential(api_key) + # DefaultAzureCredential honors AZURE_CLIENT_ID for user-assigned identities. + return DefaultAzureCredential() + + +def _secret_value(value: Any) -> str | None: + """Unwrap SecretStr values only when an SDK client needs them.""" + if value is None: + return None + if hasattr(value, "get_secret_value"): + return value.get_secret_value() + return str(value) + + +# ============================================================================= +# Retriever +# ============================================================================= + + +@register_retriever(_BACKEND_NAME) +class AzureAISearchRetriever(BaseRetriever): + """Hybrid + semantic-ranked retriever backed by Azure AI Search. + + Per AI Search index, the schema has at minimum: `id` (key), `chunk` (text, + searchable), `embedding` (Collection(Edm.Single), vectorSearchProfile), + plus `doc_id`, `file_name`, `page_number`, and `metadata`. + """ + + backend_name = _BACKEND_NAME + + def __init__(self, config: dict[str, Any] | None = None): + super().__init__(config) + cfg = _coerce_config(self.config) + self.cfg = cfg + + # Lazy embedding-client init — `llama_index.embeddings.nvidia.NVIDIAEmbedding` + # is imported at first use to keep the package installable in environments + # where llama-index isn't present (unit tests). + self._embedding = None + + # SearchClient is bound to a single index, so we cache one per collection. + self._credential = _build_search_credential(cfg) + self._client_cache: dict[str, SearchClient] = {} + + logger.info( + "AzureAISearchRetriever initialized: endpoint=%s embed_model=%s embed_endpoint=%s " + "use_hybrid=%s use_semantic_ranker=%s", + cfg.endpoint, + cfg.embed_model, + cfg.embed_endpoint, + cfg.use_hybrid, + cfg.use_semantic_ranker, + ) + + @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_endpoint), + api_key=_secret_value(self.cfg.embed_api_key), + ) + return self._embedding + + def _get_client(self, collection_name: str) -> SearchClient: + if collection_name not in self._client_cache: + self._client_cache[collection_name] = SearchClient( + endpoint=str(self.cfg.endpoint), + index_name=collection_name, + credential=self._credential, + ) + return self._client_cache[collection_name] + + # ------- abstract methods ------- + + async def retrieve( + self, + query: str, + collection_name: str, + top_k: int = 10, + filters: dict[str, Any] | None = None, + ) -> RetrievalResult: + """Run the synchronous Azure Search client without blocking the event loop.""" + 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: + """Hybrid + (optionally) semantic-ranked search. Never raises.""" + try: + # 1. Embed the query + query_vector = self.embedding.get_query_embedding(query) + + # 2. Build the search request + client = self._get_client(collection_name) + vector_query = VectorizedQuery( + vector=query_vector, + # Over-fetch on the vector side so the reranker has more + # candidates to choose from — semantic ranking only re-orders + # the input set, it doesn't reach back into the index. + k_nearest_neighbors=max(top_k * 3, 20), + fields="embedding", + ) + + search_params: dict[str, Any] = { + "vector_queries": [vector_query], + "top": top_k, + "select": ["id", "chunk", "doc_id", "file_name", "page_number", "metadata"], + } + if self.cfg.use_hybrid: + # Hybrid = include lexical text alongside the vector query. + search_params["search_text"] = query + if self.cfg.use_semantic_ranker: + search_params["query_type"] = "semantic" + search_params["semantic_configuration_name"] = _SEMANTIC_CONFIG + if filters and isinstance(filters, dict): + # Pass through a pre-built OData filter string under "$filter". + if odata := filters.get("$filter"): + search_params["filter"] = odata + + # 3. Run the search + raw_results = client.search(**search_params) + chunks = [self.normalize(hit) for hit in raw_results] + + logger.info( + "retrieve OK: collection=%s query=%r returned=%d top_k=%d", + collection_name, + query[:80], + len(chunks), + top_k, + ) + return RetrievalResult( + query=query, + backend=_BACKEND_NAME, + chunks=chunks, + success=True, + ) + + except ResourceNotFoundError: + msg = f"AI Search index {collection_name!r} not found. Has the collection been created yet?" + logger.warning(msg) + return RetrievalResult(query=query, backend=_BACKEND_NAME, chunks=[], success=False, error_message=msg) + except ClientAuthenticationError as e: + msg = f"AI Search authentication failed: {e!s}" + logger.error(msg) + return RetrievalResult(query=query, backend=_BACKEND_NAME, chunks=[], success=False, error_message=msg) + except ServiceRequestError as e: + msg = f"AI Search service unavailable: {e!s}" + logger.error(msg) + return RetrievalResult(query=query, backend=_BACKEND_NAME, chunks=[], success=False, error_message=msg) + except HttpResponseError as e: + msg = f"AI Search request failed: {e.status_code} {e.reason or e!s}" + logger.error(msg, exc_info=True) + return RetrievalResult(query=query, backend=_BACKEND_NAME, chunks=[], success=False, error_message=msg) + except Exception as e: # noqa: BLE001 — surface any unexpected error as a UI-readable string + msg = f"Unexpected error during retrieval: {e!s}" + logger.exception("retrieve unexpected error") + return RetrievalResult(query=query, backend=_BACKEND_NAME, chunks=[], success=False, error_message=msg) + + def normalize(self, raw_result: Any) -> Chunk: + """Map one AI Search hit dict to a 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")) + doc_id = raw_result.get("doc_id") + + # Score: prefer the semantic reranker score (0–4 typical range, normalise + # to 0–1) when available, otherwise the search score (already 0+ from + # hybrid RRF; clamp to 1). + rerank_score = raw_result.get("@search.reranker_score") + search_score = raw_result.get("@search.score") or 0.0 + if rerank_score is not None: + score = min(max(float(rerank_score) / 4.0, 0.0), 1.0) + else: + score = min(max(float(search_score), 0.0), 1.0) + + if page_number is not None: + display_citation = f"{file_name}, p.{page_number}" + else: + display_citation = str(file_name) + + metadata: dict[str, Any] = {} + if doc_id: + metadata["doc_id"] = doc_id + # Pass through any raw metadata blob if present + if raw_meta := raw_result.get("metadata"): + metadata["raw"] = raw_meta + + return Chunk( + chunk_id=str(chunk_id), + content=str(content), + score=score, + file_name=str(file_name), + page_number=page_number, + display_citation=display_citation, + content_type=ContentType.TEXT, + metadata=metadata, + ) + + # ------- optional ------- + + async def health_check(self) -> bool: + """Return True if the configured collection's index is reachable.""" + try: + client = self._get_client(self.cfg.collection_name) + client.get_document_count() + return True + except Exception: # noqa: BLE001 + logger.exception("Retriever health_check failed") + return False + + +# ============================================================================= +# Index schema builder +# ============================================================================= + + +# AI Search index names: lowercase letters, digits, hyphens, underscores; +# 2-128 chars; must start and end with a letter or digit; no consecutive dashes. +# (The official docs say "letters, numbers, dashes" but the service actually +# accepts underscores too, and the AI-Q frontend generates session-style +# collection names like "s_f031d9cf_...".) +_INDEX_NAME_RE = re.compile(r"^[a-z0-9](?:[a-z0-9_-]{0,126}[a-z0-9])?$") + + +def _validate_index_name(name: str) -> None: + if not _INDEX_NAME_RE.match(name) or "--" in name: + raise ValueError( + f"Invalid AI Search index name {name!r}. Names must be lowercase, " + "use only letters/digits/hyphens/underscores, be 2-128 chars, " + "start and end with a letter or digit, and not contain '--'." + ) + + +def _build_index_schema(name: str, embed_dim: int) -> SearchIndex: + """Build the canonical Azure AI Search index schema for a collection. + + Anything that queries the index must agree on field names and dimensions. + """ + return SearchIndex( + name=name, + fields=[ + SimpleField(name="id", type=SearchFieldDataType.String, key=True, filterable=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="doc_id", type=SearchFieldDataType.String, filterable=True), + SearchableField(name="file_name", filterable=True, sortable=True), + SimpleField(name="page_number", type=SearchFieldDataType.Int32, 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"), + ], + ), + semantic_search=SemanticSearch( + default_configuration_name=_SEMANTIC_CONFIG, + configurations=[ + SemanticConfiguration( + name=_SEMANTIC_CONFIG, + prioritized_fields=SemanticPrioritizedFields( + title_field=SemanticField(field_name="file_name"), + content_fields=[SemanticField(field_name="chunk")], + ), + ), + ], + ), + ) + + +# ============================================================================= +# Ingestor +# ============================================================================= + + +@register_ingestor(_BACKEND_NAME) +class AzureAISearchIngestor(BaseIngestor): + """Parse uploaded files, embed chunks, and write them to Azure AI Search. + + Submission is non-blocking — `submit_job` spawns a daemon thread and returns + a job_id immediately. `get_job_status` reflects the live state. The HTTP + layer routes file uploads here via `set_active_ingestor` (called from + register.py). + """ + + backend_name = _BACKEND_NAME + + def __init__(self, config: dict[str, Any] | None = None): + super().__init__(config) + cfg = _coerce_config(self.config) + self.cfg = cfg + + # Azure SDK clients + self._credential = _build_search_credential(cfg) + self._index_client = SearchIndexClient(endpoint=str(cfg.endpoint), credential=self._credential) + self._search_client_cache: dict[str, SearchClient] = {} + + # Lazy-init heavy llama-index/embedding deps + self._embedding = None + self._splitter = None + + # Summary LLM (LangChain wrapper) resolved by shared registration. + self._summary_llm = cfg.summary_llm + + # Job + per-collection in-memory state. AI Search is the source of + # truth for chunks; this just tracks submission lifecycle. + self._jobs_lock = threading.Lock() + self._jobs: dict[str, IngestionJobStatus] = {} + + logger.info( + "AzureAISearchIngestor initialized: endpoint=%s embed_dim=%d chunk_size=%d", + cfg.endpoint, + cfg.embed_dim, + cfg.chunk_size, + ) + + # ------- lazy clients ------- + + @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_endpoint), + api_key=_secret_value(self.cfg.embed_api_key), + ) + return self._embedding + + @property + def splitter(self): + if self._splitter is None: + from llama_index.core.node_parser import SentenceSplitter + + self._splitter = SentenceSplitter( + chunk_size=self.cfg.chunk_size, + chunk_overlap=self.cfg.chunk_overlap, + ) + return self._splitter + + def _get_search_client(self, collection_name: str) -> SearchClient: + if collection_name not in self._search_client_cache: + self._search_client_cache[collection_name] = SearchClient( + endpoint=str(self.cfg.endpoint), + index_name=collection_name, + credential=self._credential, + ) + return self._search_client_cache[collection_name] + + # ------- index management ------- + + def _ensure_index(self, collection_name: str) -> None: + """Create the AI Search index if it doesn't already exist.""" + _validate_index_name(collection_name) + try: + self._index_client.get_index(collection_name) + logger.debug("Index %r already exists", collection_name) + except ResourceNotFoundError: + schema = _build_index_schema(collection_name, self.cfg.embed_dim) + self._index_client.create_index(schema) + logger.info("Created AI Search index %r (embed_dim=%d)", collection_name, self.cfg.embed_dim) + + # ------- jobs ------- + + def submit_job( + self, + file_paths: list[str], + collection_name: str, + config: dict[str, Any] | None = None, + ) -> str: + job_id = str(uuid.uuid4()) + config = config or {} + # AI-Q's HTTP layer passes original_filenames as a parallel LIST (one + # per entry of file_paths), NOT a dict[path, name]. Other callers may + # pass a dict — handle both shapes. + original_filenames = _resolve_filenames(file_paths, config.get("original_filenames")) + + # File-progress placeholders so the UI shows them as UPLOADING immediately. + file_progress = [ + FileProgress( + file_id=str(uuid.uuid4()), + file_name=original_filenames[i], + status=FileStatus.UPLOADING, + ) + for i in range(len(file_paths)) + ] + + with self._jobs_lock: + self._jobs[job_id] = IngestionJobStatus( + job_id=job_id, + status=JobState.PENDING, + collection_name=collection_name, + backend=_BACKEND_NAME, + submitted_at=datetime.now(UTC), + total_files=len(file_paths), + processed_files=0, + file_details=file_progress, + ) + + logger.info( + "submit_job: id=%s files=%d collection=%r", + job_id, + len(file_paths), + collection_name, + ) + + thread = threading.Thread( + target=self._process_job, + args=(job_id, file_paths, collection_name, config), + daemon=True, + name=f"aiq-azure-search-ingest-{job_id[:8]}", + ) + thread.start() + return job_id + + def get_job_status(self, job_id: str) -> IngestionJobStatus: + with self._jobs_lock: + existing = self._jobs.get(job_id) + if existing is None: + return IngestionJobStatus( + job_id=job_id, + status=JobState.FAILED, + collection_name="", + backend=_BACKEND_NAME, + submitted_at=datetime.now(UTC), + error_message="Job not found", + file_details=[], + ) + return existing.model_copy(deep=True) + + def _process_job( + self, + job_id: str, + file_paths: list[str], + collection_name: str, + config: dict[str, Any], + ) -> None: + """Background ingestion worker. Updates job state in place under lock.""" + cleanup = bool(config.get("cleanup_files", self.cfg.cleanup_files)) + original_filenames = _resolve_filenames(file_paths, config.get("original_filenames")) + + # Mark PROCESSING and started_at + self._update_job(job_id, status=JobState.PROCESSING, started_at=datetime.now(UTC)) + + try: + self._ensure_index(collection_name) + except Exception as e: # noqa: BLE001 + if cleanup: + for path in file_paths: + try: + os.unlink(path) + except OSError: + pass + self._fail_job(job_id, f"Failed to ensure index {collection_name!r}: {e!s}") + return + + failed = 0 + for idx, path in enumerate(file_paths): + file_name = original_filenames[idx] + self._update_file_progress(job_id, idx, status=FileStatus.INGESTING) + try: + chunks_written = self._process_file( + path=path, + collection_name=collection_name, + file_name=file_name, + ) + self._update_file_progress( + job_id, + idx, + status=FileStatus.SUCCESS, + progress_percent=100.0, + chunks_created=chunks_written, + ) + logger.info("Ingested %s: %d chunks → %s", file_name, chunks_written, collection_name) + except Exception as e: # noqa: BLE001 + msg = self._translate_error(e) + self._update_file_progress( + job_id, + idx, + status=FileStatus.FAILED, + error_message=msg, + ) + failed += 1 + logger.exception("Failed to ingest %s", file_name) + finally: + if cleanup: + try: + os.unlink(path) + except OSError: + pass + self._update_job(job_id, processed_files=idx + 1) + + # Final status + 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=datetime.now(UTC), + ) + + def _process_file(self, *, path: str, collection_name: str, file_name: str) -> int: + """Parse → chunk → embed → upload one file. Returns chunk count.""" + from llama_index.core import SimpleDirectoryReader + + # 1. Parse. SimpleDirectoryReader returns a Document per logical chunk + # (per-page for PDFs). Metadata includes page_label. + reader = SimpleDirectoryReader(input_files=[path]) + documents = reader.load_data() + if not documents: + raise ValueError(f"No content extracted from {file_name}") + + # 2. Chunk each Document with the SentenceSplitter, preserving metadata. + nodes = self.splitter.get_nodes_from_documents(documents) + if not nodes: + raise ValueError(f"Chunking produced 0 chunks for {file_name}") + + # 3. Embed all chunk texts in one batch (NVIDIAEmbedding handles + # rate limiting + batch sizing internally). + texts = [n.get_content() for n in nodes] + embeddings = self.embedding.get_text_embedding_batch(texts) + + # 4. Build AI Search documents. doc_id is per-file, chunk_id per-chunk. + doc_id = self._make_doc_id(file_name) + docs = [] + for i, (node, vector) in enumerate(zip(nodes, embeddings, strict=True)): + chunk_id = f"{doc_id}-c{i:04d}" + page = _coerce_page_number(node.metadata.get("page_label")) + docs.append( + { + "id": chunk_id, + "chunk": node.get_content(), + "embedding": list(vector), + "doc_id": doc_id, + "file_name": file_name, + "page_number": page, + "metadata": "{}", + } + ) + + # 5. Upload. + client = self._get_search_client(collection_name) + result = client.upload_documents(documents=docs) + # Surface partial-failure as an exception so the file gets marked FAILED. + failures = [r for r in result if not r.succeeded] + if failures: + raise RuntimeError(f"AI Search rejected {len(failures)}/{len(docs)} chunks: {failures[0].error_message}") + + # 6. Summarise + register. AI-Q's intent classifier reads registered + # summaries from the system prompt's `available_documents` block, so + # this step is what makes the orchestrator route knowledge queries + # to us instead of falling through to web search. + if self.cfg.generate_summary: + full_text = "\n".join(texts) + summary = self._generate_summary(full_text, file_name) + if summary: + try: + register_summary(collection_name, file_name, summary) + logger.info("Registered summary for %s in %s", file_name, collection_name) + except Exception: # noqa: BLE001 + logger.exception("register_summary failed for %s", file_name) + + return len(docs) + + def _generate_summary(self, text: str, file_name: str) -> str | None: + """Generate a one-sentence document summary via the LangChain summary LLM. + + Returns None if no summary LLM is configured or generation fails — both + are non-fatal; the file is still ingested and queryable, just without + an entry in the agent's available_documents block. + """ + if self._summary_llm is None: + return None + # Truncate to keep the summary prompt small. + snippet = text[: self.cfg.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 + + @staticmethod + def _make_doc_id(file_name: str) -> str: + """Stable doc_id from the file name. Lowercase + hash for collision-resistance.""" + digest = hashlib.sha256(file_name.encode("utf-8")).hexdigest()[:8] + slug = re.sub(r"[^a-z0-9]+", "-", file_name.lower()).strip("-")[:32] + return f"{slug}-{digest}" if slug else digest + + # ------- job-state helpers (must be called under _jobs_lock or use it) ------- + + def _update_job(self, job_id: str, **fields: Any) -> None: + with self._jobs_lock: + existing = self._jobs.get(job_id) + if existing is None: + return + self._jobs[job_id] = existing.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=datetime.now(UTC), + ) + + def _update_file_progress(self, job_id: str, idx: int, **fields: Any) -> None: + with self._jobs_lock: + existing = self._jobs.get(job_id) + if existing is None or idx >= len(existing.file_details): + return + new_details = list(existing.file_details) + new_details[idx] = existing.file_details[idx].model_copy(update=fields) + self._jobs[job_id] = existing.model_copy(update={"file_details": new_details}) + + # ------- error translation ------- + + @staticmethod + def _translate_error(exc: Exception) -> str: + """User-readable error string for FileProgress.error_message.""" + if isinstance(exc, ResourceNotFoundError): + return f"AI Search index not found: {exc!s}" + if isinstance(exc, ClientAuthenticationError): + return f"AI Search authentication failed: {exc!s}" + if isinstance(exc, ServiceRequestError): + return f"AI Search service unavailable: {exc!s}" + if isinstance(exc, HttpResponseError): + return f"AI Search request failed ({exc.status_code}): {exc.reason or exc!s}" + return str(exc) + + # ------- collections ------- + + def create_collection( + self, + name: str, + description: str | None = None, + metadata: dict[str, Any] | None = None, + ) -> CollectionInfo: + _validate_index_name(name) + try: + self._ensure_index(name) + except ResourceExistsError: + pass + # AI Search doesn't natively store description/metadata at the index + # level — we just round-trip the request. + return CollectionInfo( + name=name, + description=description, + file_count=0, + chunk_count=0, + backend=_BACKEND_NAME, + metadata=metadata or {}, + created_at=datetime.now(UTC), + updated_at=datetime.now(UTC), + ) + + def delete_collection(self, name: str) -> bool: + try: + self._index_client.delete_index(name) + self._search_client_cache.pop(name, None) + logger.info("Deleted AI Search index %r", name) + return True + except ResourceNotFoundError: + return False + + def list_collections(self) -> list[CollectionInfo]: + out: list[CollectionInfo] = [] + for idx in self._index_client.list_indexes(): + out.append( + CollectionInfo( + name=idx.name, + file_count=0, + chunk_count=0, + backend=_BACKEND_NAME, + metadata={}, + ) + ) + return out + + def get_collection(self, name: str) -> CollectionInfo | None: + try: + idx = self._index_client.get_index(name) + except ResourceNotFoundError: + return None + return CollectionInfo( + name=idx.name, + file_count=0, + chunk_count=0, + backend=_BACKEND_NAME, + metadata={}, + ) + + # ------- files (delegating to submit_job for upload) ------- + + def upload_file( + self, + file_path: str, + collection_name: str, + metadata: dict[str, Any] | None = None, + ) -> FileInfo: + # Synchronous-ish: kick off the job and return file metadata. + # The frontend uses submit_job for actual upload; this is a single-file alias. + cfg: dict[str, Any] = {"original_filenames": [Path(file_path).name]} + if metadata: + cfg["metadata"] = metadata + self.submit_job([file_path], collection_name, cfg) + return FileInfo( + file_id=str(uuid.uuid4()), + file_name=Path(file_path).name, + collection_name=collection_name, + status=FileStatus.UPLOADING, + file_size=Path(file_path).stat().st_size if Path(file_path).is_file() else 0, + chunk_count=0, + metadata=metadata or {}, + uploaded_at=datetime.now(UTC), + ) + + def delete_file(self, file_id: str, collection_name: str) -> bool: + # file_id is the doc_id we stamped at ingest time. Delete all chunks + # whose doc_id matches, and remove the corresponding summary so the + # agent's available_documents block stays consistent. + try: + client = self._get_search_client(collection_name) + # Look up file_name from any matching chunk so we can unregister the summary. + results = client.search( + search_text="*", + filter=f"doc_id eq '{file_id}'", + select=["id", "file_name"], + top=10000, + ) + hits = list(results) + if not hits: + return False + file_name = hits[0].get("file_name") if hits else None + client.delete_documents(documents=[{"id": r["id"]} for r in hits]) + if file_name: + try: + unregister_summary(collection_name, file_name) + except Exception: # noqa: BLE001 + logger.exception("unregister_summary failed for %s", file_name) + return True + except Exception: # noqa: BLE001 + logger.exception("delete_file failed") + return False + + def list_files(self, collection_name: str) -> list[FileInfo]: + """List one FileInfo per distinct doc_id in the index. + + AI-Q's frontend polls this after ingestion to refresh the Files panel + and to decide whether to expose `knowledge_search` as an enabled data + source. Returning [] would cause uploaded files to vanish from the UI. + """ + try: + client = self._get_search_client(collection_name) + except ResourceNotFoundError: + return [] + try: + results = client.search( + search_text="*", + select=["doc_id", "file_name"], + top=10000, + ) + except ResourceNotFoundError: + return [] + except Exception: # noqa: BLE001 + logger.exception("list_files failed for %r", collection_name) + return [] + + agg: dict[str, dict[str, Any]] = {} + for hit in results: + did = hit.get("doc_id") or "unknown" + fn = hit.get("file_name") or "unknown" + entry = agg.setdefault(did, {"file_name": fn, "count": 0}) + entry["count"] += 1 + + now = datetime.now(UTC) + return [ + FileInfo( + file_id=did, + file_name=info["file_name"], + collection_name=collection_name, + status=FileStatus.SUCCESS, + file_size=None, + chunk_count=info["count"], + metadata={}, + uploaded_at=now, + ingested_at=now, + ) + for did, info in agg.items() + ] + + def get_file_status(self, file_id: str, collection_name: str) -> FileInfo | None: + """Return the FileInfo for a single doc_id (file_id == doc_id) if it exists.""" + try: + client = self._get_search_client(collection_name) + results = client.search( + search_text="*", + filter=f"doc_id eq '{file_id}'", + select=["doc_id", "file_name"], + top=10000, + ) + hits = list(results) + except ResourceNotFoundError: + return None + except Exception: # noqa: BLE001 + logger.exception("get_file_status failed for %r in %r", file_id, collection_name) + return None + + if not hits: + return None + + now = datetime.now(UTC) + return FileInfo( + file_id=file_id, + file_name=hits[0].get("file_name") or "unknown", + collection_name=collection_name, + status=FileStatus.SUCCESS, + file_size=None, + chunk_count=len(hits), + metadata={}, + uploaded_at=now, + ingested_at=now, + ) + + # ------- optional ------- + + async def health_check(self) -> bool: + try: + list(self._index_client.list_index_names()) + return True + except Exception: # noqa: BLE001 + logger.exception("Ingestor health_check failed") + return False + + +def _resolve_filenames(file_paths: list[str], raw: Any) -> list[str]: + """Normalise the `original_filenames` config field across caller shapes. + + AI-Q's HTTP route passes a parallel `list[str]` aligned to `file_paths`. + The SDK reference's example uses a `dict[str, str]` keyed by temp path. + Anything else (None, empty, mismatched length) falls back to the basename + of each file_path. + """ + if isinstance(raw, dict): + return [raw.get(p) or Path(p).name for p in file_paths] + if isinstance(raw, list): + names: list[str] = [] + for i, p in enumerate(file_paths): + entry = raw[i] if i < len(raw) else None + names.append(entry or Path(p).name) + return names + return [Path(p).name for p in file_paths] + + +def _coerce_page_number(page_label: Any) -> int | None: + """LlamaIndex sets page_label as a 1-indexed string for PDFs. Be lenient.""" + 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..4e121f17c 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 @@ -39,7 +40,7 @@ # 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"] @@ -250,6 +251,41 @@ class KnowledgeRetrievalConfig(FunctionBaseConfig, name="knowledge_retrieval"): 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=None, + description="Azure AI Search service URL (azure_ai_search only)", + ) + azure_search_api_key: SecretStr | None = Field( + default=None, + description="Optional Azure AI Search admin key; managed identity is used when omitted", + ) + azure_search_auth_mode: Literal["managed_identity", "api_key"] = Field( + default="managed_identity", + description="Authentication mode for Azure AI Search", + ) + embed_endpoint: HttpUrl = Field( + default=HttpUrl("https://integrate.api.nvidia.com/v1"), + description="OpenAI-compatible embedding endpoint (azure_ai_search only)", + ) + embed_dim: int = Field( + default=4096, + gt=0, + description="Embedding dimensions; must match existing Azure AI Search indexes", + ) + embed_api_key: SecretStr | None = Field( + default=None, + description="Optional embedding API key; NVIDIA_API_KEY is used when omitted", + ) + use_hybrid: bool = Field(default=True, description="Combine vector and keyword search (azure_ai_search only)") + use_semantic_ranker: bool = Field(default=True, description="Enable Azure semantic ranking (azure_ai_search only)") + chunk_size: int = Field(default=512, gt=0, description="Tokens per chunk (azure_ai_search only)") + chunk_overlap: int = Field(default=64, ge=0, description="Token overlap between chunks (azure_ai_search only)") + summary_max_chars: int = Field( + default=1000, + gt=0, + description="Maximum document characters sent to summary model (azure_ai_search only)", + ) @model_validator(mode="after") def validate_backend_config(self): @@ -298,6 +334,13 @@ 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") + if self.azure_search_auth_mode == "api_key" and self.azure_search_api_key is None: + raise ValueError("azure_search_auth_mode=api_key requires azure_search_api_key") + if self.chunk_overlap >= self.chunk_size: + raise ValueError("chunk_overlap must be smaller than chunk_size") return self @@ -381,11 +424,30 @@ def _setup_backend(config: KnowledgeRetrievalConfig, summary_llm_obj=None) -> tu "dask_file_transfer": config.opensearch_dask_file_transfer, "embed_model": config.embed_model, "embed_base_url": config.embed_base_url, + + 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, + "auth_mode": config.azure_search_auth_mode, + "embed_endpoint": str(config.embed_endpoint), + "embed_model": config.embed_model, + "embed_dim": config.embed_dim, + "embed_api_key": config.embed_api_key, + "use_hybrid": config.use_hybrid, + "use_semantic_ranker": config.use_semantic_ranker, + "chunk_size": config.chunk_size, + "chunk_overlap": config.chunk_overlap, + "summary_max_chars": config.summary_max_chars, + "collection_name": config.collection_name, + "cleanup_files": True, **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 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..e91be1d75 --- /dev/null +++ b/tests/knowledge_layer_tests/test_azure_ai_search.py @@ -0,0 +1,182 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from types import SimpleNamespace +from uuid import UUID + +import pytest +from azure.core.exceptions import ServiceRequestError +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 _resolve_filenames +from knowledge_layer.azure_ai_search.adapter import _validate_index_name +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 + + +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(): + with pytest.raises(ValueError, match="azure_search_endpoint"): + KnowledgeRetrievalConfig(backend="azure_ai_search") + + +def test_api_key_auth_requires_secret(): + with pytest.raises(ValueError, match="azure_search_api_key"): + KnowledgeRetrievalConfig( + backend="azure_ai_search", + azure_search_endpoint="https://example.search.windows.net", + azure_search_auth_mode="api_key", + ) + + +def test_setup_backend_preserves_secret_types(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_auth_mode="api_key", + azure_search_api_key="test-search-key", + embed_api_key="test-embed-key", + ) + + backend, backend_config = _setup_backend(config) + + assert backend == "azure_ai_search" + assert isinstance(backend_config["api_key"], SecretStr) + assert isinstance(backend_config["embed_api_key"], SecretStr) + assert backend_config["embed_dim"] == 4096 + + +@pytest.mark.parametrize( + ("hit", "expected"), + [ + ({"@search.reranker_score": 2.0, "@search.score": 0.9}, 0.5), + ({"@search.reranker_score": 8.0}, 1.0), + ({"@search.reranker_score": -1.0}, 0.0), + ({"@search.score": 0.75}, 0.75), + ], +) +def test_normalize_clamps_and_prefers_reranker_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_populates_citation_and_metadata(): + retriever = AzureAISearchRetriever.__new__(AzureAISearchRetriever) + cited = retriever.normalize( + { + "id": "chunk-1", + "chunk": "text", + "file_name": "report.pdf", + "page_number": "3", + "doc_id": "doc-1", + "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 == {"doc_id": "doc-1", "raw": '{"section":"intro"}'} + assert fallback.content == "" + assert fallback.file_name == "unknown" + assert fallback.page_number is None + assert fallback.display_citation == "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} + ) + result = RetrievalResult(query="query", backend="azure_ai_search", chunks=[chunk]) + + formatted = _format_results(result, "query") + + assert "Source: report.pdf" in formatted + assert "Citation: report.pdf, p.3" in formatted + + +@pytest.mark.asyncio +async def test_retrieve_builds_hybrid_semantic_search_request(): + class FakeEmbedding: + def get_query_embedding(self, query): + assert query == "hello" + return [0.1, 0.2] + + 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.cfg = SimpleNamespace(use_hybrid=True, use_semantic_ranker=True) + retriever._embedding = FakeEmbedding() + retriever._get_client = lambda collection_name: client + + result = await retriever.retrieve("hello", "session-1", top_k=5, filters={"$filter": "doc_id eq 'doc-1'"}) + + assert result.success + assert [chunk.content for chunk in result.chunks] == ["answer"] + assert client.kwargs["search_text"] == "hello" + assert client.kwargs["query_type"] == "semantic" + assert client.kwargs["semantic_configuration_name"] == "default-semantic" + assert client.kwargs["filter"] == "doc_id eq 'doc-1'" + assert client.kwargs["top"] == 5 + assert client.kwargs["select"] == ["id", "chunk", "doc_id", "file_name", "page_number", "metadata"] + vector_query = client.kwargs["vector_queries"][0] + assert vector_query.vector == [0.1, 0.2] + assert vector_query.k_nearest_neighbors == 20 + assert vector_query.fields == "embedding" + + +def test_index_schema_uses_requested_vector_dimensions(): + schema = _build_index_schema("session-1", 4096) + embedding = next(field for field in schema.fields if field.name == "embedding") + + assert embedding.vector_search_dimensions == 4096 + assert schema.semantic_search.default_configuration_name == "default-semantic" + + +def test_frontend_collection_names_are_valid(): + _validate_index_name("s_f031d9cf_123") + + with pytest.raises(ValueError, match="Invalid AI Search index name"): + _validate_index_name("Invalid Name") + + +def test_filename_and_page_normalization(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 + + +def test_service_connection_error_is_user_readable(): + message = AzureAISearchIngestor._translate_error(ServiceRequestError("connection refused")) + + assert message == "AI Search service unavailable: connection refused" diff --git a/uv.lock b/uv.lock index 5a0264046..b40d34564 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 = [ @@ -591,6 +591,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" @@ -1368,6 +1421,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" @@ -2207,6 +2269,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" @@ -2457,12 +2528,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" }, @@ -2474,6 +2549,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" }, @@ -2503,16 +2587,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.5.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" }, @@ -2520,6 +2610,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" }, @@ -2528,7 +2619,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" @@ -3120,6 +3211,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" @@ -3492,6 +3600,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" @@ -5948,23 +6082,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 = [ @@ -5980,23 +6114,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 = [ @@ -6233,6 +6367,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" From 8b1bd282fc3e1f5c3f1f564b427618a972d21876 Mon Sep 17 00:00:00 2001 From: Harmke Alkemade Date: Fri, 3 Jul 2026 13:37:33 +0200 Subject: [PATCH 02/27] fix(knowledge): harden Azure AI Search configuration and ingestion Signed-off-by: Harmke Alkemade --- .../customization/configuration-reference.md | 33 +- docs/source/customization/knowledge-layer.md | 36 +- docs/source/examples/azure-ai-search.md | 34 +- docs/source/index.md | 1 + .../knowledge_layer/KNOWLEDGE-LAYER-SETUP.md | 9 + .../src/azure_ai_search/README.md | 60 +- .../src/azure_ai_search/adapter.py | 1442 ++++++++++------- sources/knowledge_layer/src/register.py | 42 +- .../run_adapter_compliance.py | 19 +- .../test_azure_ai_search.py | 632 +++++++- 10 files changed, 1597 insertions(+), 711 deletions(-) diff --git a/docs/source/customization/configuration-reference.md b/docs/source/customization/configuration-reference.md index 667b62be7..27f23170c 100644 --- a/docs/source/customization/configuration-reference.md +++ b/docs/source/customization/configuration-reference.md @@ -217,7 +217,8 @@ functions: ### `knowledge_retrieval` -Semantic search over ingested documents. Supports two backends: LlamaIndex (local ChromaDB) and Foundational RAG (hosted NVIDIA RAG Blueprint). +Semantic search over ingested documents. Supports LlamaIndex (local ChromaDB), Foundational RAG (hosted +NVIDIA RAG Blueprint), and Azure AI Search. ```yaml functions: @@ -247,9 +248,24 @@ functions: # 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} + use_hybrid: true + use_semantic_ranker: true +``` + +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`. + | Parameter | Type | Default | Description | |-----------|------|---------|-------------| -| `backend` | `str` | `llamaindex` | Backend type: `llamaindex` or `foundational_rag`. | +| `backend` | `str` | `llamaindex` | Backend type: `llamaindex`, `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. | @@ -260,6 +276,19 @@ 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_auth_mode` | `str` | Auto | Uses `api_key` when `AZURE_SEARCH_API_KEY` is set; otherwise `managed_identity`. | +| `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` | Namespace prefix for AI-Q-owned indexes. | +| `embed_endpoint` | `URL` | `AIQ_EMBED_BASE_URL` or NVIDIA API | OpenAI-compatible embedding endpoint. | +| `embed_model` | `str` | `AIQ_EMBED_MODEL` or `nvidia/nv-embed-v1` | Embedding model used for Azure ingestion and retrieval. | +| `embed_dim` | `int` | `AIQ_EMBED_DIM` or `4096` | Embedding dimensions; must match the model and existing index schema. | +| `embed_api_key` | `SecretStr` | `None` | Optional embedding API key; falls back to `NVIDIA_API_KEY`. | +| `use_hybrid` | `bool` | `true` | Combine lexical and vector retrieval. | +| `use_semantic_ranker` | `bool` | `true` | Apply Azure semantic ranking; requires `use_hybrid: true`. | +| `chunk_size` | `int` | `512` | Tokens per Azure-ingested chunk. | +| `chunk_overlap` | `int` | `64` | Token overlap; must be smaller than `chunk_size`. | +| `summary_max_chars` | `int` | `1000` | Maximum document characters sent to the summary model. | ### `intent_classifier` diff --git a/docs/source/customization/knowledge-layer.md b/docs/source/customization/knowledge-layer.md index 1e07ebf61..e9475524e 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 | Hybrid and semantic 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. @@ -52,6 +53,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. --- @@ -70,6 +73,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 ``` > **New to Knowledge Layer?** Start with `llamaindex` - it requires no external services and works on macOS and Linux. @@ -152,6 +156,31 @@ functions: timeout: 120 ``` +**Azure AI Search (Managed Service)** +```yaml +functions: + knowledge_search: + _type: knowledge_retrieval + backend: azure_ai_search + collection_name: my_docs + use_hybrid: true + use_semantic_ranker: true +``` + +Set `AZURE_SEARCH_ENDPOINT` and `NVIDIA_API_KEY` in the environment. Setting +`AZURE_SEARCH_API_KEY` selects key authentication; otherwise Azure +`DefaultAzureCredential` is used. 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. + +Azure maps logical collection names to collision-safe physical indexes under `azure_search_index_prefix`. Only +indexes containing an AI-Q ownership and schema marker are listed, queried, or deleted. Legacy indexes named directly +after collections are ignored; re-ingest those collections after enabling this backend. + +Upload responses return canonical UUID file IDs. Re-uploading the same filename writes and verifies the new generation +before removing the old generation. 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. + #### Multimodal Extraction (LlamaIndex Only) By default, LlamaIndex ingests text only and uses the NVIDIA hosted embedding models. When `AIQ_EXTRACT_IMAGES` or `AIQ_EXTRACT_CHARTS` is enabled, a Vision Language Model (VLM) is used during ingestion to caption embedded images and extract structured data from charts (axis labels, data points, chart type). This makes visual content in PDFs searchable and retrievable alongside text. The VLM is only invoked at ingestion time, not at query time. @@ -212,10 +241,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) | +| **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 both backends). Types like HTML, JSON, CSV, and images are supported by the backends but the frontend upload flow does not handle them yet -- this is a separate task. +> **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: @@ -396,6 +428,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) | | `COLLECTION_NAME` | All | Default collection name | diff --git a/docs/source/examples/azure-ai-search.md b/docs/source/examples/azure-ai-search.md index c4bc37d4b..510721df7 100644 --- a/docs/source/examples/azure-ai-search.md +++ b/docs/source/examples/azure-ai-search.md @@ -16,6 +16,15 @@ Install the backend dependency: uv pip install -e "sources/knowledge_layer[azure_ai_search]" ``` +Set the environment used by the adapter: + +```bash +export AZURE_SEARCH_ENDPOINT=https://.search.windows.net +export NVIDIA_API_KEY= +# Optional; omit to use DefaultAzureCredential. +export AZURE_SEARCH_API_KEY= +``` + Replace the `knowledge_search` block in a web configuration such as `configs/config_web_default_llamaindex.yml`: @@ -26,13 +35,6 @@ functions: backend: azure_ai_search collection_name: ${COLLECTION_NAME:-aiq_default} top_k: 5 - - azure_search_endpoint: https://.search.windows.net - azure_search_auth_mode: managed_identity - - embed_endpoint: https://integrate.api.nvidia.com/v1 - embed_model: nvidia/nv-embed-v1 - embed_dim: 4096 use_hybrid: true use_semantic_ranker: true @@ -41,13 +43,21 @@ functions: summary_db: ${AIQ_SUMMARY_DB:-sqlite+aiosqlite:///./summaries.db} ``` -For API-key authentication, set `azure_search_auth_mode: api_key` and add -`azure_search_api_key: ${AZURE_SEARCH_API_KEY}`. Managed identity uses -`DefaultAzureCredential`; set `AZURE_CLIENT_ID` to select a user-assigned -identity. If `embed_api_key` is omitted, the NVIDIA embedding client reads -`NVIDIA_API_KEY`. +Explicit YAML options override environment defaults. `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`. Existing indexes must use the configured `embed_dim`. Delete and re-ingest a collection when changing embedding dimensions. 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 only lists or mutates indexes carrying its AI-Q ownership marker. +Logical collection names map to collision-safe physical names under +`azure_search_index_prefix`; un-namespaced indexes from earlier versions are +ignored and must be re-ingested. File IDs returned by upload are authoritative +for status and delete operations. Same-name uploads replace the prior file only +after the new generation has been fully indexed. diff --git a/docs/source/index.md b/docs/source/index.md index b57ff8fec..e6a534626 100644 --- a/docs/source/index.md +++ b/docs/source/index.md @@ -140,6 +140,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 350d169d9..9ea42047b 100644 --- a/sources/knowledge_layer/KNOWLEDGE-LAYER-SETUP.md +++ b/sources/knowledge_layer/KNOWLEDGE-LAYER-SETUP.md @@ -44,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 AI-Q-owned, collision-safe indexes in a managed Azure AI Search service. Canonical + UUID file IDs support status and deletion, while same-name uploads replace the prior generation after successful + indexing. 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, @@ -1040,6 +1043,12 @@ 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` | +| `AIQ_AZURE_SEARCH_INDEX_PREFIX` | azure_ai_search | Prefix for AI-Q-owned indexes (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: `4096`) | | `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/src/azure_ai_search/README.md b/sources/knowledge_layer/src/azure_ai_search/README.md index dda36060c..126b45b0a 100644 --- a/sources/knowledge_layer/src/azure_ai_search/README.md +++ b/sources/knowledge_layer/src/azure_ai_search/README.md @@ -17,25 +17,21 @@ 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} - - azure_search_endpoint: https://.search.windows.net - azure_search_auth_mode: managed_identity - # For key authentication instead: - # azure_search_auth_mode: api_key - # azure_search_api_key: ${AZURE_SEARCH_API_KEY} - - embed_endpoint: https://integrate.api.nvidia.com/v1 - embed_model: nvidia/nv-embed-v1 - embed_dim: 4096 - # NVIDIAEmbedding reads NVIDIA_API_KEY when this field is omitted. - # embed_api_key: ${NVIDIA_API_KEY} - use_hybrid: true use_semantic_ranker: true top_k: 5 @@ -47,20 +43,42 @@ functions: summary_db: ${AIQ_SUMMARY_DB:-sqlite+aiosqlite:///./summaries.db} ``` -Managed identity uses `DefaultAzureCredential`. Set `AZURE_CLIENT_ID` when a -user-assigned identity should be selected. The identity needs permission to -create and delete indexes and to read, write, and delete index documents. +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`. + +When `AZURE_SEARCH_API_KEY` is absent, the adapter uses +`DefaultAzureCredential`. Set `AZURE_CLIENT_ID` when a user-assigned identity +should be selected. The identity needs permission to create and delete indexes +and to read, write, and delete index documents. The adapter parses PDF, DOCX, TXT, and Markdown uploads with LlamaIndex, -creates one Azure AI Search index per AI-Q collection, and performs vector or -hybrid retrieval with optional semantic ranking. AI-Q frontend conversations -use their conversation ID as the collection name, keeping uploads and -WebSocket retrieval in the same index. +creates one namespaced Azure AI Search index per AI-Q collection, and performs +vector or hybrid retrieval with optional semantic ranking. Logical collection +names are sanitized and combined with a stable hash, preventing collisions and +protecting unrelated indexes in a shared service. Only indexes containing the +AI-Q ownership/schema marker are visible or mutable through this backend. + +Upload responses return canonical UUID file IDs used by job progress, list, +status, and delete operations. Re-uploading the same filename writes and +verifies a new generation before deleting old chunks. 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 any existing index. Changing from a 2048-dimensional model to `nvidia/nv-embed-v1` at 4096 dimensions requires deleting and re-ingesting the old collection. The adapter -does not alter an existing index schema. +validates ownership, fields, vector profile, dimensions, and semantic +configuration before use; it does not alter an incompatible schema. Indexes +created by the earlier un-namespaced implementation are deliberately ignored, +so re-ingest those collections. 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 diff --git a/sources/knowledge_layer/src/azure_ai_search/adapter.py b/sources/knowledge_layer/src/azure_ai_search/adapter.py index 0fba51c04..83c3ff6a8 100644 --- a/sources/knowledge_layer/src/azure_ai_search/adapter.py +++ b/sources/knowledge_layer/src/azure_ai_search/adapter.py @@ -4,22 +4,24 @@ from __future__ import annotations import asyncio -import hashlib +import json import logging import os import re import threading 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 import MatchConditions from azure.core.credentials import AzureKeyCredential from azure.core.exceptions import ClientAuthenticationError from azure.core.exceptions import HttpResponseError -from azure.core.exceptions import ResourceExistsError +from azure.core.exceptions import ResourceModifiedError from azure.core.exceptions import ResourceNotFoundError from azure.core.exceptions import ServiceRequestError from azure.identity import DefaultAzureCredential @@ -49,42 +51,81 @@ 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" _SEMANTIC_CONFIG = "default-semantic" - - -def _coerce_config( - config: dict[str, Any] | None, -) -> SimpleNamespace: - """Expose the validated KnowledgeRetrievalConfig values as attributes.""" - if not config: - raise ValueError("Azure AI Search configuration is required") - return SimpleNamespace(**config) +_SCHEMA_VERSION = 1 +_MARKER_PREFIX = "aiq.azure_ai_search:" +_MAX_INDEX_NAME_LENGTH = 128 +_MAX_BATCH_ACTIONS = 1000 +_MAX_BATCH_BYTES = 16 * 1024 * 1024 +_PAGE_SIZE = 1000 +_DELETE_ATTEMPTS = 3 + +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.""" + provided = config or {} + values: dict[str, Any] = { + "endpoint": os.environ.get("AZURE_SEARCH_ENDPOINT"), + "api_key": os.environ.get("AZURE_SEARCH_API_KEY"), + "embed_endpoint": os.environ.get("AIQ_EMBED_BASE_URL", "https://integrate.api.nvidia.com/v1"), + "embed_model": os.environ.get("AIQ_EMBED_MODEL", "nvidia/nv-embed-v1"), + "embed_dim": int(os.environ.get("AIQ_EMBED_DIM", "4096")), + "embed_api_key": None, + "use_hybrid": True, + "use_semantic_ranker": True, + "chunk_size": 512, + "chunk_overlap": 64, + "summary_max_chars": 1000, + "collection_name": "default", + "cleanup_files": True, + "generate_summary": False, + "summary_llm": None, + "index_prefix": os.environ.get("AIQ_AZURE_SEARCH_INDEX_PREFIX", "aiq"), + "start_ttl_cleanup": True, + } + values.update(provided) + values["auth_mode"] = provided.get( + "auth_mode", + "api_key" if values["api_key"] else "managed_identity", + ) + if not values["endpoint"]: + raise ValueError("Azure AI Search configuration requires `endpoint`") + if values["auth_mode"] not in {"managed_identity", "api_key"}: + raise ValueError("Azure AI Search auth_mode must be 'managed_identity' or 'api_key'") + if values["chunk_overlap"] >= values["chunk_size"]: + raise ValueError("chunk_overlap must be smaller than chunk_size") + if not values["use_hybrid"] and values["use_semantic_ranker"]: + raise ValueError("use_semantic_ranker=true requires use_hybrid=true") + return SimpleNamespace(**values) def _build_search_credential(cfg: SimpleNamespace): - """Pick the right Azure SDK credential based on auth_mode.""" + """Pick the Azure SDK credential without exposing secret values.""" if cfg.auth_mode == "api_key": api_key = _secret_value(cfg.api_key) if api_key is None: raise ValueError("auth_mode=api_key requires the `api_key` field to be set") return AzureKeyCredential(api_key) - # DefaultAzureCredential honors AZURE_CLIENT_ID for user-assigned identities. return DefaultAzureCredential() def _secret_value(value: Any) -> str | None: - """Unwrap SecretStr values only when an SDK client needs them.""" if value is None: return None if hasattr(value, "get_secret_value"): @@ -92,46 +133,284 @@ def _secret_value(value: Any) -> str | None: return str(value) -# ============================================================================= -# Retriever -# ============================================================================= +def _utc_now() -> datetime: + return datetime.now(UTC) -@register_retriever(_BACKEND_NAME) -class AzureAISearchRetriever(BaseRetriever): - """Hybrid + semantic-ranked retriever backed by Azure AI Search. +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 - Per AI Search index, the schema has at minimum: `id` (key), `chunk` (text, - searchable), `embedding` (Collection(Edm.Single), vectorSearchProfile), - plus `doc_id`, `file_name`, `page_number`, and `metadata`. - """ - backend_name = _BACKEND_NAME +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 __init__(self, config: dict[str, Any] | None = None): - super().__init__(config) - cfg = _coerce_config(self.config) - self.cfg = cfg - # Lazy embedding-client init — `llama_index.embeddings.nvidia.NVIDIAEmbedding` - # is imported at first use to keep the package installable in environments - # where llama-index isn't present (unit tests). - self._embedding = None +def _index_name_for_collection(collection_name: str, prefix: str = "aiq") -> str: + """Map a logical collection to an official, collision-safe Azure index name.""" + prefix_part = _sanitize_index_part(prefix, "aiq")[:48].rstrip("-") or "aiq" + collection_part = _sanitize_index_part(collection_name) + suffix = uuid.uuid5(uuid.NAMESPACE_URL, f"{prefix}\0{collection_name}").hex[:12] + available = _MAX_INDEX_NAME_LENGTH - len(prefix_part) - len(suffix) - 2 + collection_part = collection_part[:available].rstrip("-") or "default" + return f"{prefix_part}-{collection_part}-{suffix}" + + +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: + return _MARKER_PREFIX + json.dumps(marker, separators=(",", ":"), sort_keys=True) + + +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( + collection_name: str, + cfg: SimpleNamespace, + description: str | None = None, + metadata: dict[str, Any] | None = None, +) -> dict[str, Any]: + now = _utc_now().isoformat() + marker = { + "backend": _BACKEND_NAME, + "schema_version": _SCHEMA_VERSION, + "collection_name": collection_name, + "description": description, + "metadata": metadata or {}, + "embedding_model": cfg.embed_model, + "embedding_dim": cfg.embed_dim, + "created_at": now, + "updated_at": now, + } + _encode_marker(marker) # Validate JSON serializability before any service mutation. + return marker + + +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 _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, + ), + 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="page_number", type=SearchFieldDataType.Int32, filterable=True, sortable=True), + SimpleField(name="chunk_index", type=SearchFieldDataType.Int32, filterable=True, sortable=True), + SimpleField(name="file_size", type=SearchFieldDataType.Int64, filterable=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")], + ), + semantic_search=SemanticSearch( + default_configuration_name=_SEMANTIC_CONFIG, + configurations=[ + SemanticConfiguration( + name=_SEMANTIC_CONFIG, + prioritized_fields=SemanticPrioritizedFields( + title_field=SemanticField(field_name="file_name"), + content_fields=[SemanticField(field_name="chunk")], + ), + ), + ], + ), + ) + - # SearchClient is bound to a single index, so we cache one per collection. - self._credential = _build_search_credential(cfg) - self._client_cache: dict[str, SearchClient] = {} - - logger.info( - "AzureAISearchRetriever initialized: endpoint=%s embed_model=%s embed_endpoint=%s " - "use_hybrid=%s use_semantic_ranker=%s", - cfg.endpoint, - cfg.embed_model, - cfg.embed_endpoint, - cfg.use_hybrid, - cfg.use_semantic_ranker, +def _validate_index_schema(index: SearchIndex, collection_name: str, 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("collection_name") != collection_name: + raise RuntimeError( + f"Azure AI Search index {index.name!r} belongs to collection {marker.get('collection_name')!r}, " + f"not {collection_name!r}" + ) + if 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} embedding configuration does not match " + f"{cfg.embed_model!r}/{cfg.embed_dim}" ) + fields = {field.name: field for field in index.fields} + required_types = { + "id": SearchFieldDataType.String, + "chunk": SearchFieldDataType.String, + "embedding": SearchFieldDataType.Collection(SearchFieldDataType.Single), + "file_id": SearchFieldDataType.String, + "file_name": SearchFieldDataType.String, + "page_number": SearchFieldDataType.Int32, + "chunk_index": SearchFieldDataType.Int32, + "file_size": SearchFieldDataType.Int64, + "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["file_id"].filterable or not fields["file_name"].filterable: + raise RuntimeError(f"Azure AI Search index {index.name!r} requires filterable file 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") + semantic_names = { + semantic.name for semantic in (index.semantic_search.configurations if index.semantic_search else []) + } + if cfg.use_semantic_ranker and ( + index.semantic_search is None + or index.semantic_search.default_configuration_name != _SEMANTIC_CONFIG + or _SEMANTIC_CONFIG not in semantic_names + ): + raise RuntimeError(f"Azure AI Search index {index.name!r} semantic configuration does not match") + return marker + + +class _AzureIndexMixin: + cfg: SimpleNamespace + _credential: Any + _embedding: Any + _index_client: SearchIndexClient + _search_clients: dict[str, SearchClient] + + 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_clients = {} + self._embedding = None + @property def embedding(self): if self._embedding is None: @@ -144,16 +423,37 @@ def embedding(self): ) return self._embedding - def _get_client(self, collection_name: str) -> SearchClient: - if collection_name not in self._client_cache: - self._client_cache[collection_name] = SearchClient( - endpoint=str(self.cfg.endpoint), - index_name=collection_name, - credential=self._credential, + def _physical_index_name(self, collection_name: str) -> str: + name = _index_name_for_collection(collection_name, self.cfg.index_prefix) + _validate_index_name(name) + return name + + def _get_search_client(self, collection_name: str) -> SearchClient: + if collection_name not in self._search_clients: + self._search_clients[collection_name] = self._index_client.get_search_client( + self._physical_index_name(collection_name) ) - return self._client_cache[collection_name] + return self._search_clients[collection_name] + + +@register_retriever(_BACKEND_NAME) +class AzureAISearchRetriever(_AzureIndexMixin, BaseRetriever): + """Hybrid and semantic-ranked retriever backed by owned Azure indexes.""" + + backend_name = _BACKEND_NAME + + def __init__(self, config: dict[str, Any] | None = None): + super().__init__(config) + self._initialize_azure(self.config) + self._validated: set[str] = set() - # ------- abstract methods ------- + def _get_client(self, collection_name: str) -> SearchClient: + index_name = self._physical_index_name(collection_name) + if collection_name not in self._validated: + index = self._index_client.get_index(index_name) + _validate_index_schema(index, collection_name, self.cfg) + self._validated.add(collection_name) + return self._get_search_client(collection_name) async def retrieve( self, @@ -162,7 +462,6 @@ async def retrieve( top_k: int = 10, filters: dict[str, Any] | None = None, ) -> RetrievalResult: - """Run the synchronous Azure Search client without blocking the event loop.""" return await asyncio.to_thread(self._retrieve_sync, query, collection_name, top_k, filters) def _retrieve_sync( @@ -172,300 +471,144 @@ def _retrieve_sync( top_k: int, filters: dict[str, Any] | None, ) -> RetrievalResult: - """Hybrid + (optionally) semantic-ranked search. Never raises.""" try: - # 1. Embed the query query_vector = self.embedding.get_query_embedding(query) - - # 2. Build the search request client = self._get_client(collection_name) vector_query = VectorizedQuery( vector=query_vector, - # Over-fetch on the vector side so the reranker has more - # candidates to choose from — semantic ranking only re-orders - # the input set, it doesn't reach back into the index. k_nearest_neighbors=max(top_k * 3, 20), fields="embedding", ) - search_params: dict[str, Any] = { "vector_queries": [vector_query], "top": top_k, - "select": ["id", "chunk", "doc_id", "file_name", "page_number", "metadata"], + "select": ["id", "chunk", "file_id", "file_name", "page_number", "metadata"], } if self.cfg.use_hybrid: - # Hybrid = include lexical text alongside the vector query. search_params["search_text"] = query if self.cfg.use_semantic_ranker: search_params["query_type"] = "semantic" search_params["semantic_configuration_name"] = _SEMANTIC_CONFIG - if filters and isinstance(filters, dict): - # Pass through a pre-built OData filter string under "$filter". - if odata := filters.get("$filter"): - search_params["filter"] = odata - - # 3. Run the search - raw_results = client.search(**search_params) - chunks = [self.normalize(hit) for hit in raw_results] - - logger.info( - "retrieve OK: collection=%s query=%r returned=%d top_k=%d", - collection_name, - query[:80], - len(chunks), - top_k, - ) - return RetrievalResult( - query=query, - backend=_BACKEND_NAME, - chunks=chunks, - success=True, - ) + if filters and isinstance(filters, dict) and (odata_filter := filters.get("$filter")): + search_params["filter"] = odata_filter + chunks = [self.normalize(hit) for hit in client.search(**search_params)] + return RetrievalResult(query=query, backend=_BACKEND_NAME, chunks=chunks, success=True) except ResourceNotFoundError: - msg = f"AI Search index {collection_name!r} not found. Has the collection been created yet?" - logger.warning(msg) - return RetrievalResult(query=query, backend=_BACKEND_NAME, chunks=[], success=False, error_message=msg) - except ClientAuthenticationError as e: - msg = f"AI Search authentication failed: {e!s}" - logger.error(msg) - return RetrievalResult(query=query, backend=_BACKEND_NAME, chunks=[], success=False, error_message=msg) - except ServiceRequestError as e: - msg = f"AI Search service unavailable: {e!s}" - logger.error(msg) - return RetrievalResult(query=query, backend=_BACKEND_NAME, chunks=[], success=False, error_message=msg) - except HttpResponseError as e: - msg = f"AI Search request failed: {e.status_code} {e.reason or e!s}" - logger.error(msg, exc_info=True) - return RetrievalResult(query=query, backend=_BACKEND_NAME, chunks=[], success=False, error_message=msg) - except Exception as e: # noqa: BLE001 — surface any unexpected error as a UI-readable string - msg = f"Unexpected error during retrieval: {e!s}" - logger.exception("retrieve unexpected error") - return RetrievalResult(query=query, backend=_BACKEND_NAME, chunks=[], success=False, error_message=msg) + 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: - """Map one AI Search hit dict to a 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")) - doc_id = raw_result.get("doc_id") - - # Score: prefer the semantic reranker score (0–4 typical range, normalise - # to 0–1) when available, otherwise the search score (already 0+ from - # hybrid RRF; clamp to 1). rerank_score = raw_result.get("@search.reranker_score") search_score = raw_result.get("@search.score") or 0.0 - if rerank_score is not None: - score = min(max(float(rerank_score) / 4.0, 0.0), 1.0) - else: - score = min(max(float(search_score), 0.0), 1.0) - - if page_number is not None: - display_citation = f"{file_name}, p.{page_number}" - else: - display_citation = str(file_name) - - metadata: dict[str, Any] = {} - if doc_id: - metadata["doc_id"] = doc_id - # Pass through any raw metadata blob if present - if raw_meta := raw_result.get("metadata"): - metadata["raw"] = raw_meta - + score = float(rerank_score) / 4.0 if rerank_score is not None else 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=display_citation, + display_citation=f"{file_name}, p.{page_number}" if page_number else str(file_name), content_type=ContentType.TEXT, metadata=metadata, ) - # ------- optional ------- - async def health_check(self) -> bool: - """Return True if the configured collection's index is reachable.""" - try: - client = self._get_client(self.cfg.collection_name) - client.get_document_count() + 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("Retriever health_check failed") + logger.exception("Azure AI Search retriever health_check failed") return False -# ============================================================================= -# Index schema builder -# ============================================================================= - - -# AI Search index names: lowercase letters, digits, hyphens, underscores; -# 2-128 chars; must start and end with a letter or digit; no consecutive dashes. -# (The official docs say "letters, numbers, dashes" but the service actually -# accepts underscores too, and the AI-Q frontend generates session-style -# collection names like "s_f031d9cf_...".) -_INDEX_NAME_RE = re.compile(r"^[a-z0-9](?:[a-z0-9_-]{0,126}[a-z0-9])?$") - - -def _validate_index_name(name: str) -> None: - if not _INDEX_NAME_RE.match(name) or "--" in name: - raise ValueError( - f"Invalid AI Search index name {name!r}. Names must be lowercase, " - "use only letters/digits/hyphens/underscores, be 2-128 chars, " - "start and end with a letter or digit, and not contain '--'." - ) - - -def _build_index_schema(name: str, embed_dim: int) -> SearchIndex: - """Build the canonical Azure AI Search index schema for a collection. - - Anything that queries the index must agree on field names and dimensions. - """ - return SearchIndex( - name=name, - fields=[ - SimpleField(name="id", type=SearchFieldDataType.String, key=True, filterable=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="doc_id", type=SearchFieldDataType.String, filterable=True), - SearchableField(name="file_name", filterable=True, sortable=True), - SimpleField(name="page_number", type=SearchFieldDataType.Int32, 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"), - ], - ), - semantic_search=SemanticSearch( - default_configuration_name=_SEMANTIC_CONFIG, - configurations=[ - SemanticConfiguration( - name=_SEMANTIC_CONFIG, - prioritized_fields=SemanticPrioritizedFields( - title_field=SemanticField(field_name="file_name"), - content_fields=[SemanticField(field_name="chunk")], - ), - ), - ], - ), - ) - - -# ============================================================================= -# Ingestor -# ============================================================================= - - @register_ingestor(_BACKEND_NAME) -class AzureAISearchIngestor(BaseIngestor): - """Parse uploaded files, embed chunks, and write them to Azure AI Search. - - Submission is non-blocking — `submit_job` spawns a daemon thread and returns - a job_id immediately. `get_job_status` reflects the live state. The HTTP - layer routes file uploads here via `set_active_ingestor` (called from - register.py). - """ +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) - cfg = _coerce_config(self.config) - self.cfg = cfg - - # Azure SDK clients - self._credential = _build_search_credential(cfg) - self._index_client = SearchIndexClient(endpoint=str(cfg.endpoint), credential=self._credential) - self._search_client_cache: dict[str, SearchClient] = {} - - # Lazy-init heavy llama-index/embedding deps - self._embedding = None + self._initialize_azure(self.config) self._splitter = None - - # Summary LLM (LangChain wrapper) resolved by shared registration. - self._summary_llm = cfg.summary_llm - - # Job + per-collection in-memory state. AI Search is the source of - # truth for chunks; this just tracks submission lifecycle. - self._jobs_lock = threading.Lock() + self._summary_llm = self.cfg.summary_llm + self._jobs_lock = threading.RLock() self._jobs: dict[str, IngestionJobStatus] = {} - - logger.info( - "AzureAISearchIngestor initialized: endpoint=%s embed_dim=%d chunk_size=%d", - cfg.endpoint, - cfg.embed_dim, - cfg.chunk_size, - ) - - # ------- lazy clients ------- - - @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_endpoint), - api_key=_secret_value(self.cfg.embed_api_key), - ) - return self._embedding + self._files: dict[str, FileInfo] = {} + 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=self.cfg.chunk_size, - chunk_overlap=self.cfg.chunk_overlap, - ) + self._splitter = SentenceSplitter(chunk_size=self.cfg.chunk_size, chunk_overlap=self.cfg.chunk_overlap) return self._splitter - def _get_search_client(self, collection_name: str) -> SearchClient: - if collection_name not in self._search_client_cache: - self._search_client_cache[collection_name] = SearchClient( - endpoint=str(self.cfg.endpoint), - index_name=collection_name, - credential=self._credential, - ) - return self._search_client_cache[collection_name] + def _get_owned_index(self, collection_name: str) -> tuple[SearchIndex, dict[str, Any]]: + index = self._index_client.get_index(self._physical_index_name(collection_name)) + return index, _validate_index_schema(index, collection_name, self.cfg) - # ------- index management ------- - - def _ensure_index(self, collection_name: str) -> None: - """Create the AI Search index if it doesn't already exist.""" - _validate_index_name(collection_name) + def _ensure_index( + self, + collection_name: str, + description: str | None = None, + metadata: dict[str, Any] | None = None, + ) -> tuple[SearchIndex, dict[str, Any]]: + index_name = self._physical_index_name(collection_name) try: - self._index_client.get_index(collection_name) - logger.debug("Index %r already exists", collection_name) + return self._get_owned_index(collection_name) except ResourceNotFoundError: - schema = _build_index_schema(collection_name, self.cfg.embed_dim) - self._index_client.create_index(schema) - logger.info("Created AI Search index %r (embed_dim=%d)", collection_name, self.cfg.embed_dim) + pass + + marker = _new_marker(collection_name, self.cfg, description, metadata) + schema = _build_index_schema(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(index_name) + except ResourceNotFoundError: + raise create_error + return index, _validate_index_schema(index, collection_name, self.cfg) + + def _update_marker(self, collection_name: str, **updates: Any) -> dict[str, Any]: + for attempt in range(3): + index, marker = self._get_owned_index(collection_name) + marker.update(updates) + index.description = _encode_marker(marker) + try: + self._index_client.create_or_update_index(index, match_condition=MatchConditions.IfNotModified) + return marker + except (ResourceModifiedError, HttpResponseError) as error: + if getattr(error, "status_code", None) != 412 or attempt == 2: + raise + raise RuntimeError(f"Failed to update metadata for collection {collection_name!r}") - # ------- jobs ------- + def _update_collection_timestamp(self, collection_name: str) -> None: + self._update_marker(collection_name, updated_at=_utc_now().isoformat()) def submit_job( self, @@ -474,21 +617,41 @@ def submit_job( config: dict[str, Any] | None = None, ) -> str: job_id = str(uuid.uuid4()) - config = config or {} - # AI-Q's HTTP layer passes original_filenames as a parallel LIST (one - # per entry of file_paths), NOT a dict[path, name]. Other callers may - # pass a dict — handle both shapes. - original_filenames = _resolve_filenames(file_paths, config.get("original_filenames")) - - # File-progress placeholders so the UI shows them as UPLOADING immediately. - file_progress = [ - FileProgress( - file_id=str(uuid.uuid4()), - file_name=original_filenames[i], - status=FileStatus.UPLOADING, - ) - for i in range(len(file_paths)) - ] + 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 {} + _encode_marker({"metadata": file_metadata}) + + 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] = [] + 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)) + with self._jobs_lock: + self._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._jobs[job_id] = IngestionJobStatus( @@ -496,42 +659,33 @@ def submit_job( status=JobState.PENDING, collection_name=collection_name, backend=_BACKEND_NAME, - submitted_at=datetime.now(UTC), - total_files=len(file_paths), - processed_files=0, + submitted_at=submitted_at, + total_files=len(validated), file_details=file_progress, ) - logger.info( - "submit_job: id=%s files=%d collection=%r", - job_id, - len(file_paths), - collection_name, - ) - - thread = threading.Thread( + threading.Thread( target=self._process_job, - args=(job_id, file_paths, collection_name, config), + args=(job_id, [path for path, _ in validated], collection_name, job_config), daemon=True, name=f"aiq-azure-search-ingest-{job_id[:8]}", - ) - thread.start() + ).start() return job_id def get_job_status(self, job_id: str) -> IngestionJobStatus: with self._jobs_lock: - existing = self._jobs.get(job_id) - if existing is None: - return IngestionJobStatus( - job_id=job_id, - status=JobState.FAILED, - collection_name="", - backend=_BACKEND_NAME, - submitted_at=datetime.now(UTC), - error_message="Job not found", - file_details=[], - ) - return existing.model_copy(deep=True) + 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, @@ -540,144 +694,209 @@ def _process_job( collection_name: str, config: dict[str, Any], ) -> None: - """Background ingestion worker. Updates job state in place under lock.""" cleanup = bool(config.get("cleanup_files", self.cfg.cleanup_files)) - original_filenames = _resolve_filenames(file_paths, config.get("original_filenames")) - - # Mark PROCESSING and started_at - self._update_job(job_id, status=JobState.PROCESSING, started_at=datetime.now(UTC)) - + self._update_job(job_id, status=JobState.PROCESSING, started_at=_utc_now()) try: self._ensure_index(collection_name) - except Exception as e: # noqa: BLE001 + except Exception as error: # noqa: BLE001 + self._fail_job(job_id, f"Failed to ensure collection {collection_name!r}: {error!s}") if cleanup: - for path in file_paths: - try: - os.unlink(path) - except OSError: - pass - self._fail_job(job_id, f"Failed to ensure index {collection_name!r}: {e!s}") + self._cleanup_paths(file_paths) return failed = 0 - for idx, path in enumerate(file_paths): - file_name = original_filenames[idx] - self._update_file_progress(job_id, idx, status=FileStatus.INGESTING) + 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: - chunks_written = self._process_file( + chunk_count = self._process_file( path=path, collection_name=collection_name, - file_name=file_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, - idx, + index, status=FileStatus.SUCCESS, progress_percent=100.0, - chunks_created=chunks_written, - ) - logger.info("Ingested %s: %d chunks → %s", file_name, chunks_written, collection_name) - except Exception as e: # noqa: BLE001 - msg = self._translate_error(e) - self._update_file_progress( - job_id, - idx, - status=FileStatus.FAILED, - error_message=msg, + chunks_created=chunk_count, ) + except Exception as error: # noqa: BLE001 failed += 1 - logger.exception("Failed to ingest %s", file_name) + message = self._translate_error(error) + self._update_file_progress(job_id, index, status=FileStatus.FAILED, error_message=message) + logger.exception("Failed to ingest %s", detail.file_name) finally: if cleanup: - try: - os.unlink(path) - except OSError: - pass - self._update_job(job_id, processed_files=idx + 1) + self._cleanup_paths([path]) + self._update_job(job_id, processed_files=index + 1) - # Final status 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=datetime.now(UTC), - ) + self._update_job(job_id, status=JobState.COMPLETED, completed_at=_utc_now()) - def _process_file(self, *, path: str, collection_name: str, file_name: str) -> int: - """Parse → chunk → embed → upload one file. Returns chunk count.""" + 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], + ) -> int: from llama_index.core import SimpleDirectoryReader - # 1. Parse. SimpleDirectoryReader returns a Document per logical chunk - # (per-page for PDFs). Metadata includes page_label. - reader = SimpleDirectoryReader(input_files=[path]) - documents = reader.load_data() + old_file_ids = self._find_file_ids_by_name(file_name, collection_name) - {file_id} + documents = SimpleDirectoryReader(input_files=[path]).load_data() if not documents: raise ValueError(f"No content extracted from {file_name}") - - # 2. Chunk each Document with the SentenceSplitter, preserving metadata. nodes = self.splitter.get_nodes_from_documents(documents) if not nodes: raise ValueError(f"Chunking produced 0 chunks for {file_name}") - - # 3. Embed all chunk texts in one batch (NVIDIAEmbedding handles - # rate limiting + batch sizing internally). - texts = [n.get_content() for n in nodes] + texts = [node.get_content() for node in nodes] embeddings = self.embedding.get_text_embedding_batch(texts) - - # 4. Build AI Search documents. doc_id is per-file, chunk_id per-chunk. - doc_id = self._make_doc_id(file_name) - docs = [] - for i, (node, vector) in enumerate(zip(nodes, embeddings, strict=True)): - chunk_id = f"{doc_id}-c{i:04d}" - page = _coerce_page_number(node.metadata.get("page_label")) - docs.append( + 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": chunk_id, + "id": f"{file_id}-{chunk_index:08d}", "chunk": node.get_content(), "embedding": list(vector), - "doc_id": doc_id, + "file_id": file_id, "file_name": file_name, - "page_number": page, - "metadata": "{}", + "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, } ) - # 5. Upload. client = self._get_search_client(collection_name) - result = client.upload_documents(documents=docs) - # Surface partial-failure as an exception so the file gets marked FAILED. - failures = [r for r in result if not r.succeeded] - if failures: - raise RuntimeError(f"AI Search rejected {len(failures)}/{len(docs)} chunks: {failures[0].error_message}") - - # 6. Summarise + register. AI-Q's intent classifier reads registered - # summaries from the system prompt's `available_documents` block, so - # this step is what makes the orchestrator route knowledge queries - # to us instead of falling through to web search. - if self.cfg.generate_summary: - full_text = "\n".join(texts) - summary = self._generate_summary(full_text, file_name) - if summary: + self._upload_documents(client, search_documents) + self._update_collection_timestamp(collection_name) + for old_file_id in old_file_ids: + self._delete_file_documents(old_file_id, collection_name) + with self._jobs_lock: + self._files.pop(old_file_id, None) + + summary = self._generate_summary("\n".join(texts), file_name) if self.cfg.generate_summary else None + if summary: + register_summary(collection_name, file_name, summary) + elif old_file_ids: + unregister_summary(collection_name, file_name) + + with self._jobs_lock: + tracked = self._files.get(file_id) + if tracked: + tracked.status = FileStatus.SUCCESS + tracked.chunk_count = len(search_documents) + tracked.ingested_at = ingested_at + if summary: + tracked.metadata["summary"] = summary + return len(search_documents) + + 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: - register_summary(collection_name, file_name, summary) - logger.info("Registered summary for %s in %s", file_name, collection_name) - except Exception: # noqa: BLE001 - logger.exception("register_summary failed for %s", file_name) + 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 - return len(docs) + def _find_file_ids_by_name(self, file_name: str, collection_name: str) -> set[str]: + client = self._get_search_client(collection_name) + filter_text = f"file_name eq {_odata_literal(file_name)}" + return { + str(hit["file_id"]) + for hit in self._iter_documents(client, filter_text=filter_text, select=["id", "file_id"]) + if hit.get("file_id") + } + + def _delete_file_documents(self, file_id: str, collection_name: str) -> int: + client = self._get_search_client(collection_name) + filter_text = f"file_id eq {_odata_literal(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: - """Generate a one-sentence document summary via the LangChain summary LLM. - - Returns None if no summary LLM is configured or generation fails — both - are non-fatal; the file is still ingested and queryable, just without - an entry in the agent's available_documents block. - """ if self._summary_llm is None: return None - # Truncate to keep the summary prompt small. snippet = text[: self.cfg.summary_max_chars] prompt = f"Summarise the following document ({file_name}) in one sentence (max 30 words):\n\n{snippet}" try: @@ -688,55 +907,50 @@ def _generate_summary(self, text: str, file_name: str) -> str | None: logger.exception("Summary generation failed for %s", file_name) return None - @staticmethod - def _make_doc_id(file_name: str) -> str: - """Stable doc_id from the file name. Lowercase + hash for collision-resistance.""" - digest = hashlib.sha256(file_name.encode("utf-8")).hexdigest()[:8] - slug = re.sub(r"[^a-z0-9]+", "-", file_name.lower()).strip("-")[:32] - return f"{slug}-{digest}" if slug else digest - - # ------- job-state helpers (must be called under _jobs_lock or use it) ------- - def _update_job(self, job_id: str, **fields: Any) -> None: with self._jobs_lock: - existing = self._jobs.get(job_id) - if existing is None: - return - self._jobs[job_id] = existing.model_copy(update=fields) + 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=datetime.now(UTC), - ) + self._update_job(job_id, status=JobState.FAILED, error_message=error_message, completed_at=_utc_now()) - def _update_file_progress(self, job_id: str, idx: int, **fields: Any) -> None: + def _update_file_progress(self, job_id: str, index: int, **fields: Any) -> None: with self._jobs_lock: - existing = self._jobs.get(job_id) - if existing is None or idx >= len(existing.file_details): + job = self._jobs.get(job_id) + if job is None or index >= len(job.file_details): return - new_details = list(existing.file_details) - new_details[idx] = existing.file_details[idx].model_copy(update=fields) - self._jobs[job_id] = existing.model_copy(update={"file_details": new_details}) + 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 = _utc_now() + self._jobs[job_id] = job.model_copy(update={"file_details": details}) - # ------- error translation ------- + @staticmethod + def _cleanup_paths(paths: list[str]) -> None: + for path in paths: + try: + os.unlink(path) + except OSError: + pass @staticmethod - def _translate_error(exc: Exception) -> str: - """User-readable error string for FileProgress.error_message.""" - if isinstance(exc, ResourceNotFoundError): - return f"AI Search index not found: {exc!s}" - if isinstance(exc, ClientAuthenticationError): - return f"AI Search authentication failed: {exc!s}" - if isinstance(exc, ServiceRequestError): - return f"AI Search service unavailable: {exc!s}" - if isinstance(exc, HttpResponseError): - return f"AI Search request failed ({exc.status_code}): {exc.reason or exc!s}" - return str(exc) - - # ------- collections ------- + 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, @@ -744,224 +958,212 @@ def create_collection( description: str | None = None, metadata: dict[str, Any] | None = None, ) -> CollectionInfo: - _validate_index_name(name) - try: - self._ensure_index(name) - except ResourceExistsError: - pass - # AI Search doesn't natively store description/metadata at the index - # level — we just round-trip the request. - return CollectionInfo( - name=name, - description=description, - file_count=0, - chunk_count=0, - backend=_BACKEND_NAME, - metadata=metadata or {}, - created_at=datetime.now(UTC), - updated_at=datetime.now(UTC), - ) + index, marker = self._ensure_index(name, description, metadata) + if description is not None or metadata: + marker = self._update_marker( + name, + description=description if description is not None else marker.get("description"), + metadata={**(marker.get("metadata") or {}), **(metadata or {})}, + updated_at=_utc_now().isoformat(), + ) + return self._collection_info(name, index, marker) def delete_collection(self, name: str) -> bool: try: - self._index_client.delete_index(name) - self._search_client_cache.pop(name, None) - logger.info("Deleted AI Search index %r", name) - return True + index, _marker = self._get_owned_index(name) + except ResourceNotFoundError: + return False + self._index_client.delete_index(index.name) + try: + self._index_client.get_index(index.name) except ResourceNotFoundError: + pass + else: return False + self._search_clients.pop(name, None) + with self._jobs_lock: + self._files = {file_id: info for file_id, info in self._files.items() if info.collection_name != name} + clear_collection_summaries(name) + return True def list_collections(self) -> list[CollectionInfo]: - out: list[CollectionInfo] = [] - for idx in self._index_client.list_indexes(): - out.append( - CollectionInfo( - name=idx.name, - file_count=0, - chunk_count=0, - backend=_BACKEND_NAME, - metadata={}, - ) - ) - return out + collections: list[CollectionInfo] = [] + for index in self._index_client.list_indexes(): + marker = _decode_marker(index.description) + if not marker or marker.get("backend") != _BACKEND_NAME or marker.get("schema_version") != _SCHEMA_VERSION: + continue + collection_name = marker.get("collection_name") + if not isinstance(collection_name, str): + continue + try: + _validate_index_schema(index, collection_name, self.cfg) + collections.append(self._collection_info(collection_name, index, marker)) + except Exception: # noqa: BLE001 + logger.exception("Skipping invalid AI-Q Azure AI Search index %s", index.name) + return collections def get_collection(self, name: str) -> CollectionInfo | None: try: - idx = self._index_client.get_index(name) + index, marker = self._get_owned_index(name) except ResourceNotFoundError: return None + return self._collection_info(name, index, marker) + + def _collection_info(self, name: str, index: SearchIndex, marker: dict[str, Any]) -> CollectionInfo: + client = self._get_search_client(name) + files = self.list_files(name) return CollectionInfo( - name=idx.name, - file_count=0, - chunk_count=0, + name=name, + description=marker.get("description"), + file_count=len(files), + chunk_count=client.get_document_count(), backend=_BACKEND_NAME, - metadata={}, + metadata={ + **(marker.get("metadata") or {}), + "index_name": index.name, + "embedding_model": marker.get("embedding_model"), + "embedding_dim": marker.get("embedding_dim"), + }, + created_at=_parse_timestamp(marker.get("created_at")), + updated_at=_parse_timestamp(marker.get("updated_at")), ) - # ------- files (delegating to submit_job for upload) ------- - def upload_file( self, file_path: str, collection_name: str, metadata: dict[str, Any] | None = None, ) -> FileInfo: - # Synchronous-ish: kick off the job and return file metadata. - # The frontend uses submit_job for actual upload; this is a single-file alias. - cfg: dict[str, Any] = {"original_filenames": [Path(file_path).name]} - if metadata: - cfg["metadata"] = metadata - self.submit_job([file_path], collection_name, cfg) - return FileInfo( - file_id=str(uuid.uuid4()), - file_name=Path(file_path).name, - collection_name=collection_name, - status=FileStatus.UPLOADING, - file_size=Path(file_path).stat().st_size if Path(file_path).is_file() else 0, - chunk_count=0, - metadata=metadata or {}, - uploaded_at=datetime.now(UTC), + 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: - # file_id is the doc_id we stamped at ingest time. Delete all chunks - # whose doc_id matches, and remove the corresponding summary so the - # agent's available_documents block stays consistent. - try: - client = self._get_search_client(collection_name) - # Look up file_name from any matching chunk so we can unregister the summary. - results = client.search( - search_text="*", - filter=f"doc_id eq '{file_id}'", - select=["id", "file_name"], - top=10000, - ) - hits = list(results) - if not hits: - return False - file_name = hits[0].get("file_name") if hits else None - client.delete_documents(documents=[{"id": r["id"]} for r in hits]) - if file_name: - try: - unregister_summary(collection_name, file_name) - except Exception: # noqa: BLE001 - logger.exception("unregister_summary failed for %s", file_name) - return True - except Exception: # noqa: BLE001 - logger.exception("delete_file failed") + info = self.get_file_status(file_id, collection_name) + if info is None: + return False + deleted = self._delete_file_documents(file_id, collection_name) + if deleted == 0 and info.status not in {FileStatus.FAILED, FileStatus.UPLOADING}: return False + with self._jobs_lock: + self._files.pop(file_id, None) + if deleted: + remaining_same_name = any( + item.file_name == info.file_name and item.file_id != file_id + for item in self.list_files(collection_name) + ) + if not remaining_same_name: + unregister_summary(collection_name, info.file_name) + self._update_collection_timestamp(collection_name) + return True def list_files(self, collection_name: str) -> list[FileInfo]: - """List one FileInfo per distinct doc_id in the index. - - AI-Q's frontend polls this after ingestion to refresh the Files panel - and to decide whether to expose `knowledge_search` as an enabled data - source. Returning [] would cause uploaded files to vanish from the UI. - """ try: + self._get_owned_index(collection_name) client = self._get_search_client(collection_name) - except ResourceNotFoundError: - return [] - try: - results = client.search( - search_text="*", - select=["doc_id", "file_name"], - top=10000, + hits = self._iter_documents( + client, + select=[ + "id", + "file_id", + "file_name", + "file_size", + "uploaded_at", + "ingested_at", + "metadata", + ], ) + aggregated: dict[str, dict[str, Any]] = {} + for hit in hits: + file_id = str(hit.get("file_id") or "") + if not file_id: + continue + entry = aggregated.setdefault( + file_id, + { + "file_name": hit.get("file_name") or "unknown", + "file_size": hit.get("file_size"), + "uploaded_at": _parse_timestamp(hit.get("uploaded_at")), + "ingested_at": _parse_timestamp(hit.get("ingested_at")), + "metadata": _parse_metadata(hit.get("metadata")), + "chunk_count": 0, + }, + ) + entry["chunk_count"] += 1 except ResourceNotFoundError: return [] except Exception: # noqa: BLE001 - logger.exception("list_files failed for %r", collection_name) + logger.exception("Failed to list files for %r", collection_name) return [] - agg: dict[str, dict[str, Any]] = {} - for hit in results: - did = hit.get("doc_id") or "unknown" - fn = hit.get("file_name") or "unknown" - entry = agg.setdefault(did, {"file_name": fn, "count": 0}) - entry["count"] += 1 - - now = datetime.now(UTC) - return [ - FileInfo( - file_id=did, - file_name=info["file_name"], + files = { + file_id: FileInfo( + file_id=file_id, + file_name=entry["file_name"], collection_name=collection_name, status=FileStatus.SUCCESS, - file_size=None, - chunk_count=info["count"], - metadata={}, - uploaded_at=now, - ingested_at=now, + file_size=entry["file_size"], + chunk_count=entry["chunk_count"], + uploaded_at=entry["uploaded_at"], + ingested_at=entry["ingested_at"], + metadata=entry["metadata"], ) - for did, info in agg.items() - ] + for file_id, entry in aggregated.items() + } + with self._jobs_lock: + for file_id, tracked in self._files.items(): + if tracked.collection_name == collection_name and ( + tracked.status != FileStatus.SUCCESS or file_id not in files + ): + files[file_id] = tracked.model_copy(deep=True) + return sorted(files.values(), key=lambda item: (item.file_name, item.file_id)) def get_file_status(self, file_id: str, collection_name: str) -> FileInfo | None: - """Return the FileInfo for a single doc_id (file_id == doc_id) if it exists.""" - try: - client = self._get_search_client(collection_name) - results = client.search( - search_text="*", - filter=f"doc_id eq '{file_id}'", - select=["doc_id", "file_name"], - top=10000, - ) - hits = list(results) - except ResourceNotFoundError: - return None - except Exception: # noqa: BLE001 - logger.exception("get_file_status failed for %r in %r", file_id, collection_name) - return None - - if not hits: - return None - - now = datetime.now(UTC) - return FileInfo( - file_id=file_id, - file_name=hits[0].get("file_name") or "unknown", - collection_name=collection_name, - status=FileStatus.SUCCESS, - file_size=None, - chunk_count=len(hits), - metadata={}, - uploaded_at=now, - ingested_at=now, - ) - - # ------- optional ------- + with self._jobs_lock: + tracked = self._files.get(file_id) + if tracked and tracked.collection_name == collection_name and tracked.status != FileStatus.SUCCESS: + return tracked.model_copy(deep=True) + return next((file_info for file_info in self.list_files(collection_name) if file_info.file_id == file_id), None) async def health_check(self) -> bool: - try: + 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("Ingestor health_check failed") + logger.exception("Azure AI Search ingestor health_check failed") return False def _resolve_filenames(file_paths: list[str], raw: Any) -> list[str]: - """Normalise the `original_filenames` config field across caller shapes. - - AI-Q's HTTP route passes a parallel `list[str]` aligned to `file_paths`. - The SDK reference's example uses a `dict[str, str]` keyed by temp path. - Anything else (None, empty, mismatched length) falls back to the basename - of each file_path. - """ if isinstance(raw, dict): - return [raw.get(p) or Path(p).name for p in file_paths] + return [raw.get(path) or Path(path).name for path in file_paths] if isinstance(raw, list): - names: list[str] = [] - for i, p in enumerate(file_paths): - entry = raw[i] if i < len(raw) else None - names.append(entry or Path(p).name) - return names - return [Path(p).name for p in file_paths] + 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: - """LlamaIndex sets page_label as a 1-indexed string for PDFs. Be lenient.""" if page_label is None: return None try: diff --git a/sources/knowledge_layer/src/register.py b/sources/knowledge_layer/src/register.py index 4e121f17c..d2718c9ff 100644 --- a/sources/knowledge_layer/src/register.py +++ b/sources/knowledge_layer/src/register.py @@ -39,6 +39,16 @@ logger = logging.getLogger(__name__) +def _url_from_env(name: str, default: str | None = None) -> HttpUrl | None: + value = os.environ.get(name, 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", "azure_ai_search"] OpenSearchAuthType = Literal["none", "basic", "sigv4"] @@ -253,25 +263,34 @@ class KnowledgeRetrievalConfig(FunctionBaseConfig, name="knowledge_retrieval"): ) # Azure AI Search options azure_search_endpoint: HttpUrl | None = Field( - default=None, - description="Azure AI Search service URL (azure_ai_search only)", + default_factory=lambda: _url_from_env("AZURE_SEARCH_ENDPOINT"), + description="Azure AI Search service URL; defaults to AZURE_SEARCH_ENDPOINT", ) azure_search_api_key: SecretStr | None = Field( - default=None, - description="Optional Azure AI Search admin key; managed identity is used when omitted", + 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_auth_mode: Literal["managed_identity", "api_key"] = Field( - default="managed_identity", - description="Authentication mode for Azure AI Search", + default_factory=lambda: "api_key" if os.environ.get("AZURE_SEARCH_API_KEY") else "managed_identity", + description="Authentication mode; defaults to API key when AZURE_SEARCH_API_KEY is set", + ) + azure_search_index_prefix: str = Field( + default_factory=lambda: os.environ.get("AIQ_AZURE_SEARCH_INDEX_PREFIX", "aiq"), + min_length=1, + description="Prefix for AI-Q-owned indexes; defaults to AIQ_AZURE_SEARCH_INDEX_PREFIX or aiq", ) embed_endpoint: HttpUrl = Field( - default=HttpUrl("https://integrate.api.nvidia.com/v1"), - description="OpenAI-compatible embedding endpoint (azure_ai_search only)", + default_factory=lambda: _url_from_env("AIQ_EMBED_BASE_URL", "https://integrate.api.nvidia.com/v1"), + description="Embedding endpoint; defaults to AIQ_EMBED_BASE_URL", + ) + embed_model: str = Field( + default_factory=lambda: os.environ.get("AIQ_EMBED_MODEL", "nvidia/nv-embed-v1"), + description="Embedding model; defaults to AIQ_EMBED_MODEL", ) embed_dim: int = Field( - default=4096, + default_factory=lambda: int(os.environ.get("AIQ_EMBED_DIM", "4096")), gt=0, - description="Embedding dimensions; must match existing Azure AI Search indexes", + description="Embedding dimensions; defaults to AIQ_EMBED_DIM and must match existing indexes", ) embed_api_key: SecretStr | None = Field( default=None, @@ -341,6 +360,8 @@ def validate_backend_config(self): raise ValueError("azure_search_auth_mode=api_key requires azure_search_api_key") if self.chunk_overlap >= self.chunk_size: raise ValueError("chunk_overlap must be smaller than chunk_size") + if not self.use_hybrid and self.use_semantic_ranker: + raise ValueError("use_semantic_ranker=true requires use_hybrid=true") return self @@ -432,6 +453,7 @@ def _setup_backend(config: KnowledgeRetrievalConfig, summary_llm_obj=None) -> tu "endpoint": str(config.azure_search_endpoint), "api_key": config.azure_search_api_key, "auth_mode": config.azure_search_auth_mode, + "index_prefix": config.azure_search_index_prefix, "embed_endpoint": str(config.embed_endpoint), "embed_model": config.embed_model, "embed_dim": config.embed_dim, diff --git a/tests/knowledge_layer_tests/run_adapter_compliance.py b/tests/knowledge_layer_tests/run_adapter_compliance.py index 99ef916e2..ff3b7e6b2 100755 --- a/tests/knowledge_layer_tests/run_adapter_compliance.py +++ b/tests/knowledge_layer_tests/run_adapter_compliance.py @@ -18,7 +18,12 @@ # Quick mode - registration check only (no files/services needed) 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 +<<<<<<< HEAD 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}' +>>>>>>> 8351174 (fix(knowledge): harden Azure AI Search configuration and ingestion) # Full mode - complete ingestion + retrieval test python tests/knowledge_layer_tests/run_adapter_compliance.py --backend llamaindex @@ -94,7 +99,11 @@ def _import_backend(self): backend_imports = { "llamaindex": "knowledge_layer.llamaindex", "foundational_rag": "knowledge_layer.foundational_rag", +<<<<<<< HEAD "opensearch": "knowledge_layer.opensearch", +======= + "azure_ai_search": "knowledge_layer.azure_ai_search", +>>>>>>> 8351174 (fix(knowledge): harden Azure AI Search configuration and ingestion) } module_name = backend_imports.get(self.backend.lower()) @@ -390,8 +399,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 +465,14 @@ def main(): ) parser.add_argument( +<<<<<<< HEAD "--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, azure_ai_search)", +>>>>>>> 8351174 (fix(knowledge): harden Azure AI Search configuration and ingestion) ) 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 index e91be1d75..3daf4a611 100644 --- a/tests/knowledge_layer_tests/test_azure_ai_search.py +++ b/tests/knowledge_layer_tests/test_azure_ai_search.py @@ -1,17 +1,31 @@ # 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.exceptions import ResourceNotFoundError from azure.core.exceptions import ServiceRequestError +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 _decode_marker +from knowledge_layer.azure_ai_search.adapter import _encode_marker +from knowledge_layer.azure_ai_search.adapter import _index_name_for_collection +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 @@ -23,6 +37,181 @@ 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 + + +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 + + +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.fail_upload_ids: set[str] = set() + self.delete_failures_remaining: dict[str, int] = {} + + 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: + self.documents.pop(key, None) + results.append(FakeIndexingResult(key, succeeded, None if succeeded else "retry")) + return results + + def search(self, search_text="*", filter=None, select=None, order_by=None, top=None, **kwargs): + del search_text, order_by, kwargs + self.search_filters.append(filter) + file_id = _literal(filter, "file_id") + file_name = _literal(filter, "file_name") + after_id = _literal(filter, "id", "gt") + documents = sorted(self.documents.values(), key=lambda item: item["id"]) + if file_id is not None: + documents = [document for document in documents if document.get("file_id") == file_id] + if file_name is not None: + documents = [document for document in documents if document.get("file_name") == file_name] + if after_id is not None: + documents = [document for document in documents if document["id"] > after_id] + if top is not None: + documents = documents[:top] + if select: + return [{key: document.get(key) for key in select} for document in documents] + return [dict(document) for document in documents] + + 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", + "auth_mode": "api_key", + "api_key": SecretStr("test-key"), + "embed_model": "test-embed", + "embed_dim": 4, + "embed_endpoint": "https://integrate.api.nvidia.com/v1", + "use_hybrid": True, + "use_semantic_ranker": True, + "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._get_search_client = lambda collection_name: search_client + return ingestor, search_client + + +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(): @@ -32,9 +221,40 @@ def test_backend_registered_and_implements_sdk_contracts(): assert issubclass(AzureAISearchRetriever, BaseRetriever) -def test_config_requires_endpoint(): +def test_config_requires_endpoint_and_valid_hybrid_semantic_combination(): with pytest.raises(ValueError, match="azure_search_endpoint"): KnowledgeRetrievalConfig(backend="azure_ai_search") + with pytest.raises(ValueError, match="use_semantic_ranker"): + KnowledgeRetrievalConfig( + backend="azure_ai_search", + azure_search_endpoint="https://example.search.windows.net", + use_hybrid=False, + use_semantic_ranker=True, + ) + + +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["auth_mode"] == "api_key" + assert backend_config["api_key"].get_secret_value() == "env-search-key" + assert backend_config["index_prefix"] == "env-aiq" + assert backend_config["embed_endpoint"] == "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" + assert adapter_config.auth_mode == "api_key" def test_api_key_auth_requires_secret(): @@ -46,7 +266,7 @@ def test_api_key_auth_requires_secret(): ) -def test_setup_backend_preserves_secret_types(monkeypatch): +def test_setup_backend_preserves_secrets_and_prefix(monkeypatch): monkeypatch.setenv("KNOWLEDGE_RETRIEVER_BACKEND", "llamaindex") monkeypatch.setenv("KNOWLEDGE_INGESTOR_BACKEND", "llamaindex") config = KnowledgeRetrievalConfig( @@ -55,6 +275,7 @@ def test_setup_backend_preserves_secret_types(monkeypatch): azure_search_auth_mode="api_key", azure_search_api_key="test-search-key", embed_api_key="test-embed-key", + azure_search_index_prefix="tenant-aiq", ) backend, backend_config = _setup_backend(config) @@ -62,7 +283,55 @@ def test_setup_backend_preserves_secret_types(monkeypatch): assert backend == "azure_ai_search" assert isinstance(backend_config["api_key"], SecretStr) assert isinstance(backend_config["embed_api_key"], SecretStr) - assert backend_config["embed_dim"] == 4096 + assert backend_config["index_prefix"] == "tenant-aiq" + + +def test_index_names_are_official_stable_and_collision_safe(): + names = ["Tenant A", "tenant-a", "tenant/a", "TENANT+A"] + indexes = [_index_name_for_collection(name, "AIQ Prod") for name in names] + + assert len(set(indexes)) == len(names) + assert indexes[0] == _index_name_for_collection("Tenant A", "AIQ Prod") + assert all(len(index) <= 128 for index in indexes) + assert all(re.fullmatch(r"[a-z0-9-]+", index) for index in indexes) + _validate_index_name(indexes[0]) + 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, use_semantic_ranker=True) + marker = _new_marker("docs", cfg, "Docs", {"tenant": "alpha"}) + index = _build_index_schema("aiq-docs-123456789abc", 4, _encode_marker(marker)) + + assert _validate_index_schema(index, "docs", cfg)["metadata"] == {"tenant": "alpha"} + index.description = "unmanaged" + with pytest.raises(RuntimeError, match="not owned"): + _validate_index_schema(index, "docs", 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, "docs", 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_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( @@ -78,19 +347,18 @@ def test_normalize_clamps_and_prefers_reranker_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_populates_citation_and_metadata(): +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", - "doc_id": "doc-1", "metadata": '{"section":"intro"}', } ) @@ -98,11 +366,9 @@ def test_normalize_populates_citation_and_metadata(): assert cited.display_citation == "report.pdf, p.3" assert cited.content_type is ContentType.TEXT - assert cited.metadata == {"doc_id": "doc-1", "raw": '{"section":"intro"}'} + assert cited.metadata == {"section": "intro", "file_id": "file-1"} assert fallback.content == "" assert fallback.file_name == "unknown" - assert fallback.page_number is None - assert fallback.display_citation == "unknown" UUID(fallback.chunk_id) @@ -110,9 +376,7 @@ 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} ) - result = RetrievalResult(query="query", backend="azure_ai_search", chunks=[chunk]) - - formatted = _format_results(result, "query") + 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 @@ -120,11 +384,6 @@ def test_shared_formatter_retains_source_and_citation_lines(): @pytest.mark.asyncio async def test_retrieve_builds_hybrid_semantic_search_request(): - class FakeEmbedding: - def get_query_embedding(self, query): - assert query == "hello" - return [0.1, 0.2] - class FakeClient: def search(self, **kwargs): self.kwargs = kwargs @@ -133,50 +392,335 @@ def search(self, **kwargs): client = FakeClient() retriever = AzureAISearchRetriever.__new__(AzureAISearchRetriever) retriever.cfg = SimpleNamespace(use_hybrid=True, use_semantic_ranker=True) - retriever._embedding = FakeEmbedding() + retriever._embedding = FakeEmbedding(2) retriever._get_client = lambda collection_name: client - result = await retriever.retrieve("hello", "session-1", top_k=5, filters={"$filter": "doc_id eq 'doc-1'"}) + result = await retriever.retrieve("hello", "session-1", top_k=5, filters={"$filter": "file_id eq 'file-1'"}) assert result.success - assert [chunk.content for chunk in result.chunks] == ["answer"] assert client.kwargs["search_text"] == "hello" assert client.kwargs["query_type"] == "semantic" - assert client.kwargs["semantic_configuration_name"] == "default-semantic" - assert client.kwargs["filter"] == "doc_id eq 'doc-1'" - assert client.kwargs["top"] == 5 - assert client.kwargs["select"] == ["id", "chunk", "doc_id", "file_name", "page_number", "metadata"] - vector_query = client.kwargs["vector_queries"][0] - assert vector_query.vector == [0.1, 0.2] - assert vector_query.k_nearest_neighbors == 20 - assert vector_query.fields == "embedding" + assert client.kwargs["filter"] == "file_id eq 'file-1'" + assert client.kwargs["select"] == ["id", "chunk", "file_id", "file_name", "page_number", "metadata"] -def test_index_schema_uses_requested_vector_dimensions(): - schema = _build_index_schema("session-1", 4096) - embedding = next(field for field in schema.fields if field.name == "embedding") +def test_submit_job_uses_one_canonical_file_id(monkeypatch, tmp_path): + class NoStartThread: + def __init__(self, *args, **kwargs): + pass - assert embedding.vector_search_dimensions == 4096 - assert schema.semantic_search.default_configuration_name == "default-semantic" + 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"]}) + job = ingestor.get_job_status(job_id) + file_id = job.file_details[0].file_id + + assert file_id in ingestor._files + assert ingestor._files[file_id].file_name == "original.txt" + + +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) -def test_frontend_collection_names_are_valid(): - _validate_index_name("s_f031d9cf_123") + assert client.documents == {} + assert client.delete_batches - with pytest.raises(ValueError, match="Invalid AI Search index name"): - _validate_index_name("Invalid Name") +def test_delete_retries_per_document_failures(): + ingestor, client = _ingestor() + client.documents["chunk-1"] = {"id": "chunk-1"} + client.delete_failures_remaining["chunk-1"] = 1 -def test_filename_and_page_normalization(tmp_path): + 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): + client.documents[f"chunk-{index}"] = { + "id": f"chunk-{index}", + "file_id": f"file-{index}", + "file_name": f"file-{index}.txt", + "file_size": index, + "uploaded_at": datetime.now(UTC), + "ingested_at": datetime.now(UTC), + "metadata": "{}", + } + + 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_replacement_removes_old_generation(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", + "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 = 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 set(client.documents) == {"new-00000000"} + assert client.documents["new-00000000"]["metadata"] == '{"tenant":"alpha"}' + + +def test_failed_new_generation_preserves_old_generation(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", + "file_id": "old", + "file_name": "document.txt", + } + client.fail_upload_ids.add("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={}, + ) + + assert set(client.documents) == {"old-00000000"} + + +def test_replacement_cleanup_failure_retains_complete_new_generation(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", + "file_id": "old", + "file_name": "document.txt", + } + client.delete_failures_remaining["old-00000000"] = 10 + + with pytest.raises(RuntimeError, match="rejected delete"): + 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={}, + ) + + assert set(client.documents) == {"old-00000000", "new-00000000"} + + +def test_metadata_counts_and_status_round_trip(): + ingestor, client = _ingestor() + ingestor.create_collection("docs", metadata={"tenant": "alpha"}) + now = datetime.now(UTC) + for index in range(3): + client.documents[f"file-1-{index:08d}"] = { + "id": f"file-1-{index:08d}", + "file_id": "file-1", + "file_name": "report.pdf", + "file_size": 42, + "uploaded_at": now, + "ingested_at": now, + "metadata": '{"section":"finance"}', + } + + 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" + + +def test_failed_uploads_remain_visible(): + ingestor, _client = _ingestor() + ingestor.create_collection("docs") + ingestor._files["failed"] = azure_adapter.FileInfo( + file_id="failed", + file_name="failed.txt", + collection_name="docs", + status=FileStatus.FAILED, + error_message="bad file", + ) + + 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() + ingestor.create_collection("old") + ingestor.create_collection("new") + old_name = ingestor._physical_index_name("old") + new_name = ingestor._physical_index_name("new") + old_marker = _decode_marker(ingestor._index_client.indexes[old_name].description) + new_marker = _decode_marker(ingestor._index_client.indexes[new_name].description) + old_marker["updated_at"] = (datetime.now(UTC) - timedelta(hours=25)).isoformat() + new_marker["updated_at"] = datetime.now(UTC).isoformat() + ingestor._index_client.indexes[old_name].description = _encode_marker(old_marker) + ingestor._index_client.indexes[new_name].description = _encode_marker(new_marker) + 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 old_name not in ingestor._index_client.indexes + assert new_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() + 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() + ingestor.create_collection("docs") + client.documents["file-1-00000000"] = { + "id": "file-1-00000000", + "file_id": "file-1", + "file_name": "report.pdf", + "file_size": 1, + "uploaded_at": datetime.now(UTC), + "ingested_at": datetime.now(UTC), + "metadata": "{}", + } + client.delete_failures_remaining["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 == [] + + +@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 schema.semantic_search.default_configuration_name == "default-semantic" + + +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 - - -def test_service_connection_error_is_user_readable(): - message = AzureAISearchIngestor._translate_error(ServiceRequestError("connection refused")) - - assert message == "AI Search service unavailable: connection refused" + assert ( + AzureAISearchIngestor._translate_error(ServiceRequestError("connection refused")) + == "AI Search service unavailable: connection refused" + ) From 620ea5fe2f6c7e86ca770d9a3d26c2065ac3e45a Mon Sep 17 00:00:00 2001 From: Harmke Alkemade Date: Fri, 3 Jul 2026 14:04:52 +0200 Subject: [PATCH 03/27] fix(knowledge): align Azure embedding URL config name Signed-off-by: Harmke Alkemade --- docs/source/customization/configuration-reference.md | 2 +- sources/knowledge_layer/src/azure_ai_search/adapter.py | 4 ++-- sources/knowledge_layer/src/register.py | 6 +++--- tests/knowledge_layer_tests/test_azure_ai_search.py | 4 ++-- 4 files changed, 8 insertions(+), 8 deletions(-) diff --git a/docs/source/customization/configuration-reference.md b/docs/source/customization/configuration-reference.md index 27f23170c..799e219c2 100644 --- a/docs/source/customization/configuration-reference.md +++ b/docs/source/customization/configuration-reference.md @@ -280,7 +280,7 @@ environment. `AZURE_SEARCH_API_KEY` is optional; when absent, the adapter uses | `azure_search_auth_mode` | `str` | Auto | Uses `api_key` when `AZURE_SEARCH_API_KEY` is set; otherwise `managed_identity`. | | `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` | Namespace prefix for AI-Q-owned indexes. | -| `embed_endpoint` | `URL` | `AIQ_EMBED_BASE_URL` or NVIDIA API | OpenAI-compatible embedding endpoint. | +| `embed_base_url` | `URL` | `AIQ_EMBED_BASE_URL` or NVIDIA API | OpenAI-compatible embedding base URL. | | `embed_model` | `str` | `AIQ_EMBED_MODEL` or `nvidia/nv-embed-v1` | Embedding model used for Azure ingestion and retrieval. | | `embed_dim` | `int` | `AIQ_EMBED_DIM` or `4096` | Embedding dimensions; must match the model and existing index schema. | | `embed_api_key` | `SecretStr` | `None` | Optional embedding API key; falls back to `NVIDIA_API_KEY`. | diff --git a/sources/knowledge_layer/src/azure_ai_search/adapter.py b/sources/knowledge_layer/src/azure_ai_search/adapter.py index 83c3ff6a8..5c4469b6e 100644 --- a/sources/knowledge_layer/src/azure_ai_search/adapter.py +++ b/sources/knowledge_layer/src/azure_ai_search/adapter.py @@ -83,7 +83,7 @@ def _coerce_config(config: dict[str, Any] | None) -> SimpleNamespace: values: dict[str, Any] = { "endpoint": os.environ.get("AZURE_SEARCH_ENDPOINT"), "api_key": os.environ.get("AZURE_SEARCH_API_KEY"), - "embed_endpoint": os.environ.get("AIQ_EMBED_BASE_URL", "https://integrate.api.nvidia.com/v1"), + "embed_base_url": os.environ.get("AIQ_EMBED_BASE_URL", "https://integrate.api.nvidia.com/v1"), "embed_model": os.environ.get("AIQ_EMBED_MODEL", "nvidia/nv-embed-v1"), "embed_dim": int(os.environ.get("AIQ_EMBED_DIM", "4096")), "embed_api_key": None, @@ -418,7 +418,7 @@ def embedding(self): self._embedding = NVIDIAEmbedding( model=self.cfg.embed_model, - base_url=str(self.cfg.embed_endpoint), + base_url=str(self.cfg.embed_base_url), api_key=_secret_value(self.cfg.embed_api_key), ) return self._embedding diff --git a/sources/knowledge_layer/src/register.py b/sources/knowledge_layer/src/register.py index d2718c9ff..99b1112af 100644 --- a/sources/knowledge_layer/src/register.py +++ b/sources/knowledge_layer/src/register.py @@ -279,9 +279,9 @@ class KnowledgeRetrievalConfig(FunctionBaseConfig, name="knowledge_retrieval"): min_length=1, description="Prefix for AI-Q-owned indexes; defaults to AIQ_AZURE_SEARCH_INDEX_PREFIX or aiq", ) - embed_endpoint: HttpUrl = Field( + embed_base_url: HttpUrl = Field( default_factory=lambda: _url_from_env("AIQ_EMBED_BASE_URL", "https://integrate.api.nvidia.com/v1"), - description="Embedding endpoint; defaults to AIQ_EMBED_BASE_URL", + description="Embedding base URL; defaults to AIQ_EMBED_BASE_URL", ) embed_model: str = Field( default_factory=lambda: os.environ.get("AIQ_EMBED_MODEL", "nvidia/nv-embed-v1"), @@ -454,7 +454,7 @@ def _setup_backend(config: KnowledgeRetrievalConfig, summary_llm_obj=None) -> tu "api_key": config.azure_search_api_key, "auth_mode": config.azure_search_auth_mode, "index_prefix": config.azure_search_index_prefix, - "embed_endpoint": str(config.embed_endpoint), + "embed_base_url": str(config.embed_base_url), "embed_model": config.embed_model, "embed_dim": config.embed_dim, "embed_api_key": config.embed_api_key, diff --git a/tests/knowledge_layer_tests/test_azure_ai_search.py b/tests/knowledge_layer_tests/test_azure_ai_search.py index 3daf4a611..fe90b7e56 100644 --- a/tests/knowledge_layer_tests/test_azure_ai_search.py +++ b/tests/knowledge_layer_tests/test_azure_ai_search.py @@ -185,7 +185,7 @@ def _config(**overrides): "api_key": SecretStr("test-key"), "embed_model": "test-embed", "embed_dim": 4, - "embed_endpoint": "https://integrate.api.nvidia.com/v1", + "embed_base_url": "https://integrate.api.nvidia.com/v1", "use_hybrid": True, "use_semantic_ranker": True, "start_ttl_cleanup": False, @@ -250,7 +250,7 @@ def test_config_uses_shared_environment_defaults(monkeypatch): assert backend_config["auth_mode"] == "api_key" assert backend_config["api_key"].get_secret_value() == "env-search-key" assert backend_config["index_prefix"] == "env-aiq" - assert backend_config["embed_endpoint"] == "https://embed.example.com/v1" + 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" From e8046d218fab260fb4af8094527220e6df7b25db Mon Sep 17 00:00:00 2001 From: Harmke Alkemade Date: Fri, 3 Jul 2026 14:28:55 +0200 Subject: [PATCH 04/27] fix(knowledge): make Azure semantic ranking opt-in Signed-off-by: Harmke Alkemade --- docs/source/customization/configuration-reference.md | 4 ++-- docs/source/customization/knowledge-layer.md | 6 +++++- docs/source/examples/azure-ai-search.md | 6 +++++- sources/knowledge_layer/src/azure_ai_search/README.md | 6 +++++- sources/knowledge_layer/src/azure_ai_search/adapter.py | 2 +- sources/knowledge_layer/src/register.py | 7 ++++++- tests/knowledge_layer_tests/test_azure_ai_search.py | 7 +++++++ 7 files changed, 31 insertions(+), 7 deletions(-) diff --git a/docs/source/customization/configuration-reference.md b/docs/source/customization/configuration-reference.md index 799e219c2..fc1569a85 100644 --- a/docs/source/customization/configuration-reference.md +++ b/docs/source/customization/configuration-reference.md @@ -256,7 +256,7 @@ functions: backend: azure_ai_search collection_name: ${COLLECTION_NAME:-test_collection} use_hybrid: true - use_semantic_ranker: true + use_semantic_ranker: false ``` This example reads `AZURE_SEARCH_ENDPOINT` and `NVIDIA_API_KEY` from the @@ -285,7 +285,7 @@ environment. `AZURE_SEARCH_API_KEY` is optional; when absent, the adapter uses | `embed_dim` | `int` | `AIQ_EMBED_DIM` or `4096` | Embedding dimensions; must match the model and existing index schema. | | `embed_api_key` | `SecretStr` | `None` | Optional embedding API key; falls back to `NVIDIA_API_KEY`. | | `use_hybrid` | `bool` | `true` | Combine lexical and vector retrieval. | -| `use_semantic_ranker` | `bool` | `true` | Apply Azure semantic ranking; requires `use_hybrid: true`. | +| `use_semantic_ranker` | `bool` | `false` | Apply Azure semantic ranking when supported; requires `use_hybrid: true`. | | `chunk_size` | `int` | `512` | Tokens per Azure-ingested chunk. | | `chunk_overlap` | `int` | `64` | Token overlap; must be smaller than `chunk_size`. | | `summary_max_chars` | `int` | `1000` | Maximum document characters sent to the summary model. | diff --git a/docs/source/customization/knowledge-layer.md b/docs/source/customization/knowledge-layer.md index e9475524e..e3d2fb5d8 100644 --- a/docs/source/customization/knowledge-layer.md +++ b/docs/source/customization/knowledge-layer.md @@ -164,7 +164,7 @@ functions: backend: azure_ai_search collection_name: my_docs use_hybrid: true - use_semantic_ranker: true + use_semantic_ranker: false ``` Set `AZURE_SEARCH_ENDPOINT` and `NVIDIA_API_KEY` in the environment. Setting @@ -173,6 +173,10 @@ Set `AZURE_SEARCH_ENDPOINT` and `NVIDIA_API_KEY` in the environment. Setting LlamaIndex backend through `AIQ_EMBED_BASE_URL` and `AIQ_EMBED_MODEL`; set `AIQ_EMBED_DIM` when changing the model dimensions. +Semantic ranking is disabled by default because support depends on the Azure AI +Search service. Set `use_semantic_ranker: true` only when semantic ranking is +enabled; it also requires `use_hybrid: true`. + Azure maps logical collection names to collision-safe physical indexes under `azure_search_index_prefix`. Only indexes containing an AI-Q ownership and schema marker are listed, queried, or deleted. Legacy indexes named directly after collections are ignored; re-ingest those collections after enabling this backend. diff --git a/docs/source/examples/azure-ai-search.md b/docs/source/examples/azure-ai-search.md index 510721df7..02feeba49 100644 --- a/docs/source/examples/azure-ai-search.md +++ b/docs/source/examples/azure-ai-search.md @@ -36,7 +36,7 @@ functions: collection_name: ${COLLECTION_NAME:-aiq_default} top_k: 5 use_hybrid: true - use_semantic_ranker: true + use_semantic_ranker: false generate_summary: true summary_model: summary_llm @@ -50,6 +50,10 @@ 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`. +Semantic ranking is opt-in because availability depends on the Azure AI Search +service. Set `use_semantic_ranker: true` only when semantic ranking is enabled; +it also requires `use_hybrid: true`. + Existing indexes must use the configured `embed_dim`. Delete and re-ingest a collection when changing embedding dimensions. Frontend WebSocket queries use the conversation ID as the collection; direct API tests must supply equivalent diff --git a/sources/knowledge_layer/src/azure_ai_search/README.md b/sources/knowledge_layer/src/azure_ai_search/README.md index 126b45b0a..739e0892b 100644 --- a/sources/knowledge_layer/src/azure_ai_search/README.md +++ b/sources/knowledge_layer/src/azure_ai_search/README.md @@ -33,7 +33,7 @@ functions: backend: azure_ai_search collection_name: ${COLLECTION_NAME:-aiq_default} use_hybrid: true - use_semantic_ranker: true + use_semantic_ranker: false top_k: 5 chunk_size: 512 chunk_overlap: 64 @@ -49,6 +49,10 @@ 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`. +Semantic ranking is disabled by default because support depends on the Azure AI +Search service. Set `use_semantic_ranker: true` only when semantic ranking is +enabled; it also requires `use_hybrid: true`. + When `AZURE_SEARCH_API_KEY` is absent, the adapter uses `DefaultAzureCredential`. Set `AZURE_CLIENT_ID` when a user-assigned identity should be selected. The identity needs permission to create and delete indexes diff --git a/sources/knowledge_layer/src/azure_ai_search/adapter.py b/sources/knowledge_layer/src/azure_ai_search/adapter.py index 5c4469b6e..886cb5b61 100644 --- a/sources/knowledge_layer/src/azure_ai_search/adapter.py +++ b/sources/knowledge_layer/src/azure_ai_search/adapter.py @@ -88,7 +88,7 @@ def _coerce_config(config: dict[str, Any] | None) -> SimpleNamespace: "embed_dim": int(os.environ.get("AIQ_EMBED_DIM", "4096")), "embed_api_key": None, "use_hybrid": True, - "use_semantic_ranker": True, + "use_semantic_ranker": False, "chunk_size": 512, "chunk_overlap": 64, "summary_max_chars": 1000, diff --git a/sources/knowledge_layer/src/register.py b/sources/knowledge_layer/src/register.py index 99b1112af..8322ad60e 100644 --- a/sources/knowledge_layer/src/register.py +++ b/sources/knowledge_layer/src/register.py @@ -297,7 +297,12 @@ class KnowledgeRetrievalConfig(FunctionBaseConfig, name="knowledge_retrieval"): description="Optional embedding API key; NVIDIA_API_KEY is used when omitted", ) use_hybrid: bool = Field(default=True, description="Combine vector and keyword search (azure_ai_search only)") - use_semantic_ranker: bool = Field(default=True, description="Enable Azure semantic ranking (azure_ai_search only)") + use_semantic_ranker: bool = Field( + default=False, + description=( + "Enable Azure semantic ranking when supported by the Azure AI Search service (azure_ai_search only)" + ), + ) chunk_size: int = Field(default=512, gt=0, description="Tokens per chunk (azure_ai_search only)") chunk_overlap: int = Field(default=64, ge=0, description="Token overlap between chunks (azure_ai_search only)") summary_max_chars: int = Field( diff --git a/tests/knowledge_layer_tests/test_azure_ai_search.py b/tests/knowledge_layer_tests/test_azure_ai_search.py index fe90b7e56..02f5cf418 100644 --- a/tests/knowledge_layer_tests/test_azure_ai_search.py +++ b/tests/knowledge_layer_tests/test_azure_ai_search.py @@ -224,6 +224,11 @@ def test_backend_registered_and_implements_sdk_contracts(): def test_config_requires_endpoint_and_valid_hybrid_semantic_combination(): 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 not config.use_semantic_ranker with pytest.raises(ValueError, match="use_semantic_ranker"): KnowledgeRetrievalConfig( backend="azure_ai_search", @@ -253,8 +258,10 @@ def test_config_uses_shared_environment_defaults(monkeypatch): 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 not backend_config["use_semantic_ranker"] assert adapter_config.endpoint == "https://env.search.windows.net" assert adapter_config.auth_mode == "api_key" + assert not adapter_config.use_semantic_ranker def test_api_key_auth_requires_secret(): From 900f8c8378ad4418a5a79c1789207f7762dc97b4 Mon Sep 17 00:00:00 2001 From: Harmke Alkemade Date: Fri, 3 Jul 2026 14:37:16 +0200 Subject: [PATCH 05/27] fix(knowledge): avoid summary storage when disabled Signed-off-by: Harmke Alkemade --- .../src/azure_ai_search/adapter.py | 7 ++-- .../test_azure_ai_search.py | 38 +++++++++++++++++-- 2 files changed, 39 insertions(+), 6 deletions(-) diff --git a/sources/knowledge_layer/src/azure_ai_search/adapter.py b/sources/knowledge_layer/src/azure_ai_search/adapter.py index 886cb5b61..25a996194 100644 --- a/sources/knowledge_layer/src/azure_ai_search/adapter.py +++ b/sources/knowledge_layer/src/azure_ai_search/adapter.py @@ -798,7 +798,7 @@ def _process_file( summary = self._generate_summary("\n".join(texts), file_name) if self.cfg.generate_summary else None if summary: register_summary(collection_name, file_name, summary) - elif old_file_ids: + elif self.cfg.generate_summary and old_file_ids: unregister_summary(collection_name, file_name) with self._jobs_lock: @@ -983,7 +983,8 @@ def delete_collection(self, name: str) -> bool: self._search_clients.pop(name, None) with self._jobs_lock: self._files = {file_id: info for file_id, info in self._files.items() if info.collection_name != name} - clear_collection_summaries(name) + if self.cfg.generate_summary: + clear_collection_summaries(name) return True def list_collections(self) -> list[CollectionInfo]: @@ -1068,7 +1069,7 @@ def delete_file(self, file_id: str, collection_name: str) -> bool: item.file_name == info.file_name and item.file_id != file_id for item in self.list_files(collection_name) ) - if not remaining_same_name: + if self.cfg.generate_summary and not remaining_same_name: unregister_summary(collection_name, info.file_name) self._update_collection_timestamp(collection_name) return True diff --git a/tests/knowledge_layer_tests/test_azure_ai_search.py b/tests/knowledge_layer_tests/test_azure_ai_search.py index 02f5cf418..09da621ae 100644 --- a/tests/knowledge_layer_tests/test_azure_ai_search.py +++ b/tests/knowledge_layer_tests/test_azure_ai_search.py @@ -500,6 +500,12 @@ def test_pagination_is_exhaustive_and_file_ids_are_odata_escaped(monkeypatch): def test_same_name_replacement_removes_old_generation(monkeypatch, tmp_path): _install_reader(monkeypatch) + cleared = [] + monkeypatch.setattr( + azure_adapter, + "unregister_summary", + lambda collection, file_name: cleared.append((collection, file_name)), + ) path = tmp_path / "document.txt" path.write_text("new content", encoding="utf-8") ingestor, client = _ingestor() @@ -533,6 +539,7 @@ def test_same_name_replacement_removes_old_generation(monkeypatch, tmp_path): assert count == 1 assert set(client.documents) == {"new-00000000"} assert client.documents["new-00000000"]["metadata"] == '{"tenant":"alpha"}' + assert cleared == [] def test_failed_new_generation_preserves_old_generation(monkeypatch, tmp_path): @@ -637,7 +644,7 @@ def test_failed_uploads_remain_visible(): def test_ttl_deletes_only_expired_owned_collection_and_clears_summary(monkeypatch): - ingestor, _client = _ingestor() + ingestor, _client = _ingestor(generate_summary=True) ingestor.create_collection("old") ingestor.create_collection("new") old_name = ingestor._physical_index_name("old") @@ -662,7 +669,7 @@ def test_ttl_deletes_only_expired_owned_collection_and_clears_summary(monkeypatc def test_collection_summary_clears_only_after_confirmed_delete(monkeypatch): - ingestor, _client = _ingestor() + ingestor, _client = _ingestor(generate_summary=True) ingestor.create_collection("docs") cleared = [] monkeypatch.setattr(azure_adapter, "clear_collection_summaries", cleared.append) @@ -674,7 +681,7 @@ def test_collection_summary_clears_only_after_confirmed_delete(monkeypatch): def test_file_summary_clears_only_after_confirmed_delete(monkeypatch): - ingestor, client = _ingestor() + ingestor, client = _ingestor(generate_summary=True) ingestor.create_collection("docs") client.documents["file-1-00000000"] = { "id": "file-1-00000000", @@ -695,6 +702,31 @@ def test_file_summary_clears_only_after_confirmed_delete(monkeypatch): assert cleared == [] +def test_summary_cleanup_is_skipped_when_summaries_are_disabled(monkeypatch): + ingestor, client = _ingestor() + ingestor.create_collection("docs") + client.documents["file-1-00000000"] = { + "id": "file-1-00000000", + "file_id": "file-1", + "file_name": "report.pdf", + "file_size": 1, + "uploaded_at": datetime.now(UTC), + "ingested_at": datetime.now(UTC), + "metadata": "{}", + } + 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 == [] + + @pytest.mark.asyncio async def test_health_checks_run_sync_sdk_off_event_loop(monkeypatch): ingestor, _client = _ingestor() From 359516adc60550343c9a88c96cd3875abf1127f9 Mon Sep 17 00:00:00 2001 From: Harmke Alkemade Date: Fri, 3 Jul 2026 15:13:28 +0200 Subject: [PATCH 06/27] test(knowledge): tolerate Azure Search deletion latency Signed-off-by: Harmke Alkemade --- .../knowledge_layer_tests/run_adapter_compliance.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/tests/knowledge_layer_tests/run_adapter_compliance.py b/tests/knowledge_layer_tests/run_adapter_compliance.py index ff3b7e6b2..c44bb46c4 100755 --- a/tests/knowledge_layer_tests/run_adapter_compliance.py +++ b/tests/knowledge_layer_tests/run_adapter_compliance.py @@ -406,12 +406,12 @@ def _test_delete_file(self): if not result: return False, f"delete_file returned False for '{filename}'" - # Verify deletion - remaining = self.ingestor.list_files(self.collection_name) - remaining_names = [f.file_name for f in remaining] - - if filename in remaining_names: - return False, f"File '{filename}' still exists after deletion" + # Azure AI Search applies document deletions asynchronously. + deadline = time.monotonic() + min(self.timeout, 30) + while filename in {file.file_name for file in self.ingestor.list_files(self.collection_name)}: + if time.monotonic() >= deadline: + return False, f"File '{filename}' still exists after deletion" + time.sleep(1) return True, f"Deleted '{filename}'" From bdb3ba0e98eba752aacc83d77c4909b73f73ba85 Mon Sep 17 00:00:00 2001 From: Harmke Alkemade Date: Fri, 3 Jul 2026 15:13:39 +0200 Subject: [PATCH 07/27] docs(knowledge): document Azure Search managed identity roles Signed-off-by: Harmke Alkemade --- docs/source/customization/knowledge-layer.md | 4 +++- docs/source/examples/azure-ai-search.md | 15 +++++++++++++++ .../knowledge_layer/src/azure_ai_search/README.md | 6 ++++-- 3 files changed, 22 insertions(+), 3 deletions(-) diff --git a/docs/source/customization/knowledge-layer.md b/docs/source/customization/knowledge-layer.md index e3d2fb5d8..07b2d8508 100644 --- a/docs/source/customization/knowledge-layer.md +++ b/docs/source/customization/knowledge-layer.md @@ -169,7 +169,9 @@ functions: Set `AZURE_SEARCH_ENDPOINT` and `NVIDIA_API_KEY` in the environment. Setting `AZURE_SEARCH_API_KEY` selects key authentication; otherwise Azure -`DefaultAzureCredential` is used. Embedding defaults can be shared with the +`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. diff --git a/docs/source/examples/azure-ai-search.md b/docs/source/examples/azure-ai-search.md index 02feeba49..191778973 100644 --- a/docs/source/examples/azure-ai-search.md +++ b/docs/source/examples/azure-ai-search.md @@ -25,6 +25,21 @@ export NVIDIA_API_KEY= export AZURE_SEARCH_API_KEY= ``` +## 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, inspect, update, and delete AI-Q collection indexes. | +| `Search Index Data Contributor` | Upload, query, and delete index documents. | + +Assign the roles at the search-service scope because AI-Q creates one index per +logical collection. The principal ID is the object ID of the system-assigned or +user-assigned managed identity running AI-Q. + + Replace the `knowledge_search` block in a web configuration such as `configs/config_web_default_llamaindex.yml`: diff --git a/sources/knowledge_layer/src/azure_ai_search/README.md b/sources/knowledge_layer/src/azure_ai_search/README.md index 739e0892b..7aff59ff8 100644 --- a/sources/knowledge_layer/src/azure_ai_search/README.md +++ b/sources/knowledge_layer/src/azure_ai_search/README.md @@ -55,8 +55,10 @@ enabled; it also requires `use_hybrid: true`. When `AZURE_SEARCH_API_KEY` is absent, the adapter uses `DefaultAzureCredential`. Set `AZURE_CLIENT_ID` when a user-assigned identity -should be selected. The identity needs permission to create and delete indexes -and to read, write, and delete index documents. +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 AI-Q creates one index per collection. The adapter parses PDF, DOCX, TXT, and Markdown uploads with LlamaIndex, creates one namespaced Azure AI Search index per AI-Q collection, and performs From ea2566f8063c0a19d358e7e92bcbd9fc2e9a3165 Mon Sep 17 00:00:00 2001 From: Harmke Alkemade Date: Tue, 7 Jul 2026 10:41:06 +0200 Subject: [PATCH 08/27] fix(knowledge): resolve adapter compliance merge markers Signed-off-by: Harmke Alkemade --- .../run_adapter_compliance.py | 24 ++++++------------- 1 file changed, 7 insertions(+), 17 deletions(-) diff --git a/tests/knowledge_layer_tests/run_adapter_compliance.py b/tests/knowledge_layer_tests/run_adapter_compliance.py index c44bb46c4..2e86a1757 100755 --- a/tests/knowledge_layer_tests/run_adapter_compliance.py +++ b/tests/knowledge_layer_tests/run_adapter_compliance.py @@ -18,12 +18,9 @@ # Quick mode - registration check only (no files/services needed) 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 -<<<<<<< HEAD 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}' ->>>>>>> 8351174 (fix(knowledge): harden Azure AI Search configuration and ingestion) # Full mode - complete ingestion + retrieval test python tests/knowledge_layer_tests/run_adapter_compliance.py --backend llamaindex @@ -99,11 +96,8 @@ def _import_backend(self): backend_imports = { "llamaindex": "knowledge_layer.llamaindex", "foundational_rag": "knowledge_layer.foundational_rag", -<<<<<<< HEAD "opensearch": "knowledge_layer.opensearch", -======= "azure_ai_search": "knowledge_layer.azure_ai_search", ->>>>>>> 8351174 (fix(knowledge): harden Azure AI Search configuration and ingestion) } module_name = backend_imports.get(self.backend.lower()) @@ -406,12 +400,12 @@ def _test_delete_file(self): if not result: return False, f"delete_file returned False for '{filename}'" - # Azure AI Search applies document deletions asynchronously. - deadline = time.monotonic() + min(self.timeout, 30) - while filename in {file.file_name for file in self.ingestor.list_files(self.collection_name)}: - if time.monotonic() >= deadline: - return False, f"File '{filename}' still exists after deletion" - time.sleep(1) + # Verify deletion + remaining = self.ingestor.list_files(self.collection_name) + remaining_names = [f.file_name for f in remaining] + + if filename in remaining_names: + return False, f"File '{filename}' still exists after deletion" return True, f"Deleted '{filename}'" @@ -465,14 +459,10 @@ def main(): ) parser.add_argument( -<<<<<<< HEAD - "--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, azure_ai_search)", ->>>>>>> 8351174 (fix(knowledge): harden Azure AI Search configuration and ingestion) + 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: {})") From 57574c3ff3d4a37ee24782938705496ce896e8b7 Mon Sep 17 00:00:00 2001 From: Harmke Alkemade Date: Tue, 7 Jul 2026 11:44:09 +0200 Subject: [PATCH 09/27] fix(knowledge): align shared backend configuration Signed-off-by: Harmke Alkemade --- .../customization/configuration-reference.md | 4 ++-- sources/knowledge_layer/KNOWLEDGE-LAYER-SETUP.md | 2 +- .../src/azure_ai_search/adapter.py | 4 ++-- sources/knowledge_layer/src/register.py | 16 ++++++---------- .../test_azure_ai_search.py | 16 ++++++++++++++++ 5 files changed, 27 insertions(+), 15 deletions(-) diff --git a/docs/source/customization/configuration-reference.md b/docs/source/customization/configuration-reference.md index fc1569a85..0d1b6bf1e 100644 --- a/docs/source/customization/configuration-reference.md +++ b/docs/source/customization/configuration-reference.md @@ -281,8 +281,8 @@ environment. `AZURE_SEARCH_API_KEY` is optional; when absent, the adapter uses | `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` | Namespace prefix for AI-Q-owned indexes. | | `embed_base_url` | `URL` | `AIQ_EMBED_BASE_URL` or NVIDIA API | OpenAI-compatible embedding base URL. | -| `embed_model` | `str` | `AIQ_EMBED_MODEL` or `nvidia/nv-embed-v1` | Embedding model used for Azure ingestion and retrieval. | -| `embed_dim` | `int` | `AIQ_EMBED_DIM` or `4096` | Embedding dimensions; must match the model and existing index schema. | +| `embed_model` | `str` | `AIQ_EMBED_MODEL` or `nvidia/llama-nemotron-embed-vl-1b-v2` | Embedding model used for Azure ingestion and retrieval. | +| `embed_dim` | `int` | `AIQ_EMBED_DIM` or `2048` | Embedding dimensions; must match the model and existing index schema. | | `embed_api_key` | `SecretStr` | `None` | Optional embedding API key; falls back to `NVIDIA_API_KEY`. | | `use_hybrid` | `bool` | `true` | Combine lexical and vector retrieval. | | `use_semantic_ranker` | `bool` | `false` | Apply Azure semantic ranking when supported; requires `use_hybrid: true`. | diff --git a/sources/knowledge_layer/KNOWLEDGE-LAYER-SETUP.md b/sources/knowledge_layer/KNOWLEDGE-LAYER-SETUP.md index 9ea42047b..bca738c67 100644 --- a/sources/knowledge_layer/KNOWLEDGE-LAYER-SETUP.md +++ b/sources/knowledge_layer/KNOWLEDGE-LAYER-SETUP.md @@ -1048,7 +1048,7 @@ Configuration values are resolved in the following order (highest to lowest prio | `AIQ_AZURE_SEARCH_INDEX_PREFIX` | azure_ai_search | Prefix for AI-Q-owned indexes (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: `4096`) | +| `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/src/azure_ai_search/adapter.py b/sources/knowledge_layer/src/azure_ai_search/adapter.py index 25a996194..bacfaa043 100644 --- a/sources/knowledge_layer/src/azure_ai_search/adapter.py +++ b/sources/knowledge_layer/src/azure_ai_search/adapter.py @@ -84,8 +84,8 @@ def _coerce_config(config: dict[str, Any] | None) -> SimpleNamespace: "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", "https://integrate.api.nvidia.com/v1"), - "embed_model": os.environ.get("AIQ_EMBED_MODEL", "nvidia/nv-embed-v1"), - "embed_dim": int(os.environ.get("AIQ_EMBED_DIM", "4096")), + "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")), "embed_api_key": None, "use_hybrid": True, "use_semantic_ranker": False, diff --git a/sources/knowledge_layer/src/register.py b/sources/knowledge_layer/src/register.py index 8322ad60e..42a65031d 100644 --- a/sources/knowledge_layer/src/register.py +++ b/sources/knowledge_layer/src/register.py @@ -279,16 +279,8 @@ class KnowledgeRetrievalConfig(FunctionBaseConfig, name="knowledge_retrieval"): min_length=1, description="Prefix for AI-Q-owned indexes; defaults to AIQ_AZURE_SEARCH_INDEX_PREFIX or aiq", ) - embed_base_url: HttpUrl = Field( - default_factory=lambda: _url_from_env("AIQ_EMBED_BASE_URL", "https://integrate.api.nvidia.com/v1"), - description="Embedding base URL; defaults to AIQ_EMBED_BASE_URL", - ) - embed_model: str = Field( - default_factory=lambda: os.environ.get("AIQ_EMBED_MODEL", "nvidia/nv-embed-v1"), - description="Embedding model; defaults to AIQ_EMBED_MODEL", - ) embed_dim: int = Field( - default_factory=lambda: int(os.environ.get("AIQ_EMBED_DIM", "4096")), + default_factory=lambda: int(os.environ.get("AIQ_EMBED_DIM", "2048")), gt=0, description="Embedding dimensions; defaults to AIQ_EMBED_DIM and must match existing indexes", ) @@ -450,6 +442,8 @@ def _setup_backend(config: KnowledgeRetrievalConfig, summary_llm_obj=None) -> tu "dask_file_transfer": config.opensearch_dask_file_transfer, "embed_model": config.embed_model, "embed_base_url": config.embed_base_url, + **summary_config, + } elif backend == "azure_ai_search": import knowledge_layer.azure_ai_search.adapter # noqa: F401 @@ -474,7 +468,9 @@ def _setup_backend(config: KnowledgeRetrievalConfig, summary_llm_obj=None) -> tu } else: - raise ValueError(f"Unknown backend: {backend}. Use 'llamaindex', 'foundational_rag', 'opensearch', or 'azure_ai_search'.") + 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 diff --git a/tests/knowledge_layer_tests/test_azure_ai_search.py b/tests/knowledge_layer_tests/test_azure_ai_search.py index 09da621ae..bc7a7cafd 100644 --- a/tests/knowledge_layer_tests/test_azure_ai_search.py +++ b/tests/knowledge_layer_tests/test_azure_ai_search.py @@ -264,6 +264,22 @@ def test_config_uses_shared_environment_defaults(monkeypatch): assert not adapter_config.use_semantic_ranker +def test_shared_embedding_defaults_match_adapter(monkeypatch): + 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 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_api_key_auth_requires_secret(): with pytest.raises(ValueError, match="azure_search_api_key"): KnowledgeRetrievalConfig( From 9fc164ff68ef84606d476f53eeaadf424a0dc0c7 Mon Sep 17 00:00:00 2001 From: Harmke Alkemade Date: Tue, 7 Jul 2026 11:55:16 +0200 Subject: [PATCH 10/27] fix(knowledge): use shared NVIDIA embedding credentials Signed-off-by: Harmke Alkemade --- docs/source/customization/configuration-reference.md | 1 - sources/knowledge_layer/src/azure_ai_search/adapter.py | 2 -- sources/knowledge_layer/src/register.py | 5 ----- tests/knowledge_layer_tests/test_azure_ai_search.py | 4 +--- 4 files changed, 1 insertion(+), 11 deletions(-) diff --git a/docs/source/customization/configuration-reference.md b/docs/source/customization/configuration-reference.md index 0d1b6bf1e..8e228e852 100644 --- a/docs/source/customization/configuration-reference.md +++ b/docs/source/customization/configuration-reference.md @@ -283,7 +283,6 @@ environment. `AZURE_SEARCH_API_KEY` is optional; when absent, the adapter uses | `embed_base_url` | `URL` | `AIQ_EMBED_BASE_URL` or NVIDIA API | OpenAI-compatible embedding base URL. | | `embed_model` | `str` | `AIQ_EMBED_MODEL` or `nvidia/llama-nemotron-embed-vl-1b-v2` | Embedding model used for Azure ingestion and retrieval. | | `embed_dim` | `int` | `AIQ_EMBED_DIM` or `2048` | Embedding dimensions; must match the model and existing index schema. | -| `embed_api_key` | `SecretStr` | `None` | Optional embedding API key; falls back to `NVIDIA_API_KEY`. | | `use_hybrid` | `bool` | `true` | Combine lexical and vector retrieval. | | `use_semantic_ranker` | `bool` | `false` | Apply Azure semantic ranking when supported; requires `use_hybrid: true`. | | `chunk_size` | `int` | `512` | Tokens per Azure-ingested chunk. | diff --git a/sources/knowledge_layer/src/azure_ai_search/adapter.py b/sources/knowledge_layer/src/azure_ai_search/adapter.py index bacfaa043..539b62a4e 100644 --- a/sources/knowledge_layer/src/azure_ai_search/adapter.py +++ b/sources/knowledge_layer/src/azure_ai_search/adapter.py @@ -86,7 +86,6 @@ def _coerce_config(config: dict[str, Any] | None) -> SimpleNamespace: "embed_base_url": os.environ.get("AIQ_EMBED_BASE_URL", "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")), - "embed_api_key": None, "use_hybrid": True, "use_semantic_ranker": False, "chunk_size": 512, @@ -419,7 +418,6 @@ def embedding(self): self._embedding = NVIDIAEmbedding( model=self.cfg.embed_model, base_url=str(self.cfg.embed_base_url), - api_key=_secret_value(self.cfg.embed_api_key), ) return self._embedding diff --git a/sources/knowledge_layer/src/register.py b/sources/knowledge_layer/src/register.py index 42a65031d..78801acca 100644 --- a/sources/knowledge_layer/src/register.py +++ b/sources/knowledge_layer/src/register.py @@ -284,10 +284,6 @@ class KnowledgeRetrievalConfig(FunctionBaseConfig, name="knowledge_retrieval"): gt=0, description="Embedding dimensions; defaults to AIQ_EMBED_DIM and must match existing indexes", ) - embed_api_key: SecretStr | None = Field( - default=None, - description="Optional embedding API key; NVIDIA_API_KEY is used when omitted", - ) use_hybrid: bool = Field(default=True, description="Combine vector and keyword search (azure_ai_search only)") use_semantic_ranker: bool = Field( default=False, @@ -456,7 +452,6 @@ def _setup_backend(config: KnowledgeRetrievalConfig, summary_llm_obj=None) -> tu "embed_base_url": str(config.embed_base_url), "embed_model": config.embed_model, "embed_dim": config.embed_dim, - "embed_api_key": config.embed_api_key, "use_hybrid": config.use_hybrid, "use_semantic_ranker": config.use_semantic_ranker, "chunk_size": config.chunk_size, diff --git a/tests/knowledge_layer_tests/test_azure_ai_search.py b/tests/knowledge_layer_tests/test_azure_ai_search.py index bc7a7cafd..f3378fd61 100644 --- a/tests/knowledge_layer_tests/test_azure_ai_search.py +++ b/tests/knowledge_layer_tests/test_azure_ai_search.py @@ -296,8 +296,7 @@ def test_setup_backend_preserves_secrets_and_prefix(monkeypatch): backend="azure_ai_search", azure_search_endpoint="https://example.search.windows.net", azure_search_auth_mode="api_key", - azure_search_api_key="test-search-key", - embed_api_key="test-embed-key", + azure_search_api_key="test-search-key", # pragma: allowlist secret azure_search_index_prefix="tenant-aiq", ) @@ -305,7 +304,6 @@ def test_setup_backend_preserves_secrets_and_prefix(monkeypatch): assert backend == "azure_ai_search" assert isinstance(backend_config["api_key"], SecretStr) - assert isinstance(backend_config["embed_api_key"], SecretStr) assert backend_config["index_prefix"] == "tenant-aiq" From 3f2f5f3b0e74c4a012b269efb9a1c85fce8c3d8d Mon Sep 17 00:00:00 2001 From: Harmke Alkemade Date: Tue, 7 Jul 2026 12:07:33 +0200 Subject: [PATCH 11/27] fix(knowledge): require Azure Search SDK 11.6 Signed-off-by: Harmke Alkemade --- sources/knowledge_layer/pyproject.toml | 2 +- uv.lock | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/sources/knowledge_layer/pyproject.toml b/sources/knowledge_layer/pyproject.toml index 435a84779..4823d87a5 100644 --- a/sources/knowledge_layer/pyproject.toml +++ b/sources/knowledge_layer/pyproject.toml @@ -63,7 +63,7 @@ opensearch = [ ] azure_ai_search = [ "azure-identity>=1.15.0,<1.26", - "azure-search-documents>=11.5.0,<11.7", + "azure-search-documents>=11.6.0,<11.7", "docx2txt>=0.8", "llama-index-core>=0.11.0", "llama-index-embeddings-nvidia>=0.2.0", diff --git a/uv.lock b/uv.lock index b40d34564..97bbc1257 100644 --- a/uv.lock +++ b/uv.lock @@ -2588,7 +2588,7 @@ 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.5.0,<11.7" }, + { 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" }, From 9453c871ff14b3f2e1cb1a73289090cbefcf74e9 Mon Sep 17 00:00:00 2001 From: Harmke Alkemade Date: Tue, 7 Jul 2026 12:53:18 +0200 Subject: [PATCH 12/27] test(knowledge): isolate Azure Search environment Signed-off-by: Harmke Alkemade --- tests/knowledge_layer_tests/test_azure_ai_search.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/tests/knowledge_layer_tests/test_azure_ai_search.py b/tests/knowledge_layer_tests/test_azure_ai_search.py index f3378fd61..46dbc8f53 100644 --- a/tests/knowledge_layer_tests/test_azure_ai_search.py +++ b/tests/knowledge_layer_tests/test_azure_ai_search.py @@ -221,7 +221,10 @@ def test_backend_registered_and_implements_sdk_contracts(): assert issubclass(AzureAISearchRetriever, BaseRetriever) -def test_config_requires_endpoint_and_valid_hybrid_semantic_combination(): +def test_config_requires_endpoint_and_valid_hybrid_semantic_combination(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( @@ -280,7 +283,9 @@ def test_shared_embedding_defaults_match_adapter(monkeypatch): assert adapter_config.embed_dim == config.embed_dim -def test_api_key_auth_requires_secret(): +def test_api_key_auth_requires_secret(monkeypatch): + monkeypatch.delenv("AZURE_SEARCH_API_KEY", raising=False) + with pytest.raises(ValueError, match="azure_search_api_key"): KnowledgeRetrievalConfig( backend="azure_ai_search", @@ -296,7 +301,7 @@ def test_setup_backend_preserves_secrets_and_prefix(monkeypatch): backend="azure_ai_search", azure_search_endpoint="https://example.search.windows.net", azure_search_auth_mode="api_key", - azure_search_api_key="test-search-key", # pragma: allowlist secret + azure_search_api_key="test-search-key", # pragma: allowlist secret azure_search_index_prefix="tenant-aiq", ) From 1e50fea332a5d56e1ce546c38c21ddb7d2bc4a2d Mon Sep 17 00:00:00 2001 From: Harmke Alkemade Date: Tue, 7 Jul 2026 13:07:28 +0200 Subject: [PATCH 13/27] fix(knowledge): preserve caller-owned upload files Signed-off-by: Harmke Alkemade --- .../src/azure_ai_search/adapter.py | 2 +- sources/knowledge_layer/src/register.py | 2 +- .../test_azure_ai_search.py | 23 +++++++++++++++++++ 3 files changed, 25 insertions(+), 2 deletions(-) diff --git a/sources/knowledge_layer/src/azure_ai_search/adapter.py b/sources/knowledge_layer/src/azure_ai_search/adapter.py index 539b62a4e..40e392826 100644 --- a/sources/knowledge_layer/src/azure_ai_search/adapter.py +++ b/sources/knowledge_layer/src/azure_ai_search/adapter.py @@ -92,7 +92,7 @@ def _coerce_config(config: dict[str, Any] | None) -> SimpleNamespace: "chunk_overlap": 64, "summary_max_chars": 1000, "collection_name": "default", - "cleanup_files": True, + "cleanup_files": False, "generate_summary": False, "summary_llm": None, "index_prefix": os.environ.get("AIQ_AZURE_SEARCH_INDEX_PREFIX", "aiq"), diff --git a/sources/knowledge_layer/src/register.py b/sources/knowledge_layer/src/register.py index 78801acca..f514ba9b3 100644 --- a/sources/knowledge_layer/src/register.py +++ b/sources/knowledge_layer/src/register.py @@ -458,7 +458,7 @@ def _setup_backend(config: KnowledgeRetrievalConfig, summary_llm_obj=None) -> tu "chunk_overlap": config.chunk_overlap, "summary_max_chars": config.summary_max_chars, "collection_name": config.collection_name, - "cleanup_files": True, + "cleanup_files": False, **summary_config, } diff --git a/tests/knowledge_layer_tests/test_azure_ai_search.py b/tests/knowledge_layer_tests/test_azure_ai_search.py index 46dbc8f53..b647f08d0 100644 --- a/tests/knowledge_layer_tests/test_azure_ai_search.py +++ b/tests/knowledge_layer_tests/test_azure_ai_search.py @@ -451,6 +451,29 @@ def start(self): assert ingestor._files[file_id].file_name == "original.txt" +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_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)) From 9ffa7882ce7331d581cb7868b211e67cdbd72285 Mon Sep 17 00:00:00 2001 From: Harmke Alkemade Date: Tue, 7 Jul 2026 13:22:03 +0200 Subject: [PATCH 14/27] fix(knowledge): handle empty embedding URL environment Signed-off-by: Harmke Alkemade --- sources/knowledge_layer/src/azure_ai_search/adapter.py | 2 +- sources/knowledge_layer/src/register.py | 2 +- tests/knowledge_layer_tests/test_azure_ai_search.py | 3 +++ 3 files changed, 5 insertions(+), 2 deletions(-) diff --git a/sources/knowledge_layer/src/azure_ai_search/adapter.py b/sources/knowledge_layer/src/azure_ai_search/adapter.py index 40e392826..e71a3c962 100644 --- a/sources/knowledge_layer/src/azure_ai_search/adapter.py +++ b/sources/knowledge_layer/src/azure_ai_search/adapter.py @@ -83,7 +83,7 @@ def _coerce_config(config: dict[str, Any] | None) -> SimpleNamespace: 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", "https://integrate.api.nvidia.com/v1"), + "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")), "use_hybrid": True, diff --git a/sources/knowledge_layer/src/register.py b/sources/knowledge_layer/src/register.py index f514ba9b3..3847bcaa5 100644 --- a/sources/knowledge_layer/src/register.py +++ b/sources/knowledge_layer/src/register.py @@ -40,7 +40,7 @@ def _url_from_env(name: str, default: str | None = None) -> HttpUrl | None: - value = os.environ.get(name, default) + value = os.environ.get(name) or default return HttpUrl(value) if value else None diff --git a/tests/knowledge_layer_tests/test_azure_ai_search.py b/tests/knowledge_layer_tests/test_azure_ai_search.py index b647f08d0..a30dccd9c 100644 --- a/tests/knowledge_layer_tests/test_azure_ai_search.py +++ b/tests/knowledge_layer_tests/test_azure_ai_search.py @@ -268,6 +268,7 @@ def test_config_uses_shared_environment_defaults(monkeypatch): 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) @@ -277,6 +278,8 @@ def test_shared_embedding_defaults_match_adapter(monkeypatch): ) 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 From a5114da3652dd430be0f2baac92f71de1b3e24e5 Mon Sep 17 00:00:00 2001 From: Harmke Alkemade Date: Tue, 7 Jul 2026 13:25:08 +0200 Subject: [PATCH 15/27] fix(knowledge): derive Azure auth mode from configured key Signed-off-by: Harmke Alkemade --- sources/knowledge_layer/src/register.py | 8 +++++--- tests/knowledge_layer_tests/test_azure_ai_search.py | 4 +++- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/sources/knowledge_layer/src/register.py b/sources/knowledge_layer/src/register.py index 3847bcaa5..9d9b6b474 100644 --- a/sources/knowledge_layer/src/register.py +++ b/sources/knowledge_layer/src/register.py @@ -270,9 +270,9 @@ class KnowledgeRetrievalConfig(FunctionBaseConfig, name="knowledge_retrieval"): 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_auth_mode: Literal["managed_identity", "api_key"] = Field( - default_factory=lambda: "api_key" if os.environ.get("AZURE_SEARCH_API_KEY") else "managed_identity", - description="Authentication mode; defaults to API key when AZURE_SEARCH_API_KEY is set", + azure_search_auth_mode: Literal["managed_identity", "api_key"] | None = Field( + default=None, + description="Authentication mode; derived from azure_search_api_key when omitted", ) azure_search_index_prefix: str = Field( default_factory=lambda: os.environ.get("AIQ_AZURE_SEARCH_INDEX_PREFIX", "aiq"), @@ -349,6 +349,8 @@ def validate_backend_config(self): elif backend == "azure_ai_search": if self.azure_search_endpoint is None: raise ValueError("azure_ai_search requires azure_search_endpoint") + if self.azure_search_auth_mode is None: + self.azure_search_auth_mode = "api_key" if self.azure_search_api_key is not None else "managed_identity" if self.azure_search_auth_mode == "api_key" and self.azure_search_api_key is None: raise ValueError("azure_search_auth_mode=api_key requires azure_search_api_key") if self.chunk_overlap >= self.chunk_size: diff --git a/tests/knowledge_layer_tests/test_azure_ai_search.py b/tests/knowledge_layer_tests/test_azure_ai_search.py index a30dccd9c..755428cc8 100644 --- a/tests/knowledge_layer_tests/test_azure_ai_search.py +++ b/tests/knowledge_layer_tests/test_azure_ai_search.py @@ -231,6 +231,7 @@ def test_config_requires_endpoint_and_valid_hybrid_semantic_combination(monkeypa backend="azure_ai_search", azure_search_endpoint="https://example.search.windows.net", ) + assert config.azure_search_auth_mode == "managed_identity" assert not config.use_semantic_ranker with pytest.raises(ValueError, match="use_semantic_ranker"): KnowledgeRetrievalConfig( @@ -303,7 +304,6 @@ def test_setup_backend_preserves_secrets_and_prefix(monkeypatch): config = KnowledgeRetrievalConfig( backend="azure_ai_search", azure_search_endpoint="https://example.search.windows.net", - azure_search_auth_mode="api_key", azure_search_api_key="test-search-key", # pragma: allowlist secret azure_search_index_prefix="tenant-aiq", ) @@ -311,6 +311,8 @@ def test_setup_backend_preserves_secrets_and_prefix(monkeypatch): backend, backend_config = _setup_backend(config) assert backend == "azure_ai_search" + assert config.azure_search_auth_mode == "api_key" + assert backend_config["auth_mode"] == "api_key" assert isinstance(backend_config["api_key"], SecretStr) assert backend_config["index_prefix"] == "tenant-aiq" From e160c127cb6eadf4c3a861e5ca55ee1cba74abe9 Mon Sep 17 00:00:00 2001 From: Harmke Alkemade Date: Tue, 7 Jul 2026 13:31:41 +0200 Subject: [PATCH 16/27] refactor(knowledge): infer Azure Search credentials Signed-off-by: Harmke Alkemade --- .../customization/configuration-reference.md | 1 - .../src/azure_ai_search/adapter.py | 14 ++--------- sources/knowledge_layer/src/register.py | 9 ------- .../test_azure_ai_search.py | 24 ++++++++----------- 4 files changed, 12 insertions(+), 36 deletions(-) diff --git a/docs/source/customization/configuration-reference.md b/docs/source/customization/configuration-reference.md index 8e228e852..cb2a74b85 100644 --- a/docs/source/customization/configuration-reference.md +++ b/docs/source/customization/configuration-reference.md @@ -277,7 +277,6 @@ environment. `AZURE_SEARCH_API_KEY` is optional; when absent, the adapter uses | `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_auth_mode` | `str` | Auto | Uses `api_key` when `AZURE_SEARCH_API_KEY` is set; otherwise `managed_identity`. | | `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` | Namespace prefix for AI-Q-owned indexes. | | `embed_base_url` | `URL` | `AIQ_EMBED_BASE_URL` or NVIDIA API | OpenAI-compatible embedding base URL. | diff --git a/sources/knowledge_layer/src/azure_ai_search/adapter.py b/sources/knowledge_layer/src/azure_ai_search/adapter.py index e71a3c962..b38d0feea 100644 --- a/sources/knowledge_layer/src/azure_ai_search/adapter.py +++ b/sources/knowledge_layer/src/azure_ai_search/adapter.py @@ -99,14 +99,8 @@ def _coerce_config(config: dict[str, Any] | None) -> SimpleNamespace: "start_ttl_cleanup": True, } values.update(provided) - values["auth_mode"] = provided.get( - "auth_mode", - "api_key" if values["api_key"] else "managed_identity", - ) if not values["endpoint"]: raise ValueError("Azure AI Search configuration requires `endpoint`") - if values["auth_mode"] not in {"managed_identity", "api_key"}: - raise ValueError("Azure AI Search auth_mode must be 'managed_identity' or 'api_key'") if values["chunk_overlap"] >= values["chunk_size"]: raise ValueError("chunk_overlap must be smaller than chunk_size") if not values["use_hybrid"] and values["use_semantic_ranker"]: @@ -116,12 +110,8 @@ def _coerce_config(config: dict[str, Any] | None) -> SimpleNamespace: def _build_search_credential(cfg: SimpleNamespace): """Pick the Azure SDK credential without exposing secret values.""" - if cfg.auth_mode == "api_key": - api_key = _secret_value(cfg.api_key) - if api_key is None: - raise ValueError("auth_mode=api_key requires the `api_key` field to be set") - return AzureKeyCredential(api_key) - return DefaultAzureCredential() + api_key = _secret_value(cfg.api_key) + return AzureKeyCredential(api_key) if api_key else DefaultAzureCredential() def _secret_value(value: Any) -> str | None: diff --git a/sources/knowledge_layer/src/register.py b/sources/knowledge_layer/src/register.py index 9d9b6b474..5eef5cf59 100644 --- a/sources/knowledge_layer/src/register.py +++ b/sources/knowledge_layer/src/register.py @@ -270,10 +270,6 @@ class KnowledgeRetrievalConfig(FunctionBaseConfig, name="knowledge_retrieval"): 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_auth_mode: Literal["managed_identity", "api_key"] | None = Field( - default=None, - description="Authentication mode; derived from azure_search_api_key when omitted", - ) azure_search_index_prefix: str = Field( default_factory=lambda: os.environ.get("AIQ_AZURE_SEARCH_INDEX_PREFIX", "aiq"), min_length=1, @@ -349,10 +345,6 @@ def validate_backend_config(self): elif backend == "azure_ai_search": if self.azure_search_endpoint is None: raise ValueError("azure_ai_search requires azure_search_endpoint") - if self.azure_search_auth_mode is None: - self.azure_search_auth_mode = "api_key" if self.azure_search_api_key is not None else "managed_identity" - if self.azure_search_auth_mode == "api_key" and self.azure_search_api_key is None: - raise ValueError("azure_search_auth_mode=api_key requires azure_search_api_key") if self.chunk_overlap >= self.chunk_size: raise ValueError("chunk_overlap must be smaller than chunk_size") if not self.use_hybrid and self.use_semantic_ranker: @@ -449,7 +441,6 @@ def _setup_backend(config: KnowledgeRetrievalConfig, summary_llm_obj=None) -> tu backend_config = { "endpoint": str(config.azure_search_endpoint), "api_key": config.azure_search_api_key, - "auth_mode": config.azure_search_auth_mode, "index_prefix": config.azure_search_index_prefix, "embed_base_url": str(config.embed_base_url), "embed_model": config.embed_model, diff --git a/tests/knowledge_layer_tests/test_azure_ai_search.py b/tests/knowledge_layer_tests/test_azure_ai_search.py index 755428cc8..bbc7414b1 100644 --- a/tests/knowledge_layer_tests/test_azure_ai_search.py +++ b/tests/knowledge_layer_tests/test_azure_ai_search.py @@ -10,8 +10,10 @@ 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 @@ -181,7 +183,6 @@ def get_nodes_from_documents(self, documents): def _config(**overrides): config = { "endpoint": "https://example.search.windows.net", - "auth_mode": "api_key", "api_key": SecretStr("test-key"), "embed_model": "test-embed", "embed_dim": 4, @@ -231,7 +232,7 @@ def test_config_requires_endpoint_and_valid_hybrid_semantic_combination(monkeypa backend="azure_ai_search", azure_search_endpoint="https://example.search.windows.net", ) - assert config.azure_search_auth_mode == "managed_identity" + assert config.azure_search_api_key is None assert not config.use_semantic_ranker with pytest.raises(ValueError, match="use_semantic_ranker"): KnowledgeRetrievalConfig( @@ -256,7 +257,6 @@ def test_config_uses_shared_environment_defaults(monkeypatch): assert backend == "azure_ai_search" assert backend_config["endpoint"] == "https://env.search.windows.net/" - assert backend_config["auth_mode"] == "api_key" 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" @@ -264,7 +264,6 @@ def test_config_uses_shared_environment_defaults(monkeypatch): assert backend_config["embed_dim"] == 8 assert not backend_config["use_semantic_ranker"] assert adapter_config.endpoint == "https://env.search.windows.net" - assert adapter_config.auth_mode == "api_key" assert not adapter_config.use_semantic_ranker @@ -287,15 +286,14 @@ def test_shared_embedding_defaults_match_adapter(monkeypatch): assert adapter_config.embed_dim == config.embed_dim -def test_api_key_auth_requires_secret(monkeypatch): - monkeypatch.delenv("AZURE_SEARCH_API_KEY", raising=False) +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}) - with pytest.raises(ValueError, match="azure_search_api_key"): - KnowledgeRetrievalConfig( - backend="azure_ai_search", - azure_search_endpoint="https://example.search.windows.net", - azure_search_auth_mode="api_key", - ) + 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): @@ -311,8 +309,6 @@ def test_setup_backend_preserves_secrets_and_prefix(monkeypatch): backend, backend_config = _setup_backend(config) assert backend == "azure_ai_search" - assert config.azure_search_auth_mode == "api_key" - assert backend_config["auth_mode"] == "api_key" assert isinstance(backend_config["api_key"], SecretStr) assert backend_config["index_prefix"] == "tenant-aiq" From 8f9f19ea742f1bcd0ef2adeef4c08101528a5947 Mon Sep 17 00:00:00 2001 From: Harmke Alkemade Date: Tue, 7 Jul 2026 13:37:38 +0200 Subject: [PATCH 17/27] docs(knowledge): clarify Azure Search setup Signed-off-by: Harmke Alkemade --- docs/source/customization/knowledge-layer.md | 1 + docs/source/examples/azure-ai-search.md | 8 ++++++++ 2 files changed, 9 insertions(+) diff --git a/docs/source/customization/knowledge-layer.md b/docs/source/customization/knowledge-layer.md index 07b2d8508..68bc66a2d 100644 --- a/docs/source/customization/knowledge-layer.md +++ b/docs/source/customization/knowledge-layer.md @@ -157,6 +157,7 @@ functions: ``` **Azure AI Search (Managed Service)** + ```yaml functions: knowledge_search: diff --git a/docs/source/examples/azure-ai-search.md b/docs/source/examples/azure-ai-search.md index 191778973..f5ad2f8ca 100644 --- a/docs/source/examples/azure-ai-search.md +++ b/docs/source/examples/azure-ai-search.md @@ -10,6 +10,14 @@ 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 [index limit](https://learn.microsoft.com/azure/search/search-limits-quotas-capacity) covers + the expected number of logical collections. AI-Q creates one physical index per collection. +- 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 From 4c546fc21faced6e0ec94131d4d913b67fb0552a Mon Sep 17 00:00:00 2001 From: Harmke Alkemade Date: Tue, 7 Jul 2026 13:45:42 +0200 Subject: [PATCH 18/27] fix(knowledge): reject raw Azure OData filters Signed-off-by: Harmke Alkemade --- .../knowledge_layer/src/azure_ai_search/adapter.py | 11 ++++++++--- .../knowledge_layer_tests/test_azure_ai_search.py | 14 ++++++++++++-- 2 files changed, 20 insertions(+), 5 deletions(-) diff --git a/sources/knowledge_layer/src/azure_ai_search/adapter.py b/sources/knowledge_layer/src/azure_ai_search/adapter.py index b38d0feea..2088f8733 100644 --- a/sources/knowledge_layer/src/azure_ai_search/adapter.py +++ b/sources/knowledge_layer/src/azure_ai_search/adapter.py @@ -459,6 +459,14 @@ def _retrieve_sync( top_k: int, filters: dict[str, Any] | None, ) -> RetrievalResult: + if filters and isinstance(filters, dict) and filters.get("$filter"): + return RetrievalResult( + query=query, + backend=_BACKEND_NAME, + chunks=[], + success=False, + error_message="Raw Azure AI Search $filter expressions are not supported", + ) try: query_vector = self.embedding.get_query_embedding(query) client = self._get_client(collection_name) @@ -477,9 +485,6 @@ def _retrieve_sync( if self.cfg.use_semantic_ranker: search_params["query_type"] = "semantic" search_params["semantic_configuration_name"] = _SEMANTIC_CONFIG - if filters and isinstance(filters, dict) and (odata_filter := filters.get("$filter")): - search_params["filter"] = odata_filter - chunks = [self.normalize(hit) for hit in client.search(**search_params)] return RetrievalResult(query=query, backend=_BACKEND_NAME, chunks=chunks, success=True) except ResourceNotFoundError: diff --git a/tests/knowledge_layer_tests/test_azure_ai_search.py b/tests/knowledge_layer_tests/test_azure_ai_search.py index bbc7414b1..974818be5 100644 --- a/tests/knowledge_layer_tests/test_azure_ai_search.py +++ b/tests/knowledge_layer_tests/test_azure_ai_search.py @@ -422,15 +422,25 @@ def search(self, **kwargs): retriever._embedding = FakeEmbedding(2) retriever._get_client = lambda collection_name: client - result = await retriever.retrieve("hello", "session-1", top_k=5, filters={"$filter": "file_id eq 'file-1'"}) + result = await retriever.retrieve("hello", "session-1", top_k=5) assert result.success assert client.kwargs["search_text"] == "hello" assert client.kwargs["query_type"] == "semantic" - assert client.kwargs["filter"] == "file_id eq 'file-1'" assert client.kwargs["select"] == ["id", "chunk", "file_id", "file_name", "page_number", "metadata"] +@pytest.mark.asyncio +async def test_retrieve_rejects_raw_odata_filter(): + 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={"$filter": "file_id ne ''"}) + + assert not result.success + assert result.error_message == "Raw Azure AI Search $filter expressions are not supported" + + def test_submit_job_uses_one_canonical_file_id(monkeypatch, tmp_path): class NoStartThread: def __init__(self, *args, **kwargs): From 94c8bf893de36f95f96779d1a50531f19e528d77 Mon Sep 17 00:00:00 2001 From: Harmke Alkemade Date: Tue, 7 Jul 2026 15:12:59 +0200 Subject: [PATCH 19/27] refactor(knowledge): share Azure Search index across collections Signed-off-by: Harmke Alkemade --- .secrets.baseline | 4 +- .../customization/configuration-reference.md | 17 +- docs/source/customization/knowledge-layer.md | 31 +- docs/source/examples/azure-ai-search.md | 45 +- docs/source/examples/index.md | 2 +- .../knowledge_layer/KNOWLEDGE-LAYER-SETUP.md | 13 +- sources/knowledge_layer/README.md | 2 +- .../src/azure_ai_search/README.md | 40 +- .../src/azure_ai_search/adapter.py | 751 ++++++++++-------- sources/knowledge_layer/src/register.py | 33 +- .../test_azure_ai_search.py | 468 +++++++---- 11 files changed, 841 insertions(+), 565 deletions(-) diff --git a/.secrets.baseline b/.secrets.baseline index a398b41d2..46682e4d3 100644 --- a/.secrets.baseline +++ b/.secrets.baseline @@ -241,7 +241,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": [ @@ -290,5 +290,5 @@ } ] }, - "generated_at": "2026-07-02T22:31:04Z" + "generated_at": "2026-07-07T12:40:18Z" } diff --git a/docs/source/customization/configuration-reference.md b/docs/source/customization/configuration-reference.md index cb2a74b85..92a8b3964 100644 --- a/docs/source/customization/configuration-reference.md +++ b/docs/source/customization/configuration-reference.md @@ -217,8 +217,8 @@ functions: ### `knowledge_retrieval` -Semantic search over ingested documents. Supports LlamaIndex (local ChromaDB), Foundational RAG (hosted -NVIDIA RAG Blueprint), and Azure AI Search. +Semantic search over ingested documents. Supports LlamaIndex (local ChromaDB), OpenSearch, Foundational RAG +(hosted NVIDIA RAG Blueprint), and Azure AI Search. ```yaml functions: @@ -255,8 +255,6 @@ functions: _type: knowledge_retrieval backend: azure_ai_search collection_name: ${COLLECTION_NAME:-test_collection} - use_hybrid: true - use_semantic_ranker: false ``` This example reads `AZURE_SEARCH_ENDPOINT` and `NVIDIA_API_KEY` from the @@ -265,8 +263,8 @@ environment. `AZURE_SEARCH_API_KEY` is optional; when absent, the adapter uses | Parameter | Type | Default | Description | |-----------|------|---------|-------------| -| `backend` | `str` | `llamaindex` | Backend type: `llamaindex`, `foundational_rag`, or `azure_ai_search`. | -| `collection_name` | `str` | `default` | Name of the document collection/index. | +| `backend` | `str` | `llamaindex` | Backend type: `llamaindex`, `opensearch`, `foundational_rag`, or `azure_ai_search`. | +| `collection_name` | `str` | `default` | Name of the logical document collection. | | `top_k` | `int` | `5` | Number of results to return per query. | | `generate_summary` | `bool` | `false` | Generate one-sentence summaries for ingested documents. | | `summary_model` | `str` | `None` | LLM reference from `llms` section. Required when `generate_summary: true`. | @@ -278,15 +276,10 @@ environment. `AZURE_SEARCH_API_KEY` is optional; when absent, the adapter uses | `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` | Namespace prefix for AI-Q-owned indexes. | +| `azure_search_index_prefix` | `str` | `AIQ_AZURE_SEARCH_INDEX_PREFIX` or `aiq` | Deployment-unique namespace for the shared AI-Q index. | | `embed_base_url` | `URL` | `AIQ_EMBED_BASE_URL` or NVIDIA API | OpenAI-compatible embedding base URL. | | `embed_model` | `str` | `AIQ_EMBED_MODEL` or `nvidia/llama-nemotron-embed-vl-1b-v2` | Embedding model used for Azure ingestion and retrieval. | | `embed_dim` | `int` | `AIQ_EMBED_DIM` or `2048` | Embedding dimensions; must match the model and existing index schema. | -| `use_hybrid` | `bool` | `true` | Combine lexical and vector retrieval. | -| `use_semantic_ranker` | `bool` | `false` | Apply Azure semantic ranking when supported; requires `use_hybrid: true`. | -| `chunk_size` | `int` | `512` | Tokens per Azure-ingested chunk. | -| `chunk_overlap` | `int` | `64` | Token overlap; must be smaller than `chunk_size`. | -| `summary_max_chars` | `int` | `1000` | Maximum document characters sent to the summary model. | ### `intent_classifier` diff --git a/docs/source/customization/knowledge-layer.md b/docs/source/customization/knowledge-layer.md index 68bc66a2d..75a4e4e00 100644 --- a/docs/source/customization/knowledge-layer.md +++ b/docs/source/customization/knowledge-layer.md @@ -43,7 +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 | Hybrid and semantic retrieval | +| `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. @@ -164,8 +164,6 @@ functions: _type: knowledge_retrieval backend: azure_ai_search collection_name: my_docs - use_hybrid: true - use_semantic_ranker: false ``` Set `AZURE_SEARCH_ENDPOINT` and `NVIDIA_API_KEY` in the environment. Setting @@ -174,19 +172,20 @@ Set `AZURE_SEARCH_ENDPOINT` and `NVIDIA_API_KEY` in the environment. Setting 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. - -Semantic ranking is disabled by default because support depends on the Azure AI -Search service. Set `use_semantic_ranker: true` only when semantic ranking is -enabled; it also requires `use_hybrid: true`. - -Azure maps logical collection names to collision-safe physical indexes under `azure_search_index_prefix`. Only -indexes containing an AI-Q ownership and schema marker are listed, queried, or deleted. Legacy indexes named directly -after collections are ignored; re-ingest those collections after enabling this backend. - -Upload responses return canonical UUID file IDs. Re-uploading the same filename writes and verifies the new generation -before removing the old generation. 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. +`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. Schema-version-1 +indexes are ignored and require re-ingestion. + +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. #### Multimodal Extraction (LlamaIndex Only) diff --git a/docs/source/examples/azure-ai-search.md b/docs/source/examples/azure-ai-search.md index f5ad2f8ca..ce975b1c6 100644 --- a/docs/source/examples/azure-ai-search.md +++ b/docs/source/examples/azure-ai-search.md @@ -13,8 +13,8 @@ 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 [index limit](https://learn.microsoft.com/azure/search/search-limits-quotas-capacity) covers - the expected number of logical collections. AI-Q creates one physical index per collection. +- 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. @@ -40,12 +40,11 @@ Search service and grant the workload identity both of these built-in roles: | Role | Used for | |------|----------| -| `Search Service Contributor` | Create, inspect, update, and delete AI-Q collection indexes. | +| `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 because AI-Q creates one index per -logical collection. The principal ID is the object ID of the system-assigned or -user-assigned managed identity running AI-Q. +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. Replace the `knowledge_search` block in a web configuration such as @@ -58,8 +57,6 @@ functions: backend: azure_ai_search collection_name: ${COLLECTION_NAME:-aiq_default} top_k: 5 - use_hybrid: true - use_semantic_ranker: false generate_summary: true summary_model: summary_llm @@ -71,20 +68,18 @@ 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`. - -Semantic ranking is opt-in because availability depends on the Azure AI Search -service. Set `use_semantic_ranker: true` only when semantic ranking is enabled; -it also requires `use_hybrid: true`. - -Existing indexes must use the configured `embed_dim`. Delete and re-ingest a -collection when changing embedding dimensions. 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 only lists or mutates indexes carrying its AI-Q ownership marker. -Logical collection names map to collision-safe physical names under -`azure_search_index_prefix`; un-namespaced indexes from earlier versions are -ignored and must be re-ingested. File IDs returned by upload are authoritative -for status and delete operations. Same-name uploads replace the prior file only -after the new generation has been fully indexed. +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. Schema-version-1 indexes are ignored and must be re-ingested. 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 1360fe58d..65058cb9a 100644 --- a/docs/source/examples/index.md +++ b/docs/source/examples/index.md @@ -12,7 +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 and semantic document retrieval | `config_web_default_llamaindex.yml` base | +| [Azure AI Search Knowledge Layer](./azure-ai-search.md) | Managed hybrid document retrieval | `config_web_default_llamaindex.yml` base | | [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 Modal sandbox execution for quantitative research workflows | `config_domain_routing_and_skills.yml` | diff --git a/sources/knowledge_layer/KNOWLEDGE-LAYER-SETUP.md b/sources/knowledge_layer/KNOWLEDGE-LAYER-SETUP.md index bca738c67..ae7ece23d 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,7 +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 | Hybrid and semantic retrieval | +| `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. @@ -44,9 +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 AI-Q-owned, collision-safe indexes in a managed Azure AI Search service. Canonical - UUID file IDs support status and deletion, while same-name uploads replace the prior generation after successful - indexing. See [`src/azure_ai_search/README.md`](src/azure_ai_search/README.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, @@ -188,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: @@ -1045,7 +1046,7 @@ Configuration values are resolved in the following order (highest to lowest prio | `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` | -| `AIQ_AZURE_SEARCH_INDEX_PREFIX` | azure_ai_search | Prefix for AI-Q-owned indexes (default: `aiq`) | +| `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`) | diff --git a/sources/knowledge_layer/README.md b/sources/knowledge_layer/README.md index e1812a239..cc2dd1542 100644 --- a/sources/knowledge_layer/README.md +++ b/sources/knowledge_layer/README.md @@ -27,7 +27,7 @@ uv pip install -e "sources/knowledge_layer[azure_ai_search]" | `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 and semantic search | +| `azure_ai_search` | Azure AI Search | Managed hybrid search | ## Usage diff --git a/sources/knowledge_layer/src/azure_ai_search/README.md b/sources/knowledge_layer/src/azure_ai_search/README.md index 7aff59ff8..766c15d31 100644 --- a/sources/knowledge_layer/src/azure_ai_search/README.md +++ b/sources/knowledge_layer/src/azure_ai_search/README.md @@ -32,11 +32,7 @@ functions: _type: knowledge_retrieval backend: azure_ai_search collection_name: ${COLLECTION_NAME:-aiq_default} - use_hybrid: true - use_semantic_ranker: false top_k: 5 - chunk_size: 512 - chunk_overlap: 64 generate_summary: true summary_model: summary_llm @@ -47,44 +43,38 @@ 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`. - -Semantic ranking is disabled by default because support depends on the Azure AI -Search service. Set `use_semantic_ranker: true` only when semantic ranking is -enabled; it also requires `use_hybrid: true`. +`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 AI-Q creates one index per collection. +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 AI-Q collection, and performs -vector or hybrid retrieval with optional semantic ranking. Logical collection -names are sanitized and combined with a stable hash, preventing collisions and -protecting unrelated indexes in a shared service. Only indexes containing the -AI-Q ownership/schema marker are visible or mutable through this backend. +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. Re-uploading the same filename writes and -verifies a new generation before deleting old chunks. Upload and delete -requests stay below Azure's 1,000-action and 16 MiB limits, and every -per-document result is checked. +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 any existing index. +`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 requires deleting and re-ingesting the old collection. The adapter -validates ownership, fields, vector profile, dimensions, and semantic -configuration before use; it does not alter an incompatible schema. Indexes -created by the earlier un-namespaced implementation are deliberately ignored, -so re-ingest those collections. +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 diff --git a/sources/knowledge_layer/src/azure_ai_search/adapter.py b/sources/knowledge_layer/src/azure_ai_search/adapter.py index 2088f8733..7df83f345 100644 --- a/sources/knowledge_layer/src/azure_ai_search/adapter.py +++ b/sources/knowledge_layer/src/azure_ai_search/adapter.py @@ -9,6 +9,7 @@ import os import re import threading +import time import uuid from collections.abc import Iterator from datetime import UTC @@ -17,11 +18,9 @@ from types import SimpleNamespace from typing import Any -from azure.core import MatchConditions from azure.core.credentials import AzureKeyCredential from azure.core.exceptions import ClientAuthenticationError from azure.core.exceptions import HttpResponseError -from azure.core.exceptions import ResourceModifiedError from azure.core.exceptions import ResourceNotFoundError from azure.core.exceptions import ServiceRequestError from azure.identity import DefaultAzureCredential @@ -33,10 +32,6 @@ 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 SemanticConfiguration -from azure.search.documents.indexes.models import SemanticField -from azure.search.documents.indexes.models import SemanticPrioritizedFields -from azure.search.documents.indexes.models import SemanticSearch from azure.search.documents.indexes.models import SimpleField from azure.search.documents.indexes.models import VectorSearch from azure.search.documents.indexes.models import VectorSearchAlgorithmMetric @@ -64,14 +59,24 @@ logger = logging.getLogger(__name__) _BACKEND_NAME = "azure_ai_search" -_SEMANTIC_CONFIG = "default-semantic" _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")) @@ -79,18 +84,12 @@ def _coerce_config(config: dict[str, Any] | None) -> SimpleNamespace: """Apply adapter defaults so direct factory usage matches YAML usage.""" - provided = config or {} 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")), - "use_hybrid": True, - "use_semantic_ranker": False, - "chunk_size": 512, - "chunk_overlap": 64, - "summary_max_chars": 1000, "collection_name": "default", "cleanup_files": False, "generate_summary": False, @@ -98,13 +97,9 @@ def _coerce_config(config: dict[str, Any] | None) -> SimpleNamespace: "index_prefix": os.environ.get("AIQ_AZURE_SEARCH_INDEX_PREFIX", "aiq"), "start_ttl_cleanup": True, } - values.update(provided) + 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`") - if values["chunk_overlap"] >= values["chunk_size"]: - raise ValueError("chunk_overlap must be smaller than chunk_size") - if not values["use_hybrid"] and values["use_semantic_ranker"]: - raise ValueError("use_semantic_ranker=true requires use_hybrid=true") return SimpleNamespace(**values) @@ -143,14 +138,16 @@ def _sanitize_index_part(value: str, fallback: str = "default") -> str: return normalized or fallback -def _index_name_for_collection(collection_name: str, prefix: str = "aiq") -> str: - """Map a logical collection to an official, collision-safe Azure index name.""" - prefix_part = _sanitize_index_part(prefix, "aiq")[:48].rstrip("-") or "aiq" - collection_part = _sanitize_index_part(collection_name) - suffix = uuid.uuid5(uuid.NAMESPACE_URL, f"{prefix}\0{collection_name}").hex[:12] - available = _MAX_INDEX_NAME_LENGTH - len(prefix_part) - len(suffix) - 2 - collection_part = collection_part[:available].rstrip("-") or "default" - return f"{prefix_part}-{collection_part}-{suffix}" +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: @@ -162,7 +159,10 @@ def _validate_index_name(name: str) -> None: def _encode_marker(marker: dict[str, Any]) -> str: - return _MARKER_PREFIX + json.dumps(marker, separators=(",", ":"), sort_keys=True) + 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: @@ -175,28 +175,24 @@ def _decode_marker(description: str | None) -> dict[str, Any] | None: return marker if isinstance(marker, dict) else None -def _new_marker( - collection_name: str, - cfg: SimpleNamespace, - description: str | None = None, - metadata: dict[str, Any] | None = None, -) -> dict[str, Any]: - now = _utc_now().isoformat() +def _new_marker(cfg: SimpleNamespace) -> dict[str, Any]: marker = { "backend": _BACKEND_NAME, "schema_version": _SCHEMA_VERSION, - "collection_name": collection_name, - "description": description, - "metadata": metadata or {}, + "index_prefix": cfg.index_prefix, "embedding_model": cfg.embed_model, "embedding_dim": cfg.embed_dim, - "created_at": now, - "updated_at": now, + "created_at": _utc_now().isoformat(), } - _encode_marker(marker) # Validate JSON serializability before any service mutation. + _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("'", "''") + "'" @@ -206,6 +202,26 @@ def _and_filter(*filters: str | None) -> str | None: 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) @@ -279,6 +295,14 @@ def _build_index_schema(name: str, embed_dim: int, description: str | None = Non 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", @@ -289,9 +313,16 @@ def _build_index_schema(name: str, embed_dim: int, description: str | None = Non ), 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), @@ -310,48 +341,41 @@ def _build_index_schema(name: str, embed_dim: int, description: str | None = Non ], profiles=[VectorSearchProfile(name="hnsw-profile", algorithm_configuration_name="hnsw-default")], ), - semantic_search=SemanticSearch( - default_configuration_name=_SEMANTIC_CONFIG, - configurations=[ - SemanticConfiguration( - name=_SEMANTIC_CONFIG, - prioritized_fields=SemanticPrioritizedFields( - title_field=SemanticField(field_name="file_name"), - content_fields=[SemanticField(field_name="chunk")], - ), - ), - ], - ), ) -def _validate_index_schema(index: SearchIndex, collection_name: str, cfg: SimpleNamespace) -> dict[str, Any]: +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("collection_name") != collection_name: - raise RuntimeError( - f"Azure AI Search index {index.name!r} belongs to collection {marker.get('collection_name')!r}, " - f"not {collection_name!r}" - ) - if 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} embedding configuration does not match " - f"{cfg.embed_model!r}/{cfg.embed_dim}" - ) + 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, @@ -366,23 +390,16 @@ def _validate_index_schema(index: SearchIndex, collection_name: str, cfg: Simple ) 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["file_id"].filterable or not fields["file_name"].filterable: - raise RuntimeError(f"Azure AI Search index {index.name!r} requires filterable file identity fields") + 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") - semantic_names = { - semantic.name for semantic in (index.semantic_search.configurations if index.semantic_search else []) - } - if cfg.use_semantic_ranker and ( - index.semantic_search is None - or index.semantic_search.default_configuration_name != _SEMANTIC_CONFIG - or _SEMANTIC_CONFIG not in semantic_names - ): - raise RuntimeError(f"Azure AI Search index {index.name!r} semantic configuration does not match") return marker @@ -391,13 +408,15 @@ class _AzureIndexMixin: _credential: Any _embedding: Any _index_client: SearchIndexClient - _search_clients: dict[str, SearchClient] + _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_clients = {} + self._search_client = None + self._index_validated = False self._embedding = None @property @@ -411,37 +430,74 @@ def embedding(self): ) return self._embedding - def _physical_index_name(self, collection_name: str) -> str: - name = _index_name_for_collection(collection_name, self.cfg.index_prefix) + 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, collection_name: str) -> SearchClient: - if collection_name not in self._search_clients: - self._search_clients[collection_name] = self._index_client.get_search_client( - self._physical_index_name(collection_name) - ) - return self._search_clients[collection_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 and semantic-ranked retriever backed by owned Azure indexes.""" + """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) - self._validated: set[str] = set() def _get_client(self, collection_name: str) -> SearchClient: - index_name = self._physical_index_name(collection_name) - if collection_name not in self._validated: - index = self._index_client.get_index(index_name) - _validate_index_schema(index, collection_name, self.cfg) - self._validated.add(collection_name) - return self._get_search_client(collection_name) + 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, @@ -468,23 +524,20 @@ def _retrieve_sync( error_message="Raw Azure AI Search $filter expressions are not supported", ) try: - query_vector = self.embedding.get_query_embedding(query) 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"], } - if self.cfg.use_hybrid: - search_params["search_text"] = query - if self.cfg.use_semantic_ranker: - search_params["query_type"] = "semantic" - search_params["semantic_configuration_name"] = _SEMANTIC_CONFIG chunks = [self.normalize(hit) for hit in client.search(**search_params)] return RetrievalResult(query=query, backend=_BACKEND_NAME, chunks=chunks, success=True) except ResourceNotFoundError: @@ -505,9 +558,8 @@ def normalize(self, raw_result: Any) -> Chunk: 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")) - rerank_score = raw_result.get("@search.reranker_score") search_score = raw_result.get("@search.score") or 0.0 - score = float(rerank_score) / 4.0 if rerank_score is not None else float(search_score) + 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"): @@ -549,6 +601,8 @@ def __init__(self, config: dict[str, Any] | None = None): 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) @@ -557,51 +611,142 @@ def splitter(self): if self._splitter is None: from llama_index.core.node_parser import SentenceSplitter - self._splitter = SentenceSplitter(chunk_size=self.cfg.chunk_size, chunk_overlap=self.cfg.chunk_overlap) + self._splitter = SentenceSplitter(chunk_size=_CHUNK_SIZE, chunk_overlap=_CHUNK_OVERLAP) return self._splitter - def _get_owned_index(self, collection_name: str) -> tuple[SearchIndex, dict[str, Any]]: - index = self._index_client.get_index(self._physical_index_name(collection_name)) - return index, _validate_index_schema(index, collection_name, self.cfg) + def _write_document(self, document: dict[str, Any]) -> None: + self._upload_documents(self._get_validated_search_client(), [document]) - def _ensure_index( + 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, - ) -> tuple[SearchIndex, dict[str, Any]]: - index_name = self._physical_index_name(collection_name) - try: - return self._get_owned_index(collection_name) - except ResourceNotFoundError: - pass + *, + 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 - marker = _new_marker(collection_name, self.cfg, description, metadata) - schema = _build_index_schema(index_name, self.cfg.embed_dim, _encode_marker(marker)) + @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: - index = self._index_client.create_index(schema) - except Exception as create_error: # noqa: BLE001 - try: - index = self._index_client.get_index(index_name) - except ResourceNotFoundError: - raise create_error - return index, _validate_index_schema(index, collection_name, self.cfg) - - def _update_marker(self, collection_name: str, **updates: Any) -> dict[str, Any]: - for attempt in range(3): - index, marker = self._get_owned_index(collection_name) - marker.update(updates) - index.description = _encode_marker(marker) - try: - self._index_client.create_or_update_index(index, match_condition=MatchConditions.IfNotModified) - return marker - except (ResourceModifiedError, HttpResponseError) as error: - if getattr(error, "status_code", None) != 412 or attempt == 2: - raise - raise RuntimeError(f"Failed to update metadata for collection {collection_name!r}") + 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 _update_collection_timestamp(self, collection_name: str) -> None: - self._update_marker(collection_name, updated_at=_utc_now().isoformat()) + 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, @@ -614,7 +759,7 @@ def submit_job( 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 {} - _encode_marker({"metadata": file_metadata}) + json.dumps(file_metadata, separators=(",", ":"), sort_keys=True) if not validated: with self._jobs_lock: @@ -630,21 +775,30 @@ def submit_job( ) return job_id + self._ensure_index() + 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") + submitted_at = _utc_now() file_progress: list[FileProgress] = [] 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)) + info = 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[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}, - ) + self._files[file_id] = info + self._write_file_manifest(info) with self._jobs_lock: self._jobs[job_id] = IngestionJobStatus( @@ -689,10 +843,9 @@ def _process_job( ) -> None: cleanup = bool(config.get("cleanup_files", self.cfg.cleanup_files)) self._update_job(job_id, status=JobState.PROCESSING, started_at=_utc_now()) - try: - self._ensure_index(collection_name) - except Exception as error: # noqa: BLE001 - self._fail_job(job_id, f"Failed to ensure collection {collection_name!r}: {error!s}") + collection = self._get_collection_manifest(collection_name) + if collection is None or collection.get("status") != _COLLECTION_ACTIVE: + self._fail_job(job_id, f"Collection {collection_name!r} is unavailable") if cleanup: self._cleanup_paths(file_paths) return @@ -704,7 +857,7 @@ def _process_job( tracked = self._files[detail.file_id] self._update_file_progress(job_id, index, status=FileStatus.INGESTING) try: - chunk_count = self._process_file( + chunk_count, summary, ingested_at = self._process_file( path=path, collection_name=collection_name, file_id=detail.file_id, @@ -719,11 +872,24 @@ def _process_job( 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) - self._update_file_progress(job_id, index, status=FileStatus.FAILED, error_message=message) + 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: @@ -745,10 +911,9 @@ def _process_file( file_size: int, uploaded_at: datetime, metadata: dict[str, Any], - ) -> int: + ) -> tuple[int, str | None, datetime]: from llama_index.core import SimpleDirectoryReader - old_file_ids = self._find_file_ids_by_name(file_name, collection_name) - {file_id} documents = SimpleDirectoryReader(input_files=[path]).load_data() if not documents: raise ValueError(f"No content extracted from {file_name}") @@ -766,7 +931,9 @@ def _process_file( for chunk_index, (node, vector) in enumerate(zip(nodes, embeddings, strict=True)): search_documents.append( { - "id": f"{file_id}-{chunk_index:08d}", + "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, @@ -780,29 +947,16 @@ def _process_file( } ) - client = self._get_search_client(collection_name) + client = self._get_validated_search_client() self._upload_documents(client, search_documents) - self._update_collection_timestamp(collection_name) - for old_file_id in old_file_ids: - self._delete_file_documents(old_file_id, collection_name) - with self._jobs_lock: - self._files.pop(old_file_id, None) + 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 - if summary: - register_summary(collection_name, file_name, summary) - elif self.cfg.generate_summary and old_file_ids: - unregister_summary(collection_name, file_name) - - with self._jobs_lock: - tracked = self._files.get(file_id) - if tracked: - tracked.status = FileStatus.SUCCESS - tracked.chunk_count = len(search_documents) - tracked.ingested_at = ingested_at - if summary: - tracked.metadata["summary"] = summary - return len(search_documents) + 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] = [] @@ -866,18 +1020,9 @@ def _iter_documents( raise RuntimeError("Azure AI Search pagination did not advance") last_id = next_id - def _find_file_ids_by_name(self, file_name: str, collection_name: str) -> set[str]: - client = self._get_search_client(collection_name) - filter_text = f"file_name eq {_odata_literal(file_name)}" - return { - str(hit["file_id"]) - for hit in self._iter_documents(client, filter_text=filter_text, select=["id", "file_id"]) - if hit.get("file_id") - } - def _delete_file_documents(self, file_id: str, collection_name: str) -> int: - client = self._get_search_client(collection_name) - filter_text = f"file_id eq {_odata_literal(file_id)}" + 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"]) @@ -890,7 +1035,7 @@ def _delete_file_documents(self, file_id: str, collection_name: str) -> int: def _generate_summary(self, text: str, file_name: str) -> str | None: if self._summary_llm is None: return None - snippet = text[: self.cfg.summary_max_chars] + 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) @@ -910,6 +1055,9 @@ 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 _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): @@ -922,8 +1070,19 @@ def _update_file_progress(self, job_id: str, index: int, **fields: Any) -> None: tracked.error_message = details[index].error_message tracked.chunk_count = details[index].chunks_created if tracked.status == FileStatus.SUCCESS: - tracked.ingested_at = _utc_now() + 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: @@ -951,76 +1110,71 @@ def create_collection( description: str | None = None, metadata: dict[str, Any] | None = None, ) -> CollectionInfo: - index, marker = self._ensure_index(name, description, metadata) - if description is not None or metadata: - marker = self._update_marker( - name, - description=description if description is not None else marker.get("description"), - metadata={**(marker.get("metadata") or {}), **(metadata or {})}, - updated_at=_utc_now().isoformat(), - ) - return self._collection_info(name, index, marker) + 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: - try: - index, _marker = self._get_owned_index(name) - except ResourceNotFoundError: + manifest = self._get_collection_manifest(name) + if manifest is None: return False - self._index_client.delete_index(index.name) - try: - self._index_client.get_index(index.name) - except ResourceNotFoundError: - pass - else: - return False - self._search_clients.pop(name, None) + 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]: - collections: list[CollectionInfo] = [] - for index in self._index_client.list_indexes(): - marker = _decode_marker(index.description) - if not marker or marker.get("backend") != _BACKEND_NAME or marker.get("schema_version") != _SCHEMA_VERSION: - continue - collection_name = marker.get("collection_name") - if not isinstance(collection_name, str): - continue - try: - _validate_index_schema(index, collection_name, self.cfg) - collections.append(self._collection_info(collection_name, index, marker)) - except Exception: # noqa: BLE001 - logger.exception("Skipping invalid AI-Q Azure AI Search index %s", index.name) - return collections - - def get_collection(self, name: str) -> CollectionInfo | None: try: - index, marker = self._get_owned_index(name) + 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 None - return self._collection_info(name, index, marker) + return [] - def _collection_info(self, name: str, index: SearchIndex, marker: dict[str, Any]) -> CollectionInfo: - client = self._get_search_client(name) - files = self.list_files(name) - return CollectionInfo( - name=name, - description=marker.get("description"), - file_count=len(files), - chunk_count=client.get_document_count(), - backend=_BACKEND_NAME, - metadata={ - **(marker.get("metadata") or {}), - "index_name": index.name, - "embedding_model": marker.get("embedding_model"), - "embedding_dim": marker.get("embedding_dim"), - }, - created_at=_parse_timestamp(marker.get("created_at")), - updated_at=_parse_timestamp(marker.get("updated_at")), - ) + 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, @@ -1052,88 +1206,65 @@ 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 - deleted = self._delete_file_documents(file_id, collection_name) - if deleted == 0 and info.status not in {FileStatus.FAILED, FileStatus.UPLOADING}: - 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)) + self._wait_for_search_state( + _record_filter(_RECORD_FILE, collection_name, file_id=file_id), + present=False, + label=f"file {file_id!r}", + ) with self._jobs_lock: self._files.pop(file_id, None) - if deleted: - remaining_same_name = any( - item.file_name == info.file_name and item.file_id != file_id + + if self.cfg.generate_summary: + remaining = [ + item for item in self.list_files(collection_name) - ) - if self.cfg.generate_summary and not remaining_same_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._update_collection_timestamp(collection_name) return True def list_files(self, collection_name: str) -> list[FileInfo]: + if collection_name in self._deleted_collections: + return [] try: - self._get_owned_index(collection_name) - client = self._get_search_client(collection_name) - hits = self._iter_documents( - client, - select=[ - "id", - "file_id", - "file_name", - "file_size", - "uploaded_at", - "ingested_at", - "metadata", - ], + 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), ) - aggregated: dict[str, dict[str, Any]] = {} - for hit in hits: - file_id = str(hit.get("file_id") or "") - if not file_id: - continue - entry = aggregated.setdefault( - file_id, - { - "file_name": hit.get("file_name") or "unknown", - "file_size": hit.get("file_size"), - "uploaded_at": _parse_timestamp(hit.get("uploaded_at")), - "ingested_at": _parse_timestamp(hit.get("ingested_at")), - "metadata": _parse_metadata(hit.get("metadata")), - "chunk_count": 0, - }, - ) - entry["chunk_count"] += 1 + 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 [] - - files = { - file_id: FileInfo( - file_id=file_id, - file_name=entry["file_name"], - collection_name=collection_name, - status=FileStatus.SUCCESS, - file_size=entry["file_size"], - chunk_count=entry["chunk_count"], - uploaded_at=entry["uploaded_at"], - ingested_at=entry["ingested_at"], - metadata=entry["metadata"], - ) - for file_id, entry in aggregated.items() - } - with self._jobs_lock: - for file_id, tracked in self._files.items(): - if tracked.collection_name == collection_name and ( - tracked.status != FileStatus.SUCCESS or file_id not in files - ): - files[file_id] = tracked.model_copy(deep=True) - return sorted(files.values(), key=lambda item: (item.file_name, item.file_id)) + 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: - with self._jobs_lock: - tracked = self._files.get(file_id) - if tracked and tracked.collection_name == collection_name and tracked.status != FileStatus.SUCCESS: - return tracked.model_copy(deep=True) - return next((file_info for file_info in self.list_files(collection_name) if file_info.file_id == file_id), 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: diff --git a/sources/knowledge_layer/src/register.py b/sources/knowledge_layer/src/register.py index 5eef5cf59..409d2c119 100644 --- a/sources/knowledge_layer/src/register.py +++ b/sources/knowledge_layer/src/register.py @@ -255,7 +255,7 @@ 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"), @@ -271,29 +271,15 @@ class KnowledgeRetrievalConfig(FunctionBaseConfig, name="knowledge_retrieval"): description="Optional Azure AI Search admin key; defaults to AZURE_SEARCH_API_KEY", ) azure_search_index_prefix: str = Field( - default_factory=lambda: os.environ.get("AIQ_AZURE_SEARCH_INDEX_PREFIX", "aiq"), + default_factory=lambda: _env_value("AIQ_AZURE_SEARCH_INDEX_PREFIX", default="aiq"), min_length=1, - description="Prefix for AI-Q-owned indexes; defaults to AIQ_AZURE_SEARCH_INDEX_PREFIX or aiq", + description="Unique deployment namespace for the shared AI-Q index", ) embed_dim: int = Field( - default_factory=lambda: int(os.environ.get("AIQ_EMBED_DIM", "2048")), + default_factory=lambda: _env_int("AIQ_EMBED_DIM", 2048), gt=0, description="Embedding dimensions; defaults to AIQ_EMBED_DIM and must match existing indexes", ) - use_hybrid: bool = Field(default=True, description="Combine vector and keyword search (azure_ai_search only)") - use_semantic_ranker: bool = Field( - default=False, - description=( - "Enable Azure semantic ranking when supported by the Azure AI Search service (azure_ai_search only)" - ), - ) - chunk_size: int = Field(default=512, gt=0, description="Tokens per chunk (azure_ai_search only)") - chunk_overlap: int = Field(default=64, ge=0, description="Token overlap between chunks (azure_ai_search only)") - summary_max_chars: int = Field( - default=1000, - gt=0, - description="Maximum document characters sent to summary model (azure_ai_search only)", - ) @model_validator(mode="after") def validate_backend_config(self): @@ -345,10 +331,6 @@ def validate_backend_config(self): elif backend == "azure_ai_search": if self.azure_search_endpoint is None: raise ValueError("azure_ai_search requires azure_search_endpoint") - if self.chunk_overlap >= self.chunk_size: - raise ValueError("chunk_overlap must be smaller than chunk_size") - if not self.use_hybrid and self.use_semantic_ranker: - raise ValueError("use_semantic_ranker=true requires use_hybrid=true") return self @@ -445,11 +427,6 @@ def _setup_backend(config: KnowledgeRetrievalConfig, summary_llm_obj=None) -> tu "embed_base_url": str(config.embed_base_url), "embed_model": config.embed_model, "embed_dim": config.embed_dim, - "use_hybrid": config.use_hybrid, - "use_semantic_ranker": config.use_semantic_ranker, - "chunk_size": config.chunk_size, - "chunk_overlap": config.chunk_overlap, - "summary_max_chars": config.summary_max_chars, "collection_name": config.collection_name, "cleanup_files": False, **summary_config, @@ -550,7 +527,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/test_azure_ai_search.py b/tests/knowledge_layer_tests/test_azure_ai_search.py index 974818be5..819e14c49 100644 --- a/tests/knowledge_layer_tests/test_azure_ai_search.py +++ b/tests/knowledge_layer_tests/test_azure_ai_search.py @@ -20,9 +20,8 @@ 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 _decode_marker from knowledge_layer.azure_ai_search.adapter import _encode_marker -from knowledge_layer.azure_ai_search.adapter import _index_name_for_collection +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 @@ -49,6 +48,15 @@ def __init__(self, key: str, succeeded: bool = True, error_message: str | None = 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 @@ -62,8 +70,12 @@ def __init__(self): 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) @@ -86,28 +98,60 @@ def delete_documents(self, documents: list[dict]): if remaining > 0: self.delete_failures_remaining[key] = remaining - 1 if succeeded: - self.documents.pop(key, None) + 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 search_text, order_by, kwargs + del order_by self.search_filters.append(filter) - file_id = _literal(filter, "file_id") - file_name = _literal(filter, "file_name") + 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 = sorted(self.documents.values(), key=lambda item: item["id"]) - if file_id is not None: - documents = [document for document in documents if document.get("file_id") == file_id] - if file_name is not None: - documents = [document for document in documents if document.get("file_name") == file_name] + 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: - return [{key: document.get(key) for key in select} for document in documents] - return [dict(document) for document in documents] + 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) @@ -187,8 +231,6 @@ def _config(**overrides): "embed_model": "test-embed", "embed_dim": 4, "embed_base_url": "https://integrate.api.nvidia.com/v1", - "use_hybrid": True, - "use_semantic_ranker": True, "start_ttl_cleanup": False, "index_prefix": "aiq-test", } @@ -200,10 +242,38 @@ def _ingestor(**overrides): ingestor = AzureAISearchIngestor(_config(**overrides)) ingestor._index_client = FakeIndexClient() search_client = FakeSearchClient() - ingestor._get_search_client = lambda collection_name: search_client + 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): @@ -222,7 +292,7 @@ def test_backend_registered_and_implements_sdk_contracts(): assert issubclass(AzureAISearchRetriever, BaseRetriever) -def test_config_requires_endpoint_and_valid_hybrid_semantic_combination(monkeypatch): +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) @@ -233,14 +303,8 @@ def test_config_requires_endpoint_and_valid_hybrid_semantic_combination(monkeypa azure_search_endpoint="https://example.search.windows.net", ) assert config.azure_search_api_key is None - assert not config.use_semantic_ranker - with pytest.raises(ValueError, match="use_semantic_ranker"): - KnowledgeRetrievalConfig( - backend="azure_ai_search", - azure_search_endpoint="https://example.search.windows.net", - use_hybrid=False, - use_semantic_ranker=True, - ) + 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): @@ -262,9 +326,20 @@ def test_config_uses_shared_environment_defaults(monkeypatch): 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 not backend_config["use_semantic_ranker"] assert adapter_config.endpoint == "https://env.search.windows.net" - assert not adapter_config.use_semantic_ranker + + +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): @@ -313,33 +388,35 @@ def test_setup_backend_preserves_secrets_and_prefix(monkeypatch): assert backend_config["index_prefix"] == "tenant-aiq" -def test_index_names_are_official_stable_and_collision_safe(): - names = ["Tenant A", "tenant-a", "tenant/a", "TENANT+A"] - indexes = [_index_name_for_collection(name, "AIQ Prod") for name in names] +def test_index_name_is_stable_and_changes_with_embedding_configuration(): + index = _index_name_for_config("AIQ Prod", "test/embed", 2048) - assert len(set(indexes)) == len(names) - assert indexes[0] == _index_name_for_collection("Tenant A", "AIQ Prod") - assert all(len(index) <= 128 for index in indexes) - assert all(re.fullmatch(r"[a-z0-9-]+", index) for index in indexes) - _validate_index_name(indexes[0]) + 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, use_semantic_ranker=True) - marker = _new_marker("docs", cfg, "Docs", {"tenant": "alpha"}) + 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, "docs", cfg)["metadata"] == {"tenant": "alpha"} + assert _validate_index_schema(index, cfg)["schema_version"] == 2 + assert "metadata" not in marker + assert "collection" not in marker index.description = "unmanaged" with pytest.raises(RuntimeError, match="not owned"): - _validate_index_schema(index, "docs", cfg) + _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, "docs", cfg) + _validate_index_schema(index, cfg) def test_create_collection_handles_race_and_ignores_unmanaged_indexes(): @@ -353,6 +430,45 @@ def test_create_collection_handles_race_and_ignores_unmanaged_indexes(): 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_legacy_and_foreign_indexes_are_ignored(): + ingestor, _client = _ingestor() + legacy_marker = _new_marker(ingestor.cfg) + legacy_marker["schema_version"] = 1 + ingestor._index_client.indexes["aiq-test-docs-v1"] = _build_index_schema( + "aiq-test-docs-v1", + ingestor.cfg.embed_dim, + _encode_marker(legacy_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 @@ -364,13 +480,13 @@ def test_create_collection_propagates_real_create_failure(): @pytest.mark.parametrize( ("hit", "expected"), [ - ({"@search.reranker_score": 2.0, "@search.score": 0.9}, 0.5), - ({"@search.reranker_score": 8.0}, 1.0), - ({"@search.reranker_score": -1.0}, 0.0), + ({"@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_and_prefers_reranker_score(hit, expected): +def test_normalize_clamps_search_score(hit, expected): chunk = AzureAISearchRetriever.__new__(AzureAISearchRetriever).normalize( {"id": "chunk-1", "chunk": "text", "file_name": "report.pdf", **hit} ) @@ -410,7 +526,7 @@ def test_shared_formatter_retains_source_and_citation_lines(): @pytest.mark.asyncio -async def test_retrieve_builds_hybrid_semantic_search_request(): +async def test_retrieve_builds_scoped_hybrid_search_request(): class FakeClient: def search(self, **kwargs): self.kwargs = kwargs @@ -418,7 +534,6 @@ def search(self, **kwargs): client = FakeClient() retriever = AzureAISearchRetriever.__new__(AzureAISearchRetriever) - retriever.cfg = SimpleNamespace(use_hybrid=True, use_semantic_ranker=True) retriever._embedding = FakeEmbedding(2) retriever._get_client = lambda collection_name: client @@ -426,7 +541,10 @@ def search(self, **kwargs): assert result.success assert client.kwargs["search_text"] == "hello" - assert client.kwargs["query_type"] == "semantic" + 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"] @@ -454,12 +572,17 @@ def start(self): path.write_text("content", encoding="utf-8") ingestor, _client = _ingestor() - job_id = ingestor.submit_job([str(path)], "docs", {"original_filenames": ["original.txt"]}) + 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 file_id in ingestor._files assert ingestor._files[file_id].file_name == "original.txt" + assert ingestor.get_file_status(file_id, "docs").metadata["large"] == "x" * 10_000 def test_upload_file_preserves_caller_owned_source_by_default(monkeypatch, tmp_path): @@ -485,6 +608,37 @@ def start(self): 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)) @@ -535,15 +689,7 @@ def test_pagination_is_exhaustive_and_file_ids_are_odata_escaped(monkeypatch): ingestor, client = _ingestor() ingestor.create_collection("docs") for index in range(5): - client.documents[f"chunk-{index}"] = { - "id": f"chunk-{index}", - "file_id": f"file-{index}", - "file_name": f"file-{index}.txt", - "file_size": index, - "uploaded_at": datetime.now(UTC), - "ingested_at": datetime.now(UTC), - "metadata": "{}", - } + _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 ''" @@ -551,14 +697,8 @@ def test_pagination_is_exhaustive_and_file_ids_are_odata_escaped(monkeypatch): assert any("file_id eq 'x'' or file_id ne '''''" in (item or "") for item in client.search_filters) -def test_same_name_replacement_removes_old_generation(monkeypatch, tmp_path): +def test_same_name_upload_keeps_existing_file_chunks(monkeypatch, tmp_path): _install_reader(monkeypatch) - cleared = [] - monkeypatch.setattr( - azure_adapter, - "unregister_summary", - lambda collection, file_name: cleared.append((collection, file_name)), - ) path = tmp_path / "document.txt" path.write_text("new content", encoding="utf-8") ingestor, client = _ingestor() @@ -567,6 +707,8 @@ def test_same_name_replacement_removes_old_generation(monkeypatch, tmp_path): 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", @@ -579,7 +721,7 @@ def test_same_name_replacement_removes_old_generation(monkeypatch, tmp_path): "metadata": "{}", } - count = ingestor._process_file( + count, summary, _ingested_at = ingestor._process_file( path=str(path), collection_name="docs", file_id="new", @@ -590,12 +732,15 @@ def test_same_name_replacement_removes_old_generation(monkeypatch, tmp_path): ) assert count == 1 - assert set(client.documents) == {"new-00000000"} - assert client.documents["new-00000000"]["metadata"] == '{"tenant":"alpha"}' - assert cleared == [] + 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_new_generation_preserves_old_generation(monkeypatch, tmp_path): +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") @@ -605,10 +750,12 @@ def test_failed_new_generation_preserves_old_generation(monkeypatch, tmp_path): 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("new-00000001") + client.fail_upload_ids.add("chunk-new-00000001") with pytest.raises(RuntimeError, match="rejected upload"): ingestor._process_file( @@ -621,51 +768,33 @@ def test_failed_new_generation_preserves_old_generation(monkeypatch, tmp_path): metadata={}, ) - assert set(client.documents) == {"old-00000000"} + chunk_ids = {document["id"] for document in client.documents.values() if document.get("record_type") == "chunk"} + assert chunk_ids == {"old-00000000"} -def test_replacement_cleanup_failure_retains_complete_new_generation(monkeypatch, tmp_path): - _install_reader(monkeypatch) - path = tmp_path / "document.txt" - path.write_text("new content", encoding="utf-8") - ingestor, client = _ingestor() +def test_duplicate_file_manifests_are_independent(): + ingestor, _client = _ingestor() ingestor.create_collection("docs") - ingestor._embedding = FakeEmbedding() - ingestor._splitter = FakeSplitter([FakeNode("new")]) - client.documents["old-00000000"] = { - "id": "old-00000000", - "file_id": "old", - "file_name": "document.txt", - } - client.delete_failures_remaining["old-00000000"] = 10 + _write_file_manifest(ingestor, "old", file_name="document.txt") + _write_file_manifest(ingestor, "new", file_name="document.txt") - with pytest.raises(RuntimeError, match="rejected delete"): - 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={}, - ) - - assert set(client.documents) == {"old-00000000", "new-00000000"} + 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"}) - now = datetime.now(UTC) + _write_file_manifest(ingestor, "file-1", chunk_count=3) for index in range(3): - client.documents[f"file-1-{index:08d}"] = { - "id": f"file-1-{index:08d}", + 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", - "file_size": 42, - "uploaded_at": now, - "ingested_at": now, - "metadata": '{"section":"finance"}', } files = ingestor.list_files("docs") @@ -678,18 +807,50 @@ def test_metadata_counts_and_status_round_trip(): 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_failed_uploads_remain_visible(): ingestor, _client = _ingestor() ingestor.create_collection("docs") - ingestor._files["failed"] = azure_adapter.FileInfo( - file_id="failed", - file_name="failed.txt", - collection_name="docs", - status=FileStatus.FAILED, - error_message="bad file", - ) + _write_file_manifest(ingestor, "failed", file_name="failed.txt", status=FileStatus.FAILED) files = ingestor.list_files("docs") @@ -697,17 +858,11 @@ def test_failed_uploads_remain_visible(): def test_ttl_deletes_only_expired_owned_collection_and_clears_summary(monkeypatch): - ingestor, _client = _ingestor(generate_summary=True) + ingestor, client = _ingestor(generate_summary=True) ingestor.create_collection("old") ingestor.create_collection("new") - old_name = ingestor._physical_index_name("old") - new_name = ingestor._physical_index_name("new") - old_marker = _decode_marker(ingestor._index_client.indexes[old_name].description) - new_marker = _decode_marker(ingestor._index_client.indexes[new_name].description) - old_marker["updated_at"] = (datetime.now(UTC) - timedelta(hours=25)).isoformat() - new_marker["updated_at"] = datetime.now(UTC).isoformat() - ingestor._index_client.indexes[old_name].description = _encode_marker(old_marker) - ingestor._index_client.indexes[new_name].description = _encode_marker(new_marker) + 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) @@ -715,8 +870,9 @@ def test_ttl_deletes_only_expired_owned_collection_and_clears_summary(monkeypatc ingestor._cleanup_expired_collections() - assert old_name not in ingestor._index_client.indexes - assert new_name in ingestor._index_client.indexes + 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"] @@ -736,16 +892,15 @@ def test_collection_summary_clears_only_after_confirmed_delete(monkeypatch): def test_file_summary_clears_only_after_confirmed_delete(monkeypatch): ingestor, client = _ingestor(generate_summary=True) ingestor.create_collection("docs") - client.documents["file-1-00000000"] = { - "id": "file-1-00000000", + _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", - "file_size": 1, - "uploaded_at": datetime.now(UTC), - "ingested_at": datetime.now(UTC), - "metadata": "{}", } - client.delete_failures_remaining["file-1-00000000"] = 10 + client.delete_failures_remaining["chunk-file-1-00000000"] = 10 cleared = [] monkeypatch.setattr(azure_adapter, "unregister_summary", lambda collection, file_name: cleared.append(file_name)) @@ -756,17 +911,9 @@ def test_file_summary_clears_only_after_confirmed_delete(monkeypatch): def test_summary_cleanup_is_skipped_when_summaries_are_disabled(monkeypatch): - ingestor, client = _ingestor() + ingestor, _client = _ingestor() ingestor.create_collection("docs") - client.documents["file-1-00000000"] = { - "id": "file-1-00000000", - "file_id": "file-1", - "file_name": "report.pdf", - "file_size": 1, - "uploaded_at": datetime.now(UTC), - "ingested_at": datetime.now(UTC), - "metadata": "{}", - } + _write_file_manifest(ingestor, "file-1") cleared = [] monkeypatch.setattr( azure_adapter, @@ -780,6 +927,40 @@ def test_summary_cleanup_is_skipped_when_summaries_are_disabled(monkeypatch): 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() @@ -802,7 +983,16 @@ def test_index_schema_uses_requested_dimensions_and_fields(): assert fields["embedding"].vector_search_dimensions == 4096 assert fields["file_id"].filterable assert fields["id"].sortable - assert schema.semantic_search.default_configuration_name == "default-semantic" + 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): From e74b76e23e6ddb6ff7903aa28b06eae50164f7c4 Mon Sep 17 00:00:00 2001 From: Kyle Zheng Date: Fri, 10 Jul 2026 16:22:58 -0700 Subject: [PATCH 20/27] fix(knowledge): preserve Azure key across server reload Signed-off-by: Kyle Zheng --- sources/knowledge_layer/src/register.py | 3 ++- tests/knowledge_layer_tests/test_azure_ai_search.py | 13 +++++++++++++ 2 files changed, 15 insertions(+), 1 deletion(-) diff --git a/sources/knowledge_layer/src/register.py b/sources/knowledge_layer/src/register.py index 409d2c119..02948bc24 100644 --- a/sources/knowledge_layer/src/register.py +++ b/sources/knowledge_layer/src/register.py @@ -34,6 +34,7 @@ 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__) @@ -266,7 +267,7 @@ class KnowledgeRetrievalConfig(FunctionBaseConfig, name="knowledge_retrieval"): default_factory=lambda: _url_from_env("AZURE_SEARCH_ENDPOINT"), description="Azure AI Search service URL; defaults to AZURE_SEARCH_ENDPOINT", ) - azure_search_api_key: SecretStr | None = Field( + 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", ) diff --git a/tests/knowledge_layer_tests/test_azure_ai_search.py b/tests/knowledge_layer_tests/test_azure_ai_search.py index 819e14c49..bbae0f86b 100644 --- a/tests/knowledge_layer_tests/test_azure_ai_search.py +++ b/tests/knowledge_layer_tests/test_azure_ai_search.py @@ -388,6 +388,19 @@ def test_setup_backend_preserves_secrets_and_prefix(monkeypatch): 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) From b86340a3a8b24090ecab9b961312a30285ed4b7c Mon Sep 17 00:00:00 2001 From: Harmke Alkemade Date: Mon, 13 Jul 2026 13:33:23 +0200 Subject: [PATCH 21/27] fix(knowledge): align Azure Search schema checks with v1 Signed-off-by: Harmke Alkemade --- docs/source/customization/knowledge-layer.md | 3 +-- docs/source/examples/azure-ai-search.md | 5 ++--- .../knowledge_layer_tests/test_azure_ai_search.py | 14 +++++++------- 3 files changed, 10 insertions(+), 12 deletions(-) diff --git a/docs/source/customization/knowledge-layer.md b/docs/source/customization/knowledge-layer.md index 75a4e4e00..55ce47ff0 100644 --- a/docs/source/customization/knowledge-layer.md +++ b/docs/source/customization/knowledge-layer.md @@ -179,8 +179,7 @@ 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. Schema-version-1 -indexes are ignored and require re-ingestion. +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 diff --git a/docs/source/examples/azure-ai-search.md b/docs/source/examples/azure-ai-search.md index ce975b1c6..3f738f4f1 100644 --- a/docs/source/examples/azure-ai-search.md +++ b/docs/source/examples/azure-ai-search.md @@ -80,6 +80,5 @@ 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. Schema-version-1 indexes are ignored and must be re-ingested. File IDs -returned by upload are authoritative for status and delete operations; -same-name uploads coexist independently. +overlap. File IDs returned by upload are authoritative for status and +delete operations; same-name uploads coexist independently. diff --git a/tests/knowledge_layer_tests/test_azure_ai_search.py b/tests/knowledge_layer_tests/test_azure_ai_search.py index bbae0f86b..9c190d2ce 100644 --- a/tests/knowledge_layer_tests/test_azure_ai_search.py +++ b/tests/knowledge_layer_tests/test_azure_ai_search.py @@ -419,7 +419,7 @@ def test_marker_and_schema_validation_reject_unowned_and_mismatched_indexes(): marker = _new_marker(cfg) index = _build_index_schema("aiq-docs-123456789abc", 4, _encode_marker(marker)) - assert _validate_index_schema(index, cfg)["schema_version"] == 2 + assert _validate_index_schema(index, cfg)["schema_version"] == 1 assert "metadata" not in marker assert "collection" not in marker index.description = "unmanaged" @@ -468,14 +468,14 @@ def test_create_waits_for_collection_manifest_search_visibility(monkeypatch): assert [item.name for item in ingestor.list_collections()] == ["docs"] -def test_legacy_and_foreign_indexes_are_ignored(): +def test_mismatched_and_foreign_indexes_are_ignored(): ingestor, _client = _ingestor() - legacy_marker = _new_marker(ingestor.cfg) - legacy_marker["schema_version"] = 1 - ingestor._index_client.indexes["aiq-test-docs-v1"] = _build_index_schema( - "aiq-test-docs-v1", + 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(legacy_marker), + _encode_marker(mismatched_marker), ) ingestor._index_client.indexes["foreign"] = SearchIndex(name="foreign", fields=[], description="foreign") From 1813715ba6b8f35b5047c3df3bab51ecc4f6673a Mon Sep 17 00:00:00 2001 From: Harmke Alkemade Date: Mon, 13 Jul 2026 13:47:36 +0200 Subject: [PATCH 22/27] fix(knowledge): reject unsupported Azure Search filters Signed-off-by: Harmke Alkemade --- sources/knowledge_layer/src/azure_ai_search/adapter.py | 4 ++-- tests/knowledge_layer_tests/test_azure_ai_search.py | 9 +++++---- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/sources/knowledge_layer/src/azure_ai_search/adapter.py b/sources/knowledge_layer/src/azure_ai_search/adapter.py index 7df83f345..489ac2d47 100644 --- a/sources/knowledge_layer/src/azure_ai_search/adapter.py +++ b/sources/knowledge_layer/src/azure_ai_search/adapter.py @@ -515,13 +515,13 @@ def _retrieve_sync( top_k: int, filters: dict[str, Any] | None, ) -> RetrievalResult: - if filters and isinstance(filters, dict) and filters.get("$filter"): + if filters: return RetrievalResult( query=query, backend=_BACKEND_NAME, chunks=[], success=False, - error_message="Raw Azure AI Search $filter expressions are not supported", + error_message="Azure AI Search metadata filters are not supported", ) try: client = self._get_client(collection_name) diff --git a/tests/knowledge_layer_tests/test_azure_ai_search.py b/tests/knowledge_layer_tests/test_azure_ai_search.py index 9c190d2ce..cbe22c946 100644 --- a/tests/knowledge_layer_tests/test_azure_ai_search.py +++ b/tests/knowledge_layer_tests/test_azure_ai_search.py @@ -550,7 +550,7 @@ def search(self, **kwargs): retriever._embedding = FakeEmbedding(2) retriever._get_client = lambda collection_name: client - result = await retriever.retrieve("hello", "session-1", top_k=5) + result = await retriever.retrieve("hello", "session-1", top_k=5, filters={}) assert result.success assert client.kwargs["search_text"] == "hello" @@ -562,14 +562,15 @@ def search(self, **kwargs): @pytest.mark.asyncio -async def test_retrieve_rejects_raw_odata_filter(): +@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={"$filter": "file_id ne ''"}) + result = await retriever.retrieve("hello", "session-1", filters=filters) assert not result.success - assert result.error_message == "Raw Azure AI Search $filter expressions are not supported" + assert result.error_message == "Azure AI Search metadata filters are not supported" def test_submit_job_uses_one_canonical_file_id(monkeypatch, tmp_path): From 4809562ef0d817ed4bc26e3a3cd275b0dd92213b Mon Sep 17 00:00:00 2001 From: Harmke Alkemade Date: Mon, 13 Jul 2026 14:00:03 +0200 Subject: [PATCH 23/27] fix(knowledge): finalize Azure file deletion before consistency wait Signed-off-by: Harmke Alkemade --- .../src/azure_ai_search/adapter.py | 10 ++--- .../test_azure_ai_search.py | 41 +++++++++++++++++++ 2 files changed, 46 insertions(+), 5 deletions(-) diff --git a/sources/knowledge_layer/src/azure_ai_search/adapter.py b/sources/knowledge_layer/src/azure_ai_search/adapter.py index 489ac2d47..6f5bfc85a 100644 --- a/sources/knowledge_layer/src/azure_ai_search/adapter.py +++ b/sources/knowledge_layer/src/azure_ai_search/adapter.py @@ -1215,11 +1215,6 @@ def delete_file(self, file_id: str, collection_name: str) -> bool: [_record_id(_RECORD_FILE, collection_name, file_id)], ) self._deleted_files.add((collection_name, file_id)) - self._wait_for_search_state( - _record_filter(_RECORD_FILE, collection_name, file_id=file_id), - present=False, - label=f"file {file_id!r}", - ) with self._jobs_lock: self._files.pop(file_id, None) @@ -1235,6 +1230,11 @@ def delete_file(self, file_id: str, collection_name: str) -> bool: 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]: diff --git a/tests/knowledge_layer_tests/test_azure_ai_search.py b/tests/knowledge_layer_tests/test_azure_ai_search.py index cbe22c946..b532ba11c 100644 --- a/tests/knowledge_layer_tests/test_azure_ai_search.py +++ b/tests/knowledge_layer_tests/test_azure_ai_search.py @@ -861,6 +861,47 @@ def test_delete_waits_for_stale_file_manifest_to_disappear(monkeypatch): 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") From aebac162973dc7cba596e8eec57849e598306a37 Mon Sep 17 00:00:00 2001 From: Harmke Alkemade Date: Mon, 13 Jul 2026 14:39:12 +0200 Subject: [PATCH 24/27] docs: add Azure AI Search E2E validation setup Signed-off-by: Harmke Alkemade --- configs/config_web_azure_ai_search.yml | 211 ++++++++++++++++++ deploy/.env.example | 10 + .../customization/configuration-reference.md | 1 + docs/source/examples/azure-ai-search.md | 56 ++--- docs/source/examples/index.md | 2 +- .../knowledge_layer/KNOWLEDGE-LAYER-SETUP.md | 1 + .../run_adapter_compliance.py | 9 + 7 files changed, 263 insertions(+), 27 deletions(-) create mode 100644 configs/config_web_azure_ai_search.yml diff --git a/configs/config_web_azure_ai_search.yml b/configs/config_web_azure_ai_search.yml new file mode 100644 index 000000000..ad1d1c64b --- /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 and Paper search tools by default + +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 319d8ae4c..e38ac3bef 100644 --- a/deploy/.env.example +++ b/deploy/.env.example @@ -166,6 +166,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 92a8b3964..5dfc47d62 100644 --- a/docs/source/customization/configuration-reference.md +++ b/docs/source/customization/configuration-reference.md @@ -536,6 +536,7 @@ The repository includes several pre-built configurations: |------|------|----------| | `configs/config_cli_default.yml` | CLI | Web search, paper search, clarifier with plan approval | | `configs/config_web_default_llamaindex.yml` | Web API | LlamaIndex knowledge retrieval, web search, paper search | +| `configs/config_web_azure_ai_search.yml` | Web API | Azure AI Search knowledge retrieval and web search | | `configs/config_web_frag.yml` | Web API | Foundational RAG knowledge retrieval, web search, paper search | ## Related diff --git a/docs/source/examples/azure-ai-search.md b/docs/source/examples/azure-ai-search.md index 3f738f4f1..e05cb7a13 100644 --- a/docs/source/examples/azure-ai-search.md +++ b/docs/source/examples/azure-ai-search.md @@ -24,13 +24,23 @@ Install the backend dependency: uv pip install -e "sources/knowledge_layer[azure_ai_search]" ``` -Set the environment used by the adapter: +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 -export AZURE_SEARCH_ENDPOINT=https://.search.windows.net -export NVIDIA_API_KEY= -# Optional; omit to use DefaultAzureCredential. -export AZURE_SEARCH_API_KEY= +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 @@ -46,30 +56,24 @@ Search service and grant the workload identity both of these built-in roles: 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: -Replace the `knowledge_search` block in a web configuration such as -`configs/config_web_default_llamaindex.yml`: - -```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} +```bash +./scripts/start_e2e.sh --config_file configs/config_web_azure_ai_search.yml ``` -Explicit YAML options override environment defaults. `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. +`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 diff --git a/docs/source/examples/index.md b/docs/source/examples/index.md index 65058cb9a..5f9b733e5 100644 --- a/docs/source/examples/index.md +++ b/docs/source/examples/index.md @@ -12,7 +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_default_llamaindex.yml` base | +| [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 Modal sandbox execution for quantitative research workflows | `config_domain_routing_and_skills.yml` | diff --git a/sources/knowledge_layer/KNOWLEDGE-LAYER-SETUP.md b/sources/knowledge_layer/KNOWLEDGE-LAYER-SETUP.md index ae7ece23d..af6d9fe69 100644 --- a/sources/knowledge_layer/KNOWLEDGE-LAYER-SETUP.md +++ b/sources/knowledge_layer/KNOWLEDGE-LAYER-SETUP.md @@ -1046,6 +1046,7 @@ Configuration values are resolved in the following order (highest to lowest prio | `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 | diff --git a/tests/knowledge_layer_tests/run_adapter_compliance.py b/tests/knowledge_layer_tests/run_adapter_compliance.py index 2e86a1757..52b051a69 100755 --- a/tests/knowledge_layer_tests/run_adapter_compliance.py +++ b/tests/knowledge_layer_tests/run_adapter_compliance.py @@ -27,6 +27,15 @@ 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 From 8cd3733d89e03b766d6713e77a8e498cf086c8c4 Mon Sep 17 00:00:00 2001 From: Harmke Alkemade Date: Mon, 13 Jul 2026 15:10:19 +0200 Subject: [PATCH 25/27] fix(knowledge): register Azure ingestion jobs before writes Signed-off-by: Harmke Alkemade --- .../src/azure_ai_search/adapter.py | 86 ++++++++++++++----- .../test_azure_ai_search.py | 79 ++++++++++++++++- 2 files changed, 144 insertions(+), 21 deletions(-) diff --git a/sources/knowledge_layer/src/azure_ai_search/adapter.py b/sources/knowledge_layer/src/azure_ai_search/adapter.py index 6f5bfc85a..a9c2db159 100644 --- a/sources/knowledge_layer/src/azure_ai_search/adapter.py +++ b/sources/knowledge_layer/src/azure_ai_search/adapter.py @@ -775,19 +775,13 @@ def submit_job( ) return job_id - self._ensure_index() - 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") - 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)) - info = FileInfo( + files[file_id] = FileInfo( file_id=file_id, file_name=file_name, collection_name=collection_name, @@ -796,11 +790,9 @@ def submit_job( uploaded_at=submitted_at, metadata={**file_metadata, "job_id": job_id}, ) - with self._jobs_lock: - self._files[file_id] = info - self._write_file_manifest(info) with self._jobs_lock: + self._files.update(files) self._jobs[job_id] = IngestionJobStatus( job_id=job_id, status=JobState.PENDING, @@ -811,12 +803,19 @@ def submit_job( file_details=file_progress, ) - 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() + 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: @@ -843,11 +842,34 @@ def _process_job( ) -> None: cleanup = bool(config.get("cleanup_files", self.cfg.cleanup_files)) self._update_job(job_id, status=JobState.PROCESSING, started_at=_utc_now()) - collection = self._get_collection_manifest(collection_name) - if collection is None or collection.get("status") != _COLLECTION_ACTIVE: - self._fail_job(job_id, f"Collection {collection_name!r} is unavailable") + 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 @@ -1054,6 +1076,30 @@ def _update_job(self, job_id: str, **fields: Any) -> None: 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) diff --git a/tests/knowledge_layer_tests/test_azure_ai_search.py b/tests/knowledge_layer_tests/test_azure_ai_search.py index b532ba11c..576910af0 100644 --- a/tests/knowledge_layer_tests/test_azure_ai_search.py +++ b/tests/knowledge_layer_tests/test_azure_ai_search.py @@ -39,6 +39,7 @@ 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: @@ -594,9 +595,85 @@ def start(self): 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.get_file_status(file_id, "docs").metadata["large"] == "x" * 10_000 + 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): From 6b34c48f848f6fbfd2bfe3adeb9892860d839cf1 Mon Sep 17 00:00:00 2001 From: Kyle Zheng Date: Mon, 13 Jul 2026 10:07:04 -0700 Subject: [PATCH 26/27] docs(config): clarify optional Azure paper search Signed-off-by: Kyle Zheng --- configs/config_web_azure_ai_search.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/configs/config_web_azure_ai_search.yml b/configs/config_web_azure_ai_search.yml index ad1d1c64b..32cb949f6 100644 --- a/configs/config_web_azure_ai_search.yml +++ b/configs/config_web_azure_ai_search.yml @@ -1,7 +1,7 @@ # This is the Web mode configuration for Azure AI Search. # It has the following features: # - Knowledge retrieval using Azure AI Search -# - Web search and Paper search tools by default +# - Web search tools by default; Paper search is optional general: use_uvloop: true From 5a00cf6a71e2fadd596ba0bbadb65d57c5c6f056 Mon Sep 17 00:00:00 2001 From: Kyle Zheng Date: Mon, 13 Jul 2026 10:10:45 -0700 Subject: [PATCH 27/27] docs(knowledge): remove duplicate OpenSearch mention Signed-off-by: Kyle Zheng --- docs/source/customization/configuration-reference.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/source/customization/configuration-reference.md b/docs/source/customization/configuration-reference.md index b50f3972c..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. Supports LlamaIndex (local ChromaDB), OpenSearch, Foundational RAG -(hosted NVIDIA RAG Blueprint), Azure AI Search 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: