fixed issues with embedding generation on client and ADF side

Author: jcodella

.env.template

@@ -7,7 +7,7 @@ COSMOS_DB_AUTOSCALE_MAX_RU=1000
 AI_FOUNDRY_ENDPOINT=https://<your-account>.openai.azure.com/
 AI_FOUNDRY_API_KEY=
 EMBEDDING_MODEL=text-embedding-3-large
-EMBEDDING_DIMENSIONS=1535
+EMBEDDING_DIMENSIONS=1536
 EMBEDDING_DATA_TYPE=float32
 EMBEDDING_DISTANCE_FUNCTION=cosine
 FULL_TEXT_LANGUAGE=en-US

agent_memory_toolkit/async_memory.py

@@ -1,6 +1,6 @@
 """AsyncAgentMemory: Async version of AgentMemory using azure.cosmos.aio and AsyncAzureOpenAI."""
 
-import uuid
+import uuid, os
 from datetime import datetime, timezone
 from typing import Any, Optional
 
@@ -28,6 +28,7 @@ def __init__(
         ai_foundry_credential: Optional["TokenCredential"] = None,
         ai_foundry_api_key: Optional[str] = None,
         embedding_model: str = "text-embedding-3-large",
+        embedding_dimensions: Optional[int] = None,
         adf_endpoint: Optional[str] = None,
         adf_key: Optional[str] = None,
         use_default_credential: bool = True,
@@ -57,6 +58,9 @@ def __init__(
         self.ai_foundry_credential = ai_foundry_credential
         self.ai_foundry_api_key = ai_foundry_api_key
         self.embedding_model = embedding_model
+        self.embedding_dimensions = embedding_dimensions or int(
+            _os.environ.get("EMBEDDING_DIMENSIONS", "0") or "0"
+        ) or None
         self._embeddings_client = None
 
         self.adf_endpoint = adf_endpoint
@@ -292,9 +296,13 @@ async def _get_embedding(self, text: str) -> list[float]:
                     azure_ad_token_provider=token_provider,
                 )
 
-        response = await self._embeddings_client.embeddings.create(
-            input=[text], model=self.embedding_model,
-        )
+        kwargs: dict[str, Any] = {
+            "input": [text],
+            "model": self.embedding_model,
+        }
+        if self.embedding_dimensions:
+            kwargs["dimensions"] = self.embedding_dimensions
+        response = await self._embeddings_client.embeddings.create(**kwargs)
         return response.data[0].embedding
 
     # ------------------------------------------------------------------

agent_memory_toolkit/memory.py

@@ -1,6 +1,6 @@
 """AgentMemory: A class for managing agent memories locally and (eventually) in Cosmos DB."""
 
-import uuid
+import uuid, os
 from datetime import datetime, timezone
 from typing import Any, Optional
 
@@ -97,6 +97,7 @@ def __init__(
         ai_foundry_credential: Optional["TokenCredential"] = None,
         ai_foundry_api_key: Optional[str] = None,
         embedding_model: str = "text-embedding-3-large",
+        embedding_dimensions: Optional[int] = None,
         adf_endpoint: Optional[str] = None,
         adf_key: Optional[str] = None,
         use_default_credential: bool = True,
@@ -131,6 +132,9 @@ def __init__(
         self.ai_foundry_credential = ai_foundry_credential
         self.ai_foundry_api_key = ai_foundry_api_key
         self.embedding_model = embedding_model
+        self.embedding_dimensions = embedding_dimensions or int(
+            os.environ.get("EMBEDDING_DIMENSIONS", "0") or "0"
+        ) or None
         self._embeddings_client = None
 
         # Azure Durable Functions configuration
@@ -457,10 +461,13 @@ def _get_embedding(self, text: str) -> list[float]:
                     azure_ad_token_provider=token_provider,
                 )
 
-        response = self._embeddings_client.embeddings.create(
-            input=[text],
-            model=self.embedding_model,
-        )
+        kwargs: dict[str, Any] = {
+            "input": [text],
+            "model": self.embedding_model,
+        }
+        if self.embedding_dimensions:
+            kwargs["dimensions"] = self.embedding_dimensions
+        response = self._embeddings_client.embeddings.create(**kwargs)
         return response.data[0].embedding
 
     # ------------------------------------------------------------------

azure_functions/activities.py

@@ -51,10 +51,11 @@ def _get_embeddings_client():
 
         endpoint = os.environ["AI_FOUNDRY_ENDPOINT"]
         api_key = os.environ.get("AI_FOUNDRY_API_KEY")
+        api_version = os.environ.get("AI_FOUNDRY_API_VERSION", "2024-12-01-preview")
 
         if api_key:
             _embeddings_client = AzureOpenAI(
-                api_version="2024-12-01-preview",
+                api_version=api_version,
                 azure_endpoint=endpoint,
                 api_key=api_key,
             )
@@ -66,7 +67,7 @@ def _get_embeddings_client():
                 "https://cognitiveservices.azure.com/.default",
             )
             _embeddings_client = AzureOpenAI(
-                api_version="2024-12-01-preview",
+                api_version=api_version,
                 azure_endpoint=endpoint,
                 azure_ad_token_provider=token_provider,
             )
@@ -84,10 +85,11 @@ def _get_chat_client():
 
         endpoint = os.environ["AI_FOUNDRY_ENDPOINT"]
         api_key = os.environ.get("AI_FOUNDRY_API_KEY")
+        api_version = os.environ.get("AI_FOUNDRY_API_VERSION", "2024-12-01-preview")
 
         if api_key:
             _chat_client = AzureOpenAI(
-                api_version="2024-12-01-preview",
+                api_version=api_version,
                 azure_endpoint=endpoint,
                 api_key=api_key,
             )
@@ -99,7 +101,7 @@ def _get_chat_client():
                 "https://cognitiveservices.azure.com/.default",
             )
             _chat_client = AzureOpenAI(
-                api_version="2024-12-01-preview",
+                api_version=api_version,
                 azure_endpoint=endpoint,
                 azure_ad_token_provider=token_provider,
             )