diff --git a/python/packages/azure-documentdb/README.md b/python/packages/azure-documentdb/README.md index 671891ba821..f83dd85e720 100644 --- a/python/packages/azure-documentdb/README.md +++ b/python/packages/azure-documentdb/README.md @@ -33,7 +33,13 @@ file, then process environment. Connection strings are held in Agent Framework Connector-created clients enforce TLS, disable retryable writes, and set an application name. They are closed by `aclose()` or an async context manager. An injected PyMongo `AsyncMongoClient`, `AsyncDatabase`, or `AsyncCollection` -remains caller-owned and bypasses settings loading. +remains caller-owned and takes precedence over connection strings, client +options, and selected `.env` files. An account-level injected client resolves a +missing `database_name` from `AZURE_DOCUMENTDB_DATABASE_NAME`. For a collection +created from an injected client or database, `AZURE_DOCUMENTDB_COLLECTION_NAME` +is used only when neither `collection_name` nor the model definition supplies a +name. Connection settings and selected `.env` files are never read for these +injected-object paths. The identity running `ensure_collection_exists()` needs permission to create collections and indexes. Each filtered data field must declare diff --git a/python/packages/azure-documentdb/agent_framework_azure_documentdb/_vector_store.py b/python/packages/azure-documentdb/agent_framework_azure_documentdb/_vector_store.py index c142bc85fa5..b3db9b7f65c 100644 --- a/python/packages/azure-documentdb/agent_framework_azure_documentdb/_vector_store.py +++ b/python/packages/azure-documentdb/agent_framework_azure_documentdb/_vector_store.py @@ -23,7 +23,12 @@ ) from agent_framework._telemetry import FeatureIndex, 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 bson import BSON from bson.errors import InvalidDocument @@ -108,6 +113,35 @@ def _validate_collection_name(name: Any) -> str: return name +class _AzureDocumentDBDatabaseSettings(TypedDict, total=False): + database_name: str | None + + +class _AzureDocumentDBCollectionSettings(TypedDict, total=False): + collection_name: str | None + + +def _resolve_resource_name(value: str | None, kind: str) -> str: + if value is None: + if kind == "database": + settings = load_settings( + _AzureDocumentDBDatabaseSettings, + env_prefix="AZURE_DOCUMENTDB_", + required_fields=["database_name"], + ) + value = settings.get("database_name") + else: + settings = load_settings( + _AzureDocumentDBCollectionSettings, + env_prefix="AZURE_DOCUMENTDB_", + required_fields=["collection_name"], + ) + value = settings.get("collection_name") + if kind == "database": + return _validate_database_name(value) + return _validate_collection_name(value) + + def _validate_storage_name(name: str, *, indexed: bool) -> str: encoded = name.encode("utf-8") if not name or "\0" in name or "." in name or name.startswith("$"): @@ -163,29 +197,18 @@ def _create_client( env_file_encoding: str | None, ) -> _Client: if database is not None: - if any( - value is not None - for value in ( - connection_string, - database_name, - client, - client_options, - env_file_path, - env_file_encoding, - ) - ): - raise ValueError("database cannot be combined with a client or connection settings.") + if client is not None: + raise ValueError("Provide at most one of client or database.") if not isinstance(database, AsyncDatabase): raise TypeError("database must be a PyMongo AsyncDatabase.") + resolved_name = _validate_database_name(database.name) + if database_name is not None and database_name != resolved_name: + raise ValueError("database_name must match the injected database.") return _Client(database.client, database, owned=False) if client is not None: - if any(value is not None for value in (connection_string, client_options, env_file_path, env_file_encoding)): - raise ValueError("client cannot be combined with connection settings or client_options.") if not isinstance(client, AsyncMongoClient): raise TypeError("client must be a PyMongo AsyncMongoClient.") - if database_name is None: - raise ValueError("database_name is required with an injected client.") - resolved_database_name = _validate_database_name(database_name) + resolved_database_name = _resolve_resource_name(database_name, "database") return _Client(client, client[resolved_database_name], owned=False) settings = load_settings( @@ -610,38 +633,41 @@ def __init__( Args: record_type: Registered model type, or ``dict`` with an explicit definition. connection_string: MongoDB connection string for a connector-owned client. - database_name: Database selected for a created or injected client. - client: Caller-owned PyMongo async client. Requires ``database_name``. - database: Caller-owned resolved PyMongo async database. - collection: Caller-owned resolved PyMongo async collection. + Ignored when a client, database, or collection is injected. + database_name: Database selected for a created or injected client. An injected account + client falls back to ``AZURE_DOCUMENTDB_DATABASE_NAME`` when this is omitted. + client: Caller-owned PyMongo async client taking precedence over connection settings. + database: Caller-owned resolved PyMongo async database taking precedence over connection settings. + collection: Caller-owned resolved PyMongo async collection taking precedence over connection settings. client_options: Created-client options limited to ``tls``, ``retryWrites``, and ``appname``. + Ignored when a client, database, or collection is injected. definition: Explicit dictionary schema. - collection_name: Collection name overriding the model definition. + collection_name: Collection name overriding the model definition. With an injected client + or database, ``AZURE_DOCUMENTDB_COLLECTION_NAME`` is the fallback when neither an + explicit value nor the model definition supplies a name. embedding_generator: Default local embedding generator. - env_file_path: Optional selected settings file. - env_file_encoding: Encoding of the selected settings file. + env_file_path: Optional selected settings file, ignored for injected objects. + env_file_encoding: Encoding of the selected settings file, ignored for injected objects. """ if collection is not None: if not isinstance(collection, AsyncCollection): raise TypeError("collection must be a PyMongo AsyncCollection.") + if client is not None or database is not None: + raise ValueError("Provide at most one of client, database, or collection.") if collection_name is not None and collection_name != collection.name: raise ValueError("collection_name must match the injected collection.") - if any( - value is not None - for value in ( - connection_string, - database_name, - client, - database, - client_options, - env_file_path, - env_file_encoding, - ) - ): - raise ValueError("collection cannot be combined with a client, database, or connection settings.") resolved_collection_name = collection.name else: resolved_collection_name = collection_name + if collection is not None or client is not None or database 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_collection_name + if resolved_collection_name is not None + else getattr(registered_definition, "collection_name", None) + ) + resolved_collection_name = _resolve_resource_name(configured_name, "collection") super().__init__( record_type, definition=definition, @@ -1116,13 +1142,16 @@ def __init__( Args: connection_string: MongoDB connection string for a connector-owned client. - database_name: Database selected for a created or injected client. - client: Caller-owned PyMongo async client. Requires ``database_name``. - database: Caller-owned resolved PyMongo async database. + Ignored when a client or database is injected. + database_name: Database selected for a created or injected client. An injected account + client falls back to ``AZURE_DOCUMENTDB_DATABASE_NAME`` when this is omitted. + client: Caller-owned PyMongo async client taking precedence over connection settings. + database: Caller-owned resolved PyMongo async database taking precedence over connection settings. client_options: Created-client options limited to ``tls``, ``retryWrites``, and ``appname``. + Ignored when a client or database is injected. embedding_generator: Default embedding generator for child collections. - env_file_path: Optional selected settings file. - env_file_encoding: Encoding of the selected settings file. + env_file_path: Optional selected settings file, ignored for injected objects. + env_file_encoding: Encoding of the selected settings file, ignored for injected objects. """ super().__init__( embedding_generator=embedding_generator, diff --git a/python/packages/azure-documentdb/tests/azure_documentdb/test_settings.py b/python/packages/azure-documentdb/tests/azure_documentdb/test_settings.py index 56323e42cbf..52dae4dff1d 100644 --- a/python/packages/azure-documentdb/tests/azure_documentdb/test_settings.py +++ b/python/packages/azure-documentdb/tests/azure_documentdb/test_settings.py @@ -2,11 +2,12 @@ from __future__ import annotations -from typing import get_type_hints +from typing import Any, get_type_hints from unittest.mock import AsyncMock, MagicMock, patch import pytest -from agent_framework import SecretString, load_settings +from agent_framework import SecretString, VectorStoreCollectionDefinition, load_settings, register_vectorstoremodel +from agent_framework.exceptions import SettingNotFoundError from pymongo import AsyncMongoClient from pymongo.asynchronous.database import AsyncDatabase @@ -129,20 +130,157 @@ async def test_injected_database_and_collection_bypass_settings(definition_facto @pytest.mark.parametrize( - "kwargs", - [ - {"connection_string": "mongodb://explicit"}, - {"client_options": {"tls": True}}, - {"env_file_path": "selected.env"}, - {"env_file_encoding": "utf-16"}, - ], + "target", + ["store-account", "store-database", "collection-account", "collection-database", "collection-collection"], ) -def test_injected_client_rejects_connection_settings(kwargs, pymongo_objects): - with ( - patch.object(module, "load_settings", side_effect=AssertionError("settings must be bypassed")), - pytest.raises(ValueError, match="cannot be combined"), - ): - AzureDocumentDBStore(client=pymongo_objects[0], database_name="db", **kwargs) +async def test_injected_objects_ignore_connection_settings( + target, definition_factory, pymongo_objects, monkeypatch, tmp_path +): + client, database, collection = pymongo_objects + monkeypatch.setenv("AZURE_DOCUMENTDB_CONNECTION_STRING", "") + monkeypatch.setenv("AZURE_DOCUMENTDB_DATABASE_NAME", "x" * 64) + monkeypatch.setenv("AZURE_DOCUMENTDB_COLLECTION_NAME", "x" * 256) + kwargs: dict[str, Any] = { + "connection_string": "mongodb://ignored", + "client_options": {"tls": False}, + "env_file_path": str(tmp_path / "missing.env"), + "env_file_encoding": "invalid-encoding", + } + kind = target.split("-")[1] + if kind == "account": + kwargs.update(client=client, database_name="explicit_db") + elif kind == "database": + kwargs.update(database=database, database_name="test_database") + else: + kwargs.update(collection=collection, database_name="ignored", collection_name="documents") + if target.startswith("collection") and kind != "collection": + kwargs["collection_name"] = "explicit_collection" + + with patch.object(module, "load_settings", side_effect=AssertionError("settings must be bypassed")): + resource: Any + if target.startswith("store"): + resource = AzureDocumentDBStore(**kwargs) + else: + resource = AzureDocumentDBCollection(dict, definition=definition_factory(), **kwargs) + await resource.aclose() + + assert not resource.managed_client + client.close.assert_not_awaited() + + +@pytest.mark.parametrize("target", ["store", "collection"]) +@pytest.mark.parametrize("explicit", [False, True]) +def test_injected_account_resolves_database_independently( + target, explicit, definition_factory, pymongo_objects, monkeypatch +): + client, database, _ = pymongo_objects + monkeypatch.setenv("AZURE_DOCUMENTDB_DATABASE_NAME", "environment_db") + kwargs: dict[str, Any] = {"client": client} + if explicit: + kwargs["database_name"] = "explicit_db" + resource: Any + if target == "store": + resource = AzureDocumentDBStore(**kwargs) + else: + resource = AzureDocumentDBCollection(dict, definition=definition_factory(), **kwargs) + + expected = "explicit_db" if explicit else "environment_db" + client.__getitem__.assert_called_once_with(expected) + assert resource._client.database is database + + +@pytest.mark.parametrize("kind", ["account", "database"]) +@pytest.mark.parametrize("explicit", [False, True]) +def test_injected_collection_resolves_collection_independently( + kind, explicit, definition_factory, pymongo_objects, monkeypatch +): + client, database, _ = pymongo_objects + monkeypatch.setenv("AZURE_DOCUMENTDB_COLLECTION_NAME", "environment_collection") + definition = VectorStoreCollectionDefinition(definition_factory().fields) + kwargs = {"client": client, "database_name": "db"} if kind == "account" else {"database": database} + resource = AzureDocumentDBCollection( + dict, + definition=definition, + collection_name="explicit_collection" if explicit else None, + **kwargs, + ) + + assert resource.collection_name == ("explicit_collection" if explicit else "environment_collection") + + +@pytest.mark.parametrize("env_name", [None, "environment_collection"]) +def test_injected_collection_uses_registered_definition_before_environment( + env_name, definition_factory, pymongo_objects, monkeypatch +): + if env_name is None: + monkeypatch.delenv("AZURE_DOCUMENTDB_COLLECTION_NAME", raising=False) + else: + monkeypatch.setenv("AZURE_DOCUMENTDB_COLLECTION_NAME", env_name) + definition = VectorStoreCollectionDefinition(definition_factory().fields, collection_name="registered_collection") + resource = AzureDocumentDBCollection( + dict, + definition=definition, + client=pymongo_objects[0], + database_name="db", + ) + + assert resource.collection_name == "registered_collection" + assert resource.definition is definition + + +@pytest.mark.parametrize("kind", ["account", "database"]) +def test_registered_model_definition_precedes_environment_without_explicit_definition( + kind, definition_factory, pymongo_objects, monkeypatch +): + class RegisteredRecord: + pass + + definition = VectorStoreCollectionDefinition(definition_factory().fields, collection_name="registered_collection") + register_vectorstoremodel(RegisteredRecord, definition=definition, decoder=lambda data: RegisteredRecord()) + monkeypatch.setenv("AZURE_DOCUMENTDB_COLLECTION_NAME", "environment_collection") + client, database, _ = pymongo_objects + kwargs: dict[str, Any] = {"client": client, "database_name": "db"} if kind == "account" else {"database": database} + + resource = AzureDocumentDBCollection(RegisteredRecord, **kwargs) + + assert resource.collection_name == "registered_collection" + assert resource.definition is definition + + +@pytest.mark.parametrize("target", ["store-account", "collection-account", "collection-database"]) +def test_injected_objects_report_missing_selector(target, definition_factory, pymongo_objects, monkeypatch): + monkeypatch.delenv("AZURE_DOCUMENTDB_DATABASE_NAME", raising=False) + monkeypatch.delenv("AZURE_DOCUMENTDB_COLLECTION_NAME", raising=False) + client, database, _ = pymongo_objects + kwargs = {"database": database} if target.endswith("database") else {"client": client} + selector = "AZURE_DOCUMENTDB_DATABASE_NAME" if target == "store-account" else "AZURE_DOCUMENTDB_COLLECTION_NAME" + with pytest.raises(SettingNotFoundError, match=selector) as exc: + if target.startswith("store"): + AzureDocumentDBStore(**kwargs) + else: + AzureDocumentDBCollection( + dict, + definition=VectorStoreCollectionDefinition(definition_factory().fields), + **kwargs, + ) + parameter = "database_name" if target == "store-account" else "collection_name" + assert parameter in str(exc.value) + + +@pytest.mark.parametrize("target", ["store-database", "collection-database", "collection-collection"]) +def test_injected_proxy_rejects_mismatched_name(target, definition_factory, pymongo_objects): + _, database, collection = pymongo_objects + kwargs: dict[str, Any] = {"collection": collection, "collection_name": "other"} + message = "collection_name must match the injected collection." + if target.endswith("database"): + kwargs = {"database": database, "database_name": "other"} + message = "database_name must match the injected database." + with pytest.raises(ValueError) as exc: + if target.startswith("store"): + AzureDocumentDBStore(**kwargs) + else: + AzureDocumentDBCollection(dict, definition=definition_factory(), **kwargs) + assert str(exc.value) == message async def test_store_children_reuse_resolved_client(definition_factory, pymongo_objects):