Python: Prefer injected Cosmos vector store clients over connection settings

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,

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: