Resolving comments and adding some code improvements

Author: Aayush Kataria

agent_memory_toolkit/_utils.py

@@ -128,6 +128,53 @@ def _resolve_embedding_dimensions(val: Optional[int]) -> int:
     return parsed
 
 
+_ALLOWED_EMBEDDING_DATA_TYPES = ("float32", "uint8", "int8")
+_ALLOWED_DISTANCE_FUNCTIONS = ("cosine", "dotproduct", "euclidean")
+
+
+def _resolve_embedding_data_type(val: Optional[str]) -> str:
+    """Resolve embedding data type from explicit value or ``AI_FOUNDRY_EMBEDDING_DATA_TYPE`` env var.
+
+    Defaults to ``float32``. Raises :class:`ConfigurationError` for unknown values.
+    """
+    raw = (val if val is not None else os.environ.get("AI_FOUNDRY_EMBEDDING_DATA_TYPE") or "float32").strip()
+    if raw not in _ALLOWED_EMBEDDING_DATA_TYPES:
+        raise ConfigurationError(
+            message=(
+                f"Invalid configuration for embedding_data_type: must be one of "
+                f"{_ALLOWED_EMBEDDING_DATA_TYPES}, got {raw!r}"
+            ),
+            parameter="embedding_data_type",
+        )
+    return raw
+
+
+def _resolve_distance_function(val: Optional[str]) -> str:
+    """Resolve distance function from explicit value or ``AI_FOUNDRY_EMBEDDING_DISTANCE_FUNCTION`` env var.
+
+    Defaults to ``cosine``. Raises :class:`ConfigurationError` for unknown values.
+    """
+    raw = (val if val is not None else os.environ.get("AI_FOUNDRY_EMBEDDING_DISTANCE_FUNCTION") or "cosine").strip()
+    if raw not in _ALLOWED_DISTANCE_FUNCTIONS:
+        raise ConfigurationError(
+            message=(
+                f"Invalid configuration for distance_function: must be one of "
+                f"{_ALLOWED_DISTANCE_FUNCTIONS}, got {raw!r}"
+            ),
+            parameter="distance_function",
+        )
+    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.
+
+    Defaults to ``en-US``. Empty values fall back to the default.
+    """
+    raw = (val if val is not None else os.environ.get("COSMOS_DB_FULL_TEXT_LANGUAGE") or "en-US").strip()
+    return raw or "en-US"
+
+
 def _resolve_cosmos_throughput_mode(val: Optional[str]) -> str:
     """Resolve throughput mode from explicit value or env var.
 

agent_memory_toolkit/aio/cosmos_memory_client.py

@@ -26,7 +26,10 @@
     _make_memory,
     _resolve_cosmos_provisioning_autoscale_max_ru,
     _resolve_cosmos_throughput_mode,
+    _resolve_distance_function,
+    _resolve_embedding_data_type,
     _resolve_embedding_dimensions,
+    _resolve_full_text_language,
     _validate_connection,
     _validate_hybrid_search,
 )
@@ -482,9 +485,9 @@ async def create_memory_store(
             lease_partition_key = PartitionKey(path="/id")
             vec_policy, idx_policy, ft_policy = _container_policies(
                 embedding_dimensions=embedding_dimensions or self._embedding_dimensions or 1536,
-                embedding_data_type=embedding_data_type or "float32",
-                distance_function=distance_function or "cosine",
-                full_text_language=full_text_language or "en-US",
+                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),
             )
             offer_throughput = _cosmos_container_offer_throughput(
                 throughput_mode=self._cosmos_throughput_mode,

agent_memory_toolkit/cosmos_memory_client.py

@@ -22,7 +22,10 @@
     _make_memory,
     _resolve_cosmos_provisioning_autoscale_max_ru,
     _resolve_cosmos_throughput_mode,
+    _resolve_distance_function,
+    _resolve_embedding_data_type,
     _resolve_embedding_dimensions,
+    _resolve_full_text_language,
     _validate_connection,
     _validate_hybrid_search,
 )
@@ -456,9 +459,9 @@ def create_memory_store(
             lease_partition_key = PartitionKey(path="/id")
             vec_policy, idx_policy, ft_policy = _container_policies(
                 embedding_dimensions=embedding_dimensions or self._embedding_dimensions or 1536,
-                embedding_data_type=embedding_data_type or "float32",
-                distance_function=distance_function or "cosine",
-                full_text_language=full_text_language or "en-US",
+                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),
             )
             offer_throughput = _cosmos_container_offer_throughput(
                 throughput_mode=self._cosmos_throughput_mode,

infra/modules/cosmos.bicep

@@ -113,6 +113,12 @@ resource memoriesContainer 'Microsoft.DocumentDB/databaseAccounts/sqlDatabases/c
           {
             path: '/embedding/?'
           }
+          {
+            path: '/source_memory_ids/*'
+          }
+          {
+            path: '/supersedes_ids/*'
+          }
           {
             path: '/"_etag"/?'
           }

tests/unit/test_utils.py

@@ -6,9 +6,12 @@
     DEFAULT_TTL_BY_TYPE,
     _build_container_kwargs,
     _make_memory,
+    _resolve_distance_function,
+    _resolve_embedding_data_type,
+    _resolve_full_text_language,
     compute_content_hash,
 )
-from agent_memory_toolkit.exceptions import ValidationError
+from agent_memory_toolkit.exceptions import ConfigurationError, ValidationError
 
 
 def test_build_container_kwargs_includes_required_fields_and_extras():
@@ -180,3 +183,50 @@ def test_make_memory_invalid_role():
 def test_make_memory_invalid_type():
     with pytest.raises(ValidationError):
         _make_memory(user_id="u1", role="user", content="test", memory_type="invalid")
+
+
+def test_resolve_embedding_data_type_defaults(monkeypatch):
+    monkeypatch.delenv("AI_FOUNDRY_EMBEDDING_DATA_TYPE", raising=False)
+    assert _resolve_embedding_data_type(None) == "float32"
+
+
+def test_resolve_embedding_data_type_from_env(monkeypatch):
+    monkeypatch.setenv("AI_FOUNDRY_EMBEDDING_DATA_TYPE", "int8")
+    assert _resolve_embedding_data_type(None) == "int8"
+
+
+def test_resolve_embedding_data_type_explicit_overrides_env(monkeypatch):
+    monkeypatch.setenv("AI_FOUNDRY_EMBEDDING_DATA_TYPE", "int8")
+    assert _resolve_embedding_data_type("uint8") == "uint8"
+
+
+def test_resolve_embedding_data_type_invalid_raises(monkeypatch):
+    monkeypatch.setenv("AI_FOUNDRY_EMBEDDING_DATA_TYPE", "bogus")
+    with pytest.raises(ConfigurationError):
+        _resolve_embedding_data_type(None)
+
+
+def test_resolve_distance_function_defaults(monkeypatch):
+    monkeypatch.delenv("AI_FOUNDRY_EMBEDDING_DISTANCE_FUNCTION", raising=False)
+    assert _resolve_distance_function(None) == "cosine"
+
+
+def test_resolve_distance_function_from_env(monkeypatch):
+    monkeypatch.setenv("AI_FOUNDRY_EMBEDDING_DISTANCE_FUNCTION", "dotproduct")
+    assert _resolve_distance_function(None) == "dotproduct"
+
+
+def test_resolve_distance_function_invalid_raises(monkeypatch):
+    monkeypatch.setenv("AI_FOUNDRY_EMBEDDING_DISTANCE_FUNCTION", "manhattan")
+    with pytest.raises(ConfigurationError):
+        _resolve_distance_function(None)
+
+
+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"
+
+
+def test_resolve_full_text_language_from_env(monkeypatch):
+    monkeypatch.setenv("COSMOS_DB_FULL_TEXT_LANGUAGE", "fr-FR")
+    assert _resolve_full_text_language(None) == "fr-FR"