diff --git a/headroom/memory/adapters/embedders.py b/headroom/memory/adapters/embedders.py index 8be40a742..1df384028 100644 --- a/headroom/memory/adapters/embedders.py +++ b/headroom/memory/adapters/embedders.py @@ -14,7 +14,7 @@ import asyncio import logging import os from functools import cached_property -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, cast import numpy as np @@ -290,6 +290,7 @@ class OnnxLocalEmbedder: DEFAULT_DIMENSION = 384 DEFAULT_MAX_TOKENS = 256 ONNX_REPO = "Qdrant/all-MiniLM-L6-v2-onnx" + MAX_BATCH_SIZE = 2 def __init__(self, max_length: int = 256) -> None: self._max_length = max_length @@ -330,17 +331,12 @@ class OnnxLocalEmbedder: logger.info("ONNX embedding model loaded (384-dim, no torch)") - def _embed_single(self, text: str) -> np.ndarray: - """Embed a single text string.""" - assert self._session is not None - assert self._tokenizer is not None - - if not text or not text.strip(): - return np.zeros(self.DEFAULT_DIMENSION, dtype=np.float32) - - encoded = self._tokenizer.encode(text) - input_ids = np.array([encoded.ids], dtype=np.int64) - attention_mask = np.array([encoded.attention_mask], dtype=np.int64) + def _build_feeds( + self, + input_ids: np.ndarray, + attention_mask: np.ndarray, + ) -> dict[str, np.ndarray]: + """Build ONNX feeds for a token batch.""" token_type_ids = np.zeros_like(input_ids, dtype=np.int64) feeds: dict[str, np.ndarray] = {} @@ -352,16 +348,37 @@ class OnnxLocalEmbedder: elif "token_type_ids" in name: feeds[name] = token_type_ids - outputs = self._session.run(None, feeds) - token_embeddings = outputs[0] # (1, seq_len, 384) + return feeds + + def _embed_many(self, texts: list[str]) -> np.ndarray: + """Embed multiple non-empty text strings in one ONNX pass.""" + assert self._session is not None + assert self._tokenizer is not None + + encodings = self._tokenizer.encode_batch(texts) + input_ids = np.array([encoding.ids for encoding in encodings], dtype=np.int64) + attention_mask = np.array( + [encoding.attention_mask for encoding in encodings], dtype=np.int64 + ) + + outputs = self._session.run(None, self._build_feeds(input_ids, attention_mask)) + token_embeddings = outputs[0] # (batch, seq_len, 384) # Mean pooling over non-padding tokens mask_expanded = attention_mask[:, :, np.newaxis].astype(np.float32) summed = np.sum(token_embeddings * mask_expanded, axis=1) counts = np.clip(mask_expanded.sum(axis=1), a_min=1e-9, a_max=None) - embedding = summed / counts + embeddings = summed / counts - return _normalize_embedding(embedding[0]) + return _normalize_embeddings_batch(embeddings) + + def _embed_single(self, text: str) -> np.ndarray: + """Embed a single text string.""" + if not text or not text.strip(): + return np.zeros(self.DEFAULT_DIMENSION, dtype=np.float32) + + embedding = self._embed_many([text])[0] + return cast(np.ndarray, embedding) async def embed(self, text: str) -> np.ndarray: """Generate an embedding for a single text.""" @@ -369,18 +386,40 @@ class OnnxLocalEmbedder: if self._session is None: await asyncio.get_event_loop().run_in_executor(None, self._load_model) - return await asyncio.get_event_loop().run_in_executor(None, self._embed_single, text) + loop = asyncio.get_event_loop() + embedding = await loop.run_in_executor(None, self._embed_single, text) + return cast(np.ndarray, embedding) async def embed_batch(self, texts: list[str]) -> list[np.ndarray]: """Generate embeddings for multiple texts.""" + if not texts: + return [] + async with self._lock: if self._session is None: await asyncio.get_event_loop().run_in_executor(None, self._load_model) - results = [] - for text in texts: - emb = await asyncio.get_event_loop().run_in_executor(None, self._embed_single, text) - results.append(emb) + non_empty_indices: list[int] = [] + non_empty_texts: list[str] = [] + for i, text in enumerate(texts): + if text and text.strip(): + non_empty_indices.append(i) + non_empty_texts.append(text) + + results: list[np.ndarray] = [ + np.zeros(self.dimension, dtype=np.float32) for _ in range(len(texts)) + ] + if not non_empty_texts: + return results + + loop = asyncio.get_event_loop() + for start in range(0, len(non_empty_texts), self.MAX_BATCH_SIZE): + batch_texts = non_empty_texts[start : start + self.MAX_BATCH_SIZE] + batch_indices = non_empty_indices[start : start + self.MAX_BATCH_SIZE] + embeddings = await loop.run_in_executor(None, self._embed_many, batch_texts) + for idx, embedding in zip(batch_indices, embeddings): + results[idx] = embedding + return results @property diff --git a/headroom/memory/adapters/sqlite_vector.py b/headroom/memory/adapters/sqlite_vector.py index fa65705ec..e67310320 100644 --- a/headroom/memory/adapters/sqlite_vector.py +++ b/headroom/memory/adapters/sqlite_vector.py @@ -24,8 +24,8 @@ import struct from dataclasses import dataclass from datetime import datetime from pathlib import Path -from threading import RLock -from typing import TYPE_CHECKING, Any +from threading import RLock, get_ident +from typing import TYPE_CHECKING, Any, cast import numpy as np @@ -37,6 +37,8 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) +_SQLITE_QUERY_CHUNK_SIZE = 500 + # sqlite-vec availability check _SQLITE_VEC_AVAILABLE: bool | None = None _sqlite_vec_module: Any = None @@ -226,11 +228,12 @@ class SQLiteVectorIndex: self._db_path = Path(db_path) self._page_cache_size_kb = page_cache_size_kb self._lock = RLock() + self._connections: dict[int, sqlite3.Connection] = {} self._init_db() - def _get_conn(self) -> sqlite3.Connection: - """Get a database connection with sqlite-vec loaded.""" + def _create_conn(self) -> sqlite3.Connection: + """Create a SQLite connection with sqlite-vec loaded.""" conn = sqlite3.connect(str(self._db_path)) conn.row_factory = sqlite3.Row @@ -245,6 +248,98 @@ class SQLiteVectorIndex: return conn + def _get_conn(self) -> sqlite3.Connection: + """Get a cached per-thread SQLite connection with sqlite-vec loaded.""" + thread_id = get_ident() + conn = self._connections.get(thread_id) + if conn is None: + conn = self._create_conn() + self._connections[thread_id] = conn + return conn + + def _close_cached_connections(self) -> None: + """Close all cached SQLite connections.""" + for conn in self._connections.values(): + try: + conn.close() + except sqlite3.Error: + logger.debug("Failed to close cached sqlite-vec connection", exc_info=True) + self._connections.clear() + + @staticmethod + def _chunked(items: list[Any], chunk_size: int = _SQLITE_QUERY_CHUNK_SIZE) -> list[list[Any]]: + """Split a list into SQLite-friendly chunks.""" + return [items[i : i + chunk_size] for i in range(0, len(items), chunk_size)] + + def _select_rowids_by_memory_ids( + self, + conn: sqlite3.Connection, + memory_ids: list[str], + ) -> dict[str, int]: + """Fetch rowids for the given memory IDs.""" + rowids: dict[str, int] = {} + for chunk in self._chunked(memory_ids): + placeholders = ", ".join("?" for _ in chunk) + rows = conn.execute( + f"SELECT rowid, memory_id FROM vec_metadata WHERE memory_id IN ({placeholders})", + chunk, + ).fetchall() + for row in rows: + rowids[str(row["memory_id"])] = int(row["rowid"]) + return rowids + + def _prepare_memory_for_index(self, memory: Memory) -> tuple[np.ndarray, VectorMetadata]: + """Validate a memory and prepare it for indexing.""" + if memory.embedding is None: + raise ValueError(f"Memory {memory.id} has no embedding") + + embedding = np.asarray(memory.embedding, dtype=np.float32) + if embedding.shape[0] != self._dimension: + raise ValueError( + f"Embedding dimension {embedding.shape[0]} does not match " + f"index dimension {self._dimension}" + ) + + return embedding, VectorMetadata.from_memory(memory) + + def _metadata_insert_params(self, memory_id: str, metadata: VectorMetadata) -> tuple[Any, ...]: + """Build INSERT parameters for vector metadata.""" + return ( + memory_id, + metadata.user_id, + metadata.session_id, + metadata.agent_id, + metadata.importance, + metadata.created_at.isoformat(), + metadata.valid_until.isoformat() if metadata.valid_until else None, + json.dumps(metadata.entity_refs), + metadata.content, + json.dumps(metadata.metadata or {}), + ) + + def _metadata_update_params(self, metadata: VectorMetadata, rowid: int) -> tuple[Any, ...]: + """Build UPDATE parameters for vector metadata.""" + return ( + metadata.user_id, + metadata.session_id, + metadata.agent_id, + metadata.importance, + metadata.created_at.isoformat(), + metadata.valid_until.isoformat() if metadata.valid_until else None, + json.dumps(metadata.entity_refs), + metadata.content, + json.dumps(metadata.metadata or {}), + rowid, + ) + + @staticmethod + def _cursor_lastrowid(cursor: sqlite3.Cursor) -> int: + """Return a non-null SQLite cursor lastrowid.""" + rowid = cursor.lastrowid + if rowid is None: + raise RuntimeError("sqlite-vec insert did not produce a rowid") + return cast(int, rowid) + def _init_db(self) -> None: """Initialize the database schema.""" with self._get_conn() as conn: @@ -320,37 +415,21 @@ class SQLiteVectorIndex: Raises: ValueError: If memory has no embedding or wrong dimension. """ - if memory.embedding is None: - raise ValueError(f"Memory {memory.id} has no embedding") - - embedding = np.asarray(memory.embedding, dtype=np.float32) - if embedding.shape[0] != self._dimension: - raise ValueError( - f"Embedding dimension {embedding.shape[0]} does not match " - f"index dimension {self._dimension}" - ) - - metadata = VectorMetadata.from_memory(memory) + embedding, metadata = self._prepare_memory_for_index(memory) with self._lock: with self._get_conn() as conn: - # Check if already exists existing = conn.execute( "SELECT rowid FROM vec_metadata WHERE memory_id = ?", (memory.id,), ).fetchone() if existing: - # Update existing entry - rowid = existing[0] - - # Update vector + rowid = int(existing[0]) conn.execute( "UPDATE vec_embeddings SET embedding = ? WHERE rowid = ?", (self._serialize_f32(embedding), rowid), ) - - # Update metadata conn.execute( """ UPDATE vec_metadata SET @@ -359,22 +438,9 @@ class SQLiteVectorIndex: entity_refs = ?, content = ?, metadata_json = ? WHERE rowid = ? """, - ( - metadata.user_id, - metadata.session_id, - metadata.agent_id, - metadata.importance, - metadata.created_at.isoformat(), - metadata.valid_until.isoformat() if metadata.valid_until else None, - json.dumps(metadata.entity_refs), - metadata.content, - json.dumps(metadata.metadata or {}), - rowid, - ), + self._metadata_update_params(metadata, rowid), ) else: - # Insert new entry - # First insert metadata to get rowid cursor = conn.execute( """ INSERT INTO vec_metadata ( @@ -383,22 +449,9 @@ class SQLiteVectorIndex: entity_refs, content, metadata_json ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, - ( - memory.id, - metadata.user_id, - metadata.session_id, - metadata.agent_id, - metadata.importance, - metadata.created_at.isoformat(), - metadata.valid_until.isoformat() if metadata.valid_until else None, - json.dumps(metadata.entity_refs), - metadata.content, - json.dumps(metadata.metadata or {}), - ), + self._metadata_insert_params(memory.id, metadata), ) - rowid = cursor.lastrowid - - # Insert vector with matching rowid + rowid = self._cursor_lastrowid(cursor) conn.execute( "INSERT INTO vec_embeddings (rowid, embedding) VALUES (?, ?)", (rowid, self._serialize_f32(embedding)), @@ -415,15 +468,120 @@ class SQLiteVectorIndex: Returns: Number of memories indexed. """ - indexed = 0 + prepared: list[tuple[str, np.ndarray, VectorMetadata]] = [] for memory in memories: - if memory.embedding is not None: - try: - await self.index(memory) - indexed += 1 - except ValueError: - pass - return indexed + try: + embedding, metadata = self._prepare_memory_for_index(memory) + except ValueError: + continue + prepared.append((memory.id, embedding, metadata)) + + if not prepared: + return 0 + + memory_ids = [memory_id for memory_id, _, _ in prepared] + + with self._lock: + with self._get_conn() as conn: + if len(set(memory_ids)) != len(memory_ids): + existing_rowids = self._select_rowids_by_memory_ids( + conn, list(dict.fromkeys(memory_ids)) + ) + for memory_id, embedding, metadata in prepared: + rowid = existing_rowids.get(memory_id) + if rowid is None: + cursor = conn.execute( + """ + INSERT INTO vec_metadata ( + memory_id, user_id, session_id, agent_id, + importance, created_at, valid_until, + entity_refs, content, metadata_json + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + """, + self._metadata_insert_params(memory_id, metadata), + ) + rowid = self._cursor_lastrowid(cursor) + existing_rowids[memory_id] = rowid + conn.execute( + "INSERT INTO vec_embeddings (rowid, embedding) VALUES (?, ?)", + (rowid, self._serialize_f32(embedding)), + ) + else: + conn.execute( + "UPDATE vec_embeddings SET embedding = ? WHERE rowid = ?", + (self._serialize_f32(embedding), rowid), + ) + conn.execute( + """ + UPDATE vec_metadata SET + user_id = ?, session_id = ?, agent_id = ?, + importance = ?, created_at = ?, valid_until = ?, + entity_refs = ?, content = ?, metadata_json = ? + WHERE rowid = ? + """, + self._metadata_update_params(metadata, rowid), + ) + + conn.commit() + return len(prepared) + + existing_rowids = self._select_rowids_by_memory_ids(conn, memory_ids) + metadata_updates: list[tuple[Any, ...]] = [] + vector_updates: list[tuple[bytes, int]] = [] + metadata_inserts: list[tuple[Any, ...]] = [] + new_memory_ids: list[str] = [] + new_vectors: list[tuple[str, bytes]] = [] + + for memory_id, embedding, metadata in prepared: + rowid = existing_rowids.get(memory_id) + serialized = self._serialize_f32(embedding) + if rowid is None: + metadata_inserts.append(self._metadata_insert_params(memory_id, metadata)) + new_memory_ids.append(memory_id) + new_vectors.append((memory_id, serialized)) + else: + vector_updates.append((serialized, rowid)) + metadata_updates.append(self._metadata_update_params(metadata, rowid)) + + if vector_updates: + conn.executemany( + "UPDATE vec_embeddings SET embedding = ? WHERE rowid = ?", + vector_updates, + ) + if metadata_updates: + conn.executemany( + """ + UPDATE vec_metadata SET + user_id = ?, session_id = ?, agent_id = ?, + importance = ?, created_at = ?, valid_until = ?, + entity_refs = ?, content = ?, metadata_json = ? + WHERE rowid = ? + """, + metadata_updates, + ) + if metadata_inserts: + conn.executemany( + """ + INSERT INTO vec_metadata ( + memory_id, user_id, session_id, agent_id, + importance, created_at, valid_until, + entity_refs, content, metadata_json + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + """, + metadata_inserts, + ) + inserted_rowids = self._select_rowids_by_memory_ids(conn, new_memory_ids) + conn.executemany( + "INSERT INTO vec_embeddings (rowid, embedding) VALUES (?, ?)", + [ + (inserted_rowids[memory_id], serialized) + for memory_id, serialized in new_vectors + ], + ) + + conn.commit() + + return len(prepared) async def remove(self, memory_id: str) -> bool: """Remove a memory from the index. @@ -463,11 +621,30 @@ class SQLiteVectorIndex: Returns: Number removed. """ - removed = 0 - for memory_id in memory_ids: - if await self.remove(memory_id): - removed += 1 - return removed + unique_ids = list(dict.fromkeys(memory_ids)) + if not unique_ids: + return 0 + + with self._lock: + with self._get_conn() as conn: + rowids_by_memory_id = self._select_rowids_by_memory_ids(conn, unique_ids) + rowids = list(rowids_by_memory_id.values()) + if not rowids: + return 0 + + for rowid_chunk in self._chunked(rowids): + placeholders = ", ".join("?" for _ in rowid_chunk) + conn.execute( + f"DELETE FROM vec_embeddings WHERE rowid IN ({placeholders})", + rowid_chunk, + ) + conn.execute( + f"DELETE FROM vec_metadata WHERE rowid IN ({placeholders})", + rowid_chunk, + ) + + conn.commit() + return len(rowids) async def search(self, filter: VectorFilter) -> list[VectorSearchResult]: """Search for similar vectors. @@ -746,4 +923,5 @@ class SQLiteVectorIndex: async def close(self) -> None: """Close the index (cleanup).""" - pass # Connection-per-request pattern, nothing to close + with self._lock: + self._close_cached_connections() diff --git a/headroom/memory/mcp_server.py b/headroom/memory/mcp_server.py index 298e60ff2..e767c8efa 100644 --- a/headroom/memory/mcp_server.py +++ b/headroom/memory/mcp_server.py @@ -135,14 +135,16 @@ async def _warm_up_backend(backend: LocalBackend, user_id: str) -> None: if not all_memories: return - indexed = 0 - for mem in all_memories: - if mem.embedding is None: - mem.embedding = await hm._embedder.embed(mem.content) - await hm._store.save(mem) - await hm._vector_index.index(mem) - indexed += 1 + memories_missing_embeddings = [mem for mem in all_memories if mem.embedding is None] + if memories_missing_embeddings: + embeddings = await hm._embedder.embed_batch( + [mem.content for mem in memories_missing_embeddings] + ) + for mem, embedding in zip(memories_missing_embeddings, embeddings): + mem.embedding = embedding + await hm._store.save_batch(memories_missing_embeddings) + indexed = await hm._vector_index.index_batch(all_memories) logger.info(f"Memory MCP: indexed {indexed} memories into vector store") diff --git a/tests/test_memory/test_hierarchical.py b/tests/test_memory/test_hierarchical.py index e9e92556f..5394fedaf 100644 --- a/tests/test_memory/test_hierarchical.py +++ b/tests/test_memory/test_hierarchical.py @@ -809,3 +809,56 @@ class TestLocalEmbedder: def test_dimension_property(self, embedder): """Test that dimension property returns correct value.""" assert embedder.dimension == 384 + + +class TestOnnxLocalEmbedder: + """Tests for OnnxLocalEmbedder batching behavior.""" + + @pytest.mark.asyncio + async def test_embed_batch_uses_batched_onnx_inference(self): + """Test that non-empty inputs share ONNX batch inference.""" + from headroom.memory.adapters.embedders import OnnxLocalEmbedder + + class FakeEncoding: + def __init__(self, ids: list[int], attention_mask: list[int]) -> None: + self.ids = ids + self.attention_mask = attention_mask + + class FakeTokenizer: + def encode_batch(self, texts: list[str]) -> list[FakeEncoding]: + encodings = [] + for i, text in enumerate(texts, start=1): + token = len(text) + i + encodings.append(FakeEncoding([token, token + 1, 0], [1, 1, 0])) + return encodings + + class FakeSession: + def __init__(self) -> None: + self.run_calls = 0 + + def run(self, _output_names, feeds): + self.run_calls += 1 + input_ids = feeds["input_ids"] + batch_size, seq_len = input_ids.shape + token_embeddings = np.zeros((batch_size, seq_len, 384), dtype=np.float32) + token_embeddings[:, :, 0] = input_ids + token_embeddings[:, :, 1] = input_ids * 0.5 + return [token_embeddings] + + embedder = OnnxLocalEmbedder() + embedder.MAX_BATCH_SIZE = 8 + embedder._session = FakeSession() + embedder._tokenizer = FakeTokenizer() + embedder._input_names = ["input_ids", "attention_mask", "token_type_ids"] + + embeddings = await embedder.embed_batch(["alpha", " ", "beta", "gamma"]) + + assert len(embeddings) == 4 + assert embedder._session.run_calls == 1 + assert np.array_equal(embeddings[1], np.zeros(384, dtype=np.float32)) + assert embeddings[0].shape == (384,) + assert embeddings[2].shape == (384,) + assert embeddings[3].shape == (384,) + assert not np.allclose(embeddings[0], 0.0) + assert not np.allclose(embeddings[2], 0.0) + assert not np.allclose(embeddings[3], 0.0) diff --git a/tests/test_memory/test_mcp_server.py b/tests/test_memory/test_mcp_server.py new file mode 100644 index 000000000..85207d2f4 --- /dev/null +++ b/tests/test_memory/test_mcp_server.py @@ -0,0 +1,63 @@ +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import numpy as np +import pytest + +pytest.importorskip("mcp") + +from headroom.memory.mcp_server import _warm_up_backend +from headroom.memory.models import Memory + + +@pytest.mark.asyncio +async def test_warm_up_backend_batches_embedding_and_indexing() -> None: + """Warm-up should batch missing embeddings and vector indexing.""" + warmup_embedding = np.ones(384, dtype=np.float32) + batch_embeddings = [ + np.full(384, 2.0, dtype=np.float32), + np.full(384, 3.0, dtype=np.float32), + ] + + embedder = SimpleNamespace( + embed=AsyncMock(return_value=warmup_embedding), + embed_batch=AsyncMock(return_value=batch_embeddings), + ) + store = SimpleNamespace(save_batch=AsyncMock()) + vector_index = SimpleNamespace(index_batch=AsyncMock(return_value=3)) + + memory_without_embedding_a = Memory(content="First", user_id="alice") + memory_with_embedding = Memory( + content="Second", + user_id="alice", + embedding=np.full(384, 5.0, dtype=np.float32), + ) + memory_without_embedding_b = Memory(content="Third", user_id="alice") + memories = [ + memory_without_embedding_a, + memory_with_embedding, + memory_without_embedding_b, + ] + + backend = SimpleNamespace( + _ensure_initialized=AsyncMock(), + _hierarchical_memory=SimpleNamespace( + _embedder=embedder, + _store=store, + _vector_index=vector_index, + ), + get_user_memories=AsyncMock(return_value=memories), + ) + + await _warm_up_backend(backend, "alice") + + backend._ensure_initialized.assert_awaited_once() + backend.get_user_memories.assert_awaited_once_with("alice", limit=500) + embedder.embed.assert_awaited_once_with("warmup") + embedder.embed_batch.assert_awaited_once_with(["First", "Third"]) + store.save_batch.assert_awaited_once_with( + [memory_without_embedding_a, memory_without_embedding_b] + ) + vector_index.index_batch.assert_awaited_once_with(memories) + assert np.array_equal(memory_without_embedding_a.embedding, batch_embeddings[0]) + assert np.array_equal(memory_without_embedding_b.embedding, batch_embeddings[1]) diff --git a/tests/test_sqlite_vector_index.py b/tests/test_sqlite_vector_index.py index c67958ed0..1b2ca2190 100644 --- a/tests/test_sqlite_vector_index.py +++ b/tests/test_sqlite_vector_index.py @@ -444,6 +444,34 @@ class TestSQLiteVectorIndexEdgeCases: assert indexed == 10 assert index.size == 10 + @pytest.mark.asyncio + async def test_batch_index_uses_single_connection(self, index, monkeypatch): + """Test batch indexing reuses a single sqlite-vec connection.""" + np.random.seed(42) + memories = [ + Memory( + content=f"Content {i}", + user_id="alice", + embedding=np.random.randn(384).astype(np.float32), + ) + for i in range(10) + ] + + original_get_conn = index._get_conn + conn_calls = 0 + + def counting_get_conn(): + nonlocal conn_calls + conn_calls += 1 + return original_get_conn() + + monkeypatch.setattr(index, "_get_conn", counting_get_conn) + + indexed = await index.index_batch(memories) + + assert indexed == 10 + assert conn_calls == 1 + @pytest.mark.asyncio async def test_batch_remove(self, index): """Test batch removal.""" @@ -466,3 +494,32 @@ class TestSQLiteVectorIndexEdgeCases: assert removed == 2 assert index.size == 3 + + @pytest.mark.asyncio + async def test_batch_remove_uses_single_connection(self, index, monkeypatch): + """Test batch removal reuses a single sqlite-vec connection.""" + np.random.seed(42) + memories = [] + for i in range(5): + memory = Memory( + content=f"Content {i}", + user_id="alice", + embedding=np.random.randn(384).astype(np.float32), + ) + await index.index(memory) + memories.append(memory) + + original_get_conn = index._get_conn + conn_calls = 0 + + def counting_get_conn(): + nonlocal conn_calls + conn_calls += 1 + return original_get_conn() + + monkeypatch.setattr(index, "_get_conn", counting_get_conn) + + removed = await index.remove_batch([memories[0].id, memories[2].id, "nonexistent"]) + + assert removed == 2 + assert conn_calls == 1