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 index 15e8a044fd4..04f9b282fe6 100644 --- a/python/packages/azure-cosmos/agent_framework_azure_cosmos/_vector_store.py +++ b/python/packages/azure-cosmos/agent_framework_azure_cosmos/_vector_store.py @@ -25,7 +25,12 @@ ) 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._vectors import ( + _VECTOR_MODEL_REGISTRY, # pyright: ignore[reportPrivateUsage] + EmbeddingClient, + SearchType, + Vector, +) from agent_framework.exceptions import IntegrationException, IntegrationInvalidResponseException from azure.core.credentials import TokenCredential from azure.core.credentials_async import AsyncTokenCredential @@ -102,6 +107,23 @@ class AzureCosmosSettings(TypedDict, total=False): key: SecretString | None +class _CosmosResourceSettings(TypedDict, total=False): + name: str | None + + +def _resolve_resource_name(value: str | None, kind: str) -> str: + if value is None: + settings = load_settings( + _CosmosResourceSettings, + env_prefix=f"AZURE_COSMOS_{kind.upper()}_", + required_fields=["name"], + ) + value = settings.get("name") + if not isinstance(value, str): + raise TypeError(f"{kind} name must be a string.") + return _validate_resource_name(value, kind) + + def _validate_operation_options(options: Mapping[str, Any] | None, allowed: set[str]) -> dict[str, Any]: result = dict(options or {}) if unknown := result.keys() - allowed: @@ -245,15 +267,11 @@ def _connection_from_clients( 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.") + if database_name is not None and database_name != resolved_name: + raise ValueError("database_name must match the injected database_client id.") return _CosmosConnection( cosmos_client=None, database_client=database_client, @@ -262,14 +280,10 @@ def _connection_from_clients( 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"), + database_name=_resolve_resource_name(database_name, "database"), owns_client=False, create_database=create_database, ) @@ -750,16 +764,6 @@ def __init__( 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: @@ -801,6 +805,14 @@ def __init__( ) owns_connection = True + if container_client is not None or cosmos_client is not None or database_client is not None: + registration = _VECTOR_MODEL_REGISTRY.get(record_type) + registered_definition = definition or (registration.definition if registration is not None else None) + configured_name = ( + resolved_name if resolved_name is not None else getattr(registered_definition, "collection_name", None) + ) + resolved_name = _resolve_resource_name(configured_name, "container") + super().__init__( record_type, definition=definition, 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 index 8819799ba9d..a001b55f578 100644 --- a/python/packages/azure-cosmos/tests/azure_cosmos/test_vector_store.py +++ b/python/packages/azure-cosmos/tests/azure_cosmos/test_vector_store.py @@ -7,6 +7,7 @@ import importlib import math from collections.abc import AsyncIterator, Callable, Mapping +from dataclasses import dataclass from typing import Any, cast from unittest.mock import AsyncMock, MagicMock, patch from weakref import ref @@ -19,6 +20,7 @@ VectorStoreCollectionDefinition, VectorStoreField, load_settings, + register_vectorstoremodel, ) from agent_framework.exceptions import IntegrationException, IntegrationInvalidResponseException, SettingNotFoundError from azure.cosmos.exceptions import CosmosHttpResponseError, CosmosResourceExistsError, CosmosResourceNotFoundError @@ -218,6 +220,41 @@ async def test_owned_store_closes_client() -> None: client.close.assert_awaited_once() +@pytest.mark.parametrize("target", ["store", "collection"]) +async def test_owned_clients_resolve_settings_and_close( + target: str, 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://env.documents.azure.com/") + monkeypatch.setenv("AZURE_COSMOS_KEY", "env-key") + client = MagicMock() + client.close = AsyncMock() + resource: CosmosStore | CosmosCollection[Any] + with patch.object(vector_store_module, "CosmosClient", return_value=client) as factory: + if target == "store": + resource = CosmosStore(database_name="explicit-db", env_file_path=str(env_file)) + assert resource.database_name == "explicit-db" + else: + definition = VectorStoreCollectionDefinition(_definition().fields) + resource = CosmosCollection( + dict, definition=definition, database_name="explicit-db", env_file_path=str(env_file) + ) + assert resource.collection_name == "file-container" + assert resource._connection is not None + assert resource._connection.database_name == "explicit-db" + assert factory.call_args.kwargs["url"] == "https://file.documents.azure.com/" + assert factory.call_args.kwargs["credential"] == "file-key" + await resource.close() + await resource.close() + client.close.assert_awaited_once() + + async def test_injected_clients_bypass_settings_and_remain_open() -> None: client = MagicMock() client.close = AsyncMock() @@ -237,39 +274,207 @@ async def test_injected_clients_bypass_settings_and_remain_open() -> None: @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"), - ], + "target", ["store-account", "store-database", "collection-account", "collection-database", "collection-container"] ) -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) +@pytest.mark.parametrize("file_option", ["missing", "invalid-encoding"]) +async def test_injected_clients_ignore_connection_settings(target: str, file_option: str, tmp_path: Any) -> None: + client = MagicMock() + client.id = "injected" + client.close = AsyncMock() + env_file = tmp_path / "ignored.env" + if file_option == "invalid-encoding": + env_file.write_bytes(b"\xff") + kwargs: dict[str, Any] = { + "endpoint": "ignored", + "credential": object(), + "create_database": True, + "env_file_path": str(env_file), + "env_file_encoding": "invalid-encoding", + } + kind = target.split("-")[1] + if kind == "account": + kwargs["database_name"] = "explicit-db" + kwargs[{"account": "cosmos_client", "database": "database_client", "container": "container_client"}[kind]] = client + with patch.object(vector_store_module, "CosmosClient") as factory: + if target.startswith("store"): + store = CosmosStore(**kwargs) + assert store.database_name == ("explicit-db" if kind == "account" else "injected") + await store.close() + else: + collection = CosmosCollection( + dict, + definition=_definition(), + collection_name=None if kind == "container" else "explicit-container", + **kwargs, + ) + assert collection.collection_name == ("injected" if kind == "container" else "explicit-container") + await collection.close() + factory.assert_not_called() + client.close.assert_not_called() -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, - ) +@pytest.mark.parametrize("target", ["store-database", "collection-database", "collection-container"]) +def test_injected_proxy_rejects_mismatched_name(target: str) -> None: + client = MagicMock() + client.id = "injected" + kwargs: dict[str, Any] = ( + {"database_client": client, "database_name": "other"} + if target.endswith("database") + else {"container_client": client, "collection_name": "other"} + ) + if target.endswith("database"): + message = "database_name must match the injected database_client id." + else: + message = "collection_name must match the injected container_client id." + with pytest.raises(ValueError) as exc: + if target.startswith("store"): + CosmosStore(**kwargs) + else: + CosmosCollection(dict, definition=_definition(), **kwargs) + assert str(exc.value) == message + + +@pytest.mark.parametrize("target", ["store-database", "collection-database", "collection-container"]) +def test_injected_proxy_accepts_matching_name(target: str) -> None: + client = MagicMock() + client.id = "injected" + kwargs: dict[str, Any] = ( + {"database_client": client, "database_name": "injected"} + if target.endswith("database") + else {"container_client": client, "collection_name": "injected"} + ) + if target.startswith("store"): + assert CosmosStore(**kwargs).database_name == "injected" + else: + collection = CosmosCollection(dict, definition=_definition(), **kwargs) + if target.endswith("container"): + assert collection.collection_name == "injected" + else: + assert collection._connection is not None + assert collection._connection.database_name == "injected" + + +@pytest.mark.parametrize("target", ["store", "collection"]) +@pytest.mark.parametrize("explicit", [False, True]) +def test_injected_account_resolves_database_independently( + target: str, explicit: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("AZURE_COSMOS_DATABASE_NAME", "env-db") + kwargs: dict[str, Any] = {"cosmos_client": MagicMock()} + if explicit: + kwargs["database_name"] = "explicit-db" + if target == "store": + assert CosmosStore(**kwargs).database_name == ("explicit-db" if explicit else "env-db") + else: + collection = CosmosCollection(dict, definition=_definition(), **kwargs) + assert collection._connection is not None + assert collection._connection.database_name == ("explicit-db" if explicit else "env-db") + + +@pytest.mark.parametrize("kind", ["account", "database"]) +@pytest.mark.parametrize("explicit", [False, True]) +def test_injected_collection_resolves_container_independently( + kind: str, explicit: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("AZURE_COSMOS_CONTAINER_NAME", "env-container") + client = MagicMock() + client.id = "db" + kwargs: dict[str, Any] = {"database_name": "db"} if kind == "account" else {} + kwargs["cosmos_client" if kind == "account" else "database_client"] = client + definition = VectorStoreCollectionDefinition(_definition().fields) + collection = CosmosCollection( + dict, definition=definition, collection_name="explicit-container" if explicit else None, **kwargs + ) + assert collection.collection_name == ("explicit-container" if explicit else "env-container") + + +@pytest.mark.parametrize("kind", ["account", "database"]) +@pytest.mark.parametrize("env_name", [None, "env-container"]) +def test_injected_clients_use_registered_collection_name( + kind: str, env_name: str | None, monkeypatch: pytest.MonkeyPatch +) -> None: + @dataclass + class RegisteredRecord: + key: str + vector: list[float] | None = None + + definition = VectorStoreCollectionDefinition( + [_definition().key_field, _definition().vector_fields[0]], collection_name="registered-items" + ) + register_vectorstoremodel(RegisteredRecord, definition=definition) + assert not hasattr(RegisteredRecord, "__vectorstoremodel_definition__") + if env_name is None: + monkeypatch.delenv("AZURE_COSMOS_CONTAINER_NAME", raising=False) + else: + monkeypatch.setenv("AZURE_COSMOS_CONTAINER_NAME", env_name) + client = MagicMock() + client.id = "db" + kwargs: dict[str, Any] = ( + {"cosmos_client": client, "database_name": "db"} if kind == "account" else {"database_client": client} + ) + collection = CosmosCollection(RegisteredRecord, **kwargs) + assert collection.collection_name == "registered-items" + assert collection.definition is definition + + +@pytest.mark.parametrize( + "target", ["store-account", "collection-account", "collection-account-container", "collection-database"] +) +def test_injected_clients_report_missing_selector(target: str, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("AZURE_COSMOS_DATABASE_NAME", raising=False) + monkeypatch.delenv("AZURE_COSMOS_CONTAINER_NAME", raising=False) + client = MagicMock() + client.id = "db" + kwargs: dict[str, Any] = {"database_client" if target.endswith("database") else "cosmos_client": client} + if target.endswith("container"): + kwargs["database_name"] = "db" + selector = "AZURE_COSMOS_DATABASE_NAME" if target.endswith("account") else "AZURE_COSMOS_CONTAINER_NAME" + with pytest.raises((ValueError, SettingNotFoundError), match=selector): + if target.startswith("store"): + CosmosStore(**kwargs) + else: + CosmosCollection(dict, definition=VectorStoreCollectionDefinition(_definition().fields), **kwargs) + + +@pytest.mark.parametrize( + "target", ["store-account", "store-database", "collection-account", "collection-database", "collection-container"] +) +def test_injected_clients_ignore_invalid_ambient_settings(target: str, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("AZURE_COSMOS_ENDPOINT", "") + monkeypatch.setenv("AZURE_COSMOS_KEY", "") + monkeypatch.setenv("AZURE_COSMOS_DATABASE_NAME", "x" * 256) + monkeypatch.setenv("AZURE_COSMOS_CONTAINER_NAME", "x" * 256) + client = MagicMock() + client.id = "injected" + kind = target.split("-")[1] + kwargs: dict[str, Any] = {"database_name": "db"} if kind == "account" else {} + kwargs[{"account": "cosmos_client", "database": "database_client", "container": "container_client"}[kind]] = client + with patch.object(vector_store_module, "CosmosClient") as factory: + if target.startswith("store"): + store = CosmosStore(**kwargs) + assert store.database_name == ("db" if kind == "account" else "injected") + else: + collection = CosmosCollection(dict, definition=_definition(), **kwargs) + assert collection.collection_name == ("injected" if kind == "container" else "items") + factory.assert_not_called() + + +@pytest.mark.parametrize( + "clients", + [ + ("cosmos_client", "database_client"), + ("cosmos_client", "container_client"), + ("database_client", "container_client"), + ("cosmos_client", "database_client", "container_client"), + ], +) +def test_rejects_multiple_injected_clients(clients: tuple[str, ...]) -> None: + kwargs = {name: MagicMock() for name in clients} + with pytest.raises(ValueError, match="at most one"): + CosmosCollection(dict, definition=_definition(), **kwargs) + if "container_client" not in clients: + with pytest.raises(ValueError, match="at most one"): + CosmosStore(**kwargs) def test_missing_settings_and_invalid_credential(monkeypatch: pytest.MonkeyPatch) -> None: