Merge branch 'main' into fix/vertexai-py-batch-size

gemini-embedding-* accepts one input text per predict request, but
embed_batch() chunked at 250 for every model while defaulting to
gemini-embedding-001. Derive the chunk size from the model instead.

mem0-ts/src/oss/src/memory/index.ts

@@ -149,27 +149,6 @@ function rejectTopLevelEntityParams(
   }
 }
 
-const ENTITY_ID_ALIASES: Record<string, string> = {
-  userId: "user_id",
-  agentId: "agent_id",
-  runId: "run_id",
-};
-
-/** Rewrites camelCase entity ids inside filters to the snake_case keys vector stores index on. */
-function normalizeEntityFilterKeys(
-  filters: Record<string, any>,
-): Record<string, any> {
-  const normalized = { ...filters };
-  for (const [camelKey, snakeKey] of Object.entries(ENTITY_ID_ALIASES)) {
-    if (!(camelKey in normalized)) continue;
-    if (normalized[snakeKey] === undefined) {
-      normalized[snakeKey] = normalized[camelKey];
-    }
-    delete normalized[camelKey];
-  }
-  return normalized;
-}
-
 /**
  * Validates and normalizes an entity ID.
  * - Coerces non-string ids (e.g. numeric database keys) to string
@@ -1378,20 +1357,16 @@ export class Memory {
     // receive `agent_id: undefined` / `run_id: undefined` and fail
     // (Qdrant rejects the malformed match, pgvector binds NULL, Redis
     // emits a literal "undefined" string in TAG filters).
-    const requestedFilters = normalizeEntityFilterKeys(config.filters ?? {});
     const normalizedFilters: Record<string, any> = config.filters
       ? Object.fromEntries(
           Object.entries({
-            ...requestedFilters,
-            user_id: validateAndTrimEntityId(
-              requestedFilters.user_id,
-              "user_id",
-            ),
+            ...config.filters,
+            user_id: validateAndTrimEntityId(config.filters.user_id, "user_id"),
             agent_id: validateAndTrimEntityId(
-              requestedFilters.agent_id,
+              config.filters.agent_id,
               "agent_id",
             ),
-            run_id: validateAndTrimEntityId(requestedFilters.run_id, "run_id"),
+            run_id: validateAndTrimEntityId(config.filters.run_id, "run_id"),
           }).filter(([, v]) => v !== undefined),
         )
       : {};
@@ -1878,16 +1853,12 @@ export class Memory {
     // Validate and trim entity IDs in filters. Drop keys that resolve to
     // undefined so downstream vector stores don't receive
     // `agent_id: undefined` / `run_id: undefined` and fail.
-    const requestedFilters = normalizeEntityFilterKeys(config.filters || {});
     const filters: Record<string, any> = Object.fromEntries(
       Object.entries({
-        ...requestedFilters,
-        user_id: validateAndTrimEntityId(requestedFilters.user_id, "user_id"),
-        agent_id: validateAndTrimEntityId(
-          requestedFilters.agent_id,
-          "agent_id",
-        ),
-        run_id: validateAndTrimEntityId(requestedFilters.run_id, "run_id"),
+        ...(config.filters || {}),
+        user_id: validateAndTrimEntityId(config.filters?.user_id, "user_id"),
+        agent_id: validateAndTrimEntityId(config.filters?.agent_id, "agent_id"),
+        run_id: validateAndTrimEntityId(config.filters?.run_id, "run_id"),
       }).filter(([, v]) => v !== undefined),
     );
 

mem0-ts/src/oss/tests/memory.validation.test.ts

@@ -281,66 +281,6 @@ describe("Memory Input Validation", () => {
     });
   });
 
-  describe("camelCase entity IDs in filters", () => {
-    it("maps camelCase filter keys to snake_case in search", async () => {
-      const searchSpy = jest
-        .spyOn((memory as any).vectorStore, "search")
-        .mockResolvedValue([]);
-
-      await memory.search("q", {
-        filters: { userId: "alice", agentId: "bot", runId: "run-1" },
-      });
-
-      const passedFilters = searchSpy.mock.calls[0][2] as Record<string, any>;
-      expect(passedFilters).toMatchObject({
-        user_id: "alice",
-        agent_id: "bot",
-        run_id: "run-1",
-      });
-      expect(passedFilters.userId).toBeUndefined();
-      expect(passedFilters.agentId).toBeUndefined();
-      expect(passedFilters.runId).toBeUndefined();
-
-      searchSpy.mockRestore();
-    });
-
-    it("maps camelCase filter keys to snake_case in getAll", async () => {
-      const listSpy = jest
-        .spyOn((memory as any).vectorStore, "list")
-        .mockResolvedValue([[], 0]);
-
-      await memory.getAll({ filters: { userId: "alice" } });
-
-      const passedFilters = listSpy.mock.calls[0][0] as Record<string, any>;
-      expect(passedFilters.user_id).toBe("alice");
-      expect(passedFilters.userId).toBeUndefined();
-
-      listSpy.mockRestore();
-    });
-
-    it("prefers an explicit snake_case value over its camelCase alias", async () => {
-      const listSpy = jest
-        .spyOn((memory as any).vectorStore, "list")
-        .mockResolvedValue([[], 0]);
-
-      await memory.getAll({
-        filters: { user_id: "snake", userId: "camel" },
-      });
-
-      const passedFilters = listSpy.mock.calls[0][0] as Record<string, any>;
-      expect(passedFilters.user_id).toBe("snake");
-      expect(passedFilters.userId).toBeUndefined();
-
-      listSpy.mockRestore();
-    });
-
-    it("validates camelCase filter values like their snake_case counterparts", async () => {
-      await expect(
-        memory.search("q", { filters: { userId: "user 123" } }),
-      ).rejects.toThrow("Invalid user_id: cannot contain whitespace");
-    });
-  });
-
   describe("search() filter entity ID validation", () => {
     it("should throw error when user_id in filters is whitespace-only", async () => {
       await expect(

mem0/embeddings/vertexai.py

@@ -8,6 +8,14 @@ from mem0.embeddings.base import EmbeddingBase
 from mem0.utils.gcp_auth import GCPAuthenticator
 
 
+def _max_instances_per_request(model: str) -> int:
+    # Vertex AI's per-request input cap depends on the model family: gemini-embedding-*
+    # accepts exactly one input text per predict() call, while text-embedding-004/005 and
+    # text-multilingual-embedding-002 accept up to 250.
+    # https://docs.cloud.google.com/vertex-ai/docs/samples/aiplatform-sdk-embedding
+    return 1 if model.startswith("gemini-embedding") else 250
+
+
 class VertexAIEmbedding(EmbeddingBase):
     def __init__(self, config: Optional[BaseEmbedderConfig] = None):
         super().__init__(config)
@@ -72,8 +80,9 @@ class VertexAIEmbedding(EmbeddingBase):
                 raise ValueError(f"Invalid memory action: {memory_action}")
             embedding_type = self.embedding_types[memory_action]
         all_embeddings = []
-        for i in range(0, len(texts), 250):
-            chunk = texts[i : i + 250]
+        chunk_size = _max_instances_per_request(self.config.model)
+        for i in range(0, len(texts), chunk_size):
+            chunk = texts[i : i + chunk_size]
             inputs = [TextEmbeddingInput(text=t, task_type=embedding_type) for t in chunk]
             results = self.model.get_embeddings(texts=inputs, output_dimensionality=self.config.embedding_dims)
             all_embeddings.extend(r.values for r in results)

tests/embeddings/test_vertexai_embeddings.py

@@ -163,7 +163,9 @@ def test_invalid_memory_action(mock_text_embedding_model, mock_config):
 
 @patch("mem0.embeddings.vertexai.TextEmbeddingModel")
 def test_embed_batch_single_call(mock_text_embedding_model, mock_os_environ, mock_config):
-    mock_config.return_value.model = "gemini-embedding-001"
+    # text-embedding-005 accepts up to 250 inputs per request, so 2 texts fit in one call.
+    # (gemini-embedding-* only accepts 1 input per request; see test_embed_batch_default_model_sends_one_text_per_call.)
+    mock_config.return_value.model = "text-embedding-005"
     mock_config.return_value.embedding_dims = 256
 
     config = mock_config()
@@ -198,7 +200,9 @@ def test_embed_batch_empty_list(mock_text_embedding_model, mock_os_environ, mock
 
 @patch("mem0.embeddings.vertexai.TextEmbeddingModel")
 def test_embed_batch_count_mismatch_raises(mock_text_embedding_model, mock_os_environ, mock_config):
-    mock_config.return_value.model = "gemini-embedding-001"
+    # Use a 250-cap model so both texts land in a single chunk/call, keeping the mismatch
+    # (1 embedding returned for that one call) triggered by that single call.
+    mock_config.return_value.model = "text-embedding-005"
     mock_config.return_value.embedding_dims = 256
 
     config = mock_config()
@@ -261,8 +265,8 @@ def test_embed_batch_invalid_memory_action_raises(mock_text_embedding_model, moc
 
 @patch("mem0.embeddings.vertexai.TextEmbeddingModel")
 def test_embed_batch_chunking_triggers_two_api_calls(mock_text_embedding_model, mock_os_environ, mock_config):
-    """300 texts must produce exactly 2 get_embeddings calls (chunks of 250 and 50)."""
-    mock_config.return_value.model = "gemini-embedding-001"
+    """text-embedding-005 accepts up to 250 inputs/request: 300 texts -> 2 calls (chunks of 250 and 50)."""
+    mock_config.return_value.model = "text-embedding-005"
     mock_config.return_value.embedding_dims = 256
 
     config = mock_config()
@@ -278,3 +282,25 @@ def test_embed_batch_chunking_triggers_two_api_calls(mock_text_embedding_model,
 
     assert mock_text_embedding_model.from_pretrained.return_value.get_embeddings.call_count == 2
     assert len(result) == 300
+
+
+@patch("mem0.embeddings.vertexai.TextEmbeddingModel")
+def test_embed_batch_default_model_sends_one_text_per_call(mock_text_embedding_model, mock_os_environ, mock_config):
+    """gemini-embedding-001 (default) accepts 1 input/request: 300 texts -> 300 calls of 1 text each."""
+    mock_config.return_value.model = "gemini-embedding-001"
+    mock_config.return_value.embedding_dims = 256
+
+    config = mock_config()
+    embedder = VertexAIEmbedding(config)
+
+    def make_single_response(texts, output_dimensionality):
+        assert len(texts) == 1
+        return [Mock(values=[0.1, 0.2]) for _ in texts]
+
+    mock_text_embedding_model.from_pretrained.return_value.get_embeddings.side_effect = make_single_response
+
+    texts = [f"text {i}" for i in range(300)]
+    result = embedder.embed_batch(texts)
+
+    assert mock_text_embedding_model.from_pretrained.return_value.get_embeddings.call_count == 300
+    assert len(result) == 300