Compare commits

..
Author SHA1 Message Date
alvinfeiandClaude Opus 5.5 302fd30af1 feat(sdk): re-embed vector stores and the graph into a new embedding space
Add hy_memory.reembed for switching the embedding model while writes are
paused. Source and target are named explicitly; a space's identity is the
model name plus the width.

- vector_space_names, plan_reembed, reembed_vector_space (resumable,
  cancellable at batch boundaries, widths validated before any upsert, target
  wins on existing ids, source never written), verify_reembed,
  drop_vector_space.
- scan_vector_spaces and diff_spaces for MemHub's startup reconcile of stray
  spaces.
- reembed_graph builds a new Kuzu DB at the target width, copies every node
  and edge, recomputes Memory/beh/Topic/HierarchyTopic embeddings and rebuilds
  the HNSW index; the source graph is opened read-only. verify_graph_reembed
  and drop_graph.

EmbedService.embed_batch now chunks to embedder.batch_size (default 64),
honours Retry-After on 429, takes a per-call max_retries, and rejects vectors
whose width differs from embedding_dims.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
2026-11-01 18:33:52 +08:00
feitianxiang f5ea0af75d Merge pull request #251 from Tencent-Hunyuan/fix/hermes-injected-content
feat(memhub): land #250 on main — multi-user Runtime, hy-memory 2.0.9
2026-11-01 14:39:43 +08:00
12 changed files with 2937 additions and 26 deletions

No files matched your search

+1
View File
@@ -40,6 +40,7 @@ MEMORY_EMBEDDING_DIMS=1024
# MEMORY_EMBEDDER_TIMEOUT=15
# MEMORY_EMBEDDER_MAX_RETRIES=2
# MEMORY_EMBEDDER_RETRY_DELAY=0.3
# MEMORY_EMBEDDER_BATCH_SIZE=64
# MEMORY_EMBEDDER_EXTRA_HEADERS={}
# MEMORY_EMBEDDER_EXTRA_BODY={}
+40
View File
@@ -4,6 +4,46 @@ Hy Memory and MemHub are versioned and released independently.
## Unreleased
- Hy Memory: new `hy_memory.reembed` module for switching the embedding
model. Writers must be paused for the whole run. Nothing is derived from
the provider or endpoint: source and target are named explicitly, and the
identity of a space is the model name plus the width (`space_identity`).
Only Chroma and Kuzu are supported.
- `vector_space_names(config, base_name=None)` returns every physical
collection a client opens for the six vector roles: memory, entities,
coding keys and working facts, each in the working and coding variants.
- `reembed_vector_space(source, target, embedding=...)` pages every raw
source collection, covering all users and all layers including L1_RAW. It
embeds the text the write path embedded (`reembed_text`: the canonical
state's `state_retrieval_text`, a schema's `retrieval_text`, otherwise the
document). It then writes each record into the target with the same id,
document and metadata. Each batch's widths are checked before anything is
written. A cancel stops at a batch boundary (`ReembedCancelled`). A rerun
skips the ids the target already holds, so it resumes, and a merge never
overwrites the target. The source is never written.
- `plan_reembed`, `verify_reembed` and `drop_vector_space` complete the
vector side. `verify_reembed` checks ids, metadata, per-user counts and
the width, and runs a self-recall sample. `drop_vector_space` refuses
collections that a live client in this process holds open.
- `scan_vector_spaces(persist_directory)` is read-only. It lists every
space on disk with its roles, counts, stored width, per-user counts and
the models recorded in metadata. `diff_spaces(source, target)` returns the
ids that the source holds and the target lacks, whatever their status.
- `reembed_graph(source_path, target_path, embedding=..., target_dims=...)`
builds a new Kuzu graph at the new width. It copies every node and edge
and recomputes `Memory.embedding`, `Memory.beh_embedding`, and the Topic
and HierarchyTopic embeddings, then rebuilds the HNSW index. The source
graph is opened read-only and kept for rollback. A run that fails or is
cancelled leaves no target behind, and a rerun starts over.
`verify_graph_reembed` and `drop_graph` complete the graph side.
- Hy Memory: `EmbedService.embed_batch` now splits large batches into
requests of `embedder.batch_size` texts each (`MEMORY_EMBEDDER_BATCH_SIZE`,
default 64); before, it sent them all in one request. On a 429 it waits for
the `Retry-After` time plus jitter, capped at 60 s. It takes a per-call
`max_retries`. A returned vector whose width differs from the configured
`embedding_dims` now raises `EmbeddingDimensionMismatchError` and is not
retried.
## MemHub 1.0.5 — 2026-10-28
Build 36, on `hy-memory==2.0.8`. Linux x64/arm64 Runtime and macOS arm64
+1
View File
@@ -48,6 +48,7 @@
| `MEMORY_EMBEDDER_TIMEOUT` | `15` | 单次超时,秒 |
| `MEMORY_EMBEDDER_MAX_RETRIES` | `2` | 最大重试次数 |
| `MEMORY_EMBEDDER_RETRY_DELAY` | `0.3` | 重试间隔,秒 |
| `MEMORY_EMBEDDER_BATCH_SIZE` | `64` | `embed_batch` 单次请求最多发送的文本条数,超出按此切块 |
| `MEMORY_EMBEDDER_EXTRA_HEADERS` | — | JSON object,附加 HTTP headers |
| `MEMORY_EMBEDDER_EXTRA_BODY` | — | JSON object,附加请求 body |
+53
View File
@@ -55,6 +55,33 @@ except Exception:
from .client import HyMemoryClient
from .inspector import MemoryInspector
from .runtime import SharedRuntime
from .reembed import (
EmbeddingSpaceIdentity,
GraphReembedPlan,
GraphReembedProgress,
GraphReembedReport,
GraphVerification,
ReembedCancelled,
ReembedPlan,
ReembedProgress,
ReembedReport,
ReembedVerification,
ScannedSpace,
SpaceDiff,
VectorSpace,
diff_spaces,
drop_graph,
drop_vector_space,
plan_reembed,
reembed_graph,
reembed_text,
reembed_vector_space,
scan_vector_spaces,
space_identity,
vector_space_names,
verify_graph_reembed,
verify_reembed,
)
from .snapshot import (
DataPlaneRestore,
DataPlaneSnapshot,
@@ -129,6 +156,32 @@ __all__ = [
"DataPlaneSnapshotRejected",
"restore_local_data_plane",
"snapshot_local_data_plane",
# 换 embedding 模型:把向量库和图重新 embed 进新的 embedding 空间
"EmbeddingSpaceIdentity",
"GraphReembedPlan",
"GraphReembedProgress",
"GraphReembedReport",
"GraphVerification",
"ReembedCancelled",
"ReembedPlan",
"ReembedProgress",
"ReembedReport",
"ReembedVerification",
"ScannedSpace",
"SpaceDiff",
"VectorSpace",
"diff_spaces",
"drop_graph",
"drop_vector_space",
"plan_reembed",
"reembed_graph",
"reembed_text",
"reembed_vector_space",
"scan_vector_spaces",
"space_identity",
"vector_space_names",
"verify_graph_reembed",
"verify_reembed",
# 配置
"MemoryConfig",
# 公共扩展协议 / 精确读取类型
+5
View File
@@ -260,6 +260,7 @@ class EmbedderConfig:
max_retries: int = None # 最大重试次数
retry_delay: float = None # 重试间隔(秒)
timeout: int = None # 请求超时(秒)
batch_size: int = None # 单次 HTTP 请求最多带多少条 input(embed_batch 按此切块)
# 内部平台扩展(extra_headers / extra_body 会透传给 OpenAI client)
extra_headers: Optional[Dict[str, str]] = None
@@ -294,6 +295,10 @@ class EmbedderConfig:
# 请求切断并制造重试风暴。SDK 内层重试已关闭,15s 是每个显式 attempt
# 的真实上限;部署可通过 MEMORY_EMBEDDER_TIMEOUT 覆盖。
self.timeout = _get_env_int("MEMORY_EMBEDDER_TIMEOUT", 15)
if self.batch_size is None:
# 大批量(如重新 embed 整个库)不能一次 HTTP 发几千条:多数服务对
# input 数组有上限,超了直接 400。按 batch_size 切块逐块发送。
self.batch_size = _get_env_int("MEMORY_EMBEDDER_BATCH_SIZE", 64)
if self.extra_headers is None:
val = os.getenv("MEMORY_EMBEDDER_EXTRA_HEADERS")
if val:
+105 -19
View File
@@ -38,10 +38,12 @@ import importlib.util
import json
from json import JSONDecodeError
import logging
import random
import threading
import time
from ..config import MemoryConfig
from ..embedding import EmbeddingDimensionMismatchError
from .workload import embedding_slot
from ..utils.log_setup import get_request_id
@@ -187,6 +189,41 @@ def _redact_headers(headers: Optional[Dict[str, Any]]) -> Optional[Dict[str, Any
return out
# Retry-After 的上限:服务端给出超长等待(或恶意值)时不让单次调用挂几十分钟。
_RETRY_AFTER_MAX_S = 60.0
def _retry_after_seconds(exc: Exception) -> Optional[float]:
"""从 429 异常的响应头里取 Retry-After(秒数或 HTTP 日期),取不到返回 None。"""
response = getattr(exc, "response", None)
headers = getattr(response, "headers", None)
if not headers:
return None
try:
raw = headers.get("retry-after") or headers.get("Retry-After")
except Exception:
return None
if raw is None:
return None
raw = str(raw).strip()
try:
seconds = float(raw)
except ValueError:
try:
from email.utils import parsedate_to_datetime
from datetime import datetime, timezone
when = parsedate_to_datetime(raw)
if when.tzinfo is None:
when = when.replace(tzinfo=timezone.utc)
seconds = (when - datetime.now(timezone.utc)).total_seconds()
except Exception:
return None
if seconds != seconds: # NaN
return None
return min(max(0.0, seconds), _RETRY_AFTER_MAX_S)
def _embeddings_api_url(base_url: str) -> str:
"""OpenAI SDK posts to ``{base_url}/embeddings`` (base_url usually ends with /v1)."""
base = (base_url or "").rstrip("/")
@@ -205,7 +242,8 @@ class EmbedServiceConfig:
api_key: str = ""
base_url: str = ""
embedding_dims: int = 1536
batch_size: int = 100
# embed_batch 单次 HTTP 最多带多少条 input;超出按块顺序发送
batch_size: int = 64
enable_cache: bool = True
cache_size: int = 10000
max_retries: int = 5
@@ -287,6 +325,7 @@ class EmbedService:
max_retries=embedder.max_retries if embedder.max_retries is not None else 5,
retry_delay=embedder.retry_delay if embedder.retry_delay is not None else 3.0,
timeout=embedder.timeout if embedder.timeout is not None else 60,
batch_size=max(1, int(getattr(embedder, "batch_size", None) or 64)),
queue_batch_size=int(os.getenv("MEMORY_EMBED_BATCH_SIZE", "8")),
queue_batch_window_ms=float(os.getenv("MEMORY_EMBED_BATCH_WINDOW_MS", "100")),
extra_headers=embedder.extra_headers,
@@ -640,6 +679,8 @@ class EmbedService:
self,
texts: List[str],
use_cache: bool = True,
*,
max_retries: Optional[int] = None,
) -> List[List[float]]:
"""
批量向量化
@@ -647,17 +688,23 @@ class EmbedService:
Args:
texts: 文本列表
use_cache: 是否使用缓存
max_retries: 本次调用的重试次数;None 用配置值。长批量任务(重新
embed 整个库)可以给更宽的重试预算,不影响在线路径。
Returns:
向量列表
"""
vectors, _metas = await self.embed_batch_with_meta(texts, use_cache=use_cache)
vectors, _metas = await self.embed_batch_with_meta(
texts, use_cache=use_cache, max_retries=max_retries,
)
return vectors
async def embed_batch_with_meta(
self,
texts: List[str],
use_cache: bool = True,
*,
max_retries: Optional[int] = None,
) -> Tuple[List[List[float]], List[Dict[str, Any]]]:
"""批量向量化,同时返回每条的 provider metadata(非 vector)。"""
if not texts:
@@ -682,7 +729,12 @@ class EmbedService:
indices_to_embed.append(i)
if texts_to_embed:
vectors, batch_metas = await self._embed_texts_with_meta(texts_to_embed)
# max_retries 只在显式给出时透传:保持 `_embed_texts_with_meta(texts)`
# 这个既有调用形状不变(测试与下游替身按它打桩)。
retry_kwargs = {} if max_retries is None else {"max_retries": max_retries}
vectors, batch_metas = await self._embed_texts_with_meta(
texts_to_embed, **retry_kwargs,
)
for j, vector in enumerate(vectors):
idx = indices_to_embed[j]
@@ -700,9 +752,9 @@ class EmbedService:
return results, metas
async def _embed_texts_with_meta(
self, texts: List[str]
self, texts: List[str], *, max_retries: Optional[int] = None,
) -> Tuple[List[List[float]], List[Dict[str, Any]]]:
"""实际的向量化调用:**一次 HTTP 带走整批**;返回 (vectors, metas)。
"""实际的向量化调用:**一次 HTTP 带走一块**;返回 (vectors, metas)。
早期实现是 `for text in texts: await _embed_openai([text])` 逐条串行,
原意是"避免 batch size 限制",但代价极大:线上一批 8 条 × ~5.6s 串行
@@ -719,21 +771,34 @@ class EmbedService:
截断在这里做(发送前逐条过 `_truncate_for_embedding`):所有 embedding
路径(embed / embed_batch / embed_queued / *_with_meta)最终都落到这个
方法 —— 放在各入口会漏。理由与阈值见模块头部 `_EMBED_MAX_TOKENS`。
切块:一次带走整批只适用于小批。几千条的 input 数组会被多数服务直接
400(input 上限),所以按 `batch_size`(默认 64)切块,块之间顺序发送。
不超过一块时行为与切块前逐字相同(`batch_mode="single_request"`)。
"""
if not texts:
return [], []
safe_texts = [_truncate_for_embedding(t) for t in texts]
vectors, meta = await self._embed_openai(safe_texts)
base_meta = dict(meta or {})
base_meta["batch_mode"] = "single_request"
base_meta["batch_input_count"] = len(safe_texts)
chunk = max(1, int(self._embed_config.batch_size or 64))
vectors: List[List[float]] = []
all_metas: List[Dict[str, Any]] = []
for idx in range(len(safe_texts)):
m = dict(base_meta)
m["batch_index"] = idx
all_metas.append(m)
chunk_count = (len(safe_texts) + chunk - 1) // chunk
for start in range(0, len(safe_texts), chunk):
part = safe_texts[start:start + chunk]
retry_kwargs = {} if max_retries is None else {"max_retries": max_retries}
part_vectors, meta = await self._embed_openai(part, **retry_kwargs)
base_meta = dict(meta or {})
base_meta["batch_mode"] = "single_request" if chunk_count == 1 else "chunked"
base_meta["batch_input_count"] = len(part)
if chunk_count > 1:
base_meta["batch_chunk_index"] = start // chunk
base_meta["batch_chunk_count"] = chunk_count
for idx in range(len(part)):
m = dict(base_meta)
m["batch_index"] = start + idx
all_metas.append(m)
vectors.extend(part_vectors)
self._last_embed_meta = all_metas[-1] if all_metas else {}
return vectors, all_metas
@@ -783,7 +848,7 @@ class EmbedService:
)
async def _embed_openai(
self, texts: List[str]
self, texts: List[str], *, max_retries: Optional[int] = None,
) -> Tuple[List[List[float]], Dict[str, Any]]:
"""
使用 OpenAI 兼容接口向量化
@@ -873,7 +938,9 @@ class EmbedService:
"endpoint_overridden": bool(_ov),
}
_max_retries = max(0, int(self._embed_config.max_retries or 0))
_max_retries = max(0, int(
(self._embed_config.max_retries or 0) if max_retries is None else max_retries
))
_retry_delay = max(0.0, float(self._embed_config.retry_delay or 0.0))
last_exc: Optional[Exception] = None
@@ -919,6 +986,17 @@ class EmbedService:
f"got={len(vectors)}"
)
dims = [len(v) for v in vectors]
expected_dims = self._embed_config.embedding_dims
if expected_dims:
bad = [d for d in dims if d != int(expected_dims)]
if bad:
# 宽度不对的向量写进集合会把集合的维度钉死在错误值上(Chroma
# 首次 add 定维),或在已有集合上整批失败。确定性错误,不重试。
raise EmbeddingDimensionMismatchError(
"embedding dimension mismatch: "
f"expected={int(expected_dims)} got={bad[0]} "
f"model={self._embed_config.model}"
)
meta = {
"http_status": raw.status_code,
"response_headers": resp_headers,
@@ -954,6 +1032,8 @@ class EmbedService:
},
)
return vectors, meta
except EmbeddingDimensionMismatchError:
raise
except Exception as e:
last_exc = e
_err_name = type(e).__name__
@@ -987,12 +1067,18 @@ class EmbedService:
if _retryable and attempt < _max_retries:
# 逐出可能卡死的 client,下个 attempt 重建干净连接(恢复原防御语义)
self._evict_client(_eff_base_url, _eff_api_key, _eff_timeout)
_delay = _retry_delay
_retry_after = _retry_after_seconds(e) if _status == 429 else None
if _retry_after is not None:
# 限流时服务端明确告诉了要等多久:按它等,再加一点抖动,
# 避免多个 worker 在同一时刻一起醒来再次撞限流。
_delay = max(_delay, _retry_after) + random.uniform(0, 0.25 * _retry_after)
logger.warning(
f"[embed] attempt {attempt + 1}/{_max_retries + 1} failed "
f"({_err_name}): {_err_msg[:200]}; retry in {_retry_delay}s"
f"({_err_name}): {_err_msg[:200]}; retry in {_delay:.2f}s"
)
if _retry_delay > 0:
await asyncio.sleep(_retry_delay)
if _delay > 0:
await asyncio.sleep(_delay)
continue
self._log_embed_call(
status="error",
+18
View File
@@ -22,6 +22,7 @@ import logging
import json
import threading
import time
import weakref
from ..models.memory import MemoryNode, MemoryLayer, MemoryStatus
from ..config import MemoryConfig
@@ -77,6 +78,21 @@ async def _run_in_kuzu_pool(func, *args, **kwargs):
return await run_in_executor_guarded(_kuzu_executor, _measured)
# 本进程里已打开、尚未 close 的 store。drop_graph 据此拒绝删掉一个正被
# 活着的 client 使用的图库。WeakSet:store 被回收后自动消失。
_LIVE_STORES: "weakref.WeakSet[KuzuGraphStore]" = weakref.WeakSet()
def open_graph_paths() -> List[str]:
"""Real paths of the graph databases live stores in this process hold open."""
import os
return [
os.path.realpath(str(store._db_path))
for store in list(_LIVE_STORES)
if store._db is not None
]
def kuzu_executor_snapshot() -> Dict[str, Any]:
with _kuzu_stats_lock:
return {
@@ -350,6 +366,7 @@ class KuzuGraphStore(GraphStoreBase):
)
self._db = db
self._conn = conn
_LIVE_STORES.add(self)
except asyncio.TimeoutError:
logger.warning(
f"GraphStore init timed out after 5s (Kuzu file lock contention). "
@@ -2441,6 +2458,7 @@ class KuzuGraphStore(GraphStoreBase):
async def close(self) -> None:
"""关闭连接"""
conn, db = self._conn, self._db
_LIVE_STORES.discard(self)
self._conn = None
self._db = None
self._available = False
+14 -6
View File
@@ -89,6 +89,17 @@ def passes_created_after(
logger = logging.getLogger(__name__)
def physical_collection_name(base_name: Optional[str], dims: Optional[int]) -> str:
"""The collection name a store opens for ``base_name`` at ``dims``.
The dims suffix keeps two embedding widths from sharing one collection.
`hy_memory.reembed.vector_space_names` uses this too, so callers never
re-derive the rule.
"""
base = base_name or "agent_memories"
return f"{base}_{dims}" if dims else base
class VectorStoreBase(ABC):
"""
向量存储抽象基类
@@ -98,13 +109,10 @@ class VectorStoreBase(ABC):
def __init__(self, config: MemoryConfig):
self.config = config
base_name = config.vector_store.collection_name or "agent_memories"
dims = config.vector_store.embedding_dims
# 自动拼接维度后缀,避免不同 embedding 模型写入同一 collection 导致维度冲突
if dims:
self._collection_name = f"{base_name}_{dims}"
else:
self._collection_name = base_name
self._collection_name = physical_collection_name(
config.vector_store.collection_name, config.vector_store.embedding_dims,
)
_supports_retrieval_clauses: bool = False
_supports_narrow_episode_queries: bool = False
+29 -1
View File
@@ -9,10 +9,12 @@ Agent Memory V2 - VectorStore (Chroma)
- 同步 API → _run_in_vdb_pool() 包装为异步
"""
from typing import Optional, List, Dict, Any
from typing import Optional, List, Dict, Any, Set, Tuple
import concurrent.futures
import functools
import logging
import os
import weakref
from ..models.memory import (
MemoryNode, MemoryLayer, MemoryStatus,
@@ -33,6 +35,30 @@ _vdb_executor = concurrent.futures.ThreadPoolExecutor(
)
# 本进程里已 initialize、尚未 close 的 store。drop_vector_space 据此拒绝删掉
# 一个正被活着的 client 使用的集合。WeakSet:store 被回收后自动消失。
_LIVE_STORES: "weakref.WeakSet[ChromaVectorStore]" = weakref.WeakSet()
def _location_key(vs_config) -> Tuple[str, str]:
if getattr(vs_config, "host", None):
return ("http", f"{vs_config.host}:{vs_config.port or 8000}")
return ("local", os.path.realpath(str(vs_config.persist_directory or "")))
def open_collection_names(location: Tuple[str, str]) -> Set[str]:
"""Collections that live stores in this process hold open at ``location``."""
names: Set[str] = set()
for store in list(_LIVE_STORES):
if store._client is None or _location_key(store.config.vector_store) != location:
continue
names.add(store._collection_name)
if store._entity_collection is not None:
from ..pipelines._retrieval.config import entity_collection_name
names.add(entity_collection_name(store._collection_name))
return names
async def _run_in_vdb_pool(func, *args, **kwargs):
"""Run a synchronous Chroma call in the dedicated VDB executor."""
call = functools.partial(func, *args, **kwargs) if args or kwargs else func
@@ -88,6 +114,7 @@ class ChromaVectorStore(VectorStoreBase):
return client, collection
self._client, self._collection = await _run_in_vdb_pool(_init)
_LIVE_STORES.add(self)
logger.debug(f"VectorStore initialized (Chroma), collection={self._collection_name}")
# ================================================================
@@ -561,6 +588,7 @@ class ChromaVectorStore(VectorStoreBase):
async def close(self) -> None:
"""关闭客户端"""
# Chroma PersistentClient 无需显式关闭
_LIVE_STORES.discard(self)
self._client = None
self._collection = None
self._entity_collection = None
+1777
View File
File diff suppressed because it is too large. Load diff
+642
View File
@@ -0,0 +1,642 @@
"""Re-embed into a new embedding space: vector stores, scan, diff and merge.
The switch model: writes are paused, every vector is computed again with the
new model into a new space, the source is never touched. The graph half lives
in ``test_reembed_graph.py``.
"""
import hashlib
import json
import math
import random
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from hy_memory import (
EmbeddingDimensionMismatchError,
MemoryConfig,
ReembedCancelled,
VectorSpace,
diff_spaces,
drop_vector_space,
plan_reembed,
reembed_text,
reembed_vector_space,
scan_vector_spaces,
space_identity,
vector_space_names,
verify_reembed,
)
from hy_memory.config import EmbedderConfig
from hy_memory.core.embed_service import EmbedService
from hy_memory.models.memory import MemoryLayer, MemoryNode, MemoryStatus
from hy_memory.reembed import VECTOR_ROLES, parse_collection_name
class FakeEmbedder:
"""Deterministic vectors: same (model, text) -> same unit vector."""
def __init__(self, model: str, dims: int, *, wrong_width_after: int = -1):
self.model = model
self.dims = dims
self.calls = []
self.texts = 0
self.wrong_width_after = wrong_width_after
def vector(self, text: str):
seed = int(hashlib.sha256(f"{self.model}\0{text}".encode()).hexdigest()[:16], 16)
rng = random.Random(seed)
values = [rng.gauss(0.0, 1.0) for _ in range(self.dims)]
norm = math.sqrt(sum(v * v for v in values)) or 1.0
return [v / norm for v in values]
async def embed_batch(self, texts, use_cache=True):
self.calls.append(list(texts))
out = []
for text in texts:
self.texts += 1
vector = self.vector(text)
if 0 <= self.wrong_width_after < self.texts:
vector = vector + [0.0]
out.append(vector)
return out
async def embed(self, text, use_cache=True):
return self.vector(text)
async def embed_queued(self, text, use_cache=True):
return self.vector(text)
def get_embedding_dims(self):
return self.dims
def _client(path):
import chromadb
return chromadb.PersistentClient(path=str(path))
def _space(path, base, dims, model=None):
return VectorSpace(base_name=base, dims=dims, persist_directory=str(path), model=model)
def _custom_state():
return {
"memory_kind": "canonical_state",
"state_key": "home_city",
"state_value": "Berlin",
"state_retrieval_text": "user lives in Berlin (home city)",
}
async def _seed(path, base, dims, *, model="old-model", users=("alice", "bob"), per_user=7):
"""Write a realistic source space through the SDK's own stores."""
from hy_memory.data.vector_store_chroma import ChromaVectorStore
old = FakeEmbedder(model, dims)
cfg = MemoryConfig()
cfg.vector_store.provider = "chroma"
cfg.vector_store.persist_directory = str(path)
cfg.vector_store.collection_name = base
cfg.vector_store.embedding_dims = dims
store = ChromaVectorStore(cfg)
await store.initialize()
nodes = []
for user in users:
for i in range(per_user):
layer = [MemoryLayer.L1_RAW, MemoryLayer.L2_FACT, MemoryLayer.L3_SUMMARY][i % 3]
node = MemoryNode(
user_id=user, agent_id="a", session_id=f"s{i % 2}",
layer=layer, content=f"{user} fact {i}",
status=MemoryStatus.SUPERSEDED if i == 4 else MemoryStatus.ACTIVE,
tags=["t1", f"t{i}"],
custom={"k": i, "nested": {"x": [1, 2]}},
)
if i == 4:
node.is_latest = False
node.superseded_by = ["later"]
if i == 5:
node.custom = _custom_state()
node.embedding = old.vector(node.content)
nodes.append(node)
await store.upsert_batch(nodes)
entities = await store._get_entity_collection()
entities.add(
ids=["e1", "e2"],
embeddings=[old.vector("Berlin"), old.vector("Dylan")],
documents=["Berlin", "Dylan"],
metadatas=[{"user_id": "alice", "label": "GPE"}, {"user_id": "bob", "label": "PERSON"}],
)
await store.close()
client = _client(path)
names = vector_space_names(_space(path, base, dims))
for role, doc in (("coding_keys", "fix the flaky test"), ("working_facts_cv", "repo: uses uv")):
coll = client.get_or_create_collection(names[role], metadata={"hnsw:space": "cosine"})
coll.add(
ids=[f"{role}:0"],
embeddings=[old.vector(doc)],
documents=[doc],
metadatas=[{"user_id": "alice", "layer": "l1_raw",
"custom": json.dumps({"coding_key_text": doc})}],
)
return names
def _snapshot(path, names):
client = _client(path)
out = {}
for role, name in names.items():
try:
coll = client.get_collection(name)
except Exception:
continue
got = coll.get(include=["documents", "metadatas", "embeddings"])
out[role] = sorted(
(rid, doc, json.dumps(meta, sort_keys=True), [round(float(x), 6) for x in emb])
for rid, doc, meta, emb in zip(got["ids"], got["documents"], got["metadatas"], got["embeddings"])
)
return hashlib.sha256(json.dumps(out, sort_keys=True).encode()).hexdigest()
def _records(path, name):
got = _client(path).get_collection(name).get(include=["documents", "metadatas", "embeddings"])
return {
rid: (doc, meta, list(emb))
for rid, doc, meta, emb in zip(got["ids"], got["documents"], got["metadatas"], got["embeddings"])
}
# ----------------------------------------------------------------------
# naming
# ----------------------------------------------------------------------
class TestVectorSpaceNames:
@pytest.mark.parametrize("dims", [512, None])
async def test_names_equal_what_the_stores_open(self, tmp_path, dims):
from hy_memory.coding.store import CodingMemoryStore
from hy_memory.data.vector_store_chroma import ChromaVectorStore
cfg = MemoryConfig()
cfg.vector_store.provider = "chroma"
cfg.vector_store.persist_directory = str(tmp_path / "v")
cfg.vector_store.collection_name = "memhub_abc123"
cfg.vector_store.embedding_dims = dims
cfg.history.db_path = str(tmp_path / "history.sqlite3")
store = ChromaVectorStore(cfg)
await store.initialize()
await store._get_entity_collection()
coding = []
for variant in ("working", "coding"):
cs = CodingMemoryStore(cfg, FakeEmbedder("m", dims or 8), variant=variant)
await cs.initialize()
coding.append(cs)
opened = {c.name for c in _client(tmp_path / "v").list_collections()}
names = vector_space_names(cfg)
assert set(names) == set(VECTOR_ROLES)
assert set(names.values()) == opened
suffix = f"_{dims}" if dims else ""
assert names["memory"] == f"memhub_abc123{suffix}"
assert names["coding_keys_cv"] == f"memhub_abc123_coding_keys_cv{suffix}"
await store.close()
for cs in coding:
await cs.close()
def test_explicit_base_name_and_space(self, tmp_path):
cfg = MemoryConfig()
cfg.vector_store.provider = "chroma"
cfg.vector_store.embedding_dims = 1024
names = vector_space_names(cfg, base_name="staging_x")
assert names["memory"] == "staging_x_1024"
assert names["entities"] == "staging_x_1024_entities"
assert names["working_facts"] == "staging_x_working_facts_1024"
assert vector_space_names(_space(tmp_path, "b", 512))["memory"] == "b_512"
def test_parse_collection_name_round_trips(self):
for dims in (512, None):
for role, name in vector_space_names(VectorSpace("memhub_9f", dims)).items():
assert parse_collection_name(name) == ("memhub_9f", dims, role)
assert parse_collection_name("memhub_9f_512_tag_index") is None
def test_space_identity_is_model_and_width_only(self):
a = space_identity("BAAI/bge-small-zh-v1.5", 512)
b = space_identity(" BAAI/bge-small-zh-v1.5 ", 512)
assert a == b and a.key == "BAAI/bge-small-zh-v1.5#512"
assert a.digest != space_identity("BAAI/bge-small-zh-v1.5", 1024).digest
with pytest.raises(ValueError):
space_identity("", 512)
def test_non_chroma_provider_is_not_supported(self):
cfg = MemoryConfig()
cfg.vector_store.provider = "qdrant"
with pytest.raises(NotImplementedError):
vector_space_names(cfg)
class TestReembedText:
def test_canonical_state_uses_state_retrieval_text(self):
meta = {"content": "住在柏林", "custom": json.dumps(_custom_state())}
assert reembed_text("memory", "住在柏林", meta) == "user lives in Berlin (home city)"
def test_retrieval_text_when_present(self):
meta = {"custom": json.dumps({"retrieval_text": " pattern: plans ahead "})}
assert reembed_text("memory", "画像", meta) == "pattern: plans ahead"
def test_document_otherwise_with_content_fallback(self):
assert reembed_text("memory", "plain fact", {"custom": "{}"}) == "plain fact"
assert reembed_text("memory", None, {"content": "from metadata"}) == "from metadata"
assert reembed_text("memory", "doc", {"custom": "not json"}) == "doc"
def test_coding_key_and_working_fact_use_the_document(self):
coding = {"custom": json.dumps({"coding_key_text": "fix flaky test",
"retrieval_text": "ignored outside memory"})}
assert reembed_text("coding_keys", "fix flaky test", coding) == "fix flaky test"
assert reembed_text("working_facts", "repo: uses uv", {}) == "repo: uses uv"
# ----------------------------------------------------------------------
# re-embed
# ----------------------------------------------------------------------
@pytest.mark.parametrize(
"src_dims,tgt_dims,src_model,tgt_model",
[(512, 512, "old-model", "new-model"), (512, 1024, "old", "new"), (1024, 512, "old", "new")],
)
async def test_reembed_copies_everything_and_never_touches_source(
tmp_path, src_dims, tgt_dims, src_model, tgt_model,
):
path = tmp_path / "v"
names = await _seed(path, "memhub_src", src_dims, model=src_model)
before = _snapshot(path, names)
source = _space(path, "memhub_src", src_dims)
target = _space(path, "memhub_dst", tgt_dims, model=tgt_model)
new = FakeEmbedder(tgt_model, tgt_dims)
plan = await plan_reembed(source, target)
assert plan.per_role["memory"] == 14 and plan.per_role["entities"] == 2
assert plan.total_vectors == 18 == plan.estimated_vectors
assert plan.source_dims == src_dims and plan.target_dims == tgt_dims
seen = []
report = await reembed_vector_space(
source, target, embedding=new, batch_size=4, on_progress=seen.append,
)
assert report.embedded == 18 and report.reused == 0
assert seen and seen[-1].done == seen[-1].total
assert _snapshot(path, names) == before # source unchanged, byte for byte
target_names = vector_space_names(target)
for role in ("memory", "entities", "coding_keys", "working_facts_cv"):
src = _records(path, names[role])
dst = _records(path, target_names[role])
assert set(src) == set(dst)
for rid, (doc, meta, _) in src.items():
assert dst[rid][0] == doc
assert dst[rid][1] == meta # status, is_latest, superseded_by, timestamps, custom
assert len(dst[rid][2]) == tgt_dims
assert dst[rid][2] == pytest.approx(new.vector(reembed_text(role, doc, meta)), abs=1e-5)
memory = _records(path, target_names["memory"])
layers = {meta["layer"] for _, meta, _ in memory.values()}
users = {meta["user_id"] for _, meta, _ in memory.values()}
assert "l1_raw" in layers and users == {"alice", "bob"}
statuses = {meta["status"] for _, meta, _ in memory.values()}
assert statuses == {"active", "superseded"}
state = [r for r in memory.values() if "canonical_state" in r[1]["custom"]]
assert state and state[0][2] == pytest.approx(
new.vector("user lives in Berlin (home city)"), abs=1e-5,
)
meta = _client(path).get_collection(target_names["memory"]).metadata
assert meta["hnsw:space"] == "cosine" and meta["hy_embedding_model"] == tgt_model
verified = await verify_reembed(source, target, embedding=new)
assert verified.ok, verified.problems
assert verified.per_user == {"alice": 10, "bob": 8}
assert verified.self_recall_samples == 18 and verified.self_recall_hits == 18
async def test_resume_after_kill_only_embeds_the_rest(tmp_path):
path = tmp_path / "v"
await _seed(path, "src", 512, per_user=20)
source, target = _space(path, "src", 512), _space(path, "dst", 1024)
class Killed(Exception):
pass
first = FakeEmbedder("new", 1024)
batches = {"n": 0}
def _kill():
batches["n"] += 1
if batches["n"] > 3:
raise Killed()
return False
with pytest.raises(Killed):
await reembed_vector_space(
source, target, embedding=first, batch_size=5, max_inflight=1,
roles=["memory"], cancel=_kill,
)
assert first.texts == 15
second = FakeEmbedder("new", 1024)
report = await reembed_vector_space(
source, target, embedding=second, batch_size=5, max_inflight=1, roles=["memory"],
)
assert second.texts == 40 - 15
assert report.per_role["memory"].reused == 15
# The end state equals an uninterrupted run: every source record, verbatim,
# with the vector of its own text.
resumed = _records(path, "dst_1024")
src = _records(path, "src_512")
assert set(resumed) == set(src)
for rid, (doc, meta, vector) in resumed.items():
assert (doc, meta) == src[rid][:2]
assert vector == pytest.approx(second.vector(reembed_text("memory", doc, meta)), abs=1e-5)
assert (await verify_reembed(source, target, roles=["memory"], samples=0)).ok
async def test_cancel_at_batch_boundary_leaves_no_partial_batch(tmp_path):
path = tmp_path / "v"
await _seed(path, "src", 512)
calls = {"n": 0}
def _cancel():
calls["n"] += 1
return calls["n"] > 2
with pytest.raises(ReembedCancelled) as info:
await reembed_vector_space(
_space(path, "src", 512), _space(path, "dst", 512),
embedding=FakeEmbedder("new", 512), batch_size=3, max_inflight=1,
roles=["memory"], cancel=_cancel,
)
assert info.value.report.per_role["memory"].embedded == 6
assert _client(path).get_collection("dst_512").count() == 6
async def test_wrong_width_raises_before_any_upsert(tmp_path):
path = tmp_path / "v"
await _seed(path, "src", 512)
bad = FakeEmbedder("new", 1024, wrong_width_after=2)
with pytest.raises(EmbeddingDimensionMismatchError):
await reembed_vector_space(
_space(path, "src", 512), _space(path, "dst", 1024),
embedding=bad, batch_size=4, roles=["memory"],
)
assert _client(path).get_collection("dst_1024").count() == 0
declared_wrong = FakeEmbedder("new", 768)
with pytest.raises(EmbeddingDimensionMismatchError):
await reembed_vector_space(
_space(path, "src", 512), _space(path, "dst", 1024), embedding=declared_wrong,
)
async def test_refuses_to_write_into_the_source_collection(tmp_path):
path = tmp_path / "v"
await _seed(path, "src", 512)
with pytest.raises(ValueError):
await reembed_vector_space(
_space(path, "src", 512), _space(path, "src", 512),
embedding=FakeEmbedder("new", 512),
)
async def test_verify_detects_missing_extra_drift_and_width(tmp_path):
path = tmp_path / "v"
await _seed(path, "src", 512)
source, target = _space(path, "src", 512), _space(path, "dst", 512)
new = FakeEmbedder("new", 512)
await reembed_vector_space(source, target, embedding=new, roles=["memory"])
coll = _client(path).get_collection("dst_512")
got = coll.get(include=["documents", "metadatas"])
victim, drifted = got["ids"][0], got["ids"][1]
coll.delete(ids=[victim])
meta = dict(got["metadatas"][1])
meta["status"] = "archived"
coll.update(ids=[drifted], metadatas=[meta])
coll.add(ids=["stranger"], embeddings=[new.vector("x")], documents=["x"], metadatas=[{"user_id": "eve"}])
result = await verify_reembed(source, target, embedding=new, roles=["memory"], samples=5)
assert not result.ok
counts = result.per_role["memory"]
assert (counts["missing"], counts["extra"], counts["mismatched"]) == (1, 1, 1)
assert any("eve" in p for p in result.problems)
wide = _space(path, "dst", 1024)
wide_coll = _client(path).get_or_create_collection("dst_1024", metadata={"hnsw:space": "cosine"})
src = _records(path, "src_512")
wide_coll.add(
ids=list(src),
embeddings=[[0.0] * 512 for _ in src],
documents=[v[0] for v in src.values()],
metadatas=[v[1] for v in src.values()],
)
result = await verify_reembed(source, wide, roles=["memory"], samples=0)
assert result.per_role["memory"]["bad_dims"] == len(src) and not result.ok
async def test_drop_vector_space_refuses_live_collections(tmp_path):
from hy_memory.data.vector_store_chroma import ChromaVectorStore
path = tmp_path / "v"
await _seed(path, "src", 512)
cfg = MemoryConfig()
cfg.vector_store.provider = "chroma"
cfg.vector_store.persist_directory = str(path)
cfg.vector_store.collection_name = "src"
cfg.vector_store.embedding_dims = 512
live = ChromaVectorStore(cfg)
await live.initialize()
with pytest.raises(RuntimeError):
await drop_vector_space(_space(path, "src", 512))
await live.close()
dropped = await drop_vector_space(_space(path, "src", 512))
assert set(dropped) == {"src_512", "src_512_entities", "src_coding_keys_512", "src_working_facts_cv_512"}
assert not {c.name for c in _client(path).list_collections()}
# ----------------------------------------------------------------------
# scan, diff, merge
# ----------------------------------------------------------------------
async def _two_spaces(path):
"""A live 1024 space and a stray 512 space sharing some ids."""
from hy_memory.data.vector_store_chroma import ChromaVectorStore
async def _store(base, dims, model, contents):
cfg = MemoryConfig()
cfg.vector_store.provider = "chroma"
cfg.vector_store.persist_directory = str(path)
cfg.vector_store.collection_name = base
cfg.vector_store.embedding_dims = dims
store = ChromaVectorStore(cfg)
await store.initialize()
emb = FakeEmbedder(model, dims)
nodes = []
for node_id, user, content, status in contents:
node = MemoryNode(
node_id=node_id, user_id=user, agent_id="a", session_id="s",
layer=MemoryLayer.L2_FACT, content=content, status=status,
)
node.embedding = emb.vector(content)
nodes.append(node)
await store.upsert_batch(nodes)
await store.close()
ids = [f"00000000-0000-0000-0000-00000000000{i}" for i in range(6)]
await _store("memhub_old", 512, "bge", [
(ids[0], "alice", "shared one", MemoryStatus.ACTIVE),
(ids[1], "alice", "shared two (stray wording)", MemoryStatus.ACTIVE),
(ids[2], "bob", "only in stray", MemoryStatus.ACTIVE),
(ids[3], "bob", "only in stray, deleted", MemoryStatus.SHADOW),
])
await _store("memhub_new", 1024, "e5", [
(ids[0], "alice", "shared one", MemoryStatus.ACTIVE),
(ids[1], "alice", "shared two (live wording)", MemoryStatus.ARCHIVED),
(ids[4], "carol", "only in live", MemoryStatus.ACTIVE),
])
return ids
async def test_scan_diff_and_merge(tmp_path):
path = tmp_path / "v"
ids = await _two_spaces(path)
spaces = await scan_vector_spaces(path)
by_base = {(s.base_name, s.dims): s for s in spaces}
assert set(by_base) == {("memhub_old", 512), ("memhub_new", 1024)}
old, live = by_base[("memhub_old", 512)], by_base[("memhub_new", 1024)]
assert old.roles["memory"].count == 4 and old.width == 512
assert live.roles["memory"].count == 3 and live.width == 1024
assert old.per_user == {"alice": 2, "bob": 2}
assert live.per_user == {"alice": 2, "carol": 1}
stray, target = old.space(), _space(path, "memhub_new", 1024, model="e5")
diff = await diff_spaces(stray, target)
assert set(diff.missing["memory"]) == {ids[2], ids[3]} # status ignored
assert diff.total_missing == 2
live_before = _records(path, "memhub_new_1024")
stray_before = _snapshot(path, stray.names())
merger = FakeEmbedder("e5", 1024)
report = await reembed_vector_space(stray, target, embedding=merger)
assert merger.texts == 2 and report.embedded == 2
assert report.reused == 1 and report.kept == 1
after = _records(path, "memhub_new_1024")
assert set(after) == set(live_before) | {ids[2], ids[3]}
for rid, record in live_before.items(): # the target wins
assert after[rid][:2] == record[:2]
assert after[rid][2] == pytest.approx(record[2], abs=1e-6)
assert _snapshot(path, stray.names()) == stray_before
assert (await diff_spaces(stray, target)).empty
again = FakeEmbedder("e5", 1024)
second = await reembed_vector_space(stray, target, embedding=again)
assert again.texts == 0 and second.embedded == 0
assert _records(path, "memhub_new_1024").keys() == after.keys()
rescanned = {(s.base_name, s.dims): s for s in await scan_vector_spaces(path)}
assert rescanned[("memhub_new", 1024)].roles["memory"].count == 5
async def test_scan_of_missing_directory_is_empty(tmp_path):
assert await scan_vector_spaces(tmp_path / "nope") == []
# ----------------------------------------------------------------------
# EmbedService fixes
# ----------------------------------------------------------------------
def _fake_raw(n, dims):
items = []
for i in range(n):
item = MagicMock()
item.embedding = [0.1] * dims
item.index = i
items.append(item)
resp = MagicMock()
resp.model = "m"
resp.usage = None
resp.data = items
raw = MagicMock()
raw.status_code = 200
raw.headers = {}
raw.parse.return_value = resp
return raw
def _service(dims=8, **embedder):
cfg = MemoryConfig(embedder=EmbedderConfig(
provider="openai", model="m", api_key="sk", base_url="https://e.invalid/v1",
embedding_dims=dims, max_retries=2, retry_delay=0.0, **embedder,
))
return EmbedService(cfg)
async def test_embed_batch_chunks_to_batch_size():
svc = _service()
sizes = []
async def _create(**kwargs):
sizes.append(len(kwargs["input"]))
return _fake_raw(len(kwargs["input"]), 8)
client = MagicMock()
client.embeddings.with_raw_response.create = AsyncMock(side_effect=_create)
with patch("openai.AsyncOpenAI", return_value=client):
vectors = await svc.embed_batch([f"t{i}" for i in range(1000)], use_cache=False)
assert len(vectors) == 1000
assert len(sizes) == math.ceil(1000 / 64) and max(sizes) == 64
small = _service(batch_size=10)
sizes.clear()
with patch("openai.AsyncOpenAI", return_value=client):
_, metas = await small.embed_batch_with_meta([f"u{i}" for i in range(25)], use_cache=False)
assert sizes == [10, 10, 5]
assert [m["batch_index"] for m in metas] == list(range(25))
async def test_429_honours_retry_after():
svc = _service()
class RateLimited(Exception):
status_code = 429
def __init__(self):
super().__init__("rate limited")
self.response = MagicMock()
self.response.headers = {"retry-after": "2"}
client = MagicMock()
client.embeddings.with_raw_response.create = AsyncMock(
side_effect=[RateLimited(), _fake_raw(1, 8)],
)
sleeps = []
async def _sleep(seconds):
sleeps.append(seconds)
with patch("openai.AsyncOpenAI", return_value=client), \
patch("hy_memory.core.embed_service.asyncio.sleep", _sleep):
vectors = await svc.embed_batch(["x"], use_cache=False)
assert len(vectors) == 1
assert len(sleeps) == 1 and 2.0 <= sleeps[0] <= 2.5
async def test_per_call_max_retries():
svc = _service()
class Down(Exception):
status_code = 503
client = MagicMock()
client.embeddings.with_raw_response.create = AsyncMock(side_effect=Down("down"))
with patch("openai.AsyncOpenAI", return_value=client), pytest.raises(Down):
await svc.embed_batch(["x"], use_cache=False, max_retries=5)
assert client.embeddings.with_raw_response.create.await_count == 6
async def test_returned_width_is_validated_and_not_retried():
svc = _service(dims=8)
client = MagicMock()
client.embeddings.with_raw_response.create = AsyncMock(return_value=_fake_raw(2, 7))
with patch("openai.AsyncOpenAI", return_value=client), \
pytest.raises(EmbeddingDimensionMismatchError):
await svc.embed_batch(["a", "b"], use_cache=False)
assert client.embeddings.with_raw_response.create.await_count == 1
+252
View File
@@ -0,0 +1,252 @@
"""Re-embed the Kuzu graph into a new embedding space.
The source graph is opened read-only and never modified. A new graph is built
at the target path with the target width: every node and edge is copied, the
embedding columns are recomputed from their text, the HNSW index is rebuilt.
Runs in CI: kuzu is a core dependency.
"""
import hashlib
import json
import os
import pytest
from hy_memory import (
EmbeddingDimensionMismatchError,
MemoryConfig,
ReembedCancelled,
drop_graph,
reembed_graph,
verify_graph_reembed,
)
from hy_memory.models.memory import MemoryLayer, MemoryNode
from hy_memory.reembed import hierarchy_topic_label, plan_graph_reembed
from .test_reembed import FakeEmbedder
def _cfg(path, dims):
cfg = MemoryConfig()
cfg.graph_store.provider = "kuzu"
cfg.graph_store.db_path = str(path)
cfg.vector_store.embedding_dims = dims
return cfg
async def _open(path, dims):
from hy_memory.data.graph_store_kuzu import KuzuGraphStore
store = KuzuGraphStore(_cfg(path, dims))
await store.initialize()
assert store._available
return store
def _node(content, *, layer=MemoryLayer.L2_FACT, custom=None, user="alice"):
return MemoryNode(
user_id=user, agent_id="a", session_id="s1",
layer=layer, content=content, custom=custom or {},
)
async def _seed_graph(path, dims, model="old"):
"""A graph with every node and edge type the SDK writes."""
old = FakeEmbedder(model, dims)
gs = await _open(path, dims)
schema = _node(
"画像:喜欢提前规划", layer=MemoryLayer.L6_SCHEMA,
custom={"schema_type": "basic", "retrieval_text": "plans trips weeks ahead",
"_has_beh": True, "beh_abstraction": "prepares early under uncertainty"},
)
schema._graph_embedding = old.vector("plans trips weeks ahead")
core = _node("core: anticipates", layer=MemoryLayer.L6_SCHEMA, custom={"schema_type": "abstract"})
core._graph_embedding = old.vector("core: anticipates")
stale_beh = _node("beh without text", layer=MemoryLayer.L6_SCHEMA, custom={})
stale_beh._graph_embedding = old.vector("beh without text")
facts = [_node(f"fact {i}", user="bob" if i % 2 else "alice") for i in range(5)]
for node in [schema, core, stale_beh, *facts]:
await gs.upsert_memory_node(node)
await gs.update_embedding(schema.node_id, beh_embedding=old.vector("prepares early under uncertainty"))
await gs.update_embedding(stale_beh.node_id, beh_embedding=old.vector("lost text"))
await gs.add_cross_abstracts_to(schema.node_id, core.node_id)
ik = facts[0].get_isolation_key()
await gs.add_edge(facts[0].node_id, facts[1].node_id, "RELATED_TO", {"relation_type": "co_occurs", "weight": 0.7})
await gs.add_temporal_edge(facts[0].node_id, facts[2].node_id, topic_id="t")
for fact in facts[:3]:
await gs.add_topic_tag(fact.node_id, "travel", ik, embed_service=old)
await gs.add_topic_tag(facts[3].node_id, "work", ik) # no embed service: NULL embedding
await gs.add_topic_path(
facts[0].node_id, ["Travel", "Japan trip"], ik, "alice",
node_types=["topic", "series"], embed_service=old,
)
await gs.add_topic_path(facts[1].node_id, ["Work", "Standups"], ik, "alice", embed_service=old)
await gs.add_entity_mention(facts[0].node_id, "Kyoto", "GPE", ik)
await gs.ensure_vdbref("raw-1", "l1_raw")
await gs.add_derived_from(schema.node_id, "raw-1")
await gs.close()
return old, {"schema": schema, "core": core, "stale_beh": stale_beh, "facts": facts}
def _digest(path):
"""Hash of the Kuzu DB file and any sidecar files next to it."""
h = hashlib.sha256()
for name in sorted(os.listdir(path.parent)):
if name == path.name or name.startswith(path.name + "."):
h.update(name.encode())
with open(path.parent / name, "rb") as handle:
h.update(handle.read())
return h.hexdigest()
def _read(path, query):
import kuzu
db = kuzu.Database(str(path), read_only=True)
conn = kuzu.Connection(db)
try:
result = conn.execute(query)
rows = []
while result.has_next():
rows.append(result.get_next())
return rows
finally:
conn.close()
db.close()
@pytest.mark.parametrize(
"src_dims,tgt_dims,src_model,tgt_model",
[(512, 512, "old", "new"), (512, 1024, "old", "new"), (1024, 512, "old", "new")],
)
async def test_reembed_graph_copies_everything_at_the_new_width(
tmp_path, src_dims, tgt_dims, src_model, tgt_model,
):
source, target = tmp_path / "graph", tmp_path / "graph-new"
_, nodes = await _seed_graph(source, src_dims, model=src_model)
before = _digest(source)
plan = await plan_graph_reembed(source, target_dims=tgt_dims)
assert plan.source_width == src_dims and plan.target_width == tgt_dims
assert plan.memory_nodes == 8 and plan.topics == 2 and plan.hierarchy_topics == 4
assert plan.memory_embeddings == 3 and plan.beh_embeddings == 2
assert plan.topic_embeddings == 1 and plan.hierarchy_topic_embeddings == 4
new = FakeEmbedder(tgt_model, tgt_dims)
seen = []
report = await reembed_graph(
source, target, embedding=new, target_dims=tgt_dims, batch_size=3, on_progress=seen.append,
)
assert _digest(source) == before # the source graph is untouched
assert not os.path.exists(str(target) + ".reembed-partial")
assert report.source_width == src_dims and report.target_width == tgt_dims
assert report.node_counts == plan.node_counts
assert report.rel_counts == plan.rel_counts
assert report.embedded == {
"Memory.embedding": 3, "Memory.beh_embedding": 1,
"Topic.embedding": 1, "HierarchyTopic.embedding": 4,
}
assert report.cleared == {"Memory.beh_embedding": 1} # beh vector with no beh_abstraction
assert report.hierarchy_labels_exact == 4 and report.hierarchy_labels_inexact == ()
assert report.vector_indexes == ("memory_content_idx",)
assert seen
check = await verify_graph_reembed(source, target, target_dims=tgt_dims)
assert check.ok, check.problems
assert check.vector_index and check.target_width == tgt_dims
assert all(src == dst for src, dst in check.node_counts.values())
assert all(src == dst for src, dst in check.rel_counts.values())
schema = nodes["schema"]
rows = _read(target, f"MATCH (m:Memory {{node_id: '{schema.node_id}'}}) RETURN m.embedding, m.beh_embedding, m.custom_json, m.content;")
emb, beh, custom, content = rows[0]
assert len(emb) == tgt_dims
assert emb == pytest.approx(new.vector("plans trips weeks ahead"), abs=1e-5)
assert beh == pytest.approx(new.vector("prepares early under uncertainty"), abs=1e-5)
assert json.loads(custom)["retrieval_text"] == "plans trips weeks ahead" and content == schema.content
rows = _read(target, f"MATCH (m:Memory {{node_id: '{nodes['core'].node_id}'}}) RETURN m.embedding;")
assert rows[0][0] == pytest.approx(new.vector("core: anticipates"), abs=1e-5)
rows = _read(target, f"MATCH (m:Memory {{node_id: '{nodes['facts'][0].node_id}'}}) RETURN m.embedding IS NULL;")
assert rows[0][0] is True # a NULL source vector stays NULL
topics = dict(_read(target, "MATCH (t:Topic) RETURN t.name, t.embedding;"))
assert topics["travel"] == pytest.approx(new.vector("travel"), abs=1e-5)
assert topics["work"] is None
for subject, path_json, embedding in _read(target, "MATCH (t:HierarchyTopic) RETURN t.subject, t.path_json, t.embedding;"):
label, exact = hierarchy_topic_label(subject, path_json)
assert exact and embedding == pytest.approx(new.vector(label), abs=1e-5)
# The rebuilt graph serves the SDK: vector search on the new index works.
gs = await _open(target, tgt_dims)
try:
hits = await gs.vector_search(
query_embedding=new.vector("plans trips weeks ahead"),
isolation_key=schema.get_isolation_key(), limit=1,
)
assert hits and hits[0]["node_id"] == schema.node_id
width = await gs._run("CALL table_info('Memory') RETURN *;")
assert dict((r[1], r[2]) for r in width)["embedding"] == f"FLOAT[{tgt_dims}]"
finally:
await gs.close()
async def test_source_labels_rebuild_exactly(tmp_path):
"""The stored HierarchyTopic vectors were computed from the rebuilt label."""
source = tmp_path / "graph"
old, _ = await _seed_graph(source, 64)
rows = _read(source, "MATCH (t:HierarchyTopic) RETURN t.subject, t.path_json, t.embedding;")
assert len(rows) == 4
for subject, path_json, embedding in rows:
label, exact = hierarchy_topic_label(subject, path_json)
assert exact
assert embedding == pytest.approx(old.vector(label), abs=1e-6)
def test_legacy_path_json_is_flagged_inexact():
label, exact = hierarchy_topic_label("alice", json.dumps(["Travel", "Japan"]), "Japan")
assert label == "alice: topic:Travel > facet:Japan" and not exact
label, exact = hierarchy_topic_label("alice", "", "Japan")
assert label == "alice: Japan" and not exact
async def test_cancel_and_failure_leave_no_target(tmp_path):
source, target = tmp_path / "graph", tmp_path / "graph-new"
await _seed_graph(source, 64)
with pytest.raises(ReembedCancelled):
await reembed_graph(
source, target, embedding=FakeEmbedder("new", 32), target_dims=32,
cancel=lambda: True,
)
with pytest.raises(EmbeddingDimensionMismatchError):
await reembed_graph(
source, target, embedding=FakeEmbedder("new", 32, wrong_width_after=1), target_dims=32,
)
assert not os.path.exists(target)
assert not os.path.exists(str(target) + ".reembed-partial")
# A restart starts over and succeeds.
report = await reembed_graph(source, target, embedding=FakeEmbedder("new", 32), target_dims=32)
assert report.node_counts["Memory"] == 8
with pytest.raises(FileExistsError):
await reembed_graph(source, target, embedding=FakeEmbedder("new", 32), target_dims=32)
again = await reembed_graph(
source, target, embedding=FakeEmbedder("new", 32), target_dims=32, overwrite=True,
)
assert again.rel_counts == report.rel_counts
async def test_refuses_open_source_and_drops_only_closed_graphs(tmp_path):
source, target = tmp_path / "graph", tmp_path / "graph-new"
await _seed_graph(source, 64)
live = await _open(source, 64)
try:
with pytest.raises(RuntimeError):
await reembed_graph(source, target, embedding=FakeEmbedder("new", 64), target_dims=64)
with pytest.raises(RuntimeError):
await drop_graph(source)
finally:
await live.close()
with pytest.raises(ValueError):
await reembed_graph(source, source, embedding=FakeEmbedder("new", 64), target_dims=64)
await reembed_graph(source, target, embedding=FakeEmbedder("new", 64), target_dims=64)
assert await drop_graph(target) is True
assert not os.path.exists(target)
assert await drop_graph(target) is False
assert os.path.exists(source)