Make Cosmos vector index type configurable (default diskANN) (#24)

* Make Cosmos vector index type configurable (default diskANN)

Add vector_index_type to create_memory_store (sync + async) and the AI_FOUNDRY_EMBEDDING_VECTOR_INDEX_TYPE env var, resolved via _resolve_vector_index_type with allow-list (diskANN, quantizedFlat, flat). Previously the memories container's vector index was hard-coded to diskANN, which blocked accounts without the DiskANN capability (e.g. the classic Cosmos DB emulator) and emulator-backed integration test pipelines.

* Address review: append vector_index_type at end of create_memory_store

Move the new optional vector_index_type parameter to the end of the
create_memory_store signature (sync + async) so it does not shift the
positional order of full_text_language/throughput_mode/autoscale_max_ru,
preserving backwards compatibility for positional callers.

Author: TheovanKraay

.env.template

@@ -55,6 +55,11 @@ AI_FOUNDRY_EMBEDDING_DEPLOYMENT_NAME=text-embedding-3-large
 AI_FOUNDRY_EMBEDDING_DIMENSIONS=1536
 AI_FOUNDRY_EMBEDDING_DATA_TYPE=float32
 AI_FOUNDRY_EMBEDDING_DISTANCE_FUNCTION=cosine
+# Optional. Vector index type for the memories container: diskANN (default),
+# quantizedFlat, or flat. diskANN requires the DiskANN capability on the Cosmos
+# DB account; use quantizedFlat or flat for accounts without it (e.g. the
+# classic Cosmos DB emulator).
+AI_FOUNDRY_EMBEDDING_VECTOR_INDEX_TYPE=diskANN
 COSMOS_DB_FULL_TEXT_LANGUAGE=en-US
 
 AI_FOUNDRY_CHAT_DEPLOYMENT_NAME=<your-model-deployment>

CHANGELOG.md

@@ -1,5 +1,16 @@
 ## Release History
 
+### Unreleased
+
+#### Other Changes
+* The memories container's vector index type is now configurable instead of being
+  hard-coded to `diskANN`. Set it via the `vector_index_type` argument to
+  `create_memory_store(...)` or the `AI_FOUNDRY_EMBEDDING_VECTOR_INDEX_TYPE`
+  environment variable. Allowed values are `diskANN` (default), `quantizedFlat`,
+  and `flat`. This lets the toolkit run against Cosmos DB accounts without the
+  DiskANN capability (for example the classic Cosmos DB emulator), enabling
+  emulator-backed integration test pipelines.
+
 ### 0.1.0b2 (2026-06-03)
 
 #### Bugs Fixed

azure/cosmos/agent_memory/_utils.py

@@ -175,6 +175,7 @@ def _resolve_embedding_dimensions(val: Optional[int]) -> int:
 
 _ALLOWED_EMBEDDING_DATA_TYPES = ("float32", "uint8", "int8")
 _ALLOWED_DISTANCE_FUNCTIONS = ("cosine", "dotproduct", "euclidean")
+_ALLOWED_VECTOR_INDEX_TYPES = ("diskANN", "quantizedFlat", "flat")
 
 
 def _resolve_embedding_data_type(val: Optional[str]) -> str:
@@ -211,6 +212,27 @@ def _resolve_distance_function(val: Optional[str]) -> str:
     return raw
 
 
+def _resolve_vector_index_type(val: Optional[str]) -> str:
+    """Resolve vector index type from explicit value or ``AI_FOUNDRY_EMBEDDING_VECTOR_INDEX_TYPE`` env var.
+
+    Defaults to ``diskANN``. Raises :class:`ConfigurationError` for unknown values.
+
+    ``diskANN`` requires the Cosmos DB account to have the DiskANN vector index
+    capability enabled. Accounts that do not (for example the classic Cosmos DB
+    emulator) can use ``quantizedFlat`` or ``flat`` instead.
+    """
+    raw = (val if val is not None else os.environ.get("AI_FOUNDRY_EMBEDDING_VECTOR_INDEX_TYPE") or "diskANN").strip()
+    if raw not in _ALLOWED_VECTOR_INDEX_TYPES:
+        raise ConfigurationError(
+            message=(
+                f"Invalid configuration for vector_index_type: must be one of "
+                f"{_ALLOWED_VECTOR_INDEX_TYPES}, got {raw!r}"
+            ),
+            parameter="vector_index_type",
+        )
+    return raw
+
+
 def _resolve_full_text_language(val: Optional[str]) -> str:
     """Resolve full-text language from explicit value or ``COSMOS_DB_FULL_TEXT_LANGUAGE`` env var.
 
@@ -368,6 +390,7 @@ def _container_policies(
     embedding_data_type: str,
     distance_function: str,
     full_text_language: str,
+    vector_index_type: str = "diskANN",
 ) -> tuple[dict, dict, dict]:
     """Build the vector, indexing, and full-text policies for container creation."""
     vector_embedding_policy = {
@@ -388,7 +411,7 @@ def _container_policies(
             {"path": "/source_memory_ids/*"},
             {"path": "/supersedes_ids/*"},
         ],
-        "vectorIndexes": [{"path": "/embedding", "type": "diskANN"}],
+        "vectorIndexes": [{"path": "/embedding", "type": vector_index_type}],
         "fullTextIndexes": [{"path": "/content"}],
         # Procedural synthesis selects TOP N by (salience DESC, created_at ASC, id ASC).
         # Cosmos requires a composite index for multi-property ORDER BY; without it the

azure/cosmos/agent_memory/aio/cosmos_memory_client.py

@@ -18,6 +18,7 @@
     _resolve_distance_function,
     _resolve_embedding_data_type,
     _resolve_full_text_language,
+    _resolve_vector_index_type,
     _validate_connection,
 )
 from azure.cosmos.agent_memory.aio.auto_trigger import maybe_trigger_steps
@@ -263,6 +264,7 @@ async def create_memory_store(
         full_text_language: Optional[str] = None,
         throughput_mode: Optional[str] = None,
         autoscale_max_ru: Optional[int] = None,
+        vector_index_type: Optional[str] = None,
     ) -> None:
         """Create the Cosmos DB database and memory/counter/lease containers."""
         self._cosmos_endpoint = endpoint or self._cosmos_endpoint
@@ -310,6 +312,7 @@ async def create_memory_store(
                 embedding_data_type=_resolve_embedding_data_type(embedding_data_type),
                 distance_function=_resolve_distance_function(distance_function),
                 full_text_language=_resolve_full_text_language(full_text_language),
+                vector_index_type=_resolve_vector_index_type(vector_index_type),
             )
             self._memories_container_client = await db.create_container_if_not_exists(
                 **_build_container_kwargs(

azure/cosmos/agent_memory/cosmos_memory_client.py

@@ -19,6 +19,7 @@
     _resolve_distance_function,
     _resolve_embedding_data_type,
     _resolve_full_text_language,
+    _resolve_vector_index_type,
     _validate_connection,
 )
 from .auto_trigger import maybe_trigger_steps
@@ -235,6 +236,7 @@ def create_memory_store(
         full_text_language: Optional[str] = None,
         throughput_mode: Optional[str] = None,
         autoscale_max_ru: Optional[int] = None,
+        vector_index_type: Optional[str] = None,
     ) -> None:
         """Create the Cosmos DB database and memory/counter/lease containers."""
         self._cosmos_endpoint = endpoint or self._cosmos_endpoint
@@ -281,6 +283,7 @@ def create_memory_store(
                 embedding_data_type=_resolve_embedding_data_type(embedding_data_type),
                 distance_function=_resolve_distance_function(distance_function),
                 full_text_language=_resolve_full_text_language(full_text_language),
+                vector_index_type=_resolve_vector_index_type(vector_index_type),
             )
             self._memories_container_client = db.create_container_if_not_exists(
                 **_build_container_kwargs(

tests/unit/test_utils.py

@@ -5,10 +5,12 @@
 from azure.cosmos.agent_memory._utils import (
     DEFAULT_TTL_BY_TYPE,
     _build_container_kwargs,
+    _container_policies,
     _make_memory,
     _resolve_distance_function,
     _resolve_embedding_data_type,
     _resolve_full_text_language,
+    _resolve_vector_index_type,
     compute_content_hash,
 )
 from azure.cosmos.agent_memory.exceptions import ConfigurationError, ValidationError
@@ -226,6 +228,48 @@ def test_resolve_distance_function_invalid_raises(monkeypatch):
         _resolve_distance_function(None)
 
 
+def test_resolve_vector_index_type_defaults(monkeypatch):
+    monkeypatch.delenv("AI_FOUNDRY_EMBEDDING_VECTOR_INDEX_TYPE", raising=False)
+    assert _resolve_vector_index_type(None) == "diskANN"
+
+
+def test_resolve_vector_index_type_from_env(monkeypatch):
+    monkeypatch.setenv("AI_FOUNDRY_EMBEDDING_VECTOR_INDEX_TYPE", "quantizedFlat")
+    assert _resolve_vector_index_type(None) == "quantizedFlat"
+
+
+def test_resolve_vector_index_type_explicit_overrides_env(monkeypatch):
+    monkeypatch.setenv("AI_FOUNDRY_EMBEDDING_VECTOR_INDEX_TYPE", "quantizedFlat")
+    assert _resolve_vector_index_type("flat") == "flat"
+
+
+def test_resolve_vector_index_type_invalid_raises(monkeypatch):
+    monkeypatch.setenv("AI_FOUNDRY_EMBEDDING_VECTOR_INDEX_TYPE", "hnsw")
+    with pytest.raises(ConfigurationError):
+        _resolve_vector_index_type(None)
+
+
+def test_container_policies_defaults_to_diskann():
+    _, indexing_policy, _ = _container_policies(
+        embedding_dimensions=1536,
+        embedding_data_type="float32",
+        distance_function="cosine",
+        full_text_language="en-US",
+    )
+    assert indexing_policy["vectorIndexes"] == [{"path": "/embedding", "type": "diskANN"}]
+
+
+def test_container_policies_uses_supplied_vector_index_type():
+    _, indexing_policy, _ = _container_policies(
+        embedding_dimensions=1536,
+        embedding_data_type="float32",
+        distance_function="cosine",
+        full_text_language="en-US",
+        vector_index_type="quantizedFlat",
+    )
+    assert indexing_policy["vectorIndexes"] == [{"path": "/embedding", "type": "quantizedFlat"}]
+
+
 def test_resolve_full_text_language_defaults(monkeypatch):
     monkeypatch.delenv("COSMOS_DB_FULL_TEXT_LANGUAGE", raising=False)
     assert _resolve_full_text_language(None) == "en-US"