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/jest.integration.config.js

@@ -2,6 +2,8 @@
 module.exports = {
   ...require("./jest.config"),
   testMatch: ["**/integration/**/*.test.ts"],
+  globalSetup: "<rootDir>/src/client/tests/integration/global-setup.ts",
+  globalTeardown: "<rootDir>/src/client/tests/integration/global-teardown.ts",
   // Run integration tests serially to avoid rate limiting and race conditions
   maxWorkers: 1,
 };

mem0-ts/src/client/tests/integration/global-setup.ts

@@ -0,0 +1,40 @@
+/**
+ * Jest global setup for integration tests.
+ *
+ * Runs a full project cleanup before any integration test starts,
+ * then waits 10 seconds for the async cleanup to propagate.
+ */
+import { MemoryClient } from "../../mem0";
+
+export default async function globalSetup() {
+  const apiKey = process.env.MEM0_API_KEY;
+  if (!apiKey) return; // skip if no key — tests will be skipped too
+
+  const client = new MemoryClient({ apiKey });
+  await client.ping();
+
+  console.log("[integration] Running pre-test cleanup...");
+
+  // Full project wipe — all four filters set explicitly
+  try {
+    await client.deleteAll({
+      userId: "*",
+      agentId: "*",
+      appId: "*",
+      runId: "*",
+    });
+  } catch {
+    // ignore — may 404 if no data exists
+  }
+
+  try {
+    await client.deleteUsers();
+  } catch {
+    // ignore — may throw "No entities to delete"
+  }
+
+  // Wait 10 seconds for async cleanup to propagate
+  console.log("[integration] Waiting 10s for cleanup to propagate...");
+  await new Promise((r) => setTimeout(r, 10_000));
+  console.log("[integration] Pre-test cleanup done.");
+}

mem0-ts/src/client/tests/integration/global-teardown.ts

@@ -0,0 +1,36 @@
+/**
+ * Jest global teardown for integration tests.
+ *
+ * Runs a full project cleanup after all integration tests complete
+ * so no test data is left behind.
+ */
+import { MemoryClient } from "../../mem0";
+
+export default async function globalTeardown() {
+  const apiKey = process.env.MEM0_API_KEY;
+  if (!apiKey) return;
+
+  const client = new MemoryClient({ apiKey });
+  await client.ping();
+
+  console.log("[integration] Running post-test cleanup...");
+
+  try {
+    await client.deleteAll({
+      userId: "*",
+      agentId: "*",
+      appId: "*",
+      runId: "*",
+    });
+  } catch {
+    // ignore
+  }
+
+  try {
+    await client.deleteUsers();
+  } catch {
+    // ignore
+  }
+
+  console.log("[integration] Post-test cleanup done.");
+}

mem0-ts/src/client/tests/integration/helpers.ts

@@ -52,17 +52,13 @@ export async function withRetry<T>(
 }
 
 /**
- * Poll getAll until a user's memories reach minCount and the ID set
- * stops changing between polls.
- *
- * The Mem0 API processes memories asynchronously, and consolidating a
- * new fact replaces a memory (delete + add) instead of editing it in
- * place. Returning on the first non-empty read hands callers IDs the
- * pipeline then invalidates, so compare IDs rather than length: a
- * replacement leaves the count unchanged.
+ * Poll getAll until memories appear for a user.
+ * The Mem0 API processes memories asynchronously — after add()
+ * we need to wait for them to be available.
  *
  * Polls every 15 seconds with a maximum of 4 retries to avoid
- * hitting rate limits. Throws if minCount is never reached.
+ * hitting rate limits. Throws if results aren't available after
+ * all retries.
  */
 export async function waitForMemories(
   client: MemoryClient,
@@ -70,32 +66,18 @@ export async function waitForMemories(
   minCount: number,
   maxRetries = 4,
 ): Promise<Memory[]> {
-  let memories: Memory[] = [];
-  let previousIds = "";
-
   for (let attempt = 1; attempt <= maxRetries; attempt++) {
     const response = await withRetry(() =>
       client.getAll({ filters: { user_id: userId } }),
     );
-    memories = response.results ?? [];
-
-    const ids = memories
-      .map((m) => m.id)
-      .sort()
-      .join(",");
-    if (memories.length >= minCount && ids === previousIds) {
+    const memories = response.results ?? [];
+    if (memories.length >= minCount) {
       return memories;
     }
-    previousIds = ids;
-
     if (attempt < maxRetries) {
       await new Promise((r) => setTimeout(r, 15_000));
     }
   }
-
-  if (memories.length >= minCount) {
-    return memories;
-  }
   throw new Error(
     `waitForMemories: expected at least ${minCount} memories for user "${userId}" but did not get them after ${maxRetries} attempts`,
   );
@@ -220,3 +202,32 @@ export async function cleanupTestUser(
     // ignore
   }
 }
+
+/**
+ * Full project wipe — deletes all memories and all entities.
+ * Equivalent to Python SDK's:
+ *   client.delete_all(user_id="*", agent_id="*", app_id="*", run_id="*")
+ *
+ * Used as cleanup before and after integration test runs so tests
+ * start from a clean slate and don't leave data behind.
+ */
+export async function fullProjectCleanup(client: MemoryClient): Promise<void> {
+  // Delete all memories — all four filters set explicitly
+  try {
+    await client.deleteAll({
+      userId: "*",
+      agentId: "*",
+      appId: "*",
+      runId: "*",
+    });
+  } catch {
+    // ignore — may 404 if no data exists
+  }
+
+  // Delete all entities (users, agents, apps, runs)
+  try {
+    await client.deleteUsers();
+  } catch {
+    // ignore — may throw "No entities to delete"
+  }
+}

mem0-ts/src/client/tests/integration/management.test.ts

@@ -100,7 +100,6 @@ describeIntegration("MemoryClient Integration — Users & Project", () => {
     const hookName = `test-hook-${randomUUID().slice(0, 8)}`;
     const hookUrl = `https://example.com/webhook/${randomUUID().slice(0, 8)}`;
     const updatedName = `updated-hook-${randomUUID().slice(0, 8)}`;
-    const updatedUrl = `https://example.com/updated/${randomUUID().slice(0, 8)}`;
 
     afterAll(async () => {
       if (createdWebhookId) {
@@ -169,7 +168,7 @@ describeIntegration("MemoryClient Integration — Users & Project", () => {
         client.updateWebhook({
           webhookId: createdWebhookId,
           name: updatedName,
-          url: updatedUrl,
+          url: "https://example.com/updated",
           eventTypes: [WebhookEvent.MEMORY_DELETED],
         }),
       );

mem0-ts/src/client/tests/waitForMemories.test.ts

@@ -1,69 +0,0 @@
-import { waitForMemories } from "./integration/helpers";
-
-const clientReturning = (polls: Array<Array<{ id: string }>>) => {
-  let call = 0;
-  return {
-    getAll: jest.fn(async () => ({
-      results: polls[Math.min(call++, polls.length - 1)],
-    })),
-  } as any;
-};
-
-const settle = async <T>(pending: Promise<T>): Promise<T> => {
-  pending.catch(() => {});
-  for (let i = 0; i < 4; i++) {
-    await jest.advanceTimersByTimeAsync(15_000);
-  }
-  return pending;
-};
-
-describe("waitForMemories", () => {
-  beforeEach(() => jest.useFakeTimers());
-  afterEach(() => jest.useRealTimers());
-
-  test("waits for the id set to stop changing", async () => {
-    const client = clientReturning([
-      [{ id: "a" }],
-      [{ id: "a" }, { id: "b" }],
-      [{ id: "a" }, { id: "b" }],
-    ]);
-
-    const memories = await settle(waitForMemories(client, "user", 1));
-
-    expect(memories.map((m) => m.id)).toEqual(["a", "b"]);
-    expect(client.getAll).toHaveBeenCalledTimes(3);
-  });
-
-  test("keeps polling when consolidation replaces a memory at a constant count", async () => {
-    const client = clientReturning([
-      [{ id: "a" }, { id: "b" }],
-      [{ id: "a" }, { id: "c" }],
-      [{ id: "a" }, { id: "c" }],
-    ]);
-
-    const memories = await settle(waitForMemories(client, "user", 1));
-
-    expect(memories.map((m) => m.id)).toEqual(["a", "c"]);
-  });
-
-  test("returns the last read when the set never settles but minCount is met", async () => {
-    const client = clientReturning([
-      [{ id: "a" }],
-      [{ id: "b" }],
-      [{ id: "c" }],
-      [{ id: "d" }],
-    ]);
-
-    const memories = await settle(waitForMemories(client, "user", 1));
-
-    expect(memories.map((m) => m.id)).toEqual(["d"]);
-  });
-
-  test("throws when minCount is never reached", async () => {
-    const client = clientReturning([[]]);
-
-    await expect(settle(waitForMemories(client, "user", 1))).rejects.toThrow(
-      /expected at least 1 memories/,
-    );
-  });
-});

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