diff --git a/python/AGENTS.md b/python/AGENTS.md index 7f5a870ebd..6cdac7c462 100644 --- a/python/AGENTS.md +++ b/python/AGENTS.md @@ -112,7 +112,7 @@ python/ - [foundry](packages/foundry/README.md) - Microsoft Foundry chat, agent, memory, and embedding integrations - [azure-contentunderstanding](packages/azure-contentunderstanding/AGENTS.md) - Azure Content Understanding context provider - [azure-ai-search](packages/azure-ai-search/AGENTS.md) - Azure AI Search RAG -- [azure-cosmos](packages/azure-cosmos/AGENTS.md) - Azure Cosmos DB-backed history provider +- [azure-cosmos](packages/azure-cosmos/AGENTS.md) - Azure Cosmos DB for NoSQL vector, history, and checkpoint integrations Durable Task and Azure Functions integrations are maintained in the [Durable Agent Framework extension](https://github.com/microsoft/agent-framework-durable-extension). diff --git a/python/packages/azure-cosmos/AGENTS.md b/python/packages/azure-cosmos/AGENTS.md index 9bb7f76da9..b151077c3d 100644 --- a/python/packages/azure-cosmos/AGENTS.md +++ b/python/packages/azure-cosmos/AGENTS.md @@ -1,10 +1,14 @@ # Azure Cosmos DB Package (agent-framework-azure-cosmos) -Azure Cosmos DB history provider integration for Agent Framework. +Azure Cosmos DB for NoSQL integrations for Agent Framework. ## Main Classes - **`CosmosHistoryProvider`** - Persistent conversation history storage backed by Azure Cosmos DB +- **`CosmosCheckpointStorage`** - Workflow checkpoint storage backed by Azure Cosmos DB +- **`CosmosCollection`** - Vector collection using NoSQL `VectorDistance` queries +- **`CosmosStore`** - Factory and database administration for vector collections +- **`AzureCosmosSettings`** - Shared vector connector settings shape ## Usage @@ -21,10 +25,17 @@ provider = CosmosHistoryProvider( Container name is configured on the provider. `session_id` is used as the partition key. +Vector collections require an application-provided string key stored as `id` and a +container partition key of `/id`. Container vector and indexing policies are created +only by `ensure_collection_exists()` and validated before use. Existing incompatible +containers are never updated or recreated. +Supported vector element types are float32, int8, and uint8. Euclidean search is supported, +but Euclidean score thresholds fail before service I/O. + ## Import Path ```python -from agent_framework.azure import CosmosHistoryProvider +from agent_framework.azure import CosmosCollection, CosmosHistoryProvider, CosmosStore # or directly: -from agent_framework_azure_cosmos import CosmosHistoryProvider +from agent_framework_azure_cosmos import CosmosCollection, CosmosHistoryProvider, CosmosStore ``` diff --git a/python/packages/azure-cosmos/README.md b/python/packages/azure-cosmos/README.md index e458bcd9eb..72b2977450 100644 --- a/python/packages/azure-cosmos/README.md +++ b/python/packages/azure-cosmos/README.md @@ -1,126 +1,171 @@ -# Get Started with Microsoft Agent Framework Azure Cosmos DB +# Microsoft Agent Framework integrations for Azure Cosmos DB -Please install this package via pip: +Use Azure Cosmos DB for NoSQL as an Agent Framework vector store, conversation history +provider, or workflow checkpoint store. + +## Contents + +- `CosmosCollection` and `CosmosStore` for vector CRUD and `VectorDistance` search +- `CosmosHistoryProvider` for persistent conversation history +- `CosmosCheckpointStorage` for durable workflow checkpoints + +## Install ```bash pip install agent-framework-azure-cosmos --pre ``` -## Azure Cosmos DB History Provider +The package requires Python 3.10 or later, an Azure Cosmos DB for NoSQL account, +and `azure-cosmos` 4.7 or later. Vector collections also require the account's +[NoSQL vector search capability](https://learn.microsoft.com/azure/cosmos-db/vector-search) +to be enabled before use. -The Azure Cosmos DB integration provides `CosmosHistoryProvider` for persistent conversation history storage. +## Authentication and settings -### Basic Usage Example +Pass a caller-owned Azure credential or account key. If constructor values are omitted, +the package reads an explicitly selected `.env` file and then these environment variables: -```python -from azure.identity.aio import DefaultAzureCredential -from agent_framework_azure_cosmos import CosmosHistoryProvider +| Variable | Purpose | +| --- | --- | +| `AZURE_COSMOS_ENDPOINT` | Azure Cosmos DB account endpoint | +| `AZURE_COSMOS_DATABASE_NAME` | Database name | +| `AZURE_COSMOS_CONTAINER_NAME` | Container name for direct collection/provider use | +| `AZURE_COSMOS_KEY` | Account key; omit when passing an Azure credential | -provider = CosmosHistoryProvider( +Explicit constructor values take precedence over a selected `.env` file, which takes +precedence over process environment variables. Injected asynchronous `CosmosClient`, +`DatabaseProxy`, and `ContainerProxy` objects bypass connection settings and remain +caller-owned. + +## Vector store + +Vector collections require: + +- one application-provided string key with storage name `id`; +- a single Hash partition key path `/id`; +- top-level dense vector fields; and +- the container vector and indexing policies derived from the collection definition. + +This makes `get()` and `delete()` unambiguous point operations because the item ID is also +its partition key. IDs must contain 1-1,023 UTF-8 bytes and cannot contain `/`, `\`, +`?`, or `#`; the connector never encodes them. Custom and hierarchical partition keys +are not supported. + +```python +from dataclasses import dataclass +from typing import Annotated + +from agent_framework import VectorStoreField, vectorstoremodel +from agent_framework.azure import CosmosStore +from azure.identity.aio import AzureCliCredential + + +@vectorstoremodel(collection_name="documents") +@dataclass +class Document: + id: Annotated[str, VectorStoreField("key", storage_name="id")] + text: Annotated[str, VectorStoreField("data", is_indexed=True)] + embedding: Annotated[ + list[float], + VectorStoreField( + "vector", + dimensions=1536, + distance_function="cosine_similarity", + ), + ] + + +async with CosmosStore( endpoint="https://.documents.azure.com:443/", - credential=DefaultAzureCredential(), database_name="agent-framework", - container_name="chat-history", -) + credential=AzureCliCredential(), +) as store: + collection = store.get_collection(Document) + await collection.ensure_collection_exists() + await collection.upsert( + [Document(id="doc-1", text="Vector search", embedding=[1.0] + [0.0] * 1535)], + generate_vectors=False, + ) + results = await collection.search(vector=[1.0] + [0.0] * 1535, top=3) + async for result in results: + print(result["record"].text, result["score"]) ``` -Credentials follow the same pattern used by other Azure connectors in the repository: - -- Pass a credential object (for example `DefaultAzureCredential`) -- Or pass a key string directly -- Or set `AZURE_COSMOS_KEY` in the environment +The default vector index is `quantizedFlat`, which supports up to 4,096 dimensions and +uses quantized rather than exact-float recall. Below 1,000 indexed vectors, Azure Cosmos DB +falls back to a full scan. For exact search, explicitly set `index_kind="flat"` on vector +fields with at most 505 dimensions. -Container naming behavior: +Supported vector element types are `float32`, `int8`, and `uint8`. Supported metrics are +cosine similarity, dot product, and Euclidean distance. Scores are returned unchanged: +higher is better for cosine and dot product, while lower is better for Euclidean distance. +Score thresholds are supported for cosine and dot product; Euclidean search is supported +without `score_threshold` because direct `VectorDistance` predicates do not reliably apply +Euclidean cutoffs. +After results are consumed, `SearchResults.metadata` contains the bounded request charge, +last activity ID, and whether the last response page advertised more results; continuation +tokens are not exposed. -- Container name is configured on the provider (`container_name` or `AZURE_COSMOS_CONTAINER_NAME`) -- `session_id` is used as the Cosmos partition key for reads/writes +Vector index tuning is available through the field's `azure_cosmos` provider annotations: -See `samples/02-agents/conversations/cosmos_history_provider.py` for a runnable example. +```python +VectorStoreField( + "vector", + dimensions=1536, + index_kind="disk_ann", + provider_annotations={ + "azure_cosmos": { + "quantizer_type": "spherical", + "quantization_byte_size": 256, + "indexing_search_list_size": 200, + } + }, +) +``` -## Cosmos DB Workflow Checkpoint Storage +Search operation options support `search_list_size_multiplier`, +`quantized_vector_list_multiplier`, `filter_priority`, and `brute_force`. +Values and vectors are always sent as query parameters. -`CosmosCheckpointStorage` implements the `CheckpointStorage` protocol, enabling -durable workflow checkpointing backed by Azure Cosmos DB NoSQL. Workflows can be -paused and resumed across process restarts by persisting checkpoint state in Cosmos DB. +Writes and deletes span `/id` partitions and are not multi-item transactions. The connector +validates the complete input batch before the first request, but a service failure can still +leave an operation partially applied. Retry with stable application keys for idempotent +upserts and deletes. -### Basic Usage +See the [Azure Cosmos DB vector search documentation](https://learn.microsoft.com/azure/cosmos-db/vector-search) +and the [Agent Framework vector-store samples](https://github.com/microsoft/agent-framework/tree/main/python/samples/02-agents/vector_stores). -#### Managed Identity / RBAC (recommended for production) +## Conversation history ```python -from azure.identity.aio import DefaultAzureCredential -from agent_framework import WorkflowBuilder -from agent_framework_azure_cosmos import CosmosCheckpointStorage +from agent_framework.azure import CosmosHistoryProvider +from azure.identity.aio import AzureCliCredential -checkpoint_storage = CosmosCheckpointStorage( +provider = CosmosHistoryProvider( endpoint="https://.documents.azure.com:443/", - credential=DefaultAzureCredential(), + credential=AzureCliCredential(), database_name="agent-framework", - container_name="workflow-checkpoints", + container_name="chat-history", ) ``` -#### Account Key +`CosmosHistoryProvider` stores each conversation under its `session_id` partition key. +See the [conversation sample](https://github.com/microsoft/agent-framework/blob/main/python/samples/02-agents/conversations/cosmos_history_provider.py). + +## Workflow checkpoints ```python from agent_framework_azure_cosmos import CosmosCheckpointStorage +from azure.identity.aio import AzureCliCredential -checkpoint_storage = CosmosCheckpointStorage( +storage = CosmosCheckpointStorage( endpoint="https://.documents.azure.com:443/", - credential="", + credential=AzureCliCredential(), database_name="agent-framework", container_name="workflow-checkpoints", ) ``` -#### Then use with a workflow - -```python -from agent_framework import WorkflowBuilder - -# Build a workflow with checkpointing enabled -workflow = WorkflowBuilder( - start_executor=start, - checkpoint_storage=checkpoint_storage, -).build() - -# Run the workflow — checkpoints are automatically saved after each superstep -result = await workflow.run(message="input data") - -# Resume from a checkpoint -latest = await checkpoint_storage.get_latest(workflow_name=workflow.name) -if latest: - resumed = await workflow.run(checkpoint_id=latest.checkpoint_id) -``` - -### Authentication Options - -`CosmosCheckpointStorage` supports the same authentication modes as `CosmosHistoryProvider`: - -- **Managed identity / RBAC** (recommended): Pass `DefaultAzureCredential()`, - `ManagedIdentityCredential()`, or any Azure `TokenCredential` -- **Account key**: Pass a key string via `credential` parameter -- **Environment variables**: Set `AZURE_COSMOS_ENDPOINT`, `AZURE_COSMOS_DATABASE_NAME`, - `AZURE_COSMOS_CONTAINER_NAME`, and `AZURE_COSMOS_KEY` (key not required when using - Azure credentials) -- **Pre-created client**: Pass an existing `CosmosClient` or `ContainerProxy` - -### Database and Container Setup - -The database and container are created automatically on first use (via -`create_database_if_not_exists` and `create_container_if_not_exists`). The container -uses `/workflow_name` as the partition key. You can also pre-create them in the Azure -portal with this partition key configuration. - -### Environment Variables - -| Variable | Description | -|---|---| -| `AZURE_COSMOS_ENDPOINT` | Cosmos DB account endpoint | -| `AZURE_COSMOS_DATABASE_NAME` | Database name | -| `AZURE_COSMOS_CONTAINER_NAME` | Container name | -| `AZURE_COSMOS_KEY` | Account key (optional if using Azure credentials) | - -See `samples/03-workflows/checkpoint/cosmos_workflow_checkpointing.py` for a standalone example, -or `samples/03-workflows/checkpoint/cosmos_workflow_checkpointing_foundry.py` for an end-to-end -example with Microsoft Foundry agents. +`CosmosCheckpointStorage` uses `/workflow_name` as its partition key and creates its +database and container on first use. See the +[checkpoint sample](https://github.com/microsoft/agent-framework/blob/main/python/samples/03-workflows/checkpoint/cosmos_workflow_checkpointing.py). diff --git a/python/packages/azure-cosmos/agent_framework_azure_cosmos/__init__.py b/python/packages/azure-cosmos/agent_framework_azure_cosmos/__init__.py index 66373b0f1d..1fc45f54a2 100644 --- a/python/packages/azure-cosmos/agent_framework_azure_cosmos/__init__.py +++ b/python/packages/azure-cosmos/agent_framework_azure_cosmos/__init__.py @@ -1,17 +1,55 @@ # Copyright (c) Microsoft. All rights reserved. +import importlib import importlib.metadata +from typing import TYPE_CHECKING, Any + +import agent_framework from ._checkpoint_storage import CosmosCheckpointStorage from ._history_provider import CosmosHistoryProvider +if TYPE_CHECKING: + from ._vector_store import AzureCosmosSettings, CosmosCollection, CosmosStore # pyright: ignore[reportUnusedImport] + +_VECTOR_EXPORTS = frozenset({"AzureCosmosSettings", "CosmosCollection", "CosmosStore"}) +_HAS_VECTOR_CORE = all( + hasattr(agent_framework, name) for name in ("BaseVectorCollection", "BaseVectorSearch", "BaseVectorStore") +) + try: __version__ = importlib.metadata.version(__name__) except importlib.metadata.PackageNotFoundError: __version__ = "0.0.0" # Fallback for development mode -__all__ = [ - "CosmosCheckpointStorage", - "CosmosHistoryProvider", - "__version__", -] + +def __getattr__(name: str) -> Any: + if name not in _VECTOR_EXPORTS: + raise AttributeError(f"Module {__name__!r} has no attribute {name!r}.") + try: + return getattr(importlib.import_module("._vector_store", __name__), name) + except ImportError as exc: + raise ImportError( + "Azure Cosmos DB vector APIs require agent-framework-core with vector-store support." + ) from exc + + +def __dir__() -> list[str]: + return sorted((*globals(), *_VECTOR_EXPORTS)) + + +if _HAS_VECTOR_CORE: + __all__ = [ + "AzureCosmosSettings", + "CosmosCheckpointStorage", + "CosmosCollection", + "CosmosHistoryProvider", + "CosmosStore", + "__version__", + ] +else: + __all__ = [ + "CosmosCheckpointStorage", + "CosmosHistoryProvider", + "__version__", + ] diff --git a/python/packages/azure-cosmos/agent_framework_azure_cosmos/_vector_store.py b/python/packages/azure-cosmos/agent_framework_azure_cosmos/_vector_store.py new file mode 100644 index 0000000000..13141c589d --- /dev/null +++ b/python/packages/azure-cosmos/agent_framework_azure_cosmos/_vector_store.py @@ -0,0 +1,1425 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Azure Cosmos DB for NoSQL vector collections and stores.""" + +from __future__ import annotations + +import json +import math +import re +from collections.abc import AsyncIterable, AsyncIterator, Callable, Mapping, Sequence +from contextlib import suppress +from typing import Any, ClassVar, Generic, TypeAlias, cast +from weakref import WeakSet + +from agent_framework import ( + BaseVectorCollection, + BaseVectorSearch, + BaseVectorStore, + FilterGroup, + SearchResults, + SecretString, + VectorStoreCollectionDefinition, + VectorStoreField, + load_settings, +) +from agent_framework._telemetry import get_user_agent, mark_feature_used +from agent_framework._vector_filters import FilterExpression +from agent_framework._vectors import EmbeddingClient, SearchType, Vector +from agent_framework.exceptions import IntegrationException, IntegrationInvalidResponseException +from azure.core.credentials import TokenCredential +from azure.core.credentials_async import AsyncTokenCredential +from azure.core.serialization import AzureJSONEncoder +from azure.cosmos import PartitionKey +from azure.cosmos.aio import ContainerProxy, CosmosClient, DatabaseProxy +from azure.cosmos.exceptions import ( + CosmosHttpResponseError, + CosmosResourceExistsError, + CosmosResourceNotFoundError, +) +from typing_extensions import TypedDict, TypeVar + +from ._feature_usage import FeatureIndex + +ModelT = TypeVar("ModelT", default=Any) +AzureCredentialTypes: TypeAlias = TokenCredential | AsyncTokenCredential + +_ITEM_SIZE_LIMIT = 2 * 1024 * 1024 +_ID_BYTE_LIMIT = 1023 +_MAX_SAFE_INTEGER = 2**53 - 1 +_MAX_JSON_DEPTH = 128 +_VECTOR_NAME = re.compile(r"[A-Za-z_][A-Za-z0-9_]*\Z") +_VECTOR_TYPES = { + None: "float32", + "float": "float32", + "float32": "float32", + "int8": "int8", + "uint8": "uint8", +} +_VECTOR_RANGES: dict[str, tuple[int | float, int | float, bool]] = { + "float32": (-3.4028234663852886e38, 3.4028234663852886e38, False), + "int8": (-128, 127, True), + "uint8": (0, 255, True), +} +_DISTANCES = { + "DEFAULT": "cosine", + "cosine": "cosine", + "cosine_similarity": "cosine", + "dotproduct": "dotproduct", + "dot_prod": "dotproduct", + "euclidean": "euclidean", + "euclidean_distance": "euclidean", +} +_INDEX_KINDS = { + "default": "quantizedFlat", + "flat": "flat", + "quantized_flat": "quantizedFlat", + "quantizedFlat": "quantizedFlat", + "disk_ann": "diskANN", + "diskANN": "diskANN", +} +_FIELD_ANNOTATIONS = { + "data_type", + "quantizer_type", + "quantization_byte_size", + "indexing_search_list_size", +} +_SEARCH_OPTIONS = { + "search_list_size_multiplier", + "quantized_vector_list_multiplier", + "filter_priority", + "brute_force", +} +_FILTER_SCALAR_TYPES = (str, bool, int, float) + + +class AzureCosmosSettings(TypedDict, total=False): + """Cosmos connection settings resolved from explicit values, a selected .env file, or the environment.""" + + endpoint: str | None + database_name: str | None + container_name: str | None + key: SecretString | None + + +def _validate_operation_options(options: Mapping[str, Any] | None, allowed: set[str]) -> dict[str, Any]: + result = dict(options or {}) + if unknown := result.keys() - allowed: + raise ValueError(f"Unsupported Azure Cosmos DB option(s): {', '.join(sorted(unknown))}.") + return result + + +def _validate_resource_name(value: str, kind: str) -> str: + if not isinstance(value, str) or not value or len(value) > 255: + raise ValueError(f"Cosmos {kind} name must contain 1-255 characters.") + return value + + +def _validate_credential(credential: str | SecretString | AzureCredentialTypes | None) -> None: + if ( + credential is not None + and not isinstance(credential, (str, SecretString)) + and not callable(getattr(credential, "get_token", None)) + ): + raise TypeError("credential must be a key string, SecretString, TokenCredential, or AsyncTokenCredential.") + + +def _load_connection_settings( + *, + endpoint: str | None, + database_name: str | None, + container_name: str | None, + credential: str | SecretString | AzureCredentialTypes | None, + require_container: bool, + env_file_path: str | None, + env_file_encoding: str | None, +) -> tuple[CosmosClient, str, str | None]: + _validate_credential(credential) + required_fields = ["endpoint", "database_name"] + if require_container: + required_fields.append("container_name") + if credential is None: + required_fields.append("key") + settings = load_settings( + AzureCosmosSettings, + env_prefix="AZURE_COSMOS_", + required_fields=required_fields, + endpoint=endpoint, + database_name=database_name, + container_name=container_name, + key=credential if isinstance(credential, (str, SecretString)) else None, + env_file_path=env_file_path, + env_file_encoding=env_file_encoding, + ) + resolved_endpoint = settings.get("endpoint") + resolved_database = settings.get("database_name") + resolved_container = settings.get("container_name") + if not isinstance(resolved_endpoint, str): + raise TypeError("endpoint must be a string.") + if not isinstance(resolved_database, str): + raise TypeError("database_name must be a string.") + _validate_resource_name(resolved_database, "database") + if resolved_container is not None: + _validate_resource_name(resolved_container, "container") + resolved_credential: str | AzureCredentialTypes + if isinstance(credential, SecretString): + resolved_credential = credential.get_secret_value() + elif credential is not None: + resolved_credential = credential + else: + key = settings.get("key") + if not isinstance(key, SecretString): + raise TypeError("key must be a string or SecretString.") + resolved_credential = key.get_secret_value() + client = CosmosClient( + url=resolved_endpoint, + credential=resolved_credential, # type: ignore[arg-type] + user_agent_suffix=get_user_agent(), + ) + return client, resolved_database, resolved_container + + +class _CosmosConnection: + """Share one database proxy while closing only a client created by this connector.""" + + def __init__( + self, + *, + cosmos_client: CosmosClient | None, + database_client: DatabaseProxy | None, + database_name: str, + owns_client: bool, + create_database: bool, + ) -> None: + self.cosmos_client = cosmos_client + self.database_client = database_client + self.database_name = database_name + self.owns_client = owns_client + self.create_database = create_database + self.closed = False + self._database_ready = False + + def ensure_open(self) -> None: + if self.closed: + raise RuntimeError("Azure Cosmos DB vector store is closed.") + + async def get_database(self) -> DatabaseProxy: + self.ensure_open() + if self.database_client is not None: + if not self._database_ready: + await self.database_client.read() + self._database_ready = True + return self.database_client + if self.cosmos_client is None: + raise RuntimeError("Cosmos client is not initialized.") + if self.create_database: + try: + self.database_client = await self.cosmos_client.create_database(id=self.database_name) + except CosmosResourceExistsError: + self.database_client = self.cosmos_client.get_database_client(self.database_name) + await self.database_client.read() + else: + self.database_client = self.cosmos_client.get_database_client(self.database_name) + await self.database_client.read() + self._database_ready = True + return self.database_client + + async def close(self) -> None: + if not self.closed: + self.closed = True + if self.owns_client and self.cosmos_client is not None: + await self.cosmos_client.close() + + +def _connection_from_clients( + *, + endpoint: str | None, + database_name: str | None, + credential: str | SecretString | AzureCredentialTypes | None, + cosmos_client: CosmosClient | None, + database_client: DatabaseProxy | None, + create_database: bool, + env_file_path: str | None, + env_file_encoding: str | None, +) -> _CosmosConnection: + if cosmos_client is not None and database_client is not None: + raise ValueError("Provide at most one of cosmos_client or database_client.") + if database_client is not None: + if any(value is not None for value in (endpoint, database_name, credential, env_file_path, env_file_encoding)): + raise ValueError( + "database_client cannot be combined with endpoint, database_name, credential, or env_file options." + ) + if create_database: + raise ValueError("create_database cannot be used with an injected database_client.") + resolved_name = getattr(database_client, "id", None) + if not isinstance(resolved_name, str) or not resolved_name: + raise ValueError("Injected database_client must expose a non-empty string id.") + return _CosmosConnection( + cosmos_client=None, + database_client=database_client, + database_name=resolved_name, + owns_client=False, + create_database=False, + ) + if cosmos_client is not None: + if any(value is not None for value in (endpoint, credential, env_file_path, env_file_encoding)): + raise ValueError("cosmos_client cannot be combined with endpoint, credential, or env_file options.") + if database_name is None: + raise ValueError("database_name is required with an injected cosmos_client.") + return _CosmosConnection( + cosmos_client=cosmos_client, + database_client=None, + database_name=_validate_resource_name(database_name, "database"), + owns_client=False, + create_database=create_database, + ) + client, resolved_database, _ = _load_connection_settings( + endpoint=endpoint, + database_name=database_name, + container_name=None, + credential=credential, + require_container=False, + env_file_path=env_file_path, + env_file_encoding=env_file_encoding, + ) + return _CosmosConnection( + cosmos_client=client, + database_client=None, + database_name=resolved_database, + owns_client=True, + create_database=create_database, + ) + + +def _property_access(name: str) -> str: + if not isinstance(name, str) or not name or "\x00" in name: + raise ValueError("Cosmos field storage names must be non-empty strings without NUL.") + return f"c[{json.dumps(name, ensure_ascii=True)}]" + + +def _object_projection(names: Sequence[str]) -> str: + return "{" + ", ".join(f"{json.dumps(name, ensure_ascii=True)}: {_property_access(name)}" for name in names) + "}" + + +def _policy_path(name: str, suffix: str = "") -> str: + if _VECTOR_NAME.fullmatch(name): + return f"/{name}{suffix}" + escaped = json.dumps(name, ensure_ascii=True) + return f"/{escaped}{suffix}" + + +def _normalize_policy_path(path: Any) -> str: + if not isinstance(path, str): + raise ValueError("Cosmos policy paths must be strings.") + if not path.startswith('/"'): + return path + escaped = False + for index in range(2, len(path)): + char = path[index] + if char == '"' and not escaped: + segment = json.loads(path[1 : index + 1]) + suffix = path[index + 1 :] + if isinstance(segment, str) and _VECTOR_NAME.fullmatch(segment): + return f"/{segment}{suffix}" + return f"/{json.dumps(segment, ensure_ascii=True)}{suffix}" + escaped = char == "\\" and not escaped + if char != "\\": + escaped = False + raise ValueError(f"Invalid Cosmos policy path '{path}'.") + + +def _field_annotations(field: VectorStoreField) -> dict[str, Any]: + raw = field.provider_annotations.get("azure_cosmos") + if raw is None: + return {} + if not isinstance(raw, Mapping): + raise TypeError("Vector field azure_cosmos annotations must be a mapping.") + options = dict(cast(Mapping[str, Any], raw)) + if any(not isinstance(name, str) for name in options): + raise TypeError("Vector field azure_cosmos annotation names must be strings.") + if unknown := options.keys() - _FIELD_ANNOTATIONS: + raise ValueError(f"Unsupported Azure Cosmos DB field annotation(s): {', '.join(sorted(unknown))}.") + return options + + +def _prepare_vector_config(field: VectorStoreField) -> dict[str, Any]: + storage_name = field.storage_name or field.name + if not _VECTOR_NAME.fullmatch(storage_name): + raise ValueError("Cosmos vector storage names must be top-level ASCII identifiers.") + options = _field_annotations(field) + declared_type = options.pop("data_type", field.type_) + if declared_type is not None and not isinstance(declared_type, str): + raise TypeError("Cosmos vector data_type must be a string.") + if declared_type not in _VECTOR_TYPES: + raise NotImplementedError(f"Cosmos vector field '{field.name}' requires float32, int8, or uint8 elements.") + data_type = _VECTOR_TYPES[declared_type] + try: + distance = _DISTANCES[field.distance_function or "DEFAULT"] + except KeyError: + raise NotImplementedError(f"Unsupported Cosmos vector distance '{field.distance_function}'.") from None + try: + index_kind = _INDEX_KINDS[field.index_kind or "default"] + except KeyError: + raise NotImplementedError(f"Unsupported Cosmos vector index kind '{field.index_kind}'.") from None + dimensions = field.dimensions + maximum = 505 if index_kind == "flat" else 4096 + if dimensions is None or dimensions > maximum: + raise ValueError(f"Cosmos {index_kind} vector field '{field.name}' supports at most {maximum} dimensions.") + index_options: dict[str, Any] = {} + quantizer_type = options.pop("quantizer_type", None) + if quantizer_type is not None: + if index_kind == "flat": + raise ValueError("quantizer_type is only supported by quantizedFlat and diskANN indexes.") + if quantizer_type not in ("product", "spherical"): + raise ValueError("quantizer_type must be 'product' or 'spherical'.") + index_options["quantizerType"] = quantizer_type + quantization_bytes = options.pop("quantization_byte_size", None) + if quantization_bytes is not None: + if index_kind == "flat": + raise ValueError("quantization_byte_size is only supported by quantizedFlat and diskANN indexes.") + maximum_quantization_bytes = min(512, dimensions) + if type(quantization_bytes) is not int or not 4 <= quantization_bytes <= maximum_quantization_bytes: + raise ValueError(f"quantization_byte_size must be an integer between 4 and {maximum_quantization_bytes}.") + index_options["quantizationByteSize"] = quantization_bytes + indexing_list_size = options.pop("indexing_search_list_size", None) + if indexing_list_size is not None: + if index_kind != "diskANN": + raise ValueError("indexing_search_list_size is only supported by diskANN indexes.") + if type(indexing_list_size) is not int or not 25 <= indexing_list_size <= 500: + raise ValueError("indexing_search_list_size must be an integer between 25 and 500.") + index_options["indexingSearchListSize"] = indexing_list_size + return { + "field": field, + "storage_name": storage_name, + "path": f"/{storage_name}", + "data_type": data_type, + "distance": distance, + "index_kind": index_kind, + "index_options": index_options, + } + + +def _prepare_schema( + definition: VectorStoreCollectionDefinition, +) -> tuple[dict[str, Any], dict[str, Any], dict[str, dict[str, Any]]]: + key = definition.key_field + if key.storage_name != "id" and key.name != "id": + raise ValueError("Cosmos vector collections require the key storage name 'id'.") + if definition.key_field_storage_name != "id": + raise ValueError("Cosmos vector collections require the key storage name 'id'.") + if key.type_ != "str": + raise ValueError("Cosmos vector collections require string keys.") + if key.is_auto_generated: + raise NotImplementedError("Cosmos vector collections require application-provided keys.") + if not definition.vector_fields: + raise ValueError("Cosmos vector collections require at least one vector field.") + configs: dict[str, dict[str, Any]] = {} + embeddings: list[dict[str, Any]] = [] + vector_indexes: list[dict[str, Any]] = [] + excluded_paths: list[dict[str, str]] = [{"path": "/_etag/?"}] + for field in definition.fields: + if field.field_type != "vector": + if field.provider_annotations.get("azure_cosmos") is not None: + raise ValueError("azure_cosmos field annotations are supported only on vector fields.") + if field.is_full_text_indexed: + raise NotImplementedError("CosmosCollection does not provide full-text or hybrid search.") + if field.field_type == "data" and field.is_indexed is False: + excluded_paths.append({"path": _policy_path(field.storage_name or field.name, "/*")}) + continue + config = _prepare_vector_config(field) + configs[field.name] = config + embeddings.append({ + "path": config["path"], + "dataType": config["data_type"], + "distanceFunction": config["distance"], + "dimensions": field.dimensions, + }) + vector_indexes.append({ + "path": config["path"], + "type": config["index_kind"], + **config["index_options"], + }) + excluded_paths.append({"path": f"{config['path']}/*"}) + vector_policy = {"vectorEmbeddings": embeddings} + indexing_policy = { + "indexingMode": "consistent", + "automatic": True, + "includedPaths": [{"path": "/*"}], + "excludedPaths": excluded_paths, + "vectorIndexes": vector_indexes, + } + return vector_policy, indexing_policy, configs + + +def _required_indexed_paths(definition: VectorStoreCollectionDefinition) -> set[str]: + paths: set[str] = set() + for field in definition.data_fields: + if field.is_indexed is False: + continue + storage_name = field.storage_name or field.name + paths.add(_normalize_policy_path(_policy_path(storage_name, "/*"))) + paths.add(_normalize_policy_path(_policy_path(storage_name, "/?"))) + return paths + + +def _policy_entries(policy: Mapping[str, Any], name: str) -> list[Mapping[str, Any]]: + entries = policy.get(name) + if not isinstance(entries, list): + raise ValueError(f"Existing Cosmos policy has an invalid '{name}' value.") + typed_entries: list[Any] = entries # pyright: ignore[reportUnknownVariableType] + if not all(isinstance(entry, Mapping) for entry in typed_entries): + raise ValueError(f"Existing Cosmos policy has an invalid '{name}' value.") + return [cast(Mapping[str, Any], entry) for entry in typed_entries] + + +def _validate_existing_policies( + properties: Mapping[str, Any], + *, + vector_policy: Mapping[str, Any], + indexing_policy: Mapping[str, Any], + required_indexed_paths: set[str], +) -> None: + partition = properties.get("partitionKey") + if not isinstance(partition, Mapping): + raise ValueError("Existing Cosmos container is missing its partition-key policy.") + typed_partition = cast(Mapping[str, Any], partition) + paths = typed_partition.get("paths") + kind = typed_partition.get("kind") + if paths != ["/id"] or not isinstance(kind, str) or kind.lower() != "hash": + raise ValueError("Existing Cosmos container must use the single Hash partition key path '/id'.") + + actual_vector = properties.get("vectorEmbeddingPolicy") + if not isinstance(actual_vector, Mapping): + raise ValueError("Existing Cosmos container has no vector embedding policy.") + expected_embeddings = { + ( + _normalize_policy_path(entry.get("path")), + str(entry.get("dataType", "")).lower(), + str(entry.get("distanceFunction", "")).lower(), + entry.get("dimensions"), + ) + for entry in _policy_entries(vector_policy, "vectorEmbeddings") + } + actual_embeddings = { + ( + _normalize_policy_path(entry.get("path")), + str(entry.get("dataType", "")).lower(), + str(entry.get("distanceFunction", "")).lower(), + entry.get("dimensions"), + ) + for entry in _policy_entries(cast(Mapping[str, Any], actual_vector), "vectorEmbeddings") + } + if actual_embeddings != expected_embeddings: + raise ValueError("Existing Cosmos vector embedding policy is incompatible with the collection definition.") + + actual_indexing = properties.get("indexingPolicy") + if not isinstance(actual_indexing, Mapping): + raise ValueError("Existing Cosmos container has no indexing policy.") + typed_indexing = cast(Mapping[str, Any], actual_indexing) + mode = typed_indexing.get("indexingMode", "consistent") + if not isinstance(mode, str) or mode.lower() != "consistent" or typed_indexing.get("automatic", True) is not True: + raise ValueError("Existing Cosmos container must use automatic consistent indexing.") + included = {_normalize_policy_path(entry.get("path")) for entry in _policy_entries(typed_indexing, "includedPaths")} + if "/*" not in included: + raise ValueError("Existing Cosmos indexing policy must include the root path '/*'.") + expected_excluded = { + _normalize_policy_path(entry.get("path")) for entry in _policy_entries(indexing_policy, "excludedPaths") + } + actual_excluded = { + _normalize_policy_path(entry.get("path")) for entry in _policy_entries(typed_indexing, "excludedPaths") + } + if not expected_excluded <= actual_excluded: + raise ValueError("Existing Cosmos indexing policy does not exclude all required vector paths.") + if actual_excluded & required_indexed_paths: + raise ValueError("Existing Cosmos indexing policy excludes a data field that must remain indexed.") + + expected_indexes = { + _normalize_policy_path(entry.get("path")): entry for entry in _policy_entries(indexing_policy, "vectorIndexes") + } + actual_indexes = { + _normalize_policy_path(entry.get("path")): entry for entry in _policy_entries(typed_indexing, "vectorIndexes") + } + if actual_indexes.keys() != expected_indexes.keys(): + raise ValueError("Existing Cosmos vector index paths are incompatible with the collection definition.") + for path, expected in expected_indexes.items(): + actual = actual_indexes[path] + if str(actual.get("type", "")).lower() != str(expected.get("type", "")).lower(): + raise ValueError(f"Existing Cosmos vector index type for '{path}' is incompatible.") + for option in ("quantizerType", "quantizationByteSize", "indexingSearchListSize"): + if option in expected and actual.get(option) != expected[option]: + raise ValueError(f"Existing Cosmos vector index option '{option}' for '{path}' is incompatible.") + actual_quantizer = actual.get("quantizerType") + if "quantizerType" not in expected and actual_quantizer not in (None, "product"): + raise ValueError(f"Existing Cosmos vector index quantizer for '{path}' is incompatible.") + + +def _validate_key(value: Any) -> str: + if not isinstance(value, str): + raise TypeError("Cosmos item keys must be strings.") + size = len(value.encode("utf-8")) + if not value or size > _ID_BYTE_LIMIT or any(char in value for char in "/\\?#"): + raise ValueError("Cosmos item keys must contain 1-1023 UTF-8 bytes and cannot contain '/', '\\', '?', or '#'.") + return value + + +def _validate_json(value: Any, *, path: str, depth: int = 0) -> None: + if depth > _MAX_JSON_DEPTH: + raise ValueError(f"{path} exceeds Cosmos DB's maximum JSON nesting depth of {_MAX_JSON_DEPTH}.") + if value is None or isinstance(value, (str, bool)): + return + if isinstance(value, int): + if not -_MAX_SAFE_INTEGER <= value <= _MAX_SAFE_INTEGER: + raise ValueError(f"{path} must fit exactly in an IEEE 754 binary64 JSON number.") + return + if isinstance(value, float): + if not math.isfinite(value): + raise ValueError(f"{path} must contain only finite JSON numbers.") + return + if isinstance(value, list): + items: list[Any] = value # pyright: ignore[reportUnknownVariableType] + for index, item in enumerate(items): + _validate_json(item, path=f"{path}[{index}]", depth=depth + 1) + return + if isinstance(value, dict): + mapping: dict[Any, Any] = value # pyright: ignore[reportUnknownVariableType] + for key, item in mapping.items(): + if not isinstance(key, str): + raise TypeError(f"{path} object keys must be strings.") + _validate_json(item, path=f"{path}.{key}", depth=depth + 1) + return + raise TypeError(f"{path} must be JSON-compatible, not {type(value).__name__}.") + + +def _validate_data_field(field: VectorStoreField, value: Any) -> None: + if value is None or field.type_ is None: + return + kind = field.type_ + valid = ( + (kind == "str" and isinstance(value, str)) + or (kind == "bool" and type(value) is bool) + or (kind == "int" and type(value) is int) + or (kind == "float" and type(value) in (int, float)) + or (kind in ("list", "tuple", "set", "Sequence") and isinstance(value, list)) + or (kind == "dict" and isinstance(value, dict)) + ) + if not valid: + raise TypeError(f"Cosmos field '{field.name}' requires a value of declared type '{kind}'.") + + +def _validate_vector(value: Any, config: dict[str, Any]) -> None: + field = cast(VectorStoreField, config["field"]) + data_type = cast(str, config["data_type"]) + if not isinstance(value, Sequence) or isinstance(value, (str, bytes, bytearray)): + raise TypeError(f"Cosmos vector '{field.name}' must be a dense numeric sequence.") + typed_value: Sequence[Any] = value # pyright: ignore[reportUnknownVariableType] + if len(typed_value) != field.dimensions: + raise ValueError(f"Cosmos vector '{field.name}' requires exactly {field.dimensions} dimensions.") + minimum, maximum, integral = _VECTOR_RANGES[data_type] + for item in typed_value: + if isinstance(item, bool) or not isinstance(item, (int, float)) or (integral and type(item) is not int): + raise TypeError( + f"Cosmos {data_type} vector '{field.name}' requires " + f"{'integer' if integral else 'numeric'} elements without booleans." + ) + if not math.isfinite(item) or not minimum <= item <= maximum: + raise ValueError( + f"Cosmos {data_type} vector '{field.name}' values must be finite and within [{minimum}, {maximum}]." + ) + + +def _filter_value_compatible(field: VectorStoreField, value: Any) -> bool: + if field.field_type == "key": + return isinstance(value, str) + return ( + field.type_ is None + or (field.type_ == "str" and isinstance(value, str)) + or (field.type_ == "bool" and type(value) is bool) + or (field.type_ == "int" and type(value) is int) + or (field.type_ == "float" and type(value) in (int, float)) + ) + + +def _validate_filter_scalar(value: Any, *, allow_none: bool = False) -> None: + if value is None and allow_none: + return + if not isinstance(value, _FILTER_SCALAR_TYPES): + raise NotImplementedError("Cosmos filters support only scalar string, boolean, and numeric literals.") + if isinstance(value, float) and not math.isfinite(value): + raise ValueError("Cosmos numeric filter values must be finite.") + if isinstance(value, int) and not isinstance(value, bool) and not -_MAX_SAFE_INTEGER <= value <= _MAX_SAFE_INTEGER: + raise ValueError("Cosmos integer filter values must fit exactly in IEEE 754 binary64.") + + +def _add_parameter(parameters: list[dict[str, Any]], value: Any, prefix: str = "filter") -> str: + name = f"@{prefix}_{len(parameters)}" + parameters.append({"name": name, "value": value}) + return name + + +def _query_metadata_hook(metadata: dict[str, Any]) -> Callable[[Mapping[str, str], Any], None]: + def response_hook(headers: Mapping[str, str], _: Any) -> None: + charge = headers.get("x-ms-request-charge") + if charge is not None: + try: + metadata["request_charge"] = float(metadata["request_charge"]) + float(charge) + except (TypeError, ValueError) as exc: + raise IntegrationInvalidResponseException( + "Cosmos query returned an invalid request-charge header." + ) from exc + activity_id = headers.get("x-ms-activity-id") + if activity_id: + metadata["activity_id"] = activity_id + metadata["has_more_results"] = bool(headers.get("x-ms-continuation")) + + return response_hook + + +async def _skip_results( + results: AsyncIterable[Mapping[str, Any]], + skip: int, +) -> AsyncIterator[Mapping[str, Any]]: + index = 0 + async for result in results: + if index < skip: + index += 1 + continue + yield result + + +class CosmosCollection(BaseVectorCollection[str, ModelT], BaseVectorSearch[str, ModelT], Generic[ModelT]): + """An Azure Cosmos DB for NoSQL container with vector indexing and search.""" + + supported_key_types: ClassVar[set[str] | None] = {"str"} + supported_vector_types: ClassVar[set[str] | None] = {"float", "float32", "int8", "uint8"} + supported_search_types: ClassVar[set[SearchType]] = {"vector"} + + def __init__( + self, + record_type: type[ModelT], + *, + definition: VectorStoreCollectionDefinition | None = None, + collection_name: str | None = None, + embedding_generator: EmbeddingClient | None = None, + endpoint: str | None = None, + database_name: str | None = None, + credential: str | SecretString | AzureCredentialTypes | None = None, + cosmos_client: CosmosClient | None = None, + database_client: DatabaseProxy | None = None, + container_client: ContainerProxy | None = None, + create_database: bool = False, + env_file_path: str | None = None, + env_file_encoding: str | None = None, + _connection: _CosmosConnection | None = None, + ) -> None: + """Configure a Cosmos vector collection without contacting the service. + + Args: + record_type: Registered vector model type, or dict with an explicit definition. + definition: Optional collection definition for dictionary records. + collection_name: Container name, or ``AZURE_COSMOS_CONTAINER_NAME``. + embedding_generator: Optional local embedding client. + endpoint: Cosmos account endpoint, or ``AZURE_COSMOS_ENDPOINT``. + database_name: Database name, or ``AZURE_COSMOS_DATABASE_NAME``. + credential: Caller-owned Azure credential or key, falling back to ``AZURE_COSMOS_KEY``. + cosmos_client: Caller-owned asynchronous Cosmos account client. + database_client: Caller-owned asynchronous Cosmos database proxy. + container_client: Caller-owned asynchronous Cosmos container proxy. + create_database: Allow explicit database creation during the first service operation. + env_file_path: Optional settings file used only when no SDK client is injected. + env_file_encoding: Settings file encoding. + """ + if sum(client is not None for client in (cosmos_client, database_client, container_client)) > 1: + raise ValueError("Provide at most one of cosmos_client, database_client, or container_client.") + if _connection is not None and any( + value is not None + for value in ( + endpoint, + database_name, + credential, + cosmos_client, + database_client, + container_client, + env_file_path, + env_file_encoding, + ) + ): + raise ValueError("A store-provided connection cannot be combined with connection settings or clients.") + if _connection is not None and create_database: + raise ValueError("Store-created collections inherit database ownership from the store.") + + resolved_name = collection_name + owns_connection = False + if container_client is not None: + if ( + any( + value is not None + for value in (endpoint, database_name, credential, env_file_path, env_file_encoding) + ) + or create_database + ): + raise ValueError( + "container_client cannot be combined with connection, database, env_file, or creation options." + ) + client_name = getattr(container_client, "id", None) + if isinstance(client_name, str) and client_name: + if resolved_name is not None and resolved_name != client_name: + raise ValueError("collection_name must match the injected container_client id.") + resolved_name = client_name + self._connection = None + elif _connection is not None: + self._connection = _connection + elif cosmos_client is not None or database_client is not None: + self._connection = _connection_from_clients( + endpoint=endpoint, + database_name=database_name, + credential=credential, + cosmos_client=cosmos_client, + database_client=database_client, + create_database=create_database, + env_file_path=env_file_path, + env_file_encoding=env_file_encoding, + ) + else: + registered_definition = definition or getattr(record_type, "__vectorstoremodel_definition__", None) + configured_name = resolved_name or getattr(registered_definition, "collection_name", None) + client, resolved_database, settings_name = _load_connection_settings( + endpoint=endpoint, + database_name=database_name, + container_name=configured_name, + credential=credential, + require_container=True, + env_file_path=env_file_path, + env_file_encoding=env_file_encoding, + ) + resolved_name = settings_name + self._connection = _CosmosConnection( + cosmos_client=client, + database_client=None, + database_name=resolved_database, + owns_client=True, + create_database=create_database, + ) + owns_connection = True + + super().__init__( + record_type, + definition=definition, + collection_name=resolved_name, + embedding_generator=embedding_generator, + managed_client=owns_connection, + ) + _validate_resource_name(self.collection_name, "container") + self._vector_policy, self._indexing_policy, self._vector_configs = _prepare_schema(self.definition) + self._container_client = container_client + self._container_validated = False + self._owns_connection = owns_connection + self._closed = False + self._on_close: Callable[[CosmosCollection[ModelT]], None] | None = None + self._on_delete: Callable[[str], None] | None = None + + def _require_open(self) -> None: + if self._closed: + raise RuntimeError("Cosmos collection is closed.") + if self._connection is not None: + self._connection.ensure_open() + mark_feature_used(FeatureIndex.AZURE_COSMOS) + + async def _get_database(self) -> DatabaseProxy: + self._require_open() + if self._connection is None: + raise NotImplementedError("ContainerProxy injection does not provide database lifecycle operations.") + return await self._connection.get_database() + + async def _read_and_validate_container(self, container: ContainerProxy) -> ContainerProxy: + properties = await container.read() + _validate_existing_policies( + cast(Mapping[str, Any], properties), + vector_policy=self._vector_policy, + indexing_policy=self._indexing_policy, + required_indexed_paths=_required_indexed_paths(self.definition), + ) + self._container_validated = True + return container + + async def _get_container(self) -> ContainerProxy: + self._require_open() + if self._container_client is None: + database = await self._get_database() + self._container_client = database.get_container_client(self.collection_name) + if not self._container_validated: + await self._read_and_validate_container(self._container_client) + return self._container_client + + async def collection_exists(self, *, operation_options: Mapping[str, Any] | None = None) -> bool: + """Check whether the configured container exists.""" + _validate_operation_options(operation_options, set()) + self._require_open() + container = self._container_client + if container is None: + database = await self._get_database() + container = database.get_container_client(self.collection_name) + try: + await container.read() + except CosmosResourceNotFoundError: + return False + return True + + async def ensure_collection_exists(self, *, operation_options: Mapping[str, Any] | None = None) -> None: + """Create an absent container and validate an existing container without updating it.""" + _validate_operation_options(operation_options, set()) + self._require_open() + if self._container_client is not None and self._connection is None: + await self._read_and_validate_container(self._container_client) + return + database = await self._get_database() + container = database.get_container_client(self.collection_name) + try: + await container.read() + except CosmosResourceNotFoundError: + try: + container = await database.create_container( + id=self.collection_name, + partition_key=PartitionKey(path="/id"), + indexing_policy=self._indexing_policy, + vector_embedding_policy=self._vector_policy, + ) + except CosmosResourceExistsError: + container = database.get_container_client(self.collection_name) + self._container_client = await self._read_and_validate_container(container) + + async def ensure_collection_deleted(self, *, operation_options: Mapping[str, Any] | None = None) -> None: + """Delete the configured container when a database proxy is available.""" + _validate_operation_options(operation_options, set()) + database = await self._get_database() + with suppress(CosmosResourceNotFoundError): + await database.delete_container(self.collection_name) + self._invalidate_container() + if self._on_delete is not None: + self._on_delete(self.collection_name) + + def _invalidate_container(self) -> None: + self._container_client = None + self._container_validated = False + + def _serialize_dicts_to_store_models( + self, + records: Sequence[dict[str, Any]], + *, + context: Mapping[str, Any] | None = None, + ) -> Sequence[Any]: + del context + storage_fields = {field.storage_name or field.name: field for field in self.definition.fields} + for index, record in enumerate(records): + missing = storage_fields.keys() - record.keys() + if missing: + raise ValueError( + f"Cosmos record at position {index} is missing field(s): {', '.join(sorted(missing))}." + ) + _validate_key(record["id"]) + for storage_name, field in storage_fields.items(): + value = record[storage_name] + if field.field_type == "vector": + if value is not None: + _validate_vector(value, self._vector_configs[field.name]) + else: + _validate_data_field(field, value) + _validate_json(value, path=f"record[{index}].{storage_name}") + encoded = json.dumps( + record, + cls=AzureJSONEncoder, + ensure_ascii=True, + allow_nan=False, + separators=(",", ":"), + ).encode("utf-8") + if len(encoded) > _ITEM_SIZE_LIMIT: + raise ValueError(f"Cosmos record at position {index} exceeds the 2 MiB item limit.") + return records + + def _project_item(self, item: Mapping[str, Any], include_vectors: bool) -> dict[str, Any]: + names = self.definition.get_storage_names(include_vector_fields=include_vectors) + try: + return {name: item[name] for name in names} + except KeyError as exc: + raise IntegrationInvalidResponseException( + f"Cosmos response is missing required field '{exc.args[0]}'." + ) from exc + + async def _inner_upsert( + self, + records: Sequence[Any], + *, + operation_options: Mapping[str, Any] | None = None, + ) -> Sequence[str]: + _validate_operation_options(operation_options, set()) + container = await self._get_container() + typed_records = cast(Sequence[dict[str, Any]], records) + keys = [cast(str, record["id"]) for record in typed_records] + for index, record in enumerate(typed_records): + try: + await container.upsert_item(body=record) + except CosmosHttpResponseError as exc: + raise IntegrationException( + f"Cosmos upsert partially completed {index}/{len(records)} records; " + f"the failure corresponds to input index {index}. Retry with the same application keys." + ) from exc + return keys + + def _resolve_filter_field(self, expression: Any) -> VectorStoreField: + if "." in expression.field_name: + raise NotImplementedError("CosmosCollection does not support nested portable filter paths.") + field = self.definition.try_get_field(expression.field_name) + if field is None: + raise ValueError(f"Unknown Cosmos filter field '{expression.field_name}'.") + if field.field_type == "vector": + raise NotImplementedError("Cosmos vector fields cannot be used in portable filters.") + if field.field_type != "key" and field.is_indexed is False: + raise ValueError(f"Cosmos filter field '{field.name}' is excluded from indexing.") + return field + + def _translate_filter( + self, + expression: FilterExpression, + parameters: list[dict[str, Any]], + ) -> str: + if isinstance(expression, FilterGroup): + children = [self._translate_filter(child, parameters) for child in expression.filters] + if expression.operator == "not": + return f"(NOT {children[0]})" + delimiter = " AND " if expression.operator == "and" else " OR " + return "(" + delimiter.join(children) + ")" + + field = self._resolve_filter_field(expression) + access = _property_access(field.storage_name or field.name) + operator = expression.operator + value = expression.value + if operator == "exists": + return f"IS_DEFINED({access})" + if operator == "is_null": + return f"(IS_DEFINED({access}) AND IS_NULL({access}))" + if operator == "is_not_null": + return f"(IS_DEFINED({access}) AND NOT IS_NULL({access}))" + if operator in ("eq", "ne"): + _validate_filter_scalar(value) + if not _filter_value_compatible(field, value): + return "false" if operator == "eq" else f"IS_DEFINED({access})" + parameter = _add_parameter(parameters, value) + if operator == "eq": + return f"(IS_DEFINED({access}) AND {access} = {parameter})" + return f"(IS_DEFINED({access}) AND (IS_NULL({access}) OR {access} != {parameter}))" + if operator in ("gt", "gte", "lt", "lte", "between"): + if field.type_ not in ("str", "int", "float"): + raise NotImplementedError("Cosmos ordered filters require a declared string or numeric field.") + values = list(value) if operator == "between" else [value] + for item in values: + _validate_filter_scalar(item) + if not _filter_value_compatible(field, item): + raise TypeError(f"Cosmos ordered filter value is incompatible with field '{field.name}'.") + if operator == "between": + lower = _add_parameter(parameters, values[0]) + upper = _add_parameter(parameters, values[1]) + return ( + f"(IS_DEFINED({access}) AND NOT IS_NULL({access}) AND {access} >= {lower} AND {access} <= {upper})" + ) + parameter = _add_parameter(parameters, value) + sql_operator = {"gt": ">", "gte": ">=", "lt": "<", "lte": "<="}[operator] + return f"(IS_DEFINED({access}) AND NOT IS_NULL({access}) AND {access} {sql_operator} {parameter})" + if operator in ("in", "not_in"): + values = list(cast(Sequence[Any], value)) + for item in values: + _validate_filter_scalar(item, allow_none=True) + if not values: + return "false" if operator == "in" else f"(IS_DEFINED({access}) AND NOT IS_NULL({access}))" + parameter = _add_parameter(parameters, values) + contained = f"ARRAY_CONTAINS({parameter}, {access})" + if operator == "in": + return f"(IS_DEFINED({access}) AND NOT IS_NULL({access}) AND {contained})" + return f"(IS_DEFINED({access}) AND NOT IS_NULL({access}) AND NOT {contained})" + if operator in ("contains", "contains_any", "contains_all"): + if field.type_ not in ("list", "tuple", "set", "Sequence"): + raise NotImplementedError("Cosmos collection membership filters require a declared collection field.") + values = [value] if operator == "contains" else list(cast(Sequence[Any], value)) + for item in values: + _validate_filter_scalar(item, allow_none=True) + if not values: + return f"IS_ARRAY({access})" if operator == "contains_all" else "false" + clauses = [f"ARRAY_CONTAINS({access}, {_add_parameter(parameters, item)})" for item in values] + delimiter = " AND " if operator == "contains_all" else " OR " + return f"(IS_ARRAY({access}) AND ({delimiter.join(clauses)}))" + if operator in ("starts_with", "ends_with", "contains_text"): + if field.type_ != "str": + raise NotImplementedError("Cosmos string filters require a declared string field.") + if not isinstance(value, str): + raise TypeError("Cosmos string filter values must be strings.") + parameter = _add_parameter(parameters, value) + function = { + "starts_with": "STARTSWITH", + "ends_with": "ENDSWITH", + "contains_text": "CONTAINS", + }[operator] + return f"(IS_STRING({access}) AND {function}({access}, {parameter}))" + raise NotImplementedError(f"CosmosCollection does not support filter operator '{operator}'.") + + def _prepare_filter(self, expression: FilterExpression | None) -> tuple[str | None, list[dict[str, Any]]]: + if expression is None: + return None, [] + parameters: list[dict[str, Any]] = [] + return self._translate_filter(expression, parameters), parameters + + def _prepare_order_by(self, order_by: Mapping[str, bool] | None) -> str | None: + if not order_by: + return None + if len(order_by) > 1: + raise NotImplementedError("CosmosCollection supports one order_by field without a composite index.") + name, ascending = next(iter(order_by.items())) + if not isinstance(ascending, bool): + raise TypeError(f"Order direction for field '{name}' must be a boolean.") + field = self.definition.try_get_field(name) + if field is None: + raise ValueError(f"Unknown Cosmos order_by field '{name}'.") + if field.field_type == "vector" or (field.field_type != "key" and field.is_indexed is False): + raise ValueError(f"Cosmos order_by field '{name}' must be an indexed key or data field.") + return f"{_property_access(field.storage_name or field.name)} {'ASC' if ascending else 'DESC'}" + + async def _inner_get( + self, + *, + keys: Sequence[str] | None = None, + filter: FilterExpression | None = None, + top: int = 10, + skip: int = 0, + order_by: Mapping[str, bool] | None = None, + include_vectors: bool = False, + operation_options: Mapping[str, Any] | None = None, + ) -> Sequence[Any]: + _validate_operation_options(operation_options, set()) + if keys is not None: + if order_by or skip: + raise ValueError("Cosmos point reads preserve input order and cannot use order_by or skip.") + validated_keys = [_validate_key(key) for key in keys] + container = await self._get_container() + records: list[dict[str, Any]] = [] + for key in validated_keys: + try: + item = await container.read_item(item=key, partition_key=key) + except CosmosResourceNotFoundError: + continue + records.append(self._project_item(cast(Mapping[str, Any], item), include_vectors)) + return records + + where, parameters = self._prepare_filter(filter) + order = self._prepare_order_by(order_by) + names = self.definition.get_storage_names(include_vector_fields=include_vectors) + # Every composed expression comes from the validated collection definition. + query = f"SELECT VALUE {_object_projection(names)} FROM c" # nosec B608 # ruff: ignore[hardcoded-sql-expression] + if where is not None: + query += f" WHERE {where}" + if order is not None: + query += f" ORDER BY {order}" + parameters.extend([ + {"name": "@skip", "value": skip}, + {"name": "@top", "value": top}, + ]) + query += " OFFSET @skip LIMIT @top" + container = await self._get_container() + items = cast( + AsyncIterable[Mapping[str, Any]], + container.query_items(query=query, parameters=parameters), + ) + return [dict(item) async for item in items] + + async def _inner_delete( + self, + keys: Sequence[str], + *, + operation_options: Mapping[str, Any] | None = None, + ) -> None: + _validate_operation_options(operation_options, set()) + validated_keys = [_validate_key(key) for key in keys] + container = await self._get_container() + for index, key in enumerate(validated_keys): + try: + await container.delete_item(item=key, partition_key=key) + except CosmosResourceNotFoundError: + continue + except CosmosHttpResponseError as exc: + raise IntegrationException( + f"Cosmos delete partially completed {index}/{len(keys)} records; " + f"the failure corresponds to input index {index}. Retry with the same application keys." + ) from exc + + def _prepare_search_options( + self, + config: dict[str, Any], + operation_options: Mapping[str, Any] | None, + parameters: list[dict[str, Any]], + ) -> str: + options = _validate_operation_options(operation_options, _SEARCH_OPTIONS) + brute_force = options.pop("brute_force", False) + if not isinstance(brute_force, bool): + raise TypeError("brute_force must be a boolean.") + native_options: dict[str, Any] = {} + for name, native_name in ( + ("search_list_size_multiplier", "searchListSizeMultiplier"), + ("quantized_vector_list_multiplier", "quantizedVectorListMultiplier"), + ): + value = options.pop(name, None) + if value is not None: + if type(value) is not int or value <= 0: + raise ValueError(f"{name} must be a positive integer.") + if name == "search_list_size_multiplier" and config["index_kind"] != "diskANN": + raise ValueError("search_list_size_multiplier requires a diskANN vector index.") + if config["index_kind"] == "flat": + raise ValueError(f"{name} is not supported by a flat vector index.") + native_options[native_name] = value + filter_priority = options.pop("filter_priority", None) + if filter_priority is not None: + if ( + type(filter_priority) not in (int, float) + or not math.isfinite(filter_priority) + or not 0 <= filter_priority <= 1 + ): + raise ValueError("filter_priority must be a finite number between 0 and 1.") + if config["index_kind"] != "diskANN": + raise ValueError("filter_priority requires a diskANN vector index.") + native_options["filterPriority"] = filter_priority + if not native_options and not brute_force: + return "" + brute_force_parameter = _add_parameter(parameters, brute_force, "brute_force") + options_parameter = _add_parameter(parameters, native_options, "vector_options") + return f", {brute_force_parameter}, {options_parameter}" + + async def _inner_search( + self, + *, + search_type: SearchType, + filter: FilterExpression | None = None, + values: Any | None = None, + vector: Vector | None = None, + top: int = 3, + skip: int = 0, + include_vectors: bool = False, + vector_property_name: str | None = None, + additional_property_name: str | None = None, + score_threshold: float | None = None, + operation_options: Mapping[str, Any] | None = None, + ) -> SearchResults[Any]: + del values + if search_type != "vector": + raise NotImplementedError("CosmosCollection supports vector search only.") + if additional_property_name is not None: + raise ValueError("additional_property_name is only supported for keyword-hybrid search.") + field = self.definition.try_get_vector_field(vector_property_name) + if field is None: + raise ValueError("Cosmos vector search requires a configured vector field.") + if vector is None: + raise NotImplementedError("CosmosCollection does not support server-side embedding generation.") + config = self._vector_configs[field.name] + if score_threshold is not None and config["distance"] == "euclidean": + raise NotImplementedError( + "Euclidean score_threshold is not supported because direct Cosmos VectorDistance predicates " + "do not reliably enforce Euclidean cutoffs. Omit score_threshold to use Euclidean search." + ) + _validate_vector(vector, config) + parameters: list[dict[str, Any]] = [{"name": "@vector", "value": list(vector)}] + option_arguments = self._prepare_search_options(config, operation_options, parameters) + distance = f"VectorDistance({_property_access(config['storage_name'])}, @vector{option_arguments})" + where, filter_parameters = self._prepare_filter(filter) + parameters.extend(filter_parameters) + if score_threshold is not None: + if type(score_threshold) not in (int, float) or not math.isfinite(score_threshold): + raise ValueError("score_threshold must be a finite number.") + threshold_parameter = _add_parameter(parameters, score_threshold, "threshold") + threshold = f"{distance} >= {threshold_parameter}" + where = threshold if where is None else f"({where} AND {threshold})" + names = self.definition.get_storage_names(include_vector_fields=include_vectors) + result_projection = '{"record": ' + _object_projection(names) + f', "score": {distance}' + "}" + parameters.append({"name": "@top", "value": top + skip}) + query = f"SELECT TOP @top VALUE {result_projection} FROM c" # nosec B608 # ruff: ignore[hardcoded-sql-expression] + if where is not None: + query += f" WHERE {where}" + query += f" ORDER BY {distance}" + metadata: dict[str, Any] = { + "score_kind": (f"{config['distance']}_{'distance' if config['distance'] == 'euclidean' else 'similarity'}"), + "score_direction": "lower_is_better" if config["distance"] == "euclidean" else "higher_is_better", + "request_charge": 0.0, + "activity_id": None, + "has_more_results": False, + } + container = await self._get_container() + items = cast( + AsyncIterable[Mapping[str, Any]], + container.query_items( + query=query, + parameters=parameters, + response_hook=_query_metadata_hook(metadata), + ), + ) + return SearchResults(_skip_results(items, skip), metadata=metadata) + + def _get_record_from_result(self, result: Any) -> Any: + if not isinstance(result, Mapping): + raise IntegrationInvalidResponseException("Cosmos vector search returned an invalid record projection.") + typed_result = cast(Mapping[str, Any], result) + record = typed_result.get("record") + if not isinstance(record, Mapping): + raise IntegrationInvalidResponseException("Cosmos vector search returned an invalid record projection.") + return cast(Mapping[str, Any], record) + + def _get_score_from_result(self, result: Any) -> float | None: + if not isinstance(result, Mapping): + raise IntegrationInvalidResponseException("Cosmos vector search returned an invalid result.") + score = cast(Mapping[str, Any], result).get("score") + if not isinstance(score, (int, float)) or isinstance(score, bool) or not math.isfinite(score): + raise IntegrationInvalidResponseException("Cosmos vector search returned a missing or invalid score.") + return float(score) + + async def close(self) -> None: + """Close only a Cosmos client created by this collection.""" + if not self._closed: + self._closed = True + try: + if self._owns_connection and self._connection is not None: + await self._connection.close() + finally: + on_close, self._on_close = self._on_close, None + self._on_delete = None + if on_close is not None: + on_close(self) + + async def __aexit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None: + await self.close() + + +class CosmosStore(BaseVectorStore): + """Factory and database administration for Azure Cosmos DB for NoSQL vector containers.""" + + def __init__( + self, + *, + endpoint: str | None = None, + database_name: str | None = None, + credential: str | SecretString | AzureCredentialTypes | None = None, + cosmos_client: CosmosClient | None = None, + database_client: DatabaseProxy | None = None, + create_database: bool = False, + embedding_generator: EmbeddingClient | None = None, + env_file_path: str | None = None, + env_file_encoding: str | None = None, + ) -> None: + """Create a store from settings or a caller-owned asynchronous SDK client. + + Args: + endpoint: Cosmos account endpoint, or ``AZURE_COSMOS_ENDPOINT``. + database_name: Database name, or ``AZURE_COSMOS_DATABASE_NAME``. + credential: Caller-owned Azure credential or key, falling back to ``AZURE_COSMOS_KEY``. + cosmos_client: Caller-owned asynchronous Cosmos account client. + database_client: Caller-owned asynchronous Cosmos database proxy. + create_database: Allow explicit database creation during the first service operation. + embedding_generator: Default local embedding client for child collections. + env_file_path: Optional settings file used only when no SDK client is injected. + env_file_encoding: Settings file encoding. + """ + connection = _connection_from_clients( + endpoint=endpoint, + database_name=database_name, + credential=credential, + cosmos_client=cosmos_client, + database_client=database_client, + create_database=create_database, + env_file_path=env_file_path, + env_file_encoding=env_file_encoding, + ) + super().__init__(embedding_generator=embedding_generator, managed_client=connection.owns_client) + self._connection = connection + self.database_name = connection.database_name + self._collections: WeakSet[CosmosCollection[Any]] = WeakSet() + self._closed = False + + def _require_open(self) -> None: + if self._closed: + raise RuntimeError("Cosmos store is closed.") + self._connection.ensure_open() + mark_feature_used(FeatureIndex.AZURE_COSMOS) + + def get_collection( + self, + record_type: type[ModelT], + *, + definition: VectorStoreCollectionDefinition | None = None, + collection_name: str | None = None, + embedding_generator: EmbeddingClient | None = None, + ) -> CosmosCollection[ModelT]: + """Create a child collection that borrows the store's resolved database connection.""" + self._require_open() + collection = CosmosCollection( + record_type, + definition=definition, + collection_name=collection_name, + embedding_generator=embedding_generator or self.embedding_generator, + _connection=self._connection, + ) + self._collections.add(collection) + collection._on_close = self._collections.discard # pyright: ignore[reportPrivateUsage] + collection._on_delete = self._invalidate_collections # pyright: ignore[reportPrivateUsage] + return collection + + def _invalidate_collections(self, collection_name: str) -> None: + for collection in list(self._collections): + if collection.collection_name == collection_name: + collection._invalidate_container() # pyright: ignore[reportPrivateUsage] + + async def list_collection_names( + self, + *, + operation_options: Mapping[str, Any] | None = None, + ) -> Sequence[str]: + """List every container in the configured database.""" + _validate_operation_options(operation_options, set()) + self._require_open() + database = await self._connection.get_database() + return [ + item["id"] + async for item in database.list_containers() # pyright: ignore[reportUnknownMemberType] + if isinstance(item, Mapping) and isinstance(item.get("id"), str) + ] + + async def collection_exists( + self, + collection_name: str, + *, + operation_options: Mapping[str, Any] | None = None, + ) -> bool: + """Check one container directly without listing the database.""" + _validate_operation_options(operation_options, set()) + self._require_open() + database = await self._connection.get_database() + try: + await database.get_container_client(collection_name).read() + except CosmosResourceNotFoundError: + return False + return True + + async def _inner_ensure_collection_deleted( + self, + collection_name: str, + *, + operation_options: Mapping[str, Any] | None = None, + ) -> None: + _validate_operation_options(operation_options, set()) + self._require_open() + database = await self._connection.get_database() + with suppress(CosmosResourceNotFoundError): + await database.delete_container(collection_name) + self._invalidate_collections(collection_name) + + async def close(self) -> None: + """Close child collections and the owned Cosmos client, leaving injected objects open.""" + if not self._closed: + self._closed = True + for collection in list(self._collections): + await collection.close() + self._collections.clear() + await self._connection.close() + + async def __aexit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None: + await self.close() diff --git a/python/packages/azure-cosmos/agent_framework_azure_cosmos/py.typed b/python/packages/azure-cosmos/agent_framework_azure_cosmos/py.typed new file mode 100644 index 0000000000..e69de29bb2 diff --git a/python/packages/azure-cosmos/pyproject.toml b/python/packages/azure-cosmos/pyproject.toml index 46850c9e69..2729d40326 100644 --- a/python/packages/azure-cosmos/pyproject.toml +++ b/python/packages/azure-cosmos/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "agent-framework-azure-cosmos" -description = "Azure Cosmos DB history provider integration for Microsoft Agent Framework." +description = "Azure Cosmos DB for NoSQL integrations for Microsoft Agent Framework." authors = [{ name = "Microsoft", email = "af-support@microsoft.com"}] readme = "README.md" requires-python = ">=3.10" @@ -24,7 +24,7 @@ classifiers = [ ] dependencies = [ "agent-framework-core>=1.15.0,<2", - "azure-cosmos>=4.3.0,<5", + "azure-cosmos>=4.7.0,<5", "six>=1.17.0,<2", ] @@ -45,10 +45,12 @@ addopts = "-ra -q -r fEX" asyncio_mode = "auto" asyncio_default_fixture_loop_scope = "function" filterwarnings = [ + "ignore::agent_framework._feature_stage.ExperimentalWarning", "ignore:Support for class-based `config` is deprecated:DeprecationWarning:pydantic.*" ] timeout = 120 markers = [ + "flaky: marks tests that depend on external services", "integration: marks tests as integration tests that require external services", ] diff --git a/python/packages/azure-cosmos/tests/azure_cosmos/test_vector_store.py b/python/packages/azure-cosmos/tests/azure_cosmos/test_vector_store.py new file mode 100644 index 0000000000..108b12f378 --- /dev/null +++ b/python/packages/azure-cosmos/tests/azure_cosmos/test_vector_store.py @@ -0,0 +1,1277 @@ +# Copyright (c) Microsoft. All rights reserved. + +from __future__ import annotations + +import copy +import gc +import importlib +import math +from collections.abc import AsyncIterator, Callable, Mapping +from typing import Any, cast +from unittest.mock import AsyncMock, MagicMock, patch +from weakref import ref + +import pytest +from agent_framework import ( + Filter, + FilterGroup, + SecretString, + VectorStoreCollectionDefinition, + VectorStoreField, + load_settings, +) +from agent_framework.exceptions import IntegrationException, IntegrationInvalidResponseException, SettingNotFoundError +from azure.cosmos.exceptions import CosmosHttpResponseError, CosmosResourceExistsError, CosmosResourceNotFoundError + +import agent_framework_azure_cosmos as cosmos_package +import agent_framework_azure_cosmos._vector_store as vector_store_module +from agent_framework_azure_cosmos import AzureCosmosSettings, CosmosCollection, CosmosStore +from agent_framework_azure_cosmos._vector_store import ( + _CosmosConnection, + _normalize_policy_path, + _property_access, + _query_metadata_hook, + _required_indexed_paths, + _validate_existing_policies, + _validate_key, +) + +pytestmark = pytest.mark.filterwarnings("ignore::agent_framework._feature_stage.ExperimentalWarning") + + +def test_package_vector_exports_are_lazy() -> None: + with patch.object(importlib, "import_module", wraps=importlib.import_module) as importer: + package = importlib.reload(cosmos_package) + importer.assert_not_called() + assert "CosmosCollection" not in vars(package) + assert package.CosmosCollection is CosmosCollection + assert {"AzureCosmosSettings", "CosmosCollection", "CosmosStore"} <= set(dir(package)) + assert {"AzureCosmosSettings", "CosmosCollection", "CosmosStore"} <= set(package.__all__) + namespace: dict[str, Any] = {} + exec("from agent_framework_azure_cosmos import *", namespace) + assert namespace["CosmosCollection"] is CosmosCollection + assert namespace["CosmosStore"] is CosmosStore + assert namespace["CosmosHistoryProvider"] is package.CosmosHistoryProvider + export_name = "CosmosStore" + with ( + patch.object(importlib, "import_module", side_effect=ImportError("missing vector core")), + pytest.raises(ImportError, match="core with vector-store support"), + ): + getattr(package, export_name) + missing_name = "not_an_export" + with pytest.raises(AttributeError): + getattr(package, missing_name) + + +def _async_items(items: list[Any]) -> AsyncIterator[Any]: + async def iterator() -> AsyncIterator[Any]: + for item in items: + yield item + + return iterator() + + +def _not_found() -> CosmosResourceNotFoundError: + return CosmosResourceNotFoundError(message="missing") + + +def _http_error() -> CosmosHttpResponseError: + return CosmosHttpResponseError(message="failed") + + +def _definition( + *, + index_kind: str = "default", + distance: str = "cosine_similarity", + dimensions: int = 3, + vector_type: str = "float32", + annotations: Mapping[str, Any] | None = None, + second_vector: bool = False, +) -> VectorStoreCollectionDefinition: + fields = [ + VectorStoreField("key", name="key", storage_name="id", type_="str"), + VectorStoreField("data", name="text", storage_name="content", type_="str", is_indexed=True), + VectorStoreField("data", name="count", type_="int", is_indexed=True), + VectorStoreField("data", name="active", type_="bool", is_indexed=True), + VectorStoreField("data", name="tags", type_="list", is_indexed=True), + VectorStoreField("data", name="optional", type_="str", is_indexed=True), + VectorStoreField( + "vector", + name="vector", + storage_name="embedding", + type_=vector_type, + dimensions=dimensions, + index_kind=index_kind, + distance_function=distance, + provider_annotations=annotations, + ), + ] + if second_vector: + fields.append( + VectorStoreField( + "vector", + name="other_vector", + storage_name="otherEmbedding", + type_="uint8", + dimensions=dimensions, + index_kind="disk_ann", + distance_function="euclidean_distance", + ) + ) + return VectorStoreCollectionDefinition(fields, collection_name="items") + + +def _record(key: str = "one", *, include_vectors: bool = True) -> dict[str, Any]: + record: dict[str, Any] = { + "id": key, + "content": "hello", + "count": 2, + "active": True, + "tags": ["a", "b"], + "optional": None, + } + if include_vectors: + record["embedding"] = [1.0, 0.0, 0.0] + return record + + +def _container_properties(collection: CosmosCollection[Any]) -> dict[str, Any]: + return { + "partitionKey": {"paths": ["/id"], "kind": "Hash"}, + "vectorEmbeddingPolicy": copy.deepcopy(collection._vector_policy), + "indexingPolicy": copy.deepcopy(collection._indexing_policy), + } + + +def _collection( + *, + definition: VectorStoreCollectionDefinition | None = None, + query_results: list[Any] | None = None, +) -> tuple[CosmosCollection[dict[str, Any]], MagicMock]: + container = MagicMock() + container.id = "items" + container.read = AsyncMock() + container.upsert_item = AsyncMock(return_value={}) + container.read_item = AsyncMock() + container.delete_item = AsyncMock(return_value=None) + container.query_items = MagicMock(return_value=_async_items(query_results or [])) + collection = CosmosCollection( + dict, + definition=definition or _definition(), + container_client=container, + ) + container.read.return_value = _container_properties(collection) + return collection, container + + +def test_settings_precedence(tmp_path: Any, monkeypatch: pytest.MonkeyPatch) -> None: + env_file = tmp_path / "cosmos.env" + env_file.write_text( + "AZURE_COSMOS_ENDPOINT=https://file.documents.azure.com/\n" + "AZURE_COSMOS_DATABASE_NAME=file-db\n" + "AZURE_COSMOS_CONTAINER_NAME=file-container\n" + "AZURE_COSMOS_KEY=file-key\n" + ) + monkeypatch.setenv("AZURE_COSMOS_ENDPOINT", "https://process.documents.azure.com/") + monkeypatch.setenv("AZURE_COSMOS_DATABASE_NAME", "process-db") + settings = load_settings( + AzureCosmosSettings, + env_prefix="AZURE_COSMOS_", + endpoint="https://explicit.documents.azure.com/", + env_file_path=str(env_file), + ) + assert settings["endpoint"] == "https://explicit.documents.azure.com/" + assert settings["database_name"] == "file-db" + assert settings["container_name"] == "file-container" + key = settings["key"] + assert isinstance(key, SecretString) + assert key.get_secret_value() == "file-key" + + +def test_owned_store_uses_masked_key_and_closes_only_client() -> None: + client = MagicMock() + client.close = AsyncMock() + with patch.object(vector_store_module, "CosmosClient", return_value=client) as factory: + store = CosmosStore( + endpoint="https://account.documents.azure.com/", + database_name="db", + credential=SecretString("secret"), + ) + assert store.database_name == "db" + factory.assert_called_once() + assert factory.call_args.kwargs["credential"] == "secret" + assert "secret" not in repr(SecretString("secret")) + + +async def test_owned_store_closes_client() -> None: + client = MagicMock() + client.close = AsyncMock() + with patch.object(vector_store_module, "CosmosClient", return_value=client): + store = CosmosStore( + endpoint="https://account.documents.azure.com/", + database_name="db", + credential="key", + ) + await store.close() + client.close.assert_awaited_once() + await store.close() + client.close.assert_awaited_once() + + +async def test_injected_clients_bypass_settings_and_remain_open() -> None: + client = MagicMock() + client.close = AsyncMock() + database = MagicMock() + database.id = "db" + database.read = AsyncMock(return_value={}) + with patch.object(vector_store_module, "load_settings", side_effect=AssertionError("must not load settings")): + client_store = CosmosStore(cosmos_client=client, database_name="db") + database_store = CosmosStore(database_client=database) + collection, container = _collection() + await collection.collection_exists() + await client_store.close() + await database_store.close() + await collection.close() + client.close.assert_not_awaited() + container.close.assert_not_called() + + +@pytest.mark.parametrize( + "kwargs,match", + [ + ({"cosmos_client": MagicMock()}, "database_name"), + ({"cosmos_client": MagicMock(), "database_client": MagicMock()}, "at most one"), + ({"database_client": MagicMock(), "database_name": "db"}, "cannot be combined"), + ({"database_client": MagicMock(), "create_database": True}, "create_database"), + ({"cosmos_client": MagicMock(), "database_name": "db", "endpoint": "https://x"}, "cannot be combined"), + ], +) +def test_store_rejects_conflicting_injection(kwargs: dict[str, Any], match: str) -> None: + if "database_client" in kwargs: + kwargs["database_client"].id = "db" + with pytest.raises(ValueError, match=match): + CosmosStore(**kwargs) + + +def test_container_injection_rejects_settings_and_name_mismatch() -> None: + container = MagicMock() + container.id = "actual" + with pytest.raises(ValueError, match="cannot be combined"): + CosmosCollection( + dict, + definition=_definition(), + container_client=container, + endpoint="https://account.documents.azure.com/", + ) + with pytest.raises(ValueError, match="must match"): + CosmosCollection( + dict, + definition=_definition(), + collection_name="other", + container_client=container, + ) + + +def test_missing_settings_and_invalid_credential(monkeypatch: pytest.MonkeyPatch) -> None: + for name in ( + "AZURE_COSMOS_ENDPOINT", + "AZURE_COSMOS_DATABASE_NAME", + "AZURE_COSMOS_CONTAINER_NAME", + "AZURE_COSMOS_KEY", + ): + monkeypatch.delenv(name, raising=False) + with pytest.raises(SettingNotFoundError): + CosmosStore() + with pytest.raises(TypeError, match="credential"): + CosmosStore( + endpoint="https://account.documents.azure.com/", + database_name="db", + credential=cast(Any, object()), + ) + + +def test_default_and_multi_vector_schema() -> None: + collection, _ = _collection(definition=_definition(second_vector=True)) + embeddings = collection._vector_policy["vectorEmbeddings"] + indexes = collection._indexing_policy["vectorIndexes"] + assert embeddings == [ + { + "path": "/embedding", + "dataType": "float32", + "distanceFunction": "cosine", + "dimensions": 3, + }, + { + "path": "/otherEmbedding", + "dataType": "uint8", + "distanceFunction": "euclidean", + "dimensions": 3, + }, + ] + assert indexes == [ + {"path": "/embedding", "type": "quantizedFlat"}, + {"path": "/otherEmbedding", "type": "diskANN"}, + ] + assert {"path": "/embedding/*"} in collection._indexing_policy["excludedPaths"] + assert {"path": "/otherEmbedding/*"} in collection._indexing_policy["excludedPaths"] + + +def test_vector_schema_tuning_annotations() -> None: + definition = _definition( + index_kind="disk_ann", + vector_type="float", + dimensions=512, + annotations={ + "azure_cosmos": { + "data_type": "int8", + "quantizer_type": "spherical", + "quantization_byte_size": 256, + "indexing_search_list_size": 200, + } + }, + ) + collection, _ = _collection(definition=definition) + assert collection._vector_policy["vectorEmbeddings"][0]["dataType"] == "int8" + assert collection._indexing_policy["vectorIndexes"] == [ + { + "path": "/embedding", + "type": "diskANN", + "quantizerType": "spherical", + "quantizationByteSize": 256, + "indexingSearchListSize": 200, + } + ] + + +def test_declared_data_storage_names_are_escaped_and_unindexed_paths_are_excluded() -> None: + storage_name = 'content"] WHERE true --' + definition = VectorStoreCollectionDefinition( + [ + VectorStoreField("key", name="key", storage_name="id", type_="str"), + VectorStoreField("data", name="text", storage_name=storage_name, type_="str", is_indexed=False), + VectorStoreField("vector", name="vector", storage_name="embedding", dimensions=3, type_="float"), + ], + collection_name="items", + ) + collection, _ = _collection(definition=definition) + assert _property_access(storage_name) == 'c["content\\"] WHERE true --"]' + assert {"path": '/"content\\"] WHERE true --"/*'} in collection._indexing_policy["excludedPaths"] + with pytest.raises(ValueError, match="excluded from indexing"): + collection._prepare_filter(Filter("text", "eq", "value")) + + +@pytest.mark.parametrize( + "definition_factory,match", + [ + ( + lambda: VectorStoreCollectionDefinition( + [ + VectorStoreField("key", name="key", type_="str"), + VectorStoreField("vector", name="vector", dimensions=3, type_="float"), + ], + collection_name="items", + ), + "storage name 'id'", + ), + ( + lambda: VectorStoreCollectionDefinition( + [ + VectorStoreField("key", name="key", storage_name="id", type_="int"), + VectorStoreField("vector", name="vector", dimensions=3, type_="float"), + ], + collection_name="items", + ), + "Key field type", + ), + ( + lambda: VectorStoreCollectionDefinition( + [VectorStoreField("key", name="key", storage_name="id", type_="str")], + collection_name="items", + ), + "at least one vector", + ), + ( + lambda: VectorStoreCollectionDefinition( + [ + VectorStoreField("key", name="key", storage_name="id", type_="str"), + VectorStoreField("vector", name="nested", storage_name="a.b", dimensions=3, type_="float"), + ], + collection_name="items", + ), + "top-level ASCII", + ), + ], +) +def test_invalid_collection_schema( + definition_factory: Callable[[], VectorStoreCollectionDefinition], + match: str, +) -> None: + container = MagicMock() + container.id = "items" + with pytest.raises(ValueError, match=match): + CosmosCollection(dict, definition=definition_factory(), container_client=container) + + +@pytest.mark.parametrize( + "field_kwargs,exception,match", + [ + ({"index_kind": "flat", "dimensions": 506}, ValueError, "at most 505"), + ({"index_kind": "disk_ann", "dimensions": 4097}, ValueError, "at most 4096"), + ({"index_kind": "hnsw"}, NotImplementedError, "index kind"), + ({"distance": "cosine_distance"}, NotImplementedError, "distance"), + ({"vector_type": "float16"}, ValueError, "Vector field"), + ({"vector_type": "float64"}, ValueError, "Vector field"), + ( + {"vector_type": "float", "annotations": {"azure_cosmos": {"data_type": "float16"}}}, + NotImplementedError, + "float32, int8, or uint8", + ), + ( + {"annotations": {"azure_cosmos": {"unknown": 1}}}, + ValueError, + "annotation", + ), + ( + {"annotations": {"azure_cosmos": {"data_type": 1}}}, + TypeError, + "data_type", + ), + ( + {"index_kind": "flat", "annotations": {"azure_cosmos": {"quantizer_type": "product"}}}, + ValueError, + "quantizer_type", + ), + ( + {"index_kind": "quantized_flat", "annotations": {"azure_cosmos": {"indexing_search_list_size": 100}}}, + ValueError, + "diskANN", + ), + ( + { + "index_kind": "disk_ann", + "dimensions": 1536, + "annotations": {"azure_cosmos": {"quantization_byte_size": 513}}, + }, + ValueError, + "between 4 and 512", + ), + ( + { + "index_kind": "disk_ann", + "dimensions": 1536, + "annotations": {"azure_cosmos": {"quantization_byte_size": 3}}, + }, + ValueError, + "between 4 and 512", + ), + ( + { + "index_kind": "disk_ann", + "dimensions": 8, + "annotations": {"azure_cosmos": {"quantization_byte_size": 9}}, + }, + ValueError, + "between 4 and 8", + ), + ( + { + "index_kind": "disk_ann", + "annotations": {"azure_cosmos": {"indexing_search_list_size": 24}}, + }, + ValueError, + "between 25 and 500", + ), + ], +) +def test_invalid_vector_schema(field_kwargs: dict[str, Any], exception: type[Exception], match: str) -> None: + with pytest.raises(exception, match=match): + _collection(definition=_definition(**field_kwargs)) + + +@pytest.mark.parametrize( + "vector_type,annotations,match", + [ + ("float16", None, "Vector field"), + ("float", {"azure_cosmos": {"data_type": "float16"}}, "float32, int8, or uint8"), + ("float16", {"azure_cosmos": {"data_type": "float32"}}, "Vector field"), + ], +) +def test_float16_schema_rejected_before_container_io( + vector_type: str, + annotations: Mapping[str, Any] | None, + match: str, +) -> None: + container = MagicMock() + container.id = "items" + container.read = AsyncMock() + with pytest.raises((ValueError, NotImplementedError), match=match): + CosmosCollection( + dict, + definition=_definition(vector_type=vector_type, annotations=annotations), + container_client=container, + ) + container.read.assert_not_awaited() + container.query_items.assert_not_called() + + +def test_vector_schema_tuning_boundaries() -> None: + collection, _ = _collection( + definition=_definition( + index_kind="disk_ann", + dimensions=4, + annotations={ + "azure_cosmos": { + "quantization_byte_size": 4, + "indexing_search_list_size": 25, + } + }, + ) + ) + assert collection._indexing_policy["vectorIndexes"] == [ + { + "path": "/embedding", + "type": "diskANN", + "quantizationByteSize": 4, + "indexingSearchListSize": 25, + } + ] + + +def test_policy_path_normalization_and_semantic_defaults() -> None: + collection, _ = _collection( + definition=_definition( + index_kind="disk_ann", + annotations={"azure_cosmos": {"indexing_search_list_size": 100}}, + second_vector=True, + ) + ) + properties = _container_properties(collection) + properties["vectorEmbeddingPolicy"]["vectorEmbeddings"].reverse() + properties["indexingPolicy"]["excludedPaths"] = [ + {"path": '/"_etag"/?'}, + {"path": '/"embedding"/*'}, + {"path": '/"otherEmbedding"/*'}, + ] + properties["indexingPolicy"]["vectorIndexes"][0]["quantizerType"] = "product" + properties["indexingPolicy"]["vectorIndexes"].reverse() + assert _normalize_policy_path('/"embedding"/*') == "/embedding/*" + assert _normalize_policy_path('/"a/b"/*') == '/"a/b"/*' + assert _normalize_policy_path('/"a/b"/*') != "/a/b/*" + _validate_existing_policies( + properties, + vector_policy=collection._vector_policy, + indexing_policy=collection._indexing_policy, + required_indexed_paths=_required_indexed_paths(collection.definition), + ) + + +def test_quoted_policy_path_does_not_match_nested_path() -> None: + definition = VectorStoreCollectionDefinition( + [ + VectorStoreField("key", name="key", storage_name="id", type_="str"), + VectorStoreField("data", name="value", storage_name="a/b", type_="str", is_indexed=False), + VectorStoreField("vector", name="vector", storage_name="embedding", dimensions=3, type_="float"), + ], + collection_name="items", + ) + collection, _ = _collection(definition=definition) + properties = _container_properties(collection) + excluded_paths = properties["indexingPolicy"]["excludedPaths"] + excluded_paths[1] = {"path": "/a/b/*"} + with pytest.raises(ValueError, match="exclude all required"): + _validate_existing_policies( + properties, + vector_policy=collection._vector_policy, + indexing_policy=collection._indexing_policy, + required_indexed_paths=_required_indexed_paths(collection.definition), + ) + + +@pytest.mark.parametrize( + "mutate,match", + [ + (lambda p: p.update({"partitionKey": {"paths": ["/tenant"], "kind": "Hash"}}), "partition"), + (lambda p: p.pop("vectorEmbeddingPolicy"), "embedding policy"), + ( + lambda p: p["vectorEmbeddingPolicy"]["vectorEmbeddings"][0].update({"dimensions": 99}), + "embedding policy", + ), + (lambda p: p["indexingPolicy"].update({"automatic": False}), "automatic"), + (lambda p: p["indexingPolicy"].update({"includedPaths": []}), "root path"), + (lambda p: p["indexingPolicy"].update({"excludedPaths": [{"path": "/_etag/?"}]}), "exclude"), + ( + lambda p: p["indexingPolicy"]["excludedPaths"].append({"path": "/content/*"}), + "must remain indexed", + ), + (lambda p: p["indexingPolicy"]["vectorIndexes"][0].update({"type": "flat"}), "index type"), + ( + lambda p: p["indexingPolicy"]["vectorIndexes"][0].update({"quantizerType": "spherical"}), + "quantizer", + ), + ], +) +def test_incompatible_existing_policy(mutate: Any, match: str) -> None: + collection, _ = _collection() + properties = _container_properties(collection) + mutate(properties) + with pytest.raises(ValueError, match=match): + _validate_existing_policies( + properties, + vector_policy=collection._vector_policy, + indexing_policy=collection._indexing_policy, + required_indexed_paths=_required_indexed_paths(collection.definition), + ) + + +@pytest.mark.parametrize( + "key,exception", + [ + ("", ValueError), + ("a/b", ValueError), + ("a\\b", ValueError), + ("a?b", ValueError), + ("a#b", ValueError), + ("a" * 1024, ValueError), + (1, TypeError), + ], +) +def test_key_validation(key: Any, exception: type[Exception]) -> None: + with pytest.raises(exception): + _validate_key(key) + assert _validate_key("cafe-\N{SNOWMAN}") == "cafe-\N{SNOWMAN}" + + +@pytest.mark.parametrize( + "vector,vector_type,exception", + [ + ([1.0, 2.0], "float32", ValueError), + ([1.0, math.inf, 0.0], "float32", ValueError), + ([1.0, True, 0.0], "float32", TypeError), + ([1, 2.0, 3], "int8", TypeError), + ([1, 2, 128], "int8", ValueError), + (b"123", "float32", TypeError), + ], +) +async def test_full_batch_vector_preflight(vector: Any, vector_type: str, exception: type[Exception]) -> None: + collection, container = _collection(definition=_definition(vector_type=vector_type)) + records = [_record("valid"), {**_record("late"), "embedding": vector}] + if vector_type in ("int8", "uint8"): + records[0]["embedding"] = [1, 0, 0] + with pytest.raises(exception): + await collection.upsert(records, generate_vectors=False) + container.read.assert_not_awaited() + container.upsert_item.assert_not_awaited() + + +async def test_full_batch_item_size_and_json_preflight() -> None: + collection, container = _collection() + oversized = {**_record("large"), "content": "x" * (2 * 1024 * 1024)} + with pytest.raises(ValueError, match="2 MiB"): + await collection.upsert([_record(), oversized], generate_vectors=False) + with pytest.raises(ValueError, match="IEEE 754"): + await collection.upsert([{**_record(), "count": 2**53}], generate_vectors=False) + with pytest.raises(TypeError, match="declared type"): + await collection.upsert([{**_record(), "active": 1}], generate_vectors=False) + container.upsert_item.assert_not_awaited() + + +@pytest.mark.parametrize( + "expression,clause,values", + [ + (Filter("optional", "exists"), 'IS_DEFINED(c["optional"])', []), + ( + Filter("optional", "is_null"), + '(IS_DEFINED(c["optional"]) AND IS_NULL(c["optional"]))', + [], + ), + ( + Filter("optional", "is_not_null"), + '(IS_DEFINED(c["optional"]) AND NOT IS_NULL(c["optional"]))', + [], + ), + ( + Filter("text", "eq", "O'Reilly"), + '(IS_DEFINED(c["content"]) AND c["content"] = @filter_0)', + ["O'Reilly"], + ), + ( + Filter("text", "ne", "x"), + '(IS_DEFINED(c["content"]) AND (IS_NULL(c["content"]) OR c["content"] != @filter_0))', + ["x"], + ), + (Filter("count", "eq", True), "false", []), + (Filter("count", "ne", True), 'IS_DEFINED(c["count"])', []), + ( + Filter("count", "gt", 1), + '(IS_DEFINED(c["count"]) AND NOT IS_NULL(c["count"]) AND c["count"] > @filter_0)', + [1], + ), + ( + Filter("count", "between", [1, 3]), + ( + '(IS_DEFINED(c["count"]) AND NOT IS_NULL(c["count"]) ' + 'AND c["count"] >= @filter_0 AND c["count"] <= @filter_1)' + ), + [1, 3], + ), + ( + Filter("text", "in", ["a", "b"]), + '(IS_DEFINED(c["content"]) AND NOT IS_NULL(c["content"]) AND ARRAY_CONTAINS(@filter_0, c["content"]))', + [["a", "b"]], + ), + (Filter("text", "in", []), "false", []), + ( + Filter("text", "not_in", []), + '(IS_DEFINED(c["content"]) AND NOT IS_NULL(c["content"]))', + [], + ), + ( + Filter("tags", "contains", "a"), + '(IS_ARRAY(c["tags"]) AND (ARRAY_CONTAINS(c["tags"], @filter_0)))', + ["a"], + ), + (Filter("tags", "contains_any", []), "false", []), + (Filter("tags", "contains_all", []), 'IS_ARRAY(c["tags"])', []), + ( + Filter("tags", "contains_all", ["a", "b"]), + '(IS_ARRAY(c["tags"]) AND (ARRAY_CONTAINS(c["tags"], @filter_0) AND ARRAY_CONTAINS(c["tags"], @filter_1)))', + ["a", "b"], + ), + ( + Filter("tags", "contains_any", [None]), + '(IS_ARRAY(c["tags"]) AND (ARRAY_CONTAINS(c["tags"], @filter_0)))', + [None], + ), + ( + Filter("text", "in", [None, "a"]), + '(IS_DEFINED(c["content"]) AND NOT IS_NULL(c["content"]) AND ARRAY_CONTAINS(@filter_0, c["content"]))', + [[None, "a"]], + ), + ( + Filter("text", "starts_with", "he"), + '(IS_STRING(c["content"]) AND STARTSWITH(c["content"], @filter_0))', + ["he"], + ), + ( + Filter("text", "ends_with", "lo"), + '(IS_STRING(c["content"]) AND ENDSWITH(c["content"], @filter_0))', + ["lo"], + ), + ( + Filter("text", "contains_text", "ell"), + '(IS_STRING(c["content"]) AND CONTAINS(c["content"], @filter_0))', + ["ell"], + ), + ], +) +def test_filter_translation_matches_portable_semantics( + expression: Filter, + clause: str, + values: list[Any], +) -> None: + collection, _ = _collection() + actual, parameters = collection._prepare_filter(expression) + assert actual == clause + assert actual is not None + assert [parameter["value"] for parameter in parameters] == values + if values == ["O'Reilly"]: + assert "O'Reilly" not in actual + + +def test_filter_groups_and_rejections() -> None: + collection, _ = _collection() + group = FilterGroup( + "and", + [ + Filter("count", "gte", 1), + FilterGroup("not", [Filter("active", "eq", False)]), + ], + ) + clause, parameters = collection._prepare_filter(group) + assert clause == ( + '((IS_DEFINED(c["count"]) AND NOT IS_NULL(c["count"]) AND c["count"] >= @filter_0) ' + 'AND (NOT (IS_DEFINED(c["active"]) AND c["active"] = @filter_1)))' + ) + assert [item["value"] for item in parameters] == [1, False] + for expression, exception in ( + (Filter("tags", "eq", {"a": 1}), NotImplementedError), + (Filter("text", "in", [["nested"]]), NotImplementedError), + (Filter("vector", "eq", 1), NotImplementedError), + (Filter("text.child", "eq", "x"), NotImplementedError), + (Filter("text", "provider.unknown", "x"), NotImplementedError), + (Filter("active", "gt", True), NotImplementedError), + ): + with pytest.raises(exception): + collection._prepare_filter(expression) + + +async def test_get_rejects_unsupported_filter_before_io() -> None: + collection, container = _collection() + with pytest.raises(NotImplementedError, match="nested"): + await collection.get(filter=Filter("text.child", "eq", "x")) + container.read.assert_not_awaited() + container.query_items.assert_not_called() + + +async def test_ensure_existing_and_create_absent() -> None: + definition = _definition() + database = MagicMock() + database.id = "db" + database.read = AsyncMock(return_value={}) + existing = MagicMock() + existing.read = AsyncMock() + database.get_container_client.return_value = existing + database.create_container = AsyncMock() + collection = CosmosCollection( + dict, + definition=definition, + database_client=database, + ) + existing.read.return_value = _container_properties(collection) + await collection.ensure_collection_exists() + database.create_container.assert_not_awaited() + + existing.read.side_effect = [_not_found(), _container_properties(collection)] + created = MagicMock() + created.read = AsyncMock(return_value=_container_properties(collection)) + database.create_container.return_value = created + collection._container_client = None + collection._container_validated = False + await collection.ensure_collection_exists() + await_args = database.create_container.await_args + assert await_args is not None + kwargs = await_args.kwargs + assert kwargs["partition_key"].path == "/id" + assert kwargs["vector_embedding_policy"] == collection._vector_policy + assert kwargs["indexing_policy"] == collection._indexing_policy + + +async def test_ensure_handles_create_race_and_revalidates_winner() -> None: + database = MagicMock() + database.id = "db" + database.read = AsyncMock(return_value={}) + winner = MagicMock() + winner.read = AsyncMock() + database.get_container_client.return_value = winner + database.create_container = AsyncMock(side_effect=CosmosResourceExistsError(message="race")) + collection = CosmosCollection(dict, definition=_definition(), database_client=database) + winner.read.side_effect = [_not_found(), _container_properties(collection)] + await collection.ensure_collection_exists() + assert winner.read.await_count == 2 + + +async def test_database_creation_is_explicit() -> None: + client = MagicMock() + database = MagicMock() + database.id = "db" + database.read = AsyncMock(return_value={}) + client.create_database = AsyncMock(return_value=database) + client.get_database_client.return_value = database + connection = _CosmosConnection( + cosmos_client=client, + database_client=None, + database_name="db", + owns_client=False, + create_database=True, + ) + assert await connection.get_database() is database + client.create_database.assert_awaited_once_with(id="db") + assert await connection.get_database() is database + client.create_database.assert_awaited_once() + + +async def test_database_creation_race_uses_existing_database() -> None: + client = MagicMock() + database = MagicMock() + database.read = AsyncMock(return_value={}) + client.create_database = AsyncMock(side_effect=CosmosResourceExistsError(message="race")) + client.get_database_client.return_value = database + connection = _CosmosConnection( + cosmos_client=client, + database_client=None, + database_name="db", + owns_client=False, + create_database=True, + ) + assert await connection.get_database() is database + database.read.assert_awaited_once() + + +async def test_collection_exists_and_delete_lifecycle() -> None: + database = MagicMock() + database.id = "db" + database.read = AsyncMock(return_value={}) + container = MagicMock() + container.read = AsyncMock(return_value={}) + database.get_container_client.return_value = container + database.delete_container = AsyncMock(side_effect=_not_found()) + collection = CosmosCollection(dict, definition=_definition(), database_client=database) + assert await collection.collection_exists() + container.read.side_effect = _not_found() + assert not await collection.collection_exists() + await collection.ensure_collection_deleted() + database.delete_container.assert_awaited_once_with("items") + + +async def test_upsert_point_get_and_delete() -> None: + collection, container = _collection() + records = [_record("one"), _record("two")] + assert await collection.upsert(records, generate_vectors=False) == ["one", "two"] + assert [call.kwargs["body"]["id"] for call in container.upsert_item.await_args_list] == ["one", "two"] + + container.read_item.side_effect = [_record("one"), _not_found(), _record("two")] + found = await collection.get(["one", "missing", "two"]) + assert [item["key"] for item in found] == ["one", "two"] + assert all("vector" not in item for item in found) + calls = container.read_item.await_args_list + assert [(call.kwargs["item"], call.kwargs["partition_key"]) for call in calls] == [ + ("one", "one"), + ("missing", "missing"), + ("two", "two"), + ] + await collection.delete(["one", "missing"]) + assert [(call.kwargs["item"], call.kwargs["partition_key"]) for call in container.delete_item.await_args_list] == [ + ("one", "one"), + ("missing", "missing"), + ] + + +async def test_partial_write_and_delete_failures_are_aligned() -> None: + collection, container = _collection() + container.upsert_item.side_effect = [{}, _http_error()] + with pytest.raises(IntegrationException, match="1/2.*input index 1"): + await collection.upsert([_record("one"), _record("two")], generate_vectors=False) + container.delete_item.side_effect = [None, _http_error()] + with pytest.raises(IntegrationException, match="1/2.*input index 1"): + await collection.delete(["one", "two"]) + + +async def test_filtered_get_query_is_parameterized_and_bounded() -> None: + collection, container = _collection(query_results=[_record("one", include_vectors=False)]) + found = await collection.get( + filter=Filter("text", "eq", "hello"), + top=4, + skip=2, + order_by={"count": False}, + ) + assert [item["key"] for item in found] == ["one"] + kwargs = container.query_items.call_args.kwargs + query = kwargs["query"] + assert "hello" not in query + assert 'WHERE (IS_DEFINED(c["content"]) AND c["content"] = @filter_0)' in query + assert 'ORDER BY c["count"] DESC OFFSET @skip LIMIT @top' in query + assert kwargs["parameters"] == [ + {"name": "@filter_0", "value": "hello"}, + {"name": "@skip", "value": 2}, + {"name": "@top", "value": 4}, + ] + + +def test_order_by_validation() -> None: + collection, _ = _collection() + with pytest.raises(NotImplementedError, match="one order_by"): + collection._prepare_order_by({"text": True, "count": False}) + with pytest.raises(TypeError, match="boolean"): + collection._prepare_order_by(cast(Mapping[str, bool], {"text": 1})) + with pytest.raises(ValueError, match="Unknown"): + collection._prepare_order_by({"missing": True}) + with pytest.raises(ValueError, match="indexed"): + collection._prepare_order_by({"vector": True}) + + +async def test_cosine_search_query_threshold_filter_order_and_metadata() -> None: + collection, container = _collection( + query_results=[{"record": _record("one", include_vectors=False), "score": 0.75}] + ) + + def query_items(**kwargs: Any) -> AsyncIterator[Any]: + kwargs["response_hook"]( + { + "x-ms-request-charge": "3.5", + "x-ms-activity-id": "activity", + "x-ms-continuation": "", + }, + None, + ) + return _async_items([{"record": _record("one", include_vectors=False), "score": 0.75}]) + + container.query_items.side_effect = query_items + results = await collection.search( + vector=[1.0, 0.0, 0.0], + filter=Filter("count", "gte", 1), + score_threshold=0.7, + top=2, + ) + rows = [row async for row in results] + assert rows == [ + { + "record": { + "key": "one", + "text": "hello", + "count": 2, + "active": True, + "tags": ["a", "b"], + "optional": None, + }, + "score": 0.75, + } + ] + query_kwargs = container.query_items.call_args.kwargs + query = query_kwargs["query"] + assert query.startswith("SELECT TOP @top VALUE") + assert ( + 'WHERE ((IS_DEFINED(c["count"]) AND NOT IS_NULL(c["count"]) AND c["count"] >= @filter_0) AND VectorDistance' + ) in query + assert " >= @threshold_" in query + assert query.index("WHERE") < query.index("ORDER BY") + assert " OFFSET " not in query + assert [item["value"] for item in query_kwargs["parameters"]] == [ + [1.0, 0.0, 0.0], + 1, + 0.7, + 2, + ] + assert results.metadata == { + "score_kind": "cosine_similarity", + "score_direction": "higher_is_better", + "request_charge": 3.5, + "activity_id": "activity", + "has_more_results": False, + } + + +async def test_euclidean_search_uses_bounded_top_and_lazily_skips() -> None: + collection, container = _collection( + definition=_definition(distance="euclidean_distance"), + query_results=[ + {"record": _record("skipped", include_vectors=False), "score": 0.0}, + {"record": _record("kept", include_vectors=False), "score": 0.25}, + ], + ) + results = await collection.search( + vector=[1.0, 0.0, 0.0], + top=2, + skip=1, + ) + rows = [row async for row in results] + assert rows == [ + { + "record": { + "key": "kept", + "text": "hello", + "count": 2, + "active": True, + "tags": ["a", "b"], + "optional": None, + }, + "score": 0.25, + } + ] + kwargs = container.query_items.call_args.kwargs + query = kwargs["query"] + assert query.startswith("SELECT TOP @top VALUE") + assert "@threshold_" not in query + assert "OFFSET" not in query + assert {"name": "@top", "value": 3} in kwargs["parameters"] + assert results.metadata is not None + assert results.metadata["score_kind"] == "euclidean_distance" + assert results.metadata["score_direction"] == "lower_is_better" + + +async def test_euclidean_threshold_rejected_before_io() -> None: + collection, container = _collection(definition=_definition(distance="euclidean_distance")) + with pytest.raises(NotImplementedError, match="Omit score_threshold"): + await collection.search( + vector=[1.0, 0.0, 0.0], + score_threshold=0.5, + ) + container.read.assert_not_awaited() + container.query_items.assert_not_called() + + +@pytest.mark.parametrize("distance", ["cosine_similarity", "dot_prod"]) +async def test_similarity_threshold_is_inclusive_and_parameterized(distance: str) -> None: + collection, container = _collection( + definition=_definition(distance=distance), + query_results=[{"record": _record("one", include_vectors=False), "score": 0.5}], + ) + results = await collection.search( + vector=[1.0, 0.0, 0.0], + score_threshold=0.5, + ) + assert [row async for row in results][0]["score"] == 0.5 + kwargs = container.query_items.call_args.kwargs + assert " >= @threshold_" in kwargs["query"] + assert {"name": "@threshold_1", "value": 0.5} in kwargs["parameters"] + + +async def test_search_options_are_validated_and_parameterized() -> None: + collection, container = _collection(definition=_definition(index_kind="disk_ann")) + results = await collection.search( + vector=[1.0, 0.0, 0.0], + operation_options={ + "search_list_size_multiplier": 10, + "quantized_vector_list_multiplier": 5, + "filter_priority": 0.5, + "brute_force": True, + }, + ) + assert [row async for row in results] == [] + kwargs = container.query_items.call_args.kwargs + query = kwargs["query"] + assert "searchListSizeMultiplier" not in query + assert "@brute_force_" in query + options = next(item["value"] for item in kwargs["parameters"] if item["name"].startswith("@vector_options_")) + assert options == { + "searchListSizeMultiplier": 10, + "quantizedVectorListMultiplier": 5, + "filterPriority": 0.5, + } + + +@pytest.mark.parametrize( + "options,match", + [ + ({"unknown": 1}, "Unsupported"), + ({"brute_force": 1}, "boolean"), + ({"search_list_size_multiplier": 0}, "positive"), + ({"filter_priority": 2}, "between 0 and 1"), + ], +) +async def test_invalid_search_options(options: dict[str, Any], match: str) -> None: + collection, container = _collection(definition=_definition(index_kind="disk_ann")) + with pytest.raises((TypeError, ValueError), match=match): + await collection.search(vector=[1.0, 0.0, 0.0], operation_options=options) + container.query_items.assert_not_called() + + +async def test_flat_index_rejects_approximate_search_options() -> None: + collection, _ = _collection(definition=_definition(index_kind="flat")) + with pytest.raises(ValueError, match="flat"): + await collection.search( + vector=[1.0, 0.0, 0.0], + operation_options={"quantized_vector_list_multiplier": 2}, + ) + + +async def test_search_rejects_server_vectorization_and_invalid_score() -> None: + collection, _ = _collection() + with pytest.raises(NotImplementedError, match="server-side embedding"): + await collection.search("text") + with pytest.raises(ValueError, match="finite"): + await collection.search(vector=[1.0, 0.0, 0.0], score_threshold=math.inf) + with pytest.raises(ValueError, match="keyword-hybrid"): + await collection.search( + vector=[1.0, 0.0, 0.0], + additional_property_name="text", + ) + with pytest.raises(IntegrationInvalidResponseException, match="invalid score"): + collection._get_score_from_result({"record": {}, "score": True}) + with pytest.raises(IntegrationInvalidResponseException, match="record projection"): + collection._get_record_from_result({"record": 1}) + + +def test_metadata_hook_is_bounded_and_rejects_bad_charge() -> None: + metadata: dict[str, Any] = { + "request_charge": 0.0, + "activity_id": None, + "has_more_results": False, + } + hook = _query_metadata_hook(metadata) + with pytest.raises(IntegrationInvalidResponseException, match="request-charge"): + hook( + { + "x-ms-request-charge": "bad", + "x-ms-activity-id": "activity", + "x-ms-continuation": "opaque", + }, + None, + ) + assert "opaque" not in repr(metadata) + + +async def test_store_children_share_database_and_do_not_reload_settings() -> None: + database = MagicMock() + database.id = "db" + database.read = AsyncMock(return_value={}) + database.list_containers = MagicMock(return_value=_async_items([{"id": "one"}, {"id": 2}, {"id": "two"}])) + existing = MagicMock() + existing.read = AsyncMock(return_value={}) + database.get_container_client.return_value = existing + database.delete_container = AsyncMock(side_effect=_not_found()) + with patch.object(vector_store_module, "load_settings", side_effect=AssertionError("must not reload settings")): + store = CosmosStore(database_client=database) + collection = store.get_collection(dict, definition=_definition()) + assert await store.list_collection_names() == ["one", "two"] + assert await store.collection_exists("one") + existing.read.side_effect = _not_found() + assert not await store.collection_exists("missing") + await store.ensure_collection_deleted("missing") + assert collection._connection is store._connection + await store.close() + assert collection._closed + database.close.assert_not_called() + + +async def test_store_does_not_retain_abandoned_collection_handles() -> None: + database = MagicMock() + database.id = "db" + database.read = AsyncMock(return_value={}) + store = CosmosStore(database_client=database) + collection = store.get_collection(dict, definition=_definition()) + collection_ref = ref(collection) + assert len(store._collections) == 1 + del collection + gc.collect() + assert collection_ref() is None + assert len(store._collections) == 0 + await store.close() + + +async def test_store_delete_invalidates_same_name_children_only() -> None: + database = MagicMock() + database.id = "db" + database.read = AsyncMock(return_value={}) + lifecycle_container = MagicMock() + lifecycle_container.read = AsyncMock(return_value={}) + database.get_container_client.return_value = lifecycle_container + database.delete_container = AsyncMock(return_value=None) + store = CosmosStore(database_client=database) + first = store.get_collection(dict, definition=_definition()) + sibling = store.get_collection(dict, definition=_definition()) + other = store.get_collection(dict, definition=_definition(), collection_name="other") + for collection in (first, sibling, other): + collection._container_client = MagicMock() + collection._container_validated = True + + await store.ensure_collection_deleted("items") + + assert first._container_client is None + assert not first._container_validated + assert sibling._container_client is None + assert not sibling._container_validated + assert other._container_client is not None + assert other._container_validated + await store.close() + + +async def test_child_delete_invalidates_same_name_sibling() -> None: + database = MagicMock() + database.id = "db" + database.read = AsyncMock(return_value={}) + database.delete_container = AsyncMock(return_value=None) + store = CosmosStore(database_client=database) + first = store.get_collection(dict, definition=_definition()) + sibling = store.get_collection(dict, definition=_definition()) + first._container_client = MagicMock() + first._container_validated = True + sibling._container_client = MagicMock() + sibling._container_validated = True + + await first.ensure_collection_deleted() + + assert first._container_client is None + assert not first._container_validated + assert sibling._container_client is None + assert not sibling._container_validated + await store.close() diff --git a/python/packages/azure-cosmos/tests/azure_cosmos/test_vector_store_integration.py b/python/packages/azure-cosmos/tests/azure_cosmos/test_vector_store_integration.py new file mode 100644 index 0000000000..b51ffe849d --- /dev/null +++ b/python/packages/azure-cosmos/tests/azure_cosmos/test_vector_store_integration.py @@ -0,0 +1,322 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Explicitly authorized cloud test for a unique, disposable Cosmos database.""" + +from __future__ import annotations + +import os +from contextlib import suppress +from typing import cast +from uuid import uuid4 + +import pytest +from agent_framework import Filter, FilterGroup, VectorStoreCollectionDefinition, VectorStoreField +from azure.cosmos.aio import CosmosClient +from azure.cosmos.exceptions import CosmosResourceNotFoundError +from azure.identity.aio import AzureCliCredential + +from agent_framework_azure_cosmos import CosmosCollection, CosmosStore + +_USE_AZURE_CLI = os.getenv("AZURE_COSMOS_VECTOR_TEST_USE_AZURE_CLI") == "1" +_CREATE_DATABASE = os.getenv("AZURE_COSMOS_VECTOR_TEST_CREATE_DATABASE", "1") == "1" + +pytestmark = [ + pytest.mark.flaky, + pytest.mark.integration, + pytest.mark.timeout(240), + pytest.mark.filterwarnings("ignore::agent_framework._feature_stage.ExperimentalWarning"), + pytest.mark.skipif( + os.getenv("AZURE_COSMOS_VECTOR_TESTS") != "1" + or not os.getenv("AZURE_COSMOS_ENDPOINT") + or (not _USE_AZURE_CLI and not os.getenv("AZURE_COSMOS_KEY")), + reason="Requires explicit disposable-database authorization and a vector-enabled Cosmos NoSQL account.", + ), +] + + +async def test_disposable_vector_database_crud_filter_and_search() -> None: + """Use a bounded exact index because no high-volume live-test cost is authorized by default.""" + database_name = os.getenv("AZURE_COSMOS_VECTOR_TEST_DATABASE_NAME", f"af-vector-{uuid4().hex}") + container_name = os.getenv("AZURE_COSMOS_VECTOR_TEST_CONTAINER_NAME", f"items-{uuid4().hex}") + default_container_name = os.getenv( + "AZURE_COSMOS_VECTOR_TEST_DEFAULT_CONTAINER_NAME", + f"items-default-{uuid4().hex}", + ) + definition = VectorStoreCollectionDefinition( + [ + VectorStoreField("key", name="key", storage_name="id", type_="str"), + VectorStoreField("data", name="label", type_="str", is_indexed=True), + VectorStoreField( + "vector", + name="vector", + storage_name="embedding", + type_="float32", + dimensions=3, + index_kind="flat", + distance_function="cosine_similarity", + ), + ], + collection_name=container_name, + ) + credential: str | AzureCliCredential + cli_credential: AzureCliCredential | None = None + if _USE_AZURE_CLI: + cli_credential = AzureCliCredential( + tenant_id=os.getenv("AZURE_COSMOS_VECTOR_TEST_TENANT_ID", ""), + ) + credential = cli_credential + else: + credential = os.environ["AZURE_COSMOS_KEY"] + client = CosmosClient( + os.environ["AZURE_COSMOS_ENDPOINT"], + credential=credential, # pyrefly: ignore[bad-argument-type] + ) + store = CosmosStore( + cosmos_client=client, + database_name=database_name, + create_database=_CREATE_DATABASE, + ) + collection: CosmosCollection[dict] | None = None + default_collection: CosmosCollection[dict] | None = None + try: + collection = store.get_collection(dict, definition=definition) + await collection.ensure_collection_exists() + records = [ + {"id": "one", "label": "first", "embedding": [1.0, 0.0, 0.0]}, + {"id": "two", "label": "second", "embedding": [0.0, 1.0, 0.0]}, + {"id": "three", "label": "third", "embedding": [0.0, 0.0, 1.0]}, + ] + assert await collection.upsert(records, generate_vectors=False) == ["one", "two", "three"] + assert [item["key"] for item in await collection.get(["one", "three"])] == ["one", "three"] + assert [item["key"] for item in await collection.get(filter=Filter("label", "eq", "second"))] == ["two"] + results = await collection.search(vector=[1.0, 0.0, 0.0], top=1) + rows = [row async for row in results] + assert rows[0]["record"]["key"] == "one" + await collection.delete(["one", "two", "three"]) + assert await collection.get(["one", "two", "three"]) == [] + await collection.ensure_collection_deleted() + + default_definition = VectorStoreCollectionDefinition( + [ + VectorStoreField("key", name="key", storage_name="id", type_="str"), + VectorStoreField("data", name="category", type_="str", is_indexed=True), + VectorStoreField("data", name="count", type_="int", is_indexed=True), + VectorStoreField("data", name="active", type_="bool", is_indexed=True), + VectorStoreField("data", name="tags", type_="list", is_indexed=True), + VectorStoreField("data", name="optional", type_="str", is_indexed=True), + VectorStoreField( + "vector", + name="cosine", + type_="float32", + dimensions=3, + distance_function="cosine_similarity", + ), + VectorStoreField( + "vector", + name="dot", + type_="float32", + dimensions=3, + index_kind="flat", + distance_function="dot_prod", + ), + VectorStoreField( + "vector", + name="euclidean", + type_="int8", + dimensions=3, + index_kind="flat", + distance_function="euclidean_distance", + ), + VectorStoreField( + "vector", + name="byte_vector", + type_="uint8", + dimensions=3, + index_kind="flat", + distance_function="cosine_similarity", + ), + ], + collection_name=default_container_name, + ) + default_collection = store.get_collection(dict, definition=default_definition) + await default_collection.ensure_collection_exists() + await default_collection.ensure_collection_exists() + + multi_records = [ + { + "id": "one", + "category": "a", + "count": 1, + "active": True, + "tags": ["x", None], + "optional": None, + "cosine": [1.0, 0.0, 0.0], + "dot": [2.0, 0.0, 0.0], + "euclidean": [1, 0, 0], + "byte_vector": [255, 0, 0], + }, + { + "id": "two", + "category": "b", + "count": 2, + "active": False, + "tags": [], + "optional": "value", + "cosine": [0.0, 1.0, 0.0], + "dot": [1.0, 0.0, 0.0], + "euclidean": [2, 0, 0], + "byte_vector": [0, 255, 0], + }, + { + "id": "three", + "category": "c", + "count": 3, + "active": True, + "tags": ["y"], + "optional": "other", + "cosine": [-1.0, 0.0, 0.0], + "dot": [0.0, 1.0, 0.0], + "euclidean": [3, 0, 0], + "byte_vector": [0, 0, 255], + }, + ] + assert await default_collection.upsert(multi_records, generate_vectors=False) == ["one", "two", "three"] + + sdk_container = client.get_database_client(database_name).get_container_client(default_container_name) + properties = await sdk_container.read() + assert properties["partitionKey"]["paths"] == ["/id"] + assert {item["dataType"] for item in properties["vectorEmbeddingPolicy"]["vectorEmbeddings"]} == { + "float32", + "int8", + "uint8", + } + assert {item["type"] for item in properties["indexingPolicy"]["vectorIndexes"]} == { + "quantizedFlat", + "flat", + } + + cosine_results = [ + item + async for item in await default_collection.search( + vector=[1.0, 0.0, 0.0], + vector_property_name="cosine", + top=3, + operation_options={"quantized_vector_list_multiplier": 5}, + ) + ] + assert [item["record"]["key"] for item in cosine_results] == ["one", "two", "three"] + assert [item["score"] for item in cosine_results] == pytest.approx([1.0, 0.0, -1.0]) + + dot_results = [ + item + async for item in await default_collection.search( + vector=[1.0, 0.0, 0.0], + vector_property_name="dot", + top=3, + ) + ] + assert [item["record"]["key"] for item in dot_results] == ["one", "two", "three"] + assert [item["score"] for item in dot_results] == pytest.approx([2.0, 1.0, 0.0]) + + euclidean_results = [ + item + async for item in await default_collection.search( + vector=[1, 0, 0], + vector_property_name="euclidean", + top=3, + ) + ] + assert [item["record"]["key"] for item in euclidean_results] == ["one", "two", "three"] + assert [item["score"] for item in euclidean_results] == pytest.approx([0.0, 1.0, 2.0]) + + cosine_threshold = [ + item + async for item in await default_collection.search( + vector=[1.0, 0.0, 0.0], + vector_property_name="cosine", + filter=Filter("active", "eq", True), + score_threshold=0.5, + top=3, + ) + ] + assert [item["record"]["key"] for item in cosine_threshold] == ["one"] + dot_threshold = [ + item + async for item in await default_collection.search( + vector=[1.0, 0.0, 0.0], + vector_property_name="dot", + score_threshold=1.5, + top=3, + ) + ] + assert [item["record"]["key"] for item in dot_threshold] == ["one"] + with pytest.raises(NotImplementedError, match="Omit score_threshold"): + await default_collection.search( + vector=[1, 0, 0], + vector_property_name="euclidean", + score_threshold=1.0, + top=3, + ) + + assert [item["key"] for item in await default_collection.get(filter=Filter("optional", "is_null"))] == ["one"] + assert [ + item["key"] for item in await default_collection.get(filter=Filter("tags", "contains_any", [None])) + ] == ["one"] + assert {item["key"] for item in await default_collection.get(filter=Filter("tags", "contains_all", []))} == { + "one", + "two", + "three", + } + assert await default_collection.get(filter=Filter("count", "eq", True)) == [] + assert await default_collection.get(filter=Filter("active", "eq", 1)) == [] + + missing_optional = {**multi_records[0], "id": "missing"} + missing_optional.pop("optional") + await sdk_container.upsert_item(missing_optional) + exists_clause, exists_parameters = default_collection._prepare_filter(Filter("optional", "exists")) + missing_clause, missing_parameters = default_collection._prepare_filter( + FilterGroup("not", [Filter("optional", "exists")]) + ) + assert exists_clause is not None + assert missing_clause is not None + defined_ids = { + cast(str, item) + async for item in sdk_container.query_items( + query=f"SELECT VALUE c.id FROM c WHERE {exists_clause}", # noqa: S608 + parameters=exists_parameters, + ) + } + missing_ids = [ + cast(str, item) + async for item in sdk_container.query_items( + query=f"SELECT VALUE c.id FROM c WHERE {missing_clause}", # noqa: S608 + parameters=missing_parameters, + ) + ] + assert defined_ids == {"one", "two", "three"} + assert missing_ids == ["missing"] + + late_valid = {**multi_records[0], "id": "late-valid"} + late_invalid = {**multi_records[1], "id": "late-invalid", "euclidean": [1, 2.5, 3]} + with pytest.raises(TypeError, match="integer elements"): + await default_collection.upsert([late_valid, late_invalid], generate_vectors=False) + assert await default_collection.get(["late-valid"]) == [] + + assert cosine_results[0]["record"].keys().isdisjoint({"cosine", "dot", "euclidean", "byte_vector"}) + assert cosine_results[0]["score"] == pytest.approx(1.0) + await default_collection.delete(["one", "two", "three", "missing"]) + assert await default_collection.get(["one", "two", "three", "missing"]) == [] + finally: + try: + if _CREATE_DATABASE: + with suppress(CosmosResourceNotFoundError): + await client.delete_database(database_name) + else: + for candidate in (default_collection, collection): + if candidate is not None: + await candidate.ensure_collection_deleted() + finally: + await store.close() + await client.close() + if cli_credential is not None: + await cli_credential.close() diff --git a/python/packages/core/agent_framework/azure/__init__.py b/python/packages/core/agent_framework/azure/__init__.py index 50e7209493..d6e198492b 100644 --- a/python/packages/core/agent_framework/azure/__init__.py +++ b/python/packages/core/agent_framework/azure/__init__.py @@ -16,7 +16,10 @@ "AzureAISearchContextProvider": ("agent_framework_azure_ai_search", "agent-framework-azure-ai-search"), "AzureAISearchSettings": ("agent_framework_azure_ai_search", "agent-framework-azure-ai-search"), "AzureAISearchStore": ("agent_framework_azure_ai_search", "agent-framework-azure-ai-search"), + "AzureCosmosSettings": ("agent_framework_azure_cosmos", "agent-framework-azure-cosmos"), + "CosmosCollection": ("agent_framework_azure_cosmos", "agent-framework-azure-cosmos"), "CosmosHistoryProvider": ("agent_framework_azure_cosmos", "agent-framework-azure-cosmos"), + "CosmosStore": ("agent_framework_azure_cosmos", "agent-framework-azure-cosmos"), "DurableAIAgent": ("agent_framework_durabletask", "agent-framework-durabletask"), "DurableAIAgentClient": ("agent_framework_durabletask", "agent-framework-durabletask"), "DurableAIAgentOrchestrationContext": ("agent_framework_durabletask", "agent-framework-durabletask"), diff --git a/python/packages/core/agent_framework/azure/__init__.pyi b/python/packages/core/agent_framework/azure/__init__.pyi index 577a4ee86b..fa03370db5 100644 --- a/python/packages/core/agent_framework/azure/__init__.pyi +++ b/python/packages/core/agent_framework/azure/__init__.pyi @@ -10,6 +10,7 @@ from agent_framework_azure_ai_search import ( AzureAISearchStore, ) from agent_framework_azure_cosmos import CosmosHistoryProvider +from agent_framework_azure_cosmos._vector_store import AzureCosmosSettings, CosmosCollection, CosmosStore from agent_framework_azurefunctions import AgentFunctionApp, WorkflowHitlContext from agent_framework_durabletask import ( AgentCallbackContext, @@ -29,7 +30,10 @@ __all__ = [ "AzureAISearchContextProvider", "AzureAISearchSettings", "AzureAISearchStore", + "AzureCosmosSettings", + "CosmosCollection", "CosmosHistoryProvider", + "CosmosStore", "DurableAIAgent", "DurableAIAgentClient", "DurableAIAgentOrchestrationContext", diff --git a/python/packages/core/tests/core/test_azure_namespace.py b/python/packages/core/tests/core/test_azure_namespace.py index d4e4b8317b..50e0a1ba60 100644 --- a/python/packages/core/tests/core/test_azure_namespace.py +++ b/python/packages/core/tests/core/test_azure_namespace.py @@ -4,9 +4,17 @@ import agent_framework.azure as azure -CosmosHistoryProvider = pytest.importorskip("agent_framework_azure_cosmos").CosmosHistoryProvider +azure_cosmos = pytest.importorskip("agent_framework_azure_cosmos") def test_azure_namespace_exposes_cosmos_history_provider() -> None: - assert azure.CosmosHistoryProvider is CosmosHistoryProvider - assert "CosmosHistoryProvider" in dir(azure) + assert azure.CosmosHistoryProvider is azure_cosmos.CosmosHistoryProvider + assert azure.CosmosCollection is azure_cosmos.CosmosCollection + assert azure.CosmosStore is azure_cosmos.CosmosStore + assert azure.AzureCosmosSettings is azure_cosmos.AzureCosmosSettings + assert { + "AzureCosmosSettings", + "CosmosCollection", + "CosmosHistoryProvider", + "CosmosStore", + } <= set(dir(azure)) diff --git a/python/uv.lock b/python/uv.lock index 53965727d2..8c03328fb4 100644 --- a/python/uv.lock +++ b/python/uv.lock @@ -305,7 +305,7 @@ dependencies = [ [package.metadata] requires-dist = [ { name = "agent-framework-core", editable = "packages/core" }, - { name = "azure-cosmos", specifier = ">=4.3.0,<5" }, + { name = "azure-cosmos", specifier = ">=4.7.0,<5" }, { name = "six", specifier = ">=1.17.0,<2" }, ]