mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-10 14:27:00 -04:00
fix(memory): batch onnx embeddings and sqlite-vec ops
Make the ONNX + sqlite-vec memory path truly batched. Batch ONNX embed_batch calls, batch sqlite-vec index/remove work under a single cached connection, and update MCP warm-up to use batch embed/save/index flows. Add focused regression tests for ONNX batching, sqlite-vec single-connection batch behavior, and MCP warm-up batching. Skip the MCP-specific test when optional MCP dependencies are not installed. Refs #240
This commit is contained in:
Kayzo
parent
4c9c29c421
commit
f5cea7c51e
6 changed files with 485 additions and 93 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
63
tests/test_memory/test_mcp_server.py
Normal file
63
tests/test_memory/test_mcp_server.py
Normal file
|
|
@ -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])
|
||||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue