Compare commits
27
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
7df371bb10 | ||
|
|
f38c370fbd | ||
|
|
e87d6305be | ||
|
|
7a78fca807 | ||
|
|
1520b6de1f | ||
|
|
3742b3e453 | ||
|
|
98bc52b9f4 | ||
|
|
30fea9394f | ||
|
|
e966e913b6 | ||
|
|
ba6248b90b | ||
|
|
6cad02bc82 | ||
|
|
fdc57822b5 | ||
|
|
4e1748f65a | ||
|
|
b20331aa39 | ||
|
|
104fbcce09 | ||
|
|
302fd30af1 | ||
|
|
8d0a069950 | ||
|
|
cdb44658cf | ||
|
|
2ab04291a9 | ||
|
|
3772a69c6b | ||
|
|
80d3b55723 | ||
|
|
b5ec0861f0 | ||
|
|
47e938da1b | ||
|
|
f5ea0af75d | ||
|
|
1c2f48579f | ||
|
|
d29615abeb | ||
|
|
0965519e16 |
No files matched your search
@@ -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={}
|
||||
|
||||
|
||||
@@ -4,6 +4,75 @@ Hy Memory and MemHub are versioned and released independently.
|
||||
|
||||
## Unreleased
|
||||
|
||||
## Hy Memory 2.0.10 — 2026-11-01
|
||||
|
||||
- Published
|
||||
[`hy-memory==2.0.10`](https://pypi.org/project/hy-memory/2.0.10/).
|
||||
- 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.
|
||||
- `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.
|
||||
- Raw conversations are stored already trimmed. MemHub now keeps every raw
|
||||
forever and lite mode stores nothing else, so tool output and long pasted
|
||||
logs no longer stay in full. One rule (`hy_memory/utils/raw_trim.py`) is
|
||||
applied to each message when the raw is written:
|
||||
- `user` and `system` turns are kept in full;
|
||||
- `assistant` turns and other speakers (`[John]:`) longer than 500
|
||||
characters, after stripping, are cut to the first 500 plus `...`. This is
|
||||
the limit history injection used before;
|
||||
- `tool` turns keep their head and tail, and the middle becomes
|
||||
`[...N bytes omitted by truncate...]`, with the coding limit
|
||||
`coding.tool_result_max_bytes` (`MEMORY_CODING_TOOL_RESULT_MAX_BYTES`,
|
||||
default 2048).
|
||||
|
||||
The rule is idempotent. It covers chat `L1_RAW` in lite, pro and ultra (lite
|
||||
keeps tool turns, because it has no coding route) and the `memory_history`
|
||||
ADD record. The coding/working `coding_trace` was already trimmed when it
|
||||
was written. Its tool and assistant truncations are now idempotent, so
|
||||
reading the history back no longer rewrites the omitted-count markers. The
|
||||
text embedded for `L1_RAW` is the trimmed text. New raws carry
|
||||
`custom.raw_trimmed=true`.
|
||||
- Extractor history ("last k messages") now joins the stored raws directly
|
||||
and does not trim them again. Raws written before this change have no
|
||||
`raw_trimmed` flag, so they are still trimmed with the same rule when they
|
||||
are read. The extractor still receives the current batch in full (tool
|
||||
turns are still removed on the chat route).
|
||||
|
||||
## MemHub 1.0.5 — 2026-10-28
|
||||
|
||||
Build 36, on `hy-memory==2.0.8`. Linux x64/arm64 Runtime and macOS arm64
|
||||
|
||||
@@ -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 |
|
||||
|
||||
|
||||
@@ -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",
|
||||
# 公共扩展协议 / 精确读取类型
|
||||
|
||||
+11
-2
@@ -3089,9 +3089,18 @@ class HyMemoryClient:
|
||||
if self._history_store is not None:
|
||||
try:
|
||||
isolation_key = self._build_isolation_key(uid, agent_id)
|
||||
# 从 content 或 messages 中提取用于记录的文本
|
||||
# 从 content 或 messages 中提取用于记录的文本。messages 与 L1_RAW
|
||||
# 同一规则先裁剪再落审计库(工具输出 / 长回复不原样长期保存)。
|
||||
from .utils.raw_trim import trim_raw_messages
|
||||
history_content = content or "; ".join(
|
||||
f"{m.role}: {m.content}" for m in chat_messages
|
||||
f"{m.role}: {m.content}"
|
||||
for m in trim_raw_messages(
|
||||
chat_messages,
|
||||
tool_max_bytes=getattr(
|
||||
getattr(self._config, "coding", None),
|
||||
"tool_result_max_bytes", None,
|
||||
),
|
||||
)
|
||||
)
|
||||
await self._history_store.record_add(
|
||||
memory_id=result.memory_id,
|
||||
|
||||
@@ -11,6 +11,7 @@ Coding Memory 预处理工具(规则,零 LLM)
|
||||
使用方式与 scene 路由见 docs/scene_coding_usage.md。
|
||||
"""
|
||||
|
||||
import re
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
from ..pipelines.base import ChatMessage
|
||||
@@ -63,6 +64,13 @@ DEFAULT_MAX_TOOL_RESULT_BYTES = 2048
|
||||
DEFAULT_HEAD_BYTES = 1024
|
||||
DEFAULT_TAIL_BYTES = 512
|
||||
|
||||
# 本模块自己产出的省略标记。截断结果会再次经过截断(落库前一次、读历史时又一次),
|
||||
# 认出自己的标记才能保证幂等:已截断的文本再截一次原样返回,而不是把省略计数改写。
|
||||
_TOOL_OMIT_MARKER_RE = re.compile(r"\n\n\[\.\.\.\d+ bytes omitted by truncate\.\.\.\]\n\n")
|
||||
# 标记本身的字节数上限(数字位数有限),用于判断「已截断且只超出一个标记的长度」。
|
||||
_TOOL_OMIT_MARKER_MAX_BYTES = 128
|
||||
_ASSISTANT_OMIT_MARKER_RE = re.compile(r" \[\.\.\.\d+ chars omitted\.\.\.\] ")
|
||||
|
||||
|
||||
def truncate_tool_result_text(
|
||||
text: str,
|
||||
@@ -81,6 +89,12 @@ def truncate_tool_result_text(
|
||||
raw = text.encode("utf-8")
|
||||
if len(raw) <= max_bytes:
|
||||
return text
|
||||
# 已截断过(max_bytes 很小时标记本身会让结果略超上限):原样返回,保证幂等。
|
||||
if (
|
||||
len(raw) <= max_bytes + _TOOL_OMIT_MARKER_MAX_BYTES
|
||||
and _TOOL_OMIT_MARKER_RE.search(text)
|
||||
):
|
||||
return text
|
||||
# 自定义 max_bytes 可能小于默认 head+tail(1024+512)。按原比例缩放内容窗口,
|
||||
# 给省略标记预留空间,不能因为窗口和过大而把超限原文直接返回。
|
||||
marker_budget = min(96, max(24, max_bytes // 4))
|
||||
@@ -161,6 +175,10 @@ def truncate_assistant_text(
|
||||
text = "" if text is None else str(text)
|
||||
if len(text) <= head_chars + tail_chars:
|
||||
return text
|
||||
# 已截断过:中段恰好是本函数的省略标记 → 原样返回,保证幂等(否则会把
|
||||
# 省略计数改写成标记自身的长度)。
|
||||
if _ASSISTANT_OMIT_MARKER_RE.fullmatch(text[head_chars:len(text) - tail_chars]):
|
||||
return text
|
||||
omitted = len(text) - head_chars - tail_chars
|
||||
return (
|
||||
text[:head_chars]
|
||||
|
||||
@@ -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
@@ -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",
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -280,6 +280,22 @@ class MemoryWriter(WritePipeline):
|
||||
custom["source_raw_memory_ids"] = source_ids
|
||||
return custom
|
||||
|
||||
def _flat_trimmed_raw(self, request: WriteRequest) -> str:
|
||||
"""把 request.messages 按 raw 统一规则裁剪后拼成 L1_RAW 文本。
|
||||
|
||||
只裁落库副本;request.messages 不动(本次 extractor 输入仍用完整 messages)。
|
||||
"""
|
||||
import copy as _copy
|
||||
from ..utils.raw_trim import trim_raw_messages
|
||||
|
||||
_coding_cfg = getattr(getattr(self, "config", None), "coding", None)
|
||||
_trimmed_req = _copy.copy(request)
|
||||
_trimmed_req.messages = trim_raw_messages(
|
||||
request.messages,
|
||||
tool_max_bytes=getattr(_coding_cfg, "tool_result_max_bytes", None),
|
||||
)
|
||||
return _trimmed_req.get_flat_content()
|
||||
|
||||
@staticmethod
|
||||
def _attach_client_meta(custom: Dict[str, Any], request: Optional[WriteRequest]) -> None:
|
||||
"""给新建节点挂上调用方 meta(单份 list)。
|
||||
@@ -2554,9 +2570,14 @@ class MemoryWriter(WritePipeline):
|
||||
try:
|
||||
# 参数校验
|
||||
content = request.content
|
||||
# raw 在写入时就按统一规则裁好(见 utils/raw_trim.py):工具输出、长
|
||||
# assistant 回复不再原样长期存储;历史拼接直接用存下来的文本,不再二次裁剪。
|
||||
# request.messages 保持原样——本次 extractor 输入仍是完整 messages。
|
||||
_raw_trimmed = False
|
||||
if not content and request.has_messages():
|
||||
content = request.get_flat_content()
|
||||
content = self._flat_trimmed_raw(request)
|
||||
request.content = content
|
||||
_raw_trimmed = True
|
||||
|
||||
if not content:
|
||||
response.error_code = 400
|
||||
@@ -2627,6 +2648,9 @@ class MemoryWriter(WritePipeline):
|
||||
vector_store = await self._get_vector_store()
|
||||
layer_enum = MemoryLayer.from_string(suggested_layer)
|
||||
_raw_custom: Dict[str, Any] = {}
|
||||
if _raw_trimmed:
|
||||
from ..utils.raw_trim import RAW_TRIMMED_KEY
|
||||
_raw_custom[RAW_TRIMMED_KEY] = True
|
||||
self._attach_client_meta(_raw_custom, request)
|
||||
mem_node = MemoryNode(
|
||||
user_id=request.user_id,
|
||||
@@ -2929,13 +2953,16 @@ class MemoryWriter(WritePipeline):
|
||||
agent_id: Optional[str] = None,
|
||||
exclude_memory_id: str = "",
|
||||
max_turns: int = 20,
|
||||
max_chars_assistant: int = 500,
|
||||
max_chars_assistant: Optional[int] = None,
|
||||
) -> str:
|
||||
"""
|
||||
获取该用户最近的 L1_RAW 对话记录,按轮次拆分后拼成 history context。
|
||||
|
||||
每条 L1_RAW 的 content 是多轮对话([user]: ...\n[assistant]: ...)。
|
||||
拆成单条 message 后,user 消息完整保留,assistant 消息截取前 max_chars_assistant 字符。
|
||||
raw 在写入时已按 utils/raw_trim.py 的规则裁好(custom[RAW_TRIMMED_KEY]),
|
||||
这里拆成单条 message 后直接拼接,不再二次裁剪。该规则之前写入的旧 raw
|
||||
没有这个标记,仍按同一(幂等)规则裁剪后再拼,避免旧的长回复 / 工具输出
|
||||
撑大 prompt。max_chars_assistant 只作用于这些旧 raw,默认同写入规则。
|
||||
最终取最近 max_turns 条 message(注意单位是 message 条数,非对话轮数;
|
||||
调用方按 k 轮 × 2 传入)。
|
||||
|
||||
@@ -2977,23 +3004,33 @@ class MemoryWriter(WritePipeline):
|
||||
# 避免把正文里形如 "[注]: xxx" 的普通文本误判成消息边界。
|
||||
import re as _re
|
||||
from .base import parse_user_name_line
|
||||
_KNOWN_ROLES = {"user", "assistant", "system", "tool"}
|
||||
# 带方括号的任意 label(含对方真人名字)
|
||||
_labeled_prefix = _re.compile(r"^\[([^\]\n]{1,32})\]:\s?")
|
||||
# 无方括号的已知 role(历史数据容错,沿用旧正则的 `\[?...\]?` 宽容度)
|
||||
_bare_role_prefix = _re.compile(r"^(user|assistant|system|tool):\s?")
|
||||
|
||||
# (label, content) 元组列表;label 为 role 或对方真人名字
|
||||
parsed: List[Tuple[str, str]] = []
|
||||
from ..utils.raw_trim import (
|
||||
DEFAULT_RAW_ASSISTANT_MAX_CHARS,
|
||||
RAW_TRIMMED_KEY,
|
||||
trim_raw_body,
|
||||
)
|
||||
if max_chars_assistant is None:
|
||||
max_chars_assistant = DEFAULT_RAW_ASSISTANT_MAX_CHARS
|
||||
_coding_cfg = getattr(getattr(self, "config", None), "coding", None)
|
||||
_tool_max_bytes = getattr(_coding_cfg, "tool_result_max_bytes", None)
|
||||
|
||||
# (label, content, 是否旧 raw 需兼容裁剪) 元组列表;label 为 role 或对方真人名字
|
||||
parsed: List[Tuple[str, str, bool]] = []
|
||||
user_name: str = ""
|
||||
for node in nodes:
|
||||
content = node.content or ""
|
||||
_legacy = not bool((node.custom or {}).get(RAW_TRIMMED_KEY))
|
||||
for _line_idx, line in enumerate(content.split("\n")):
|
||||
m = _labeled_prefix.match(line) or _bare_role_prefix.match(line)
|
||||
if m:
|
||||
label = m.group(1)
|
||||
body = line[m.end():]
|
||||
parsed.append((label, body))
|
||||
parsed.append((label, body, _legacy))
|
||||
continue
|
||||
# 名字行只可能是每个 node content 的第一行(get_flat_content 插在文首)。
|
||||
# 逐 node 判断而非用 `not parsed`:parsed 跨 node 累积,第二个 node
|
||||
@@ -3006,18 +3043,18 @@ class MemoryWriter(WritePipeline):
|
||||
continue
|
||||
if parsed:
|
||||
# 续行:拼回上一条 message(保留换行)
|
||||
prev_label, prev_body = parsed[-1]
|
||||
parsed[-1] = (prev_label, prev_body + "\n" + line)
|
||||
prev_label, prev_body, prev_legacy = parsed[-1]
|
||||
parsed[-1] = (prev_label, prev_body + "\n" + line, prev_legacy)
|
||||
else:
|
||||
# 开头就是无前缀行(异常/老数据):当成 user 消息兜底
|
||||
parsed.append(("user", line))
|
||||
parsed.append(("user", line, _legacy))
|
||||
|
||||
# 取最近 max_turns 条 message(按 message 截断,绝不从消息中间切开)
|
||||
recent = parsed[-max_turns:] if len(parsed) > max_turns else parsed
|
||||
if not recent:
|
||||
return ""
|
||||
|
||||
# 格式化:非 user 侧整条截断(含对方真人),user 完整保留
|
||||
# 格式化:新 raw 直接拼接;旧 raw 按写入同一规则兼容裁剪
|
||||
all_messages: List[str] = []
|
||||
if user_name:
|
||||
from .base import render_user_name_line
|
||||
@@ -3025,17 +3062,18 @@ class MemoryWriter(WritePipeline):
|
||||
|
||||
_name_line = render_user_name_line(
|
||||
user_name,
|
||||
is_chinese("\n".join(b for _, b in recent)),
|
||||
is_chinese("\n".join(b for _, b, _ in recent)),
|
||||
)
|
||||
if _name_line:
|
||||
all_messages.append(_name_line)
|
||||
for label, body in recent:
|
||||
for label, body, legacy in recent:
|
||||
if legacy:
|
||||
body = trim_raw_body(
|
||||
label, body,
|
||||
assistant_max_chars=max_chars_assistant,
|
||||
tool_max_bytes=_tool_max_bytes,
|
||||
)
|
||||
body = body.strip()
|
||||
# 截断口径:assistant + 对方真人(非已知 role 的 label)。
|
||||
# user 完整保留;system / tool 沿用历史行为(不截断)。
|
||||
_truncate = label == "assistant" or label not in _KNOWN_ROLES
|
||||
if _truncate and len(body) > max_chars_assistant:
|
||||
body = body[:max_chars_assistant] + "..."
|
||||
all_messages.append(f"[{label}]: {body}")
|
||||
|
||||
return "\n".join(all_messages)
|
||||
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,91 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
原始对话(raw)落库前的统一裁剪规则(规则,零 LLM)。
|
||||
|
||||
raw 会被长期保存(MemHub 已取消 raw 的保留期限;lite 模式下 raw 是唯一存储),
|
||||
不裁剪会把工具输出、大段粘贴日志原样存下来。因此 raw 在**写入时**就裁好,
|
||||
读取方(extractor 的「最近 k 条」历史注入)直接拼接,不再二次裁剪。
|
||||
|
||||
规则按说话人标签(ChatMessage.speaker_label() / L1_RAW 里的 ``[label]``):
|
||||
|
||||
- ``user``:完整保留;
|
||||
- ``system``:完整保留(与历史注入原口径一致);
|
||||
- ``tool``:头尾保留、中段替换为 ``[...N bytes omitted by truncate...]``,
|
||||
上限为 ``coding.tool_result_max_bytes``(env ``MEMORY_CODING_TOOL_RESULT_MAX_BYTES``,
|
||||
默认 2048),与 coding 链路同一个函数;
|
||||
- 其余(``assistant``、对方真人名字等非已知 role 的 label):strip 后超过
|
||||
``assistant_max_chars``(默认 500,即历史注入原来的 ``max_chars_assistant``)
|
||||
时截成前 N 个字符 + ``...``。
|
||||
|
||||
所有规则确定且幂等:对已裁剪的文本再裁一次结果不变。
|
||||
"""
|
||||
|
||||
from typing import List, Optional
|
||||
|
||||
from ..pipelines.base import ChatMessage
|
||||
|
||||
# 历史注入原 `_get_recent_history(max_chars_assistant=500)` 的默认值,现为唯一来源。
|
||||
DEFAULT_RAW_ASSISTANT_MAX_CHARS = 500
|
||||
|
||||
# L1_RAW 节点 custom 上的标记:内容已按本模块规则在写入时裁剪。
|
||||
# 历史拼接据此区分「新 raw 直接拼接」与「本规则之前写入的旧 raw 需兼容裁剪」。
|
||||
RAW_TRIMMED_KEY = "raw_trimmed"
|
||||
|
||||
# 不裁剪的标签(tool 单独处理)
|
||||
_KEEP_LABELS = frozenset({"user", "system"})
|
||||
|
||||
|
||||
def trim_raw_body(
|
||||
label: str,
|
||||
body: str,
|
||||
*,
|
||||
assistant_max_chars: int = DEFAULT_RAW_ASSISTANT_MAX_CHARS,
|
||||
tool_max_bytes: Optional[int] = None,
|
||||
) -> str:
|
||||
"""按说话人标签裁剪单条消息正文。幂等。"""
|
||||
if not isinstance(body, str):
|
||||
body = "" if body is None else str(body)
|
||||
if label in _KEEP_LABELS:
|
||||
return body
|
||||
if label == "tool":
|
||||
from ..coding.preproc import (
|
||||
DEFAULT_MAX_TOOL_RESULT_BYTES,
|
||||
truncate_tool_result_text,
|
||||
)
|
||||
|
||||
return truncate_tool_result_text(
|
||||
body, max_bytes=int(tool_max_bytes or DEFAULT_MAX_TOOL_RESULT_BYTES),
|
||||
)
|
||||
stripped = body.strip()
|
||||
if len(stripped) > assistant_max_chars:
|
||||
return stripped[:assistant_max_chars] + "..."
|
||||
return body
|
||||
|
||||
|
||||
def trim_raw_messages(
|
||||
messages: List[ChatMessage],
|
||||
*,
|
||||
assistant_max_chars: int = DEFAULT_RAW_ASSISTANT_MAX_CHARS,
|
||||
tool_max_bytes: Optional[int] = None,
|
||||
) -> List[ChatMessage]:
|
||||
"""返回裁剪后的新列表;不修改入参(messages 仍要原样给 extractor)。"""
|
||||
out: List[ChatMessage] = []
|
||||
for m in messages:
|
||||
new_content = trim_raw_body(
|
||||
m.speaker_label(),
|
||||
m.content,
|
||||
assistant_max_chars=assistant_max_chars,
|
||||
tool_max_bytes=tool_max_bytes,
|
||||
)
|
||||
if new_content == m.content:
|
||||
out.append(m)
|
||||
continue
|
||||
out.append(ChatMessage(
|
||||
role=m.role,
|
||||
content=new_content,
|
||||
name=m.name,
|
||||
tool_calls=list(m.tool_calls),
|
||||
tool_call_id=m.tool_call_id,
|
||||
tool_name=m.tool_name,
|
||||
))
|
||||
return out
|
||||
+2
-2
@@ -13,8 +13,8 @@
|
||||
接口先补 SDK。Runtime writers 停止后的 App 启动迁移可直接操作底层存储,
|
||||
但不得把迁移专用方法加入公共 Client。
|
||||
3. MemHub **不存 memory 实体**;**可以存 `memory_id`** 做组织/回执。原文进入 **SDK raw**。
|
||||
4. SDK raw 只保留滚动三天;一条 raw 对应一个完整 turn,只含真人 user 与
|
||||
assistant,tool 不持久化。提取后的高层 memory 不受该窗口影响。
|
||||
4. SDK raw 全部保留,任何模式(lite / pro / ultra)都不删除;一条 raw 对应一个
|
||||
完整 turn,只含真人 user 与 assistant,tool 不持久化。
|
||||
5. `source_raw_ids` 是 SDK 内部溯源关系,不是产品 Evidence 合同;记忆详情不提供
|
||||
“原始依据”或按 raw ID 展开原文的入口。
|
||||
|
||||
|
||||
+2
-2
@@ -18,7 +18,7 @@ If you run into a problem, please open an Issue with the version, platform,repro
|
||||
|
||||
## 发行与平台合同
|
||||
|
||||
源码当前固定使用公开 PyPI 正式版 `hy-memory==2.0.9`。MemHub 源码版本为 `1.0.5`
|
||||
源码当前固定使用公开 PyPI 正式版 `hy-memory==2.0.10`。MemHub 源码版本为 `1.0.5`
|
||||
(build 36),GitHub 仓库为
|
||||
[`Tencent-Hunyuan/Hy-Memory`](https://github.com/Tencent-Hunyuan/Hy-Memory),npm 包名为
|
||||
`memhub-cli`。当前代码架构见
|
||||
@@ -118,7 +118,7 @@ MemHub 的 SQLite 数据默认位于 `~/.one-memory`,不在 npm 包或 Runtime
|
||||
|
||||
## Hy Memory
|
||||
|
||||
SDK 依赖固定为公开 PyPI 正式版 `hy-memory==2.0.9`。SDK 配置、权限边界及 search
|
||||
SDK 依赖固定为公开 PyPI 正式版 `hy-memory==2.0.10`。SDK 配置、权限边界及 search
|
||||
返回格式见 [`MEMORY.md`](MEMORY.md)。HY search 返回的 `content` 已包含 SDK 统一生成的
|
||||
逻辑时间和演化上下文;调用方不应再次转换时间戳或重复拼接 `When`。
|
||||
|
||||
|
||||
@@ -2352,7 +2352,7 @@
|
||||
{
|
||||
"SPDXID": "SPDXRef-Package-hy-memory",
|
||||
"name": "hy-memory",
|
||||
"versionInfo": "2.0.9",
|
||||
"versionInfo": "2.0.10",
|
||||
"downloadLocation": "NOASSERTION",
|
||||
"filesAnalyzed": false,
|
||||
"licenseConcluded": "Apache-2.0",
|
||||
|
||||
@@ -9,7 +9,7 @@ Per-package license texts are not reproduced here yet; the license
|
||||
column is filled in as part of release review, and the SBOM records
|
||||
`NOASSERTION` until then rather than guessing.
|
||||
|
||||
## hy-memory 2.0.9
|
||||
## hy-memory 2.0.10
|
||||
|
||||
License: Apache-2.0
|
||||
|
||||
|
||||
@@ -315,6 +315,8 @@ describe("catalog exit codes", () => {
|
||||
"models",
|
||||
"memories",
|
||||
"profile",
|
||||
"settings",
|
||||
"evolve",
|
||||
"graph-width",
|
||||
"host",
|
||||
"recall",
|
||||
|
||||
+424
-28
@@ -122,6 +122,14 @@ import {
|
||||
doctorExitCode,
|
||||
renderDoctorText,
|
||||
} from "./doctor-report.js";
|
||||
import {
|
||||
activateWithEmbeddingSwitch,
|
||||
cleanupSpace,
|
||||
listSpaces,
|
||||
type SwitchIo,
|
||||
switchCancel,
|
||||
switchStatus,
|
||||
} from "./embedding-switch-flow.js";
|
||||
import {
|
||||
EXIT_FAILED,
|
||||
EXIT_OK,
|
||||
@@ -188,6 +196,25 @@ import {
|
||||
removeManagedCliLink,
|
||||
rollbackManagedCliLink,
|
||||
} from "./managed-cli-link.js";
|
||||
import {
|
||||
type FeaturesDocument,
|
||||
featuresExitCode,
|
||||
parseFeaturesCommand,
|
||||
renderFeaturesText,
|
||||
runFeaturesCommand,
|
||||
} from "./memory-features-flow.js";
|
||||
import {
|
||||
type EvolveDocument,
|
||||
evolveExitCode,
|
||||
type ModeDocument,
|
||||
modeExitCode,
|
||||
parseEvolveCommand,
|
||||
parseSettingsCommand,
|
||||
renderEvolveText,
|
||||
renderModeText,
|
||||
runEvolveCommand,
|
||||
runModeCommand,
|
||||
} from "./memory-mode-flow.js";
|
||||
import {
|
||||
addMemory,
|
||||
DEFAULT_USER_ID,
|
||||
@@ -229,7 +256,6 @@ import {
|
||||
renderModelReportText,
|
||||
} from "./model-report.js";
|
||||
import {
|
||||
activateModelProfile,
|
||||
editModelProfile,
|
||||
type ModelWriteDocument,
|
||||
modelWriteExitCode,
|
||||
@@ -325,6 +351,8 @@ export interface RunControls {
|
||||
watchSignal?: AbortSignal;
|
||||
watchWait?: (milliseconds: number, signal: AbortSignal) => Promise<void>;
|
||||
watchNow?: () => string;
|
||||
/** `memhub evolve --wait` polls on this; tests inject an instant one. */
|
||||
sleep?: (milliseconds: number) => Promise<void>;
|
||||
readStdin?: () => Promise<Buffer | null>;
|
||||
stdinIsTTY?: boolean;
|
||||
/** Injectable onboarding prompter so the wizard can be driven in tests
|
||||
@@ -1682,6 +1710,18 @@ async function runCommand(
|
||||
);
|
||||
if (apiKey === undefined) return usage();
|
||||
}
|
||||
if (
|
||||
parsed.kind === "spaces" &&
|
||||
parsed.cleanup !== undefined &&
|
||||
!parsed.confirmed
|
||||
) {
|
||||
return emitConfirmPreview(asJson, {
|
||||
command: "models spaces",
|
||||
target: parsed.cleanup,
|
||||
effect: `删除已停用的向量空间 ${parsed.cleanup}(不可恢复;正在使用的空间不能删除)`,
|
||||
retry: `${PRODUCT} models spaces --cleanup ${parsed.cleanup} --yes`,
|
||||
});
|
||||
}
|
||||
if (parsed.kind === "remove" && !parsed.confirmed) {
|
||||
return emitConfirmPreview(asJson, {
|
||||
command: "models remove",
|
||||
@@ -1704,6 +1744,50 @@ async function runCommand(
|
||||
return connected.exitCode;
|
||||
}
|
||||
const { reader, operations } = connected;
|
||||
// Ctrl-C while following a switch stops watching, never the switch.
|
||||
const interrupt = new AbortController();
|
||||
const switchIo: SwitchIo = {
|
||||
reader,
|
||||
operations,
|
||||
emit: (text) => {
|
||||
emit(text);
|
||||
},
|
||||
asJson,
|
||||
interactive:
|
||||
!asJson && (controls.stdinIsTTY ?? process.stdin.isTTY) === true,
|
||||
confirm: async (question) => {
|
||||
const prompter =
|
||||
controls.wizardPrompter ?? (await createTerminalPrompter());
|
||||
try {
|
||||
return /^(y|yes)$/i.test((await prompter.ask(question)).trim());
|
||||
} finally {
|
||||
if ("close" in prompter && typeof prompter.close === "function") {
|
||||
prompter.close();
|
||||
}
|
||||
}
|
||||
},
|
||||
...(controls.watchWait ? { wait: controls.watchWait } : {}),
|
||||
signal: controls.watchSignal ?? interrupt.signal,
|
||||
};
|
||||
const withInterrupt = async (run: () => Promise<number>) => {
|
||||
const onInterrupt = () => interrupt.abort();
|
||||
process.once("SIGINT", onInterrupt);
|
||||
try {
|
||||
return await run();
|
||||
} finally {
|
||||
process.removeListener("SIGINT", onInterrupt);
|
||||
}
|
||||
};
|
||||
if (parsed.kind === "switch") {
|
||||
return parsed.action === "status"
|
||||
? await switchStatus(switchIo)
|
||||
: await switchCancel(switchIo);
|
||||
}
|
||||
if (parsed.kind === "spaces") {
|
||||
return parsed.cleanup === undefined
|
||||
? await listSpaces(switchIo)
|
||||
: await cleanupSpace(parsed.cleanup, switchIo);
|
||||
}
|
||||
if (parsed.kind === "add") {
|
||||
const created = await createModelProfile(
|
||||
config,
|
||||
@@ -1718,33 +1802,66 @@ async function runCommand(
|
||||
emitFlowDocument(created, asJson);
|
||||
return flowExitCode(created);
|
||||
}
|
||||
const activated = await activateModelProfile(
|
||||
String(created.receipt.profile_id),
|
||||
// Core activates only a probed profile: probe first, then the same
|
||||
// path as `models activate`, so an embedding change shows its plan.
|
||||
const profileId = String(created.receipt.profile_id);
|
||||
const probe = await probeUntilSettled(
|
||||
reader,
|
||||
operations,
|
||||
profileId,
|
||||
controls.watchWait,
|
||||
);
|
||||
const document = {
|
||||
...created,
|
||||
activation:
|
||||
activated.state === "committed"
|
||||
? { state: "committed" }
|
||||
: { state: activated.state, code: activated.code },
|
||||
};
|
||||
const probed =
|
||||
probe.state === "completed" &&
|
||||
Boolean(
|
||||
(
|
||||
(probe.receipt.result as Record<string, unknown> | undefined)
|
||||
?.capabilities as Record<string, unknown> | undefined
|
||||
)?.embedding_available,
|
||||
);
|
||||
emit(
|
||||
asJson
|
||||
? `${JSON.stringify(document)}\n`
|
||||
? `${JSON.stringify({
|
||||
...created,
|
||||
probe: {
|
||||
state: probe.state,
|
||||
...("operation_id" in probe && probe.operation_id
|
||||
? { operation_id: probe.operation_id }
|
||||
: {}),
|
||||
...("code" in probe ? { code: probe.code } : {}),
|
||||
},
|
||||
})}\n`
|
||||
: `${renderFlowText(created)}${
|
||||
activated.state === "committed"
|
||||
? "已设为当前使用的模型配置。\n"
|
||||
: `设为当前使用未完成:${activated.state}(${activated.code})\n`
|
||||
probed
|
||||
? "探测通过。\n"
|
||||
: `探测未通过(${"code" in probe ? probe.code : probe.state}),没有设为当前使用;用 ${PRODUCT} models test ${profileId} 查看原因。\n`
|
||||
}`,
|
||||
);
|
||||
return activated.state === "committed" ? EXIT_OK : EXIT_FAILED;
|
||||
if (!probed) return EXIT_FAILED;
|
||||
return await withInterrupt(() =>
|
||||
activateWithEmbeddingSwitch(
|
||||
profileId,
|
||||
{ confirmed: parsed.confirmed, follow: parsed.follow },
|
||||
switchIo,
|
||||
),
|
||||
);
|
||||
}
|
||||
if (parsed.kind === "edit" || parsed.kind === "activate") {
|
||||
const document =
|
||||
parsed.kind === "edit"
|
||||
? await editModelProfile(parsed.profileId, config, apiKey, operations)
|
||||
: await activateModelProfile(parsed.profileId, operations);
|
||||
if (parsed.kind === "activate") {
|
||||
return await withInterrupt(() =>
|
||||
activateWithEmbeddingSwitch(
|
||||
parsed.profileId,
|
||||
{ confirmed: parsed.confirmed, follow: parsed.follow },
|
||||
switchIo,
|
||||
),
|
||||
);
|
||||
}
|
||||
if (parsed.kind === "edit") {
|
||||
const document = await editModelProfile(
|
||||
parsed.profileId,
|
||||
config,
|
||||
apiKey,
|
||||
operations,
|
||||
);
|
||||
emitModelWriteDocument(document, asJson);
|
||||
return modelWriteExitCode(document);
|
||||
}
|
||||
@@ -2018,12 +2135,17 @@ async function runCommand(
|
||||
return sourceProgressExitCode(document);
|
||||
}
|
||||
if (named === "profile") {
|
||||
const profileFlags = arguments_.slice(1);
|
||||
if (
|
||||
arguments_.filter((item) => item === "--json").length > 1 ||
|
||||
arguments_.slice(1).some((value) => value !== "--json")
|
||||
profileFlags.filter((item) => item === "--json").length > 1 ||
|
||||
profileFlags.filter((item) => item === "--versions").length > 1 ||
|
||||
profileFlags.some((value) => value !== "--json" && value !== "--versions")
|
||||
) {
|
||||
return usage();
|
||||
}
|
||||
if (profileFlags.includes("--versions")) {
|
||||
return await profileVersionsCommand(environment, { asJson, userId });
|
||||
}
|
||||
const connected = await connectRuntime(environment, {
|
||||
readerTimeoutMs: 15_000,
|
||||
});
|
||||
@@ -2223,6 +2345,89 @@ async function runCommand(
|
||||
);
|
||||
return outcome.state === "committed" ? EXIT_OK : EXIT_FAILED;
|
||||
}
|
||||
if (
|
||||
named === "settings" &&
|
||||
arguments_.slice(1).find((value) => value !== "--json") === "features"
|
||||
) {
|
||||
const parsed = parseFeaturesCommand(
|
||||
arguments_
|
||||
.slice(1)
|
||||
.filter((value, index, all) => index !== all.indexOf("features")),
|
||||
);
|
||||
if (!parsed) return usage();
|
||||
const connected = await connectRuntime(environment);
|
||||
if (!connected.ok) {
|
||||
const document: FeaturesDocument = {
|
||||
schema_version: 1,
|
||||
product: PRODUCT,
|
||||
command: "settings features",
|
||||
state: "unavailable",
|
||||
code: connected.code,
|
||||
};
|
||||
emit(
|
||||
asJson ? `${JSON.stringify(document)}\n` : renderFeaturesText(document),
|
||||
);
|
||||
return connected.exitCode;
|
||||
}
|
||||
const document = await runFeaturesCommand(
|
||||
parsed,
|
||||
connected.reader,
|
||||
connected.operations,
|
||||
);
|
||||
emit(
|
||||
asJson ? `${JSON.stringify(document)}\n` : renderFeaturesText(document),
|
||||
);
|
||||
return featuresExitCode(document);
|
||||
}
|
||||
if (named === "settings") {
|
||||
const parsed = parseSettingsCommand(arguments_.slice(1));
|
||||
if (!parsed) return usage();
|
||||
const connected = await connectRuntime(environment);
|
||||
if (!connected.ok) {
|
||||
const document: ModeDocument = {
|
||||
schema_version: 1,
|
||||
product: PRODUCT,
|
||||
command: "settings mode",
|
||||
state: "unavailable",
|
||||
code: connected.code,
|
||||
};
|
||||
emit(asJson ? `${JSON.stringify(document)}\n` : renderModeText(document));
|
||||
return connected.exitCode;
|
||||
}
|
||||
const document = await runModeCommand(
|
||||
parsed,
|
||||
connected.reader,
|
||||
connected.operations,
|
||||
);
|
||||
emit(asJson ? `${JSON.stringify(document)}\n` : renderModeText(document));
|
||||
return modeExitCode(document);
|
||||
}
|
||||
if (named === "evolve") {
|
||||
const parsed = parseEvolveCommand(arguments_.slice(1));
|
||||
if (!parsed) return usage();
|
||||
const connected = await connectRuntime(environment);
|
||||
if (!connected.ok) {
|
||||
const document: EvolveDocument = {
|
||||
schema_version: 1,
|
||||
product: PRODUCT,
|
||||
command: "evolve",
|
||||
state: "unavailable",
|
||||
code: connected.code,
|
||||
};
|
||||
emit(
|
||||
asJson ? `${JSON.stringify(document)}\n` : renderEvolveText(document),
|
||||
);
|
||||
return connected.exitCode;
|
||||
}
|
||||
const document = await runEvolveCommand(parsed, {
|
||||
userId: userIdParameters(userId).user_id,
|
||||
reader: connected.reader,
|
||||
operations: connected.operations,
|
||||
sleep: controls.sleep,
|
||||
});
|
||||
emit(asJson ? `${JSON.stringify(document)}\n` : renderEvolveText(document));
|
||||
return evolveExitCode(document);
|
||||
}
|
||||
if (named === "graph-width") {
|
||||
const parsed = parseGraphWidthCommand(arguments_.slice(1));
|
||||
if (!parsed) return usage();
|
||||
@@ -3888,6 +4093,108 @@ async function emitModelReceipt(
|
||||
|
||||
/** Attach to a verified Runtime, telling a stopped one (exit 3) from one
|
||||
* that is up but unusable (exit 1). */
|
||||
/**
|
||||
* `memhub profile --versions`: the long-term profile's L0 version chains,
|
||||
* latest first, so a value that replaced a better one can be found by id.
|
||||
*/
|
||||
async function profileVersionsCommand(
|
||||
environment: NodeJS.ProcessEnv,
|
||||
options: { asJson: boolean; userId?: string },
|
||||
): Promise<number> {
|
||||
const fail = (state: string, code: string, exitCode: number): number => {
|
||||
emit(
|
||||
options.asJson
|
||||
? `${JSON.stringify({ schema_version: 2, product: PRODUCT, state, code })}\n`
|
||||
: `画像版本读取失败:${code}\n`,
|
||||
);
|
||||
return exitCode;
|
||||
};
|
||||
const connected = await connectRuntime(environment, {
|
||||
readerTimeoutMs: 15_000,
|
||||
});
|
||||
if (!connected.ok) {
|
||||
return fail("unavailable", connected.code, connected.exitCode);
|
||||
}
|
||||
const outcome = await connected.reader<unknown>(
|
||||
"/api/v2/user-profile/versions",
|
||||
userIdParameters(options.userId),
|
||||
);
|
||||
const chains = outcome.state === "ok" ? profileChains(outcome.data) : null;
|
||||
if (outcome.state !== "ok" || chains === null) {
|
||||
return fail(
|
||||
outcome.state === "ok" ? "invalid" : outcome.state,
|
||||
outcome.state === "ok" ? "response_invalid" : outcome.code,
|
||||
EXIT_FAILED,
|
||||
);
|
||||
}
|
||||
if (options.asJson) {
|
||||
emit(
|
||||
`${JSON.stringify({
|
||||
schema_version: 2,
|
||||
product: PRODUCT,
|
||||
state: "ok",
|
||||
...(outcome.data as object),
|
||||
})}\n`,
|
||||
);
|
||||
return EXIT_OK;
|
||||
}
|
||||
if (chains.length === 0) {
|
||||
emit("尚未形成用户画像版本。\n");
|
||||
return EXIT_OK;
|
||||
}
|
||||
const lines: string[] = [];
|
||||
for (const chain of chains) {
|
||||
lines.push(`画像 ${chain.memory_id}(${chain.status})`);
|
||||
for (const version of chain.versions) {
|
||||
const when =
|
||||
typeof version.gmt_created === "number"
|
||||
? new Date(version.gmt_created * 1000).toISOString()
|
||||
: "时间未知";
|
||||
lines.push(
|
||||
` ${version.current ? "*" : " "} ${version.memory_id} ${when} ${version.content}`,
|
||||
);
|
||||
}
|
||||
}
|
||||
lines.push(
|
||||
"* 表示当前生效的版本;删除错误版本:memhub memories delete <memory-id> --yes",
|
||||
);
|
||||
emit(`${lines.join("\n")}\n`);
|
||||
return EXIT_OK;
|
||||
}
|
||||
|
||||
type ProfileChainView = {
|
||||
memory_id: string;
|
||||
status: string;
|
||||
versions: {
|
||||
memory_id: string;
|
||||
content: string;
|
||||
gmt_created?: number | null;
|
||||
current: boolean;
|
||||
}[];
|
||||
};
|
||||
|
||||
function profileChains(value: unknown): ProfileChainView[] | null {
|
||||
if (!value || typeof value !== "object") return null;
|
||||
const chains = (value as { chains?: unknown }).chains;
|
||||
if (!Array.isArray(chains)) return null;
|
||||
const valid = chains.every(
|
||||
(chain) =>
|
||||
chain &&
|
||||
typeof chain === "object" &&
|
||||
typeof (chain as ProfileChainView).memory_id === "string" &&
|
||||
typeof (chain as ProfileChainView).status === "string" &&
|
||||
Array.isArray((chain as ProfileChainView).versions) &&
|
||||
(chain as ProfileChainView).versions.every(
|
||||
(version) =>
|
||||
version &&
|
||||
typeof version.memory_id === "string" &&
|
||||
typeof version.content === "string" &&
|
||||
typeof version.current === "boolean",
|
||||
),
|
||||
);
|
||||
return valid ? (chains as ProfileChainView[]) : null;
|
||||
}
|
||||
|
||||
async function connectRuntime(
|
||||
environment: NodeJS.ProcessEnv,
|
||||
options?: { readerTimeoutMs?: number },
|
||||
@@ -4155,6 +4462,8 @@ const SDK_MIGRATION_FAILURES: Record<string, string> = {
|
||||
"核对发现新版本还读不到你的全部记忆,没有切换数据源。",
|
||||
sdk_migration_model_profile_required:
|
||||
"迁移需要先有可用的模型配置,用来核对新版本能否读回记忆。请先用 memhub models setup 完成配置,再运行 memhub restart 重试。",
|
||||
sdk_migration_vector_profile_missing:
|
||||
"记忆的向量来自一个已删除的模型配置,当前配置读不回它们,所以没有迁移。请重新添加并启用那个模型配置,再运行 memhub restart 重试。",
|
||||
};
|
||||
|
||||
const MIGRATION_ISSUES_URL =
|
||||
@@ -5076,6 +5385,8 @@ type ModelsCommand =
|
||||
configFile: string;
|
||||
apiKeyFile?: string;
|
||||
activate: boolean;
|
||||
confirmed: boolean;
|
||||
follow?: boolean;
|
||||
}
|
||||
| {
|
||||
kind: "edit";
|
||||
@@ -5083,7 +5394,14 @@ type ModelsCommand =
|
||||
configFile: string;
|
||||
apiKeyFile?: string;
|
||||
}
|
||||
| { kind: "activate"; profileId: string }
|
||||
| {
|
||||
kind: "activate";
|
||||
profileId: string;
|
||||
confirmed: boolean;
|
||||
follow?: boolean;
|
||||
}
|
||||
| { kind: "switch"; action: "status" | "cancel" }
|
||||
| { kind: "spaces"; cleanup?: string; confirmed: boolean }
|
||||
| { kind: "test"; profileId: string }
|
||||
| { kind: "receipt"; operationId: string }
|
||||
| { kind: "remove"; profileId: string; confirmed: boolean };
|
||||
@@ -5124,16 +5442,46 @@ function parseModelsCommand(rest: readonly string[]): ModelsCommand | null {
|
||||
const flags = readFlags(
|
||||
tokens.slice(1),
|
||||
["--config-file", "--api-key-file"],
|
||||
["--activate"],
|
||||
["--activate", "--yes", "--wait", "--no-wait"],
|
||||
);
|
||||
const configFile = flags?.values.get("--config-file");
|
||||
if (!flags || configFile === undefined) return null;
|
||||
const follow = followFlag(flags.switches);
|
||||
if (follow === null) return null;
|
||||
const activate = flags.switches.has("--activate");
|
||||
// The switch flags only mean something when the add also activates.
|
||||
if (
|
||||
!activate &&
|
||||
["--yes", "--wait", "--no-wait"].some((flag) => flags.switches.has(flag))
|
||||
) {
|
||||
return null;
|
||||
}
|
||||
const apiKeyFile = flags.values.get("--api-key-file");
|
||||
return {
|
||||
kind: "add",
|
||||
configFile,
|
||||
...(apiKeyFile === undefined ? {} : { apiKeyFile }),
|
||||
activate: flags.switches.has("--activate"),
|
||||
activate,
|
||||
confirmed: flags.switches.has("--yes"),
|
||||
...(follow === undefined ? {} : { follow }),
|
||||
};
|
||||
}
|
||||
if (action === "switch") {
|
||||
if (tokens.length !== 2 || (target !== "status" && target !== "cancel")) {
|
||||
return null;
|
||||
}
|
||||
return { kind: "switch", action: target };
|
||||
}
|
||||
if (action === "spaces") {
|
||||
const flags = readFlags(tokens.slice(1), ["--cleanup"], ["--yes"]);
|
||||
if (!flags) return null;
|
||||
const cleanup = flags.values.get("--cleanup");
|
||||
if (cleanup === undefined && flags.switches.has("--yes")) return null;
|
||||
if (cleanup !== undefined && !OPERATION_ID.test(cleanup)) return null;
|
||||
return {
|
||||
kind: "spaces",
|
||||
...(cleanup === undefined ? {} : { cleanup }),
|
||||
confirmed: flags.switches.has("--yes"),
|
||||
};
|
||||
}
|
||||
if (target === undefined || !OPERATION_ID.test(target)) return null;
|
||||
@@ -5154,8 +5502,21 @@ function parseModelsCommand(rest: readonly string[]): ModelsCommand | null {
|
||||
};
|
||||
}
|
||||
if (action === "activate") {
|
||||
const flags = readFlags(tokens.slice(2), [], ["--confirm", "--yes"]);
|
||||
return flags ? { kind: "activate", profileId: target } : null;
|
||||
const flags = readFlags(
|
||||
tokens.slice(2),
|
||||
[],
|
||||
["--confirm", "--yes", "--wait", "--no-wait"],
|
||||
);
|
||||
if (!flags) return null;
|
||||
const follow = followFlag(flags.switches);
|
||||
if (follow === null) return null;
|
||||
return {
|
||||
kind: "activate",
|
||||
profileId: target,
|
||||
// `--confirm` stays the hidden alias of `--yes`.
|
||||
confirmed: flags.switches.has("--yes") || flags.switches.has("--confirm"),
|
||||
...(follow === undefined ? {} : { follow }),
|
||||
};
|
||||
}
|
||||
if (action === "test") {
|
||||
return tokens.length === 2 ? { kind: "test", profileId: target } : null;
|
||||
@@ -5178,6 +5539,35 @@ function parseModelsCommand(rest: readonly string[]): ModelsCommand | null {
|
||||
return null;
|
||||
}
|
||||
|
||||
/** `--wait` / `--no-wait`: undefined when neither, null when both. */
|
||||
function followFlag(switches: ReadonlySet<string>): boolean | undefined | null {
|
||||
const wait = switches.has("--wait");
|
||||
const noWait = switches.has("--no-wait");
|
||||
if (wait && noWait) return null;
|
||||
return wait ? true : noWait ? false : undefined;
|
||||
}
|
||||
|
||||
/** Start a probe and read its receipt until it settles (about two minutes). */
|
||||
async function probeUntilSettled(
|
||||
reader: CoreReader,
|
||||
operations: CoreOperations,
|
||||
profileId: string,
|
||||
wait?: (milliseconds: number, signal: AbortSignal) => Promise<void>,
|
||||
): Promise<ModelProbeFlowDocument> {
|
||||
let document = await startModelProbe(reader, operations, profileId);
|
||||
const pause =
|
||||
wait ??
|
||||
((milliseconds: number) =>
|
||||
new Promise<void>((resolve) => setTimeout(resolve, milliseconds)));
|
||||
const signal = new AbortController().signal;
|
||||
for (let attempt = 0; attempt < 120; attempt += 1) {
|
||||
if (document.state !== "pending" || !document.operation_id) break;
|
||||
await pause(1_000, signal);
|
||||
document = await readModelProbeReceipt(reader, document.operation_id);
|
||||
}
|
||||
return document;
|
||||
}
|
||||
|
||||
type SourcesCommand =
|
||||
| {
|
||||
kind: "ids";
|
||||
@@ -6705,7 +7095,13 @@ const OWN_DATA_DIR_COMMANDS = new Set([
|
||||
/** Remove one `--data-dir <absolute path>` from argv; null when it is
|
||||
* repeated, missing its value, or not an absolute path. */
|
||||
/** Memory commands that honour the global `--user-id`. */
|
||||
const USER_ID_COMMANDS = new Set(["memories", "recall", "profile", "write"]);
|
||||
const USER_ID_COMMANDS = new Set([
|
||||
"memories",
|
||||
"recall",
|
||||
"profile",
|
||||
"write",
|
||||
"evolve",
|
||||
]);
|
||||
|
||||
/**
|
||||
* Take `--user-id <id>` out of the arguments. The default user (`local`) is
|
||||
|
||||
@@ -361,9 +361,9 @@ export const COMMANDS: readonly CommandDescriptor[] = [
|
||||
{
|
||||
name: "models",
|
||||
summary:
|
||||
"管理生成模型配置。add/test 直接提交并返回操作 ID(异步,用 receipt 查结果);edit/activate 直接执行;remove 不可逆,需 --yes。<profile-id>/<operation-id> 都是 UUID;脚本模式的凭据只从受控文件读取。",
|
||||
"管理生成模型配置。add/test 直接提交并返回操作 ID(异步,用 receipt 查结果);edit 直接执行;activate 换了向量模型(模型或维度)时先显示切换计划,需 --yes 确认后重新生成全部向量,switch 查看或取消、spaces 查看或清理旧向量数据;remove 不可逆,需 --yes。<profile-id>/<operation-id> 都是 UUID;脚本模式的凭据只从受控文件读取。",
|
||||
usage:
|
||||
"memhub models list [--json]\n memhub models show <profile-id> [--json]\n memhub models add --config-file <path> [--api-key-file <path>] [--activate] [--json]\n memhub models test <profile-id> [--json]\n memhub models receipt <operation-id> [--json]\n memhub models edit <profile-id> --config-file <path> [--api-key-file <path>] [--json]\n memhub models activate <profile-id> [--json]\n memhub models remove <profile-id> [--yes] [--json]",
|
||||
"memhub models list [--json]\n memhub models show <profile-id> [--json]\n memhub models add --config-file <path> [--api-key-file <path>] [--activate [--yes] [--wait|--no-wait]] [--json]\n memhub models test <profile-id> [--json]\n memhub models receipt <operation-id> [--json]\n memhub models edit <profile-id> --config-file <path> [--api-key-file <path>] [--json]\n memhub models activate <profile-id> [--yes] [--wait|--no-wait] [--json]\n memhub models switch <status|cancel> [--json]\n memhub models spaces [--cleanup <space-id> --yes] [--json]\n memhub models remove <profile-id> [--yes] [--json]",
|
||||
alternative_usage: "memhub models setup",
|
||||
options: [
|
||||
{
|
||||
@@ -377,7 +377,17 @@ export const COMMANDS: readonly CommandDescriptor[] = [
|
||||
},
|
||||
{
|
||||
flag: "--activate",
|
||||
summary: "add 成功后把新配置设为当前使用。",
|
||||
summary:
|
||||
"add 成功后先探测,再按 activate 的规则设为当前使用(向量模型变化时同样要确认切换)。",
|
||||
},
|
||||
{
|
||||
flag: "--wait | --no-wait",
|
||||
summary:
|
||||
"activate 开始向量模型切换后是否跟随进度;终端里默认跟随,Ctrl-C 只停止显示。",
|
||||
},
|
||||
{
|
||||
flag: "--cleanup <space-id>",
|
||||
summary: "spaces:立即删除一个已停用的向量空间(不可恢复,需 --yes)。",
|
||||
},
|
||||
DESTRUCTIVE_CONFIRM_OPTION,
|
||||
JSON_OPTION,
|
||||
@@ -434,9 +444,68 @@ export const COMMANDS: readonly CommandDescriptor[] = [
|
||||
},
|
||||
{
|
||||
name: "profile",
|
||||
summary: "读取 SDK 组装后的用户长期画像(L0 + 已启用的长期画像补充)。",
|
||||
usage: "memhub profile [--user-id <id>] [--json]",
|
||||
options: [USER_ID_OPTION, JSON_OPTION],
|
||||
summary:
|
||||
"读取 SDK 组装后的用户长期画像(L0 + 已启用的长期画像补充);--versions 列出 L0 画像的版本链(最新在前,* 为当前版本)。",
|
||||
usage: "memhub profile [--versions] [--user-id <id>] [--json]",
|
||||
options: [
|
||||
{
|
||||
flag: "--versions",
|
||||
summary:
|
||||
"列出 L0 基本信息的版本链(get_user_basic_info),用于找出覆盖了正确值的错误版本。",
|
||||
},
|
||||
USER_ID_OPTION,
|
||||
JSON_OPTION,
|
||||
],
|
||||
exit_codes: RUNTIME_EXITS,
|
||||
},
|
||||
{
|
||||
name: "settings",
|
||||
summary:
|
||||
"查看或切换记忆模式(lite / pro / ultra)。lite 不调用模型,只在原始对话里检索;pro 由模型提取记忆(System 1);ultra 在 pro 之上每晚自动进化(System 2,使用知识图谱)。任何模式都保留全部原始对话。切换不迁移数据:已处理的不会重做,排队中的和新的按新模式处理;需 --yes,不加时只显示计划。环境变量 ONE_MEMORY_GRAPH_STORE 会强制 ultra。settings features 查看或开关 SDK 的可选检索与提取功能(全局设置,默认全部保持 SDK 默认值):graph_rerank 图谱重排(仅 ultra)、same_series_timeline 同系列时间线(需 graph_rerank)、episode_recall 片段召回(pro / ultra,需先开 session_episodes 写入片段)、session_episodes 片段写入、episode_summary 片段摘要(每个片段一次模型调用)、intention 意图提取(默认开)、few_shot 提取示例;set 需 --yes,不加时只显示计划;reset 恢复默认值。",
|
||||
usage:
|
||||
"memhub settings mode [lite|pro|ultra] [--yes] [--json]\n" +
|
||||
" memhub settings features [--json]\n" +
|
||||
" memhub settings features set <name>=<true|false> [...] [--yes] [--json]\n" +
|
||||
" memhub settings features reset [<name> ...] [--json]",
|
||||
options: [
|
||||
{
|
||||
flag: "mode [lite|pro|ultra]",
|
||||
summary: "不带模式只显示当前模式;带模式时先显示切换计划。",
|
||||
},
|
||||
{
|
||||
flag: "features set <name>=<true|false>",
|
||||
summary:
|
||||
"开关一个或多个 SDK 功能;功能在当前模式下不生效时会注明原因(例如 graph_rerank 需要 ultra)。",
|
||||
},
|
||||
{
|
||||
flag: "features reset [<name> ...]",
|
||||
summary: "把指定功能(不写则全部)恢复为 SDK 默认值。",
|
||||
},
|
||||
CONFIRM_OPTION,
|
||||
JSON_OPTION,
|
||||
],
|
||||
exit_codes: RUNTIME_EXITS,
|
||||
},
|
||||
{
|
||||
name: "evolve",
|
||||
summary:
|
||||
"立即为一个用户运行记忆进化(System 2,仅 ultra 模式)。在后台运行并返回运行 ID;--wait 等待结果。同一用户同时只运行一次;每晚的自动进化由 Runtime 按设置安排。",
|
||||
usage:
|
||||
"memhub evolve [--user-id <id>] [--agent-id <id>] [--max-batches <n>] [--wait] [--json]",
|
||||
options: [
|
||||
USER_ID_OPTION,
|
||||
{
|
||||
flag: "--agent-id <id>",
|
||||
summary:
|
||||
"进化哪个 agent 的记忆,默认 all-apps(MemHub 写入记忆用的 agent)。",
|
||||
},
|
||||
{
|
||||
flag: "--max-batches <n>",
|
||||
summary: "每轮最多处理的批次数(1–1000),默认由 SDK 决定。",
|
||||
},
|
||||
{ flag: "--wait", summary: "等待这次进化结束再返回。" },
|
||||
JSON_OPTION,
|
||||
],
|
||||
exit_codes: RUNTIME_EXITS,
|
||||
},
|
||||
{
|
||||
|
||||
@@ -76,6 +76,8 @@ const COMMANDS = [
|
||||
"models",
|
||||
"memories",
|
||||
"profile",
|
||||
"settings",
|
||||
"evolve",
|
||||
"graph-width",
|
||||
"host",
|
||||
"recall",
|
||||
@@ -633,6 +635,8 @@ describe("documented commands", () => {
|
||||
"receipt",
|
||||
"edit",
|
||||
"activate",
|
||||
"switch",
|
||||
"spaces",
|
||||
"remove",
|
||||
],
|
||||
// Source processing is one Freeze/Thaw switch under `sync`; the obsolete
|
||||
|
||||
@@ -97,7 +97,12 @@ export function contextExitCode(document: ContextSearchDocument): number {
|
||||
}
|
||||
|
||||
export function renderContextText(document: ContextSearchDocument): string {
|
||||
if (document.state === "unavailable") return "暂时无法查记忆,请稍后重试。\n";
|
||||
if (document.state === "unavailable") {
|
||||
// No lexical fallback: without the embedding service nothing is searched.
|
||||
return document.code === "embedding_unavailable"
|
||||
? "向量服务不可用,无法检索;请检查向量模型配置。\n"
|
||||
: "暂时无法查记忆,请稍后重试。\n";
|
||||
}
|
||||
if (document.state === "denied") return "当前应用没有查记忆权限。\n";
|
||||
if (document.state === "invalid") return "查记忆返回了无法识别的结果。\n";
|
||||
if (document.state === "rejected") {
|
||||
|
||||
@@ -8,6 +8,9 @@ import axios from "axios";
|
||||
*/
|
||||
const READ_ONLY_PATHS = [
|
||||
"/api/v1/call-events",
|
||||
// Read-only: the registry plus a scan of the vector store's collection names.
|
||||
"/api/v1/embedding-spaces",
|
||||
"/api/v1/embedding-switches/current",
|
||||
"/api/v1/import-jobs",
|
||||
"/api/v1/live-sync",
|
||||
"/api/v1/memory-processing",
|
||||
@@ -17,7 +20,14 @@ const READ_ONLY_PATHS = [
|
||||
"/api/v2/home-summary",
|
||||
"/api/v2/memory-library",
|
||||
"/api/v2/user-profile",
|
||||
"/api/v2/user-profile/versions",
|
||||
"/api/v1/settings/memory-mode",
|
||||
"/api/v1/settings/memory-features",
|
||||
"/api/v2/evolve",
|
||||
] as const;
|
||||
/** One evolve run's receipt (`memhub evolve --wait`). */
|
||||
const EVOLVE_RUN =
|
||||
/^\/api\/v2\/evolve\/[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$/;
|
||||
const SDK_MEMORY_DETAIL = /^\/api\/v2\/memories\/[A-Za-z0-9][\w.-]{0,127}$/;
|
||||
/** Status of one `POST /api/v2/memories` addition, by its request id. */
|
||||
const MEMORY_ADDITION =
|
||||
@@ -156,10 +166,23 @@ export function createCoreReader(attachment: {
|
||||
code: detailCode(response.data, response.status),
|
||||
};
|
||||
}
|
||||
return { state: "unavailable" as const, code: `http_${response.status}` };
|
||||
return {
|
||||
state: "unavailable" as const,
|
||||
code: unavailableCode(response.data, response.status),
|
||||
};
|
||||
};
|
||||
}
|
||||
|
||||
/** A 503 that names `embedding_unavailable` keeps it: the search could not
|
||||
* embed its query, which needs another fix than "try again later". */
|
||||
function unavailableCode(value: unknown, status: number): string {
|
||||
return status === 503 && detailCode(value, status) === EMBEDDING_UNAVAILABLE
|
||||
? EMBEDDING_UNAVAILABLE
|
||||
: `http_${status}`;
|
||||
}
|
||||
|
||||
export const EMBEDDING_UNAVAILABLE = "embedding_unavailable";
|
||||
|
||||
/** Core's stable `detail` code when it sent one, else the bare status. */
|
||||
function detailCode(value: unknown, status: number): string {
|
||||
const detail =
|
||||
@@ -211,7 +234,10 @@ export function createCoreContextQuery(attachment: {
|
||||
if (response.status === 404) {
|
||||
return { state: "not_found" as const, code: "http_404" };
|
||||
}
|
||||
return { state: "unavailable" as const, code: `http_${response.status}` };
|
||||
return {
|
||||
state: "unavailable" as const,
|
||||
code: unavailableCode(response.data, response.status),
|
||||
};
|
||||
};
|
||||
}
|
||||
|
||||
@@ -220,6 +246,7 @@ export function isReadOnlyPath(path: string): boolean {
|
||||
if (IMPORT_JOB.test(path)) return true;
|
||||
if (OPERATION_RECEIPT.test(path)) return true;
|
||||
if (MEMORY_ADDITION.test(path)) return true;
|
||||
if (EVOLVE_RUN.test(path)) return true;
|
||||
if (!SDK_MEMORY_DETAIL.test(path)) return false;
|
||||
const segment = path.slice("/api/v2/memories/".length);
|
||||
return !(RESERVED_MEMORY_SEGMENTS as readonly string[]).includes(segment);
|
||||
|
||||
@@ -26,12 +26,32 @@ const SOURCE_IMPORT_JOB_PATH = "/api/v1/import-jobs";
|
||||
const IMPORT_JOB_ACTION_PATH =
|
||||
/^\/api\/v1\/import-jobs\/[A-Za-z0-9][\w.-]{0,127}\/(?:pause|resume|cancel)$/;
|
||||
const GRAPH_WIDTH_PATH = "/api/v1/graph-width";
|
||||
/** `memhub evolve` starts a run; `memhub settings mode <m> --yes` sets the mode. */
|
||||
const EVOLVE_PATH = "/api/v2/evolve";
|
||||
const MEMORY_MODE_PATH = "/api/v1/settings/memory-mode";
|
||||
/** `memhub settings features set|reset` sets the SDK features. */
|
||||
const MEMORY_FEATURES_PATH = "/api/v1/settings/memory-features";
|
||||
const ONBOARDING_COMPLETE_PATH = "/api/v1/onboarding/complete";
|
||||
const EMBEDDING_SWITCH_PATHS = [
|
||||
"/api/v1/embedding-switches",
|
||||
"/api/v1/embedding-switches/plan",
|
||||
] as const;
|
||||
const EMBEDDING_SWITCH_CANCEL =
|
||||
/^\/api\/v1\/embedding-switches\/[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}\/cancel$/;
|
||||
const EMBEDDING_SPACE_PATH =
|
||||
/^\/api\/v1\/embedding-spaces\/[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$/;
|
||||
const MAX_BYTES = 4 * 1024 * 1024;
|
||||
|
||||
export type OperationOutcome<T> =
|
||||
| { state: "committed"; data: T }
|
||||
| { state: "rejected"; code: string; http_status: number }
|
||||
| {
|
||||
state: "rejected";
|
||||
code: string;
|
||||
http_status: number;
|
||||
/** Core's refusal body, when it carries more than the code (a switch
|
||||
* plan rides along with `embedding_switch_required`). */
|
||||
body?: unknown;
|
||||
}
|
||||
| { state: "unknown"; code: string };
|
||||
|
||||
export interface CoreOperations {
|
||||
@@ -97,7 +117,10 @@ export function createCoreOperations(attachment: {
|
||||
!SOURCE_AUTHORIZE_PATH.test(path) &&
|
||||
!SOURCE_REVOKE_PATH.test(path) &&
|
||||
path !== GRAPH_WIDTH_PATH &&
|
||||
path !== EVOLVE_PATH &&
|
||||
path !== ONBOARDING_COMPLETE_PATH &&
|
||||
!(EMBEDDING_SWITCH_PATHS as readonly string[]).includes(path) &&
|
||||
!EMBEDDING_SWITCH_CANCEL.test(path) &&
|
||||
!MODEL_PROFILE_ACTIVATE.test(path)
|
||||
) {
|
||||
throw new Error("core_operation_path_not_allowed");
|
||||
@@ -115,11 +138,18 @@ export function createCoreOperations(attachment: {
|
||||
if (response.status >= 200 && response.status < 300) {
|
||||
return { state: "committed", data: response.data as T };
|
||||
}
|
||||
if (response.status >= 400 && response.status < 500) {
|
||||
if (
|
||||
(response.status >= 400 && response.status < 500) ||
|
||||
// Core answers a refusal it can name with 503 (`sdk_reembed_unavailable`)
|
||||
// or 507 (`insufficient_disk_space`) on the switch routes.
|
||||
([503, 507].includes(response.status) &&
|
||||
(EMBEDDING_SWITCH_PATHS as readonly string[]).includes(path))
|
||||
) {
|
||||
return {
|
||||
state: "rejected",
|
||||
code: safeErrorCode(response.data, response.status),
|
||||
http_status: response.status,
|
||||
body: response.data,
|
||||
};
|
||||
}
|
||||
return {
|
||||
@@ -132,7 +162,11 @@ export function createCoreOperations(attachment: {
|
||||
body: Record<string, unknown>,
|
||||
signal?: AbortSignal,
|
||||
) {
|
||||
if (!SDK_MEMORY_PATH.test(path) || path.includes("?")) {
|
||||
if (
|
||||
path !== MEMORY_MODE_PATH &&
|
||||
path !== MEMORY_FEATURES_PATH &&
|
||||
(!SDK_MEMORY_PATH.test(path) || path.includes("?"))
|
||||
) {
|
||||
throw new Error("core_operation_path_not_allowed");
|
||||
}
|
||||
let response: { status: number; data: unknown };
|
||||
@@ -194,7 +228,11 @@ export function createCoreOperations(attachment: {
|
||||
};
|
||||
},
|
||||
async delete<T>(path: string, signal?: AbortSignal) {
|
||||
if (!SDK_MEMORY_PATH.test(path) && !MODEL_PROFILE_PATH.test(path)) {
|
||||
if (
|
||||
!SDK_MEMORY_PATH.test(path) &&
|
||||
!MODEL_PROFILE_PATH.test(path) &&
|
||||
!EMBEDDING_SPACE_PATH.test(path)
|
||||
) {
|
||||
throw new Error("core_operation_path_not_allowed");
|
||||
}
|
||||
let response: { status: number; data: unknown };
|
||||
|
||||
@@ -0,0 +1,488 @@
|
||||
import { beforeEach, describe, expect, it, vi } from "vitest";
|
||||
|
||||
import { run } from "./cli.js";
|
||||
import { renderContextText } from "./context-search.js";
|
||||
import { describeEmbeddingSwitchSummary } from "./embedding-switch-flow.js";
|
||||
import { composeStatusReport, renderStatusText } from "./status-report.js";
|
||||
|
||||
const transport = vi.hoisted(() => ({
|
||||
get: vi.fn(),
|
||||
post: vi.fn(),
|
||||
put: vi.fn(),
|
||||
patch: vi.fn(),
|
||||
delete: vi.fn(),
|
||||
}));
|
||||
vi.mock("axios", () => ({
|
||||
default: { create: vi.fn(() => transport), get: vi.fn() },
|
||||
}));
|
||||
|
||||
const runtime = vi.hoisted(() => ({
|
||||
inspectRuntimeStatus: vi.fn(),
|
||||
discoverRuntimeAttachment: vi.fn(),
|
||||
}));
|
||||
vi.mock("@one-memory/runtime-control", async (importOriginal) => ({
|
||||
...(await importOriginal<typeof import("@one-memory/runtime-control")>()),
|
||||
...runtime,
|
||||
}));
|
||||
|
||||
const PROFILE = "44444444-4444-4444-8444-444444444444";
|
||||
const JOB = "55555555-5555-4555-8555-555555555555";
|
||||
const SPACE = "66666666-6666-4666-8666-666666666666";
|
||||
|
||||
function plan(overrides: Record<string, unknown> = {}) {
|
||||
return {
|
||||
needed: true,
|
||||
reason: "embedding_identity_changed",
|
||||
plan_token: "token-1",
|
||||
source: {
|
||||
space_id: "77777777-7777-4777-8777-777777777777",
|
||||
model: "bge-small-zh-v1.5",
|
||||
dims: 512,
|
||||
collection_base: "memhub_aaaaaaaaaaaaaaaa",
|
||||
vectors: 12480,
|
||||
},
|
||||
target: {
|
||||
space_id: null,
|
||||
model: "text-embedding-v3",
|
||||
dims: 1024,
|
||||
collection_base: "",
|
||||
},
|
||||
target_profile_id: PROFILE,
|
||||
target_revision: 5,
|
||||
target_provider_kind: "openai_compatible",
|
||||
target_local: false,
|
||||
sources: [],
|
||||
total_vectors: 12480,
|
||||
users: 2,
|
||||
graph_vectors: null,
|
||||
estimate_s: 120,
|
||||
requests: 195,
|
||||
disk_bytes_needed: 70 * 1024 * 1024,
|
||||
disk_bytes_free: 10 * 1024 * 1024 * 1024,
|
||||
retention_days: 14,
|
||||
blockers: [],
|
||||
...overrides,
|
||||
};
|
||||
}
|
||||
|
||||
function job(
|
||||
state: string,
|
||||
percent: number,
|
||||
extra: Record<string, unknown> = {},
|
||||
) {
|
||||
return {
|
||||
job_id: JOB,
|
||||
kind: "switch",
|
||||
state,
|
||||
source: plan().source,
|
||||
target: { ...plan().target, space_id: SPACE },
|
||||
stray_space_ids: [],
|
||||
target_profile_id: PROFILE,
|
||||
target_revision: 5,
|
||||
requested_by: "cli",
|
||||
progress: {
|
||||
phase: state === "reembedding" ? "vectors" : "prepare",
|
||||
done: Math.round((12480 * percent) / 100),
|
||||
total: 12480,
|
||||
percent,
|
||||
eta_s: 60,
|
||||
},
|
||||
cancel_requested: false,
|
||||
error_code: null,
|
||||
attempt: 1,
|
||||
created_at: "T1",
|
||||
updated_at: "T1",
|
||||
finished_at: state === "done" ? "T2" : null,
|
||||
...extra,
|
||||
};
|
||||
}
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
for (const method of Object.values(transport)) method.mockReset();
|
||||
runtime.inspectRuntimeStatus.mockResolvedValue({
|
||||
state: "running",
|
||||
code: "runtime_verified",
|
||||
instanceId: "instance-1",
|
||||
host: "127.0.0.1",
|
||||
port: 8765,
|
||||
version: "1.2.3",
|
||||
});
|
||||
runtime.discoverRuntimeAttachment.mockResolvedValue({
|
||||
instanceId: "instance-1",
|
||||
apiToken: "synthetic-runtime-token",
|
||||
host: "127.0.0.1",
|
||||
port: 8765,
|
||||
servicePid: 4242,
|
||||
health: {
|
||||
status: "ok",
|
||||
service: "one-memory-core",
|
||||
version: "1.2.3",
|
||||
instance_id: "instance-1",
|
||||
},
|
||||
});
|
||||
});
|
||||
|
||||
describe("models activate with an embedding change", () => {
|
||||
it("prints the plan and exits 2 without --yes on a non-terminal", async () => {
|
||||
transport.post.mockImplementation(async (url: string) =>
|
||||
url === "/api/v1/embedding-switches/plan"
|
||||
? { status: 200, data: plan() }
|
||||
: { status: 404, data: {} },
|
||||
);
|
||||
const result = await capturedRun(["models", "activate", PROFILE]);
|
||||
expect(result.code).toBe(2);
|
||||
expect(result.stdout).toContain("将更换向量模型");
|
||||
expect(result.stdout).toContain("bge-small-zh-v1.5(512 维)");
|
||||
expect(result.stdout).toContain("text-embedding-v3(1024 维)");
|
||||
expect(result.stdout).toContain("12,480");
|
||||
expect(result.stdout).toContain("重新运行并加上 --yes 以确认");
|
||||
expect(transport.post.mock.calls.map(([url]) => url)).toEqual([
|
||||
"/api/v1/embedding-switches/plan",
|
||||
]);
|
||||
});
|
||||
|
||||
it("starts the switch with --yes and emits JSON progress lines", async () => {
|
||||
transport.post.mockImplementation(async (url: string) => {
|
||||
if (url === "/api/v1/embedding-switches/plan") {
|
||||
return { status: 200, data: plan() };
|
||||
}
|
||||
if (url === "/api/v1/embedding-switches") {
|
||||
return { status: 202, data: job("preparing", 0) };
|
||||
}
|
||||
return { status: 404, data: {} };
|
||||
});
|
||||
const answers = [
|
||||
{
|
||||
phase: "running",
|
||||
writes_paused: true,
|
||||
job: job("reembedding", 50),
|
||||
last_finished: null,
|
||||
},
|
||||
{
|
||||
phase: "idle",
|
||||
writes_paused: false,
|
||||
job: null,
|
||||
last_finished: job("done", 100),
|
||||
},
|
||||
];
|
||||
transport.get.mockImplementation(async (url: string) =>
|
||||
url === "/api/v1/embedding-switches/current"
|
||||
? { status: 200, data: { schema_version: 1, ...answers.shift() } }
|
||||
: { status: 404, data: {} },
|
||||
);
|
||||
const result = await capturedRun(
|
||||
["models", "activate", PROFILE, "--yes", "--wait", "--json"],
|
||||
{ watchWait: async () => undefined },
|
||||
);
|
||||
expect(result.code).toBe(0);
|
||||
const lines = result.stdout
|
||||
.trim()
|
||||
.split("\n")
|
||||
.map((line) => JSON.parse(line));
|
||||
expect(lines.map((line) => line.state)).toEqual([
|
||||
"preparing",
|
||||
"reembedding",
|
||||
"done",
|
||||
"done",
|
||||
]);
|
||||
expect(lines[1]).toMatchObject({ percent: 50, job_id: JOB });
|
||||
expect(lines.at(-1)).toMatchObject({ result: "switched" });
|
||||
expect(transport.post).toHaveBeenCalledWith(
|
||||
"/api/v1/embedding-switches",
|
||||
{
|
||||
profile_id: PROFILE,
|
||||
expected_revision: 5,
|
||||
plan_token: "token-1",
|
||||
requested_by: "cli",
|
||||
},
|
||||
undefined,
|
||||
);
|
||||
});
|
||||
|
||||
it("refuses a plan with blockers and starts nothing", async () => {
|
||||
transport.post.mockResolvedValue({
|
||||
status: 200,
|
||||
data: plan({ blockers: ["sdk_reembed_unavailable"] }),
|
||||
});
|
||||
const result = await capturedRun([
|
||||
"models",
|
||||
"activate",
|
||||
PROFILE,
|
||||
"--yes",
|
||||
"--json",
|
||||
]);
|
||||
expect(result.code).toBe(1);
|
||||
expect(JSON.parse(result.stdout)).toMatchObject({
|
||||
state: "blocked",
|
||||
code: "sdk_reembed_unavailable",
|
||||
});
|
||||
expect(transport.post).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it("says nothing is re-embedded when only generation changed", async () => {
|
||||
transport.post.mockImplementation(async (url: string) =>
|
||||
url === "/api/v1/embedding-switches/plan"
|
||||
? {
|
||||
status: 200,
|
||||
data: plan({ needed: false, reason: "same_embedding" }),
|
||||
}
|
||||
: { status: 200, data: profileResponse() },
|
||||
);
|
||||
const result = await capturedRun(["models", "activate", PROFILE]);
|
||||
expect(result.code).toBe(0);
|
||||
expect(result.stdout).toContain("向量模型未变化,无需重新生成向量。");
|
||||
});
|
||||
|
||||
it("reads the plan carried by a 409 embedding_switch_required", async () => {
|
||||
transport.post.mockImplementation(async (url: string) =>
|
||||
url.endsWith("/activate")
|
||||
? {
|
||||
status: 409,
|
||||
data: { detail: "embedding_switch_required", plan: plan() },
|
||||
}
|
||||
: { status: 404, data: { detail: "Not Found" } },
|
||||
);
|
||||
const result = await capturedRun(["models", "activate", PROFILE, "--json"]);
|
||||
expect(result.code).toBe(2);
|
||||
expect(JSON.parse(result.stdout)).toMatchObject({
|
||||
state: "confirmation_required",
|
||||
code: "embedding_switch_required",
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe("models switch and spaces", () => {
|
||||
it("reports no switch, then cancels the running one", async () => {
|
||||
transport.get.mockResolvedValueOnce({
|
||||
status: 200,
|
||||
data: {
|
||||
schema_version: 1,
|
||||
phase: "idle",
|
||||
writes_paused: false,
|
||||
job: null,
|
||||
last_finished: null,
|
||||
},
|
||||
});
|
||||
const idle = await capturedRun(["models", "switch", "status"]);
|
||||
expect(idle.code).toBe(0);
|
||||
expect(idle.stdout).toContain("没有进行中的向量模型切换");
|
||||
|
||||
transport.get.mockResolvedValueOnce({
|
||||
status: 200,
|
||||
data: {
|
||||
schema_version: 1,
|
||||
phase: "running",
|
||||
writes_paused: true,
|
||||
job: job("reembedding", 20),
|
||||
last_finished: null,
|
||||
},
|
||||
});
|
||||
transport.post.mockResolvedValue({
|
||||
status: 200,
|
||||
data: job("reembedding", 20, { cancel_requested: true }),
|
||||
});
|
||||
const cancelled = await capturedRun([
|
||||
"models",
|
||||
"switch",
|
||||
"cancel",
|
||||
"--json",
|
||||
]);
|
||||
expect(cancelled.code).toBe(0);
|
||||
expect(transport.post).toHaveBeenCalledWith(
|
||||
`/api/v1/embedding-switches/${JOB}/cancel`,
|
||||
{},
|
||||
undefined,
|
||||
);
|
||||
expect(JSON.parse(cancelled.stdout)).toMatchObject({
|
||||
cancel_requested: true,
|
||||
});
|
||||
});
|
||||
|
||||
it("lists spaces and needs --yes before cleaning one up", async () => {
|
||||
transport.get.mockResolvedValue({
|
||||
status: 200,
|
||||
data: {
|
||||
schema_version: 1,
|
||||
live_space_id: "x",
|
||||
spaces: [
|
||||
{
|
||||
space_id: SPACE,
|
||||
state: "retired",
|
||||
model: "bge-small-zh-v1.5",
|
||||
dims: 512,
|
||||
vectors: 10,
|
||||
keep_until: "2026-11-15T00:00:00+00:00",
|
||||
},
|
||||
{
|
||||
space_id: "88888888-8888-4888-8888-888888888888",
|
||||
state: "dropped",
|
||||
model: "m",
|
||||
dims: 8,
|
||||
},
|
||||
],
|
||||
scanned: [],
|
||||
scan_error: null,
|
||||
},
|
||||
});
|
||||
const listed = await capturedRun(["models", "spaces"]);
|
||||
expect(listed.code).toBe(0);
|
||||
expect(listed.stdout).toContain("已停用");
|
||||
expect(listed.stdout).toContain("保留到 2026-11-15");
|
||||
expect(listed.stdout).not.toContain("88888888");
|
||||
|
||||
const preview = await capturedRun(["models", "spaces", "--cleanup", SPACE]);
|
||||
expect(preview.code).toBe(0);
|
||||
expect(preview.stdout).toContain("--yes");
|
||||
expect(transport.delete).not.toHaveBeenCalled();
|
||||
|
||||
transport.delete.mockResolvedValue({
|
||||
status: 200,
|
||||
data: { schema_version: 1 },
|
||||
});
|
||||
const dropped = await capturedRun([
|
||||
"models",
|
||||
"spaces",
|
||||
"--cleanup",
|
||||
SPACE,
|
||||
"--yes",
|
||||
]);
|
||||
expect(dropped.code).toBe(0);
|
||||
expect(transport.delete).toHaveBeenCalledWith(
|
||||
`/api/v1/embedding-spaces/${SPACE}`,
|
||||
undefined,
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe("status and search wording", () => {
|
||||
it("shows the switch line in memhub status", async () => {
|
||||
const reader = async (path: string) =>
|
||||
path === "/api/v1/memory-processing"
|
||||
? {
|
||||
state: "ok" as const,
|
||||
data: {
|
||||
pending: 3,
|
||||
fresh_pending: 1,
|
||||
backlog_pending: 2,
|
||||
outside_window_pending: 0,
|
||||
running: 0,
|
||||
waiting_for_model: 0,
|
||||
completed: 9,
|
||||
failed: 0,
|
||||
embedding_switch: {
|
||||
job_id: JOB,
|
||||
kind: "switch",
|
||||
state: "reembedding",
|
||||
percent: 50,
|
||||
eta_s: 60,
|
||||
},
|
||||
},
|
||||
}
|
||||
: { state: "unavailable" as const, code: "http_503" };
|
||||
const report = await composeStatusReport({
|
||||
dataDir: "/synthetic/data",
|
||||
runtime: {
|
||||
state: "running",
|
||||
code: "runtime_verified",
|
||||
instanceId: "instance-1",
|
||||
host: "127.0.0.1",
|
||||
port: 8765,
|
||||
} as never,
|
||||
background: null,
|
||||
reader: reader as never,
|
||||
});
|
||||
const text = renderStatusText(report, "MemHub");
|
||||
expect(text).toContain("向量模型:切换中 50%,记忆整理已暂停");
|
||||
expect(
|
||||
describeEmbeddingSwitchSummary({
|
||||
kind: "merge",
|
||||
state: "reembedding",
|
||||
percent: 30,
|
||||
}),
|
||||
).toBe("正在合并旧向量数据 30%,记忆整理已暂停");
|
||||
expect(
|
||||
describeEmbeddingSwitchSummary({
|
||||
kind: "merge",
|
||||
state: "scanning",
|
||||
percent: 0,
|
||||
}),
|
||||
).toBe("正在检查向量数据…");
|
||||
});
|
||||
|
||||
it("names an unavailable embedding service instead of 'try later'", () => {
|
||||
expect(
|
||||
renderContextText({
|
||||
schema_version: 2,
|
||||
product: "memhub",
|
||||
state: "unavailable",
|
||||
code: "embedding_unavailable",
|
||||
items: [],
|
||||
}),
|
||||
).toContain("向量服务不可用");
|
||||
});
|
||||
});
|
||||
|
||||
function profileResponse() {
|
||||
return {
|
||||
profile_id: PROFILE,
|
||||
name: "Synthetic",
|
||||
provider_kind: "openai_compatible",
|
||||
base_url: "https://example.invalid/v1",
|
||||
model: "synthetic-model-2",
|
||||
embedding_model: "text-embedding-v3",
|
||||
embedding_provider_kind: "openai_compatible",
|
||||
embedding_base_url: "https://example.invalid/v1",
|
||||
embedding_dimensions: 1024,
|
||||
has_api_key: true,
|
||||
credential_state: "available",
|
||||
timeout_seconds: 30,
|
||||
max_concurrency: 1,
|
||||
context_window: null,
|
||||
disable_thinking: true,
|
||||
capabilities: {
|
||||
endpoint_reachable: true,
|
||||
model_available: true,
|
||||
structured_output: true,
|
||||
embedding_available: true,
|
||||
},
|
||||
revision: 5,
|
||||
is_active: true,
|
||||
created_at: "2026-09-19T00:00:00Z",
|
||||
updated_at: "2026-09-19T00:00:00Z",
|
||||
};
|
||||
}
|
||||
|
||||
async function capturedRun(
|
||||
arguments_: string[],
|
||||
controls: Parameters<typeof run>[2] = {},
|
||||
) {
|
||||
let stdout = "";
|
||||
let stderr = "";
|
||||
const stdoutSpy = vi
|
||||
.spyOn(process.stdout, "write")
|
||||
.mockImplementation((chunk) => {
|
||||
stdout += String(chunk);
|
||||
return true;
|
||||
});
|
||||
const stderrSpy = vi
|
||||
.spyOn(process.stderr, "write")
|
||||
.mockImplementation((chunk) => {
|
||||
stderr += String(chunk);
|
||||
return true;
|
||||
});
|
||||
try {
|
||||
return {
|
||||
code: await run(
|
||||
arguments_,
|
||||
{ ONE_MEMORY_DATA_DIR: "/synthetic/data" },
|
||||
{ stdinIsTTY: false, ...controls },
|
||||
),
|
||||
stdout,
|
||||
stderr,
|
||||
};
|
||||
} finally {
|
||||
stdoutSpy.mockRestore();
|
||||
stderrSpy.mockRestore();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,847 @@
|
||||
/**
|
||||
* Switching the embedding model from the CLI: `models activate` when the
|
||||
* embedding changes, `models switch status|cancel`, `models spaces`.
|
||||
*
|
||||
* Core owns every decision (`/api/v1/embedding-switches*`, contract
|
||||
* `docs/contracts/embedding-switch-v1.md`). A change of embedding model or
|
||||
* width is never applied by activation alone: Core plans a switch, the user
|
||||
* confirms it, Core re-embeds every memory into a new space and only then
|
||||
* makes it live. The CLI shows the plan, asks, starts the job and follows it.
|
||||
*/
|
||||
|
||||
import type { CoreReader } from "./core-client.js";
|
||||
import type {
|
||||
CoreDeleteOperations,
|
||||
CoreOperations,
|
||||
OperationOutcome,
|
||||
} from "./core-operations.js";
|
||||
import { EXIT_FAILED, EXIT_OK, EXIT_USAGE } from "./exit-codes.js";
|
||||
import {
|
||||
activateModelProfile,
|
||||
type ModelWriteDocument,
|
||||
modelProfileResult,
|
||||
renderModelWriteText,
|
||||
} from "./model-write-flows.js";
|
||||
import { PRODUCT } from "./version.js";
|
||||
|
||||
const UUID = /^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$/;
|
||||
const CODE = /^[a-z0-9_.:-]{1,128}$/;
|
||||
const TERMINAL = new Set(["done", "rolled_back"]);
|
||||
const POLL_MS = 1_000;
|
||||
|
||||
export interface SpaceSummary {
|
||||
space_id: string | null;
|
||||
model: string | null;
|
||||
dims: number | null;
|
||||
collection_base: string;
|
||||
vectors?: number | null;
|
||||
}
|
||||
|
||||
export interface EmbeddingSwitchPlan {
|
||||
needed: boolean;
|
||||
reason: string;
|
||||
plan_token: string;
|
||||
source: SpaceSummary;
|
||||
target: SpaceSummary;
|
||||
target_profile_id: string;
|
||||
target_revision: number;
|
||||
target_provider_kind: string;
|
||||
target_local: boolean;
|
||||
total_vectors: number;
|
||||
users: number | null;
|
||||
graph_vectors: number | null;
|
||||
estimate_s: number | null;
|
||||
requests: number | null;
|
||||
disk_bytes_needed: number;
|
||||
disk_bytes_free: number | null;
|
||||
retention_days: number;
|
||||
blockers: string[];
|
||||
}
|
||||
|
||||
export interface SwitchProgress {
|
||||
phase: string;
|
||||
done: number;
|
||||
total: number;
|
||||
percent: number;
|
||||
eta_s: number | null;
|
||||
}
|
||||
|
||||
export interface EmbeddingSwitchJob {
|
||||
job_id: string;
|
||||
kind: "switch" | "merge";
|
||||
state: string;
|
||||
source: SpaceSummary;
|
||||
target: SpaceSummary;
|
||||
progress: SwitchProgress;
|
||||
cancel_requested: boolean;
|
||||
error_code: string | null;
|
||||
finished_at: string | null;
|
||||
}
|
||||
|
||||
export interface SwitchCurrent {
|
||||
phase: "idle" | "scanning" | "running";
|
||||
writes_paused: boolean;
|
||||
job: EmbeddingSwitchJob | null;
|
||||
last_finished: EmbeddingSwitchJob | null;
|
||||
}
|
||||
|
||||
export interface SwitchIo {
|
||||
reader: CoreReader;
|
||||
operations: CoreOperations & CoreDeleteOperations;
|
||||
emit: (text: string) => void;
|
||||
asJson: boolean;
|
||||
/** A terminal is attached: may prompt, and follows progress by default. */
|
||||
interactive: boolean;
|
||||
confirm?: (question: string) => Promise<boolean>;
|
||||
wait?: (milliseconds: number, signal: AbortSignal) => Promise<void>;
|
||||
signal?: AbortSignal;
|
||||
}
|
||||
|
||||
export interface ActivateOptions {
|
||||
confirmed: boolean;
|
||||
/** `--wait` / `--no-wait`; undefined follows on a terminal only. */
|
||||
follow?: boolean;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Parsing: bounded projections of Core's answers, never passthrough.
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
function record(value: unknown): Record<string, unknown> | null {
|
||||
return value && typeof value === "object" && !Array.isArray(value)
|
||||
? (value as Record<string, unknown>)
|
||||
: null;
|
||||
}
|
||||
|
||||
function int(value: unknown): number | null {
|
||||
return typeof value === "number" && Number.isSafeInteger(value) && value >= 0
|
||||
? value
|
||||
: null;
|
||||
}
|
||||
|
||||
function str(value: unknown, max = 256): string | null {
|
||||
return typeof value === "string" && value.length <= max ? value : null;
|
||||
}
|
||||
|
||||
function space(value: unknown): SpaceSummary | null {
|
||||
const item = record(value);
|
||||
if (!item) return null;
|
||||
return {
|
||||
space_id: str(item.space_id, 64),
|
||||
model: str(item.model),
|
||||
dims: int(item.dims),
|
||||
collection_base: str(item.collection_base) ?? "",
|
||||
vectors: int(item.vectors),
|
||||
};
|
||||
}
|
||||
|
||||
export function parsePlan(value: unknown): EmbeddingSwitchPlan | null {
|
||||
const plan = record(value);
|
||||
if (!plan || typeof plan.needed !== "boolean") return null;
|
||||
const source = space(plan.source);
|
||||
const target = space(plan.target);
|
||||
const token = str(plan.plan_token, 128);
|
||||
const profile = str(plan.target_profile_id, 128);
|
||||
const revision = int(plan.target_revision);
|
||||
if (!source || !target || !token || !profile || revision === null) {
|
||||
return null;
|
||||
}
|
||||
const blockers = Array.isArray(plan.blockers)
|
||||
? plan.blockers.filter(
|
||||
(item): item is string => typeof item === "string" && CODE.test(item),
|
||||
)
|
||||
: [];
|
||||
return {
|
||||
needed: plan.needed,
|
||||
reason: str(plan.reason, 64) ?? "",
|
||||
plan_token: token,
|
||||
source,
|
||||
target,
|
||||
target_profile_id: profile,
|
||||
target_revision: revision,
|
||||
target_provider_kind: str(plan.target_provider_kind, 64) ?? "",
|
||||
target_local: plan.target_local === true,
|
||||
total_vectors: int(plan.total_vectors) ?? 0,
|
||||
users: int(plan.users),
|
||||
graph_vectors: int(plan.graph_vectors),
|
||||
estimate_s: int(plan.estimate_s),
|
||||
requests: int(plan.requests),
|
||||
disk_bytes_needed: int(plan.disk_bytes_needed) ?? 0,
|
||||
disk_bytes_free: int(plan.disk_bytes_free),
|
||||
retention_days: int(plan.retention_days) ?? 14,
|
||||
blockers,
|
||||
};
|
||||
}
|
||||
|
||||
export function parseJob(value: unknown): EmbeddingSwitchJob | null {
|
||||
const job = record(value);
|
||||
if (!job) return null;
|
||||
const jobId = str(job.job_id, 64);
|
||||
const state = str(job.state, 32);
|
||||
const source = space(job.source);
|
||||
const target = space(job.target);
|
||||
const progress = record(job.progress) ?? {};
|
||||
if (!jobId || !state || !source || !target) return null;
|
||||
if (job.kind !== "switch" && job.kind !== "merge") return null;
|
||||
const percent =
|
||||
typeof progress.percent === "number" && Number.isFinite(progress.percent)
|
||||
? Math.min(100, Math.max(0, progress.percent))
|
||||
: 0;
|
||||
return {
|
||||
job_id: jobId,
|
||||
kind: job.kind,
|
||||
state,
|
||||
source,
|
||||
target,
|
||||
progress: {
|
||||
phase: str(progress.phase, 32) ?? "prepare",
|
||||
done: int(progress.done) ?? 0,
|
||||
total: int(progress.total) ?? 0,
|
||||
percent,
|
||||
eta_s: int(progress.eta_s),
|
||||
},
|
||||
cancel_requested: job.cancel_requested === true,
|
||||
error_code:
|
||||
typeof job.error_code === "string" && CODE.test(job.error_code)
|
||||
? job.error_code
|
||||
: null,
|
||||
finished_at: str(job.finished_at, 64),
|
||||
};
|
||||
}
|
||||
|
||||
export function parseCurrent(value: unknown): SwitchCurrent | null {
|
||||
const current = record(value);
|
||||
if (!current) return null;
|
||||
const phase = current.phase;
|
||||
if (phase !== "idle" && phase !== "scanning" && phase !== "running") {
|
||||
return null;
|
||||
}
|
||||
return {
|
||||
phase,
|
||||
writes_paused: current.writes_paused === true,
|
||||
job: current.job == null ? null : parseJob(current.job),
|
||||
last_finished:
|
||||
current.last_finished == null ? null : parseJob(current.last_finished),
|
||||
};
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Wording
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
const BLOCKER_TEXT: Record<string, string> = {
|
||||
embedding_switch_running: "上一次向量模型切换还在进行,请等待完成或先取消。",
|
||||
embedding_probe_required: "新配置还没有通过探测;先运行 memhub models test。",
|
||||
sdk_reembed_unavailable:
|
||||
"当前运行时的 SDK 不支持重新生成向量,请先升级 MemHub。",
|
||||
embedding_switch_requires_migration:
|
||||
"还有旧版本的向量数据没有迁移完,请先完成升级迁移。",
|
||||
embedding_dimensions_unknown:
|
||||
"新配置的向量维度未知;先探测或在配置里写明 embedding_dimensions。",
|
||||
insufficient_disk_space: "磁盘空间不足,无法容纳新旧两份向量数据。",
|
||||
};
|
||||
|
||||
const FAILURE_TEXT: Record<string, string> = {
|
||||
writes_did_not_drain: "等待正在进行的记忆整理结束超时",
|
||||
embedding_service_failed: "向量服务连续请求失败",
|
||||
embedding_dimensions_mismatch: "向量服务返回的维度与探测到的不一致",
|
||||
embedding_verification_failed: "新向量没有通过校验",
|
||||
target_profile_changed: "切换期间目标配置被修改",
|
||||
embedding_switch_cancelled: "已取消",
|
||||
model_secret_unavailable: "读不到向量服务的 API Key",
|
||||
secure_storage_unavailable: "系统钥匙串不可用",
|
||||
embedding_switch_clients_busy: "记忆图谱一直被占用,无法重新生成",
|
||||
};
|
||||
|
||||
const STATE_TEXT: Record<string, string> = {
|
||||
preparing: "准备中(等待记忆整理暂停)",
|
||||
reembedding: "重新生成向量",
|
||||
verifying: "校验新向量",
|
||||
committing: "正式切换",
|
||||
failed: "正在回滚",
|
||||
done: "已完成",
|
||||
rolled_back: "已回滚",
|
||||
};
|
||||
|
||||
export function blockerText(code: string): string {
|
||||
return BLOCKER_TEXT[code] ?? code;
|
||||
}
|
||||
|
||||
function spaceLabel(item: SpaceSummary, local?: boolean): string {
|
||||
const dims = item.dims === null ? "维度未知" : `${item.dims} 维`;
|
||||
return `${item.model ?? "未知模型"}(${dims}${local ? ",本地" : ""})`;
|
||||
}
|
||||
|
||||
function duration(seconds: number | null): string {
|
||||
if (seconds === null) return "未知";
|
||||
if (seconds < 60) return `${Math.max(1, seconds)} 秒`;
|
||||
return `${Math.round(seconds / 60)} 分钟`;
|
||||
}
|
||||
|
||||
function megabytes(bytes: number | null): string {
|
||||
return bytes === null ? "未知" : `${Math.ceil(bytes / (1024 * 1024))} MB`;
|
||||
}
|
||||
|
||||
export function renderPlanText(plan: EmbeddingSwitchPlan): string {
|
||||
const lines = [
|
||||
`将更换向量模型:${spaceLabel(plan.source)} → ${spaceLabel(plan.target, plan.target_local)}`,
|
||||
`需要用新模型重新生成 ${plan.total_vectors.toLocaleString("en-US")} 条向量` +
|
||||
`${plan.users === null ? "" : `(${plan.users} 个用户)`}` +
|
||||
`,预计约 ${duration(plan.estimate_s)}` +
|
||||
`${plan.requests === null ? ",在本机计算,不产生费用" : `,约 ${plan.requests} 次请求`}` +
|
||||
`,额外磁盘约 ${megabytes(plan.disk_bytes_needed)}。`,
|
||||
];
|
||||
if (plan.graph_vectors) {
|
||||
lines.push(`记忆图谱另有 ${plan.graph_vectors} 条向量一并重新生成。`);
|
||||
}
|
||||
if (
|
||||
plan.source.dims !== null &&
|
||||
plan.source.dims === plan.target.dims &&
|
||||
plan.source.model !== plan.target.model
|
||||
) {
|
||||
lines.push("新旧模型维度相同,但向量不能通用,仍需重新生成。");
|
||||
}
|
||||
lines.push(
|
||||
"切换期间记忆整理暂停,检索继续使用旧模型;完成并校验通过后才会生效。",
|
||||
`旧的向量数据保留 ${plan.retention_days} 天,可用 memhub models spaces 清理。`,
|
||||
);
|
||||
for (const blocker of plan.blockers) {
|
||||
lines.push(`无法开始:${blockerText(blocker)}`);
|
||||
}
|
||||
return `${lines.join("\n")}\n`;
|
||||
}
|
||||
|
||||
export function renderProgressLine(job: EmbeddingSwitchJob): string {
|
||||
const what = job.kind === "merge" ? "正在合并旧向量数据" : "正在切换向量模型";
|
||||
const state = STATE_TEXT[job.state] ?? job.state;
|
||||
if (job.state === "reembedding" && job.progress.total > 0) {
|
||||
const eta =
|
||||
job.progress.eta_s === null
|
||||
? ""
|
||||
: ` 预计还需 ${duration(job.progress.eta_s)}`;
|
||||
return `${what}:${state} ${job.progress.done.toLocaleString("en-US")}/${job.progress.total.toLocaleString("en-US")} ${Math.floor(job.progress.percent)}%${eta}`;
|
||||
}
|
||||
return `${what}:${state}`;
|
||||
}
|
||||
|
||||
export function renderOutcomeText(job: EmbeddingSwitchJob): string {
|
||||
if (job.state === "done") {
|
||||
return job.kind === "merge"
|
||||
? "旧向量数据已合并,记忆整理已恢复。\n"
|
||||
: `已切换到 ${job.target.model ?? "新模型"}。\n`;
|
||||
}
|
||||
const code = job.error_code ?? "embedding_switch_failed";
|
||||
if (code === "embedding_switch_cancelled") {
|
||||
return `已取消切换,继续使用 ${job.source.model ?? "旧模型"}。\n`;
|
||||
}
|
||||
return `切换未完成,已恢复使用 ${job.source.model ?? "旧模型"},数据未受影响。原因:${FAILURE_TEXT[code] ?? code}(${code})\n`;
|
||||
}
|
||||
|
||||
/** One line for `memhub status`, from `memory-processing.embedding_switch`. */
|
||||
export function describeEmbeddingSwitchSummary(value: unknown): string | null {
|
||||
const summary = record(value);
|
||||
if (!summary) return null;
|
||||
const percent =
|
||||
typeof summary.percent === "number" && Number.isFinite(summary.percent)
|
||||
? Math.floor(summary.percent)
|
||||
: 0;
|
||||
if (summary.state === "scanning") return "正在检查向量数据…";
|
||||
if (summary.kind === "merge") {
|
||||
return `正在合并旧向量数据 ${percent}%,记忆整理已暂停`;
|
||||
}
|
||||
if (summary.state === "failed") return "切换未完成,正在恢复旧向量模型";
|
||||
return `切换中 ${percent}%,记忆整理已暂停`;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Flows
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
function document(
|
||||
command: string,
|
||||
fields: Record<string, unknown>,
|
||||
): Record<string, unknown> {
|
||||
return { schema_version: 1, product: PRODUCT, command, ...fields };
|
||||
}
|
||||
|
||||
function emitDocument(io: SwitchIo, value: Record<string, unknown>): void {
|
||||
io.emit(`${JSON.stringify(value)}\n`);
|
||||
}
|
||||
|
||||
async function requestPlan(
|
||||
profileId: string,
|
||||
io: SwitchIo,
|
||||
): Promise<
|
||||
| { state: "plan"; plan: EmbeddingSwitchPlan }
|
||||
| { state: "unsupported" }
|
||||
| { state: "failed"; code: string }
|
||||
> {
|
||||
const outcome = await io.operations.post<unknown>(
|
||||
"/api/v1/embedding-switches/plan",
|
||||
{ profile_id: profileId },
|
||||
);
|
||||
if (outcome.state === "committed") {
|
||||
const plan = parsePlan(outcome.data);
|
||||
// Not a plan: activation itself still refuses an embedding change.
|
||||
return plan ? { state: "plan", plan } : { state: "unsupported" };
|
||||
}
|
||||
if (outcome.state === "rejected") {
|
||||
// A Core before the registry has no plan route: activation as before.
|
||||
if (outcome.http_status === 404 && outcome.code === "http_404") {
|
||||
return { state: "unsupported" };
|
||||
}
|
||||
const plan = parsePlan(record(outcome.body)?.plan);
|
||||
if (plan) return { state: "plan", plan };
|
||||
return { state: "failed", code: outcome.code };
|
||||
}
|
||||
return { state: "failed", code: outcome.code };
|
||||
}
|
||||
|
||||
function emitActivation(io: SwitchIo, result: ModelWriteDocument): number {
|
||||
io.emit(
|
||||
io.asJson ? `${JSON.stringify(result)}\n` : renderModelWriteText(result),
|
||||
);
|
||||
return result.state === "committed" ? EXIT_OK : EXIT_FAILED;
|
||||
}
|
||||
|
||||
/**
|
||||
* `memhub models activate <id> [--yes] [--wait|--no-wait] [--json]`.
|
||||
*
|
||||
* Asks Core for the plan first. No embedding change: activate as before.
|
||||
* An embedding change: show the plan, confirm (`--yes`, or a prompt on a
|
||||
* terminal; exit 2 otherwise), start the switch and follow it.
|
||||
*/
|
||||
export async function activateWithEmbeddingSwitch(
|
||||
profileId: string,
|
||||
options: ActivateOptions,
|
||||
io: SwitchIo,
|
||||
): Promise<number> {
|
||||
if (!UUID.test(profileId)) {
|
||||
return emitActivation(
|
||||
io,
|
||||
await activateModelProfile(profileId, io.operations),
|
||||
);
|
||||
}
|
||||
const planned = await requestPlan(profileId, io);
|
||||
if (planned.state === "unsupported") {
|
||||
const activated = await activateOrPlan(profileId, io);
|
||||
return activated.state === "plan"
|
||||
? confirmAndSwitch(activated.plan, options, io)
|
||||
: emitActivation(io, activated.result);
|
||||
}
|
||||
if (planned.state === "failed") {
|
||||
return reject(io, "models activate", planned.code);
|
||||
}
|
||||
const plan = planned.plan;
|
||||
if (!plan.needed) {
|
||||
if (!io.asJson && plan.reason === "same_embedding") {
|
||||
io.emit("向量模型未变化,无需重新生成向量。\n");
|
||||
}
|
||||
const activated = await activateOrPlan(profileId, io);
|
||||
if (activated.state === "plan") {
|
||||
return confirmAndSwitch(activated.plan, options, io);
|
||||
}
|
||||
return emitActivation(io, activated.result);
|
||||
}
|
||||
return confirmAndSwitch(plan, options, io);
|
||||
}
|
||||
|
||||
async function activateOrPlan(
|
||||
profileId: string,
|
||||
io: SwitchIo,
|
||||
): Promise<
|
||||
| { state: "plan"; plan: EmbeddingSwitchPlan }
|
||||
| { state: "activated"; result: ModelWriteDocument }
|
||||
> {
|
||||
const result: OperationOutcome<unknown> = await io.operations.post<unknown>(
|
||||
`/api/v1/model-profiles/${encodeURIComponent(profileId)}/activate`,
|
||||
{},
|
||||
);
|
||||
if (
|
||||
result.state === "rejected" &&
|
||||
result.code === "embedding_switch_required"
|
||||
) {
|
||||
const plan = parsePlan(record(result.body)?.plan);
|
||||
if (plan) return { state: "plan", plan };
|
||||
}
|
||||
return { state: "activated", result: modelProfileResult(profileId, result) };
|
||||
}
|
||||
|
||||
function reject(io: SwitchIo, command: string, code: string): number {
|
||||
if (io.asJson) {
|
||||
emitDocument(io, document(command, { state: "rejected", code }));
|
||||
} else {
|
||||
io.emit(`${command} 未完成:${blockerText(code)}(${code})\n`);
|
||||
}
|
||||
return EXIT_FAILED;
|
||||
}
|
||||
|
||||
async function confirmAndSwitch(
|
||||
plan: EmbeddingSwitchPlan,
|
||||
options: ActivateOptions,
|
||||
io: SwitchIo,
|
||||
): Promise<number> {
|
||||
if (plan.blockers.length > 0) {
|
||||
if (io.asJson) {
|
||||
emitDocument(
|
||||
io,
|
||||
document("models activate", {
|
||||
state: "blocked",
|
||||
code: plan.blockers[0],
|
||||
plan,
|
||||
}),
|
||||
);
|
||||
} else {
|
||||
io.emit(renderPlanText(plan));
|
||||
}
|
||||
return EXIT_FAILED;
|
||||
}
|
||||
if (!options.confirmed) {
|
||||
const ask = io.interactive && !io.asJson ? io.confirm : undefined;
|
||||
if (ask === undefined) {
|
||||
if (io.asJson) {
|
||||
emitDocument(
|
||||
io,
|
||||
document("models activate", {
|
||||
state: "confirmation_required",
|
||||
code: "embedding_switch_required",
|
||||
plan,
|
||||
next_step: `${PRODUCT} models activate ${plan.target_profile_id} --yes`,
|
||||
}),
|
||||
);
|
||||
} else {
|
||||
io.emit(`${renderPlanText(plan)}重新运行并加上 --yes 以确认\n`);
|
||||
}
|
||||
return EXIT_USAGE;
|
||||
}
|
||||
io.emit(renderPlanText(plan));
|
||||
if (!(await ask("继续?[y/N] "))) {
|
||||
io.emit("未开始切换,仍在使用原来的向量模型。\n");
|
||||
return EXIT_OK;
|
||||
}
|
||||
}
|
||||
const started = await io.operations.post<unknown>(
|
||||
"/api/v1/embedding-switches",
|
||||
{
|
||||
profile_id: plan.target_profile_id,
|
||||
expected_revision: plan.target_revision,
|
||||
plan_token: plan.plan_token,
|
||||
requested_by: "cli",
|
||||
},
|
||||
);
|
||||
if (started.state !== "committed") {
|
||||
return reject(io, "models activate", started.code);
|
||||
}
|
||||
const job = parseJob(started.data);
|
||||
if (!job) return reject(io, "models activate", "embedding_switch_invalid");
|
||||
const follow = options.follow ?? io.interactive;
|
||||
if (!follow) {
|
||||
if (io.asJson) {
|
||||
emitDocument(io, document("models activate", progressFields(job)));
|
||||
} else {
|
||||
io.emit(
|
||||
`已开始切换向量模型(${job.job_id})。用 memhub models switch status 查看进度。\n`,
|
||||
);
|
||||
}
|
||||
return EXIT_OK;
|
||||
}
|
||||
return followJob(job, io, "models activate");
|
||||
}
|
||||
|
||||
function progressFields(job: EmbeddingSwitchJob): Record<string, unknown> {
|
||||
return {
|
||||
state: job.state,
|
||||
job_id: job.job_id,
|
||||
kind: job.kind,
|
||||
phase: job.progress.phase,
|
||||
percent: job.progress.percent,
|
||||
done: job.progress.done,
|
||||
total: job.progress.total,
|
||||
eta_s: job.progress.eta_s,
|
||||
...(job.error_code ? { code: job.error_code } : {}),
|
||||
};
|
||||
}
|
||||
|
||||
async function readCurrent(io: SwitchIo): Promise<SwitchCurrent | null> {
|
||||
const outcome = await io.reader<unknown>(
|
||||
"/api/v1/embedding-switches/current",
|
||||
);
|
||||
return outcome.state === "ok" ? parseCurrent(outcome.data) : null;
|
||||
}
|
||||
|
||||
/** Poll `current` until the job ends or the user stops watching. */
|
||||
async function followJob(
|
||||
started: EmbeddingSwitchJob,
|
||||
io: SwitchIo,
|
||||
command: string,
|
||||
): Promise<number> {
|
||||
const wait = io.wait ?? defaultWait;
|
||||
const signal = io.signal ?? new AbortController().signal;
|
||||
let last = "";
|
||||
let job = started;
|
||||
const show = (value: EmbeddingSwitchJob) => {
|
||||
const line = io.asJson
|
||||
? JSON.stringify(document(command, progressFields(value)))
|
||||
: renderProgressLine(value);
|
||||
if (line !== last) {
|
||||
io.emit(`${line}\n`);
|
||||
last = line;
|
||||
}
|
||||
};
|
||||
show(job);
|
||||
if (!io.asJson) {
|
||||
io.emit(
|
||||
"(Ctrl-C 只停止显示;用 `memhub models switch cancel` 取消切换)\n",
|
||||
);
|
||||
}
|
||||
while (!TERMINAL.has(job.state)) {
|
||||
try {
|
||||
await wait(POLL_MS, signal);
|
||||
} catch {
|
||||
// Aborted below.
|
||||
}
|
||||
if (signal.aborted) {
|
||||
if (!io.asJson) {
|
||||
io.emit(
|
||||
"已停止显示进度;切换仍在后台进行。用 `memhub models switch cancel` 取消。\n",
|
||||
);
|
||||
}
|
||||
return EXIT_OK;
|
||||
}
|
||||
const current = await readCurrent(io);
|
||||
if (current === null) continue;
|
||||
const next =
|
||||
current.job?.job_id === job.job_id
|
||||
? current.job
|
||||
: current.last_finished?.job_id === job.job_id
|
||||
? current.last_finished
|
||||
: null;
|
||||
if (next === null) continue;
|
||||
job = next;
|
||||
show(job);
|
||||
}
|
||||
if (io.asJson) {
|
||||
emitDocument(
|
||||
io,
|
||||
document(command, {
|
||||
...progressFields(job),
|
||||
result: job.state === "done" ? "switched" : "rolled_back",
|
||||
}),
|
||||
);
|
||||
} else {
|
||||
io.emit(renderOutcomeText(job));
|
||||
}
|
||||
return job.state === "done" ? EXIT_OK : EXIT_FAILED;
|
||||
}
|
||||
|
||||
function defaultWait(milliseconds: number, signal: AbortSignal): Promise<void> {
|
||||
return new Promise((resolve, reject) => {
|
||||
if (signal.aborted) {
|
||||
reject(new Error("aborted"));
|
||||
return;
|
||||
}
|
||||
const timer = setTimeout(() => {
|
||||
signal.removeEventListener("abort", onAbort);
|
||||
resolve();
|
||||
}, milliseconds);
|
||||
const onAbort = () => {
|
||||
clearTimeout(timer);
|
||||
reject(new Error("aborted"));
|
||||
};
|
||||
signal.addEventListener("abort", onAbort, { once: true });
|
||||
});
|
||||
}
|
||||
|
||||
/** `memhub models switch status [--json]`. */
|
||||
export async function switchStatus(io: SwitchIo): Promise<number> {
|
||||
const outcome = await io.reader<unknown>(
|
||||
"/api/v1/embedding-switches/current",
|
||||
);
|
||||
if (outcome.state !== "ok")
|
||||
return reject(io, "models switch status", outcome.code);
|
||||
const current = parseCurrent(outcome.data);
|
||||
if (!current) {
|
||||
return reject(io, "models switch status", "embedding_switch_invalid");
|
||||
}
|
||||
if (io.asJson) {
|
||||
emitDocument(
|
||||
io,
|
||||
document("models switch status", {
|
||||
state: current.phase,
|
||||
writes_paused: current.writes_paused,
|
||||
job: current.job ? progressFields(current.job) : null,
|
||||
last_finished: current.last_finished
|
||||
? progressFields(current.last_finished)
|
||||
: null,
|
||||
}),
|
||||
);
|
||||
return EXIT_OK;
|
||||
}
|
||||
if (current.job) {
|
||||
io.emit(
|
||||
`${renderProgressLine(current.job)}\n记忆整理已暂停,检索不受影响。\n`,
|
||||
);
|
||||
} else if (current.phase === "scanning") {
|
||||
io.emit("正在检查向量数据…\n");
|
||||
} else if (current.last_finished) {
|
||||
io.emit(
|
||||
`没有进行中的向量模型切换。上一次:${renderOutcomeText(current.last_finished)}`,
|
||||
);
|
||||
} else {
|
||||
io.emit("没有进行中的向量模型切换。\n");
|
||||
}
|
||||
return EXIT_OK;
|
||||
}
|
||||
|
||||
/** `memhub models switch cancel [--json]`: cancel the running job, if any. */
|
||||
export async function switchCancel(io: SwitchIo): Promise<number> {
|
||||
const current = await readCurrent(io);
|
||||
if (current === null) {
|
||||
return reject(io, "models switch cancel", "embedding_switch_unavailable");
|
||||
}
|
||||
if (!current.job) {
|
||||
return reject(io, "models switch cancel", "embedding_switch_not_found");
|
||||
}
|
||||
const outcome = await io.operations.post<unknown>(
|
||||
`/api/v1/embedding-switches/${current.job.job_id}/cancel`,
|
||||
{},
|
||||
);
|
||||
if (outcome.state !== "committed") {
|
||||
return reject(io, "models switch cancel", outcome.code);
|
||||
}
|
||||
const job = parseJob(outcome.data) ?? current.job;
|
||||
if (io.asJson) {
|
||||
emitDocument(
|
||||
io,
|
||||
document("models switch cancel", {
|
||||
...progressFields(job),
|
||||
cancel_requested: true,
|
||||
}),
|
||||
);
|
||||
} else {
|
||||
io.emit("已请求取消;当前批次完成后回滚,继续使用原来的向量模型。\n");
|
||||
}
|
||||
return EXIT_OK;
|
||||
}
|
||||
|
||||
interface SpaceRow {
|
||||
space_id: string;
|
||||
state: string;
|
||||
model: string | null;
|
||||
dims: number | null;
|
||||
vectors: number | null;
|
||||
keep_until: string | null;
|
||||
}
|
||||
|
||||
function spaceRows(value: unknown): SpaceRow[] | null {
|
||||
const body = record(value);
|
||||
if (!body || !Array.isArray(body.spaces)) return null;
|
||||
const rows: SpaceRow[] = [];
|
||||
for (const item of body.spaces) {
|
||||
const row = record(item);
|
||||
const spaceId = str(row?.space_id, 64);
|
||||
const state = str(row?.state, 16);
|
||||
if (!row || !spaceId || !state) return null;
|
||||
if (state === "dropped") continue;
|
||||
rows.push({
|
||||
space_id: spaceId,
|
||||
state,
|
||||
model: str(row.model),
|
||||
dims: int(row.dims) ?? int(row.name_dims),
|
||||
vectors: int(row.vectors),
|
||||
keep_until: str(row.keep_until, 64),
|
||||
});
|
||||
}
|
||||
return rows;
|
||||
}
|
||||
|
||||
const SPACE_STATE: Record<string, string> = {
|
||||
live: "使用中",
|
||||
building: "正在生成",
|
||||
stray: "待合并",
|
||||
retired: "已停用",
|
||||
};
|
||||
|
||||
/** `memhub models spaces [--json]`. */
|
||||
export async function listSpaces(io: SwitchIo): Promise<number> {
|
||||
const outcome = await io.reader<unknown>("/api/v1/embedding-spaces");
|
||||
if (outcome.state !== "ok") return reject(io, "models spaces", outcome.code);
|
||||
const rows = spaceRows(outcome.data);
|
||||
if (!rows) return reject(io, "models spaces", "embedding_spaces_invalid");
|
||||
const body = record(outcome.data) ?? {};
|
||||
const unregistered = Array.isArray(body.scanned)
|
||||
? body.scanned
|
||||
.map(record)
|
||||
.filter(
|
||||
(item): item is Record<string, unknown> =>
|
||||
item !== null && item.state === "unregistered",
|
||||
)
|
||||
.map((item) => ({
|
||||
collection_base: str(item.collection_base) ?? "",
|
||||
vectors: int(item.vectors) ?? 0,
|
||||
}))
|
||||
: [];
|
||||
if (io.asJson) {
|
||||
emitDocument(
|
||||
io,
|
||||
document("models spaces", {
|
||||
state: "ok",
|
||||
spaces: rows,
|
||||
unregistered,
|
||||
scan_error:
|
||||
typeof body.scan_error === "string" && CODE.test(body.scan_error)
|
||||
? body.scan_error
|
||||
: null,
|
||||
}),
|
||||
);
|
||||
return EXIT_OK;
|
||||
}
|
||||
const lines = ["向量空间:"];
|
||||
for (const row of rows) {
|
||||
const dims = row.dims === null ? "维度未知" : `${row.dims} 维`;
|
||||
const keep =
|
||||
row.state === "retired" && row.keep_until
|
||||
? `,保留到 ${row.keep_until.slice(0, 10)}`
|
||||
: "";
|
||||
lines.push(
|
||||
` ${row.space_id} ${SPACE_STATE[row.state] ?? row.state} ${row.model ?? "未知模型"}(${dims}) ${row.vectors ?? "?"} 条${keep}`,
|
||||
);
|
||||
}
|
||||
for (const item of unregistered) {
|
||||
lines.push(
|
||||
` (未登记)${item.collection_base} ${item.vectors} 条,下次启动时合并`,
|
||||
);
|
||||
}
|
||||
if (rows.some((row) => row.state === "retired")) {
|
||||
lines.push(
|
||||
"清理已停用的空间:memhub models spaces --cleanup <space-id> --yes",
|
||||
);
|
||||
}
|
||||
io.emit(`${lines.join("\n")}\n`);
|
||||
return EXIT_OK;
|
||||
}
|
||||
|
||||
/** `memhub models spaces --cleanup <space-id> --yes`: drop a retired space now. */
|
||||
export async function cleanupSpace(
|
||||
spaceId: string,
|
||||
io: SwitchIo,
|
||||
): Promise<number> {
|
||||
if (!UUID.test(spaceId)) {
|
||||
return reject(io, "models spaces", "embedding_space_not_found");
|
||||
}
|
||||
const outcome = await io.operations.delete<unknown>(
|
||||
`/api/v1/embedding-spaces/${spaceId}`,
|
||||
);
|
||||
if (outcome.state !== "committed") {
|
||||
return reject(io, "models spaces", outcome.code);
|
||||
}
|
||||
if (io.asJson) {
|
||||
emitDocument(
|
||||
io,
|
||||
document("models spaces", { state: "dropped", space_id: spaceId }),
|
||||
);
|
||||
} else {
|
||||
io.emit(`已清理旧向量数据:${spaceId}\n`);
|
||||
}
|
||||
return EXIT_OK;
|
||||
}
|
||||
@@ -510,15 +510,98 @@ describe("models add", () => {
|
||||
expect(result.stdout + result.stderr).not.toContain("SYNTHETIC_LOOSE_KEY");
|
||||
});
|
||||
|
||||
it("activates the new profile with --activate", async () => {
|
||||
it("probes the new profile before activating it with --activate", async () => {
|
||||
const configFile = await modelConfigFile();
|
||||
transport.post.mockImplementation(async (url: string) =>
|
||||
url === "/api/v1/model-profiles"
|
||||
? { status: 200, data: PROFILE }
|
||||
: url === "/api/v1/model-profiles/profile-1/activate"
|
||||
? { status: 200, data: { ...PROFILE, is_active: true } }
|
||||
: { status: 404, data: {} },
|
||||
let probeId: string | undefined;
|
||||
transport.get.mockImplementation(async (url: string) => {
|
||||
if (url === "/api/v1/model-profiles") {
|
||||
return { status: 200, data: [PROBE_PROFILE] };
|
||||
}
|
||||
const match = /^\/api\/v1\/model-probe-operations\/(.+)$/.exec(url);
|
||||
if (match) {
|
||||
return {
|
||||
status: 200,
|
||||
data:
|
||||
probeId === undefined
|
||||
? probeNotFound(match[1])
|
||||
: probeReceiptPassing(match[1]),
|
||||
};
|
||||
}
|
||||
return { status: 404, data: {} };
|
||||
});
|
||||
transport.post.mockImplementation(async (url: string, body: unknown) => {
|
||||
if (url === "/api/v1/model-profiles") {
|
||||
return { status: 200, data: PROBE_PROFILE };
|
||||
}
|
||||
if (url === `/api/v1/model-profiles/${PROBE_PROFILE.profile_id}/probe`) {
|
||||
probeId = (body as { operation_id: string }).operation_id;
|
||||
return { status: 200, data: probePending(probeId) };
|
||||
}
|
||||
if (url === "/api/v1/embedding-switches/plan") {
|
||||
return { status: 200, data: samePlan() };
|
||||
}
|
||||
if (
|
||||
url === `/api/v1/model-profiles/${PROBE_PROFILE.profile_id}/activate`
|
||||
) {
|
||||
return {
|
||||
status: 200,
|
||||
data: {
|
||||
...PROBE_PROFILE,
|
||||
is_active: true,
|
||||
embedding_model: "synthetic-embedding",
|
||||
},
|
||||
};
|
||||
}
|
||||
return { status: 404, data: {} };
|
||||
});
|
||||
const result = await capturedRun(
|
||||
["models", "add", "--config-file", configFile, "--activate", "--json"],
|
||||
{ watchWait: async () => undefined },
|
||||
);
|
||||
expect(result.code).toBe(0);
|
||||
const lines = result.stdout
|
||||
.trim()
|
||||
.split("\n")
|
||||
.map((line) => JSON.parse(line));
|
||||
expect(lines[0].probe).toMatchObject({ state: "completed" });
|
||||
expect(lines.at(-1)).toMatchObject({ state: "committed" });
|
||||
expect(transport.post.mock.calls.map(([url]) => url)).toEqual([
|
||||
"/api/v1/model-profiles",
|
||||
`/api/v1/model-profiles/${PROBE_PROFILE.profile_id}/probe`,
|
||||
"/api/v1/embedding-switches/plan",
|
||||
`/api/v1/model-profiles/${PROBE_PROFILE.profile_id}/activate`,
|
||||
]);
|
||||
});
|
||||
|
||||
it("does not activate a profile whose probe fails", async () => {
|
||||
const configFile = await modelConfigFile();
|
||||
let probeId: string | undefined;
|
||||
transport.get.mockImplementation(async (url: string) => {
|
||||
if (url === "/api/v1/model-profiles") {
|
||||
return { status: 200, data: [PROBE_PROFILE] };
|
||||
}
|
||||
const match = /^\/api\/v1\/model-probe-operations\/(.+)$/.exec(url);
|
||||
if (match) {
|
||||
return {
|
||||
status: 200,
|
||||
data:
|
||||
probeId === undefined
|
||||
? probeNotFound(match[1])
|
||||
: probeReceipt(match[1]),
|
||||
};
|
||||
}
|
||||
return { status: 404, data: {} };
|
||||
});
|
||||
transport.post.mockImplementation(async (url: string, body: unknown) => {
|
||||
if (url === "/api/v1/model-profiles") {
|
||||
return { status: 200, data: PROBE_PROFILE };
|
||||
}
|
||||
if (url.endsWith("/probe")) {
|
||||
probeId = (body as { operation_id: string }).operation_id;
|
||||
return { status: 200, data: probeReceipt(probeId) };
|
||||
}
|
||||
return { status: 404, data: {} };
|
||||
});
|
||||
const result = await capturedRun([
|
||||
"models",
|
||||
"add",
|
||||
@@ -527,14 +610,12 @@ describe("models add", () => {
|
||||
"--activate",
|
||||
"--json",
|
||||
]);
|
||||
expect(result.code).toBe(0);
|
||||
expect(JSON.parse(result.stdout).activation).toEqual({
|
||||
state: "committed",
|
||||
});
|
||||
expect(transport.post.mock.calls.map(([url]) => url)).toEqual([
|
||||
"/api/v1/model-profiles",
|
||||
"/api/v1/model-profiles/profile-1/activate",
|
||||
]);
|
||||
expect(result.code).toBe(1);
|
||||
expect(
|
||||
transport.post.mock.calls.some(([url]) =>
|
||||
String(url).endsWith("/activate"),
|
||||
),
|
||||
).toBe(false);
|
||||
});
|
||||
|
||||
it("recovers an unknown create through its operation receipt", async () => {
|
||||
@@ -752,6 +833,54 @@ function probePending(operationId: string) {
|
||||
};
|
||||
}
|
||||
|
||||
function samePlan() {
|
||||
const space = {
|
||||
space_id: null,
|
||||
model: "synthetic-embedding",
|
||||
dims: 8,
|
||||
collection_base: "",
|
||||
};
|
||||
return {
|
||||
needed: false,
|
||||
reason: "same_embedding",
|
||||
plan_token: "token-1",
|
||||
source: space,
|
||||
target: space,
|
||||
target_profile_id: PROBE_PROFILE.profile_id,
|
||||
target_revision: 3,
|
||||
target_provider_kind: "openai_compatible",
|
||||
target_local: false,
|
||||
blockers: [],
|
||||
};
|
||||
}
|
||||
|
||||
function probeReceiptPassing(operationId: string) {
|
||||
const receipt = probeReceipt(operationId);
|
||||
return {
|
||||
...receipt,
|
||||
result: {
|
||||
...receipt.result,
|
||||
stages: [
|
||||
{
|
||||
stage: "embedding",
|
||||
ok: true,
|
||||
code: "ok",
|
||||
latency_ms: 1,
|
||||
embedding_dimensions: 8,
|
||||
},
|
||||
],
|
||||
capabilities: {
|
||||
endpoint_reachable: true,
|
||||
model_available: true,
|
||||
structured_output: true,
|
||||
embedding_available: true,
|
||||
},
|
||||
embedding_width: 8,
|
||||
embedding_switch_required: false,
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
function probeNotFound(operationId: string) {
|
||||
return {
|
||||
schema_version: 1,
|
||||
@@ -819,7 +948,10 @@ function importJob() {
|
||||
};
|
||||
}
|
||||
|
||||
async function capturedRun(argv: string[]) {
|
||||
async function capturedRun(
|
||||
argv: string[],
|
||||
controls: Parameters<typeof run>[2] = {},
|
||||
) {
|
||||
let stdout = "";
|
||||
let stderr = "";
|
||||
const originalOut = process.stdout.write;
|
||||
@@ -833,7 +965,7 @@ async function capturedRun(argv: string[]) {
|
||||
return true;
|
||||
}) as typeof process.stderr.write;
|
||||
try {
|
||||
const code = await run(argv, {});
|
||||
const code = await run(argv, {}, controls);
|
||||
return { code, stdout, stderr };
|
||||
} finally {
|
||||
process.stdout.write = originalOut;
|
||||
|
||||
@@ -0,0 +1,318 @@
|
||||
import { beforeEach, describe, expect, it, vi } from "vitest";
|
||||
|
||||
import { run } from "./cli.js";
|
||||
import { parseFeaturesCommand } from "./memory-features-flow.js";
|
||||
import { composeStatusReport, renderStatusText } from "./status-report.js";
|
||||
|
||||
const transport = vi.hoisted(() => ({
|
||||
get: vi.fn(),
|
||||
post: vi.fn(),
|
||||
put: vi.fn(),
|
||||
patch: vi.fn(),
|
||||
delete: vi.fn(),
|
||||
}));
|
||||
vi.mock("axios", () => ({
|
||||
default: {
|
||||
create: vi.fn(() => transport),
|
||||
get: vi.fn(),
|
||||
post: vi.fn(),
|
||||
isAxiosError: () => false,
|
||||
},
|
||||
}));
|
||||
|
||||
const runtime = vi.hoisted(() => ({
|
||||
inspectRuntimeStatus: vi.fn(),
|
||||
discoverRuntimeAttachment: vi.fn(),
|
||||
}));
|
||||
vi.mock("@one-memory/runtime-control", async (importOriginal) => ({
|
||||
...(await importOriginal<typeof import("@one-memory/runtime-control")>()),
|
||||
...runtime,
|
||||
}));
|
||||
|
||||
const DEFAULTS: Record<string, boolean> = {
|
||||
graph_rerank: false,
|
||||
same_series_timeline: false,
|
||||
episode_recall: false,
|
||||
session_episodes: false,
|
||||
episode_summary: false,
|
||||
intention: true,
|
||||
few_shot: false,
|
||||
};
|
||||
|
||||
function featuresState(
|
||||
configured: Record<string, boolean> = {},
|
||||
mode = "pro",
|
||||
extra: Record<string, unknown> = {},
|
||||
) {
|
||||
const nonDefault: Record<string, boolean> = {};
|
||||
const features = Object.entries(DEFAULTS).map(([name, fallback]) => {
|
||||
const enabled = configured[name] ?? fallback;
|
||||
if (enabled !== fallback) nonDefault[name] = enabled;
|
||||
const needsUltra = name === "graph_rerank";
|
||||
const reason = !enabled
|
||||
? "disabled"
|
||||
: needsUltra && mode !== "ultra"
|
||||
? "requires_ultra"
|
||||
: null;
|
||||
return {
|
||||
name,
|
||||
enabled,
|
||||
configured: name in configured ? configured[name] : null,
|
||||
default: fallback,
|
||||
active: reason === null,
|
||||
reason,
|
||||
side: "read",
|
||||
requires_mode: needsUltra ? "ultra" : "pro",
|
||||
requires_feature: null,
|
||||
cost: "none",
|
||||
sdk_field: `recall.${name}`,
|
||||
env: `MEMORY_${name.toUpperCase()}`,
|
||||
description: name,
|
||||
};
|
||||
});
|
||||
return { mode, features, non_default: nonDefault, ...extra };
|
||||
}
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
for (const method of Object.values(transport)) method.mockReset();
|
||||
runtime.inspectRuntimeStatus.mockResolvedValue({
|
||||
state: "running",
|
||||
code: "runtime_verified",
|
||||
instanceId: "instance-1",
|
||||
host: "127.0.0.1",
|
||||
port: 8765,
|
||||
version: "1.2.3",
|
||||
});
|
||||
runtime.discoverRuntimeAttachment.mockResolvedValue({
|
||||
instanceId: "instance-1",
|
||||
apiToken: "synthetic-runtime-token",
|
||||
host: "127.0.0.1",
|
||||
port: 8765,
|
||||
servicePid: 4242,
|
||||
health: {
|
||||
status: "ok",
|
||||
service: "one-memory-core",
|
||||
version: "1.2.3",
|
||||
instance_id: "instance-1",
|
||||
},
|
||||
});
|
||||
transport.get.mockImplementation(async (url: string) =>
|
||||
url === "/api/v1/settings/memory-features"
|
||||
? { status: 200, data: featuresState() }
|
||||
: { status: 404, data: {} },
|
||||
);
|
||||
});
|
||||
|
||||
describe("memhub settings features", () => {
|
||||
it("lists every feature with its value and whether it acts", async () => {
|
||||
const result = await capturedRun(["settings", "features", "--json"]);
|
||||
expect(result.code).toBe(0);
|
||||
const json = JSON.parse(result.stdout);
|
||||
expect(json.state).toBe("ok");
|
||||
expect(json.features.features).toHaveLength(7);
|
||||
expect(transport.put).not.toHaveBeenCalled();
|
||||
const text = (await capturedRun(["settings", "features"])).stdout;
|
||||
expect(text).toContain("graph_rerank:关(默认)");
|
||||
expect(text).toContain("intention:开(默认)");
|
||||
expect(text).toContain("先开启 session_episodes");
|
||||
});
|
||||
|
||||
it("prints the plan and changes nothing without --yes", async () => {
|
||||
const result = await capturedRun([
|
||||
"settings",
|
||||
"features",
|
||||
"set",
|
||||
"graph_rerank=true",
|
||||
"episode_recall=true",
|
||||
]);
|
||||
expect(result.code).toBe(0);
|
||||
expect(transport.put).not.toHaveBeenCalled();
|
||||
expect(result.stdout).toContain(
|
||||
"memhub settings features set graph_rerank=true episode_recall=true --yes",
|
||||
);
|
||||
});
|
||||
|
||||
it("sets features with --yes and reports the ones inactive in this mode", async () => {
|
||||
transport.put.mockResolvedValue({
|
||||
status: 200,
|
||||
data: featuresState({ graph_rerank: true, episode_recall: true }, "pro", {
|
||||
changed: ["graph_rerank", "episode_recall"],
|
||||
}),
|
||||
});
|
||||
const result = await capturedRun([
|
||||
"settings",
|
||||
"features",
|
||||
"set",
|
||||
"graph_rerank=true",
|
||||
"episode_recall=true",
|
||||
"--yes",
|
||||
]);
|
||||
expect(result.code).toBe(0);
|
||||
expect(transport.put).toHaveBeenCalledWith(
|
||||
"/api/v1/settings/memory-features",
|
||||
{ graph_rerank: true, episode_recall: true },
|
||||
undefined,
|
||||
);
|
||||
expect(result.stdout).toContain("已更新:graph_rerank、episode_recall");
|
||||
expect(result.stdout).toContain(
|
||||
"graph_rerank:开(已设置,当前模式下不生效(需要 ultra))",
|
||||
);
|
||||
});
|
||||
|
||||
it("does nothing when the values are already set", async () => {
|
||||
transport.get.mockResolvedValue({
|
||||
status: 200,
|
||||
data: featuresState({ few_shot: true }),
|
||||
});
|
||||
const result = await capturedRun([
|
||||
"settings",
|
||||
"features",
|
||||
"set",
|
||||
"few_shot=true",
|
||||
"--yes",
|
||||
]);
|
||||
expect(result.code).toBe(0);
|
||||
expect(result.stdout).toContain("没有需要改动的设置");
|
||||
expect(transport.put).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("resets the named features, or every set one, to the default", async () => {
|
||||
transport.get.mockResolvedValue({
|
||||
status: 200,
|
||||
data: featuresState({ few_shot: true, intention: false }),
|
||||
});
|
||||
transport.put.mockResolvedValue({
|
||||
status: 200,
|
||||
data: featuresState({}, "pro", { changed: ["intention", "few_shot"] }),
|
||||
});
|
||||
expect(
|
||||
(await capturedRun(["settings", "features", "reset", "few_shot"])).code,
|
||||
).toBe(0);
|
||||
expect(transport.put).toHaveBeenLastCalledWith(
|
||||
"/api/v1/settings/memory-features",
|
||||
{ few_shot: null },
|
||||
undefined,
|
||||
);
|
||||
const all = await capturedRun(["settings", "features", "reset", "--json"]);
|
||||
expect(JSON.parse(all.stdout).state).toBe("committed");
|
||||
expect(transport.put).toHaveBeenLastCalledWith(
|
||||
"/api/v1/settings/memory-features",
|
||||
{ intention: null, few_shot: null },
|
||||
undefined,
|
||||
);
|
||||
});
|
||||
|
||||
it("refuses unknown names and malformed values before the Runtime", async () => {
|
||||
for (const bad of [
|
||||
["settings", "features", "set"],
|
||||
["settings", "features", "set", "local_rerank=true"],
|
||||
["settings", "features", "set", "graph_rerank=yes"],
|
||||
["settings", "features", "set", "graph_rerank"],
|
||||
["settings", "features", "set", "few_shot=true", "few_shot=false"],
|
||||
["settings", "features", "reset", "nope"],
|
||||
["settings", "features", "toggle"],
|
||||
]) {
|
||||
const result = await capturedRun(bad);
|
||||
expect(result.code, bad.join(" ")).toBe(2);
|
||||
}
|
||||
expect(transport.put).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("surfaces a 422 from the Runtime as rejected", async () => {
|
||||
transport.put.mockResolvedValue({
|
||||
status: 422,
|
||||
data: { detail: "invalid_request" },
|
||||
});
|
||||
const result = await capturedRun([
|
||||
"settings",
|
||||
"features",
|
||||
"set",
|
||||
"few_shot=true",
|
||||
"--yes",
|
||||
"--json",
|
||||
]);
|
||||
expect(result.code).toBe(1);
|
||||
expect(JSON.parse(result.stdout).state).toBe("rejected");
|
||||
});
|
||||
|
||||
it("parses the forms the help shows", () => {
|
||||
expect(parseFeaturesCommand([])).toEqual({ kind: "show" });
|
||||
expect(parseFeaturesCommand(["set", "intention=false", "--yes"])).toEqual({
|
||||
kind: "set",
|
||||
values: { intention: false },
|
||||
confirmed: true,
|
||||
});
|
||||
expect(parseFeaturesCommand(["reset"])).toEqual({
|
||||
kind: "reset",
|
||||
names: [],
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe("status pipeline features", () => {
|
||||
it("shows the features that differ from the defaults", async () => {
|
||||
const processing = {
|
||||
pending: 0,
|
||||
fresh_pending: 0,
|
||||
backlog_pending: 0,
|
||||
outside_window_pending: 0,
|
||||
running: 0,
|
||||
waiting_for_model: 0,
|
||||
completed: 3,
|
||||
failed: 0,
|
||||
blocked_stages: [],
|
||||
pipeline: {
|
||||
mode: "pro",
|
||||
configured_mode: "pro",
|
||||
forced_by_env: false,
|
||||
last_evolve_at: null,
|
||||
graph_build: { state: "idle", users: {} },
|
||||
evolve: { mode: "pro", schedule: null, running: [], users: {} },
|
||||
features: { episode_recall: true, intention: false },
|
||||
},
|
||||
};
|
||||
const report = await composeStatusReport({
|
||||
dataDir: "/tmp/isolated",
|
||||
runtime: { state: "running", code: "ok", host: "127.0.0.1", port: 1 },
|
||||
background: null,
|
||||
reader: (async (path: string) =>
|
||||
path === "/api/v1/memory-processing"
|
||||
? { state: "ok", data: processing }
|
||||
: { state: "unavailable", code: "x" }) as never,
|
||||
});
|
||||
expect(report.pipeline?.features).toEqual({
|
||||
episode_recall: true,
|
||||
intention: false,
|
||||
});
|
||||
const text = renderStatusText(report, "MemHub");
|
||||
expect(text).toContain("SDK 功能:episode_recall=开、intention=关");
|
||||
});
|
||||
});
|
||||
|
||||
async function capturedRun(arguments_: string[]) {
|
||||
let stdout = "";
|
||||
let stderr = "";
|
||||
const stdoutSpy = vi
|
||||
.spyOn(process.stdout, "write")
|
||||
.mockImplementation((chunk) => {
|
||||
stdout += String(chunk);
|
||||
return true;
|
||||
});
|
||||
const stderrSpy = vi
|
||||
.spyOn(process.stderr, "write")
|
||||
.mockImplementation((chunk) => {
|
||||
stderr += String(chunk);
|
||||
return true;
|
||||
});
|
||||
try {
|
||||
return {
|
||||
code: await run(arguments_, { ONE_MEMORY_DATA_DIR: "/synthetic/data" }),
|
||||
stdout,
|
||||
stderr,
|
||||
};
|
||||
} finally {
|
||||
stdoutSpy.mockRestore();
|
||||
stderrSpy.mockRestore();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,273 @@
|
||||
/**
|
||||
* `memhub settings features [set <name>=<true|false> ... | reset [<name> ...]]`.
|
||||
*
|
||||
* The SDK's optional retrieval and extraction features (graph rerank, episode
|
||||
* recall, the episode write side, intention, few-shot). One global setting in
|
||||
* the Runtime; an unset feature keeps the SDK default, so an install nobody
|
||||
* touched behaves as before. `set` prints its plan first and needs `--yes`,
|
||||
* like `settings mode`, because some features spend model calls; `reset`
|
||||
* only returns features to their defaults and applies at once.
|
||||
*/
|
||||
|
||||
import type { CoreReader } from "./core-client.js";
|
||||
import type { CorePutOperations } from "./core-operations.js";
|
||||
import { PRODUCT } from "./version.js";
|
||||
|
||||
export const MEMORY_FEATURES = [
|
||||
"graph_rerank",
|
||||
"same_series_timeline",
|
||||
"episode_recall",
|
||||
"session_episodes",
|
||||
"episode_summary",
|
||||
"intention",
|
||||
"few_shot",
|
||||
] as const;
|
||||
export type MemoryFeature = (typeof MEMORY_FEATURES)[number];
|
||||
|
||||
const FEATURE_LABEL: Record<MemoryFeature, string> = {
|
||||
graph_rerank: "图谱重排(沿知识图谱的边传播支持度;仅 ultra)",
|
||||
same_series_timeline:
|
||||
"同系列时间线(同一系列的记忆按时间放在一起;需图谱重排)",
|
||||
episode_recall:
|
||||
"片段召回(补入同一片段的记忆并多通道打分;pro / ultra,需片段写入)",
|
||||
session_episodes: "片段写入(同一会话提取的记忆归入同一片段)",
|
||||
episode_summary: "片段摘要(片段结束时总结一次;每个片段一次模型调用)",
|
||||
intention: "意图提取(随事实一起提取用户的计划与意图)",
|
||||
few_shot: "提取示例(在提取提示里附上示例;不加调用,输入更长)",
|
||||
};
|
||||
|
||||
const REASON_TEXT: Record<string, string> = {
|
||||
disabled: "未开启",
|
||||
requires_pro: "lite 模式下不生效(需要 pro 或 ultra)",
|
||||
requires_ultra: "当前模式下不生效(需要 ultra)",
|
||||
requires_graph_rerank: "需要先开启并生效 graph_rerank",
|
||||
};
|
||||
|
||||
export type FeaturesCommand =
|
||||
| { kind: "show" }
|
||||
| {
|
||||
kind: "set";
|
||||
values: Partial<Record<MemoryFeature, boolean>>;
|
||||
confirmed: boolean;
|
||||
}
|
||||
| { kind: "reset"; names: MemoryFeature[] };
|
||||
|
||||
function isFeature(value: string): value is MemoryFeature {
|
||||
return (MEMORY_FEATURES as readonly string[]).includes(value);
|
||||
}
|
||||
|
||||
/** Arguments after `settings features`; null = usage. */
|
||||
export function parseFeaturesCommand(
|
||||
values: readonly string[],
|
||||
): FeaturesCommand | null {
|
||||
const rest = values.filter((value) => value !== "--json");
|
||||
if (rest.length === 0) return { kind: "show" };
|
||||
const [verb, ...tail] = rest;
|
||||
const yes = tail.filter((value) => value === "--yes").length;
|
||||
if (yes > 1) return null;
|
||||
const positional = tail.filter((value) => value !== "--yes");
|
||||
if (verb === "set") {
|
||||
if (positional.length === 0) return null;
|
||||
const parsed: Partial<Record<MemoryFeature, boolean>> = {};
|
||||
for (const item of positional) {
|
||||
const match = /^([a-z_]+)=(true|false)$/.exec(item);
|
||||
if (!match || !isFeature(match[1]) || match[1] in parsed) return null;
|
||||
parsed[match[1]] = match[2] === "true";
|
||||
}
|
||||
return { kind: "set", values: parsed, confirmed: yes === 1 };
|
||||
}
|
||||
if (verb === "reset") {
|
||||
if (new Set(positional).size !== positional.length) return null;
|
||||
if (!positional.every(isFeature)) return null;
|
||||
return { kind: "reset", names: positional as MemoryFeature[] };
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
export interface FeatureState {
|
||||
name: MemoryFeature;
|
||||
enabled: boolean;
|
||||
configured: boolean | null;
|
||||
default: boolean;
|
||||
active: boolean;
|
||||
reason: string | null;
|
||||
cost: string;
|
||||
requires_mode: string;
|
||||
description: string;
|
||||
}
|
||||
|
||||
export interface FeaturesState {
|
||||
mode: string;
|
||||
features: FeatureState[];
|
||||
non_default: Partial<Record<MemoryFeature, boolean>>;
|
||||
changed?: MemoryFeature[];
|
||||
}
|
||||
|
||||
function featuresState(value: unknown): FeaturesState | null {
|
||||
if (!value || typeof value !== "object") return null;
|
||||
const item = value as Record<string, unknown>;
|
||||
if (
|
||||
typeof item.mode !== "string" ||
|
||||
!Array.isArray(item.features) ||
|
||||
!item.non_default ||
|
||||
typeof item.non_default !== "object"
|
||||
) {
|
||||
return null;
|
||||
}
|
||||
for (const feature of item.features) {
|
||||
if (
|
||||
!feature ||
|
||||
typeof feature !== "object" ||
|
||||
typeof (feature as FeatureState).name !== "string" ||
|
||||
typeof (feature as FeatureState).enabled !== "boolean" ||
|
||||
typeof (feature as FeatureState).active !== "boolean"
|
||||
) {
|
||||
return null;
|
||||
}
|
||||
}
|
||||
return item as unknown as FeaturesState;
|
||||
}
|
||||
|
||||
export type FeaturesDocument = {
|
||||
schema_version: 1;
|
||||
product: string;
|
||||
command: "settings features";
|
||||
} & (
|
||||
| { state: "ok"; features: FeaturesState }
|
||||
| {
|
||||
state: "preview";
|
||||
features: FeaturesState;
|
||||
target: Partial<Record<MemoryFeature, boolean>>;
|
||||
applied: false;
|
||||
next_step: string;
|
||||
}
|
||||
| { state: "unchanged"; features: FeaturesState }
|
||||
| { state: "committed"; features: FeaturesState; changed: MemoryFeature[] }
|
||||
| { state: "rejected" | "unknown" | "unavailable"; code: string }
|
||||
);
|
||||
|
||||
const BASE = {
|
||||
schema_version: 1 as const,
|
||||
product: PRODUCT,
|
||||
command: "settings features" as const,
|
||||
};
|
||||
|
||||
const PATH = "/api/v1/settings/memory-features";
|
||||
|
||||
export async function runFeaturesCommand(
|
||||
command: FeaturesCommand,
|
||||
reader: CoreReader,
|
||||
operations: CorePutOperations,
|
||||
): Promise<FeaturesDocument> {
|
||||
const read = await reader<unknown>(PATH);
|
||||
if (read.state !== "ok") {
|
||||
return { ...BASE, state: "unavailable", code: read.code };
|
||||
}
|
||||
const current = featuresState(read.data);
|
||||
if (current === null) {
|
||||
return { ...BASE, state: "unavailable", code: "response_invalid" };
|
||||
}
|
||||
if (command.kind === "show") {
|
||||
return { ...BASE, state: "ok", features: current };
|
||||
}
|
||||
let body: Record<string, unknown>;
|
||||
if (command.kind === "set") {
|
||||
const configured = new Map(
|
||||
current.features.map((feature) => [feature.name, feature.configured]),
|
||||
);
|
||||
const differs = Object.entries(command.values).some(
|
||||
([name, value]) => configured.get(name as MemoryFeature) !== value,
|
||||
);
|
||||
if (!differs) return { ...BASE, state: "unchanged", features: current };
|
||||
if (!command.confirmed) {
|
||||
const pairs = Object.entries(command.values)
|
||||
.map(([name, value]) => `${name}=${value}`)
|
||||
.join(" ");
|
||||
return {
|
||||
...BASE,
|
||||
state: "preview",
|
||||
features: current,
|
||||
target: command.values,
|
||||
applied: false,
|
||||
next_step: `${PRODUCT} settings features set ${pairs} --yes`,
|
||||
};
|
||||
}
|
||||
body = { ...command.values };
|
||||
} else {
|
||||
const names =
|
||||
command.names.length > 0
|
||||
? command.names
|
||||
: current.features
|
||||
.filter((feature) => feature.configured !== null)
|
||||
.map((feature) => feature.name);
|
||||
if (
|
||||
names.every(
|
||||
(name) =>
|
||||
current.features.find((feature) => feature.name === name)
|
||||
?.configured === null,
|
||||
)
|
||||
) {
|
||||
return { ...BASE, state: "unchanged", features: current };
|
||||
}
|
||||
body = Object.fromEntries(names.map((name) => [name, null]));
|
||||
}
|
||||
const result = await operations.put<unknown>(PATH, body);
|
||||
if (result.state !== "committed") {
|
||||
return { ...BASE, state: result.state, code: result.code };
|
||||
}
|
||||
const changed = featuresState(result.data);
|
||||
if (changed === null) {
|
||||
return { ...BASE, state: "unknown", code: "response_invalid" };
|
||||
}
|
||||
return {
|
||||
...BASE,
|
||||
state: "committed",
|
||||
features: changed,
|
||||
changed: changed.changed ?? [],
|
||||
};
|
||||
}
|
||||
|
||||
export function featuresExitCode(document: FeaturesDocument): number {
|
||||
return ["ok", "preview", "unchanged", "committed"].includes(document.state)
|
||||
? 0
|
||||
: 1;
|
||||
}
|
||||
|
||||
function featureLine(feature: FeatureState): string {
|
||||
const label = FEATURE_LABEL[feature.name] ?? feature.name;
|
||||
const value = feature.enabled ? "开" : "关";
|
||||
const origin = feature.configured === null ? "默认" : "已设置";
|
||||
const activity =
|
||||
feature.enabled && !feature.active && feature.reason
|
||||
? `,${REASON_TEXT[feature.reason] ?? feature.reason}`
|
||||
: "";
|
||||
return `${feature.name}:${value}(${origin}${activity})— ${label}`;
|
||||
}
|
||||
|
||||
export function renderFeaturesText(document: FeaturesDocument): string {
|
||||
if ("code" in document) {
|
||||
return `SDK 功能未能${document.state === "unavailable" ? "读取" : "修改"}:${document.code}\n`;
|
||||
}
|
||||
const state = document.features;
|
||||
const lines = [`当前模式:${state.mode}`];
|
||||
for (const feature of state.features) lines.push(featureLine(feature));
|
||||
lines.push(
|
||||
"片段召回(episode_recall)只对带片段的记忆有用:先开启 session_episodes,之后提取的记忆才会带片段。",
|
||||
);
|
||||
if (document.state === "unchanged") lines.push("没有需要改动的设置。");
|
||||
if (document.state === "preview") {
|
||||
lines.push("将设置:");
|
||||
for (const [name, value] of Object.entries(document.target)) {
|
||||
lines.push(`- ${name} = ${value ? "开" : "关"}`);
|
||||
}
|
||||
lines.push(`未执行任何改动。确认后运行:${document.next_step}`);
|
||||
}
|
||||
if (document.state === "committed") {
|
||||
lines.push(
|
||||
document.changed.length > 0
|
||||
? `已更新:${document.changed.join("、")}。正在进行的写入按原设置完成,之后的写入和检索使用新设置。`
|
||||
: "已保存,生效的值没有变化。",
|
||||
);
|
||||
}
|
||||
return `${lines.join("\n")}\n`;
|
||||
}
|
||||
@@ -0,0 +1,309 @@
|
||||
import { beforeEach, describe, expect, it, vi } from "vitest";
|
||||
|
||||
import { run } from "./cli.js";
|
||||
import { modeEffects, parseSettingsCommand } from "./memory-mode-flow.js";
|
||||
import { composeStatusReport, renderStatusText } from "./status-report.js";
|
||||
|
||||
const transport = vi.hoisted(() => ({
|
||||
get: vi.fn(),
|
||||
post: vi.fn(),
|
||||
put: vi.fn(),
|
||||
patch: vi.fn(),
|
||||
delete: vi.fn(),
|
||||
}));
|
||||
vi.mock("axios", () => ({
|
||||
default: {
|
||||
create: vi.fn(() => transport),
|
||||
get: vi.fn(),
|
||||
post: vi.fn(),
|
||||
isAxiosError: () => false,
|
||||
},
|
||||
}));
|
||||
|
||||
const runtime = vi.hoisted(() => ({
|
||||
inspectRuntimeStatus: vi.fn(),
|
||||
discoverRuntimeAttachment: vi.fn(),
|
||||
}));
|
||||
vi.mock("@one-memory/runtime-control", async (importOriginal) => ({
|
||||
...(await importOriginal<typeof import("@one-memory/runtime-control")>()),
|
||||
...runtime,
|
||||
}));
|
||||
|
||||
const RUN_ID = "11111111-2222-4333-8444-555555555555";
|
||||
|
||||
function modeState(mode: string, configured = mode) {
|
||||
return {
|
||||
mode,
|
||||
configured_mode: configured,
|
||||
forced_by_env: false,
|
||||
modes: ["lite", "pro", "ultra"],
|
||||
graph_build: { state: "idle", users: {} },
|
||||
};
|
||||
}
|
||||
|
||||
function receipt(state: string, extra: Record<string, unknown> = {}) {
|
||||
return {
|
||||
run_id: RUN_ID,
|
||||
user_id: "wife",
|
||||
agent_id: "all-apps",
|
||||
state,
|
||||
trigger: "manual",
|
||||
started_at: "2026-11-01T03:00:00Z",
|
||||
rounds: state === "running" ? 0 : 2,
|
||||
tasks_processed: state === "running" ? 0 : 5,
|
||||
batch_failures: 0,
|
||||
capped: false,
|
||||
elapsed_ms: 10,
|
||||
results: { schemas: ["never printed"] },
|
||||
...extra,
|
||||
};
|
||||
}
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
for (const method of Object.values(transport)) method.mockReset();
|
||||
runtime.inspectRuntimeStatus.mockResolvedValue({
|
||||
state: "running",
|
||||
code: "runtime_verified",
|
||||
instanceId: "instance-1",
|
||||
host: "127.0.0.1",
|
||||
port: 8765,
|
||||
version: "1.2.3",
|
||||
});
|
||||
runtime.discoverRuntimeAttachment.mockResolvedValue({
|
||||
instanceId: "instance-1",
|
||||
apiToken: "synthetic-runtime-token",
|
||||
host: "127.0.0.1",
|
||||
port: 8765,
|
||||
servicePid: 4242,
|
||||
health: {
|
||||
status: "ok",
|
||||
service: "one-memory-core",
|
||||
version: "1.2.3",
|
||||
instance_id: "instance-1",
|
||||
},
|
||||
});
|
||||
transport.get.mockImplementation(async (url: string) =>
|
||||
url === "/api/v1/settings/memory-mode"
|
||||
? { status: 200, data: modeState("pro") }
|
||||
: { status: 404, data: {} },
|
||||
);
|
||||
});
|
||||
|
||||
describe("memhub settings mode", () => {
|
||||
it("shows the current mode without changing it", async () => {
|
||||
const result = await capturedRun(["settings", "mode", "--json"]);
|
||||
expect(result.code).toBe(0);
|
||||
expect(JSON.parse(result.stdout).mode.mode).toBe("pro");
|
||||
expect(transport.put).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("prints the plan and changes nothing without --yes", async () => {
|
||||
const result = await capturedRun(["settings", "mode", "lite"]);
|
||||
expect(result.code).toBe(0);
|
||||
expect(transport.put).not.toHaveBeenCalled();
|
||||
expect(result.stdout).toContain("更高模式生成的记忆仍保存着");
|
||||
expect(result.stdout).toContain("memhub settings mode lite --yes");
|
||||
const json = JSON.parse(
|
||||
(await capturedRun(["settings", "mode", "ultra", "--json"])).stdout,
|
||||
);
|
||||
expect(json.state).toBe("preview");
|
||||
expect(json.applied).toBe(false);
|
||||
expect(json.effects).toEqual(["existing_memories_added_to_graph"]);
|
||||
});
|
||||
|
||||
it("changes the mode with --yes and reports what changed", async () => {
|
||||
transport.put.mockResolvedValue({
|
||||
status: 200,
|
||||
data: {
|
||||
...modeState("lite"),
|
||||
previous_mode: "pro",
|
||||
effects: ["higher_mode_memories_hidden"],
|
||||
},
|
||||
});
|
||||
const result = await capturedRun(["settings", "mode", "lite", "--yes"]);
|
||||
expect(result.code).toBe(0);
|
||||
expect(transport.put).toHaveBeenCalledWith(
|
||||
"/api/v1/settings/memory-mode",
|
||||
{ mode: "lite" },
|
||||
undefined,
|
||||
);
|
||||
expect(result.stdout).toContain("已切换");
|
||||
expect(result.stdout).toContain("原始对话:全部保留");
|
||||
});
|
||||
|
||||
it("does nothing for the mode already set, and refuses bad input", async () => {
|
||||
const same = await capturedRun(["settings", "mode", "pro", "--yes"]);
|
||||
expect(same.code).toBe(0);
|
||||
expect(same.stdout).toContain("已经是这个模式");
|
||||
expect(transport.put).not.toHaveBeenCalled();
|
||||
for (const bad of [
|
||||
["settings"],
|
||||
["settings", "mode", "max"],
|
||||
["settings", "mode", "--yes"],
|
||||
["settings", "mode", "lite", "pro"],
|
||||
["settings", "language", "en"],
|
||||
]) {
|
||||
expect((await capturedRun(bad)).code, bad.join(" ")).toBe(2);
|
||||
}
|
||||
});
|
||||
|
||||
it("plans the same effects Core reports", () => {
|
||||
expect(modeEffects("lite", "pro")).toEqual([
|
||||
"earlier_raw_not_extracted",
|
||||
"lite_search_stops",
|
||||
]);
|
||||
expect(modeEffects("ultra", "lite")).toEqual([
|
||||
"higher_mode_memories_hidden",
|
||||
"evolve_stopped",
|
||||
]);
|
||||
expect(modeEffects("pro", "pro")).toEqual([]);
|
||||
expect(parseSettingsCommand(["mode", "ultra", "--yes"])).toEqual({
|
||||
kind: "set",
|
||||
mode: "ultra",
|
||||
confirmed: true,
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe("memhub evolve", () => {
|
||||
it("starts a run for the named user and returns at once", async () => {
|
||||
transport.post.mockResolvedValue({ status: 202, data: receipt("running") });
|
||||
const result = await capturedRun([
|
||||
"evolve",
|
||||
"--user-id",
|
||||
"wife",
|
||||
"--agent-id",
|
||||
"all-apps",
|
||||
"--max-batches",
|
||||
"3",
|
||||
"--json",
|
||||
]);
|
||||
expect(result.code).toBe(0);
|
||||
expect(transport.post).toHaveBeenCalledWith(
|
||||
"/api/v2/evolve",
|
||||
{ user_id: "wife", agent_id: "all-apps", max_batches: 3 },
|
||||
undefined,
|
||||
);
|
||||
const document = JSON.parse(result.stdout);
|
||||
expect(document.state).toBe("running");
|
||||
expect(document.run.results).toBeUndefined();
|
||||
});
|
||||
|
||||
it("follows the run with --wait until it ends", async () => {
|
||||
transport.post.mockResolvedValue({ status: 202, data: receipt("running") });
|
||||
const reads = [receipt("running"), receipt("succeeded")];
|
||||
transport.get.mockImplementation(async (url: string) =>
|
||||
url === `/api/v2/evolve/${RUN_ID}`
|
||||
? { status: 200, data: reads.shift() }
|
||||
: { status: 404, data: {} },
|
||||
);
|
||||
const result = await capturedRun(["evolve", "--wait"], {
|
||||
sleep: async () => undefined,
|
||||
});
|
||||
expect(result.code).toBe(0);
|
||||
expect(transport.post).toHaveBeenCalledWith(
|
||||
"/api/v2/evolve",
|
||||
{},
|
||||
undefined,
|
||||
);
|
||||
expect(result.stdout).toContain("进化完成");
|
||||
expect(result.stdout).toContain("2 轮");
|
||||
});
|
||||
|
||||
it("names the mode when Core refuses outside ultra", async () => {
|
||||
transport.post.mockResolvedValue({
|
||||
status: 409,
|
||||
data: { detail: "memory_mode_not_ultra" },
|
||||
});
|
||||
const result = await capturedRun(["evolve"]);
|
||||
expect(result.code).toBe(1);
|
||||
expect(result.stdout).toContain("memhub settings mode ultra --yes");
|
||||
expect((await capturedRun(["evolve", "--max-batches", "0"])).code).toBe(2);
|
||||
expect((await capturedRun(["evolve", "--now"])).code).toBe(2);
|
||||
});
|
||||
});
|
||||
|
||||
describe("status pipeline", () => {
|
||||
it("reports pipeline.mode and the last evolve from memory-processing", async () => {
|
||||
const processing = {
|
||||
pending: 0,
|
||||
fresh_pending: 0,
|
||||
backlog_pending: 0,
|
||||
outside_window_pending: 0,
|
||||
running: 0,
|
||||
waiting_for_model: 0,
|
||||
completed: 3,
|
||||
failed: 0,
|
||||
blocked_stages: [],
|
||||
pipeline: {
|
||||
mode: "ultra",
|
||||
configured_mode: "ultra",
|
||||
forced_by_env: false,
|
||||
last_evolve_at: "2026-11-01T03:00:05+08:00",
|
||||
graph_build: { state: "succeeded", users: {} },
|
||||
evolve: {
|
||||
mode: "ultra",
|
||||
schedule: { enabled: true, time: "03:00" },
|
||||
running: [],
|
||||
users: {
|
||||
local: {
|
||||
last_run_at: "2026-11-01T03:00:05+08:00",
|
||||
last_result: { state: "succeeded", tasks_processed: 4 },
|
||||
last_error: null,
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
};
|
||||
const report = await composeStatusReport({
|
||||
dataDir: "/tmp/isolated",
|
||||
runtime: { state: "running", code: "ok", host: "127.0.0.1", port: 1 },
|
||||
background: null,
|
||||
reader: (async (path: string) =>
|
||||
path === "/api/v1/memory-processing"
|
||||
? { state: "ok", data: processing }
|
||||
: { state: "unavailable", code: "x" }) as never,
|
||||
});
|
||||
expect(report.pipeline?.mode).toBe("ultra");
|
||||
expect(report.pipeline?.last_evolve_at).toBe("2026-11-01T03:00:05+08:00");
|
||||
expect(report.pipeline?.evolve_users.local?.state).toBe("succeeded");
|
||||
const text = renderStatusText(report, "MemHub");
|
||||
expect(text).toContain("记忆模式:ultra");
|
||||
expect(text).toContain("每日 03:00 自动进化");
|
||||
});
|
||||
});
|
||||
|
||||
async function capturedRun(
|
||||
arguments_: string[],
|
||||
controls: Parameters<typeof run>[2] = {},
|
||||
) {
|
||||
let stdout = "";
|
||||
let stderr = "";
|
||||
const stdoutSpy = vi
|
||||
.spyOn(process.stdout, "write")
|
||||
.mockImplementation((chunk) => {
|
||||
stdout += String(chunk);
|
||||
return true;
|
||||
});
|
||||
const stderrSpy = vi
|
||||
.spyOn(process.stderr, "write")
|
||||
.mockImplementation((chunk) => {
|
||||
stderr += String(chunk);
|
||||
return true;
|
||||
});
|
||||
try {
|
||||
return {
|
||||
code: await run(
|
||||
arguments_,
|
||||
{ ONE_MEMORY_DATA_DIR: "/synthetic/data" },
|
||||
controls,
|
||||
),
|
||||
stdout,
|
||||
stderr,
|
||||
};
|
||||
} finally {
|
||||
stdoutSpy.mockRestore();
|
||||
stderrSpy.mockRestore();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,372 @@
|
||||
/**
|
||||
* `memhub settings mode` and `memhub evolve`.
|
||||
*
|
||||
* The mode (lite / pro / ultra) decides what the SDK does with every
|
||||
* conversation, and switching it migrates nothing: what one mode made is not
|
||||
* shown by a lower one, and raw kept by lite is not extracted until it is
|
||||
* imported again. So a change prints its plan first and needs `--yes`, like
|
||||
* every other command whose effect a person cannot simply take back.
|
||||
*
|
||||
* `memhub evolve` starts System 2 (ultra only) and, with `--wait`, follows
|
||||
* the run until it ends.
|
||||
*/
|
||||
|
||||
import type { CoreReader } from "./core-client.js";
|
||||
import type { CoreOperations, CorePutOperations } from "./core-operations.js";
|
||||
import { PRODUCT } from "./version.js";
|
||||
|
||||
export const MEMORY_MODES = ["lite", "pro", "ultra"] as const;
|
||||
export type MemoryMode = (typeof MEMORY_MODES)[number];
|
||||
|
||||
export type ModeEffect =
|
||||
| "earlier_raw_not_extracted"
|
||||
| "higher_mode_memories_hidden"
|
||||
| "existing_memories_added_to_graph"
|
||||
| "evolve_stopped"
|
||||
| "lite_search_stops"
|
||||
| "graph_backfill_already_running";
|
||||
|
||||
const MODE_SUMMARY: Record<MemoryMode, string> = {
|
||||
lite: "lite:不调用模型,保留全部原始对话,只在原始对话里检索",
|
||||
pro: "pro:模型从对话中提取记忆(System 1)",
|
||||
ultra: "ultra:pro + 每晚自动进化(System 2,使用知识图谱)",
|
||||
};
|
||||
|
||||
const EFFECT_TEXT: Record<ModeEffect, string> = {
|
||||
earlier_raw_not_extracted:
|
||||
"lite 下已处理完的对话不会再被提取;只有尚未处理的和之后的新对话会由模型整理成记忆,检索立即改为记忆检索",
|
||||
higher_mode_memories_hidden:
|
||||
"更高模式生成的记忆仍保存着,但在这个模式下不会显示",
|
||||
existing_memories_added_to_graph:
|
||||
"已有记忆会在后台加入知识图谱(每个用户一次)",
|
||||
evolve_stopped: "每晚的自动进化(System 2)停止",
|
||||
lite_search_stops: "lite 只在原始对话里的检索立即停止,改为检索记忆",
|
||||
graph_backfill_already_running: "知识图谱的补建已在进行中",
|
||||
};
|
||||
|
||||
/** The same rule Core applies (`memory_mode_effects`), for the plan. */
|
||||
export function modeEffects(
|
||||
previous: MemoryMode,
|
||||
next: MemoryMode,
|
||||
): ModeEffect[] {
|
||||
if (previous === next) return [];
|
||||
const rank: Record<MemoryMode, number> = { lite: 0, pro: 1, ultra: 2 };
|
||||
const effects: ModeEffect[] = [];
|
||||
if (previous === "lite") {
|
||||
effects.push("earlier_raw_not_extracted", "lite_search_stops");
|
||||
}
|
||||
if (rank[next] < rank[previous]) effects.push("higher_mode_memories_hidden");
|
||||
if (previous === "ultra") effects.push("evolve_stopped");
|
||||
if (next === "ultra") effects.push("existing_memories_added_to_graph");
|
||||
return effects;
|
||||
}
|
||||
|
||||
export type ModeCommand =
|
||||
| { kind: "show" }
|
||||
| { kind: "set"; mode: MemoryMode; confirmed: boolean };
|
||||
|
||||
/** `settings mode [lite|pro|ultra] [--yes]` after `settings`; null = usage. */
|
||||
export function parseSettingsCommand(
|
||||
values: readonly string[],
|
||||
): ModeCommand | null {
|
||||
const rest = values.filter((value) => value !== "--json");
|
||||
if (rest[0] !== "mode") return null;
|
||||
const tail = rest.slice(1);
|
||||
const confirmed = tail.includes("--yes");
|
||||
const positional = tail.filter((value) => value !== "--yes");
|
||||
if (tail.filter((value) => value === "--yes").length > 1) return null;
|
||||
if (positional.length === 0) return confirmed ? null : { kind: "show" };
|
||||
if (positional.length !== 1) return null;
|
||||
const mode = positional[0] as MemoryMode;
|
||||
if (!(MEMORY_MODES as readonly string[]).includes(mode)) return null;
|
||||
return { kind: "set", mode, confirmed };
|
||||
}
|
||||
|
||||
export interface ModeState {
|
||||
mode: MemoryMode;
|
||||
configured_mode: MemoryMode;
|
||||
forced_by_env: boolean;
|
||||
graph_build?: { state?: string } | null;
|
||||
previous_mode?: MemoryMode;
|
||||
effects?: ModeEffect[];
|
||||
}
|
||||
|
||||
function modeState(value: unknown): ModeState | null {
|
||||
if (!value || typeof value !== "object") return null;
|
||||
const item = value as Record<string, unknown>;
|
||||
if (
|
||||
!(MEMORY_MODES as readonly unknown[]).includes(item.mode) ||
|
||||
!(MEMORY_MODES as readonly unknown[]).includes(item.configured_mode) ||
|
||||
typeof item.forced_by_env !== "boolean"
|
||||
) {
|
||||
return null;
|
||||
}
|
||||
return item as unknown as ModeState;
|
||||
}
|
||||
|
||||
export type ModeDocument = {
|
||||
schema_version: 1;
|
||||
product: string;
|
||||
command: "settings mode";
|
||||
} & (
|
||||
| { state: "ok"; mode: ModeState }
|
||||
| {
|
||||
state: "preview";
|
||||
mode: ModeState;
|
||||
target: MemoryMode;
|
||||
effects: ModeEffect[];
|
||||
applied: false;
|
||||
next_step: string;
|
||||
}
|
||||
| { state: "unchanged"; mode: ModeState }
|
||||
| { state: "committed"; mode: ModeState; effects: ModeEffect[] }
|
||||
| { state: "rejected" | "unknown" | "unavailable"; code: string }
|
||||
);
|
||||
|
||||
const BASE = {
|
||||
schema_version: 1 as const,
|
||||
product: PRODUCT,
|
||||
command: "settings mode" as const,
|
||||
};
|
||||
|
||||
export async function runModeCommand(
|
||||
command: ModeCommand,
|
||||
reader: CoreReader,
|
||||
operations: CorePutOperations,
|
||||
): Promise<ModeDocument> {
|
||||
const read = await reader<unknown>("/api/v1/settings/memory-mode");
|
||||
if (read.state !== "ok") {
|
||||
return { ...BASE, state: "unavailable", code: read.code };
|
||||
}
|
||||
const current = modeState(read.data);
|
||||
if (current === null) {
|
||||
return { ...BASE, state: "unavailable", code: "response_invalid" };
|
||||
}
|
||||
if (command.kind === "show") return { ...BASE, state: "ok", mode: current };
|
||||
if (current.configured_mode === command.mode) {
|
||||
return { ...BASE, state: "unchanged", mode: current };
|
||||
}
|
||||
if (!command.confirmed) {
|
||||
return {
|
||||
...BASE,
|
||||
state: "preview",
|
||||
mode: current,
|
||||
target: command.mode,
|
||||
effects: modeEffects(current.mode, command.mode),
|
||||
applied: false,
|
||||
next_step: `${PRODUCT} settings mode ${command.mode} --yes`,
|
||||
};
|
||||
}
|
||||
const result = await operations.put<unknown>("/api/v1/settings/memory-mode", {
|
||||
mode: command.mode,
|
||||
});
|
||||
if (result.state !== "committed") {
|
||||
return { ...BASE, state: result.state, code: result.code };
|
||||
}
|
||||
const changed = modeState(result.data);
|
||||
if (changed === null) {
|
||||
return { ...BASE, state: "unknown", code: "response_invalid" };
|
||||
}
|
||||
return {
|
||||
...BASE,
|
||||
state: "committed",
|
||||
mode: changed,
|
||||
effects: changed.effects ?? [],
|
||||
};
|
||||
}
|
||||
|
||||
export function modeExitCode(document: ModeDocument): number {
|
||||
return ["ok", "preview", "unchanged", "committed"].includes(document.state)
|
||||
? 0
|
||||
: 1;
|
||||
}
|
||||
|
||||
export function renderModeText(document: ModeDocument): string {
|
||||
if ("code" in document) {
|
||||
return `记忆模式未能${document.state === "unavailable" ? "读取" : "修改"}:${document.code}\n`;
|
||||
}
|
||||
const mode = document.mode;
|
||||
const lines = [`当前模式:${MODE_SUMMARY[mode.mode]}`];
|
||||
if (mode.forced_by_env) {
|
||||
lines.push(
|
||||
`设置里是 ${mode.configured_mode};环境变量 ONE_MEMORY_GRAPH_STORE 强制使用 ultra。`,
|
||||
);
|
||||
}
|
||||
lines.push("原始对话:全部保留(任何模式都不删除)");
|
||||
if (document.state === "unchanged") {
|
||||
lines.push("已经是这个模式,未做改动。");
|
||||
}
|
||||
if (document.state === "preview") {
|
||||
lines.push(`将切换到:${MODE_SUMMARY[document.target]}`);
|
||||
for (const effect of document.effects)
|
||||
lines.push(`- ${EFFECT_TEXT[effect]}`);
|
||||
lines.push(`未执行任何改动。确认后运行:${document.next_step}`);
|
||||
}
|
||||
if (document.state === "committed") {
|
||||
lines.push("已切换。");
|
||||
for (const effect of document.effects)
|
||||
lines.push(`- ${EFFECT_TEXT[effect]}`);
|
||||
}
|
||||
return `${lines.join("\n")}\n`;
|
||||
}
|
||||
|
||||
export type EvolveCommand = {
|
||||
agentId?: string;
|
||||
maxBatches?: number;
|
||||
wait: boolean;
|
||||
};
|
||||
|
||||
/** `evolve [--agent-id <id>] [--max-batches <n>] [--wait]`; null = usage. */
|
||||
export function parseEvolveCommand(
|
||||
values: readonly string[],
|
||||
): EvolveCommand | null {
|
||||
const parsed: EvolveCommand = { wait: false };
|
||||
const rest = values.filter((value) => value !== "--json");
|
||||
for (let index = 0; index < rest.length; index += 1) {
|
||||
const flag = rest[index];
|
||||
if (flag === "--wait" && !parsed.wait) {
|
||||
parsed.wait = true;
|
||||
} else if (flag === "--agent-id" && parsed.agentId === undefined) {
|
||||
const value = rest[index + 1];
|
||||
if (!value || !/^[A-Za-z0-9._:-]{1,128}$/.test(value)) return null;
|
||||
parsed.agentId = value;
|
||||
index += 1;
|
||||
} else if (flag === "--max-batches" && parsed.maxBatches === undefined) {
|
||||
const value = Number(rest[index + 1]);
|
||||
if (!Number.isInteger(value) || value < 1 || value > 1000) return null;
|
||||
parsed.maxBatches = value;
|
||||
index += 1;
|
||||
} else {
|
||||
return null;
|
||||
}
|
||||
}
|
||||
return parsed;
|
||||
}
|
||||
|
||||
export type EvolveDocument = {
|
||||
schema_version: 1;
|
||||
product: string;
|
||||
command: "evolve";
|
||||
} & (
|
||||
| { state: "running" | "succeeded" | "failed"; run: Record<string, unknown> }
|
||||
| { state: "rejected" | "unknown" | "unavailable"; code: string }
|
||||
);
|
||||
|
||||
const RUN_ID = /^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$/;
|
||||
|
||||
function evolveRun(value: unknown): Record<string, unknown> | null {
|
||||
if (!value || typeof value !== "object") return null;
|
||||
const run = value as Record<string, unknown>;
|
||||
if (
|
||||
typeof run.run_id !== "string" ||
|
||||
!RUN_ID.test(run.run_id) ||
|
||||
!["running", "succeeded", "failed"].includes(String(run.state))
|
||||
) {
|
||||
return null;
|
||||
}
|
||||
return run;
|
||||
}
|
||||
|
||||
export async function runEvolveCommand(
|
||||
command: EvolveCommand,
|
||||
options: {
|
||||
userId?: string;
|
||||
reader: CoreReader;
|
||||
operations: CoreOperations;
|
||||
sleep?: (ms: number) => Promise<void>;
|
||||
pollMs?: number;
|
||||
},
|
||||
): Promise<EvolveDocument> {
|
||||
const base = {
|
||||
schema_version: 1 as const,
|
||||
product: PRODUCT,
|
||||
command: "evolve" as const,
|
||||
};
|
||||
const body: Record<string, unknown> = {};
|
||||
if (options.userId !== undefined) body.user_id = options.userId;
|
||||
if (command.agentId !== undefined) body.agent_id = command.agentId;
|
||||
if (command.maxBatches !== undefined) body.max_batches = command.maxBatches;
|
||||
const started = await options.operations.post<unknown>(
|
||||
"/api/v2/evolve",
|
||||
body,
|
||||
);
|
||||
if (started.state !== "committed") {
|
||||
return { ...base, state: started.state, code: started.code };
|
||||
}
|
||||
let run = evolveRun(started.data);
|
||||
if (run === null)
|
||||
return { ...base, state: "unknown", code: "response_invalid" };
|
||||
const sleep =
|
||||
options.sleep ??
|
||||
((ms: number) => new Promise((done) => setTimeout(done, ms)));
|
||||
while (command.wait && run.state === "running") {
|
||||
await sleep(options.pollMs ?? 2_000);
|
||||
const read = await options.reader<unknown>(
|
||||
`/api/v2/evolve/${run.run_id}`,
|
||||
options.userId === undefined ? {} : { user_id: options.userId },
|
||||
);
|
||||
if (read.state !== "ok") {
|
||||
return { ...base, state: "unavailable", code: read.code };
|
||||
}
|
||||
const next = evolveRun(read.data);
|
||||
if (next === null)
|
||||
return { ...base, state: "unknown", code: "response_invalid" };
|
||||
run = next;
|
||||
}
|
||||
return {
|
||||
...base,
|
||||
state: run.state as "running" | "succeeded" | "failed",
|
||||
run: projectRun(run),
|
||||
};
|
||||
}
|
||||
|
||||
/** Counts and codes; SDK result bodies stay in the Core receipt. */
|
||||
function projectRun(run: Record<string, unknown>): Record<string, unknown> {
|
||||
const keep = [
|
||||
"run_id",
|
||||
"user_id",
|
||||
"agent_id",
|
||||
"state",
|
||||
"trigger",
|
||||
"started_at",
|
||||
"finished_at",
|
||||
"rounds",
|
||||
"tasks_processed",
|
||||
"batch_failures",
|
||||
"cycle_status",
|
||||
"should_continue",
|
||||
"continuation_reason",
|
||||
"remaining_clusters",
|
||||
"remaining_actionable_clusters",
|
||||
"remaining_fresh_facts",
|
||||
"capped",
|
||||
"elapsed_ms",
|
||||
"error_code",
|
||||
"failure_kind",
|
||||
];
|
||||
return Object.fromEntries(
|
||||
keep.filter((key) => key in run).map((key) => [key, run[key]]),
|
||||
);
|
||||
}
|
||||
|
||||
export function evolveExitCode(document: EvolveDocument): number {
|
||||
return document.state === "running" || document.state === "succeeded" ? 0 : 1;
|
||||
}
|
||||
|
||||
export function renderEvolveText(document: EvolveDocument): string {
|
||||
if ("code" in document) {
|
||||
return document.code === "memory_mode_not_ultra"
|
||||
? "进化只在 ultra 模式下可用。切换:memhub settings mode ultra --yes\n"
|
||||
: `进化未能开始:${document.code}\n`;
|
||||
}
|
||||
const run = document.run;
|
||||
if (document.state === "running") {
|
||||
return `进化已开始(${run.run_id}),在后台运行;加 --wait 等待结果。\n`;
|
||||
}
|
||||
if (document.state === "failed") {
|
||||
return `进化失败(${run.run_id}):${String(run.error_code ?? "unknown")}\n`;
|
||||
}
|
||||
return `进化完成(${run.run_id}):${run.rounds} 轮,处理 ${run.tasks_processed} 项${
|
||||
run.cycle_status ? `,周期 ${run.cycle_status}` : ""
|
||||
}${run.capped ? ",已到轮数上限,可再次运行继续" : ""}\n`;
|
||||
}
|
||||
@@ -372,6 +372,9 @@ export function renderMemoryAddText(document: MemoryAddDocument): string {
|
||||
export function renderMemoryPageText(document: MemoryPageDocument): string {
|
||||
const page = document.page;
|
||||
if (page.state !== "ok") {
|
||||
if (page.state === "unavailable" && page.code === "embedding_unavailable") {
|
||||
return "向量服务不可用,无法检索;请检查向量模型配置。\n";
|
||||
}
|
||||
if (page.state === "unavailable") {
|
||||
return "后台正忙,记忆列表这次没查完,不代表没有记忆,请稍后再试\n";
|
||||
}
|
||||
|
||||
@@ -22,6 +22,10 @@ const RESULT_KEYS = new Set([
|
||||
"stages",
|
||||
"capabilities",
|
||||
]);
|
||||
const OPTIONAL_RESULT_KEYS = [
|
||||
"embedding_width",
|
||||
"embedding_switch_required",
|
||||
] as const;
|
||||
const STAGE_KEYS = new Set([
|
||||
"stage",
|
||||
"ok",
|
||||
@@ -272,9 +276,34 @@ function validateResult(
|
||||
inputRevision: number,
|
||||
): Record<string, unknown> | null {
|
||||
const result = record(value);
|
||||
// A Core with the embedding registry adds what the embedding stage measured
|
||||
// and whether using it needs an embedding switch; older Cores send neither.
|
||||
const optional = Object.fromEntries(
|
||||
OPTIONAL_RESULT_KEYS.filter((key) => key in (result ?? {})).map((key) => [
|
||||
key,
|
||||
result?.[key],
|
||||
]),
|
||||
);
|
||||
const required = result
|
||||
? Object.fromEntries(
|
||||
Object.entries(result).filter(
|
||||
([key]) => !(OPTIONAL_RESULT_KEYS as readonly string[]).includes(key),
|
||||
),
|
||||
)
|
||||
: null;
|
||||
if (
|
||||
!result ||
|
||||
!exactKeys(result, RESULT_KEYS) ||
|
||||
!required ||
|
||||
!exactKeys(required, RESULT_KEYS) ||
|
||||
!(
|
||||
optional.embedding_width === undefined ||
|
||||
optional.embedding_width === null ||
|
||||
positiveInteger(optional.embedding_width)
|
||||
) ||
|
||||
!(
|
||||
optional.embedding_switch_required === undefined ||
|
||||
typeof optional.embedding_switch_required === "boolean"
|
||||
) ||
|
||||
result.profile_id !== profileId ||
|
||||
result.input_profile_revision !== inputRevision ||
|
||||
!positiveInteger(result.profile_revision) ||
|
||||
|
||||
@@ -117,7 +117,7 @@ export async function editModelProfile(
|
||||
return modelProfileResult(profileId, result);
|
||||
}
|
||||
|
||||
function modelProfileResult(
|
||||
export function modelProfileResult(
|
||||
operationId: string,
|
||||
result:
|
||||
| Awaited<ReturnType<CoreOperations["post"]>>
|
||||
|
||||
@@ -0,0 +1,166 @@
|
||||
import { beforeEach, describe, expect, it, vi } from "vitest";
|
||||
|
||||
import { run } from "./cli.js";
|
||||
|
||||
const transport = vi.hoisted(() => ({
|
||||
get: vi.fn(),
|
||||
post: vi.fn(),
|
||||
put: vi.fn(),
|
||||
patch: vi.fn(),
|
||||
delete: vi.fn(),
|
||||
}));
|
||||
vi.mock("axios", () => ({
|
||||
default: {
|
||||
create: vi.fn(() => transport),
|
||||
get: vi.fn(),
|
||||
post: vi.fn(),
|
||||
isAxiosError: () => false,
|
||||
},
|
||||
}));
|
||||
|
||||
const runtime = vi.hoisted(() => ({
|
||||
inspectRuntimeStatus: vi.fn(),
|
||||
discoverRuntimeAttachment: vi.fn(),
|
||||
}));
|
||||
vi.mock("@one-memory/runtime-control", async (importOriginal) => ({
|
||||
...(await importOriginal<typeof import("@one-memory/runtime-control")>()),
|
||||
...runtime,
|
||||
}));
|
||||
|
||||
const VERSIONS = {
|
||||
user_id: "wife",
|
||||
request_id: "r-versions",
|
||||
elapsed_ms: 2,
|
||||
chains: [
|
||||
{
|
||||
memory_id: "l0-head",
|
||||
content: "The user's name is Fei.",
|
||||
status: "active",
|
||||
memory_at: null,
|
||||
gmt_created: 1_790_000_100,
|
||||
basic_info: { name: "Fei" },
|
||||
versions: [
|
||||
{
|
||||
memory_id: "l0-head",
|
||||
content: "The user's name is Fei.",
|
||||
memory_at: null,
|
||||
gmt_created: 1_790_000_100,
|
||||
current: true,
|
||||
},
|
||||
{
|
||||
memory_id: "l0-old",
|
||||
content: "The user's name is 费天翔.",
|
||||
memory_at: null,
|
||||
gmt_created: 1_790_000_000,
|
||||
current: false,
|
||||
},
|
||||
],
|
||||
},
|
||||
],
|
||||
};
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
for (const method of Object.values(transport)) method.mockReset();
|
||||
runtime.inspectRuntimeStatus.mockResolvedValue({
|
||||
state: "running",
|
||||
code: "runtime_verified",
|
||||
instanceId: "instance-1",
|
||||
host: "127.0.0.1",
|
||||
port: 8765,
|
||||
version: "1.2.3",
|
||||
});
|
||||
runtime.discoverRuntimeAttachment.mockResolvedValue({
|
||||
instanceId: "instance-1",
|
||||
apiToken: "synthetic-runtime-token",
|
||||
host: "127.0.0.1",
|
||||
port: 8765,
|
||||
servicePid: 4242,
|
||||
health: {
|
||||
status: "ok",
|
||||
service: "one-memory-core",
|
||||
version: "1.2.3",
|
||||
instance_id: "instance-1",
|
||||
},
|
||||
});
|
||||
transport.get.mockImplementation(async (url: string) =>
|
||||
url === "/api/v2/user-profile/versions"
|
||||
? { status: 200, data: VERSIONS }
|
||||
: { status: 404, data: {} },
|
||||
);
|
||||
});
|
||||
|
||||
describe("memhub profile --versions", () => {
|
||||
it("reads the L0 version chain for the named user", async () => {
|
||||
const result = await capturedRun([
|
||||
"profile",
|
||||
"--versions",
|
||||
"--user-id",
|
||||
"wife",
|
||||
"--json",
|
||||
]);
|
||||
expect(result.code).toBe(0);
|
||||
expect(transport.get).toHaveBeenCalledWith(
|
||||
"/api/v2/user-profile/versions",
|
||||
{ params: { user_id: "wife" } },
|
||||
);
|
||||
const document = JSON.parse(result.stdout);
|
||||
expect(document.state).toBe("ok");
|
||||
expect(document.chains[0].versions).toHaveLength(2);
|
||||
});
|
||||
|
||||
it("marks the current version in text and names no user for local", async () => {
|
||||
const result = await capturedRun(["profile", "--versions"]);
|
||||
expect(result.code).toBe(0);
|
||||
const [url, config] = transport.get.mock.calls[0] ?? [];
|
||||
expect(url).toBe("/api/v2/user-profile/versions");
|
||||
expect(config?.params?.user_id).toBeUndefined();
|
||||
expect(result.stdout).toContain("* l0-head");
|
||||
expect(result.stdout).toContain(" l0-old");
|
||||
expect(result.stdout).toContain("费天翔");
|
||||
});
|
||||
|
||||
it("refuses an unknown flag and a repeated --versions", async () => {
|
||||
expect((await capturedRun(["profile", "--version-chain"])).code).toBe(2);
|
||||
expect(
|
||||
(await capturedRun(["profile", "--versions", "--versions"])).code,
|
||||
).toBe(2);
|
||||
expect(transport.get).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("reports a malformed answer as invalid", async () => {
|
||||
transport.get.mockResolvedValue({ status: 200, data: { chains: "x" } });
|
||||
const result = await capturedRun(["profile", "--versions", "--json"]);
|
||||
expect(result.code).toBe(1);
|
||||
expect(JSON.parse(result.stdout).code).toBe("response_invalid");
|
||||
});
|
||||
});
|
||||
|
||||
async function capturedRun(arguments_: string[]) {
|
||||
let stdout = "";
|
||||
let stderr = "";
|
||||
const stdoutSpy = vi
|
||||
.spyOn(process.stdout, "write")
|
||||
.mockImplementation((chunk) => {
|
||||
stdout += String(chunk);
|
||||
return true;
|
||||
});
|
||||
const stderrSpy = vi
|
||||
.spyOn(process.stderr, "write")
|
||||
.mockImplementation((chunk) => {
|
||||
stderr += String(chunk);
|
||||
return true;
|
||||
});
|
||||
try {
|
||||
return {
|
||||
code: await run(arguments_, {
|
||||
ONE_MEMORY_DATA_DIR: "/synthetic/data",
|
||||
}),
|
||||
stdout,
|
||||
stderr,
|
||||
};
|
||||
} finally {
|
||||
stdoutSpy.mockRestore();
|
||||
stderrSpy.mockRestore();
|
||||
}
|
||||
}
|
||||
@@ -14,6 +14,11 @@ import {
|
||||
validateProcessing,
|
||||
validateSources,
|
||||
} from "./core-responses.js";
|
||||
import {
|
||||
describePipeline,
|
||||
type PipelineSummary,
|
||||
pipelineLine,
|
||||
} from "./status-report.js";
|
||||
import { PRODUCT } from "./version.js";
|
||||
|
||||
type ReadState =
|
||||
@@ -84,6 +89,8 @@ export interface ProgressDocument {
|
||||
summary_state: ReadState;
|
||||
summary: Record<string, unknown> | null;
|
||||
};
|
||||
/** The SDK mode and last evolve per user; absent on an older Core. */
|
||||
pipeline?: PipelineSummary | null;
|
||||
}
|
||||
|
||||
export async function readSources(
|
||||
@@ -283,6 +290,10 @@ export async function readProgress(
|
||||
summary:
|
||||
processing.state === "ok" ? projectProcessing(processing.data) : null,
|
||||
},
|
||||
pipeline:
|
||||
processingRead.state === "ok"
|
||||
? describePipeline(processingRead.data)
|
||||
: null,
|
||||
};
|
||||
}
|
||||
|
||||
@@ -353,6 +364,18 @@ export function renderProgressText(document: ProgressDocument): string {
|
||||
? `记忆整理:待处理 ${document.generation.summary.pending} · 进行中 ${document.generation.summary.running} · 等待模型 ${document.generation.summary.waiting_for_model} · 失败 ${document.generation.summary.failed}`
|
||||
: "记忆整理:暂时无法读取",
|
||||
);
|
||||
if (document.pipeline) {
|
||||
lines.push(`记忆模式:${pipelineLine(document.pipeline)}`);
|
||||
for (const [userId, user] of Object.entries(
|
||||
document.pipeline.evolve_users,
|
||||
)) {
|
||||
lines.push(
|
||||
` 进化 ${userId}:${user.last_run_at ?? "尚未进化"}${
|
||||
user.last_error ? `(失败:${user.last_error})` : ""
|
||||
}`,
|
||||
);
|
||||
}
|
||||
}
|
||||
if (document.imports.output_truncated) {
|
||||
lines.push("这里只显示汇总;完整记录请使用 --json。");
|
||||
}
|
||||
|
||||
@@ -18,6 +18,7 @@ import {
|
||||
validateProcessing,
|
||||
validateSources,
|
||||
} from "./core-responses.js";
|
||||
import { describeEmbeddingSwitchSummary } from "./embedding-switch-flow.js";
|
||||
import { PRODUCT } from "./version.js";
|
||||
|
||||
const MIGRATION_ISSUES_URL =
|
||||
@@ -57,6 +58,9 @@ export interface StatusDocument {
|
||||
broker?: BackgroundHealth["broker"];
|
||||
};
|
||||
sections: Record<string, Section>;
|
||||
/** R10: the SDK mode and System 2 state; null when the Runtime is not read
|
||||
* or is older than the setting. */
|
||||
pipeline?: PipelineSummary | null;
|
||||
summary: {
|
||||
state: "ok" | "partial";
|
||||
degraded: string[];
|
||||
@@ -64,6 +68,99 @@ export interface StatusDocument {
|
||||
};
|
||||
}
|
||||
|
||||
export interface PipelineSummary {
|
||||
mode: "lite" | "pro" | "ultra";
|
||||
configured_mode: "lite" | "pro" | "ultra";
|
||||
forced_by_env: boolean;
|
||||
last_evolve_at: string | null;
|
||||
evolve_schedule: { enabled: boolean; time: string } | null;
|
||||
evolve_running: number;
|
||||
evolve_users: Record<
|
||||
string,
|
||||
{
|
||||
last_run_at: string | null;
|
||||
last_error: string | null;
|
||||
state: string | null;
|
||||
}
|
||||
>;
|
||||
graph_build: string | null;
|
||||
/** SDK features (`memory_features`) whose value differs from the default. */
|
||||
features: Record<string, boolean>;
|
||||
}
|
||||
|
||||
const MODES = new Set(["lite", "pro", "ultra"]);
|
||||
|
||||
/** Projection of `pipeline` in `/api/v1/memory-processing`; null if absent or malformed. */
|
||||
export function describePipeline(value: unknown): PipelineSummary | null {
|
||||
if (!value || typeof value !== "object") return null;
|
||||
const pipeline = (value as { pipeline?: unknown }).pipeline;
|
||||
if (!pipeline || typeof pipeline !== "object") return null;
|
||||
const item = pipeline as Record<string, unknown>;
|
||||
if (
|
||||
!MODES.has(String(item.mode)) ||
|
||||
!MODES.has(String(item.configured_mode))
|
||||
) {
|
||||
return null;
|
||||
}
|
||||
const evolve =
|
||||
item.evolve && typeof item.evolve === "object"
|
||||
? (item.evolve as Record<string, unknown>)
|
||||
: {};
|
||||
const schedule =
|
||||
evolve.schedule && typeof evolve.schedule === "object"
|
||||
? (evolve.schedule as Record<string, unknown>)
|
||||
: null;
|
||||
const users: PipelineSummary["evolve_users"] = {};
|
||||
if (evolve.users && typeof evolve.users === "object") {
|
||||
for (const [userId, raw] of Object.entries(
|
||||
evolve.users as Record<string, unknown>,
|
||||
)) {
|
||||
const user =
|
||||
raw && typeof raw === "object" ? (raw as Record<string, unknown>) : {};
|
||||
const result =
|
||||
user.last_result && typeof user.last_result === "object"
|
||||
? (user.last_result as Record<string, unknown>)
|
||||
: {};
|
||||
users[userId] = {
|
||||
last_run_at:
|
||||
typeof user.last_run_at === "string" ? user.last_run_at : null,
|
||||
last_error:
|
||||
typeof user.last_error === "string" ? user.last_error : null,
|
||||
state: typeof result.state === "string" ? result.state : null,
|
||||
};
|
||||
}
|
||||
}
|
||||
const features: Record<string, boolean> = {};
|
||||
if (item.features && typeof item.features === "object") {
|
||||
for (const [name, value] of Object.entries(
|
||||
item.features as Record<string, unknown>,
|
||||
)) {
|
||||
if (typeof value === "boolean") features[name] = value;
|
||||
}
|
||||
}
|
||||
const graph =
|
||||
item.graph_build && typeof item.graph_build === "object"
|
||||
? (item.graph_build as Record<string, unknown>)
|
||||
: {};
|
||||
return {
|
||||
mode: item.mode as PipelineSummary["mode"],
|
||||
configured_mode: item.configured_mode as PipelineSummary["mode"],
|
||||
forced_by_env: item.forced_by_env === true,
|
||||
last_evolve_at:
|
||||
typeof item.last_evolve_at === "string" ? item.last_evolve_at : null,
|
||||
evolve_schedule:
|
||||
schedule &&
|
||||
typeof schedule.enabled === "boolean" &&
|
||||
typeof schedule.time === "string"
|
||||
? { enabled: schedule.enabled, time: schedule.time }
|
||||
: null,
|
||||
evolve_running: Array.isArray(evolve.running) ? evolve.running.length : 0,
|
||||
evolve_users: users,
|
||||
graph_build: typeof graph.state === "string" ? graph.state : null,
|
||||
features,
|
||||
};
|
||||
}
|
||||
|
||||
const SECTION_NAMES = [
|
||||
"sources",
|
||||
"sync",
|
||||
@@ -97,6 +194,7 @@ export async function composeStatusReport(options: {
|
||||
attachFailure?: string;
|
||||
}): Promise<StatusDocument> {
|
||||
const sections: Record<string, Section> = {};
|
||||
let pipeline: PipelineSummary | null = null;
|
||||
if (options.reader === null) {
|
||||
const state: SectionState =
|
||||
options.attachFailure === undefined ? "not_checked" : "unavailable";
|
||||
@@ -118,6 +216,8 @@ export async function composeStatusReport(options: {
|
||||
validateProcessing,
|
||||
describeProcessing,
|
||||
);
|
||||
pipeline =
|
||||
processing.state === "ok" ? describePipeline(processing.data) : null;
|
||||
sections.model = section(model, validateModelProfiles, describeModel);
|
||||
sections.memories = section(
|
||||
memories,
|
||||
@@ -169,6 +269,7 @@ export async function composeStatusReport(options: {
|
||||
},
|
||||
bridge,
|
||||
sections,
|
||||
pipeline,
|
||||
summary: {
|
||||
state:
|
||||
degraded.length === 0 && notChecked.length === 0 ? "ok" : "partial",
|
||||
@@ -205,6 +306,11 @@ export function renderStatusText(
|
||||
`${SECTION_LABELS[name]}:${sectionLine(report.sections[name])}`,
|
||||
);
|
||||
}
|
||||
if (report.pipeline) lines.push(`记忆模式:${pipelineLine(report.pipeline)}`);
|
||||
const embeddingSwitch = describeEmbeddingSwitchSummary(
|
||||
report.sections.processing?.embedding_switch,
|
||||
);
|
||||
if (embeddingSwitch !== null) lines.push(`向量模型:${embeddingSwitch}`);
|
||||
if (report.summary.state === "partial") {
|
||||
const notes = [
|
||||
report.summary.degraded.length > 0
|
||||
@@ -220,6 +326,36 @@ export function renderStatusText(
|
||||
return lines.join("\n");
|
||||
}
|
||||
|
||||
export function pipelineLine(pipeline: PipelineSummary): string {
|
||||
const parts: string[] = [pipeline.mode];
|
||||
if (pipeline.forced_by_env)
|
||||
parts.push("(ONE_MEMORY_GRAPH_STORE 强制 ultra)");
|
||||
if (pipeline.mode === "ultra") {
|
||||
const schedule = pipeline.evolve_schedule;
|
||||
parts.push(
|
||||
schedule?.enabled
|
||||
? `,每日 ${schedule.time} 自动进化`
|
||||
: ",自动进化已关闭",
|
||||
);
|
||||
parts.push(
|
||||
pipeline.last_evolve_at
|
||||
? `,上次进化 ${pipeline.last_evolve_at}`
|
||||
: ",尚未进化",
|
||||
);
|
||||
if (pipeline.evolve_running > 0)
|
||||
parts.push(`,进行中 ${pipeline.evolve_running}`);
|
||||
}
|
||||
const changed = Object.entries(pipeline.features);
|
||||
if (changed.length > 0) {
|
||||
parts.push(
|
||||
`;SDK 功能:${changed
|
||||
.map(([name, value]) => `${name}=${value ? "开" : "关"}`)
|
||||
.join("、")}`,
|
||||
);
|
||||
}
|
||||
return parts.join("");
|
||||
}
|
||||
|
||||
function runtimeLine(report: StatusDocument): string {
|
||||
if (report.runtime.state !== "running") {
|
||||
return report.runtime.state === "stopped" ? "未运行" : "暂时无法确认";
|
||||
@@ -234,7 +370,10 @@ function sectionLine(section: Section | undefined): string {
|
||||
if (section.state === "not_checked") return "未检查";
|
||||
if (section.state !== "ok") return "暂时不可用,请稍后重试";
|
||||
return Object.entries(section)
|
||||
.filter(([key]) => key !== "state" && key !== "code")
|
||||
.filter(
|
||||
([key]) =>
|
||||
key !== "state" && key !== "code" && key !== "embedding_switch",
|
||||
)
|
||||
.map(([key, value]) => `${fieldLabel(key)} ${formatValue(value)}`)
|
||||
.join(" ");
|
||||
}
|
||||
@@ -350,9 +489,30 @@ function describeProcessing(
|
||||
reason: block.reason,
|
||||
count: block.count,
|
||||
})),
|
||||
...embeddingSwitchField(summary),
|
||||
};
|
||||
}
|
||||
|
||||
/** The running switch or merge, projected; absent when none runs. */
|
||||
function embeddingSwitchField(
|
||||
summary: MemoryProcessingSummary,
|
||||
): Record<string, unknown> {
|
||||
const value = (summary as unknown as Record<string, unknown>)
|
||||
.embedding_switch;
|
||||
if (!value || typeof value !== "object" || Array.isArray(value)) return {};
|
||||
const item = value as Record<string, unknown>;
|
||||
const kind = item.kind === "merge" ? "merge" : "switch";
|
||||
const state =
|
||||
typeof item.state === "string" && /^[a-z_]{1,32}$/.test(item.state)
|
||||
? item.state
|
||||
: "unknown";
|
||||
const percent =
|
||||
typeof item.percent === "number" && Number.isFinite(item.percent)
|
||||
? Math.min(100, Math.max(0, item.percent))
|
||||
: 0;
|
||||
return { embedding_switch: { kind, state, percent } };
|
||||
}
|
||||
|
||||
/** Projection, not passthrough: a Core that leaks a key must not make us leak. */
|
||||
function describeModel(profiles: ModelProfile[]): Record<string, unknown> {
|
||||
const active = profiles.find((profile) => profile.is_active);
|
||||
|
||||
@@ -12,7 +12,7 @@ dependencies = [
|
||||
"fastapi>=0.141.1",
|
||||
"filelock>=3.32.4",
|
||||
# Public SDK version validated against MemHub's permission and recall contracts.
|
||||
"hy-memory==2.0.9",
|
||||
"hy-memory==2.0.10",
|
||||
# chromadb arrives through the HY SDK; locked here so an unrelated upgrade
|
||||
# cannot move the SDK's store. Built-in embedding consumes numpy directly.
|
||||
"chromadb==1.5.9",
|
||||
|
||||
File diff suppressed because it is too large.
Load diff
@@ -35,7 +35,7 @@ import re # noqa: E402
|
||||
import time # noqa: E402
|
||||
from collections.abc import Iterator, Sequence # noqa: E402
|
||||
from pathlib import Path # noqa: E402
|
||||
from typing import TextIO # noqa: E402
|
||||
from typing import Any, TextIO # noqa: E402
|
||||
|
||||
from one_memory.config import load_settings # noqa: E402
|
||||
from one_memory.runtime_client import RuntimeStatus, inspect_runtime # noqa: E402
|
||||
@@ -584,6 +584,7 @@ def _sdk_migration_status(options: Sequence[str], stream: TextIO) -> int:
|
||||
}
|
||||
else:
|
||||
from one_memory.hy_runtime_config import HY_MEMORY_NAMESPACE
|
||||
from one_memory.hy_sdk_runtime import builtin_retrieval_snapshot
|
||||
from one_memory.sdk_data_plane_migration import (
|
||||
MemoryAuthority,
|
||||
SdkDataPlaneMigrationStore,
|
||||
@@ -595,19 +596,27 @@ def _sdk_migration_status(options: Sequence[str], stream: TextIO) -> int:
|
||||
|
||||
authority = SdkDataPlaneMigrationStore(database_path).authority()
|
||||
if authority is MemoryAuthority.SDK:
|
||||
profile = read_active_profile(database_path)
|
||||
profile = read_active_profile(database_path) or builtin_retrieval_snapshot()
|
||||
leftovers = bool(
|
||||
legacy_tables_present(database_path)
|
||||
or abandoned_staging_present(
|
||||
data_dir=settings.data_dir,
|
||||
namespace=HY_MEMORY_NAMESPACE,
|
||||
)
|
||||
or (profile is not None and legacy_vector_collections(settings.data_dir, profile))
|
||||
or legacy_vector_collections(settings.data_dir, profile)
|
||||
)
|
||||
state = "required" if leftovers else "not_required"
|
||||
else:
|
||||
state = "required"
|
||||
document = {"schema_version": 1, "authority": authority.value, "state": state}
|
||||
from one_memory.embedding_switch import read_open_job
|
||||
|
||||
# Informational. An unfinished embedding switch never makes a
|
||||
# migration required or blocks it: the upgrade installs, `sdk-migrate`
|
||||
# leaves the switch's tables alone, and the new Core resumes the job.
|
||||
switch = read_open_job(database_path)
|
||||
if switch is not None:
|
||||
document["embedding_switch"] = switch
|
||||
from one_memory.sdk_migration_log import write_migration_log
|
||||
|
||||
write_migration_log(
|
||||
@@ -615,7 +624,15 @@ def _sdk_migration_status(options: Sequence[str], stream: TextIO) -> int:
|
||||
writer="core-sdk-migration-status",
|
||||
event="status.probed",
|
||||
outcome="succeeded",
|
||||
details={"authority": document["authority"], "state": document["state"]},
|
||||
details={
|
||||
"authority": document["authority"],
|
||||
"state": document["state"],
|
||||
**(
|
||||
{"embedding_switch": document["embedding_switch"]}
|
||||
if "embedding_switch" in document
|
||||
else {}
|
||||
),
|
||||
},
|
||||
)
|
||||
if as_json:
|
||||
stream.write(f"{json.dumps(document, ensure_ascii=False)}\n")
|
||||
@@ -695,6 +712,12 @@ def _refuse_unmigrated_home(data_dir: Path) -> None:
|
||||
raise SystemExit(SDK_MIGRATION_REQUIRED_EXIT)
|
||||
|
||||
|
||||
def _without_embedding_switch(settings: Any) -> Any:
|
||||
"""The maintenance Core's settings: it never resumes or runs a switch."""
|
||||
copy = getattr(settings, "model_copy", None)
|
||||
return copy(update={"embedding_switch_worker": False}) if callable(copy) else settings
|
||||
|
||||
|
||||
def serve(*, maintenance: bool = False) -> None:
|
||||
import uvicorn
|
||||
|
||||
@@ -739,7 +762,8 @@ def serve(*, maintenance: bool = False) -> None:
|
||||
if launcher is not None:
|
||||
ParentWatchdog(launcher).start()
|
||||
application = create_app(
|
||||
settings,
|
||||
# Only the serving Core runs or resumes an embedding switch.
|
||||
_without_embedding_switch(settings) if maintenance else settings,
|
||||
manage_runtime=not maintenance,
|
||||
)
|
||||
if maintenance:
|
||||
|
||||
@@ -45,6 +45,11 @@ class Settings(BaseSettings):
|
||||
model_request_limit: int | None = Field(default=None, ge=1, le=100_000)
|
||||
model_request_body_byte_limit: int = Field(default=24_576, ge=1)
|
||||
model_request_output_token_limit: int = Field(default=16_384, ge=1)
|
||||
# The embedding switch worker (and the startup merge). The maintenance
|
||||
# Core turns it off: only the serving Core may resume or run a switch.
|
||||
embedding_switch_worker: bool = True
|
||||
# How long a switch waits for model writes already past their claim.
|
||||
embedding_switch_drain_timeout_seconds: float = Field(default=600, gt=0)
|
||||
|
||||
@field_validator("development_origin")
|
||||
@classmethod
|
||||
|
||||
@@ -5,6 +5,7 @@ from typing import Annotated, Literal
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, model_validator
|
||||
|
||||
from one_memory.sdk_search_options import SdkSearchOptions
|
||||
from one_memory.user_scope import DEFAULT_USER_ID, UserId
|
||||
|
||||
ContextKind = Literal["raw", "memory"]
|
||||
@@ -49,12 +50,17 @@ class ContextSearchRequest(BaseModel):
|
||||
return self
|
||||
|
||||
|
||||
class SdkContextSearchRequest(ContextSearchRequest):
|
||||
class SdkContextSearchRequest(ContextSearchRequest, SdkSearchOptions):
|
||||
"""`POST /api/v2/context/search`: the three SDK recall lanes are set apart.
|
||||
|
||||
`limit` is the normal lane's top-k and the cap on returned items;
|
||||
`schema_limit` and `intention_limit` are the SDK's schema (L6) and
|
||||
proactive intention (L7) lanes. A caller's values are passed through.
|
||||
|
||||
The `SdkSearchOptions` fields reach the memory lane's `client.search`;
|
||||
the raw lane takes the floor and the source, session, agent and time
|
||||
filters (never wider than the caller's raw authority). Absent, both
|
||||
lanes make the call they always made.
|
||||
"""
|
||||
|
||||
limit: int = Field(default=10, ge=1, le=20, strict=True)
|
||||
|
||||
@@ -19,19 +19,63 @@ from one_memory.context_models import (
|
||||
SdkMemoryContextItem,
|
||||
SdkRawContextItem,
|
||||
)
|
||||
from one_memory.embedding_health import EMBEDDING_UNAVAILABLE, EmbeddingUnavailableError
|
||||
from one_memory.external_memory_access import external_access
|
||||
from one_memory.hy_runtime_config import HY_MEMORY_READER
|
||||
from one_memory.hy_sdk_runtime import sdk_relevance_floor
|
||||
from one_memory.hy_sdk_types import (
|
||||
HySdkMemoryRecord,
|
||||
parse_sdk_lane_search,
|
||||
parse_sdk_raw_search,
|
||||
parse_sdk_search,
|
||||
)
|
||||
from one_memory.imports.storage import EvidenceStore
|
||||
from one_memory.sdk_search_options import SdkSearchFilters, SearchLane
|
||||
from one_memory.semantic_deadline import SemanticMemoryHit
|
||||
from one_memory.user_scope import owns
|
||||
|
||||
CALL_EVENT_RESPONSE_BUDGET_MS = 250
|
||||
|
||||
#: Pass-through filters the raw lane honours. Scene, layers, reader and the
|
||||
#: schema/working floors describe memory lanes; the raw lane is always the
|
||||
#: chat store's L1 raw layer.
|
||||
_RAW_LANE_FILTERS = ("created_after", "session_ids", "agent_ids", "exclude_sources", "rerank")
|
||||
|
||||
|
||||
def raw_lane_sources(
|
||||
allowed: Sequence[str] | frozenset[str] | None,
|
||||
requested: Sequence[str] | None,
|
||||
) -> list[str] | None:
|
||||
"""The raw lane's source filter: the caller's, never wider than its authority.
|
||||
|
||||
`allowed` empty or None means every source (the call Core always made).
|
||||
None back means nothing is left to search.
|
||||
"""
|
||||
allowed_list = list(allowed or ())
|
||||
if requested is None:
|
||||
return allowed_list
|
||||
chosen = [
|
||||
source for source in dict.fromkeys(requested) if not allowed_list or source in allowed_list
|
||||
]
|
||||
return chosen or None
|
||||
|
||||
|
||||
def raw_lane_overrides(value: SdkSearchFilters) -> dict[str, object]:
|
||||
"""The caller's raw-lane overrides: its floor and its row filters."""
|
||||
overrides = value.filter_kwargs()
|
||||
found = {name: overrides[name] for name in _RAW_LANE_FILTERS if name in overrides}
|
||||
floor = value.lane_floor("raw")
|
||||
if floor is not None:
|
||||
found["min_score"] = floor
|
||||
return found
|
||||
|
||||
|
||||
def hit_namespace(hit: SemanticMemoryHit) -> dict[str, str]:
|
||||
"""The `get` namespace of a hit: nothing for chat, kind and variant otherwise."""
|
||||
if hit.kind in {"experience", "fact"} and hit.variant is not None:
|
||||
return {"kind": hit.kind, "variant": hit.variant}
|
||||
return {}
|
||||
|
||||
|
||||
def sdk_context_router(dependencies: Sequence[Depends]) -> APIRouter:
|
||||
"""SDK-native context: no Evidence/Governance ids or fabricated revisions."""
|
||||
@@ -67,26 +111,36 @@ def sdk_context_router(dependencies: Sequence[Depends]) -> APIRouter:
|
||||
client = await client_handle.client()
|
||||
except Exception:
|
||||
raise HTTPException(503, "sdk_context_unavailable") from None
|
||||
memory_mode = getattr(request.app.state, "memory_mode", None)
|
||||
lite = callable(memory_mode) and memory_mode() == "lite"
|
||||
|
||||
async def raw_lane() -> tuple[SdkRawContextItem, ...]:
|
||||
if "raw" not in value.kinds or not authority.raw_allowed:
|
||||
return ()
|
||||
sources = raw_lane_sources(authority.raw_sources, value.include_sources)
|
||||
if sources is None:
|
||||
return ()
|
||||
raw = await asyncio.to_thread(
|
||||
client.search,
|
||||
value.query,
|
||||
user_ids=[value.user_id],
|
||||
agent_ids=[],
|
||||
include_sources=list(authority.raw_sources or ()),
|
||||
limit=value.limit,
|
||||
min_score=0.0,
|
||||
schema_min_score=0.0,
|
||||
schema_limit=0,
|
||||
intention_limit=0,
|
||||
scene="chat",
|
||||
reader=HY_MEMORY_READER,
|
||||
layers=["l1_raw"],
|
||||
**{
|
||||
"user_ids": [value.user_id],
|
||||
"agent_ids": [],
|
||||
"include_sources": sources,
|
||||
"limit": value.limit,
|
||||
"min_score": 0.0,
|
||||
"schema_min_score": 0.0,
|
||||
"schema_limit": 0,
|
||||
"intention_limit": 0,
|
||||
"scene": "chat",
|
||||
"reader": HY_MEMORY_READER,
|
||||
"layers": ["l1_raw"],
|
||||
**raw_lane_overrides(value),
|
||||
},
|
||||
)
|
||||
hits = tuple(
|
||||
hit for hit in parse_sdk_raw_search(raw) if value.passes_floor("raw", hit.score)
|
||||
)
|
||||
hits = parse_sdk_raw_search(raw)
|
||||
return tuple(
|
||||
SdkRawContextItem(
|
||||
raw_id=hit.raw_id,
|
||||
@@ -104,25 +158,41 @@ def sdk_context_router(dependencies: Sequence[Depends]) -> APIRouter:
|
||||
)
|
||||
|
||||
async def memory_lane() -> tuple[SdkMemoryContextItem, ...]:
|
||||
if "memory" not in value.kinds or not authority.memory_allowed:
|
||||
# lite extracts nothing: its search answers raw turns, which the
|
||||
# raw lane already reads under the raw authority.
|
||||
if "memory" not in value.kinds or not authority.memory_allowed or lite:
|
||||
return ()
|
||||
overrides = value.search_overrides()
|
||||
raw = await asyncio.to_thread(
|
||||
client.search,
|
||||
value.query,
|
||||
user_ids=[value.user_id],
|
||||
agent_ids=[],
|
||||
limit=value.limit,
|
||||
min_score=sdk_relevance_floor(),
|
||||
schema_min_score=sdk_relevance_floor(),
|
||||
schema_limit=value.schema_limit,
|
||||
intention_limit=value.intention_limit,
|
||||
scene="normal",
|
||||
reader=HY_MEMORY_READER,
|
||||
**{
|
||||
"user_ids": [value.user_id],
|
||||
"agent_ids": [],
|
||||
"limit": value.limit,
|
||||
"min_score": sdk_relevance_floor(),
|
||||
"schema_min_score": sdk_relevance_floor(),
|
||||
"schema_limit": value.schema_limit,
|
||||
"intention_limit": value.intention_limit,
|
||||
"scene": "normal",
|
||||
"reader": HY_MEMORY_READER,
|
||||
**overrides,
|
||||
},
|
||||
)
|
||||
hits = (
|
||||
tuple(
|
||||
hit
|
||||
for hit in parse_sdk_lane_search(
|
||||
raw, error_code="hy_memory_search_shape_unknown"
|
||||
)
|
||||
if value.passes_floor(cast(SearchLane, hit.lane), hit.score)
|
||||
)
|
||||
if overrides
|
||||
else parse_sdk_search(raw, error_code="hy_memory_search_shape_unknown")
|
||||
)
|
||||
hits = parse_sdk_search(raw, error_code="hy_memory_search_shape_unknown")
|
||||
items: list[SdkMemoryContextItem] = []
|
||||
for hit in hits:
|
||||
record = await asyncio.to_thread(client.get, hit.memory_id)
|
||||
record = await asyncio.to_thread(client.get, hit.memory_id, **hit_namespace(hit))
|
||||
memory = None if record is None else HySdkMemoryRecord.from_sdk(record)
|
||||
# `get` takes no user: a hit is kept only if it is the caller's.
|
||||
if (
|
||||
@@ -194,6 +264,10 @@ def sdk_context_router(dependencies: Sequence[Depends]) -> APIRouter:
|
||||
)
|
||||
await record(final, timed_out=True)
|
||||
return final
|
||||
for lane in (raw_items, memory_items):
|
||||
if isinstance(lane, EmbeddingUnavailableError):
|
||||
# No lexical fallback: the query could not be embedded.
|
||||
raise HTTPException(503, EMBEDDING_UNAVAILABLE)
|
||||
partial = isinstance(raw_items, Exception) or isinstance(memory_items, Exception)
|
||||
raw_candidates = (
|
||||
() if isinstance(raw_items, Exception) else cast(tuple[SdkContextItem, ...], raw_items)
|
||||
|
||||
@@ -0,0 +1,180 @@
|
||||
"""What the embedding switch needs from the SDK's re-embed module, in Core terms.
|
||||
|
||||
The switch worker, the startup merge and the retention sweep only talk to a
|
||||
`ReembedBackend`. The production one (`hy_sdk_migration.SdkReembedBackend`)
|
||||
is the single place that calls `hy_memory`'s re-embed functions; tests use
|
||||
an in-memory one. Every space is named explicitly -- collection names and
|
||||
graph paths -- and nothing here derives a name from a provider or endpoint.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable, Mapping
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Protocol
|
||||
|
||||
from one_memory.embedding_spaces import EmbeddingSpace, space_collection_names
|
||||
from one_memory.model_profiles import ModelProfileSnapshot
|
||||
|
||||
|
||||
class ReembedUnavailableError(RuntimeError):
|
||||
"""The installed SDK has no re-embed module (before hy-memory 2.0.10)."""
|
||||
|
||||
def __init__(self, code: str = "sdk_reembed_unavailable") -> None:
|
||||
super().__init__(code)
|
||||
self.code = code
|
||||
|
||||
|
||||
class ReembedCancelledError(RuntimeError):
|
||||
"""A cancel was seen at a batch boundary; every finished batch is written."""
|
||||
|
||||
|
||||
class EmbeddingWidthError(RuntimeError):
|
||||
"""The embedding service returned vectors of another width than the target's."""
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class PhysicalSpace:
|
||||
"""One space's collections, named explicitly."""
|
||||
|
||||
collection_base: str
|
||||
name_dims: int | None
|
||||
#: The width of the vectors in it, or to be written to it.
|
||||
dims: int | None
|
||||
model: str | None = None
|
||||
#: Explicit role -> collection name; None derives them from the base.
|
||||
collections: Mapping[str, str] | None = None
|
||||
|
||||
def names(self) -> dict[str, str]:
|
||||
if self.collections is not None:
|
||||
return dict(self.collections)
|
||||
return space_collection_names(self.collection_base, self.name_dims)
|
||||
|
||||
@classmethod
|
||||
def of(cls, space: EmbeddingSpace) -> PhysicalSpace:
|
||||
return cls(
|
||||
collection_base=space.collection_base,
|
||||
name_dims=space.name_dims,
|
||||
dims=space.dims if space.dims is not None else space.name_dims,
|
||||
model=space.model,
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ScannedSpace:
|
||||
"""A space found in the vector store, grouped from its collection names."""
|
||||
|
||||
collection_base: str
|
||||
name_dims: int | None
|
||||
#: The stored width when every non-empty role agrees.
|
||||
width: int | None
|
||||
counts: Mapping[str, int]
|
||||
collections: Mapping[str, str]
|
||||
per_user: Mapping[str, int] = field(default_factory=dict)
|
||||
|
||||
@property
|
||||
def vectors(self) -> int:
|
||||
return sum(self.counts.values())
|
||||
|
||||
@property
|
||||
def key(self) -> tuple[str, int]:
|
||||
return (self.collection_base, self.name_dims or 0)
|
||||
|
||||
def physical(self) -> PhysicalSpace:
|
||||
return PhysicalSpace(
|
||||
collection_base=self.collection_base,
|
||||
name_dims=self.name_dims,
|
||||
dims=self.width if self.width is not None else self.name_dims,
|
||||
collections=dict(self.collections),
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ReembedStep:
|
||||
"""Progress of one role (vectors) or table (graph)."""
|
||||
|
||||
phase: str
|
||||
role: str
|
||||
done: int
|
||||
total: int
|
||||
embedded: int = 0
|
||||
reused: int = 0
|
||||
kept: int = 0
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class VectorCheck:
|
||||
per_role: Mapping[str, Mapping[str, int]]
|
||||
#: Target records per metadata `user_id`.
|
||||
per_user: Mapping[str, int]
|
||||
problems: tuple[str, ...]
|
||||
warnings: tuple[str, ...] = ()
|
||||
self_recall_hits: int = 0
|
||||
self_recall_samples: int = 0
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class GraphCheck:
|
||||
ok: bool
|
||||
problems: tuple[str, ...] = ()
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class EmbeddingTarget:
|
||||
"""The embedding that fills a space: the profile to reach it, and its key."""
|
||||
|
||||
profile: ModelProfileSnapshot
|
||||
space: EmbeddingSpace
|
||||
api_key: str | None
|
||||
|
||||
|
||||
Cancel = Callable[[], bool]
|
||||
OnStep = Callable[[ReembedStep], None]
|
||||
|
||||
|
||||
class ReembedBackend(Protocol):
|
||||
def unavailable_reason(self) -> str | None:
|
||||
"""None when re-embedding works; otherwise the code that says why not."""
|
||||
|
||||
async def scan(self, *, per_user: bool = False) -> tuple[ScannedSpace, ...]:
|
||||
"""Every space with collections in the home's vector store. Read-only."""
|
||||
|
||||
async def diff(self, source: PhysicalSpace, target: PhysicalSpace) -> dict[str, int]:
|
||||
"""Per role, how many ids the source holds that the target does not."""
|
||||
|
||||
async def reembed(
|
||||
self,
|
||||
source: PhysicalSpace,
|
||||
target: PhysicalSpace,
|
||||
*,
|
||||
embedding: EmbeddingTarget,
|
||||
cancel: Cancel,
|
||||
on_step: OnStep,
|
||||
batch_size: int = 64,
|
||||
) -> None:
|
||||
"""Embed every source record the target lacks into the target. Target wins."""
|
||||
|
||||
async def verify(
|
||||
self, source: PhysicalSpace, target: PhysicalSpace, *, embedding: EmbeddingTarget
|
||||
) -> VectorCheck: ...
|
||||
|
||||
async def graph_vectors(self, path: Path) -> int | None:
|
||||
"""Embeddings a graph holds, or None when there is no graph there."""
|
||||
|
||||
async def reembed_graph(
|
||||
self,
|
||||
source: Path,
|
||||
target: Path,
|
||||
*,
|
||||
embedding: EmbeddingTarget,
|
||||
target_dims: int,
|
||||
cancel: Cancel,
|
||||
on_step: OnStep,
|
||||
) -> None: ...
|
||||
|
||||
async def verify_graph(self, source: Path, target: Path, *, target_dims: int) -> GraphCheck: ...
|
||||
|
||||
async def drop(self, space: PhysicalSpace) -> None: ...
|
||||
|
||||
async def drop_graph(self, path: Path) -> None: ...
|
||||
@@ -0,0 +1,147 @@
|
||||
"""`embedding_unavailable`: a search that could not embed its query says so.
|
||||
|
||||
With a remote embedding model, every search embeds the query through that
|
||||
service. When the service is down, the SDK swallows the failure and answers
|
||||
with no hits, which reads as "you have no such memory". There is no lexical
|
||||
fallback (owner decision): the search fails with the explicit code
|
||||
`embedding_unavailable`, the API answers 503 with that code, and the App
|
||||
shows a toast.
|
||||
|
||||
The check is only made when a search came back empty or raised, so a working
|
||||
service costs nothing on searches that find something. It sends one tiny
|
||||
embedding request to the same endpoint, model and width the read client
|
||||
uses; the answer is remembered briefly so a burst of searches sends one.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Callable, Mapping
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
from one_memory.http_client import create_client
|
||||
from one_memory.hy_sdk_types import HySdkError
|
||||
from one_memory.model_profiles import ModelProfileSnapshot, ProviderKind
|
||||
|
||||
EMBEDDING_UNAVAILABLE = "embedding_unavailable"
|
||||
_OK_TTL_S = 30.0
|
||||
_FAILED_TTL_S = 10.0
|
||||
_TIMEOUT_S = 8.0
|
||||
|
||||
|
||||
def _no_transport() -> httpx.BaseTransport | None:
|
||||
return None
|
||||
|
||||
|
||||
#: Tests replace this to answer the check without a network.
|
||||
transport_factory: Callable[[], httpx.BaseTransport | None] = _no_transport
|
||||
|
||||
|
||||
class EmbeddingUnavailableError(HySdkError):
|
||||
"""The live space's embedding service did not answer; search cannot run."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__(EMBEDDING_UNAVAILABLE)
|
||||
|
||||
|
||||
class EmbeddingHealth:
|
||||
"""Whether a remote embedding answers, remembered for a few seconds."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._lock = threading.Lock()
|
||||
self._known: dict[tuple[str, str, int | None], tuple[bool, float]] = {}
|
||||
|
||||
def available(
|
||||
self,
|
||||
profile: ModelProfileSnapshot,
|
||||
*,
|
||||
endpoint: str | None,
|
||||
name_dims: int | None,
|
||||
api_key: str | None,
|
||||
) -> bool:
|
||||
provider = profile.embedding_provider_kind or profile.provider_kind
|
||||
if provider is ProviderKind.LOCAL_ONNX or endpoint is None:
|
||||
return True
|
||||
model = profile.embedding_model or ""
|
||||
key = (endpoint, model, name_dims)
|
||||
now = time.monotonic()
|
||||
with self._lock:
|
||||
known = self._known.get(key)
|
||||
if known is not None and known[1] > now:
|
||||
return known[0]
|
||||
ok = _answers(
|
||||
endpoint,
|
||||
model=model,
|
||||
name_dims=name_dims,
|
||||
api_key="ollama" if provider is ProviderKind.OLLAMA else api_key,
|
||||
)
|
||||
with self._lock:
|
||||
self._known[key] = (ok, now + (_OK_TTL_S if ok else _FAILED_TTL_S))
|
||||
return ok
|
||||
|
||||
def forget(self) -> None:
|
||||
with self._lock:
|
||||
self._known.clear()
|
||||
|
||||
|
||||
def _answers(endpoint: str, *, model: str, name_dims: int | None, api_key: str | None) -> bool:
|
||||
body: dict[str, Any] = {"model": model, "input": ["ping"]}
|
||||
if name_dims:
|
||||
body["dimensions"] = name_dims
|
||||
headers = {"Authorization": f"Bearer {api_key}"} if api_key else {}
|
||||
transport = transport_factory()
|
||||
try:
|
||||
with create_client(
|
||||
timeout=_TIMEOUT_S,
|
||||
**({"transport": transport} if transport is not None else {}),
|
||||
) as client:
|
||||
response = client.post(f"{endpoint.rstrip('/')}/embeddings", json=body, headers=headers)
|
||||
if response.status_code != 200:
|
||||
return False
|
||||
data = response.json().get("data")
|
||||
return bool(isinstance(data, list) and data and data[0].get("embedding"))
|
||||
except (httpx.HTTPError, ValueError, AttributeError):
|
||||
return False
|
||||
|
||||
|
||||
def has_hits(result: object) -> bool:
|
||||
"""Whether an SDK search result holds any hit, whatever its scene's shape."""
|
||||
if not isinstance(result, Mapping):
|
||||
return bool(result)
|
||||
for key in ("memories", "facts", "experiences"):
|
||||
value = result.get(key)
|
||||
if isinstance(value, list) and value:
|
||||
return True
|
||||
if isinstance(value, Mapping) and any(
|
||||
isinstance(item, list) and item for item in value.values()
|
||||
):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
class EmbeddingGuardedClient:
|
||||
"""An SDK client whose `search` reports a dead embedding instead of nothing.
|
||||
|
||||
Everything else is the client's own attribute, unchanged.
|
||||
"""
|
||||
|
||||
def __init__(self, client: Any, available: Callable[[], bool]) -> None:
|
||||
self._client = client
|
||||
self._available = available
|
||||
|
||||
def search(self, *args: Any, **kwargs: Any) -> Any:
|
||||
try:
|
||||
result = self._client.search(*args, **kwargs)
|
||||
except Exception:
|
||||
if not self._available():
|
||||
raise EmbeddingUnavailableError() from None
|
||||
raise
|
||||
if not has_hits(result) and not self._available():
|
||||
raise EmbeddingUnavailableError()
|
||||
return result
|
||||
|
||||
def __getattr__(self, name: str) -> Any:
|
||||
return getattr(self._client, name)
|
||||
@@ -0,0 +1,554 @@
|
||||
"""Which embedding space the home reads and writes: an explicit registry.
|
||||
|
||||
The SDK names its vector collections after `vector_store.collection_name`
|
||||
and the configured width, so the space used to follow whatever the active
|
||||
profile hashed to (`_collection_base`). Any change to that identity -- a
|
||||
probe recording the measured width, a new endpoint -- silently opened new,
|
||||
empty collections beside the old ones, and nothing carried the memories over.
|
||||
|
||||
The live space is pinned here instead, in the main Core database beside the
|
||||
model profiles, so a switch commits the new space and the profile that fills
|
||||
it in one transaction. The identity of a space is its embedding model and
|
||||
width only: a new endpoint or provider serving the same model at the same
|
||||
width keeps the space.
|
||||
|
||||
Existing homes keep their collections byte-for-byte. The first start registers
|
||||
the space the active profile (or the built-in model) already opened under its
|
||||
existing four-attribute name; only spaces created by a switch use the
|
||||
model-and-width name. A space whose width was never declared ("auto") keeps
|
||||
its unsuffixed name and only records the width it was measured at.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import sqlite3
|
||||
from collections.abc import Iterator, Mapping
|
||||
from contextlib import contextmanager
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from pathlib import Path
|
||||
from typing import Literal
|
||||
from uuid import uuid4
|
||||
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
|
||||
from one_memory.database import take_the_write_lock_first
|
||||
from one_memory.database_access import main_database
|
||||
from one_memory.embedding_assets import BUILTIN_EMBEDDING_DIMENSIONS, BUILTIN_EMBEDDING_MODEL
|
||||
from one_memory.memory_server.migrations import embedding_spaces_ddl
|
||||
from one_memory.model_profiles import ModelProfileSnapshot, ProviderKind
|
||||
|
||||
#: How long a space replaced by a switch or folded in by a merge is kept.
|
||||
RETIRED_SPACE_RETENTION = timedelta(days=14)
|
||||
|
||||
SpaceState = Literal["live", "building", "stray", "retired", "dropped"]
|
||||
SpaceNaming = Literal["legacy", "v2", "discovered"]
|
||||
|
||||
#: Every vector collection one space consists of, as the SDK names them.
|
||||
VECTOR_ROLES = (
|
||||
"memory",
|
||||
"entities",
|
||||
"coding_keys",
|
||||
"coding_keys_cv",
|
||||
"working_facts",
|
||||
"working_facts_cv",
|
||||
)
|
||||
|
||||
#: The graph directory of the first space; later spaces get their own.
|
||||
DEFAULT_GRAPH_DIR = "graph"
|
||||
|
||||
EMBEDDING_SPACES_SCHEMA = embedding_spaces_ddl()
|
||||
|
||||
|
||||
#: Bumped by every registry write in this process, so cached live rows are
|
||||
#: reread after a change made through any connection (a probe's, a commit's).
|
||||
_CHANGES = [0]
|
||||
|
||||
|
||||
def _changed() -> None:
|
||||
_CHANGES[0] += 1
|
||||
|
||||
|
||||
class EmbeddingSpaceError(RuntimeError):
|
||||
"""A registry refusal with a stable, product-facing code."""
|
||||
|
||||
def __init__(self, code: str) -> None:
|
||||
super().__init__(code)
|
||||
self.code = code
|
||||
|
||||
|
||||
class EmbeddingIdentity(BaseModel):
|
||||
"""What makes two vectors comparable: the model and the width.
|
||||
|
||||
`dims=None` is a width nobody declared or measured yet; it matches any
|
||||
width, because the collection it names was opened without one.
|
||||
"""
|
||||
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
model: str
|
||||
dims: int | None = None
|
||||
|
||||
def matches(self, other: EmbeddingIdentity) -> bool:
|
||||
if self.model != other.model:
|
||||
return False
|
||||
return self.dims is None or other.dims is None or self.dims == other.dims
|
||||
|
||||
|
||||
class EmbeddingSpace(BaseModel):
|
||||
"""One embedding space: the physical collections and graph one model fills."""
|
||||
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
space_id: str
|
||||
state: SpaceState
|
||||
#: None only for a space found on disk that no registry row named.
|
||||
model: str | None = None
|
||||
#: The vector width stored in this space, once known.
|
||||
dims: int | None = None
|
||||
#: The width the collection names and embedding requests carry. None is
|
||||
#: an "auto" space: unsuffixed names and no `dimensions` request field.
|
||||
name_dims: int | None = None
|
||||
collection_base: str
|
||||
naming: SpaceNaming = "legacy"
|
||||
graph_dir: str | None = DEFAULT_GRAPH_DIR
|
||||
provider_kind: str | None = None
|
||||
#: The profile version whose embedding fills this space; None = built-in.
|
||||
profile_id: str | None = None
|
||||
profile_revision: int | None = None
|
||||
#: Vectors counted the last time the space was scanned.
|
||||
vectors: int | None = None
|
||||
created_at: str
|
||||
live_since: str | None = None
|
||||
retired_at: str | None = None
|
||||
keep_until: str | None = None
|
||||
dropped_at: str | None = None
|
||||
|
||||
@property
|
||||
def identity(self) -> EmbeddingIdentity | None:
|
||||
if self.model is None:
|
||||
return None
|
||||
return EmbeddingIdentity(model=self.model, dims=self.dims)
|
||||
|
||||
@property
|
||||
def builtin(self) -> bool:
|
||||
return self.provider_kind == ProviderKind.LOCAL_ONNX.value and self.profile_id is None
|
||||
|
||||
def collection_names(self) -> dict[str, str]:
|
||||
return space_collection_names(self.collection_base, self.name_dims)
|
||||
|
||||
|
||||
def space_collection_names(base: str, name_dims: int | None) -> dict[str, str]:
|
||||
"""The SDK's physical collection name per role for one space.
|
||||
|
||||
Mirrors `VectorStoreBase` (the `_<dims>` suffix), the entity store and the
|
||||
coding stores. The SDK's own `vector_space_names` is checked against this
|
||||
in the migration adapter when the installed SDK provides it.
|
||||
"""
|
||||
suffix = f"_{name_dims}" if name_dims else ""
|
||||
return {
|
||||
"memory": f"{base}{suffix}",
|
||||
"entities": f"{base}{suffix}_entities",
|
||||
"coding_keys": f"{base}_coding_keys{suffix}",
|
||||
"coding_keys_cv": f"{base}_coding_keys_cv{suffix}",
|
||||
"working_facts": f"{base}_working_facts{suffix}",
|
||||
"working_facts_cv": f"{base}_working_facts_cv{suffix}",
|
||||
}
|
||||
|
||||
|
||||
def embedding_identity(profile: ModelProfileSnapshot) -> EmbeddingIdentity:
|
||||
"""The embedding identity a profile asks for: model name and width only."""
|
||||
provider = profile.embedding_provider_kind or profile.provider_kind
|
||||
local = provider is ProviderKind.LOCAL_ONNX
|
||||
model = profile.embedding_model or (BUILTIN_EMBEDDING_MODEL if local else None)
|
||||
if not model:
|
||||
raise EmbeddingSpaceError("hy_sdk_embedding_model_required")
|
||||
dims = profile.embedding_dimensions
|
||||
if local and dims is None:
|
||||
dims = BUILTIN_EMBEDDING_DIMENSIONS
|
||||
return EmbeddingIdentity(model=model, dims=dims)
|
||||
|
||||
|
||||
def v2_collection_base(identity: EmbeddingIdentity, generation: int = 0) -> str:
|
||||
"""The base name a switch gives a new space: model and width, nothing else."""
|
||||
digest = hashlib.sha256(f"{identity.model}\0{identity.dims or 'auto'}".encode()).hexdigest()
|
||||
base = f"memhub_e_{digest[:16]}"
|
||||
return f"{base}_g{generation}" if generation else base
|
||||
|
||||
|
||||
def graph_dir_for(space_id: str) -> str:
|
||||
return f"graph-{space_id[:8]}"
|
||||
|
||||
|
||||
def derived_space(profile: ModelProfileSnapshot) -> EmbeddingSpace:
|
||||
"""The space `profile` opened before the registry existed, under its old name.
|
||||
|
||||
Never stored as is: bootstrap stores it, and offline tools that run beside
|
||||
a home with no registry yet read through it.
|
||||
"""
|
||||
from one_memory.hy_sdk_runtime import legacy_collection_base
|
||||
|
||||
provider = profile.embedding_provider_kind or profile.provider_kind
|
||||
identity = embedding_identity(profile)
|
||||
builtin = profile.profile_id == _BUILTIN_PROFILE_ID
|
||||
return EmbeddingSpace(
|
||||
space_id="derived",
|
||||
state="live",
|
||||
model=identity.model,
|
||||
dims=identity.dims,
|
||||
name_dims=identity.dims,
|
||||
collection_base=legacy_collection_base(profile),
|
||||
naming="legacy",
|
||||
graph_dir=DEFAULT_GRAPH_DIR,
|
||||
provider_kind=provider.value,
|
||||
profile_id=None if builtin else profile.profile_id,
|
||||
profile_revision=None if builtin else profile.revision,
|
||||
created_at=_now(),
|
||||
)
|
||||
|
||||
|
||||
_BUILTIN_PROFILE_ID = "builtin-retrieval"
|
||||
|
||||
|
||||
def _now() -> str:
|
||||
return datetime.now(UTC).isoformat()
|
||||
|
||||
|
||||
def _row(row: sqlite3.Row | Mapping[str, object]) -> EmbeddingSpace:
|
||||
return EmbeddingSpace.model_validate(dict(row))
|
||||
|
||||
|
||||
def ensure_schema(connection: sqlite3.Connection) -> None:
|
||||
connection.executescript(EMBEDDING_SPACES_SCHEMA)
|
||||
|
||||
|
||||
def registry_present(connection: sqlite3.Connection) -> bool:
|
||||
return (
|
||||
connection.execute(
|
||||
"SELECT 1 FROM sqlite_master WHERE type='table' AND name='embedding_spaces'"
|
||||
).fetchone()
|
||||
is not None
|
||||
)
|
||||
|
||||
|
||||
def live_space_in(connection: sqlite3.Connection) -> EmbeddingSpace | None:
|
||||
if not registry_present(connection):
|
||||
return None
|
||||
previous = connection.row_factory
|
||||
connection.row_factory = sqlite3.Row
|
||||
try:
|
||||
row = connection.execute("SELECT * FROM embedding_spaces WHERE state='live'").fetchone()
|
||||
finally:
|
||||
connection.row_factory = previous
|
||||
return _row(row) if row is not None else None
|
||||
|
||||
|
||||
def read_live_space(database_path: Path) -> EmbeddingSpace | None:
|
||||
"""The live space, read only, or None for a home with no registry yet."""
|
||||
if not database_path.is_file():
|
||||
return None
|
||||
with main_database(database_path).read(row_factory=sqlite3.Row) as connection:
|
||||
return live_space_in(connection)
|
||||
|
||||
|
||||
class EmbeddingSpaceStore:
|
||||
"""The `embedding_spaces` table of the main Core database."""
|
||||
|
||||
def __init__(self, database_path: Path) -> None:
|
||||
self.database_path = database_path
|
||||
#: The live row, read once per change: every read and write asks for
|
||||
#: it. Only this process writes the registry while Core serves.
|
||||
self._live_cache: tuple[int, EmbeddingSpace] | None = None
|
||||
with self._connect(writing=True) as connection:
|
||||
ensure_schema(connection)
|
||||
|
||||
def invalidate(self) -> None:
|
||||
_changed()
|
||||
|
||||
@contextmanager
|
||||
def _connect(
|
||||
self, *, writing: bool = False, lock_first: bool = True
|
||||
) -> Iterator[sqlite3.Connection]:
|
||||
database = main_database(self.database_path)
|
||||
checkout = (
|
||||
database.write(row_factory=sqlite3.Row)
|
||||
if writing
|
||||
else database.read(row_factory=sqlite3.Row)
|
||||
)
|
||||
with checkout as connection, connection:
|
||||
if writing and lock_first:
|
||||
take_the_write_lock_first(connection)
|
||||
yield connection
|
||||
|
||||
@contextmanager
|
||||
def transaction(self) -> Iterator[sqlite3.Connection]:
|
||||
"""One `BEGIN IMMEDIATE` on the main database, committed on success."""
|
||||
_changed()
|
||||
try:
|
||||
with self._connect(writing=True, lock_first=False) as connection:
|
||||
connection.execute("BEGIN IMMEDIATE")
|
||||
yield connection
|
||||
finally:
|
||||
_changed()
|
||||
|
||||
def live(self) -> EmbeddingSpace | None:
|
||||
generation = _CHANGES[0]
|
||||
cached = self._live_cache
|
||||
if cached is not None and cached[0] == generation:
|
||||
return cached[1]
|
||||
with self._connect() as connection:
|
||||
space = live_space_in(connection)
|
||||
self._live_cache = (generation, space) if space is not None else None
|
||||
return space
|
||||
|
||||
def require_live(self) -> EmbeddingSpace:
|
||||
space = self.live()
|
||||
if space is None:
|
||||
raise EmbeddingSpaceError("embedding_space_not_bootstrapped")
|
||||
return space
|
||||
|
||||
def get(self, space_id: str) -> EmbeddingSpace | None:
|
||||
with self._connect() as connection:
|
||||
row = connection.execute(
|
||||
"SELECT * FROM embedding_spaces WHERE space_id=?", (space_id,)
|
||||
).fetchone()
|
||||
return _row(row) if row is not None else None
|
||||
|
||||
def list_spaces(self, *, include_dropped: bool = False) -> list[EmbeddingSpace]:
|
||||
with self._connect() as connection:
|
||||
rows = connection.execute(
|
||||
"SELECT * FROM embedding_spaces "
|
||||
+ ("" if include_dropped else "WHERE state != 'dropped' ")
|
||||
+ "ORDER BY created_at, space_id"
|
||||
).fetchall()
|
||||
return [_row(row) for row in rows]
|
||||
|
||||
def bootstrap(self, active: ModelProfileSnapshot | None) -> EmbeddingSpace:
|
||||
"""Register the space this home already uses, once, under its existing name.
|
||||
|
||||
`active` is the active profile, or None for the built-in model. The
|
||||
collection name is derived exactly as before the registry, so nothing
|
||||
on disk moves.
|
||||
"""
|
||||
from one_memory.hy_sdk_runtime import builtin_retrieval_snapshot
|
||||
|
||||
with self.transaction() as connection:
|
||||
existing = live_space_in(connection)
|
||||
if existing is not None:
|
||||
return existing
|
||||
derived = derived_space(active or builtin_retrieval_snapshot())
|
||||
space = derived.model_copy(
|
||||
update={"space_id": str(uuid4()), "live_since": derived.created_at}
|
||||
)
|
||||
_insert(connection, space)
|
||||
return space
|
||||
|
||||
def record_width(self, space_id: str, dims: int) -> EmbeddingSpace:
|
||||
"""A relabel: the space's stored width is now known. Names never change."""
|
||||
with self.transaction() as connection:
|
||||
record_width_in(connection, space_id, dims)
|
||||
space = self.get(space_id)
|
||||
if space is None: # pragma: no cover - the row was just updated
|
||||
raise EmbeddingSpaceError("embedding_space_not_found")
|
||||
return space
|
||||
|
||||
def record_vectors(self, counts: Mapping[str, int]) -> None:
|
||||
"""Cache the vector count per space id from the last scan (display only)."""
|
||||
if not counts:
|
||||
return
|
||||
_changed()
|
||||
with self._connect(writing=True) as connection:
|
||||
connection.executemany(
|
||||
"UPDATE embedding_spaces SET vectors=? WHERE space_id=?",
|
||||
[(count, space_id) for space_id, count in counts.items()],
|
||||
)
|
||||
|
||||
def register_stray(
|
||||
self,
|
||||
*,
|
||||
collection_base: str,
|
||||
name_dims: int | None,
|
||||
dims: int | None,
|
||||
vectors: int,
|
||||
) -> EmbeddingSpace:
|
||||
"""Name a non-empty space the scan found that no row described."""
|
||||
with self.transaction() as connection:
|
||||
row = connection.execute(
|
||||
"SELECT * FROM embedding_spaces WHERE collection_base=? "
|
||||
"AND COALESCE(name_dims,0)=? AND state != 'dropped'",
|
||||
(collection_base, name_dims or 0),
|
||||
).fetchone()
|
||||
if row is not None:
|
||||
return _row(row)
|
||||
space = EmbeddingSpace(
|
||||
space_id=str(uuid4()),
|
||||
state="stray",
|
||||
dims=dims,
|
||||
name_dims=name_dims,
|
||||
collection_base=collection_base,
|
||||
naming="discovered",
|
||||
graph_dir=None,
|
||||
vectors=vectors,
|
||||
created_at=_now(),
|
||||
)
|
||||
_insert(connection, space)
|
||||
return space
|
||||
|
||||
def mark_dropped(self, space_id: str) -> None:
|
||||
_changed()
|
||||
with self._connect(writing=True) as connection:
|
||||
connection.execute(
|
||||
"UPDATE embedding_spaces SET state='dropped', dropped_at=? "
|
||||
"WHERE space_id=? AND state IN ('retired','building')",
|
||||
(_now(), space_id),
|
||||
)
|
||||
|
||||
def expired(self, *, now: datetime | None = None) -> list[EmbeddingSpace]:
|
||||
moment = (now or datetime.now(UTC)).isoformat()
|
||||
with self._connect() as connection:
|
||||
rows = connection.execute(
|
||||
"SELECT * FROM embedding_spaces WHERE state='retired' "
|
||||
"AND keep_until IS NOT NULL AND keep_until <= ?",
|
||||
(moment,),
|
||||
).fetchall()
|
||||
return [_row(row) for row in rows]
|
||||
|
||||
|
||||
def registry_for(profiles: object) -> EmbeddingSpaceStore | None:
|
||||
"""The registry beside a profile store, or None for a store with no database."""
|
||||
path = getattr(profiles, "database_path", None)
|
||||
return EmbeddingSpaceStore(path) if isinstance(path, Path) else None
|
||||
|
||||
|
||||
def _insert(connection: sqlite3.Connection, space: EmbeddingSpace) -> None:
|
||||
values = space.model_dump()
|
||||
columns = ", ".join(values)
|
||||
marks = ", ".join("?" for _ in values)
|
||||
connection.execute(
|
||||
f"INSERT INTO embedding_spaces ({columns}) VALUES ({marks})",
|
||||
tuple(values.values()),
|
||||
)
|
||||
|
||||
|
||||
def insert_space_in(connection: sqlite3.Connection, space: EmbeddingSpace) -> None:
|
||||
_changed()
|
||||
_insert(connection, space)
|
||||
|
||||
|
||||
def record_width_in(connection: sqlite3.Connection, space_id: str, dims: int) -> None:
|
||||
_changed()
|
||||
changed = connection.execute(
|
||||
"UPDATE embedding_spaces SET dims=? WHERE space_id=? AND (dims IS NULL OR dims=?)",
|
||||
(dims, space_id, dims),
|
||||
)
|
||||
if changed.rowcount != 1:
|
||||
raise EmbeddingSpaceError("embedding_space_width_conflict")
|
||||
|
||||
|
||||
def space_in(connection: sqlite3.Connection, space_id: str) -> EmbeddingSpace | None:
|
||||
previous = connection.row_factory
|
||||
connection.row_factory = sqlite3.Row
|
||||
try:
|
||||
row = connection.execute(
|
||||
"SELECT * FROM embedding_spaces WHERE space_id=?", (space_id,)
|
||||
).fetchone()
|
||||
finally:
|
||||
connection.row_factory = previous
|
||||
return _row(row) if row is not None else None
|
||||
|
||||
|
||||
def spaces_in(connection: sqlite3.Connection) -> list[EmbeddingSpace]:
|
||||
previous = connection.row_factory
|
||||
connection.row_factory = sqlite3.Row
|
||||
try:
|
||||
rows = connection.execute("SELECT * FROM embedding_spaces").fetchall()
|
||||
finally:
|
||||
connection.row_factory = previous
|
||||
return [_row(row) for row in rows]
|
||||
|
||||
|
||||
def set_state_in(
|
||||
connection: sqlite3.Connection,
|
||||
space_id: str,
|
||||
*,
|
||||
state: SpaceState,
|
||||
expected: tuple[SpaceState, ...],
|
||||
now: datetime | None = None,
|
||||
) -> None:
|
||||
"""Move one space between states, refusing any it was not expected in."""
|
||||
_changed()
|
||||
moment = now or datetime.now(UTC)
|
||||
updates: dict[str, object] = {"state": state}
|
||||
if state == "live":
|
||||
updates["live_since"] = moment.isoformat()
|
||||
elif state == "retired":
|
||||
updates["retired_at"] = moment.isoformat()
|
||||
updates["keep_until"] = (moment + RETIRED_SPACE_RETENTION).isoformat()
|
||||
elif state == "dropped":
|
||||
updates["dropped_at"] = moment.isoformat()
|
||||
assignments = ", ".join(f"{column}=?" for column in updates)
|
||||
marks = ",".join("?" for _ in expected)
|
||||
changed = connection.execute(
|
||||
f"UPDATE embedding_spaces SET {assignments} WHERE space_id=? AND state IN ({marks})",
|
||||
(*updates.values(), space_id, *expected),
|
||||
)
|
||||
if changed.rowcount != 1:
|
||||
raise EmbeddingSpaceError("embedding_space_state_changed")
|
||||
|
||||
|
||||
def next_v2_base(
|
||||
connection: sqlite3.Connection,
|
||||
identity: EmbeddingIdentity,
|
||||
*,
|
||||
stored_names: frozenset[str],
|
||||
) -> str:
|
||||
"""A base for a new space no registry row and no stored collection uses.
|
||||
|
||||
Switching back to a model whose retired space is still kept gets a fresh
|
||||
`_g<n>` base, so stale retired vectors are never mixed with fresh ones.
|
||||
"""
|
||||
used = {space.collection_base for space in spaces_in(connection)}
|
||||
generation = 0
|
||||
while True:
|
||||
base = v2_collection_base(identity, generation)
|
||||
names = space_collection_names(base, identity.dims).values()
|
||||
if base not in used and not any(name in stored_names for name in names):
|
||||
return base
|
||||
generation += 1
|
||||
|
||||
|
||||
def probe_width_allowed(
|
||||
connection: sqlite3.Connection,
|
||||
profile_id: str,
|
||||
width: int,
|
||||
*,
|
||||
data_dir: Path,
|
||||
) -> bool:
|
||||
"""Whether a probe may record `width` on a profile without moving the space.
|
||||
|
||||
A profile that fills the live space (the active one, or the one the space
|
||||
was built with) may record its measured width only when it is the
|
||||
space's: equal to the width already known, or -- for an "auto" space
|
||||
whose width was never recorded -- equal to the width Chroma stored, or
|
||||
any width when nothing is stored yet. That is a relabel; the names never
|
||||
change. Any other width needs a switch (B2: it used to silently open a
|
||||
new, empty collection). Every other profile records what it measured.
|
||||
"""
|
||||
live = live_space_in(connection)
|
||||
if live is None:
|
||||
return True
|
||||
active = connection.execute(
|
||||
"SELECT profile_id FROM model_profiles WHERE is_active=1"
|
||||
).fetchone()
|
||||
fills_live = live.profile_id == profile_id or (active is not None and active[0] == profile_id)
|
||||
if not fills_live:
|
||||
return True
|
||||
if live.dims is not None:
|
||||
return live.dims == width
|
||||
from one_memory.sdk_vector_reattach import stored_collection_width
|
||||
|
||||
stored = stored_collection_width(data_dir, live.collection_names()["memory"])
|
||||
if stored is not None and stored != width:
|
||||
return False
|
||||
record_width_in(connection, live.space_id, width)
|
||||
return True
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,161 @@
|
||||
"""HTTP surface of the embedding switch: spaces, plans, jobs, cleanup.
|
||||
|
||||
The contract the App and the CLI build on is
|
||||
`memhub/docs/contracts/embedding-switch-v1.md`. Errors keep Core's shape,
|
||||
`{"detail": "<code>"}`; a refusal that computed a plan also carries it as
|
||||
`"plan"`, so a client can show the confirmation without asking again.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from collections.abc import Sequence
|
||||
from typing import Any, cast
|
||||
|
||||
from fastapi import APIRouter, Request, Response
|
||||
from fastapi.params import Depends
|
||||
from fastapi.responses import JSONResponse
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
from one_memory.embedding_backend import ReembedUnavailableError
|
||||
from one_memory.embedding_spaces import EmbeddingSpaceError
|
||||
from one_memory.embedding_switch import (
|
||||
EmbeddingSwitchError,
|
||||
EmbeddingSwitchPlan,
|
||||
EmbeddingSwitchService,
|
||||
)
|
||||
|
||||
#: HTTP status per refusal code; anything not listed is a 409.
|
||||
_STATUS = {
|
||||
"model_profile_not_found": 404,
|
||||
"embedding_switch_not_found": 404,
|
||||
"embedding_space_not_found": 404,
|
||||
"sdk_reembed_unavailable": 503,
|
||||
"secure_storage_unavailable": 503,
|
||||
"insufficient_disk_space": 507,
|
||||
}
|
||||
|
||||
|
||||
class EmbeddingSwitchPlanRequest(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
profile_id: str = Field(min_length=1, max_length=128)
|
||||
expected_revision: int | None = Field(default=None, ge=1)
|
||||
|
||||
|
||||
class EmbeddingSwitchStartRequest(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
profile_id: str = Field(min_length=1, max_length=128)
|
||||
expected_revision: int = Field(ge=1)
|
||||
plan_token: str = Field(min_length=1, max_length=128)
|
||||
requested_by: str = Field(default="app", pattern=r"^(app|cli)$")
|
||||
|
||||
|
||||
def refusal(code: str, *, plan: EmbeddingSwitchPlan | None = None) -> JSONResponse:
|
||||
content: dict[str, Any] = {"detail": code}
|
||||
if plan is not None:
|
||||
content["plan"] = plan.model_dump(mode="json")
|
||||
return JSONResponse(status_code=_STATUS.get(code, 409), content=content)
|
||||
|
||||
|
||||
def _service(request: Request) -> EmbeddingSwitchService:
|
||||
return cast(EmbeddingSwitchService, request.app.state.embedding_switch)
|
||||
|
||||
|
||||
def embedding_switch_router(dependencies: Sequence[Depends]) -> APIRouter:
|
||||
router = APIRouter(dependencies=dependencies)
|
||||
|
||||
@router.get("/api/v1/embedding-spaces")
|
||||
async def list_embedding_spaces(request: Request) -> dict[str, object]:
|
||||
"""The registry and what the vector store holds, space by space."""
|
||||
service = _service(request)
|
||||
spaces = await asyncio.to_thread(service.spaces.list_spaces)
|
||||
live = next((space for space in spaces if space.state == "live"), None)
|
||||
scanned: list[dict[str, object]] | None = None
|
||||
scan_error = service.backend.unavailable_reason()
|
||||
if scan_error is None:
|
||||
try:
|
||||
found = await service.backend.scan()
|
||||
except Exception as error: # noqa: BLE001 - the registry still answers
|
||||
scan_error = getattr(error, "code", None) or "embedding_space_scan_failed"
|
||||
else:
|
||||
known = {(space.collection_base, space.name_dims or 0): space for space in spaces}
|
||||
scanned = [
|
||||
{
|
||||
"collection_base": item.collection_base,
|
||||
"name_dims": item.name_dims,
|
||||
"width": item.width,
|
||||
"vectors": item.vectors,
|
||||
"counts": dict(item.counts),
|
||||
"space_id": known[item.key].space_id if item.key in known else None,
|
||||
"state": known[item.key].state if item.key in known else "unregistered",
|
||||
}
|
||||
for item in found
|
||||
]
|
||||
return {
|
||||
"schema_version": 1,
|
||||
"live_space_id": live.space_id if live is not None else None,
|
||||
"spaces": [space.model_dump(mode="json") for space in spaces],
|
||||
"scanned": scanned,
|
||||
"scan_error": scan_error,
|
||||
}
|
||||
|
||||
@router.delete("/api/v1/embedding-spaces/{space_id}", response_model=None)
|
||||
async def delete_embedding_space(space_id: str, request: Request) -> Any:
|
||||
"""Clean up a retired space now instead of after its 14 days."""
|
||||
try:
|
||||
space = await _service(request).drop_space(space_id)
|
||||
except (EmbeddingSwitchError, EmbeddingSpaceError, ReembedUnavailableError) as error:
|
||||
return refusal(error.code)
|
||||
except RuntimeError as error:
|
||||
if "held open" in str(error):
|
||||
return refusal("embedding_space_in_use")
|
||||
raise
|
||||
return {"schema_version": 1, "space": space.model_dump(mode="json")}
|
||||
|
||||
@router.post("/api/v1/embedding-switches/plan", response_model=None)
|
||||
async def plan_embedding_switch(value: EmbeddingSwitchPlanRequest, request: Request) -> Any:
|
||||
try:
|
||||
plan = await _service(request).plan(
|
||||
value.profile_id, expected_revision=value.expected_revision
|
||||
)
|
||||
except EmbeddingSwitchError as error:
|
||||
return refusal(error.code, plan=error.plan)
|
||||
return plan.model_dump(mode="json")
|
||||
|
||||
@router.post("/api/v1/embedding-switches", status_code=202, response_model=None)
|
||||
async def start_embedding_switch(value: EmbeddingSwitchStartRequest, request: Request) -> Any:
|
||||
try:
|
||||
job = await _service(request).start_switch(
|
||||
value.profile_id,
|
||||
expected_revision=value.expected_revision,
|
||||
plan_token=value.plan_token,
|
||||
requested_by=value.requested_by,
|
||||
)
|
||||
except EmbeddingSwitchError as error:
|
||||
return refusal(error.code, plan=error.plan)
|
||||
return job.model_dump(mode="json")
|
||||
|
||||
@router.get("/api/v1/embedding-switches/current")
|
||||
async def current_embedding_switch(request: Request, response: Response) -> dict[str, object]:
|
||||
"""Polled by the loading screen, the banner and `memhub models switch status`."""
|
||||
response.headers["Cache-Control"] = "no-store"
|
||||
return await asyncio.to_thread(_service(request).current)
|
||||
|
||||
@router.get("/api/v1/embedding-switches/{job_id}", response_model=None)
|
||||
async def get_embedding_switch(job_id: str, request: Request) -> Any:
|
||||
job = await asyncio.to_thread(_service(request).get_job, job_id)
|
||||
if job is None:
|
||||
return refusal("embedding_switch_not_found")
|
||||
return job.model_dump(mode="json")
|
||||
|
||||
@router.post("/api/v1/embedding-switches/{job_id}/cancel", response_model=None)
|
||||
async def cancel_embedding_switch(job_id: str, request: Request) -> Any:
|
||||
try:
|
||||
job = await asyncio.to_thread(_service(request).cancel, job_id)
|
||||
except EmbeddingSwitchError as error:
|
||||
return refusal(error.code)
|
||||
return job.model_dump(mode="json")
|
||||
|
||||
return router
|
||||
@@ -71,6 +71,16 @@ def graph_store_enabled() -> bool:
|
||||
return raw != "none"
|
||||
|
||||
|
||||
def effective_memory_mode(configured: str = "pro") -> str:
|
||||
"""The SDK mode clients are built with: the setting, unless dev forces ultra.
|
||||
|
||||
`ONE_MEMORY_GRAPH_STORE` (any value but empty or `none`) predates the
|
||||
setting and keeps working for development: it forces ultra whatever the
|
||||
configured mode. Without it the configured mode (default pro) applies.
|
||||
"""
|
||||
return "ultra" if graph_store_enabled() else configured
|
||||
|
||||
|
||||
def inspect_graph_width_mismatch(config: Any) -> GraphWidthMismatch | None:
|
||||
"""Return a mismatch without renaming or deleting anything."""
|
||||
db_path = Path(config.graph_store.db_path)
|
||||
@@ -83,11 +93,16 @@ def inspect_profile_graph_width(
|
||||
*,
|
||||
api_key: str | None,
|
||||
data_dir: Path,
|
||||
mode: str | None = None,
|
||||
graph_dir: str = "graph",
|
||||
) -> GraphWidthMismatch | None:
|
||||
"""Inspect the graph path and width selected by the active profile."""
|
||||
"""Inspect the graph path and width selected by the active profile.
|
||||
|
||||
`graph_dir` is the live embedding space's graph (`embedding_spaces`).
|
||||
"""
|
||||
del api_key
|
||||
return _inspect_graph_width(
|
||||
_profile_graph_path(data_dir),
|
||||
_profile_graph_path(data_dir, mode, graph_dir),
|
||||
_profile_embedding_width(profile),
|
||||
)
|
||||
|
||||
@@ -98,11 +113,13 @@ def decide_profile_graph_width(
|
||||
api_key: str | None,
|
||||
data_dir: Path,
|
||||
decision: GraphWidthDecision,
|
||||
mode: str | None = None,
|
||||
graph_dir: str = "graph",
|
||||
) -> GraphWidthOutcome:
|
||||
"""Apply the user's decision to the active profile's graph."""
|
||||
del api_key
|
||||
config = _GraphAdminConfig(
|
||||
graph_store=_GraphStoreConfig(str(_profile_graph_path(data_dir))),
|
||||
graph_store=_GraphStoreConfig(str(_profile_graph_path(data_dir, mode, graph_dir))),
|
||||
embedder=_EmbeddingWidthConfig(_profile_embedding_width(profile)),
|
||||
)
|
||||
return apply_graph_width_decision(
|
||||
@@ -192,6 +209,12 @@ _GRAPH_WIDTH_CACHE: dict[str, int | None] = {}
|
||||
_GRAPH_WIDTH_LOCK = threading.Lock()
|
||||
|
||||
|
||||
def forget_graph_widths() -> None:
|
||||
"""Drop every remembered width: an embedding switch moved the graph."""
|
||||
with _GRAPH_WIDTH_LOCK:
|
||||
_GRAPH_WIDTH_CACHE.clear()
|
||||
|
||||
|
||||
def _graph_embedding_width(db_path: Path) -> int | None:
|
||||
"""Read the Memory embedding width directly from Kuzu's schema."""
|
||||
try:
|
||||
@@ -240,8 +263,9 @@ def _inspect_graph_width(
|
||||
)
|
||||
|
||||
|
||||
def _profile_graph_path(data_dir: Path) -> Path:
|
||||
name = "graph" if graph_store_enabled() else "graph-disabled"
|
||||
def _profile_graph_path(data_dir: Path, mode: str | None = None, graph_dir: str = "graph") -> Path:
|
||||
ultra = mode == "ultra" if mode is not None else graph_store_enabled()
|
||||
name = graph_dir if ultra else "graph-disabled"
|
||||
return data_dir / "hy-memory" / HY_MEMORY_NAMESPACE / name
|
||||
|
||||
|
||||
|
||||
@@ -14,7 +14,7 @@ from pathlib import Path
|
||||
from typing import Literal, cast
|
||||
|
||||
HY_MEMORY_DISTRIBUTION = "hy-memory"
|
||||
HY_MEMORY_VERSION = "2.0.9"
|
||||
HY_MEMORY_VERSION = "2.0.10"
|
||||
HY_MEMORY_NAMESPACE = "governed-v2"
|
||||
DiagnosticLogLevel = Literal["error", "warning", "info", "debug"]
|
||||
|
||||
|
||||
@@ -3,20 +3,26 @@
|
||||
Runtime product code must use public ``HyMemoryClient`` methods. The startup
|
||||
migration boundary is the isolated exception: all writers are stopped, so it
|
||||
may copy, restore and repair SDK-owned local stores.
|
||||
|
||||
The embedding switch is the second such boundary: model writes are paused
|
||||
while it re-embeds one space into another (`SdkReembedBackend`), so it may
|
||||
read and fill SDK collections through the SDK's re-embed module.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import sqlite3
|
||||
from collections.abc import Iterator
|
||||
from contextlib import closing
|
||||
from contextlib import closing, contextmanager
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from hy_memory import ( # type: ignore[import-untyped]
|
||||
import hy_memory # type: ignore[import-untyped]
|
||||
from hy_memory import (
|
||||
DataPlaneRestore,
|
||||
DataPlaneSnapshot,
|
||||
restore_local_data_plane,
|
||||
@@ -497,7 +503,7 @@ def _ensure_coding_schema(path: Path) -> None:
|
||||
CodingMemoryStore,
|
||||
)
|
||||
|
||||
# Relies on the pinned SDK's internals (hy-memory==2.0.9 is pinned exactly).
|
||||
# Relies on the pinned SDK's internals (hy-memory==2.0.10 is pinned exactly).
|
||||
store = CodingMemoryStore.__new__(CodingMemoryStore)
|
||||
store._db_path = str(path)
|
||||
store._conn = None
|
||||
@@ -511,3 +517,256 @@ def _ensure_coding_schema(path: Path) -> None:
|
||||
def _batches[T](values: list[T] | tuple[T, ...]) -> Iterator[tuple[T, ...]]:
|
||||
for start in range(0, len(values), _BATCH):
|
||||
yield tuple(values[start : start + _BATCH])
|
||||
|
||||
|
||||
# The SDK functions the embedding switch needs. They arrived in hy-memory
|
||||
# 2.0.10; an older SDK reports `sdk_reembed_unavailable` instead of failing.
|
||||
_REEMBED_API = (
|
||||
"VectorSpace",
|
||||
"scan_vector_spaces",
|
||||
"diff_spaces",
|
||||
"reembed_vector_space",
|
||||
"verify_reembed",
|
||||
"drop_vector_space",
|
||||
"reembed_graph",
|
||||
"verify_graph_reembed",
|
||||
"drop_graph",
|
||||
"ReembedCancelled",
|
||||
)
|
||||
|
||||
|
||||
class SdkReembedBackend:
|
||||
"""`ReembedBackend` over `hy_memory`'s re-embed module, for one home.
|
||||
|
||||
Runs beside live SDK clients: the vector functions open Chroma with the
|
||||
SDK's own settings, so they share its client system. The graph functions
|
||||
refuse a graph a live client holds open; the switch closes them first.
|
||||
"""
|
||||
|
||||
def __init__(self, data_dir: Path) -> None:
|
||||
self.data_dir = data_dir
|
||||
self.root = data_dir.resolve() / "hy-memory" / HY_MEMORY_NAMESPACE
|
||||
self.persist_directory = self.root / "vector"
|
||||
self._chroma_opened = False
|
||||
|
||||
def unavailable_reason(self) -> str | None:
|
||||
missing = [name for name in _REEMBED_API if not hasattr(hy_memory, name)]
|
||||
return "sdk_reembed_unavailable" if missing else None
|
||||
|
||||
def _sdk(self) -> Any:
|
||||
from one_memory.embedding_backend import ReembedUnavailableError
|
||||
|
||||
reason = self.unavailable_reason()
|
||||
if reason is not None:
|
||||
raise ReembedUnavailableError(reason)
|
||||
return hy_memory
|
||||
|
||||
def _open_chroma_once(self) -> None:
|
||||
"""Open Chroma once under the client-construction lock.
|
||||
|
||||
Every later open beside the Runtime finds that system. Two first
|
||||
opens race over Chroma's unlocked cache, and a client under
|
||||
construction lost its stores that way (see `hy_sdk_runtime`).
|
||||
"""
|
||||
if self._chroma_opened or not self.persist_directory.is_dir():
|
||||
return
|
||||
import chromadb
|
||||
|
||||
from one_memory.hy_sdk_runtime import first_store_open
|
||||
|
||||
with first_store_open(self.data_dir):
|
||||
chromadb.PersistentClient(path=str(self.persist_directory))
|
||||
self._chroma_opened = True
|
||||
|
||||
def _vector_space(self, space: Any) -> Any:
|
||||
return self._sdk().VectorSpace(
|
||||
base_name=space.collection_base,
|
||||
dims=space.dims,
|
||||
persist_directory=str(self.persist_directory),
|
||||
model=space.model,
|
||||
collections=space.names(),
|
||||
)
|
||||
|
||||
async def scan(self, *, per_user: bool = False) -> tuple[Any, ...]:
|
||||
from one_memory.embedding_backend import ScannedSpace
|
||||
|
||||
if not self.persist_directory.is_dir():
|
||||
return ()
|
||||
await asyncio.to_thread(self._open_chroma_once)
|
||||
found = await self._sdk().scan_vector_spaces(str(self.persist_directory), per_user=per_user)
|
||||
return tuple(
|
||||
ScannedSpace(
|
||||
collection_base=str(item.base_name),
|
||||
name_dims=item.dims,
|
||||
width=item.width,
|
||||
counts={role: int(entry.count) for role, entry in item.roles.items()},
|
||||
collections={role: str(entry.name) for role, entry in item.roles.items()},
|
||||
per_user=dict(item.per_user) if per_user else {},
|
||||
)
|
||||
for item in found
|
||||
)
|
||||
|
||||
async def diff(self, source: Any, target: Any) -> dict[str, int]:
|
||||
await asyncio.to_thread(self._open_chroma_once)
|
||||
result = await self._sdk().diff_spaces(
|
||||
self._vector_space(source), self._vector_space(target)
|
||||
)
|
||||
return {role: len(ids) for role, ids in result.missing.items()}
|
||||
|
||||
async def reembed(
|
||||
self,
|
||||
source: Any,
|
||||
target: Any,
|
||||
*,
|
||||
embedding: Any,
|
||||
cancel: Any,
|
||||
on_step: Any,
|
||||
batch_size: int = 64,
|
||||
) -> None:
|
||||
from one_memory.embedding_backend import ReembedStep
|
||||
|
||||
await asyncio.to_thread(self._open_chroma_once)
|
||||
sdk = self._sdk()
|
||||
|
||||
def progress(item: Any) -> None:
|
||||
on_step(
|
||||
ReembedStep(
|
||||
phase="vectors",
|
||||
role=str(item.role),
|
||||
done=int(item.done),
|
||||
total=int(item.total),
|
||||
embedded=int(item.embedded),
|
||||
reused=int(item.reused),
|
||||
kept=int(getattr(item, "kept", 0)),
|
||||
)
|
||||
)
|
||||
|
||||
with self._translated():
|
||||
await sdk.reembed_vector_space(
|
||||
self._vector_space(source),
|
||||
self._vector_space(target),
|
||||
embedding=self._embedding_service(embedding),
|
||||
batch_size=batch_size,
|
||||
cancel=cancel,
|
||||
on_progress=progress,
|
||||
# The target wins: an id it already holds is never touched, so
|
||||
# a resumed run continues and a merge never overwrites live data.
|
||||
skip_existing=True,
|
||||
)
|
||||
|
||||
async def verify(self, source: Any, target: Any, *, embedding: Any) -> Any:
|
||||
from one_memory.embedding_backend import VectorCheck
|
||||
|
||||
await asyncio.to_thread(self._open_chroma_once)
|
||||
with self._translated():
|
||||
result = await self._sdk().verify_reembed(
|
||||
self._vector_space(source),
|
||||
self._vector_space(target),
|
||||
embedding=self._embedding_service(embedding),
|
||||
)
|
||||
return VectorCheck(
|
||||
per_role={role: dict(values) for role, values in result.per_role.items()},
|
||||
per_user=dict(result.per_user),
|
||||
problems=tuple(result.problems),
|
||||
warnings=tuple(result.warnings),
|
||||
self_recall_hits=int(result.self_recall_hits),
|
||||
self_recall_samples=int(result.self_recall_samples),
|
||||
)
|
||||
|
||||
async def graph_vectors(self, path: Path) -> int | None:
|
||||
if not path.exists():
|
||||
return None
|
||||
# Not re-exported at the package root; the module is the SDK's own.
|
||||
plan = await self._sdk().reembed.plan_graph_reembed(path)
|
||||
return int(plan.vectors)
|
||||
|
||||
async def reembed_graph(
|
||||
self,
|
||||
source: Path,
|
||||
target: Path,
|
||||
*,
|
||||
embedding: Any,
|
||||
target_dims: int,
|
||||
cancel: Any,
|
||||
on_step: Any,
|
||||
) -> None:
|
||||
from one_memory.embedding_backend import ReembedStep
|
||||
|
||||
def progress(item: Any) -> None:
|
||||
on_step(
|
||||
ReembedStep(
|
||||
phase="graph",
|
||||
role=str(item.table),
|
||||
done=int(item.done),
|
||||
total=int(item.total),
|
||||
embedded=int(item.embedded),
|
||||
)
|
||||
)
|
||||
|
||||
with self._translated():
|
||||
await self._sdk().reembed_graph(
|
||||
source,
|
||||
target,
|
||||
embedding=self._embedding_service(embedding),
|
||||
target_dims=target_dims,
|
||||
cancel=cancel,
|
||||
on_progress=progress,
|
||||
# A resumed switch rebuilds the copy from the untouched source.
|
||||
overwrite=True,
|
||||
)
|
||||
|
||||
async def verify_graph(self, source: Path, target: Path, *, target_dims: int) -> Any:
|
||||
from one_memory.embedding_backend import GraphCheck
|
||||
|
||||
result = await self._sdk().verify_graph_reembed(source, target, target_dims=target_dims)
|
||||
return GraphCheck(ok=bool(result.ok), problems=tuple(result.problems))
|
||||
|
||||
async def drop(self, space: Any) -> None:
|
||||
await asyncio.to_thread(self._open_chroma_once)
|
||||
await self._sdk().drop_vector_space(self._vector_space(space))
|
||||
|
||||
async def drop_graph(self, path: Path) -> None:
|
||||
await self._sdk().drop_graph(path)
|
||||
|
||||
def _embedding_service(self, embedding: Any) -> Any:
|
||||
"""The service that computes a target space's vectors.
|
||||
|
||||
The built-in model is MemHub's ONNX service, exactly what the Runtime
|
||||
client embeds with. A remote model gets the SDK's own service, built
|
||||
from the same configuration the Runtime client would use.
|
||||
"""
|
||||
from one_memory.hy_sdk_runtime import (
|
||||
LocalOnnxEmbeddingService,
|
||||
build_sdk_config,
|
||||
retrieval_view,
|
||||
)
|
||||
from one_memory.model_profiles import ProviderKind
|
||||
|
||||
profile = embedding.profile
|
||||
provider = profile.embedding_provider_kind or profile.provider_kind
|
||||
if provider is ProviderKind.LOCAL_ONNX:
|
||||
return LocalOnnxEmbeddingService()
|
||||
from hy_memory.core.embed_service import EmbedService # type: ignore[import-untyped]
|
||||
|
||||
return EmbedService(
|
||||
# The retrieval view: re-embedding never needs the generation model.
|
||||
build_sdk_config(
|
||||
retrieval_view(profile),
|
||||
embedding.space,
|
||||
api_key=embedding.api_key,
|
||||
data_dir=self.data_dir,
|
||||
)
|
||||
)
|
||||
|
||||
@contextmanager
|
||||
def _translated(self) -> Iterator[None]:
|
||||
"""SDK refusals as the Core errors the switch classifies."""
|
||||
from one_memory.embedding_backend import EmbeddingWidthError, ReembedCancelledError
|
||||
|
||||
sdk = self._sdk()
|
||||
try:
|
||||
yield
|
||||
except sdk.ReembedCancelled:
|
||||
raise ReembedCancelledError("embedding_switch_cancelled") from None
|
||||
except hy_memory.EmbeddingDimensionMismatchError as error:
|
||||
raise EmbeddingWidthError("embedding_dimensions_mismatch") from error
|
||||
@@ -11,10 +11,11 @@ import asyncio
|
||||
import hashlib
|
||||
import os
|
||||
import threading
|
||||
from collections.abc import Callable, Coroutine
|
||||
from collections.abc import Callable, Coroutine, Iterable, Iterator, Mapping
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any, ClassVar
|
||||
from typing import Any, ClassVar, Literal, cast
|
||||
|
||||
from hy_memory import ( # type: ignore[import-untyped]
|
||||
AmbiguousMemoryIdError as AmbiguousMemoryIdError, # re-export for api.py
|
||||
@@ -29,14 +30,27 @@ from one_memory.embedding_assets import (
|
||||
BUILTIN_EMBEDDING_DIMENSIONS,
|
||||
BUILTIN_EMBEDDING_MODEL,
|
||||
)
|
||||
from one_memory.embedding_spaces import (
|
||||
DEFAULT_GRAPH_DIR,
|
||||
EmbeddingSpace,
|
||||
EmbeddingSpaceError,
|
||||
derived_space,
|
||||
embedding_identity,
|
||||
)
|
||||
from one_memory.hy_graph_admin import (
|
||||
GraphWidthMismatchError,
|
||||
graph_store_enabled,
|
||||
effective_memory_mode,
|
||||
inspect_graph_width_mismatch,
|
||||
)
|
||||
from one_memory.hy_sdk_types import HySdkRuntimeError
|
||||
from one_memory.local_embedding import LocalEmbedder
|
||||
from one_memory.log_file import confine_hy_sdk_logging
|
||||
from one_memory.memory_features import (
|
||||
FEATURES,
|
||||
FEATURES_BY_NAME,
|
||||
FeatureOverrides,
|
||||
feature_overrides,
|
||||
)
|
||||
from one_memory.model_profiles import (
|
||||
ModelProfileSnapshot,
|
||||
ProviderKind,
|
||||
@@ -48,6 +62,10 @@ BUILTIN_RETRIEVAL_PROFILE_ID = "builtin-retrieval"
|
||||
_LOCAL_ONNX_PLACEHOLDER_URL = "http://127.0.0.1/local-onnx"
|
||||
DEFAULT_RELEVANCE_FLOOR = 0.35
|
||||
|
||||
#: The SDK's memory modes. A mode never changes a collection name: the space
|
||||
#: names the collections, the mode only which stores the client opens.
|
||||
MemoryMode = Literal["lite", "pro", "ultra"]
|
||||
|
||||
|
||||
@dataclass
|
||||
class _RuntimeEntry:
|
||||
@@ -170,29 +188,73 @@ class SdkClientHandle:
|
||||
api_key: str | None,
|
||||
data_dir: Path,
|
||||
http_transport_factory: Callable[[], Any] | None = None,
|
||||
space: EmbeddingSpace | None = None,
|
||||
mode: str | None = None,
|
||||
features: FeatureOverrides = (),
|
||||
) -> None:
|
||||
self._profile = profile
|
||||
self._api_key = api_key
|
||||
self._data_dir = data_dir
|
||||
self._http_transport_factory = http_transport_factory
|
||||
self._space = space
|
||||
#: The SDK mode this handle's client is built with, fixed for its life.
|
||||
self.mode = mode or effective_memory_mode()
|
||||
#: The SDK features set on its config (`memory_features`), fixed likewise.
|
||||
self.features = features
|
||||
self._lease: ClientLease | None = None
|
||||
self._constructing: asyncio.Future[ClientLease] | None = None
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
@property
|
||||
def profile(self) -> ModelProfileSnapshot:
|
||||
return self._profile
|
||||
|
||||
@property
|
||||
def space(self) -> EmbeddingSpace | None:
|
||||
return self._space
|
||||
|
||||
@property
|
||||
def api_key(self) -> str | None:
|
||||
return self._api_key
|
||||
|
||||
async def client(self) -> HyMemoryClient:
|
||||
if self._lease is not None:
|
||||
return self._lease.client
|
||||
async with self._lock:
|
||||
if self._lease is None:
|
||||
self._lease = await asyncio.to_thread(
|
||||
acquire_client,
|
||||
self._profile,
|
||||
api_key=self._api_key,
|
||||
data_dir=self._data_dir,
|
||||
http_transport_factory=self._http_transport_factory,
|
||||
constructing = asyncio.ensure_future(
|
||||
asyncio.to_thread(
|
||||
_acquire_for_handle,
|
||||
self._profile,
|
||||
self._space,
|
||||
api_key=self._api_key,
|
||||
data_dir=self._data_dir,
|
||||
http_transport_factory=self._http_transport_factory,
|
||||
mode=self.mode,
|
||||
features=self.features,
|
||||
)
|
||||
)
|
||||
self._constructing = constructing
|
||||
try:
|
||||
self._lease = await asyncio.shield(constructing)
|
||||
except asyncio.CancelledError:
|
||||
# The construction thread runs on. Keep what it builds, so
|
||||
# `close` releases it instead of leaving its runtime open.
|
||||
constructing.add_done_callback(self._keep)
|
||||
raise
|
||||
return self._lease.client
|
||||
|
||||
def _keep(self, constructing: asyncio.Future[ClientLease]) -> None:
|
||||
if constructing.cancelled() or constructing.exception() is not None:
|
||||
return
|
||||
if self._lease is None:
|
||||
self._lease = constructing.result()
|
||||
|
||||
async def close(self) -> None:
|
||||
constructing = self._constructing
|
||||
if constructing is not None and not constructing.done():
|
||||
# A construction a cancelled caller left behind: wait for its lease.
|
||||
await asyncio.gather(constructing, return_exceptions=True)
|
||||
async with self._lock:
|
||||
lease, self._lease = self._lease, None
|
||||
if lease is not None:
|
||||
@@ -217,6 +279,33 @@ def builtin_retrieval_snapshot() -> ModelProfileSnapshot:
|
||||
)
|
||||
|
||||
|
||||
def retrieval_view(profile: ModelProfileSnapshot) -> ModelProfileSnapshot:
|
||||
"""The read side of `profile`: its embedding space, and no generation.
|
||||
|
||||
Reads must open the collection writes fill (#242), and the collection is
|
||||
named by the embedding only. Generation stays the built-in placeholder, so
|
||||
a read never needs the generation model. A remote embedding still needs
|
||||
the profile's API key, which `build_sdk_config` uses as the embedding key.
|
||||
"""
|
||||
provider = profile.embedding_provider_kind or profile.provider_kind
|
||||
endpoint: str | None = None
|
||||
if provider is not ProviderKind.LOCAL_ONNX:
|
||||
endpoint = profile.embedding_base_url
|
||||
if not endpoint:
|
||||
raise HySdkRuntimeError("hy_sdk_model_endpoint_missing")
|
||||
return builtin_retrieval_snapshot().model_copy(
|
||||
update={
|
||||
"profile_id": profile.profile_id,
|
||||
"revision": profile.revision,
|
||||
"embedding_provider_kind": provider,
|
||||
"embedding_model": profile.embedding_model,
|
||||
"embedding_base_url": endpoint,
|
||||
"embedding_dimensions": profile.embedding_dimensions,
|
||||
"has_api_key": profile.has_api_key,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def sdk_relevance_floor() -> float:
|
||||
raw = os.environ.get("ONE_MEMORY_HY_RELEVANCE_FLOOR", "").strip()
|
||||
if not raw:
|
||||
@@ -228,31 +317,71 @@ def sdk_relevance_floor() -> float:
|
||||
return min(1.0, max(0.0, value))
|
||||
|
||||
|
||||
def sdk_feature_defaults() -> dict[str, bool]:
|
||||
"""Each `memory_features` feature's SDK default, as a fresh config has it.
|
||||
|
||||
That is the SDK's code default unless the process environment sets the
|
||||
SDK's own variable (`FeatureSpec.env`), exactly what an unset feature
|
||||
leaves on a client.
|
||||
"""
|
||||
config = MemoryConfig()
|
||||
return {
|
||||
spec.name: bool(getattr(getattr(config, spec.section), spec.field)) for spec in FEATURES
|
||||
}
|
||||
|
||||
|
||||
def apply_sdk_features(
|
||||
config: MemoryConfig, features: Mapping[str, bool] | Iterable[tuple[str, bool]]
|
||||
) -> MemoryConfig:
|
||||
"""Set each feature a person chose on its `MemoryConfig` field.
|
||||
|
||||
Every exposed feature is a config field the SDK reads from the config at
|
||||
call time, never from `os.environ`, so this is per client and
|
||||
deterministic. Nothing here touches `vector_store`, so the collection
|
||||
identity never depends on a feature.
|
||||
"""
|
||||
values = dict(features.items() if isinstance(features, Mapping) else features)
|
||||
for name, value in feature_overrides(values):
|
||||
spec = FEATURES_BY_NAME[name]
|
||||
setattr(getattr(config, spec.section), spec.field, value)
|
||||
return config
|
||||
|
||||
|
||||
def build_sdk_config(
|
||||
profile: ModelProfileSnapshot,
|
||||
space: EmbeddingSpace | None = None,
|
||||
mode: MemoryMode | None = None,
|
||||
*,
|
||||
api_key: str | None,
|
||||
data_dir: Path,
|
||||
features: Mapping[str, bool] | Iterable[tuple[str, bool]] = (),
|
||||
) -> MemoryConfig:
|
||||
"""Map a model snapshot to SDK models, endpoints, backends, and local paths."""
|
||||
"""Map a model snapshot and its embedding space to an SDK configuration.
|
||||
|
||||
`space` names the collections and the graph; the profile only says how to
|
||||
reach the models. A profile whose embedding identity (model and width) is
|
||||
not the space's raises `embedding_space_mismatch` instead of opening new,
|
||||
empty collections. Without `space`, the space the profile opened before
|
||||
the registry existed is used: offline tools on a home with no registry.
|
||||
`mode` never changes a name, and neither do `features`
|
||||
(`memory_features`), which are set on top; with none the config is
|
||||
exactly what it was before the setting existed.
|
||||
"""
|
||||
if mode is not None and mode not in {"lite", "pro", "ultra"}:
|
||||
raise HySdkRuntimeError("hy_sdk_memory_mode_invalid")
|
||||
generation_endpoint, generation_key = _generation_access(profile, api_key)
|
||||
embedding_provider = profile.embedding_provider_kind or profile.provider_kind
|
||||
embedding_service_is_local = embedding_provider is ProviderKind.LOCAL_ONNX
|
||||
embedding_model = profile.embedding_model or (
|
||||
BUILTIN_EMBEDDING_MODEL if embedding_service_is_local else None
|
||||
)
|
||||
if not embedding_model:
|
||||
raise HySdkRuntimeError("hy_sdk_embedding_model_required")
|
||||
dimensions: int | None
|
||||
try:
|
||||
identity = embedding_identity(profile)
|
||||
except EmbeddingSpaceError as error:
|
||||
raise HySdkRuntimeError(error.code) from None
|
||||
if embedding_service_is_local:
|
||||
embedding_endpoint = None
|
||||
embedding_key = None
|
||||
dimensions = profile.embedding_dimensions or BUILTIN_EMBEDDING_DIMENSIONS
|
||||
else:
|
||||
embedding_endpoint = _openai_endpoint(
|
||||
embedding_provider,
|
||||
profile.embedding_base_url or profile.base_url,
|
||||
)
|
||||
# Never the generation endpoint: a remote embedding has its own.
|
||||
embedding_endpoint = _openai_endpoint(embedding_provider, profile.embedding_base_url)
|
||||
embedding_key = "ollama" if embedding_provider is ProviderKind.OLLAMA else api_key
|
||||
if embedding_provider is ProviderKind.OPENAI_COMPATIBLE and not embedding_key:
|
||||
raise HySdkRuntimeError("hy_sdk_embedding_api_key_missing")
|
||||
@@ -261,13 +390,18 @@ def build_sdk_config(
|
||||
ProviderKind.OLLAMA,
|
||||
}:
|
||||
raise HySdkRuntimeError("hy_sdk_embedding_provider_unsupported")
|
||||
dimensions = profile.embedding_dimensions
|
||||
if space is None:
|
||||
space = derived_space(profile)
|
||||
elif space.identity is None or not identity.matches(space.identity):
|
||||
raise HySdkRuntimeError("embedding_space_mismatch")
|
||||
# The width the space was opened with: its names carry it, and a remote
|
||||
# request sends it as `dimensions`. An "auto" space sends none, as before.
|
||||
dimensions = space.name_dims
|
||||
embedding_model = identity.model
|
||||
|
||||
root = data_dir.resolve() / "hy-memory" / "governed-v2"
|
||||
root.mkdir(parents=True, exist_ok=True)
|
||||
collection = _collection_base(
|
||||
embedding_provider, embedding_endpoint, embedding_model, dimensions
|
||||
)
|
||||
collection = space.collection_base
|
||||
|
||||
llm: dict[str, Any] = {
|
||||
"provider": "openai",
|
||||
@@ -282,7 +416,7 @@ def build_sdk_config(
|
||||
if thinking_off:
|
||||
llm["extra_body"] = thinking_off
|
||||
|
||||
return MemoryConfig.from_dict(
|
||||
config = MemoryConfig.from_dict(
|
||||
{
|
||||
"llm": llm,
|
||||
"embedder": {
|
||||
@@ -305,7 +439,7 @@ def build_sdk_config(
|
||||
},
|
||||
"graph_store": {
|
||||
"provider": "kuzu",
|
||||
"db_path": str(root / "graph"),
|
||||
"db_path": str(root / (space.graph_dir or DEFAULT_GRAPH_DIR)),
|
||||
"url": None,
|
||||
"username": None,
|
||||
"password": None,
|
||||
@@ -319,18 +453,50 @@ def build_sdk_config(
|
||||
"coding": {"db_path": str(root / "coding.sqlite3")},
|
||||
}
|
||||
)
|
||||
return apply_sdk_features(config, features)
|
||||
|
||||
|
||||
def _acquire_for_handle(
|
||||
profile: ModelProfileSnapshot,
|
||||
space: EmbeddingSpace | None,
|
||||
*,
|
||||
api_key: str | None,
|
||||
data_dir: Path,
|
||||
http_transport_factory: Callable[[], Any] | None,
|
||||
mode: str | None,
|
||||
features: FeatureOverrides = (),
|
||||
) -> ClientLease:
|
||||
# `acquire_client` is looked up at call time, so a stand-in replaces every
|
||||
# construction. A handle without a space keeps the call it always made.
|
||||
named: dict[str, Any] = {"space": space} if space is not None else {}
|
||||
return acquire_client(
|
||||
profile,
|
||||
api_key=api_key,
|
||||
data_dir=data_dir,
|
||||
http_transport_factory=http_transport_factory,
|
||||
mode=mode,
|
||||
features=features,
|
||||
**named,
|
||||
)
|
||||
|
||||
|
||||
def acquire_client(
|
||||
profile: ModelProfileSnapshot,
|
||||
space: EmbeddingSpace | None = None,
|
||||
*,
|
||||
api_key: str | None,
|
||||
data_dir: Path,
|
||||
http_transport_factory: Callable[[], Any] | None = None,
|
||||
mode: str | None = None,
|
||||
features: FeatureOverrides = (),
|
||||
) -> ClientLease:
|
||||
"""Construct a public SDK client and return its explicit runtime lease.
|
||||
|
||||
Waits while `hold_stores_offline` holds this home's stores.
|
||||
`space` names the collections (`build_sdk_config`). `mode` is the SDK
|
||||
pipeline mode (lite / pro / ultra) and never changes a name; None keeps
|
||||
the pre-setting default (ultra when `ONE_MEMORY_GRAPH_STORE` is set, else
|
||||
pro). `features` are the `memory_features` set on its config; none keeps
|
||||
every SDK default. Waits while `hold_stores_offline` holds this home's stores.
|
||||
"""
|
||||
key = data_dir.resolve()
|
||||
with _STORE_GATE:
|
||||
@@ -338,11 +504,15 @@ def acquire_client(
|
||||
_STORE_GATE.wait()
|
||||
_OPENING[key] = _OPENING.get(key, 0) + 1
|
||||
try:
|
||||
named: dict[str, Any] = {"space": space} if space is not None else {}
|
||||
return _acquire_client(
|
||||
profile,
|
||||
api_key=api_key,
|
||||
data_dir=data_dir,
|
||||
http_transport_factory=http_transport_factory,
|
||||
mode=mode or effective_memory_mode(),
|
||||
features=features,
|
||||
**named,
|
||||
)
|
||||
finally:
|
||||
with _STORE_GATE:
|
||||
@@ -352,6 +522,20 @@ def acquire_client(
|
||||
_STORE_GATE.notify_all()
|
||||
|
||||
|
||||
@contextmanager
|
||||
def first_store_open(data_dir: Path) -> Iterator[None]:
|
||||
"""One opening of this home's stores at a time: Chroma's first open is unlocked.
|
||||
|
||||
The same lock client construction takes, for anything else that opens
|
||||
Chroma beside the Runtime (the re-embed functions).
|
||||
"""
|
||||
key = data_dir.resolve()
|
||||
with _STORE_GATE:
|
||||
constructing = _CONSTRUCTING.setdefault(key, threading.Lock())
|
||||
with constructing:
|
||||
yield
|
||||
|
||||
|
||||
def hold_stores_offline(data_dir: Path) -> bool:
|
||||
"""Hold this home's SDK stores for offline maintenance; False if a client is open.
|
||||
|
||||
@@ -377,12 +561,22 @@ def release_stores(data_dir: Path) -> None:
|
||||
|
||||
def _acquire_client(
|
||||
profile: ModelProfileSnapshot,
|
||||
space: EmbeddingSpace | None = None,
|
||||
*,
|
||||
api_key: str | None,
|
||||
data_dir: Path,
|
||||
http_transport_factory: Callable[[], Any] | None,
|
||||
mode: str = "pro",
|
||||
features: FeatureOverrides = (),
|
||||
) -> ClientLease:
|
||||
config = build_sdk_config(profile, api_key=api_key, data_dir=data_dir)
|
||||
config = build_sdk_config(
|
||||
profile,
|
||||
space,
|
||||
cast(MemoryMode, mode),
|
||||
api_key=api_key,
|
||||
data_dir=data_dir,
|
||||
features=features,
|
||||
)
|
||||
config.http_transport_factory = http_transport_factory
|
||||
mismatch = inspect_graph_width_mismatch(config)
|
||||
if mismatch is not None:
|
||||
@@ -398,7 +592,7 @@ def _acquire_client(
|
||||
with constructing:
|
||||
client = HyMemoryClient(
|
||||
config=config,
|
||||
mode="ultra" if graph_store_enabled() else "pro",
|
||||
mode=mode,
|
||||
runtime=runtime,
|
||||
embedding_service=embedding_service,
|
||||
)
|
||||
@@ -438,6 +632,33 @@ def _generation_access(
|
||||
return endpoint, api_key
|
||||
|
||||
|
||||
def embedding_endpoint(profile: ModelProfileSnapshot) -> str | None:
|
||||
"""The OpenAI-compatible base URL the SDK embeds through; None when local."""
|
||||
provider = profile.embedding_provider_kind or profile.provider_kind
|
||||
if provider is ProviderKind.LOCAL_ONNX:
|
||||
return None
|
||||
return _openai_endpoint(provider, profile.embedding_base_url)
|
||||
|
||||
|
||||
def legacy_collection_base(profile: ModelProfileSnapshot) -> str:
|
||||
"""The four-attribute name a profile's space had before the registry.
|
||||
|
||||
Existing homes are registered under exactly this name, so it must never
|
||||
change: provider, embedding endpoint, model and declared width.
|
||||
"""
|
||||
provider = profile.embedding_provider_kind or profile.provider_kind
|
||||
local = provider is ProviderKind.LOCAL_ONNX
|
||||
model = profile.embedding_model or (BUILTIN_EMBEDDING_MODEL if local else None)
|
||||
if not model:
|
||||
raise HySdkRuntimeError("hy_sdk_embedding_model_required")
|
||||
if local:
|
||||
return _collection_base(
|
||||
provider, None, model, profile.embedding_dimensions or BUILTIN_EMBEDDING_DIMENSIONS
|
||||
)
|
||||
endpoint = _openai_endpoint(provider, profile.embedding_base_url)
|
||||
return _collection_base(provider, endpoint, model, profile.embedding_dimensions)
|
||||
|
||||
|
||||
def _collection_base(
|
||||
provider: ProviderKind,
|
||||
endpoint: str | None,
|
||||
@@ -467,7 +688,7 @@ def vector_collection_bases(profile: ModelProfileSnapshot) -> tuple[str, tuple[s
|
||||
endpoint, dimensions = None, profile.embedding_dimensions or BUILTIN_EMBEDDING_DIMENSIONS
|
||||
legacy_endpoint = "local-onnx"
|
||||
else:
|
||||
endpoint = _openai_endpoint(provider, profile.embedding_base_url or profile.base_url)
|
||||
endpoint = _openai_endpoint(provider, profile.embedding_base_url)
|
||||
dimensions = profile.embedding_dimensions
|
||||
legacy_endpoint = endpoint
|
||||
legacy_model = profile.embedding_model or "bge-small-zh-v1.5"
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from collections.abc import Mapping
|
||||
from datetime import UTC, datetime
|
||||
from typing import Literal, NamedTuple, cast
|
||||
@@ -391,6 +392,8 @@ class HySdkRawRecord(BaseModel):
|
||||
memory_at: int | None = None
|
||||
gmt_created: int | None = None
|
||||
custom: dict[str, object] = Field(default_factory=dict)
|
||||
#: active, or shadow once pro extracted it; None when the SDK gave none.
|
||||
status: str | None = None
|
||||
|
||||
@classmethod
|
||||
def from_sdk(cls, raw: object) -> HySdkRawRecord:
|
||||
@@ -417,11 +420,34 @@ class HySdkRawRecord(BaseModel):
|
||||
raw.get("gmt_created") if isinstance(raw.get("gmt_created"), int) else None
|
||||
),
|
||||
custom=dict(custom),
|
||||
status=str(raw["status"]) if isinstance(raw.get("status"), str) else None,
|
||||
)
|
||||
except (KeyError, TypeError, ValueError) as error:
|
||||
raise HySdkReadError("hy_memory_raw_shape_unknown") from error
|
||||
|
||||
|
||||
_RAW_TURN_LINE = re.compile(r"^\[(user|assistant|tool|system)\]:\s?(.*)$")
|
||||
|
||||
|
||||
def raw_turn_messages(content: str) -> tuple[tuple[str, str], ...]:
|
||||
"""The `(role, text)` messages of one SDK raw turn.
|
||||
|
||||
The SDK stores a turn as `[user]: …\\n[assistant]: …`; a line without a
|
||||
role prefix continues the message above it. Content that names no role
|
||||
at all is one message of unknown role (`""`).
|
||||
"""
|
||||
messages: list[list[str]] = []
|
||||
for line in content.splitlines():
|
||||
match = _RAW_TURN_LINE.match(line)
|
||||
if match is not None:
|
||||
messages.append([match.group(1), match.group(2)])
|
||||
elif messages:
|
||||
messages[-1][1] += "\n" + line
|
||||
else:
|
||||
messages.append(["", line])
|
||||
return tuple((role, text) for role, text in messages)
|
||||
|
||||
|
||||
class HySdkRawHit(HySdkRawRecord):
|
||||
score: float = Field(ge=0, le=1)
|
||||
|
||||
@@ -516,6 +542,96 @@ class HySdkUserProfile(BaseModel):
|
||||
)
|
||||
|
||||
|
||||
class HySdkProfileVersion(BaseModel):
|
||||
"""One version on an L0 basic-info chain."""
|
||||
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
memory_id: str
|
||||
content: str
|
||||
#: Conversation time as the SDK renders it on a chain (epoch or text).
|
||||
memory_at: int | str | None = None
|
||||
gmt_created: int | None = None
|
||||
#: The chain head: the version the profile shows now.
|
||||
current: bool = False
|
||||
|
||||
|
||||
class HySdkProfileChain(BaseModel):
|
||||
"""One L0 basic-info head and its versions, latest first."""
|
||||
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
memory_id: str
|
||||
content: str
|
||||
status: str
|
||||
memory_at: int | None = None
|
||||
gmt_created: int | None = None
|
||||
#: The structured fields this head holds (`basic_info_kv`).
|
||||
basic_info: dict[str, object] = Field(default_factory=dict)
|
||||
versions: tuple[HySdkProfileVersion, ...]
|
||||
|
||||
|
||||
class HySdkProfileVersions(BaseModel):
|
||||
"""`get_user_basic_info`: the long-term profile's L0 version chains."""
|
||||
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
request_id: str
|
||||
chains: tuple[HySdkProfileChain, ...]
|
||||
elapsed_ms: float = Field(ge=0)
|
||||
|
||||
@classmethod
|
||||
def from_sdk(cls, raw: object) -> HySdkProfileVersions:
|
||||
if not isinstance(raw, Mapping):
|
||||
raise HySdkReadError("hy_memory_profile_versions_shape_unknown")
|
||||
request_id = raw.get("request_id")
|
||||
rows = raw.get("memories")
|
||||
elapsed_ms = raw.get("elapsed_ms")
|
||||
if (
|
||||
not isinstance(request_id, str)
|
||||
or not isinstance(rows, list)
|
||||
or not isinstance(elapsed_ms, int | float)
|
||||
or isinstance(elapsed_ms, bool)
|
||||
):
|
||||
raise HySdkReadError("hy_memory_profile_versions_shape_unknown")
|
||||
chains: list[HySdkProfileChain] = []
|
||||
for row in rows:
|
||||
if not isinstance(row, Mapping) or not isinstance(row.get("memory_id"), str):
|
||||
raise HySdkReadError("hy_memory_profile_versions_shape_unknown")
|
||||
head_id = str(row["memory_id"])
|
||||
chain = row.get("evolution_chain")
|
||||
entries = chain if isinstance(chain, list) and chain else [{**row, "node_id": head_id}]
|
||||
versions = tuple(
|
||||
HySdkProfileVersion(
|
||||
memory_id=str(entry.get("node_id") or entry.get("memory_id")),
|
||||
content=str(entry.get("content") or ""),
|
||||
memory_at=(
|
||||
entry.get("memory_at")
|
||||
if isinstance(entry.get("memory_at"), int | str)
|
||||
and not isinstance(entry.get("memory_at"), bool)
|
||||
else None
|
||||
),
|
||||
gmt_created=_epoch(entry.get("gmt_created")),
|
||||
current=str(entry.get("node_id") or entry.get("memory_id")) == head_id,
|
||||
)
|
||||
for entry in entries
|
||||
if isinstance(entry, Mapping) and (entry.get("node_id") or entry.get("memory_id"))
|
||||
)
|
||||
basic_info = row.get("basic_info_kv")
|
||||
chains.append(
|
||||
HySdkProfileChain(
|
||||
memory_id=head_id,
|
||||
content=str(row.get("content") or ""),
|
||||
status=str(row.get("status") or "active"),
|
||||
memory_at=_epoch(row.get("memory_at")),
|
||||
gmt_created=_epoch(row.get("gmt_created")),
|
||||
basic_info=dict(basic_info) if isinstance(basic_info, Mapping) else {},
|
||||
versions=versions,
|
||||
)
|
||||
)
|
||||
return cls(request_id=request_id, chains=tuple(chains), elapsed_ms=float(elapsed_ms))
|
||||
|
||||
|
||||
def parse_sdk_search(raw: object, *, error_code: str) -> tuple[SemanticMemoryHit, ...]:
|
||||
rows = _reply_rows(raw, error_code)
|
||||
found: list[SemanticMemoryHit] = []
|
||||
@@ -542,6 +658,120 @@ def parse_sdk_search(raw: object, *, error_code: str) -> tuple[SemanticMemoryHit
|
||||
return tuple(found)
|
||||
|
||||
|
||||
def parse_sdk_lane_search(raw: object, *, error_code: str) -> tuple[SemanticMemoryHit, ...]:
|
||||
"""Every scene's search reply as hits that name their lane and namespace.
|
||||
|
||||
`parse_sdk_search` reads the chat reply only and forgets which lane a hit
|
||||
came from. A caller may now pick the scene and a floor per lane, so this
|
||||
reads all three reply shapes:
|
||||
|
||||
- chat: `memories` as `{schema, proactive, normal}` (or one list, normal);
|
||||
- working/coding: `facts` and `experiences` of that variant;
|
||||
- all: one fused `memories` list whose rows carry `scene`, `kind` and
|
||||
their own lane's score (`lane_score`). The fused `score` is a rank
|
||||
fusion (1/(60+rank)), not a relevance, so the lane score is reported.
|
||||
|
||||
Highest score first, first occurrence of an id kept.
|
||||
"""
|
||||
if not isinstance(raw, Mapping):
|
||||
raise HySdkReadError(error_code)
|
||||
scene = raw.get("scene")
|
||||
found: list[SemanticMemoryHit] = []
|
||||
if scene == "all":
|
||||
rows = raw.get("memories")
|
||||
if not isinstance(rows, list) or any(not isinstance(row, Mapping) for row in rows):
|
||||
raise HySdkReadError(error_code)
|
||||
for row in rows:
|
||||
hit = _lane_hit(row, scene=str(row.get("scene") or "chat"), score_key="lane_score")
|
||||
if hit is not None:
|
||||
found.append(hit)
|
||||
elif scene in ("productivity", "working", "coding"):
|
||||
variant = raw.get("variant") if raw.get("variant") in {"working", "coding"} else None
|
||||
lane = variant or (scene if scene in {"working", "coding"} else "working")
|
||||
for key, kind in (("facts", "fact"), ("experiences", "experience")):
|
||||
rows = raw.get(key) or []
|
||||
if not isinstance(rows, list) or any(not isinstance(row, Mapping) for row in rows):
|
||||
raise HySdkReadError(error_code)
|
||||
for row in rows:
|
||||
hit = _lane_hit({**row, "kind": kind}, scene=lane, score_key="score")
|
||||
if hit is not None:
|
||||
found.append(hit)
|
||||
else:
|
||||
memories = raw.get("memories")
|
||||
if isinstance(memories, Mapping):
|
||||
if set(memories) != {"schema", "proactive", "normal"}:
|
||||
raise HySdkReadError(error_code)
|
||||
channels = tuple(memories.items())
|
||||
elif isinstance(memories, list):
|
||||
channels = (("normal", memories),)
|
||||
else:
|
||||
raise HySdkReadError(error_code)
|
||||
for lane, rows in channels:
|
||||
if not isinstance(rows, list) or any(not isinstance(row, Mapping) for row in rows):
|
||||
raise HySdkReadError(error_code)
|
||||
for row in rows:
|
||||
hit = _lane_hit(row, scene="chat", score_key="score", lane=str(lane))
|
||||
if hit is not None:
|
||||
found.append(hit)
|
||||
ordered: list[SemanticMemoryHit] = []
|
||||
seen: set[tuple[str, str, str | None]] = set()
|
||||
for hit in sorted(found, key=lambda value: value.score, reverse=True):
|
||||
identity = (hit.memory_id, hit.kind, hit.variant)
|
||||
if identity in seen:
|
||||
continue
|
||||
seen.add(identity)
|
||||
ordered.append(hit)
|
||||
return tuple(ordered)
|
||||
|
||||
|
||||
def _lane_hit(
|
||||
row: Mapping[str, object],
|
||||
*,
|
||||
scene: str,
|
||||
score_key: str,
|
||||
lane: str | None = None,
|
||||
) -> SemanticMemoryHit | None:
|
||||
kind = str(row.get("kind")) if row.get("kind") in {"experience", "fact"} else "chat"
|
||||
if kind == "chat":
|
||||
memory_id = row.get("memory_id") or row.get("id")
|
||||
variant = None
|
||||
hit_lane = lane or "normal"
|
||||
else:
|
||||
id_key = "experience_id" if kind == "experience" else "fact_id"
|
||||
memory_id = row.get("memory_id") or row.get(id_key)
|
||||
variant = scene if scene in {"working", "coding"} else None
|
||||
if variant is None:
|
||||
return None
|
||||
hit_lane = variant
|
||||
if not isinstance(memory_id, str) or not memory_id:
|
||||
return None
|
||||
value = row.get(score_key)
|
||||
if not isinstance(value, int | float) or isinstance(value, bool):
|
||||
value = row.get("score")
|
||||
score = (
|
||||
max(0.0, min(1.0, float(value)))
|
||||
if isinstance(value, int | float) and not isinstance(value, bool)
|
||||
else 0.0
|
||||
)
|
||||
if kind != "chat":
|
||||
return SemanticMemoryHit(
|
||||
memory_id=memory_id, score=score, lane=hit_lane, kind=kind, variant=variant
|
||||
)
|
||||
content = row.get("content")
|
||||
member_ids, formatting_complete = _formatted_member_ids(row)
|
||||
return SemanticMemoryHit(
|
||||
memory_id=memory_id,
|
||||
score=score,
|
||||
formatted_content=(
|
||||
content.strip() if isinstance(content, str) and content.strip() else None
|
||||
),
|
||||
member_ids=member_ids,
|
||||
sdk_formatted=True,
|
||||
formatting_complete=formatting_complete,
|
||||
lane=hit_lane,
|
||||
)
|
||||
|
||||
|
||||
class LibraryHit(NamedTuple):
|
||||
"""One keyword match in a library search window, with its namespace."""
|
||||
|
||||
|
||||
@@ -28,7 +28,10 @@ from one_memory.hy_sdk_types import LibraryHit
|
||||
DEFAULT_WINDOW_TTL_SECONDS = 300.0
|
||||
DEFAULT_MAX_WINDOWS = 32
|
||||
|
||||
WindowKey = tuple[str, str | None, tuple[str, ...], str, str]
|
||||
#: (query, source, statuses, scene, user, other filters as sorted name/values pairs)
|
||||
WindowKey = tuple[
|
||||
str, str | None, tuple[str, ...], str, str, tuple[tuple[str, tuple[str, ...]], ...]
|
||||
]
|
||||
|
||||
|
||||
def window_key(
|
||||
@@ -37,8 +40,12 @@ def window_key(
|
||||
statuses: list[str],
|
||||
scene: str = "all",
|
||||
user_id: str = "local",
|
||||
filters: dict[str, list[str]] | None = None,
|
||||
) -> WindowKey:
|
||||
return (query, source, tuple(sorted(set(statuses))), scene, user_id)
|
||||
extra = tuple(
|
||||
sorted((name, tuple(sorted(set(values)))) for name, values in (filters or {}).items())
|
||||
)
|
||||
return (query, source, tuple(sorted(set(statuses))), scene, user_id, extra)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
|
||||
@@ -0,0 +1,618 @@
|
||||
"""System 2 for ultra: evolve runs, their receipts, the nightly schedule, graph backfill.
|
||||
|
||||
hy-memory 2.0.9's `client.evolve(user_id, agent_id, max_batches)` is ultra
|
||||
only, holds no lock, and has no scheduler: MemHub owns all three.
|
||||
|
||||
- A run is asynchronous: `start` answers a receipt at once and the run goes
|
||||
on in the background (evolve makes model calls across the user's whole
|
||||
library and can take minutes, longer than any HTTP timeout). One run per
|
||||
user at a time: asking again while one runs returns that run. A run holds
|
||||
the write coordinator's per-user exclusivity, so it never overlaps that
|
||||
user's write batch, and loops while the SDK answers `should_continue`, up
|
||||
to `EVOLVE_MAX_ROUNDS`.
|
||||
- The schedule runs every user that has memories once a day at the
|
||||
configured local time, and catches up after the machine slept through it.
|
||||
- Switching to ultra backfills each user's existing facts into the graph
|
||||
once (`build_memory_graph`).
|
||||
|
||||
What is kept on disk is counts and stable codes only, never memory text:
|
||||
`<data_dir>/evolve-state.json`.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import threading
|
||||
from collections.abc import Awaitable, Callable, Iterable, Mapping
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from pathlib import Path
|
||||
from tempfile import NamedTemporaryFile
|
||||
from typing import Any, Literal, Protocol
|
||||
from uuid import uuid4
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
from one_memory.diagnostic_log import diagnostic_event
|
||||
from one_memory.memory_server.config_state import EvolveSchedule
|
||||
from one_memory.memory_server.durable_queue import safe_error_code
|
||||
from one_memory.sdk_write_processor import PRODUCT_AGENT_ID
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
#: Rounds one run makes while the SDK says there is more to do.
|
||||
EVOLVE_MAX_ROUNDS = 20
|
||||
#: How often the schedule looks at the clock; also bounds catch-up after wake.
|
||||
SCHEDULE_CHECK_SECONDS = 300.0
|
||||
#: Receipts kept in memory for `GET /api/v2/evolve/{run_id}`.
|
||||
_KEPT_RUNS = 64
|
||||
_STABLE_CODE = re.compile(r"[a-z][a-z0-9_]{0,95}")
|
||||
_COUNTS = (
|
||||
"remaining_clusters",
|
||||
"remaining_actionable_clusters",
|
||||
"remaining_fresh_facts",
|
||||
)
|
||||
|
||||
RunState = Literal["running", "succeeded", "failed"]
|
||||
Trigger = Literal["manual", "scheduled"]
|
||||
|
||||
|
||||
class EvolveModeError(RuntimeError):
|
||||
"""Evolve asked for outside ultra."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__("memory_mode_not_ultra")
|
||||
|
||||
|
||||
class EvolveRunReceipt(BaseModel):
|
||||
"""One evolve run. Counts and codes; `results` only while in memory."""
|
||||
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
run_id: str
|
||||
user_id: str
|
||||
agent_id: str
|
||||
max_batches: int | None = None
|
||||
trigger: Trigger = "manual"
|
||||
state: RunState = "running"
|
||||
started_at: datetime
|
||||
finished_at: datetime | None = None
|
||||
rounds: int = 0
|
||||
tasks_processed: int = 0
|
||||
batch_failures: int = 0
|
||||
cycle_status: str | None = None
|
||||
should_continue: bool | None = None
|
||||
continuation_reason: str | None = None
|
||||
remaining_clusters: int | None = None
|
||||
remaining_actionable_clusters: int | None = None
|
||||
remaining_fresh_facts: int | None = None
|
||||
#: True when the round cap stopped a run the SDK would have continued.
|
||||
capped: bool = False
|
||||
elapsed_ms: float = 0.0
|
||||
error_code: str | None = None
|
||||
failure_kind: str | None = None
|
||||
#: The last round's SDK `request_id`, for matching SDK logs.
|
||||
request_id: str | None = None
|
||||
#: The last round's SDK `results` and `cross_domain_sweeper`. Kept only
|
||||
#: in this receipt, never written to disk or diagnostics.
|
||||
results: Any = None
|
||||
cross_domain_sweeper: Any = None
|
||||
|
||||
def summary(self) -> dict[str, object]:
|
||||
"""What is persisted and shown in status: no SDK result bodies."""
|
||||
return self.model_dump(mode="json", exclude={"results", "cross_domain_sweeper"})
|
||||
|
||||
|
||||
class UserEvolveState(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
last_run_at: datetime | None = None
|
||||
last_result: dict[str, object] | None = None
|
||||
last_error: str | None = None
|
||||
|
||||
|
||||
class GraphBuildUser(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
state: Literal["pending", "running", "succeeded", "failed"] = "pending"
|
||||
memories_scanned: int | None = None
|
||||
nodes: int | None = None
|
||||
error_code: str | None = None
|
||||
|
||||
|
||||
class GraphBuildState(BaseModel):
|
||||
"""The one-time backfill of existing facts into the graph after → ultra."""
|
||||
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
state: Literal["idle", "running", "succeeded", "failed"] = "idle"
|
||||
started_at: datetime | None = None
|
||||
finished_at: datetime | None = None
|
||||
users: dict[str, GraphBuildUser] = Field(default_factory=dict)
|
||||
|
||||
|
||||
class EvolveStateDocument(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
schema_version: int = Field(default=1, ge=1, le=1)
|
||||
users: dict[str, UserEvolveState] = Field(default_factory=dict)
|
||||
#: When the schedule first ran in ultra: the "last run" of a user that
|
||||
#: never ran, so enabling ultra at noon does not evolve everyone at once.
|
||||
schedule_anchor: datetime | None = None
|
||||
graph_build: GraphBuildState = Field(default_factory=GraphBuildState)
|
||||
|
||||
|
||||
class EvolveStateStore:
|
||||
"""The only writer of `<data_dir>/evolve-state.json`."""
|
||||
|
||||
def __init__(self, data_dir: Path) -> None:
|
||||
self.path = data_dir / "evolve-state.json"
|
||||
self._lock = threading.RLock()
|
||||
self._document = self._read()
|
||||
|
||||
def document(self) -> EvolveStateDocument:
|
||||
with self._lock:
|
||||
return self._document.model_copy(deep=True)
|
||||
|
||||
def user(self, user_id: str) -> UserEvolveState:
|
||||
with self._lock:
|
||||
return self._document.users.get(user_id, UserEvolveState()).model_copy()
|
||||
|
||||
def record_run(self, receipt: EvolveRunReceipt) -> None:
|
||||
with self._lock:
|
||||
users = dict(self._document.users)
|
||||
users[receipt.user_id] = UserEvolveState(
|
||||
last_run_at=receipt.started_at,
|
||||
last_result=receipt.summary(),
|
||||
last_error=receipt.error_code,
|
||||
)
|
||||
self._save(self._document.model_copy(update={"users": users}))
|
||||
|
||||
def anchor(self, now: datetime) -> datetime:
|
||||
"""The schedule anchor, set to `now` the first time it is asked."""
|
||||
with self._lock:
|
||||
if self._document.schedule_anchor is None:
|
||||
self._save(self._document.model_copy(update={"schedule_anchor": now}))
|
||||
assert self._document.schedule_anchor is not None
|
||||
return self._document.schedule_anchor
|
||||
|
||||
def graph_build(self) -> GraphBuildState:
|
||||
with self._lock:
|
||||
return self._document.graph_build.model_copy(deep=True)
|
||||
|
||||
def set_graph_build(self, state: GraphBuildState) -> None:
|
||||
with self._lock:
|
||||
self._save(self._document.model_copy(update={"graph_build": state}))
|
||||
|
||||
def _save(self, document: EvolveStateDocument) -> None:
|
||||
payload = json.dumps(document.model_dump(mode="json"), indent=2, sort_keys=True) + "\n"
|
||||
self.path.parent.mkdir(parents=True, exist_ok=True)
|
||||
temporary: Path | None = None
|
||||
try:
|
||||
with NamedTemporaryFile(
|
||||
mode="w",
|
||||
encoding="utf-8",
|
||||
dir=self.path.parent,
|
||||
prefix=".evolve-state.",
|
||||
suffix=".tmp",
|
||||
delete=False,
|
||||
) as stream:
|
||||
temporary = Path(stream.name)
|
||||
stream.write(payload)
|
||||
stream.flush()
|
||||
os.fsync(stream.fileno())
|
||||
os.chmod(temporary, 0o600)
|
||||
os.replace(temporary, self.path)
|
||||
except OSError:
|
||||
if temporary is not None:
|
||||
temporary.unlink(missing_ok=True)
|
||||
# Status is best-effort: a run must not fail because its record did.
|
||||
logger.warning("evolve_state_write_failed")
|
||||
self._document = document
|
||||
|
||||
def _read(self) -> EvolveStateDocument:
|
||||
try:
|
||||
return EvolveStateDocument.model_validate_json(self.path.read_text(encoding="utf-8"))
|
||||
except FileNotFoundError:
|
||||
return EvolveStateDocument()
|
||||
except (OSError, ValueError):
|
||||
# A damaged status file only loses history; start a fresh one.
|
||||
logger.warning("evolve_state_unreadable")
|
||||
return EvolveStateDocument()
|
||||
|
||||
|
||||
class EvolveCoordinator(Protocol):
|
||||
"""The write coordinator surface evolve and the graph backfill use."""
|
||||
|
||||
def exclusive_for_evolve(self, user_id: str) -> Any: ...
|
||||
|
||||
def use_client(self) -> Any: ...
|
||||
|
||||
|
||||
class EvolveService:
|
||||
"""Start, serialize and record evolve runs."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
coordinator: EvolveCoordinator,
|
||||
state: EvolveStateStore,
|
||||
memory_mode: Callable[[], str],
|
||||
max_rounds: int = EVOLVE_MAX_ROUNDS,
|
||||
now: Callable[[], datetime] = lambda: datetime.now(UTC),
|
||||
) -> None:
|
||||
self.coordinator = coordinator
|
||||
self.state = state
|
||||
self.memory_mode = memory_mode
|
||||
self.max_rounds = max_rounds
|
||||
self._now = now
|
||||
self._runs: dict[str, EvolveRunReceipt] = {}
|
||||
self._tasks: dict[str, asyncio.Task[None]] = {}
|
||||
self._active: dict[str, str] = {}
|
||||
|
||||
def start(
|
||||
self,
|
||||
user_id: str,
|
||||
*,
|
||||
agent_id: str | None = None,
|
||||
max_batches: int | None = None,
|
||||
trigger: Trigger = "manual",
|
||||
) -> EvolveRunReceipt:
|
||||
"""Start a run, or return the one already running for this user."""
|
||||
if self.memory_mode() != "ultra":
|
||||
raise EvolveModeError()
|
||||
running = self._active.get(user_id)
|
||||
if running is not None:
|
||||
return self._runs[running]
|
||||
receipt = EvolveRunReceipt(
|
||||
run_id=str(uuid4()),
|
||||
user_id=user_id,
|
||||
agent_id=agent_id or PRODUCT_AGENT_ID,
|
||||
max_batches=max_batches,
|
||||
trigger=trigger,
|
||||
started_at=self._now(),
|
||||
)
|
||||
self._remember(receipt)
|
||||
self._active[user_id] = receipt.run_id
|
||||
self._tasks[receipt.run_id] = asyncio.create_task(
|
||||
self._run(receipt.run_id), name="one-memory-evolve"
|
||||
)
|
||||
return receipt
|
||||
|
||||
def get(self, run_id: str) -> EvolveRunReceipt | None:
|
||||
return self._runs.get(run_id)
|
||||
|
||||
def running(self) -> tuple[EvolveRunReceipt, ...]:
|
||||
return tuple(self._runs[run_id] for run_id in self._active.values())
|
||||
|
||||
async def wait(self, run_id: str) -> EvolveRunReceipt | None:
|
||||
task = self._tasks.get(run_id)
|
||||
if task is not None:
|
||||
await asyncio.shield(task)
|
||||
return self._runs.get(run_id)
|
||||
|
||||
async def shutdown(self) -> None:
|
||||
tasks = tuple(self._tasks.values())
|
||||
for task in tasks:
|
||||
task.cancel()
|
||||
await asyncio.gather(*tasks, return_exceptions=True)
|
||||
self._tasks.clear()
|
||||
|
||||
def _remember(self, receipt: EvolveRunReceipt) -> None:
|
||||
self._runs[receipt.run_id] = receipt
|
||||
while len(self._runs) > _KEPT_RUNS:
|
||||
oldest = next(iter(self._runs))
|
||||
if oldest in self._active.values():
|
||||
break
|
||||
del self._runs[oldest]
|
||||
|
||||
async def _run(self, run_id: str) -> None:
|
||||
receipt = self._runs[run_id]
|
||||
diagnostic_event(
|
||||
level="info",
|
||||
component="memory.evolve",
|
||||
event="memory_evolve.started",
|
||||
outcome="started",
|
||||
details={"trigger": receipt.trigger, "max_batches": receipt.max_batches},
|
||||
)
|
||||
try:
|
||||
async with (
|
||||
self.coordinator.exclusive_for_evolve(receipt.user_id),
|
||||
self.coordinator.use_client() as client,
|
||||
):
|
||||
while receipt.rounds < self.max_rounds:
|
||||
result = await asyncio.to_thread(
|
||||
client.evolve,
|
||||
receipt.user_id,
|
||||
agent_id=receipt.agent_id,
|
||||
max_batches=receipt.max_batches,
|
||||
)
|
||||
receipt = _folded(receipt, result)
|
||||
self._runs[run_id] = receipt
|
||||
if receipt.state == "failed" or not receipt.should_continue:
|
||||
break
|
||||
else:
|
||||
receipt = receipt.model_copy(update={"capped": True})
|
||||
if receipt.state == "running":
|
||||
receipt = receipt.model_copy(update={"state": "succeeded"})
|
||||
except asyncio.CancelledError:
|
||||
receipt = receipt.model_copy(
|
||||
update={"state": "failed", "error_code": "memory_evolve_cancelled"}
|
||||
)
|
||||
raise
|
||||
except Exception as error: # noqa: BLE001 - recorded as a stable code
|
||||
code = safe_error_code(error)
|
||||
receipt = receipt.model_copy(
|
||||
update={
|
||||
"state": "failed",
|
||||
"error_code": _stable(code, "hy_memory_evolve_failed"),
|
||||
"failure_kind": type(error).__name__,
|
||||
}
|
||||
)
|
||||
finally:
|
||||
receipt = receipt.model_copy(update={"finished_at": self._now()})
|
||||
self._runs[run_id] = receipt
|
||||
self._active.pop(receipt.user_id, None)
|
||||
self._tasks.pop(run_id, None)
|
||||
await asyncio.to_thread(self.state.record_run, receipt)
|
||||
diagnostic_event(
|
||||
level="info" if receipt.state == "succeeded" else "warning",
|
||||
component="memory.evolve",
|
||||
event="memory_evolve.finished",
|
||||
outcome="succeeded" if receipt.state == "succeeded" else "failed",
|
||||
details={
|
||||
"trigger": receipt.trigger,
|
||||
"rounds": receipt.rounds,
|
||||
"tasks_processed": receipt.tasks_processed,
|
||||
"batch_failures": receipt.batch_failures,
|
||||
"capped": receipt.capped,
|
||||
"elapsed_ms": round(receipt.elapsed_ms, 1),
|
||||
"error_code": receipt.error_code,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _folded(receipt: EvolveRunReceipt, result: object) -> EvolveRunReceipt:
|
||||
"""One more evolve round's result added to the run."""
|
||||
if not isinstance(result, Mapping):
|
||||
return receipt.model_copy(
|
||||
update={"state": "failed", "error_code": "hy_memory_evolve_shape_unknown"}
|
||||
)
|
||||
failures = result.get("batch_failures")
|
||||
update: dict[str, object] = {
|
||||
"rounds": receipt.rounds + 1,
|
||||
"tasks_processed": receipt.tasks_processed + _int(result.get("tasks_processed")),
|
||||
"batch_failures": receipt.batch_failures
|
||||
+ (len(failures) if isinstance(failures, list | tuple) else _int(failures)),
|
||||
"elapsed_ms": receipt.elapsed_ms + _float(result.get("elapsed_ms")),
|
||||
"should_continue": result.get("should_continue") is True,
|
||||
"cycle_status": _text(result.get("cycle_status")),
|
||||
"continuation_reason": _text(result.get("continuation_reason")),
|
||||
"request_id": _text(result.get("request_id")),
|
||||
"results": result.get("results"),
|
||||
"cross_domain_sweeper": result.get("cross_domain_sweeper"),
|
||||
}
|
||||
for name in _COUNTS:
|
||||
value = result.get(name)
|
||||
update[name] = value if isinstance(value, int) and not isinstance(value, bool) else None
|
||||
if result.get("success") is not True:
|
||||
code = result.get("error_code")
|
||||
kind = result.get("failure_kind")
|
||||
update.update(
|
||||
state="failed",
|
||||
should_continue=False,
|
||||
error_code=(
|
||||
code
|
||||
if isinstance(code, str) and _STABLE_CODE.fullmatch(code)
|
||||
else "hy_memory_evolve_rejected"
|
||||
),
|
||||
failure_kind=kind if isinstance(kind, str) and _STABLE_CODE.fullmatch(kind) else None,
|
||||
)
|
||||
return receipt.model_copy(update=update)
|
||||
|
||||
|
||||
def _stable(code: str, fallback: str) -> str:
|
||||
return code if _STABLE_CODE.fullmatch(code) else fallback
|
||||
|
||||
|
||||
def _int(value: object) -> int:
|
||||
return value if isinstance(value, int) and not isinstance(value, bool) else 0
|
||||
|
||||
|
||||
def _float(value: object) -> float:
|
||||
if isinstance(value, int | float) and not isinstance(value, bool):
|
||||
return float(value)
|
||||
return 0.0
|
||||
|
||||
|
||||
def _text(value: object) -> str | None:
|
||||
return value[:128] if isinstance(value, str) else None
|
||||
|
||||
|
||||
def evolve_due(
|
||||
now: datetime,
|
||||
schedule: EvolveSchedule,
|
||||
last_run_at: datetime | None,
|
||||
anchor: datetime,
|
||||
) -> bool:
|
||||
"""Whether a user's nightly evolve is due at `now` (aware, local time).
|
||||
|
||||
Due once the day's slot (`schedule.time`, local) has passed and the last
|
||||
run (or, for a user who never ran, the anchor) is older than that slot.
|
||||
A last run more than 24 hours old is due at once as well: a machine asleep
|
||||
at the slot catches up when it wakes, even before the next day's slot.
|
||||
"""
|
||||
if not schedule.enabled:
|
||||
return False
|
||||
local = now.astimezone()
|
||||
hour, minute = (int(part) for part in schedule.time.split(":"))
|
||||
slot = local.replace(hour=hour, minute=minute, second=0, microsecond=0)
|
||||
reference = last_run_at or anchor
|
||||
if local >= slot and reference < slot:
|
||||
return True
|
||||
return last_run_at is not None and local - last_run_at > timedelta(hours=24)
|
||||
|
||||
|
||||
class EvolveScheduler:
|
||||
"""Run evolve nightly for every user that has memories (ultra only)."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
service: EvolveService,
|
||||
schedule: Callable[[], EvolveSchedule],
|
||||
users_with_memories: Callable[[], Awaitable[Iterable[str]]],
|
||||
now: Callable[[], datetime] = lambda: datetime.now().astimezone(),
|
||||
check_seconds: float = SCHEDULE_CHECK_SECONDS,
|
||||
) -> None:
|
||||
self.service = service
|
||||
self.schedule = schedule
|
||||
self.users_with_memories = users_with_memories
|
||||
self._now = now
|
||||
self.check_seconds = check_seconds
|
||||
self._task: asyncio.Task[None] | None = None
|
||||
|
||||
def start(self) -> None:
|
||||
if self._task is None or self._task.done():
|
||||
self._task = asyncio.create_task(self._loop(), name="one-memory-evolve-schedule")
|
||||
|
||||
async def tick(self) -> tuple[str, ...]:
|
||||
"""Run every due user once, one after another; the users run."""
|
||||
schedule = self.schedule()
|
||||
if not schedule.enabled or self.service.memory_mode() != "ultra":
|
||||
return ()
|
||||
now = self._now()
|
||||
anchor = await asyncio.to_thread(self.service.state.anchor, now)
|
||||
ran: list[str] = []
|
||||
for user_id in await self.users_with_memories():
|
||||
last = self.service.state.user(user_id).last_run_at
|
||||
if not evolve_due(now, schedule, last, anchor):
|
||||
continue
|
||||
try:
|
||||
receipt = self.service.start(user_id, trigger="scheduled")
|
||||
except EvolveModeError:
|
||||
break
|
||||
await self.service.wait(receipt.run_id)
|
||||
ran.append(user_id)
|
||||
if ran:
|
||||
diagnostic_event(
|
||||
level="info",
|
||||
component="memory.evolve",
|
||||
event="memory_evolve.schedule_ran",
|
||||
outcome="succeeded",
|
||||
details={"users": len(ran)},
|
||||
)
|
||||
return tuple(ran)
|
||||
|
||||
async def _loop(self) -> None:
|
||||
while True:
|
||||
try:
|
||||
await self.tick()
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception as error: # noqa: BLE001 - the loop must survive
|
||||
logger.warning("evolve_schedule_tick_failed error_type=%s", type(error).__name__)
|
||||
await asyncio.sleep(self.check_seconds)
|
||||
|
||||
async def shutdown(self) -> None:
|
||||
if self._task is None:
|
||||
return
|
||||
self._task.cancel()
|
||||
await asyncio.gather(self._task, return_exceptions=True)
|
||||
self._task = None
|
||||
|
||||
|
||||
class MemoryGraphBackfill:
|
||||
"""After a switch to ultra, project each user's existing facts into the graph once."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
coordinator: EvolveCoordinator,
|
||||
state: EvolveStateStore,
|
||||
now: Callable[[], datetime] = lambda: datetime.now(UTC),
|
||||
) -> None:
|
||||
self.coordinator = coordinator
|
||||
self.state = state
|
||||
self._now = now
|
||||
self._task: asyncio.Task[None] | None = None
|
||||
|
||||
@property
|
||||
def running(self) -> bool:
|
||||
return self._task is not None and not self._task.done()
|
||||
|
||||
def start(self, user_ids: Iterable[str]) -> GraphBuildState:
|
||||
if self.running:
|
||||
return self.state.graph_build()
|
||||
users = tuple(dict.fromkeys(user_ids))
|
||||
state = GraphBuildState(
|
||||
state="running",
|
||||
started_at=self._now(),
|
||||
users={user_id: GraphBuildUser() for user_id in users},
|
||||
)
|
||||
self.state.set_graph_build(state)
|
||||
self._task = asyncio.create_task(self._run(users), name="one-memory-graph-backfill")
|
||||
return state
|
||||
|
||||
async def wait(self) -> None:
|
||||
if self._task is not None:
|
||||
await asyncio.shield(self._task)
|
||||
|
||||
async def _run(self, users: tuple[str, ...]) -> None:
|
||||
state = self.state.graph_build()
|
||||
for user_id in users:
|
||||
state.users[user_id] = GraphBuildUser(state="running")
|
||||
await asyncio.to_thread(self.state.set_graph_build, state.model_copy(deep=True))
|
||||
try:
|
||||
async with (
|
||||
self.coordinator.exclusive_for_evolve(user_id),
|
||||
self.coordinator.use_client() as client,
|
||||
):
|
||||
result = await asyncio.to_thread(
|
||||
client.build_memory_graph,
|
||||
user_id,
|
||||
agent_id=None,
|
||||
limit=10000,
|
||||
clear_first=False,
|
||||
)
|
||||
ok = isinstance(result, Mapping) and result.get("success") is True
|
||||
state.users[user_id] = GraphBuildUser(
|
||||
state="succeeded" if ok else "failed",
|
||||
memories_scanned=_int(result.get("memories_scanned"))
|
||||
if isinstance(result, Mapping)
|
||||
else None,
|
||||
nodes=_int(result.get("nodes")) if isinstance(result, Mapping) else None,
|
||||
error_code=None if ok else "hy_memory_graph_build_rejected",
|
||||
)
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception as error: # noqa: BLE001 - recorded as a stable code
|
||||
code = safe_error_code(error)
|
||||
state.users[user_id] = GraphBuildUser(
|
||||
state="failed",
|
||||
error_code=_stable(code, "hy_memory_graph_build_failed"),
|
||||
)
|
||||
failed = any(user.state == "failed" for user in state.users.values())
|
||||
state.state = "failed" if failed else "succeeded"
|
||||
state.finished_at = self._now()
|
||||
await asyncio.to_thread(self.state.set_graph_build, state)
|
||||
diagnostic_event(
|
||||
level="warning" if failed else "info",
|
||||
component="memory.evolve",
|
||||
event="memory_graph_backfill.finished",
|
||||
outcome="failed" if failed else "succeeded",
|
||||
details={
|
||||
"users": len(state.users),
|
||||
"failed_users": sum(user.state == "failed" for user in state.users.values()),
|
||||
},
|
||||
)
|
||||
|
||||
async def shutdown(self) -> None:
|
||||
if self._task is None:
|
||||
return
|
||||
self._task.cancel()
|
||||
await asyncio.gather(self._task, return_exceptions=True)
|
||||
self._task = None
|
||||
@@ -0,0 +1,262 @@
|
||||
"""SDK retrieval and extraction features a person may turn on or off.
|
||||
|
||||
MemHub hard-codes none of these. Each is one `MemoryConfig` field that the SDK
|
||||
reads from its config object (never from `os.environ` at call time), so
|
||||
`hy_sdk_runtime.build_sdk_config` sets it per client and a client's features
|
||||
are fixed for its life, like its mode. An unset feature keeps the SDK default,
|
||||
so with nothing set the SDK config is byte for byte what it was before this
|
||||
setting existed.
|
||||
|
||||
The setting is one global document in `config.json` (`memory_features`),
|
||||
beside `memory_mode`; management stays global, not per user. A change
|
||||
rebuilds the SDK handles the way a mode change does (#218 use count).
|
||||
|
||||
What is deliberately left out, and why, is in
|
||||
docs/contracts/memory-features-v1.md. In short: flags the SDK reads straight
|
||||
from the process environment (the System 2 schema dedup/compaction lane, the
|
||||
balanced-catalog and specific-supplement lanes, the hybrid reader's
|
||||
`MEMORY_HYBRID_V2_*` knobs) cannot be scoped to one client; flags MemHub never
|
||||
exercises (query rewrite needs `queries=[...]`, session narratives need
|
||||
`finalize_session`); flags that are dangerous or internal (graph write, local
|
||||
cross-encoder download, access write-back); and everything MemHub already
|
||||
controls elsewhere (pipeline log, trace, memory operations, telemetry, mode).
|
||||
|
||||
No SDK import here: only `hy_sdk_runtime` and `hy_sdk_migration` import it.
|
||||
`config_state` imports this module, so it imports nothing from there.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from typing import Literal
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, StrictBool
|
||||
|
||||
MemoryFeatureName = Literal[
|
||||
"graph_rerank",
|
||||
"same_series_timeline",
|
||||
"episode_recall",
|
||||
"session_episodes",
|
||||
"episode_summary",
|
||||
"intention",
|
||||
"few_shot",
|
||||
]
|
||||
#: Read features change search; write features change what extraction stores.
|
||||
FeatureSide = Literal["read", "write"]
|
||||
#: none: no extra call. extra_query: one more store query per search.
|
||||
#: llm_call: extra generation-model calls. prompt_tokens: no extra call, but
|
||||
#: a longer extraction prompt.
|
||||
FeatureCost = Literal["none", "extra_query", "llm_call", "prompt_tokens"]
|
||||
#: Why an enabled feature does nothing right now (`active: false`).
|
||||
FeatureReason = Literal["disabled", "requires_pro", "requires_ultra", "requires_graph_rerank"]
|
||||
#: A client's features as a hashable key: sorted (name, value) pairs, set ones only.
|
||||
type FeatureOverrides = tuple[tuple[str, bool], ...]
|
||||
|
||||
_MODE_RANK: dict[str, int] = {"lite": 0, "pro": 1, "ultra": 2}
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class FeatureSpec:
|
||||
name: MemoryFeatureName
|
||||
#: `MemoryConfig` section and field the value is written to.
|
||||
section: Literal["graph_store", "recall", "llm"]
|
||||
field: str
|
||||
#: The SDK's own env var for the same field; only its default comes from it.
|
||||
env: str
|
||||
side: FeatureSide
|
||||
#: Lowest mode the feature acts in.
|
||||
requires_mode: Literal["pro", "ultra"]
|
||||
#: Another feature that must be active for this one to act.
|
||||
requires_feature: MemoryFeatureName | None
|
||||
cost: FeatureCost
|
||||
description: str
|
||||
|
||||
@property
|
||||
def sdk_field(self) -> str:
|
||||
return f"{self.section}.{self.field}"
|
||||
|
||||
|
||||
FEATURES: tuple[FeatureSpec, ...] = (
|
||||
FeatureSpec(
|
||||
name="graph_rerank",
|
||||
section="graph_store",
|
||||
field="graph_rerank_enabled",
|
||||
env="MEMORY_GRAPH_RERANK_ENABLED",
|
||||
side="read",
|
||||
requires_mode="ultra",
|
||||
requires_feature=None,
|
||||
cost="extra_query",
|
||||
description=(
|
||||
"Usage A: propagate support over knowledge-graph edges inside the "
|
||||
"recall pool, so memories linked to strong hits rank higher. Needs "
|
||||
"the graph store (ultra). No model call; one graph query per search."
|
||||
),
|
||||
),
|
||||
FeatureSpec(
|
||||
name="same_series_timeline",
|
||||
section="recall",
|
||||
field="same_series_timeline_enabled",
|
||||
env="MEMORY_SAME_SERIES_TIMELINE_ENABLED",
|
||||
side="read",
|
||||
requires_mode="ultra",
|
||||
requires_feature="graph_rerank",
|
||||
cost="none",
|
||||
description=(
|
||||
"Keep memories joined by a same-series graph edge together, in time "
|
||||
"order, in the results. Uses the edges graph rerank reads, so it "
|
||||
"needs graph_rerank (ultra). No model call."
|
||||
),
|
||||
),
|
||||
FeatureSpec(
|
||||
name="episode_recall",
|
||||
section="recall",
|
||||
field="episode_recall_enabled",
|
||||
env="MEMORY_EPISODE_RECALL_ENABLED",
|
||||
side="read",
|
||||
requires_mode="pro",
|
||||
requires_feature=None,
|
||||
cost="extra_query",
|
||||
description=(
|
||||
"Usage B: pull in memories from the same episode as a hit and score "
|
||||
"them on tag-keyword overlap, content overlap and time affinity. No "
|
||||
"graph needed (pro or ultra). No model call; one store query per "
|
||||
"search. Helps only memories that carry an episode_id."
|
||||
),
|
||||
),
|
||||
FeatureSpec(
|
||||
name="session_episodes",
|
||||
section="recall",
|
||||
field="session_episode_enabled",
|
||||
env="MEMORY_SESSION_EPISODE_ENABLED",
|
||||
side="write",
|
||||
requires_mode="pro",
|
||||
requires_feature=None,
|
||||
cost="none",
|
||||
description=(
|
||||
"Episode write side: memories extracted from one conversation share "
|
||||
"one stable episode_id, instead of the default rule (a new episode "
|
||||
"after 30 minutes idle or 50 memories). No model call."
|
||||
),
|
||||
),
|
||||
FeatureSpec(
|
||||
name="episode_summary",
|
||||
section="recall",
|
||||
field="episode_summary_enabled",
|
||||
env="MEMORY_EPISODE_SUMMARY_ENABLED",
|
||||
side="write",
|
||||
requires_mode="pro",
|
||||
requires_feature=None,
|
||||
cost="llm_call",
|
||||
description=(
|
||||
"When an episode closes (the idle-gap rule), summarize it once into "
|
||||
"an episode memory. One model call per closed episode of 3 or more "
|
||||
"memories."
|
||||
),
|
||||
),
|
||||
FeatureSpec(
|
||||
name="intention",
|
||||
section="llm",
|
||||
field="enable_intention",
|
||||
env="MEMORY_ENABLE_INTENTION",
|
||||
side="write",
|
||||
requires_mode="pro",
|
||||
requires_feature=None,
|
||||
cost="none",
|
||||
description=(
|
||||
"Extract the user's stated plans and intentions (L7) along with "
|
||||
"facts. On by default; off drops them from the extraction prompt. "
|
||||
"No extra call."
|
||||
),
|
||||
),
|
||||
FeatureSpec(
|
||||
name="few_shot",
|
||||
section="llm",
|
||||
field="few_shot_enabled",
|
||||
env="MEMORY_FEW_SHOT_ENABLED",
|
||||
side="write",
|
||||
requires_mode="pro",
|
||||
requires_feature=None,
|
||||
cost="prompt_tokens",
|
||||
description=(
|
||||
"Append worked examples to the extraction and reconcile prompts. "
|
||||
"Can help a weaker model; no extra call, but every extraction "
|
||||
"sends more input tokens."
|
||||
),
|
||||
),
|
||||
)
|
||||
FEATURE_NAMES: tuple[MemoryFeatureName, ...] = tuple(spec.name for spec in FEATURES)
|
||||
FEATURES_BY_NAME: dict[str, FeatureSpec] = {spec.name: spec for spec in FEATURES}
|
||||
|
||||
|
||||
class MemoryFeatures(BaseModel):
|
||||
"""`config.json` `memory_features`, and the `PUT` body.
|
||||
|
||||
Every key is optional; absent (or null) means the SDK default. In a `PUT`
|
||||
an absent key is left as it is and null resets it. Unknown keys and
|
||||
non-boolean values are refused (422).
|
||||
"""
|
||||
|
||||
model_config = ConfigDict(frozen=True, extra="forbid", strict=True)
|
||||
|
||||
graph_rerank: StrictBool | None = None
|
||||
same_series_timeline: StrictBool | None = None
|
||||
episode_recall: StrictBool | None = None
|
||||
session_episodes: StrictBool | None = None
|
||||
episode_summary: StrictBool | None = None
|
||||
intention: StrictBool | None = None
|
||||
few_shot: StrictBool | None = None
|
||||
|
||||
def overrides(self) -> FeatureOverrides:
|
||||
"""The features a person set, as the client key."""
|
||||
return feature_overrides(
|
||||
{name: value for name, value in self.model_dump().items() if value is not None}
|
||||
)
|
||||
|
||||
def merged(self, update: MemoryFeatures) -> MemoryFeatures:
|
||||
"""This document with the keys `update` names: a value sets, null resets."""
|
||||
values = self.model_dump()
|
||||
for name in update.model_fields_set:
|
||||
values[name] = getattr(update, name)
|
||||
return MemoryFeatures(**{key: value for key, value in values.items() if value is not None})
|
||||
|
||||
|
||||
def feature_overrides(values: Mapping[str, bool]) -> FeatureOverrides:
|
||||
unknown = set(values) - set(FEATURES_BY_NAME)
|
||||
if unknown:
|
||||
raise ValueError("unknown_memory_feature")
|
||||
return tuple(sorted((name, bool(value)) for name, value in values.items()))
|
||||
|
||||
|
||||
def effective_features(
|
||||
overrides: FeatureOverrides, defaults: Mapping[str, bool]
|
||||
) -> dict[str, bool]:
|
||||
"""Every feature's value: what a person set, else the SDK default."""
|
||||
values: dict[str, bool] = {spec.name: bool(defaults.get(spec.name, False)) for spec in FEATURES}
|
||||
values.update(dict(overrides))
|
||||
return values
|
||||
|
||||
|
||||
def feature_activity(name: str, effective: Mapping[str, bool], mode: str) -> FeatureReason | None:
|
||||
"""None when the feature acts in this mode; else why it does not."""
|
||||
spec = FEATURES_BY_NAME[name]
|
||||
if not effective.get(name, False):
|
||||
return "disabled"
|
||||
if _MODE_RANK.get(mode, 0) < _MODE_RANK[spec.requires_mode]:
|
||||
return "requires_ultra" if spec.requires_mode == "ultra" else "requires_pro"
|
||||
if spec.requires_feature is not None and (
|
||||
feature_activity(spec.requires_feature, effective, mode) is not None
|
||||
):
|
||||
return "requires_graph_rerank"
|
||||
return None
|
||||
|
||||
|
||||
def non_default_features(
|
||||
overrides: FeatureOverrides, defaults: Mapping[str, bool]
|
||||
) -> dict[str, bool]:
|
||||
"""The features whose value differs from the SDK default, for status."""
|
||||
return {
|
||||
name: value
|
||||
for name, value in effective_features(overrides, defaults).items()
|
||||
if value != bool(defaults.get(name, False))
|
||||
}
|
||||
@@ -13,9 +13,11 @@ import threading
|
||||
from collections.abc import Mapping
|
||||
from pathlib import Path
|
||||
from tempfile import NamedTemporaryFile
|
||||
from typing import Literal
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, StrictBool, field_validator
|
||||
|
||||
from one_memory.memory_features import MemoryFeatures
|
||||
from one_memory.memory_queue import ExtractionWindowPreset
|
||||
|
||||
_SOURCE_ID = re.compile(r"^[a-z0-9][a-z0-9-]{0,127}$")
|
||||
@@ -27,6 +29,31 @@ class SourceConfig(BaseModel):
|
||||
sync_enabled: StrictBool = False
|
||||
|
||||
|
||||
#: The SDK pipeline mode (hy-memory `HyMemoryClient(mode=...)`):
|
||||
#: lite keeps raw conversations only, with no model call; pro extracts
|
||||
#: memories (System 1); ultra adds the knowledge graph and evolve (System 2).
|
||||
MemoryMode = Literal["lite", "pro", "ultra"]
|
||||
MEMORY_MODES: tuple[MemoryMode, ...] = ("lite", "pro", "ultra")
|
||||
_SCHEDULE_TIME = re.compile(r"^(?:[01][0-9]|2[0-3]):[0-5][0-9]$")
|
||||
|
||||
|
||||
class EvolveSchedule(BaseModel):
|
||||
"""When ultra runs evolve for every user: once a day, local time."""
|
||||
|
||||
model_config = ConfigDict(frozen=True, extra="forbid", strict=True)
|
||||
|
||||
enabled: StrictBool = True
|
||||
#: `HH:MM`, 24-hour, the machine's local time.
|
||||
time: str = "03:00"
|
||||
|
||||
@field_validator("time")
|
||||
@classmethod
|
||||
def valid_time(cls, value: str) -> str:
|
||||
if _SCHEDULE_TIME.fullmatch(value) is None:
|
||||
raise ValueError("invalid_evolve_schedule_time")
|
||||
return value
|
||||
|
||||
|
||||
class ConfigDocument(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid", strict=True)
|
||||
|
||||
@@ -39,6 +66,13 @@ class ConfigDocument(BaseModel):
|
||||
#: needs no re-entry of a key.
|
||||
sedimentation_enabled: StrictBool = True
|
||||
extraction_window_preset: ExtractionWindowPreset = ExtractionWindowPreset.THREE_MONTHS
|
||||
#: The SDK mode every client is built with. `ONE_MEMORY_GRAPH_STORE` (dev)
|
||||
#: still forces ultra over it; see `hy_graph_admin.effective_memory_mode`.
|
||||
memory_mode: MemoryMode = "pro"
|
||||
evolve_schedule: EvolveSchedule = Field(default_factory=EvolveSchedule)
|
||||
#: SDK retrieval and extraction features a person set; an unset one keeps
|
||||
#: the SDK default. See `one_memory.memory_features`.
|
||||
memory_features: MemoryFeatures = Field(default_factory=MemoryFeatures)
|
||||
|
||||
@field_validator("sources")
|
||||
@classmethod
|
||||
@@ -61,6 +95,9 @@ class ConfigFileError(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
_WRITTEN_WHEN_SET = frozenset({"memory_mode", "evolve_schedule", "memory_features"})
|
||||
|
||||
|
||||
class ConfigStore:
|
||||
"""The only writer of ``<data_dir>/config.json`` inside Core."""
|
||||
|
||||
@@ -107,6 +144,45 @@ class ConfigStore:
|
||||
self._document = updated
|
||||
return updated
|
||||
|
||||
def memory_mode(self) -> MemoryMode:
|
||||
with self._lock:
|
||||
return self._document.memory_mode
|
||||
|
||||
def set_memory_mode(self, mode: MemoryMode) -> ConfigDocument:
|
||||
if mode not in MEMORY_MODES:
|
||||
raise ValueError("invalid_memory_mode")
|
||||
with self._lock:
|
||||
updated = self._document.model_copy(update={"memory_mode": mode})
|
||||
self._write(updated)
|
||||
self._document = updated
|
||||
return updated
|
||||
|
||||
def evolve_schedule(self) -> EvolveSchedule:
|
||||
with self._lock:
|
||||
return self._document.evolve_schedule
|
||||
|
||||
def set_evolve_schedule(self, schedule: EvolveSchedule) -> ConfigDocument:
|
||||
if not isinstance(schedule, EvolveSchedule):
|
||||
raise ValueError("invalid_evolve_schedule")
|
||||
with self._lock:
|
||||
updated = self._document.model_copy(update={"evolve_schedule": schedule})
|
||||
self._write(updated)
|
||||
self._document = updated
|
||||
return updated
|
||||
|
||||
def memory_features(self) -> MemoryFeatures:
|
||||
with self._lock:
|
||||
return self._document.memory_features
|
||||
|
||||
def set_memory_features(self, features: MemoryFeatures) -> ConfigDocument:
|
||||
if not isinstance(features, MemoryFeatures):
|
||||
raise ValueError("invalid_memory_features")
|
||||
with self._lock:
|
||||
updated = self._document.model_copy(update={"memory_features": features})
|
||||
self._write(updated)
|
||||
self._document = updated
|
||||
return updated
|
||||
|
||||
def migrate_extraction_window_preset(self, preset: ExtractionWindowPreset) -> ConfigDocument:
|
||||
"""Adopt the legacy SQL value only when config.json has no value."""
|
||||
if not isinstance(preset, ExtractionWindowPreset):
|
||||
@@ -152,9 +228,16 @@ class ConfigStore:
|
||||
raise ConfigFileError("config_invalid") from error
|
||||
|
||||
def _write(self, document: ConfigDocument) -> None:
|
||||
# Settings added after 1.0.5 are written only once a person chose
|
||||
# them, so an untouched install's file stays what older builds read.
|
||||
unset = {name for name in _WRITTEN_WHEN_SET if name not in document.model_fields_set}
|
||||
values = document.model_dump(mode="json", exclude=unset)
|
||||
if "memory_features" in values:
|
||||
# Only the features a person set; an absent one is the SDK default.
|
||||
values["memory_features"] = document.memory_features.model_dump(exclude_none=True)
|
||||
payload = (
|
||||
json.dumps(
|
||||
document.model_dump(mode="json"),
|
||||
values,
|
||||
ensure_ascii=False,
|
||||
indent=2,
|
||||
sort_keys=True,
|
||||
|
||||
@@ -253,6 +253,9 @@ MODEL_NOT_READY_CODES = frozenset(
|
||||
"model_secret_unavailable",
|
||||
"secure_storage_unavailable",
|
||||
"hy_memory_client_unavailable",
|
||||
# The active profile's embedding is not the live space's: activation
|
||||
# must go through an embedding switch, so the write waits for it.
|
||||
"embedding_space_mismatch",
|
||||
}
|
||||
)
|
||||
|
||||
@@ -2563,6 +2566,11 @@ async def execute_job(
|
||||
in {
|
||||
"model_content_source_disabled",
|
||||
"model_sedimentation_disabled",
|
||||
# Evolve holds this user; the job runs once it is done.
|
||||
"memory_evolve_in_progress",
|
||||
# An embedding switch paused writes: the job stays queued and
|
||||
# sediments into the new space once the switch commits.
|
||||
"embedding_switch_in_progress",
|
||||
}
|
||||
and not hy_invoked
|
||||
):
|
||||
|
||||
@@ -1383,6 +1383,83 @@ async def _050_model_probe_operations_reserved_at(cursor: Any) -> None:
|
||||
)
|
||||
|
||||
|
||||
def embedding_spaces_ddl() -> str:
|
||||
"""The embedding space registry (`embedding_spaces.py`); the main Core SQLite only.
|
||||
|
||||
Unversioned on purpose, like the visibility tables: a migration raises
|
||||
the build's schema version, and a library this build opened would then
|
||||
refuse an older package. The table is additive and created idempotently;
|
||||
an older build ignores it and opens the space its active profile names,
|
||||
which is the registered space on every home that never switched. SQLite
|
||||
only (partial unique indexes): it lives in the main Core database.
|
||||
"""
|
||||
return """
|
||||
CREATE TABLE IF NOT EXISTS embedding_spaces (
|
||||
space_id TEXT PRIMARY KEY NOT NULL,
|
||||
state TEXT NOT NULL CHECK(state IN ('live','building','stray','retired','dropped')),
|
||||
model TEXT,
|
||||
dims INTEGER,
|
||||
name_dims INTEGER,
|
||||
collection_base TEXT NOT NULL,
|
||||
naming TEXT NOT NULL CHECK(naming IN ('legacy','v2','discovered')),
|
||||
graph_dir TEXT,
|
||||
provider_kind TEXT,
|
||||
profile_id TEXT,
|
||||
profile_revision INTEGER,
|
||||
vectors INTEGER,
|
||||
created_at TEXT NOT NULL,
|
||||
live_since TEXT,
|
||||
retired_at TEXT,
|
||||
keep_until TEXT,
|
||||
dropped_at TEXT
|
||||
);
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS embedding_spaces_one_live
|
||||
ON embedding_spaces(state) WHERE state = 'live';
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS embedding_spaces_one_physical
|
||||
ON embedding_spaces(collection_base, COALESCE(name_dims, 0))
|
||||
WHERE state != 'dropped';
|
||||
"""
|
||||
|
||||
|
||||
def embedding_switch_jobs_ddl() -> str:
|
||||
"""Embedding switch and merge jobs (`embedding_switch.py`); unversioned, as above.
|
||||
|
||||
At most one job is open at a time (the partial unique index), and it
|
||||
lives beside the profiles and the registry so the commit is one
|
||||
transaction over all three.
|
||||
"""
|
||||
return """
|
||||
CREATE TABLE IF NOT EXISTS embedding_switch_jobs (
|
||||
job_id TEXT PRIMARY KEY NOT NULL,
|
||||
kind TEXT NOT NULL CHECK(kind IN ('switch','merge')),
|
||||
state TEXT NOT NULL CHECK(state IN (
|
||||
'preparing','reembedding','verifying','committing','done','failed','rolled_back'
|
||||
)),
|
||||
source_space_id TEXT NOT NULL,
|
||||
target_space_id TEXT NOT NULL,
|
||||
stray_space_ids TEXT NOT NULL DEFAULT '[]',
|
||||
target_profile_id TEXT,
|
||||
target_revision INTEGER,
|
||||
prior_profile_id TEXT,
|
||||
requested_by TEXT NOT NULL,
|
||||
plan_json TEXT NOT NULL,
|
||||
progress_json TEXT NOT NULL DEFAULT '{}',
|
||||
verification_json TEXT,
|
||||
error_code TEXT,
|
||||
error_detail TEXT,
|
||||
cancel_requested INTEGER NOT NULL DEFAULT 0,
|
||||
owner_pid INTEGER,
|
||||
heartbeat_at TEXT,
|
||||
attempt INTEGER NOT NULL DEFAULT 1,
|
||||
created_at TEXT NOT NULL,
|
||||
updated_at TEXT NOT NULL,
|
||||
finished_at TEXT
|
||||
);
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS embedding_switch_one_open
|
||||
ON embedding_switch_jobs((1)) WHERE state NOT IN ('done','rolled_back');
|
||||
"""
|
||||
|
||||
|
||||
def model_profile_seeds_ddl(dialect: str = "mysql") -> str:
|
||||
"""Which presets have been seeded, so seeding happens once."""
|
||||
if dialect == "sqlite":
|
||||
|
||||
@@ -160,7 +160,7 @@ def pending(db: sqlite3.Connection, *, limit: int = 100) -> list[str]:
|
||||
|
||||
|
||||
def _trigger_sql() -> list[tuple[str, str]]:
|
||||
"""(table it depends on, statement) for each trigger."""
|
||||
"""(table it is on, statement) for each trigger."""
|
||||
return [
|
||||
(
|
||||
"evidence",
|
||||
@@ -303,9 +303,61 @@ def install_triggers(db: sqlite3.Connection) -> None:
|
||||
for table, statement in _trigger_sql():
|
||||
if table not in present:
|
||||
continue
|
||||
if not _body_tables(statement) <= present:
|
||||
# SQLite resolves a trigger body only when it fires, so a body
|
||||
# naming a dropped table fails the write it rides on (#233: the
|
||||
# SDK cutover drops `governed_memory_evidence`, and every grant
|
||||
# write then failed). Not installed, and an old copy is removed.
|
||||
cursor.execute(f"DROP TRIGGER IF EXISTS {_trigger_name(statement)}")
|
||||
continue
|
||||
cursor.execute(statement)
|
||||
|
||||
|
||||
def _trigger_name(statement: str) -> str:
|
||||
return statement.split("IF NOT EXISTS", 1)[1].split()[0]
|
||||
|
||||
|
||||
def _body_tables(statement: str) -> set[str]:
|
||||
"""Tables the trigger body reads or writes besides the one it is on."""
|
||||
_head, _begin, body = statement.partition("BEGIN")
|
||||
words = body.replace("(", " ").replace(")", " ").split()
|
||||
return {
|
||||
name
|
||||
for keyword, name in zip(words, words[1:], strict=False)
|
||||
if keyword.upper() in {"FROM", "JOIN", "INTO"}
|
||||
}
|
||||
|
||||
|
||||
def drop_stale_triggers(database_path: Path) -> tuple[str, ...]:
|
||||
"""Drop this module's triggers whose body names a table that is gone.
|
||||
|
||||
Startup on SDK authority does not install the cache (it is read only on
|
||||
legacy authority), but a home migrated by an earlier build still carries
|
||||
the triggers it installed. One whose body reads a dropped table fails
|
||||
every write to the table it is on, so it is dropped here. Idempotent.
|
||||
"""
|
||||
with main_database(database_path).write() as db, db:
|
||||
cursor = db.cursor()
|
||||
cursor.row_factory = sqlite3.Row
|
||||
present = {
|
||||
str(row["name"])
|
||||
for row in cursor.execute("SELECT name FROM sqlite_master WHERE type='table'")
|
||||
}
|
||||
installed = {
|
||||
str(row["name"])
|
||||
for row in cursor.execute("SELECT name FROM sqlite_master WHERE type='trigger'")
|
||||
}
|
||||
dropped = tuple(
|
||||
name
|
||||
for _table, statement in _trigger_sql()
|
||||
if (name := _trigger_name(statement)) in installed
|
||||
and not _body_tables(statement) <= present
|
||||
)
|
||||
for name in dropped:
|
||||
cursor.execute(f"DROP TRIGGER IF EXISTS {name}")
|
||||
return dropped
|
||||
|
||||
|
||||
def startup(database_path: Path) -> None:
|
||||
"""Idempotent: create the derived tables and triggers. No version change.
|
||||
|
||||
|
||||
@@ -423,7 +423,7 @@ class ModelProber:
|
||||
ModelProbeOperations(self.store).assert_no_pending(profile.profile_id)
|
||||
# Measuring embedding width may create a new profile version of its own.
|
||||
input_revision = profile.revision
|
||||
profile = self._record_embedding_width(profile, stages)
|
||||
profile, _switch_required = self._record_embedding_width(profile, stages)
|
||||
self.store.record_capabilities(
|
||||
profile.profile_id,
|
||||
expected_revision=profile.revision,
|
||||
@@ -442,32 +442,55 @@ class ModelProber:
|
||||
self,
|
||||
profile: ModelProfile,
|
||||
stages: tuple[ProbeStage, ...],
|
||||
) -> ModelProfile:
|
||||
"""Store the measured vector width on the profile.
|
||||
) -> tuple[ModelProfile, bool]:
|
||||
"""Store the measured vector width on the profile, unless it moves the space.
|
||||
|
||||
The graph store fixes its embedding column width when it is first
|
||||
created, so guessing this leaves every relation write failing later.
|
||||
On the profile that fills the live embedding space a different width
|
||||
is not recorded: it would name other collections (B2). True when the
|
||||
width was refused that way: using it takes an embedding switch, which
|
||||
the switch plan reports (the probe result keeps its shape).
|
||||
"""
|
||||
width = next(
|
||||
(
|
||||
stage.embedding_dimensions
|
||||
for stage in stages
|
||||
if stage.stage == "embedding" and stage.embedding_dimensions
|
||||
),
|
||||
None,
|
||||
)
|
||||
width = measured_width(stages)
|
||||
if width is None or width == profile.embedding_dimensions:
|
||||
return profile
|
||||
return self.store.update(
|
||||
profile.profile_id,
|
||||
ModelProfilePatch(
|
||||
expected_revision=profile.revision,
|
||||
embedding_dimensions=width,
|
||||
return profile, False
|
||||
from one_memory.embedding_spaces import EmbeddingSpaceStore, probe_width_allowed
|
||||
|
||||
spaces = EmbeddingSpaceStore(self.store.database_path)
|
||||
with spaces.transaction() as connection:
|
||||
allowed = probe_width_allowed(
|
||||
connection,
|
||||
profile.profile_id,
|
||||
width,
|
||||
data_dir=self.store.database_path.parent,
|
||||
)
|
||||
if not allowed:
|
||||
return profile, True
|
||||
return (
|
||||
self.store.update(
|
||||
profile.profile_id,
|
||||
ModelProfilePatch(
|
||||
expected_revision=profile.revision,
|
||||
embedding_dimensions=width,
|
||||
),
|
||||
self.secrets,
|
||||
),
|
||||
self.secrets,
|
||||
False,
|
||||
)
|
||||
|
||||
|
||||
def measured_width(stages: tuple[ProbeStage, ...]) -> int | None:
|
||||
return next(
|
||||
(
|
||||
stage.embedding_dimensions
|
||||
for stage in stages
|
||||
if stage.stage == "embedding" and stage.embedding_dimensions
|
||||
),
|
||||
None,
|
||||
)
|
||||
|
||||
|
||||
def _models_request(
|
||||
profile: ModelProfile,
|
||||
secret: str | None,
|
||||
@@ -564,7 +587,7 @@ def _embedding_request(
|
||||
secret: str | None,
|
||||
) -> tuple[str, str, dict[str, str], dict[str, object]]:
|
||||
provider = profile.embedding_provider_kind or profile.provider_kind
|
||||
base_url = (profile.embedding_base_url or profile.base_url or "").rstrip("/")
|
||||
base_url = (profile.embedding_base_url or "").rstrip("/")
|
||||
if provider is ProviderKind.OLLAMA:
|
||||
return (
|
||||
"POST",
|
||||
|
||||
@@ -13,6 +13,7 @@ from pydantic import BaseModel, ConfigDict, Field, field_validator
|
||||
|
||||
from one_memory.database import take_the_write_lock_first
|
||||
from one_memory.database_access import main_database
|
||||
from one_memory.embedding_spaces import probe_width_allowed
|
||||
from one_memory.memory_server.migrations import model_probe_operations_ddl
|
||||
from one_memory.model_probe import ModelProbeResult, ProbeStage
|
||||
from one_memory.model_profiles import ModelCapabilities, ModelProfileStore
|
||||
@@ -349,7 +350,19 @@ class ModelProbeOperations:
|
||||
)
|
||||
revision = profile["revision"]
|
||||
now = datetime.now(UTC).isoformat()
|
||||
if width is not None and width != profile["embedding_dimensions"]:
|
||||
# The profile that fills the live embedding space records only the
|
||||
# space's own width: another one would move writes to new, empty
|
||||
# collections (B2). Using it takes an embedding switch.
|
||||
if (
|
||||
width is not None
|
||||
and width != profile["embedding_dimensions"]
|
||||
and probe_width_allowed(
|
||||
db,
|
||||
profile["profile_id"],
|
||||
width,
|
||||
data_dir=self.store.database_path.parent,
|
||||
)
|
||||
):
|
||||
revision += 1
|
||||
db.execute(
|
||||
"UPDATE model_profiles SET embedding_dimensions=?,revision=? "
|
||||
|
||||
@@ -97,7 +97,6 @@ class ModelProfileInput(BaseModel):
|
||||
)
|
||||
embedding_provider, resolved_embedding_base_url = resolve_embedding_endpoint(
|
||||
provider_kind=self.provider_kind,
|
||||
base_url=base_url,
|
||||
embedding_model=embedding_model,
|
||||
embedding_provider_kind=embedding_provider_kind,
|
||||
embedding_base_url=embedding_base_url,
|
||||
@@ -107,6 +106,12 @@ class ModelProfileInput(BaseModel):
|
||||
ProviderKind.ANTHROPIC,
|
||||
}:
|
||||
raise ValueError("embedding provider must be OpenAI-compatible, Ollama, or local ONNX")
|
||||
if (
|
||||
embedding_provider is not ProviderKind.LOCAL_ONNX
|
||||
and resolved_embedding_base_url is None
|
||||
):
|
||||
# `ModelProfileStore.create` refuses it with its own code.
|
||||
return self
|
||||
if (
|
||||
embedding_provider is not ProviderKind.LOCAL_ONNX
|
||||
and resolved_embedding_base_url is not None
|
||||
@@ -308,11 +313,12 @@ class ModelProfileStore:
|
||||
f"""
|
||||
UPDATE {table}
|
||||
SET embedding_provider_kind = provider_kind,
|
||||
embedding_base_url = base_url
|
||||
embedding_base_url = COALESCE(embedding_base_url, base_url)
|
||||
WHERE embedding_model IS NOT NULL
|
||||
AND embedding_provider_kind IS NULL
|
||||
"""
|
||||
)
|
||||
_pin_embedding_endpoints(connection, table)
|
||||
connection.execute(
|
||||
"""
|
||||
INSERT OR IGNORE INTO model_profile_versions (
|
||||
@@ -355,11 +361,11 @@ class ModelProfileStore:
|
||||
)
|
||||
embedding_provider_kind, embedding_base_url = resolve_embedding_endpoint(
|
||||
provider_kind=value.provider_kind,
|
||||
base_url=base_url,
|
||||
embedding_model=defaulted_model,
|
||||
embedding_provider_kind=defaulted_provider,
|
||||
embedding_base_url=defaulted_base_url,
|
||||
)
|
||||
require_embedding_base_url(embedding_provider_kind, embedding_base_url)
|
||||
try:
|
||||
with self._connect(writing=True) as connection:
|
||||
connection.execute(
|
||||
@@ -612,7 +618,6 @@ class ModelProfileStore:
|
||||
)
|
||||
embedding_provider_kind, embedding_base_url = resolve_embedding_endpoint(
|
||||
provider_kind=current.provider_kind,
|
||||
base_url=base_url,
|
||||
embedding_model=embedding_model,
|
||||
embedding_provider_kind=(
|
||||
patch.embedding_provider_kind
|
||||
@@ -627,6 +632,7 @@ class ModelProfileStore:
|
||||
)
|
||||
if embedding_provider_kind in {ProviderKind.HY_DEFAULT, ProviderKind.ANTHROPIC}:
|
||||
raise ValueError("embedding provider must be OpenAI-compatible, Ollama, or local ONNX")
|
||||
require_embedding_base_url(embedding_provider_kind, embedding_base_url)
|
||||
if (
|
||||
embedding_provider_kind is not ProviderKind.LOCAL_ONNX
|
||||
and embedding_base_url is not None
|
||||
@@ -697,10 +703,16 @@ class ModelProfileStore:
|
||||
}
|
||||
)
|
||||
# User changes to service/model/credentials need a new explicit activation.
|
||||
# Probe-discovered embedding width alone retains its existing behavior.
|
||||
# Probe-discovered embedding width alone retains its existing behavior;
|
||||
# replacing one declared width with another is a new embedding
|
||||
# identity, and activating it goes through an embedding switch.
|
||||
deactivate = (
|
||||
patch.api_key is not None
|
||||
or patch.clear_api_key
|
||||
or (
|
||||
row["embedding_dimensions"] is not None
|
||||
and next_values["embedding_dimensions"] != row["embedding_dimensions"]
|
||||
)
|
||||
or any(
|
||||
row[key] != next_values[key]
|
||||
for key in (
|
||||
@@ -1167,21 +1179,77 @@ def attach_disable_thinking(
|
||||
def resolve_embedding_endpoint(
|
||||
*,
|
||||
provider_kind: ProviderKind,
|
||||
base_url: str | None,
|
||||
embedding_model: str | None,
|
||||
embedding_provider_kind: ProviderKind | None,
|
||||
embedding_base_url: str | None,
|
||||
) -> tuple[ProviderKind | None, str | None]:
|
||||
"""The embedding provider and the endpoint exactly as given.
|
||||
|
||||
The endpoint never follows the generation `base_url`: the collection name
|
||||
hashes it, so a proxy or region edit of the generation endpoint moved
|
||||
writes to a new, empty collection, and providers without embeddings were
|
||||
sent embedding requests. `require_embedding_base_url` refuses a remote
|
||||
embedding without one.
|
||||
"""
|
||||
if embedding_model is None:
|
||||
return None, None
|
||||
if embedding_provider_kind is ProviderKind.LOCAL_ONNX:
|
||||
return ProviderKind.LOCAL_ONNX, None
|
||||
if embedding_provider_kind is None:
|
||||
return provider_kind, embedding_base_url or base_url
|
||||
return (
|
||||
embedding_provider_kind,
|
||||
embedding_base_url or DEFAULT_BASE_URLS.get(embedding_provider_kind),
|
||||
)
|
||||
return embedding_provider_kind or provider_kind, embedding_base_url
|
||||
|
||||
|
||||
class EmbeddingBaseUrlRequiredError(ValueError):
|
||||
def __init__(self) -> None:
|
||||
super().__init__("embedding_base_url_required")
|
||||
|
||||
|
||||
def require_embedding_base_url(provider: ProviderKind | None, base_url: str | None) -> None:
|
||||
"""A remote embedding must name its endpoint; the UI may offer the generation one."""
|
||||
if provider not in {None, ProviderKind.LOCAL_ONNX} and not base_url:
|
||||
raise EmbeddingBaseUrlRequiredError()
|
||||
|
||||
|
||||
def pinned_embedding_base_url(
|
||||
*,
|
||||
provider_kind: ProviderKind | str,
|
||||
base_url: str | None,
|
||||
embedding_model: str | None,
|
||||
embedding_provider_kind: ProviderKind | str | None,
|
||||
embedding_base_url: str | None,
|
||||
) -> str | None:
|
||||
"""The endpoint a stored profile's embedding resolved to before it was pinned.
|
||||
|
||||
The Runtime used `embedding_base_url or base_url`, and the provider default
|
||||
only when both were empty (no collection could have existed then). Writing
|
||||
that value keeps the endpoint, and so the collection name, byte-identical.
|
||||
"""
|
||||
provider = ProviderKind(embedding_provider_kind or provider_kind)
|
||||
if embedding_model is None or embedding_base_url or provider is ProviderKind.LOCAL_ONNX:
|
||||
return embedding_base_url
|
||||
return base_url or DEFAULT_BASE_URLS.get(provider)
|
||||
|
||||
|
||||
def _pin_embedding_endpoints(connection: sqlite3.Connection, table: str) -> None:
|
||||
"""Write the resolved endpoint into every remote embedding left without one."""
|
||||
rows = connection.execute(
|
||||
f"SELECT rowid, provider_kind, base_url, embedding_model, embedding_provider_kind "
|
||||
f"FROM {table} WHERE embedding_model IS NOT NULL AND embedding_base_url IS NULL "
|
||||
"AND COALESCE(embedding_provider_kind, provider_kind) != 'local_onnx'"
|
||||
).fetchall()
|
||||
for row in rows:
|
||||
connection.execute(
|
||||
f"UPDATE {table} SET embedding_base_url = ? WHERE rowid = ?",
|
||||
(
|
||||
pinned_embedding_base_url(
|
||||
provider_kind=row["provider_kind"],
|
||||
base_url=row["base_url"],
|
||||
embedding_model=row["embedding_model"],
|
||||
embedding_provider_kind=row["embedding_provider_kind"],
|
||||
embedding_base_url=None,
|
||||
),
|
||||
row["rowid"],
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def validate_model_base_url(value: str) -> None:
|
||||
|
||||
@@ -14,6 +14,7 @@ from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
from one_memory.hy_sdk_types import HySdkMemoryRecord
|
||||
from one_memory.memory_governance import GovernedMemory, MemoryStatus
|
||||
from one_memory.sdk_search_options import SdkSearchOptions
|
||||
from one_memory.semantic_deadline import (
|
||||
SemanticMemoryHit,
|
||||
SemanticSearch,
|
||||
@@ -50,7 +51,15 @@ class TruncationBudget(StrEnum):
|
||||
TIMEOUT = "timeout_budget"
|
||||
|
||||
|
||||
class RecallRequest(BaseModel):
|
||||
class RecallRequest(SdkSearchOptions):
|
||||
"""`POST /api/v1/recall`.
|
||||
|
||||
The `SdkSearchOptions` fields, `schema_limit` and `intention_limit` pass
|
||||
through to the SDK memory search (and the floor and row filters to the
|
||||
raw fallback). Absent, recall makes the search it always made: the
|
||||
relevance floor, schema and intention lanes sized to `limit`.
|
||||
"""
|
||||
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
query: str = Field(min_length=1, max_length=2_000)
|
||||
@@ -70,6 +79,17 @@ class RecallRequest(BaseModel):
|
||||
timeout_ms: int = Field(default=2_000, ge=50, le=30_000)
|
||||
#: Whose memories to recall. Omitted, it is `local`, the App's library.
|
||||
user_id: UserId = DEFAULT_USER_ID
|
||||
#: The SDK's schema (L6) and proactive intention (L7) lanes.
|
||||
schema_limit: int | None = Field(default=None, ge=0, le=20)
|
||||
intention_limit: int | None = Field(default=None, ge=0, le=20)
|
||||
|
||||
def search_overrides(self) -> Mapping[str, object]:
|
||||
values = dict(super().search_overrides())
|
||||
if self.schema_limit is not None:
|
||||
values["schema_limit"] = self.schema_limit
|
||||
if self.intention_limit is not None:
|
||||
values["intention_limit"] = self.intention_limit
|
||||
return values
|
||||
|
||||
|
||||
class RecallHit(BaseModel):
|
||||
|
||||
@@ -6,7 +6,7 @@ import asyncio
|
||||
import json
|
||||
from collections.abc import AsyncIterator, Awaitable, Callable
|
||||
from contextlib import aclosing
|
||||
from typing import Any, Literal, Protocol
|
||||
from typing import Any, Literal, Protocol, cast
|
||||
|
||||
import httpx
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
@@ -14,7 +14,7 @@ from pydantic import BaseModel, ConfigDict
|
||||
from one_memory.http_client import create_async_client
|
||||
from one_memory.hy_runtime_config import HY_MEMORY_READER
|
||||
from one_memory.hy_sdk_runtime import sdk_relevance_floor
|
||||
from one_memory.hy_sdk_types import HySdkMemoryRecord, parse_sdk_search
|
||||
from one_memory.hy_sdk_types import HySdkMemoryRecord, parse_sdk_lane_search, parse_sdk_search
|
||||
from one_memory.memory_qa import (
|
||||
MemoryAnswerError,
|
||||
_DeferredSecretStore,
|
||||
@@ -31,6 +31,7 @@ from one_memory.model_profiles import (
|
||||
ProviderKind,
|
||||
attach_disable_thinking,
|
||||
)
|
||||
from one_memory.sdk_search_options import SdkSearchFilters, SearchLane
|
||||
from one_memory.secrets import SecretStore, SecretStoreUnavailableError
|
||||
from one_memory.user_scope import owns
|
||||
|
||||
@@ -96,22 +97,39 @@ class SdkMemoryAnswerer:
|
||||
question: str,
|
||||
*,
|
||||
user_id: str = "local",
|
||||
filters: SdkSearchFilters | None = None,
|
||||
) -> tuple[SdkAskMemory, ...]:
|
||||
"""The user's active memories nearest the question.
|
||||
|
||||
`filters` are a caller's pass-through search options; absent, the
|
||||
search is the one ask always made.
|
||||
"""
|
||||
client = await self.client_handle.client()
|
||||
overrides = filters.search_overrides() if filters is not None else {}
|
||||
raw = await asyncio.to_thread(
|
||||
client.search,
|
||||
question,
|
||||
user_ids=[user_id],
|
||||
agent_ids=[],
|
||||
limit=12,
|
||||
min_score=sdk_relevance_floor(),
|
||||
schema_min_score=sdk_relevance_floor(),
|
||||
schema_limit=12,
|
||||
intention_limit=12,
|
||||
scene="normal",
|
||||
reader=HY_MEMORY_READER,
|
||||
**{
|
||||
"user_ids": [user_id],
|
||||
"agent_ids": [],
|
||||
"limit": 12,
|
||||
"min_score": sdk_relevance_floor(),
|
||||
"schema_min_score": sdk_relevance_floor(),
|
||||
"schema_limit": 12,
|
||||
"intention_limit": 12,
|
||||
"scene": "normal",
|
||||
"reader": HY_MEMORY_READER,
|
||||
**overrides,
|
||||
},
|
||||
)
|
||||
hits = parse_sdk_search(raw, error_code="hy_memory_search_shape_unknown")
|
||||
if filters is not None and overrides:
|
||||
hits = tuple(
|
||||
hit
|
||||
for hit in parse_sdk_lane_search(raw, error_code="hy_memory_search_shape_unknown")
|
||||
if filters.passes_floor(cast(SearchLane, hit.lane), hit.score)
|
||||
)
|
||||
else:
|
||||
hits = parse_sdk_search(raw, error_code="hy_memory_search_shape_unknown")
|
||||
found: list[SdkAskMemory] = []
|
||||
seen: set[str] = set()
|
||||
for hit in hits:
|
||||
@@ -134,8 +152,9 @@ class SdkMemoryAnswerer:
|
||||
question: str,
|
||||
*,
|
||||
user_id: str = "local",
|
||||
filters: SdkSearchFilters | None = None,
|
||||
) -> SdkMemoryAnswer:
|
||||
candidates = await self.retrieve(question, user_id=user_id)
|
||||
candidates = await self.retrieve(question, user_id=user_id, filters=filters)
|
||||
if not candidates:
|
||||
return SdkMemoryAnswer(answer=_NO_GROUNDING, memories=(), mode="empty")
|
||||
profile = await self._active_profile()
|
||||
@@ -166,8 +185,9 @@ class SdkMemoryAnswerer:
|
||||
question: str,
|
||||
*,
|
||||
user_id: str = "local",
|
||||
filters: SdkSearchFilters | None = None,
|
||||
) -> AsyncIterator[dict[str, object]]:
|
||||
candidates = await self.retrieve(question, user_id=user_id)
|
||||
candidates = await self.retrieve(question, user_id=user_id, filters=filters)
|
||||
await self._require_current(candidates)
|
||||
yield {"type": "recall", "memories": candidates}
|
||||
if not candidates:
|
||||
|
||||
@@ -5,8 +5,8 @@ from __future__ import annotations
|
||||
import asyncio
|
||||
import logging
|
||||
import re
|
||||
from collections.abc import AsyncIterator, Callable, Iterable
|
||||
from contextlib import AsyncExitStack, asynccontextmanager, suppress
|
||||
from collections.abc import AsyncIterator, Callable
|
||||
from contextlib import AsyncExitStack, asynccontextmanager
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
@@ -33,16 +33,6 @@ from one_memory.semantic_deadline import SemanticMemoryHit
|
||||
_STABLE_CODE = re.compile(r"[a-z][a-z0-9_]{0,95}")
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
RAW_RETENTION_DAYS = 3
|
||||
RAW_RETENTION_CHECK_SECONDS = 6 * 60 * 60
|
||||
_RAW_STATUSES = (
|
||||
"active",
|
||||
"superseded",
|
||||
"shadow",
|
||||
"negated",
|
||||
"conflicted",
|
||||
"archived",
|
||||
)
|
||||
|
||||
|
||||
class SdkMemoryRuntime:
|
||||
@@ -56,82 +46,20 @@ class SdkMemoryRuntime:
|
||||
write_queue: DurableWriteQueue,
|
||||
extraction_settings: ExtractionSettingsStore,
|
||||
data_dir: Path,
|
||||
user_ids: Callable[[], Iterable[str]] = lambda: ("local",),
|
||||
memory_mode: Callable[[], str] = lambda: "pro",
|
||||
writes_open: Callable[[], bool] = lambda: True,
|
||||
) -> None:
|
||||
self.coordinator = coordinator
|
||||
#: Every user whose raw layer is kept to the retention window.
|
||||
self.user_ids = user_ids
|
||||
#: The SDK mode, for the graph store's path (`hy_graph_admin`).
|
||||
self.memory_mode = memory_mode
|
||||
#: False while an embedding switch pauses model writes.
|
||||
self.writes_open = writes_open
|
||||
self.read_handle = read_handle
|
||||
self.write_queue = write_queue
|
||||
self.extraction_settings = extraction_settings
|
||||
self.data_dir = data_dir
|
||||
self._raw_retention_task: asyncio.Task[None] | None = None
|
||||
|
||||
def start_raw_retention(self) -> None:
|
||||
if self._raw_retention_task is None or self._raw_retention_task.done():
|
||||
self._raw_retention_task = asyncio.create_task(
|
||||
self._raw_retention_loop(),
|
||||
name="one-memory-sdk-raw-retention",
|
||||
)
|
||||
|
||||
async def prune_expired_raw(self, *, now: datetime | None = None) -> int:
|
||||
deleted = 0
|
||||
for user_id in dict.fromkeys(await asyncio.to_thread(lambda: tuple(self.user_ids()))):
|
||||
deleted += await self._prune_expired_raw_for(user_id, now=now)
|
||||
return deleted
|
||||
|
||||
async def _prune_expired_raw_for(self, user_id: str, *, now: datetime | None) -> int:
|
||||
cutoff = (now or datetime.now(UTC)) - timedelta(days=RAW_RETENTION_DAYS)
|
||||
cutoff_timestamp = int(cutoff.timestamp())
|
||||
client = await self.read_handle.client()
|
||||
deleted = 0
|
||||
while True:
|
||||
page = await asyncio.to_thread(
|
||||
client.list_memories,
|
||||
user_id=user_id,
|
||||
layers=["l1_raw"],
|
||||
statuses=list(_RAW_STATUSES),
|
||||
order="asc",
|
||||
limit=500,
|
||||
offset=0,
|
||||
)
|
||||
bucket = page.get("vdb") if isinstance(page, dict) else None
|
||||
rows = bucket.get("memories") if isinstance(bucket, dict) else None
|
||||
if not isinstance(rows, list) or not rows:
|
||||
return deleted
|
||||
expired = [
|
||||
row
|
||||
for row in rows
|
||||
if isinstance(row, dict)
|
||||
and isinstance(row.get("memory_id"), str)
|
||||
and int(row.get("memory_at") or row.get("gmt_created") or 0) < cutoff_timestamp
|
||||
]
|
||||
if not expired:
|
||||
return deleted
|
||||
removed_page = 0
|
||||
for row in expired:
|
||||
result = await asyncio.to_thread(client.delete, row["memory_id"])
|
||||
if isinstance(result, dict) and result.get("success") is True:
|
||||
count = int(result.get("deleted_count") or 0)
|
||||
deleted += count
|
||||
removed_page += count
|
||||
if removed_page == 0:
|
||||
return deleted
|
||||
|
||||
async def _raw_retention_loop(self) -> None:
|
||||
while True:
|
||||
await asyncio.sleep(RAW_RETENTION_CHECK_SECONDS)
|
||||
try:
|
||||
removed = await self.prune_expired_raw()
|
||||
if removed:
|
||||
logger.info("sdk_raw_retention_deleted count=%d", removed)
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception as error:
|
||||
logger.warning(
|
||||
"sdk_raw_retention_failed error_type=%s",
|
||||
type(error).__name__,
|
||||
)
|
||||
# MemHub keeps every raw the SDK stores, in every mode: there is no
|
||||
# retention window (owner decision; the SDK itself never deletes raw).
|
||||
|
||||
async def semantic_recall(
|
||||
self,
|
||||
@@ -186,6 +114,8 @@ class SdkMemoryRuntime:
|
||||
)
|
||||
|
||||
async def digest(self) -> dict[str, object]:
|
||||
if not self.writes_open():
|
||||
raise HySdkError("embedding_switch_in_progress")
|
||||
async with self.coordinator.use_client() as client:
|
||||
try:
|
||||
raw = await asyncio.to_thread(
|
||||
@@ -226,11 +156,14 @@ class SdkMemoryRuntime:
|
||||
if str(error) == "model_profile_not_active":
|
||||
return None
|
||||
raise
|
||||
space = await self.coordinator.live_space()
|
||||
return await asyncio.to_thread(
|
||||
inspect_profile_graph_width,
|
||||
profile,
|
||||
api_key=secret,
|
||||
data_dir=self.data_dir,
|
||||
mode=self.memory_mode(),
|
||||
graph_dir=space.graph_dir if space is not None and space.graph_dir else "graph",
|
||||
)
|
||||
|
||||
async def decide_graph_width(
|
||||
@@ -241,12 +174,15 @@ class SdkMemoryRuntime:
|
||||
profile, secret = await self.coordinator.active_profile_credentials()
|
||||
except RuntimeError as error:
|
||||
raise HySdkError(str(error)) from error
|
||||
space = await self.coordinator.live_space()
|
||||
return await asyncio.to_thread(
|
||||
decide_profile_graph_width,
|
||||
profile,
|
||||
api_key=secret,
|
||||
data_dir=self.data_dir,
|
||||
decision=decision,
|
||||
mode=self.memory_mode(),
|
||||
graph_dir=space.graph_dir if space is not None and space.graph_dir else "graph",
|
||||
)
|
||||
|
||||
@asynccontextmanager
|
||||
@@ -256,6 +192,8 @@ class SdkMemoryRuntime:
|
||||
Callers record these against the edit or deletion instead of failing
|
||||
the request that already committed it.
|
||||
"""
|
||||
if not self.writes_open():
|
||||
raise HySdkError("embedding_switch_in_progress")
|
||||
async with AsyncExitStack() as stack:
|
||||
try:
|
||||
client = await stack.enter_async_context(self.coordinator.use_client())
|
||||
@@ -385,12 +323,7 @@ class SdkMemoryRuntime:
|
||||
return None
|
||||
|
||||
async def shutdown(self) -> None:
|
||||
if self._raw_retention_task is None:
|
||||
return
|
||||
self._raw_retention_task.cancel()
|
||||
with suppress(asyncio.CancelledError):
|
||||
await self._raw_retention_task
|
||||
self._raw_retention_task = None
|
||||
return None
|
||||
|
||||
|
||||
def _agent_ids(raw: object) -> tuple[str, ...]:
|
||||
|
||||
@@ -23,6 +23,7 @@ from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from one_memory.database_access import main_database
|
||||
from one_memory.hy_sdk_runtime import builtin_retrieval_snapshot
|
||||
from one_memory.imports.models import EvidenceRecord
|
||||
from one_memory.imports.storage import EvidenceStore
|
||||
from one_memory.memory_server.contracts import (
|
||||
@@ -51,6 +52,7 @@ from one_memory.sdk_migration_bundle import inspect_sdk_migration_bundle
|
||||
from one_memory.sdk_migration_verify import (
|
||||
MemoryReadbackFailed,
|
||||
ReadbackReport,
|
||||
expected_memory_ids,
|
||||
legacy_memory_ids_from,
|
||||
legacy_memory_ids_from_bundle,
|
||||
read_active_profile,
|
||||
@@ -61,6 +63,8 @@ from one_memory.sdk_vector_reattach import (
|
||||
ReattachedCollection,
|
||||
legacy_vector_collections,
|
||||
reattach_legacy_vectors,
|
||||
stored_memory_ids,
|
||||
vectors_belong_to_another_space,
|
||||
)
|
||||
from one_memory.target_schema import drop_legacy_tables, legacy_tables_present
|
||||
|
||||
@@ -457,13 +461,17 @@ def resume_legacy_cleanup(*, data_dir: Path, database_path: Path) -> tuple[str,
|
||||
bundle_id=marker.backup_bundle_id,
|
||||
)
|
||||
from one_memory.hy_runtime_config import HY_MEMORY_NAMESPACE
|
||||
from one_memory.memory_visibility import drop_stale_triggers
|
||||
|
||||
recover_abandoned_staging(
|
||||
data_dir=data_dir,
|
||||
database_path=database_path,
|
||||
namespace=HY_MEMORY_NAMESPACE,
|
||||
)
|
||||
return drop_legacy_tables(database_path)
|
||||
dropped = drop_legacy_tables(database_path)
|
||||
# The visibility cache's grant triggers read a table just dropped (#233).
|
||||
drop_stale_triggers(database_path)
|
||||
return dropped
|
||||
|
||||
|
||||
MIGRATION_PREFIX = "sdk-authority-v2"
|
||||
@@ -700,13 +708,14 @@ def run_startup_migration(
|
||||
step["authority"] = authority.value
|
||||
log.migration_id = authority_migration
|
||||
with log.step("preflight.profile_read") as step:
|
||||
profile = read_active_profile(database_path)
|
||||
step["profile_present"] = profile is not None
|
||||
active = read_active_profile(database_path)
|
||||
step["profile_present"] = active is not None
|
||||
# No active profile means no generation model, not no embedding: the
|
||||
# vectors are the built-in model's, which is also what the Runtime reads.
|
||||
profile = active if active is not None else builtin_retrieval_snapshot()
|
||||
if authority is MemoryAuthority.SDK:
|
||||
with log.step("preflight.leftovers_scan") as step:
|
||||
vectors_waiting = profile is not None and bool(
|
||||
legacy_vector_collections(data_dir, profile)
|
||||
)
|
||||
vectors_waiting = bool(legacy_vector_collections(data_dir, profile))
|
||||
leftovers = bool(legacy_tables_present(database_path)) or abandoned_staging_present(
|
||||
data_dir=data_dir, namespace=namespace
|
||||
)
|
||||
@@ -717,7 +726,6 @@ def run_startup_migration(
|
||||
reattached: tuple[ReattachedCollection, ...] = ()
|
||||
readback: ReadbackReport | None = None
|
||||
if vectors_waiting:
|
||||
assert profile is not None
|
||||
progress("vectors")
|
||||
with log.step("vectors.reattach") as step:
|
||||
reattached = reattach_legacy_vectors(data_dir, profile)
|
||||
@@ -769,12 +777,12 @@ def run_startup_migration(
|
||||
step["chat"] = len(legacy.chat)
|
||||
step["productivity"] = len(legacy.productivity)
|
||||
step["users"] = len(legacy.users)
|
||||
if profile is None and not legacy.empty:
|
||||
# Without the embedding profile the Runtime's collections cannot be
|
||||
# named, so neither reattached nor read back. Legacy stays intact.
|
||||
log.check("preflight.model_profile_check", "sdk_migration_model_profile_required")
|
||||
raise DataPlaneMigrationError("sdk_migration_model_profile_required")
|
||||
log.check("preflight.model_profile_check", None)
|
||||
if active is None and not legacy.empty and vectors_belong_to_another_space(data_dir, profile):
|
||||
# The memories live in the space of a profile that is no longer active,
|
||||
# so the built-in model cannot read them. Legacy stays intact.
|
||||
log.check("preflight.vector_space_check", "sdk_migration_vector_profile_missing")
|
||||
raise DataPlaneMigrationError("sdk_migration_vector_profile_missing")
|
||||
log.check("preflight.vector_space_check", None)
|
||||
with capture_fence(data_dir):
|
||||
# Builds through 19 left a staged SDK copy beside the original store.
|
||||
# Settling it before the bundle means the bundle holds the original
|
||||
@@ -830,17 +838,23 @@ def run_startup_migration(
|
||||
step.update(preparation.as_document())
|
||||
reattached = ()
|
||||
readback = None
|
||||
if profile is not None:
|
||||
progress("vectors")
|
||||
with log.step("vectors.reattach") as step:
|
||||
reattached = reattach_legacy_vectors(data_dir, profile)
|
||||
step["collections"] = len(reattached)
|
||||
if not legacy.empty:
|
||||
assert profile is not None
|
||||
# Before reattach touches the store, and after the staging
|
||||
# recovery above has settled which store is the original.
|
||||
with log.step("verify.stored_ids") as step:
|
||||
expected, missing = expected_memory_ids(legacy, stored_memory_ids(data_dir))
|
||||
step["chat_expected"] = len(expected.chat)
|
||||
step["chat_missing_before_migration"] = missing
|
||||
progress("vectors")
|
||||
with log.step("vectors.reattach") as step:
|
||||
reattached = reattach_legacy_vectors(data_dir, profile)
|
||||
step["collections"] = len(reattached)
|
||||
if not expected.empty:
|
||||
progress("verify")
|
||||
with log.step("verify.readback") as step:
|
||||
with client_for(profile) as client:
|
||||
readback = verify_readback(client, legacy)
|
||||
readback = verify_readback(
|
||||
client, expected, chat_missing_before_migration=missing
|
||||
)
|
||||
step.update(readback.as_document())
|
||||
if not readback.passed:
|
||||
store.set_phase(migration_id, phase="readback_failed", state="blocked")
|
||||
|
||||
@@ -14,12 +14,12 @@ import json
|
||||
import sqlite3
|
||||
from collections.abc import Iterator
|
||||
from contextlib import closing, contextmanager
|
||||
from dataclasses import dataclass
|
||||
from dataclasses import dataclass, replace
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from one_memory.database_access import main_database
|
||||
from one_memory.model_profiles import ModelProfileSnapshot
|
||||
from one_memory.model_profiles import ModelProfileSnapshot, pinned_embedding_base_url
|
||||
from one_memory.sdk_data_plane_migration import DataPlaneMigrationError
|
||||
|
||||
_PRODUCTIVITY_KINDS = frozenset({"coding", "working"})
|
||||
@@ -47,6 +47,8 @@ class ReadbackReport:
|
||||
chat_readable: int
|
||||
productivity_expected: int
|
||||
productivity_readable: int
|
||||
#: Confirmed in the mirror, gone from the SDK before this migration began.
|
||||
chat_missing_before_migration: int = 0
|
||||
|
||||
@property
|
||||
def passed(self) -> bool:
|
||||
@@ -63,6 +65,7 @@ class ReadbackReport:
|
||||
"chat_readable": self.chat_readable,
|
||||
"productivity_expected": self.productivity_expected,
|
||||
"productivity_readable": self.productivity_readable,
|
||||
"chat_missing_before_migration": self.chat_missing_before_migration,
|
||||
"passed": self.passed,
|
||||
}
|
||||
|
||||
@@ -115,6 +118,27 @@ def read_legacy_memory_ids(connection: sqlite3.Connection) -> LegacyMemoryIds:
|
||||
)
|
||||
|
||||
|
||||
def expected_memory_ids(
|
||||
legacy: LegacyMemoryIds, stored: frozenset[str]
|
||||
) -> tuple[LegacyMemoryIds, int]:
|
||||
"""What the readback must find, and how many confirmed chats it need not.
|
||||
|
||||
A 1.0.4 library can hold confirmed rows whose SDK node the SDK had already
|
||||
merged or deleted: no migration could read those back (658 confirmed, 657
|
||||
in the store, and the whole upgrade failed). So chat is held to what the
|
||||
store held before this run, and anything in there must still read back.
|
||||
|
||||
If none of the confirmed chats is in the store, that is not history: the
|
||||
store is empty or is the wrong one. Nothing is excused then, so the
|
||||
readback still fails on every id instead of checking nothing just before
|
||||
the mirror is dropped.
|
||||
"""
|
||||
chat = tuple(memory_id for memory_id in legacy.chat if memory_id in stored)
|
||||
if not chat:
|
||||
return legacy, 0
|
||||
return replace(legacy, chat=chat), len(legacy.chat) - len(chat)
|
||||
|
||||
|
||||
def legacy_memory_ids_from(database_path: Path) -> LegacyMemoryIds:
|
||||
with main_database(database_path).read() as connection:
|
||||
return read_legacy_memory_ids(connection)
|
||||
@@ -160,7 +184,14 @@ def read_active_profile(database_path: Path) -> ModelProfileSnapshot | None:
|
||||
model=row["model"],
|
||||
embedding_model=row["embedding_model"],
|
||||
embedding_provider_kind=row["embedding_provider_kind"],
|
||||
embedding_base_url=row["embedding_base_url"],
|
||||
# Read-only, beside a Core that may not have pinned it yet.
|
||||
embedding_base_url=pinned_embedding_base_url(
|
||||
provider_kind=row["provider_kind"],
|
||||
base_url=row["base_url"],
|
||||
embedding_model=row["embedding_model"],
|
||||
embedding_provider_kind=row["embedding_provider_kind"],
|
||||
embedding_base_url=row["embedding_base_url"],
|
||||
),
|
||||
embedding_dimensions=row["embedding_dimensions"],
|
||||
timeout_seconds=row["timeout_seconds"],
|
||||
max_concurrency=row["max_concurrency"],
|
||||
@@ -179,22 +210,49 @@ def runtime_sdk_client(
|
||||
from one_memory.model_profiles import ModelProfileStore
|
||||
from one_memory.secrets import KeyringSecretStore, SecretStoreUnavailableError
|
||||
|
||||
try:
|
||||
secret = ModelProfileStore(database_path).resolve_snapshot_secret(
|
||||
profile.profile_id, profile.revision, KeyringSecretStore()
|
||||
)
|
||||
except SecretStoreUnavailableError:
|
||||
raise DataPlaneMigrationError("sdk_migration_secure_storage_unavailable") from None
|
||||
if profile.has_api_key and secret is None:
|
||||
raise DataPlaneMigrationError("sdk_migration_model_secret_unavailable")
|
||||
lease = acquire_client(profile, api_key=secret, data_dir=data_dir)
|
||||
secret: str | None = None
|
||||
# Without a key there is nothing to resolve, and the built-in retrieval
|
||||
# snapshot has no stored version to look one up by.
|
||||
if profile.has_api_key:
|
||||
try:
|
||||
secret = ModelProfileStore(database_path).resolve_snapshot_secret(
|
||||
profile.profile_id, profile.revision, KeyringSecretStore()
|
||||
)
|
||||
except SecretStoreUnavailableError:
|
||||
raise DataPlaneMigrationError("sdk_migration_secure_storage_unavailable") from None
|
||||
if secret is None:
|
||||
raise DataPlaneMigrationError("sdk_migration_model_secret_unavailable")
|
||||
# The registry's live space when this home has one and the profile fills
|
||||
# it; otherwise the space the profile named before the registry existed.
|
||||
# A switch in progress never changes it: the profile and the live space
|
||||
# change together, at the switch's commit.
|
||||
space = live_space_for(database_path, profile)
|
||||
lease = (
|
||||
acquire_client(profile, space, api_key=secret, data_dir=data_dir)
|
||||
if space is not None
|
||||
else acquire_client(profile, api_key=secret, data_dir=data_dir)
|
||||
)
|
||||
try:
|
||||
yield lease.client
|
||||
finally:
|
||||
lease.close()
|
||||
|
||||
|
||||
def verify_readback(client: Any, legacy: LegacyMemoryIds) -> ReadbackReport:
|
||||
def live_space_for(database_path: Path, profile: ModelProfileSnapshot) -> Any:
|
||||
"""The live embedding space `profile` reads, or None to derive it from the profile."""
|
||||
from one_memory.embedding_spaces import embedding_identity, read_live_space
|
||||
|
||||
space = read_live_space(database_path)
|
||||
if space is None or space.identity is None:
|
||||
return None
|
||||
if not embedding_identity(profile).matches(space.identity):
|
||||
return None
|
||||
return space
|
||||
|
||||
|
||||
def verify_readback(
|
||||
client: Any, legacy: LegacyMemoryIds, *, chat_missing_before_migration: int = 0
|
||||
) -> ReadbackReport:
|
||||
listed: set[str] = set()
|
||||
for user_id in legacy.users or ("local",):
|
||||
offset = 0
|
||||
@@ -228,5 +286,6 @@ def verify_readback(client: Any, legacy: LegacyMemoryIds) -> ReadbackReport:
|
||||
chat_readable=sum(1 for memory_id in legacy.chat if readable(memory_id)),
|
||||
productivity_expected=len(legacy.productivity),
|
||||
productivity_readable=sum(1 for memory_id in legacy.productivity if readable(memory_id)),
|
||||
chat_missing_before_migration=chat_missing_before_migration,
|
||||
)
|
||||
|
||||
@@ -0,0 +1,233 @@
|
||||
"""The read client, kept in the live embedding space.
|
||||
|
||||
Every read used to open the built-in 512-dim collection while writes followed
|
||||
the active profile, so a custom embedding read an empty space (#242). Reads
|
||||
now open the space the registry pins as live (`embedding_spaces`), which is
|
||||
the space writes fill. The embedding that reaches it is the active profile's
|
||||
when its model and width are the space's -- a new endpoint for the same model
|
||||
keeps the space -- and otherwise the profile version that filled the space
|
||||
(the built-in model for the built-in space). A switch that is still running
|
||||
never changes this: reads keep serving the old space until its commit.
|
||||
|
||||
A change of space or profile builds a new client and retires the old one; a
|
||||
retired client is closed once no task still uses it, as `SdkWriteCoordinator`
|
||||
does for writes (#218). A use is the task that asked for the client: callers
|
||||
keep their `client = await handle.client()` shape, and the use ends with that
|
||||
task (one HTTP request, one retention pass).
|
||||
|
||||
The memory mode (lite / pro / ultra) is fixed per SDK client, so a mode change
|
||||
rebuilds and retires the same way, and so do the SDK features
|
||||
(`memory_features`): handles are keyed by (space, profile, mode, features).
|
||||
Neither mode nor features names a collection; both are separate client
|
||||
parameters.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from collections.abc import Callable
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from one_memory.embedding_health import EmbeddingGuardedClient, EmbeddingHealth
|
||||
from one_memory.embedding_spaces import (
|
||||
EmbeddingSpace,
|
||||
EmbeddingSpaceStore,
|
||||
embedding_identity,
|
||||
registry_for,
|
||||
)
|
||||
from one_memory.hy_graph_admin import effective_memory_mode
|
||||
from one_memory.hy_sdk_runtime import (
|
||||
BUILTIN_RETRIEVAL_PROFILE_ID,
|
||||
SdkClientHandle,
|
||||
builtin_retrieval_snapshot,
|
||||
embedding_endpoint,
|
||||
retrieval_view,
|
||||
)
|
||||
from one_memory.hy_sdk_types import HySdkRuntimeError
|
||||
from one_memory.memory_features import FeatureOverrides
|
||||
from one_memory.model_profiles import ModelProfileSnapshot, ModelProfileStore, ProviderKind
|
||||
from one_memory.secrets import SecretStore, SecretStoreUnavailableError
|
||||
|
||||
_BUILTIN = (BUILTIN_RETRIEVAL_PROFILE_ID, 0)
|
||||
#: (space id, profile id, revision, memory mode, features)
|
||||
_HandleKey = tuple[str, str, int, str, FeatureOverrides]
|
||||
|
||||
|
||||
class ActiveProfileReadHandle(SdkClientHandle):
|
||||
"""An `SdkClientHandle` that follows the live space and its embedding."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
profiles: ModelProfileStore,
|
||||
secrets: SecretStore,
|
||||
data_dir: Path,
|
||||
memory_mode: Callable[[], str] = effective_memory_mode,
|
||||
features: Callable[[], FeatureOverrides] = lambda: (),
|
||||
spaces: EmbeddingSpaceStore | None = None,
|
||||
health: EmbeddingHealth | None = None,
|
||||
) -> None:
|
||||
super().__init__(builtin_retrieval_snapshot(), api_key=None, data_dir=data_dir)
|
||||
self.profiles = profiles
|
||||
self.secrets = secrets
|
||||
self.data_dir = data_dir
|
||||
self.memory_mode = memory_mode
|
||||
self.current_features = features
|
||||
#: None only for a profile store without a database (test doubles).
|
||||
self.spaces = spaces if spaces is not None else registry_for(profiles)
|
||||
self.health = health or EmbeddingHealth()
|
||||
self._handles: dict[_HandleKey, SdkClientHandle] = {}
|
||||
#: Whether an active profile version reads a space, per (profile, space).
|
||||
self._matches: dict[tuple[str, int, str], bool] = {}
|
||||
#: Tasks in flight per handle; a retired handle closes at zero.
|
||||
self._uses: dict[SdkClientHandle, int] = {}
|
||||
self._retired: set[SdkClientHandle] = set()
|
||||
self._closing: set[asyncio.Task[None]] = set()
|
||||
self._handle_lock = asyncio.Lock()
|
||||
|
||||
async def client(self) -> Any:
|
||||
handle = await self.current_handle()
|
||||
return self._guarded(handle, await handle.client())
|
||||
|
||||
def _guarded(self, handle: SdkClientHandle, client: Any) -> Any:
|
||||
"""A remote embedding's searches report `embedding_unavailable` (no fallback)."""
|
||||
profile = getattr(handle, "profile", None)
|
||||
if not isinstance(profile, ModelProfileSnapshot):
|
||||
return client
|
||||
endpoint = embedding_endpoint(profile)
|
||||
if endpoint is None:
|
||||
return client
|
||||
space = handle.space
|
||||
name_dims = space.name_dims if space is not None else profile.embedding_dimensions
|
||||
return EmbeddingGuardedClient(
|
||||
client,
|
||||
lambda: self.health.available(
|
||||
profile, endpoint=endpoint, name_dims=name_dims, api_key=handle.api_key
|
||||
),
|
||||
)
|
||||
|
||||
async def current_handle(self) -> SdkClientHandle:
|
||||
"""The handle of the live space, held for the asking task."""
|
||||
space, identity = await asyncio.to_thread(self._resolve)
|
||||
mode, features = self.memory_mode(), self.current_features()
|
||||
key = (space.space_id if space is not None else "", *identity, mode, features)
|
||||
async with self._handle_lock:
|
||||
handle = self._handles.get(key)
|
||||
if handle is None:
|
||||
# Only on a change of space or profile: resolving a key can
|
||||
# reach the system keychain, so it is not done on every read.
|
||||
handle = await asyncio.to_thread(self._build, space, identity, mode, features)
|
||||
self._handles[key] = handle
|
||||
self._retired.update(value for name, value in self._handles.items() if name != key)
|
||||
self._handles = {key: handle}
|
||||
unused = {value for value in self._retired if not self._uses.get(value)}
|
||||
self._retired -= unused
|
||||
self._hold(handle)
|
||||
await asyncio.gather(*(value.close() for value in unused))
|
||||
return handle
|
||||
|
||||
async def close(self) -> None:
|
||||
async with self._handle_lock:
|
||||
handles = (*self._handles.values(), *self._retired)
|
||||
self._handles, self._retired, self._uses = {}, set(), {}
|
||||
await asyncio.gather(*(handle.close() for handle in handles), *self._closing)
|
||||
|
||||
async def close_idle(self) -> None:
|
||||
"""Retire every client; the unused close now, the rest when their task ends."""
|
||||
async with self._handle_lock:
|
||||
self._retired.update(self._handles.values())
|
||||
self._handles = {}
|
||||
unused = {value for value in self._retired if not self._uses.get(value)}
|
||||
self._retired -= unused
|
||||
await asyncio.gather(*(value.close() for value in unused))
|
||||
|
||||
def _resolve(self) -> tuple[EmbeddingSpace | None, tuple[str, int]]:
|
||||
active = self.profiles.active_identity()
|
||||
if self.spaces is None:
|
||||
return None, active or _BUILTIN
|
||||
space = self.spaces.live()
|
||||
if space is None:
|
||||
# A home this process has not bootstrapped yet: register what it
|
||||
# already uses, under the name it already has.
|
||||
space = self.spaces.bootstrap(
|
||||
self.profiles.snapshot(*active) if active is not None else None
|
||||
)
|
||||
if active is not None and self._reads(active, space):
|
||||
return space, active
|
||||
if space.profile_id is not None and space.profile_revision is not None:
|
||||
identity = (space.profile_id, space.profile_revision)
|
||||
try:
|
||||
self.profiles.snapshot(*identity)
|
||||
except KeyError:
|
||||
# Its endpoint and key went with the profile; the built-in
|
||||
# model cannot read these vectors. Say so, never read wrong.
|
||||
raise HySdkRuntimeError("embedding_space_profile_missing") from None
|
||||
return space, identity
|
||||
return space, _BUILTIN
|
||||
|
||||
def _reads(self, identity: tuple[str, int], space: EmbeddingSpace) -> bool:
|
||||
key = (*identity, space.space_id)
|
||||
known = self._matches.get(key)
|
||||
if known is None:
|
||||
if space.identity is None:
|
||||
known = False
|
||||
else:
|
||||
profile = self.profiles.snapshot(*identity)
|
||||
known = embedding_identity(profile).matches(space.identity)
|
||||
self._matches[key] = known
|
||||
return known
|
||||
|
||||
def _hold(self, handle: SdkClientHandle) -> None:
|
||||
"""Count a use of `handle` until the asking task finishes."""
|
||||
task = asyncio.current_task()
|
||||
if task is None: # pragma: no cover - always called from a coroutine
|
||||
return
|
||||
self._uses[handle] = self._uses.get(handle, 0) + 1
|
||||
task.add_done_callback(lambda _task: self._release(handle))
|
||||
|
||||
def _release(self, handle: SdkClientHandle) -> None:
|
||||
remaining = self._uses.get(handle, 0) - 1
|
||||
if remaining > 0:
|
||||
self._uses[handle] = remaining
|
||||
return
|
||||
self._uses.pop(handle, None)
|
||||
if handle in self._retired:
|
||||
self._retired.discard(handle)
|
||||
closing = asyncio.get_running_loop().create_task(handle.close())
|
||||
self._closing.add(closing)
|
||||
closing.add_done_callback(self._closing.discard)
|
||||
|
||||
def _build(
|
||||
self,
|
||||
space: EmbeddingSpace | None,
|
||||
identity: tuple[str, int],
|
||||
mode: str,
|
||||
features: FeatureOverrides = (),
|
||||
) -> SdkClientHandle:
|
||||
# Without a registry (a profile store with no database) the handle
|
||||
# opens the space its profile names, exactly as before.
|
||||
named: dict[str, Any] = {"space": space} if space is not None else {}
|
||||
if identity == _BUILTIN:
|
||||
return SdkClientHandle(
|
||||
builtin_retrieval_snapshot(),
|
||||
api_key=None,
|
||||
data_dir=self.data_dir,
|
||||
mode=mode,
|
||||
features=features,
|
||||
**named,
|
||||
)
|
||||
profile: ModelProfileSnapshot = self.profiles.snapshot(*identity)
|
||||
view = retrieval_view(profile)
|
||||
secret: str | None = None
|
||||
# Only a remote OpenAI-compatible embedding sends the shared key.
|
||||
if view.embedding_provider_kind is ProviderKind.OPENAI_COMPATIBLE:
|
||||
try:
|
||||
secret = self.profiles.resolve_snapshot_secret(*identity, self.secrets)
|
||||
except SecretStoreUnavailableError:
|
||||
raise HySdkRuntimeError("secure_storage_unavailable") from None
|
||||
if secret is None:
|
||||
raise HySdkRuntimeError("model_secret_unavailable")
|
||||
return SdkClientHandle(
|
||||
view, api_key=secret, data_dir=self.data_dir, mode=mode, features=features, **named
|
||||
)
|
||||
@@ -0,0 +1,179 @@
|
||||
"""SDK search parameters a host may pass through to `client.search`.
|
||||
|
||||
`POST /api/v2/context/search`, `POST /api/v1/recall` and the v2 ask routes
|
||||
used to fix every SDK search knob, so a host that wanted a stricter relevance
|
||||
floor filtered the answer itself. These fields are the public `search`
|
||||
parameters of hy-memory 2.0.9 that are safe to hand a caller. Every field is
|
||||
optional: absent, the route keeps exactly the call it made before.
|
||||
|
||||
Never exposed: `user_ids` (a route searches its one scoped `user_id`),
|
||||
`embed_override` and `query_embedding` (they bypass the configured embedder),
|
||||
and `request_id` (Core's own correlation id).
|
||||
|
||||
A caller's floor is also applied to what Core returns, not only handed to the
|
||||
SDK: a hit below the floor that names its lane is dropped (`passes_floor`).
|
||||
With `min_score` alone that floor covers every lane, so `min_score=0.9`
|
||||
answers only results scoring at least 0.9.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping
|
||||
from typing import Annotated, Literal
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, StrictBool, model_validator
|
||||
|
||||
#: The SDK's `MemoryLayer` values (hy-memory 2.0.9).
|
||||
MemoryLayerName = Literal[
|
||||
"l0_basic_info",
|
||||
"l1_raw",
|
||||
"l2_fact",
|
||||
"l3_summary",
|
||||
"l4_identity",
|
||||
"l5_knowledge",
|
||||
"l6_schema",
|
||||
"l7_intention",
|
||||
]
|
||||
MEMORY_LAYERS: tuple[str, ...] = (
|
||||
"l0_basic_info",
|
||||
"l1_raw",
|
||||
"l2_fact",
|
||||
"l3_summary",
|
||||
"l4_identity",
|
||||
"l5_knowledge",
|
||||
"l6_schema",
|
||||
"l7_intention",
|
||||
)
|
||||
#: The SDK's `MemoryStatus` values.
|
||||
MEMORY_STATUSES: tuple[str, ...] = (
|
||||
"active",
|
||||
"superseded",
|
||||
"negated",
|
||||
"conflicted",
|
||||
"shadow",
|
||||
"archived",
|
||||
)
|
||||
SearchScene = Literal["chat", "working", "coding", "all"]
|
||||
#: Which recall lane produced a hit, for the floor that applies to it.
|
||||
SearchLane = Literal["normal", "schema", "proactive", "working", "coding", "raw"]
|
||||
|
||||
Score = Annotated[float, Field(ge=0.0, le=1.0)]
|
||||
SourceName = Annotated[str, Field(min_length=1, max_length=128)]
|
||||
ScopeId = Annotated[str, Field(min_length=1, max_length=512)]
|
||||
|
||||
|
||||
class SdkSearchFilters(BaseModel):
|
||||
"""Filters and floors every search route accepts (recall, context, ask)."""
|
||||
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
#: Chat lane floor; with nothing more specific, every lane's.
|
||||
min_score: Score | None = None
|
||||
#: Schema (L6) lane floor; defaults to `min_score` when only that is set.
|
||||
schema_min_score: Score | None = None
|
||||
#: False skips the SDK's reranking stages for this call.
|
||||
rerank: StrictBool | None = None
|
||||
include_sources: tuple[SourceName, ...] | None = Field(default=None, max_length=64)
|
||||
exclude_sources: tuple[SourceName, ...] | None = Field(default=None, max_length=64)
|
||||
#: Unix seconds: only memories created at or after it.
|
||||
created_after: float | None = Field(default=None, ge=0)
|
||||
session_ids: tuple[ScopeId, ...] | None = Field(default=None, max_length=64)
|
||||
agent_ids: tuple[ScopeId, ...] | None = Field(default=None, max_length=64)
|
||||
#: Only these layers (pro/ultra; lite always searches raw).
|
||||
layers: tuple[MemoryLayerName, ...] | None = Field(default=None, min_length=1, max_length=8)
|
||||
|
||||
def filter_kwargs(self) -> dict[str, object]:
|
||||
"""The SDK keyword arguments this request sets, and only those."""
|
||||
values: dict[str, object] = {}
|
||||
for name in (
|
||||
"min_score",
|
||||
"schema_min_score",
|
||||
"rerank",
|
||||
"created_after",
|
||||
):
|
||||
value = getattr(self, name)
|
||||
if value is not None:
|
||||
values[name] = value
|
||||
for name in ("include_sources", "exclude_sources", "session_ids", "agent_ids", "layers"):
|
||||
value = getattr(self, name)
|
||||
if value is not None:
|
||||
values[name] = list(dict.fromkeys(value))
|
||||
if self.min_score is not None and self.schema_min_score is None:
|
||||
values["schema_min_score"] = self.min_score
|
||||
return values
|
||||
|
||||
def lane_floor(self, lane: SearchLane) -> float | None:
|
||||
"""The caller's floor for one lane, or None when the caller set none."""
|
||||
if lane == "schema":
|
||||
return self.schema_min_score if self.schema_min_score is not None else self.min_score
|
||||
return self.min_score
|
||||
|
||||
def passes_floor(self, lane: SearchLane, score: float) -> bool:
|
||||
floor = self.lane_floor(lane)
|
||||
return floor is None or score >= floor
|
||||
|
||||
def search_overrides(self) -> Mapping[str, object]:
|
||||
return self.filter_kwargs()
|
||||
|
||||
|
||||
class SdkSearchOptions(SdkSearchFilters):
|
||||
"""Every SDK search knob a recall or context caller may set."""
|
||||
|
||||
#: working/coding fact-layer floor (SDK default 0.30).
|
||||
working_fact_min_score: Score | None = None
|
||||
#: working/coding experience-layer floor (SDK default fact floor + 0.10).
|
||||
working_experience_min_score: Score | None = None
|
||||
reader: Literal["legacy", "hybrid"] | None = None
|
||||
#: chat (the default), working, coding, or all three fused.
|
||||
scene: SearchScene | None = None
|
||||
#: agentic spends one model call to classify and rewrite the query.
|
||||
search_mode: Literal["default", "agentic"] | None = None
|
||||
#: working/coding lanes only; `branch` needs `workspace_id`.
|
||||
workspace_id: ScopeId | None = None
|
||||
branch: ScopeId | None = None
|
||||
|
||||
@model_validator(mode="after")
|
||||
def branch_needs_workspace(self) -> SdkSearchOptions:
|
||||
if self.branch is not None and self.workspace_id is None:
|
||||
raise ValueError("branch requires workspace_id")
|
||||
return self
|
||||
|
||||
def search_overrides(self) -> Mapping[str, object]:
|
||||
values = self.filter_kwargs()
|
||||
for name in (
|
||||
"working_fact_min_score",
|
||||
"working_experience_min_score",
|
||||
"reader",
|
||||
"scene",
|
||||
"search_mode",
|
||||
"workspace_id",
|
||||
"branch",
|
||||
):
|
||||
value = getattr(self, name)
|
||||
if value is not None:
|
||||
values[name] = value
|
||||
if self.min_score is not None and self.working_fact_min_score is None:
|
||||
values["working_fact_min_score"] = self.min_score
|
||||
if self.min_score is not None and self.working_experience_min_score is None:
|
||||
values["working_experience_min_score"] = self.min_score
|
||||
return values
|
||||
|
||||
def lane_floor(self, lane: SearchLane) -> float | None:
|
||||
if lane in {"working", "coding"}:
|
||||
explicit = [
|
||||
value
|
||||
for value in (self.working_fact_min_score, self.working_experience_min_score)
|
||||
if value is not None
|
||||
]
|
||||
# Facts and experiences share one score field on a fused hit, so
|
||||
# the looser of the two floors is the one Core can enforce.
|
||||
return min(explicit) if explicit else self.min_score
|
||||
return super().lane_floor(lane)
|
||||
|
||||
|
||||
def split_csv(values: list[str] | None) -> list[str]:
|
||||
"""Query values given repeated (`a&a`) or comma separated (`a,b`)."""
|
||||
found: list[str] = []
|
||||
for value in values or ():
|
||||
found.extend(part.strip() for part in value.split(",") if part.strip())
|
||||
return list(dict.fromkeys(found))
|
||||
@@ -9,6 +9,7 @@ one. Runtime code only reaches vectors through the SDK. This module runs from
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
import sqlite3
|
||||
from collections.abc import Iterator
|
||||
from contextlib import closing
|
||||
@@ -21,6 +22,8 @@ from one_memory.model_profiles import ModelProfileSnapshot
|
||||
from one_memory.sdk_data_plane_migration import DataPlaneMigrationError
|
||||
|
||||
_BATCH = 500
|
||||
#: Every name a MemHub build gave an embedding space's base collection.
|
||||
_SPACE_NAME = re.compile(r"^(?:memhub_[0-9a-f]{16}|one_memory_[0-9a-f]{12})(?:_|$)")
|
||||
|
||||
|
||||
def open_offline_chroma(directory: Path) -> Any:
|
||||
@@ -51,18 +54,49 @@ def _vector_directory(data_dir: Path) -> Path:
|
||||
return data_dir / "hy-memory" / "governed-v2" / "vector"
|
||||
|
||||
|
||||
def _stored_names(data_dir: Path) -> tuple[str, ...]:
|
||||
def _stored_names(data_dir: Path, *, filled: bool = False) -> tuple[str, ...]:
|
||||
"""Collection names; with `filled`, only those holding at least one row."""
|
||||
database = _vector_directory(data_dir) / "chroma.sqlite3"
|
||||
if not database.is_file():
|
||||
return ()
|
||||
with closing(sqlite3.connect(f"file:{database}?mode=ro", uri=True)) as connection:
|
||||
try:
|
||||
rows = connection.execute("SELECT name FROM collections").fetchall()
|
||||
rows = connection.execute(
|
||||
"SELECT DISTINCT collections.name FROM collections "
|
||||
"JOIN segments ON segments.collection = collections.id "
|
||||
"JOIN embeddings ON embeddings.segment_id = segments.id"
|
||||
if filled
|
||||
else "SELECT name FROM collections"
|
||||
).fetchall()
|
||||
except sqlite3.Error:
|
||||
return ()
|
||||
return tuple(sorted(str(row[0]) for row in rows))
|
||||
|
||||
|
||||
def filled_collection_names(data_dir: Path) -> tuple[str, ...]:
|
||||
"""Every vector collection holding at least one record. Read-only.
|
||||
|
||||
Safe beside a serving Core: it reads Chroma's own SQLite file, never the
|
||||
client system the SDK shares.
|
||||
"""
|
||||
return _stored_names(data_dir, filled=True)
|
||||
|
||||
|
||||
def stored_collection_width(data_dir: Path, name: str) -> int | None:
|
||||
"""The vector width Chroma fixed for `name` on its first add; None if unset."""
|
||||
database = _vector_directory(data_dir) / "chroma.sqlite3"
|
||||
if not database.is_file():
|
||||
return None
|
||||
with closing(sqlite3.connect(f"file:{database}?mode=ro", uri=True)) as connection:
|
||||
try:
|
||||
row = connection.execute(
|
||||
"SELECT dimension FROM collections WHERE name=?", (name,)
|
||||
).fetchone()
|
||||
except sqlite3.Error:
|
||||
return None
|
||||
return int(row[0]) if row is not None and row[0] is not None else None
|
||||
|
||||
|
||||
def legacy_vector_collections(
|
||||
data_dir: Path, profile: ModelProfileSnapshot
|
||||
) -> tuple[tuple[str, str], ...]:
|
||||
@@ -80,6 +114,59 @@ def legacy_vector_collections(
|
||||
return tuple(pairs)
|
||||
|
||||
|
||||
def vectors_belong_to_another_space(data_dir: Path, profile: ModelProfileSnapshot) -> bool:
|
||||
"""Vectors exist, but only in embedding spaces `profile` does not name.
|
||||
|
||||
Counted rows, not names: every Runtime opens its read collection with
|
||||
`get_or_create`, so an empty collection under the expected base proves
|
||||
nothing. Read-only, so a preflight can ask before anything changes.
|
||||
"""
|
||||
current, legacy_bases = vector_collection_bases(profile)
|
||||
known = (current, *legacy_bases)
|
||||
filled = _stored_names(data_dir, filled=True)
|
||||
return not any(
|
||||
name == base or name.startswith(f"{base}_") for name in filled for base in known
|
||||
) and any(_SPACE_NAME.match(name) for name in filled)
|
||||
|
||||
|
||||
def stored_memory_ids(data_dir: Path) -> frozenset[str]:
|
||||
"""Every id in the vector store as found, across all embedding spaces.
|
||||
|
||||
All of them on purpose. A memory stranded under a space the Runtime will
|
||||
not read must count as present, so the readback then fails on it instead
|
||||
of excusing it. Both the point id and the SDK's `node_id` are kept: the
|
||||
SDK derives a point id from a node id that is not a UUID.
|
||||
|
||||
Callers must have stopped every SDK writer. A store that cannot be read
|
||||
raises: "unknown" must never be taken for "absent".
|
||||
"""
|
||||
if not (_vector_directory(data_dir) / "chroma.sqlite3").is_file():
|
||||
return frozenset()
|
||||
found: set[str] = set()
|
||||
try:
|
||||
client = open_offline_chroma(_vector_directory(data_dir))
|
||||
try:
|
||||
for collection in client.list_collections():
|
||||
offset = 0
|
||||
while True:
|
||||
page = collection.get(include=["metadatas"], limit=_BATCH, offset=offset)
|
||||
ids = list(page["ids"])
|
||||
if not ids:
|
||||
break
|
||||
found.update(ids)
|
||||
found.update(
|
||||
str(metadata["node_id"])
|
||||
for metadata in page["metadatas"] or ()
|
||||
if metadata and metadata.get("node_id")
|
||||
)
|
||||
offset += len(ids)
|
||||
finally:
|
||||
client.close()
|
||||
except Exception as error:
|
||||
raise DataPlaneMigrationError("sdk_migration_vector_store_unreadable") from error
|
||||
return frozenset(found)
|
||||
|
||||
|
||||
def _records(collection: Any) -> Iterator[dict[str, Any]]:
|
||||
offset = 0
|
||||
while True:
|
||||
|
||||
@@ -9,9 +9,17 @@ from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from typing import Any
|
||||
|
||||
from one_memory.hy_sdk_runtime import SdkClientHandle
|
||||
from one_memory.embedding_spaces import (
|
||||
EmbeddingSpace,
|
||||
EmbeddingSpaceStore,
|
||||
embedding_identity,
|
||||
registry_for,
|
||||
)
|
||||
from one_memory.hy_graph_admin import effective_memory_mode
|
||||
from one_memory.hy_sdk_runtime import SdkClientHandle, builtin_retrieval_snapshot
|
||||
from one_memory.imports.models import EvidenceContentKind, EvidenceOrigin, EvidenceRecord
|
||||
from one_memory.input_sanitization import sanitize_model_input
|
||||
from one_memory.memory_features import FeatureOverrides
|
||||
from one_memory.model_profiles import ModelProfileSnapshot, ModelProfileStore
|
||||
from one_memory.secrets import SecretStore, SecretStoreUnavailableError
|
||||
|
||||
@@ -31,8 +39,11 @@ class SdkWriteReceipt:
|
||||
class SdkWriteProcessor:
|
||||
"""Turn one durable job into exactly one public SDK write call."""
|
||||
|
||||
def __init__(self, client_handle: SdkClientHandle) -> None:
|
||||
def __init__(self, client_handle: SdkClientHandle, *, lite: bool = False) -> None:
|
||||
self._client_handle = client_handle
|
||||
#: lite makes no model call: `add_extracted` would still call one, so a
|
||||
#: direct memory note is stored as one raw user turn through `add`.
|
||||
self._lite = lite
|
||||
|
||||
async def process_sdk_job(
|
||||
self,
|
||||
@@ -64,7 +75,7 @@ class SdkWriteProcessor:
|
||||
# waits and retries instead of being left `unknown`.
|
||||
raise RuntimeError("hy_memory_client_unavailable") from error
|
||||
|
||||
if extracted:
|
||||
if extracted and not self._lite:
|
||||
contents = [sanitize_model_input(record.verbatim) for record in records]
|
||||
await before_hy_invocation()
|
||||
result = await asyncio.to_thread(
|
||||
@@ -81,12 +92,15 @@ class SdkWriteProcessor:
|
||||
else:
|
||||
messages = [
|
||||
{
|
||||
"role": record.role,
|
||||
"role": "user" if extracted else record.role,
|
||||
"content": sanitize_model_input(record.verbatim),
|
||||
}
|
||||
for record in records
|
||||
if record.content_kind is EvidenceContentKind.RAW_CONVERSATION
|
||||
and record.role in _MESSAGE_ROLES
|
||||
if extracted
|
||||
or (
|
||||
record.content_kind is EvidenceContentKind.RAW_CONVERSATION
|
||||
and record.role in _MESSAGE_ROLES
|
||||
)
|
||||
]
|
||||
if not messages:
|
||||
raise RuntimeError("hy_memory_task_without_inputs")
|
||||
@@ -138,6 +152,22 @@ class SdkWriteCoordinator:
|
||||
from under a write still running on it: the SDK cleared its stores and
|
||||
the SharedRuntime closed mid-`add`, the job failed with AttributeError
|
||||
after HY was invoked, and it was left `unknown`.
|
||||
|
||||
Handles are keyed by profile, memory mode and the SDK features
|
||||
(`memory_features`), so a mode or feature change retires handles the same
|
||||
way. In lite no generation model is needed: without an
|
||||
active profile the built-in local embedder writes raw turns.
|
||||
|
||||
`exclusive_for_evolve(user_id)` keeps evolve and that user's writes
|
||||
apart: it waits for the user's write in flight, and while it runs that
|
||||
user's jobs are refused before the SDK is called
|
||||
(`memory_evolve_in_progress`), so the drain puts them back.
|
||||
|
||||
Writes go to the live embedding space only. An active profile whose
|
||||
embedding (model and width) is not the live space's is refused with
|
||||
`embedding_space_mismatch`: it used to open a new, empty collection.
|
||||
While an embedding switch runs, `writes_open` is false and no write
|
||||
starts (`embedding_switch_in_progress`); the job waits in the queue.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
@@ -148,13 +178,31 @@ class SdkWriteCoordinator:
|
||||
data_dir: Any,
|
||||
transport_factory: Callable[[ModelProfileSnapshot], Callable[[], Any] | None],
|
||||
sedimentation_enabled: Callable[[], bool],
|
||||
memory_mode: Callable[[], str] = effective_memory_mode,
|
||||
memory_features: Callable[[], FeatureOverrides] = lambda: (),
|
||||
spaces: EmbeddingSpaceStore | None = None,
|
||||
writes_open: Callable[[], bool] = lambda: True,
|
||||
) -> None:
|
||||
self.profiles = profiles
|
||||
self.secrets = secrets
|
||||
self.data_dir = data_dir
|
||||
self.transport_factory = transport_factory
|
||||
self.sedimentation_enabled = sedimentation_enabled
|
||||
self._handles: dict[tuple[str, int], SdkClientHandle] = {}
|
||||
self.memory_mode = memory_mode
|
||||
self.memory_features = memory_features
|
||||
#: Keyed by (profile id, revision, space id, memory mode, features).
|
||||
self._handles: dict[tuple[str, int, str, str, FeatureOverrides], SdkClientHandle] = {}
|
||||
#: Users evolve (or the graph backfill) is running for, the writes in
|
||||
#: flight per user, and one exclusive run per user at a time. Writes
|
||||
#: of one user still run side by side; only evolve waits for them.
|
||||
self._evolving: set[str] = set()
|
||||
self._writing: dict[str, int] = {}
|
||||
self._writes_idle = asyncio.Condition()
|
||||
self._exclusive: dict[str, asyncio.Lock] = {}
|
||||
#: None only for a profile store without a database (test doubles):
|
||||
#: then every client opens the space its profile names, as before.
|
||||
self.spaces = spaces if spaces is not None else registry_for(profiles)
|
||||
self.writes_open = writes_open
|
||||
#: Uses in flight per handle; a retired handle closes at zero.
|
||||
self._uses: dict[SdkClientHandle, int] = {}
|
||||
self._retired: set[SdkClientHandle] = set()
|
||||
@@ -169,16 +217,76 @@ class SdkWriteCoordinator:
|
||||
) -> SdkWriteReceipt:
|
||||
if not self.sedimentation_enabled():
|
||||
raise RuntimeError("model_sedimentation_disabled")
|
||||
async with self._active_handle() as handle:
|
||||
return await SdkWriteProcessor(handle).process_sdk_job(
|
||||
records,
|
||||
request_id=request_id,
|
||||
before_hy_invocation=before_hy_invocation,
|
||||
)
|
||||
if not self.writes_open():
|
||||
# An embedding switch paused writes: the drain puts the job back.
|
||||
raise RuntimeError("embedding_switch_in_progress")
|
||||
user_id = records[0].user_id if records else "local"
|
||||
if user_id in self._evolving:
|
||||
# Before any SDK call: the drain returns the job to pending.
|
||||
raise RuntimeError("memory_evolve_in_progress")
|
||||
self._writing[user_id] = self._writing.get(user_id, 0) + 1
|
||||
try:
|
||||
async with self._active_handle() as handle:
|
||||
return await SdkWriteProcessor(
|
||||
handle, lite=getattr(handle, "mode", None) == "lite"
|
||||
).process_sdk_job(
|
||||
records,
|
||||
request_id=request_id,
|
||||
before_hy_invocation=before_hy_invocation,
|
||||
)
|
||||
finally:
|
||||
remaining = self._writing.get(user_id, 1) - 1
|
||||
if remaining > 0:
|
||||
self._writing[user_id] = remaining
|
||||
else:
|
||||
self._writing.pop(user_id, None)
|
||||
async with self._writes_idle:
|
||||
self._writes_idle.notify_all()
|
||||
|
||||
async def inquire_write_outcome(
|
||||
self, request_id: str
|
||||
) -> SdkWriteReceipt | None:
|
||||
async def retire_stale_handles(self) -> None:
|
||||
"""After a mode or feature change: retire handles built otherwise.
|
||||
|
||||
The idle ones close now. A handle still writing is closed when its
|
||||
last use ends (#218), so a write that started in the old mode, or
|
||||
with the old features, finishes in it.
|
||||
"""
|
||||
build = (self.memory_mode(), self.memory_features())
|
||||
async with self._lock:
|
||||
stale = {key: handle for key, handle in self._handles.items() if key[3:] != build}
|
||||
for key in stale:
|
||||
del self._handles[key]
|
||||
self._retired.update(stale.values())
|
||||
unused = {value for value in self._retired if not self._uses.get(value)}
|
||||
self._retired -= unused
|
||||
await asyncio.gather(*(value.close() for value in unused))
|
||||
|
||||
#: The name #253 introduced; a mode change is one kind of rebuild.
|
||||
retire_stale_modes = retire_stale_handles
|
||||
|
||||
def evolving(self, user_id: str) -> bool:
|
||||
"""Whether evolve holds this user, so their queued writes should wait."""
|
||||
return user_id in self._evolving
|
||||
|
||||
@asynccontextmanager
|
||||
async def exclusive_for_evolve(self, user_id: str) -> AsyncIterator[None]:
|
||||
"""Hold one user's writes off while evolve runs for them.
|
||||
|
||||
New writes of the user are refused from here on; the ones already in
|
||||
flight are waited for before the body runs. Evolve and the graph
|
||||
backfill write too, so an embedding switch refuses them as well.
|
||||
"""
|
||||
if not self.writes_open():
|
||||
raise RuntimeError("embedding_switch_in_progress")
|
||||
async with self._exclusive.setdefault(user_id, asyncio.Lock()):
|
||||
self._evolving.add(user_id)
|
||||
try:
|
||||
async with self._writes_idle:
|
||||
await self._writes_idle.wait_for(lambda: not self._writing.get(user_id))
|
||||
yield
|
||||
finally:
|
||||
self._evolving.discard(user_id)
|
||||
|
||||
async def inquire_write_outcome(self, request_id: str) -> SdkWriteReceipt | None:
|
||||
async with self._active_handle() as handle:
|
||||
return await SdkWriteProcessor(handle).inquire_write_outcome(request_id)
|
||||
|
||||
@@ -188,6 +296,19 @@ class SdkWriteCoordinator:
|
||||
self._handles, self._retired, self._uses = {}, set(), {}
|
||||
await asyncio.gather(*(handle.close() for handle in handles))
|
||||
|
||||
def in_use(self) -> int:
|
||||
"""Writes and reads holding a client right now."""
|
||||
return sum(self._uses.values())
|
||||
|
||||
async def close_idle(self) -> None:
|
||||
"""Retire every handle; the unused close now, the rest when released."""
|
||||
async with self._lock:
|
||||
self._retired.update(self._handles.values())
|
||||
self._handles = {}
|
||||
unused = {value for value in self._retired if not self._uses.get(value)}
|
||||
self._retired -= unused
|
||||
await asyncio.gather(*(value.close() for value in unused))
|
||||
|
||||
@asynccontextmanager
|
||||
async def use_client(self) -> AsyncIterator[Any]:
|
||||
"""The public SDK client of the active write profile, kept open while used."""
|
||||
@@ -198,27 +319,38 @@ class SdkWriteCoordinator:
|
||||
self,
|
||||
) -> tuple[ModelProfileSnapshot, str | None]:
|
||||
"""Resolve one stable profile snapshot for graph administration."""
|
||||
_identity, profile, secret = await asyncio.to_thread(self._read_active_credentials)
|
||||
_identity, profile, secret, _space = await asyncio.to_thread(self._read_active_credentials)
|
||||
return profile, secret
|
||||
|
||||
async def live_space(self) -> EmbeddingSpace | None:
|
||||
if self.spaces is None:
|
||||
return None
|
||||
return await asyncio.to_thread(self._live_space)
|
||||
|
||||
@asynccontextmanager
|
||||
async def _active_handle(self) -> AsyncIterator[SdkClientHandle]:
|
||||
"""Use the active profile's handle; retire the others, closing the unused."""
|
||||
identity, profile, secret = await asyncio.to_thread(
|
||||
self._read_active_credentials
|
||||
mode = self.memory_mode()
|
||||
features = self.memory_features()
|
||||
profile_identity, profile, secret, space = await asyncio.to_thread(
|
||||
self._read_active_credentials, mode
|
||||
)
|
||||
key = (*profile_identity, space.space_id if space is not None else "", mode, features)
|
||||
async with self._lock:
|
||||
handle = self._handles.get(identity)
|
||||
handle = self._handles.get(key)
|
||||
if handle is None:
|
||||
handle = SdkClientHandle(
|
||||
profile,
|
||||
api_key=secret,
|
||||
data_dir=self.data_dir,
|
||||
http_transport_factory=self.transport_factory(profile),
|
||||
mode=mode,
|
||||
features=features,
|
||||
space=space,
|
||||
)
|
||||
self._handles[identity] = handle
|
||||
self._retired.update(value for key, value in self._handles.items() if key != identity)
|
||||
self._handles = {identity: handle}
|
||||
self._handles[key] = handle
|
||||
self._retired.update(value for name, value in self._handles.items() if name != key)
|
||||
self._handles = {key: handle}
|
||||
self._uses[handle] = self._uses.get(handle, 0) + 1
|
||||
unused = {value for value in self._retired if not self._uses.get(value)}
|
||||
self._retired -= unused
|
||||
@@ -237,13 +369,51 @@ class SdkWriteCoordinator:
|
||||
self._retired.discard(handle)
|
||||
await handle.close()
|
||||
|
||||
def _live_space(self) -> EmbeddingSpace:
|
||||
assert self.spaces is not None
|
||||
space = self.spaces.live()
|
||||
if space is None:
|
||||
active = self.profiles.active_identity()
|
||||
space = self.spaces.bootstrap(
|
||||
self.profiles.snapshot(*active) if active is not None else None
|
||||
)
|
||||
return space
|
||||
|
||||
def _lite_credentials(
|
||||
self,
|
||||
) -> tuple[tuple[str, int], ModelProfileSnapshot, str | None, EmbeddingSpace | None]:
|
||||
space = self._live_space() if self.spaces is not None else None
|
||||
if space is None or space.profile_id is None or space.profile_revision is None:
|
||||
builtin = builtin_retrieval_snapshot()
|
||||
return (builtin.profile_id, builtin.revision), builtin, None, space
|
||||
identity = (space.profile_id, space.profile_revision)
|
||||
profile = self.profiles.snapshot(*identity)
|
||||
try:
|
||||
secret = self.profiles.resolve_snapshot_secret(*identity, self.secrets)
|
||||
except SecretStoreUnavailableError:
|
||||
raise RuntimeError("secure_storage_unavailable") from None
|
||||
if profile.has_api_key and secret is None:
|
||||
raise RuntimeError("model_secret_unavailable")
|
||||
return identity, profile, secret, space
|
||||
|
||||
def _read_active_credentials(
|
||||
self,
|
||||
) -> tuple[tuple[str, int], ModelProfileSnapshot, str | None]:
|
||||
mode: str = "pro",
|
||||
) -> tuple[tuple[str, int], ModelProfileSnapshot, str | None, EmbeddingSpace | None]:
|
||||
identity = self.profiles.active_identity()
|
||||
if identity is None:
|
||||
if mode == "lite":
|
||||
# lite stores raw turns with one embedding and calls no model:
|
||||
# the live space's own, so they land where reads look.
|
||||
return self._lite_credentials()
|
||||
raise RuntimeError("model_profile_not_active")
|
||||
profile = self.profiles.snapshot(*identity)
|
||||
space = self._live_space() if self.spaces is not None else None
|
||||
if space is not None and (
|
||||
space.identity is None or not embedding_identity(profile).matches(space.identity)
|
||||
):
|
||||
# Never a new, empty collection: the switch moves the space.
|
||||
raise RuntimeError("embedding_space_mismatch")
|
||||
try:
|
||||
secret = self.profiles.resolve_snapshot_secret(
|
||||
*identity,
|
||||
@@ -255,7 +425,7 @@ class SdkWriteCoordinator:
|
||||
raise RuntimeError("model_profile_changed")
|
||||
if profile.has_api_key and secret is None:
|
||||
raise RuntimeError("model_secret_unavailable")
|
||||
return identity, profile, secret
|
||||
return identity, profile, secret, space
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
|
||||
@@ -27,6 +27,11 @@ class SemanticMemoryHit:
|
||||
member_ids: tuple[str, ...] = ()
|
||||
sdk_formatted: bool = False
|
||||
formatting_complete: bool = False
|
||||
#: The recall lane that produced the hit (`sdk_search_options.SearchLane`)
|
||||
#: and the namespace to read it back from. Defaults are a chat memory.
|
||||
lane: str = "normal"
|
||||
kind: str = "chat"
|
||||
variant: str | None = None
|
||||
|
||||
def as_dict(self) -> dict[str, object]:
|
||||
return {
|
||||
@@ -140,6 +145,10 @@ async def search_within(
|
||||
return (), "failed"
|
||||
error = pending.exception()
|
||||
if error is not None:
|
||||
if getattr(error, "code", None) == "embedding_unavailable":
|
||||
# Not "HY found nothing": the query could not be embedded, and
|
||||
# the caller must say so instead of answering empty.
|
||||
raise error
|
||||
return (), "failed"
|
||||
found = pending.result()
|
||||
normalized = normalize_semantic_hits(found)
|
||||
|
||||
@@ -88,6 +88,9 @@
|
||||
| `test_diagnostic_trace_api.py` | `integration` | 覆盖 trace lookup is disabled until detailed logging is enabled、trace lookup returns only safe structured steps、log level route applies current process level 等 4 个场景。 |
|
||||
| `test_disable_thinking.py` | `contract` | 运行期生成请求须按档案可选关闭思考;默认不向厂商下发未知字段。 |
|
||||
| `test_dual_engine_helper.py` | `contract` | Tests for the dual-engine harness itself. |
|
||||
| `test_embedding_switch.py` | `contract` | Switching the embedding model, and merging stray spaces, without losing a memory. |
|
||||
| `test_embedding_switch_api.py` | `integration` | The embedding switch through Core's HTTP surface and lifespan. |
|
||||
| `test_embedding_switch_sdk_adapter.py` | `contract` | The re-embed adapter against the SDK's real re-embed module and Chroma. |
|
||||
| `test_event_loop_blocking.py` | `contract` | Proof, as a test, that blocking work on the loop freezes it — and the fix works. |
|
||||
| `test_evidence_store.py` | `contract` | 覆盖 same source message merges origins without duplication、identical text from different sources is not deduplicated、rescanning an unchanged file does not write the message again 等 26 个场景。 |
|
||||
| `test_export_probe.py` | `contract` | What a probe may say about a stranger's chat history, and what it may not. |
|
||||
@@ -159,8 +162,10 @@
|
||||
| `test_memory_citation_query_plan.py` | `contract` | Correlated citation checks must look links up by memory, not scan by 'active'. |
|
||||
| `test_memory_content_import.py` | `integration` | 覆盖 existing memory preview keeps unknown author and exact original、invalid existing memory records never leave a partial upload、memory content requires explicit mode and rejects vendor guessing 等 5 个场景。 |
|
||||
| `test_memory_evidence_freshness.py` | `contract` | 新鲜度拆除后:读路径恒为 current;证据修订 / 畸形 proof / segment ref 不再藏记忆或 503。 |
|
||||
| `test_memory_features.py` | `contract` | The SDK retrieval and extraction features MemHub lets a person turn on. |
|
||||
| `test_memory_freshness_integration_review.py` | `integration` | Independent synthetic read-path review; no lifespan, model or user database. |
|
||||
| `test_memory_library_all_scenes.py` | `integration` | The v2 memory library across chat, working and coding memories. |
|
||||
| `test_memory_mode_and_evolve.py` | `contract` | The memory mode (lite / pro / ultra), evolve and its nightly schedule. |
|
||||
| `test_memory_qa.py` | `contract` | 覆盖 search terms keep meaningful single chinese character、streamed answer emits recall then incremental grounded text、governance change during generation stops the old answer 等 14 个场景。 |
|
||||
| `test_memory_qa_api.py` | `integration` | 覆盖 ask stream endpoint emits recall deltas and final、ask stream endpoint requires the local token 场景。 |
|
||||
| `test_memory_qa_permissions.py` | `contract` | 覆盖 initial refusal precedes all secret access and network、invalid authorization result fails before secret access、authorization receives exact question and only outbound grounding 等 16 个场景。 |
|
||||
@@ -248,6 +253,7 @@
|
||||
| `test_schema_gate.py` | `contract` | 覆盖 fresh database records the build schema、future schema refuses to open、legacy database without a ledger is left untouched 场景。 |
|
||||
| `test_schema_migrations.py` | `contract` | 覆盖 evidence store adds source ordinal to previous schema、import store migrates previous job and error schema、old active job index lets a sync sit beside paused history 等 4 个场景。 |
|
||||
| `test_schema_ownership_gate.py` | `contract` | A1-03/A1-04: schema ownership, and the hard part of telling DDL from reads. |
|
||||
| `test_sdk_api_surface.py` | `contract` | The SDK surface MemHub exposes: search pass-through, layers, raw, v2 aliases. |
|
||||
| `test_sdk_client_construction.py` | `contract` | SDK clients of one home are constructed one at a time. |
|
||||
| `test_sdk_data_plane_migration.py` | `contract` | SDK-authority migration ledger, aliases and atomic cutover marker. |
|
||||
| `test_sdk_memory_qa.py` | `contract` | 覆盖 sdk ask propagates sdk search failure without legacy fallback、sdk ask revalidates memory before returning extract、sdk ask uses native sdk memory without governance fields 场景。 |
|
||||
@@ -311,4 +317,4 @@
|
||||
| `test_workspace_recall.py` | `contract` | Workspace memories were unreachable from every host; these pin the way in. |
|
||||
| `test_write_queue_lease_renewal.py` | `contract` | Rule 8: a lease must outlive the work it covers, and be renewed while it runs. |
|
||||
|
||||
共 306 个测试文件。
|
||||
共 312 个测试文件。
|
||||
@@ -14,8 +14,8 @@ from pathlib import Path
|
||||
import pytest
|
||||
|
||||
PRIMARY_MARKERS = ("unit", "contract", "integration", "performance", "live")
|
||||
EXPECTED_TEST_FILE_COUNT = 306
|
||||
EXPECTED_TEST_FILES_SHA256 = "e616dcf9fea288bf15a4445045f7b4278517f3304091c04fffd55be80308d481"
|
||||
EXPECTED_TEST_FILE_COUNT = 312
|
||||
EXPECTED_TEST_FILES_SHA256 = "7d0dc6efd88d71d5275bea8016a0ff087ff7568abcb8fd01e5f82c77a08a484f"
|
||||
|
||||
# Pure transformations, validation models, parity checks and state machines.
|
||||
# This list is intentionally small: uncertain files default to contract.
|
||||
@@ -54,6 +54,7 @@ INTEGRATION_FILES = frozenset(
|
||||
"test_backup_endpoint.py",
|
||||
"test_recall_raw_fallback.py",
|
||||
"test_core_ready_before_recovery.py",
|
||||
"test_embedding_switch_api.py",
|
||||
"test_backup_restore_api.py",
|
||||
"test_c1_import_via_api_path.py",
|
||||
"test_capture_deadline.py",
|
||||
|
||||
@@ -0,0 +1,301 @@
|
||||
"""An in-memory `ReembedBackend` for embedding switch tests.
|
||||
|
||||
It keeps collections as plain dicts named exactly as the SDK names them, so a
|
||||
test can lay out a home -- a live space, a stray 512 space left by an older
|
||||
build, a retired space -- and check what a switch or merge did to each. The
|
||||
vectors themselves are not computed: a record remembers the model and width
|
||||
it was "embedded" with, which is what the assertions need.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import re
|
||||
import shutil
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
|
||||
from one_memory.embedding_backend import (
|
||||
EmbeddingTarget,
|
||||
EmbeddingWidthError,
|
||||
GraphCheck,
|
||||
PhysicalSpace,
|
||||
ReembedCancelledError,
|
||||
ReembedStep,
|
||||
ScannedSpace,
|
||||
VectorCheck,
|
||||
)
|
||||
from one_memory.embedding_spaces import VECTOR_ROLES, space_collection_names
|
||||
|
||||
_ROLE_SUFFIX = re.compile(
|
||||
r"^(?P<base>.+?)_(?P<role>coding_keys_cv|coding_keys|working_facts_cv|working_facts)"
|
||||
r"(?:_(?P<dims>\d+))?$"
|
||||
)
|
||||
_DIMS_SUFFIX = re.compile(r"^(?P<base>.+)_(?P<dims>\d+)$")
|
||||
|
||||
|
||||
class Crash(BaseException):
|
||||
"""Stands in for the process dying: nothing in the worker catches it."""
|
||||
|
||||
|
||||
@dataclass
|
||||
class Record:
|
||||
document: str
|
||||
metadata: dict[str, object]
|
||||
width: int
|
||||
model: str
|
||||
|
||||
|
||||
@dataclass
|
||||
class FakeVectorStore:
|
||||
collections: dict[str, dict[str, Record]] = field(default_factory=dict)
|
||||
|
||||
def put(
|
||||
self,
|
||||
base: str,
|
||||
name_dims: int | None,
|
||||
records: dict[str, tuple[str, str]],
|
||||
*,
|
||||
role: str = "memory",
|
||||
width: int | None = None,
|
||||
model: str = "old-model",
|
||||
) -> None:
|
||||
"""`records` maps id -> (document, user_id)."""
|
||||
name = space_collection_names(base, name_dims)[role]
|
||||
target = self.collections.setdefault(name, {})
|
||||
for record_id, (document, user) in records.items():
|
||||
target[record_id] = Record(
|
||||
document=document,
|
||||
metadata={"user_id": user, "content": document},
|
||||
width=width or name_dims or 0,
|
||||
model=model,
|
||||
)
|
||||
|
||||
def space(self, base: str, name_dims: int | None) -> dict[str, dict[str, Record]]:
|
||||
return {
|
||||
role: dict(self.collections.get(name, {}))
|
||||
for role, name in space_collection_names(base, name_dims).items()
|
||||
if name in self.collections
|
||||
}
|
||||
|
||||
def ids(self, base: str, name_dims: int | None, role: str = "memory") -> set[str]:
|
||||
return set(self.collections.get(space_collection_names(base, name_dims)[role], {}))
|
||||
|
||||
|
||||
def parse(name: str) -> tuple[str, int | None, str] | None:
|
||||
if name.endswith("_tag_index"):
|
||||
return None
|
||||
match = _ROLE_SUFFIX.match(name)
|
||||
if match:
|
||||
dims = match.group("dims")
|
||||
return match.group("base"), int(dims) if dims else None, match.group("role")
|
||||
role = "memory"
|
||||
if name.endswith("_entities"):
|
||||
name, role = name[: -len("_entities")], "entities"
|
||||
match = _DIMS_SUFFIX.match(name)
|
||||
if match:
|
||||
return match.group("base"), int(match.group("dims")), role
|
||||
return name, None, role
|
||||
|
||||
|
||||
class FakeReembedBackend:
|
||||
"""The backend contract, over `FakeVectorStore`."""
|
||||
|
||||
def __init__(self, store: FakeVectorStore | None = None) -> None:
|
||||
self.store = store or FakeVectorStore()
|
||||
self.unavailable: str | None = None
|
||||
#: Called before every batch; a test raises `Crash` or flips a cancel here.
|
||||
self.before_batch: Callable[[str, int], None] = lambda _role, _done: None
|
||||
#: Width the "embedding service" returns; None returns the target's.
|
||||
self.returned_width: int | None = None
|
||||
#: Ids the re-embed silently loses, to prove verification catches it.
|
||||
self.lose_ids: set[str] = set()
|
||||
self.batch_size = 2
|
||||
self.reembed_calls: list[tuple[str, str]] = []
|
||||
self.dropped: list[str] = []
|
||||
self.graphs_built: list[Path] = []
|
||||
#: When set, every batch waits for it: a switch a test holds mid-way.
|
||||
self.pause: asyncio.Event | None = None
|
||||
|
||||
def unavailable_reason(self) -> str | None:
|
||||
return self.unavailable
|
||||
|
||||
async def scan(self, *, per_user: bool = False) -> tuple[ScannedSpace, ...]:
|
||||
grouped: dict[tuple[str, int | None], dict[str, str]] = {}
|
||||
for name in self.store.collections:
|
||||
parsed = parse(name)
|
||||
if parsed is None:
|
||||
continue
|
||||
base, dims, role = parsed
|
||||
grouped.setdefault((base, dims), {})[role] = name
|
||||
out = []
|
||||
for (base, dims), roles in sorted(
|
||||
grouped.items(), key=lambda kv: (kv[0][0], kv[0][1] or 0)
|
||||
):
|
||||
counts = {role: len(self.store.collections[name]) for role, name in roles.items()}
|
||||
widths = {
|
||||
record.width
|
||||
for name in roles.values()
|
||||
for record in self.store.collections[name].values()
|
||||
}
|
||||
users: dict[str, int] = {}
|
||||
if per_user:
|
||||
for name in roles.values():
|
||||
for record in self.store.collections[name].values():
|
||||
user = str(record.metadata.get("user_id") or "")
|
||||
users[user] = users.get(user, 0) + 1
|
||||
out.append(
|
||||
ScannedSpace(
|
||||
collection_base=base,
|
||||
name_dims=dims,
|
||||
width=widths.pop() if len(widths) == 1 else None,
|
||||
counts=counts,
|
||||
collections=roles,
|
||||
per_user=users,
|
||||
)
|
||||
)
|
||||
return tuple(out)
|
||||
|
||||
async def diff(self, source: PhysicalSpace, target: PhysicalSpace) -> dict[str, int]:
|
||||
target_names = target.names()
|
||||
out = {}
|
||||
for role, name in source.names().items():
|
||||
have = set(self.store.collections.get(target_names.get(role, ""), {}))
|
||||
out[role] = len(set(self.store.collections.get(name, {})) - have)
|
||||
return out
|
||||
|
||||
async def reembed(
|
||||
self,
|
||||
source: PhysicalSpace,
|
||||
target: PhysicalSpace,
|
||||
*,
|
||||
embedding: EmbeddingTarget,
|
||||
cancel: Callable[[], bool],
|
||||
on_step: Callable[[ReembedStep], None],
|
||||
batch_size: int = 64,
|
||||
) -> None:
|
||||
del batch_size
|
||||
if target.dims is None:
|
||||
raise ValueError("the target space needs an explicit width")
|
||||
self.reembed_calls.append((source.collection_base, target.collection_base))
|
||||
target_names = target.names()
|
||||
model = embedding.space.model or "?"
|
||||
for role in VECTOR_ROLES:
|
||||
name = source.names().get(role)
|
||||
if name is None or name not in self.store.collections:
|
||||
continue
|
||||
records = sorted(self.store.collections[name].items())
|
||||
destination = self.store.collections.setdefault(target_names[role], {})
|
||||
done = embedded = kept = 0
|
||||
for start in range(0, len(records), self.batch_size):
|
||||
if cancel():
|
||||
raise ReembedCancelledError("embedding_switch_cancelled")
|
||||
self.before_batch(role, done)
|
||||
if self.pause is not None:
|
||||
await self.pause.wait()
|
||||
batch = records[start : start + self.batch_size]
|
||||
width = self.returned_width or target.dims
|
||||
if width != target.dims:
|
||||
raise EmbeddingWidthError("embedding_dimensions_mismatch")
|
||||
for record_id, record in batch:
|
||||
if record_id in destination:
|
||||
kept += 1
|
||||
continue
|
||||
if record_id in self.lose_ids:
|
||||
continue
|
||||
destination[record_id] = Record(
|
||||
document=record.document,
|
||||
metadata=dict(record.metadata),
|
||||
width=width,
|
||||
model=model,
|
||||
)
|
||||
embedded += 1
|
||||
done += len(batch)
|
||||
on_step(
|
||||
ReembedStep(
|
||||
phase="vectors",
|
||||
role=role,
|
||||
done=done,
|
||||
total=len(records),
|
||||
embedded=embedded,
|
||||
kept=kept,
|
||||
)
|
||||
)
|
||||
|
||||
async def verify(
|
||||
self, source: PhysicalSpace, target: PhysicalSpace, *, embedding: EmbeddingTarget
|
||||
) -> VectorCheck:
|
||||
del embedding
|
||||
problems: list[str] = []
|
||||
per_role: dict[str, dict[str, int]] = {}
|
||||
users_source: dict[str, int] = {}
|
||||
users_target: dict[str, int] = {}
|
||||
target_names = target.names()
|
||||
for role, name in source.names().items():
|
||||
src = self.store.collections.get(name, {})
|
||||
tgt = self.store.collections.get(target_names.get(role, ""), {})
|
||||
missing = set(src) - set(tgt)
|
||||
extra = set(tgt) - set(src)
|
||||
mismatched = {
|
||||
key for key in set(src) & set(tgt) if src[key].document != tgt[key].document
|
||||
}
|
||||
bad = sum(1 for record in tgt.values() if record.width != target.dims)
|
||||
per_role[role] = {
|
||||
"source": len(src),
|
||||
"target": len(tgt),
|
||||
"missing": len(missing),
|
||||
"extra": len(extra),
|
||||
"mismatched": len(mismatched),
|
||||
"bad_dims": bad,
|
||||
}
|
||||
for label, value in (
|
||||
("missing", missing),
|
||||
("extra", extra),
|
||||
("mismatched", mismatched),
|
||||
):
|
||||
if value:
|
||||
problems.append(f"{role}: {len(value)} {label}")
|
||||
if bad:
|
||||
problems.append(f"{role}: {bad} bad dims")
|
||||
for record in src.values():
|
||||
user = str(record.metadata.get("user_id") or "")
|
||||
users_source[user] = users_source.get(user, 0) + 1
|
||||
for record in tgt.values():
|
||||
user = str(record.metadata.get("user_id") or "")
|
||||
users_target[user] = users_target.get(user, 0) + 1
|
||||
for user in sorted(set(users_source) | set(users_target)):
|
||||
if users_source.get(user, 0) != users_target.get(user, 0):
|
||||
problems.append(f"user {user!r}: counts differ")
|
||||
return VectorCheck(per_role=per_role, per_user=users_target, problems=tuple(problems))
|
||||
|
||||
async def graph_vectors(self, path: Path) -> int | None:
|
||||
return 3 if path.exists() else None
|
||||
|
||||
async def reembed_graph(
|
||||
self,
|
||||
source: Path,
|
||||
target: Path,
|
||||
*,
|
||||
embedding: EmbeddingTarget,
|
||||
target_dims: int,
|
||||
cancel: Callable[[], bool],
|
||||
on_step: Callable[[ReembedStep], None],
|
||||
) -> None:
|
||||
del embedding, cancel
|
||||
target.mkdir(parents=True, exist_ok=True)
|
||||
(target / "width").write_text(str(target_dims))
|
||||
self.graphs_built.append(target)
|
||||
on_step(ReembedStep(phase="graph", role="Memory", done=3, total=3, embedded=3))
|
||||
|
||||
async def verify_graph(self, source: Path, target: Path, *, target_dims: int) -> GraphCheck:
|
||||
ok = (target / "width").is_file() and (target / "width").read_text() == str(target_dims)
|
||||
return GraphCheck(ok=ok, problems=() if ok else ("graph width",))
|
||||
|
||||
async def drop(self, space: PhysicalSpace) -> None:
|
||||
for name in space.names().values():
|
||||
if self.store.collections.pop(name, None) is not None:
|
||||
self.dropped.append(name)
|
||||
|
||||
async def drop_graph(self, path: Path) -> None:
|
||||
shutil.rmtree(path, ignore_errors=True)
|
||||
@@ -24,7 +24,8 @@ def model_snapshot(
|
||||
model="memory-model",
|
||||
embedding_model="embedding-model",
|
||||
embedding_provider_kind=embedding_provider,
|
||||
embedding_base_url=embedding_base_url,
|
||||
# Stored profiles always name it; see `pinned_embedding_base_url`.
|
||||
embedding_base_url=embedding_base_url or base_url,
|
||||
embedding_dimensions=768,
|
||||
timeout_seconds=30,
|
||||
max_concurrency=1,
|
||||
|
||||
@@ -175,7 +175,13 @@ def test_sdk_migrate_reports_each_step_then_the_result(tmp_path, monkeypatch) ->
|
||||
|
||||
assert code == 0, output.getvalue()
|
||||
lines = [json.loads(line) for line in output.getvalue().splitlines()]
|
||||
assert [line["step"] for line in lines[:-1]] == ["bundle", "prepare", "cutover", "cleanup"]
|
||||
assert [line["step"] for line in lines[:-1]] == [
|
||||
"bundle",
|
||||
"prepare",
|
||||
"vectors",
|
||||
"cutover",
|
||||
"cleanup",
|
||||
]
|
||||
assert lines[-1]["state"] == "cutover"
|
||||
assert lines[-1]["bundle_id"] == f"{lines[-1]['migration_id']}-pre-cutover"
|
||||
assert not (tmp_path / ".sdk-migration-capture-fence.json").exists()
|
||||
|
||||
@@ -164,7 +164,16 @@ def test_metadata_repair_never_overlaps_an_sdk_client(tmp_path, monkeypatch):
|
||||
events.append("repair_finished")
|
||||
return MetadataRepairResult(None, 0, 0, 0)
|
||||
|
||||
def open_client(profile, *, api_key, data_dir, http_transport_factory):
|
||||
def open_client(
|
||||
profile,
|
||||
space=None,
|
||||
*,
|
||||
api_key,
|
||||
data_dir,
|
||||
http_transport_factory,
|
||||
mode="pro",
|
||||
features=(),
|
||||
):
|
||||
events.append("client_opened")
|
||||
return ClientLease(object(), data_dir.resolve())
|
||||
|
||||
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,383 @@
|
||||
"""The embedding switch through Core's HTTP surface and lifespan.
|
||||
|
||||
Activation refusals and the empty-home path, a whole switch driven over the
|
||||
API with writes paused and intake still queueing, the startup merge of a
|
||||
stray space before the drain starts, and `embedding_unavailable` on recall
|
||||
and context search when the remote embedding does not answer.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import time
|
||||
from types import SimpleNamespace
|
||||
from uuid import uuid4
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from httpx import ASGITransport, AsyncClient
|
||||
from pydantic import SecretStr
|
||||
|
||||
from one_memory import embedding_health
|
||||
from one_memory.api import create_app
|
||||
from one_memory.config import Settings
|
||||
from one_memory.model_profiles import ModelCapabilities, ModelProfileInput, ProviderKind
|
||||
from one_memory.secrets import InMemorySecretStore
|
||||
from tests.embedding_fakes import FakeReembedBackend
|
||||
|
||||
pytestmark = pytest.mark.fresh_sdk
|
||||
|
||||
ENDPOINT = "https://models.example/v1"
|
||||
|
||||
|
||||
def builtin_base() -> str:
|
||||
identity = "\0".join(("local_onnx", "local", "bge-small-zh-v1.5", "512"))
|
||||
return "memhub_" + hashlib.sha256(identity.encode()).hexdigest()[:16]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def backend(monkeypatch) -> FakeReembedBackend:
|
||||
fake = FakeReembedBackend()
|
||||
monkeypatch.setattr(
|
||||
"one_memory.embedding_switch.default_reembed_backend", lambda _data_dir: fake
|
||||
)
|
||||
return fake
|
||||
|
||||
|
||||
def application(tmp_path):
|
||||
return create_app(
|
||||
Settings(
|
||||
data_dir=tmp_path / "data",
|
||||
source_home=tmp_path,
|
||||
api_token=SecretStr("switch-test"),
|
||||
imports_enabled=False,
|
||||
automatic_memory_processing=False,
|
||||
live_sync_enabled=False,
|
||||
),
|
||||
secret_store=InMemorySecretStore(),
|
||||
)
|
||||
|
||||
|
||||
def client_for(app):
|
||||
return AsyncClient(
|
||||
transport=ASGITransport(app=app),
|
||||
base_url="http://test",
|
||||
headers={"Authorization": "Bearer switch-test"},
|
||||
)
|
||||
|
||||
|
||||
async def recovered(client) -> dict:
|
||||
deadline = time.monotonic() + 10
|
||||
while time.monotonic() < deadline:
|
||||
body = (await client.get("/api/v1/health")).json()
|
||||
if body["recovery"] in {"done", "failed"}:
|
||||
return body
|
||||
await asyncio.sleep(0.01)
|
||||
raise AssertionError("recovery did not finish")
|
||||
|
||||
|
||||
async def until(client, predicate, *, seconds: float = 10.0) -> dict:
|
||||
deadline = time.monotonic() + seconds
|
||||
while time.monotonic() < deadline:
|
||||
body = (await client.get("/api/v1/embedding-switches/current")).json()
|
||||
if predicate(body):
|
||||
return body
|
||||
await asyncio.sleep(0.01)
|
||||
raise AssertionError(f"switch never reached the expected state: {body}")
|
||||
|
||||
|
||||
def remote_profile(app, *, model="bge-m3", dims=1024, name="custom", endpoint=ENDPOINT):
|
||||
store, secrets = app.state.model_profile_store, app.state.secret_store
|
||||
created = store.create(
|
||||
ModelProfileInput(
|
||||
name=name,
|
||||
provider_kind=ProviderKind.OPENAI_COMPATIBLE,
|
||||
base_url=endpoint,
|
||||
model="chat",
|
||||
embedding_provider_kind=ProviderKind.OPENAI_COMPATIBLE,
|
||||
embedding_base_url=endpoint,
|
||||
embedding_model=model,
|
||||
embedding_dimensions=dims,
|
||||
api_key=SecretStr("sk-test"),
|
||||
),
|
||||
secrets,
|
||||
)
|
||||
store.record_capabilities(
|
||||
created.profile_id,
|
||||
expected_revision=created.revision,
|
||||
capabilities=ModelCapabilities(
|
||||
endpoint_reachable=True,
|
||||
model_available=True,
|
||||
structured_output=True,
|
||||
embedding_available=True,
|
||||
),
|
||||
secrets=secrets,
|
||||
)
|
||||
return created
|
||||
|
||||
|
||||
def test_activating_another_embedding_answers_with_the_plan(tmp_path, backend):
|
||||
backend.store.put(builtin_base(), 512, {"a": ("one", "local"), "b": ("two", "bob")})
|
||||
|
||||
async def scenario():
|
||||
app = application(tmp_path)
|
||||
async with app.router.lifespan_context(app), client_for(app) as client:
|
||||
await recovered(client)
|
||||
profile = remote_profile(app)
|
||||
|
||||
response = await client.post(f"/api/v1/model-profiles/{profile.profile_id}/activate")
|
||||
|
||||
assert response.status_code == 409
|
||||
body = response.json()
|
||||
assert body["detail"] == "embedding_switch_required"
|
||||
plan = body["plan"]
|
||||
assert plan["needed"] is True and plan["total_vectors"] == 2 and plan["users"] == 2
|
||||
assert plan["source"]["model"] == "bge-small-zh-v1.5"
|
||||
assert (plan["target"]["model"], plan["target"]["dims"]) == ("bge-m3", 1024)
|
||||
assert app.state.model_profile_store.active_identity() is None
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_an_empty_home_activates_another_embedding_directly(tmp_path, backend):
|
||||
async def scenario():
|
||||
app = application(tmp_path)
|
||||
async with app.router.lifespan_context(app), client_for(app) as client:
|
||||
await recovered(client)
|
||||
builtin = app.state.embedding_spaces.require_live()
|
||||
profile = remote_profile(app)
|
||||
|
||||
response = await client.post(f"/api/v1/model-profiles/{profile.profile_id}/activate")
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.json()["is_active"] is True
|
||||
live = app.state.embedding_spaces.require_live()
|
||||
assert (live.model, live.dims) == ("bge-m3", 1024)
|
||||
assert app.state.embedding_spaces.get(builtin.space_id).state == "dropped"
|
||||
# The same embedding again: a plain activation, no plan.
|
||||
again = remote_profile(app, name="same model, other chat model")
|
||||
assert (
|
||||
await client.post(f"/api/v1/model-profiles/{again.profile_id}/activate")
|
||||
).status_code == 200
|
||||
assert app.state.embedding_spaces.require_live().space_id == live.space_id
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_a_switch_over_the_api_pauses_writes_and_keeps_intake(tmp_path, backend):
|
||||
backend.store.put(
|
||||
builtin_base(),
|
||||
512,
|
||||
{f"m{i}": (f"memory {i}", "local") for i in range(6)},
|
||||
model="bge-small-zh-v1.5",
|
||||
)
|
||||
|
||||
async def scenario():
|
||||
app = application(tmp_path)
|
||||
async with app.router.lifespan_context(app), client_for(app) as client:
|
||||
await recovered(client)
|
||||
old = app.state.embedding_spaces.require_live()
|
||||
profile = remote_profile(app)
|
||||
planned = await client.post(
|
||||
"/api/v1/embedding-switches/plan", json={"profile_id": profile.profile_id}
|
||||
)
|
||||
assert planned.status_code == 200
|
||||
plan = planned.json()
|
||||
stale = await client.post(
|
||||
"/api/v1/embedding-switches",
|
||||
json={
|
||||
"profile_id": profile.profile_id,
|
||||
"expected_revision": profile.revision,
|
||||
"plan_token": "stale",
|
||||
},
|
||||
)
|
||||
assert stale.status_code == 409
|
||||
assert stale.json()["detail"] == "embedding_switch_plan_stale"
|
||||
|
||||
backend.pause = asyncio.Event()
|
||||
started = await client.post(
|
||||
"/api/v1/embedding-switches",
|
||||
json={
|
||||
"profile_id": profile.profile_id,
|
||||
"expected_revision": profile.revision,
|
||||
"plan_token": plan["plan_token"],
|
||||
"requested_by": "cli",
|
||||
},
|
||||
)
|
||||
assert started.status_code == 202, started.text
|
||||
job_id = started.json()["job_id"]
|
||||
running = await until(
|
||||
client, lambda body: (body["job"] or {}).get("state") == "reembedding"
|
||||
)
|
||||
assert running["writes_paused"] is True and running["phase"] == "running"
|
||||
|
||||
# Direct writers wait; intake keeps queueing; status shows the switch.
|
||||
edit = await client.put(
|
||||
f"/api/v2/memories/{uuid4()}", json={"content": "changed", "user_id": "local"}
|
||||
)
|
||||
assert edit.status_code == 409
|
||||
assert edit.json()["detail"] == "embedding_switch_in_progress"
|
||||
delete = await client.delete(f"/api/v2/memories/{uuid4()}")
|
||||
assert delete.status_code == 409
|
||||
added = await client.post(
|
||||
"/api/v2/memories",
|
||||
json={"request_id": str(uuid4()), "text": "请记住这件事", "scope": "all-apps"},
|
||||
)
|
||||
assert added.status_code == 201
|
||||
processing = (await client.get("/api/v1/memory-processing")).json()
|
||||
assert processing["embedding_switch"]["job_id"] == job_id
|
||||
assert (await client.get("/api/v1/health")).json()["embedding_switch"]["kind"] == (
|
||||
"switch"
|
||||
)
|
||||
other = remote_profile(app, model="text-embedding-v3", dims=768, name="other")
|
||||
refused = await client.post(f"/api/v1/model-profiles/{other.profile_id}/activate")
|
||||
assert refused.json()["detail"] == "embedding_switch_running"
|
||||
|
||||
backend.pause.set()
|
||||
done = await until(
|
||||
client, lambda body: (body["last_finished"] or {}).get("job_id") == job_id
|
||||
)
|
||||
assert done["last_finished"]["state"] == "done"
|
||||
assert done["writes_paused"] is False
|
||||
assert app.state.model_profile_store.active_identity()[0] == profile.profile_id
|
||||
|
||||
spaces = (await client.get("/api/v1/embedding-spaces")).json()
|
||||
states = {space["space_id"]: space["state"] for space in spaces["spaces"]}
|
||||
assert states[old.space_id] == "retired"
|
||||
assert spaces["live_space_id"] == done["last_finished"]["target"]["space_id"]
|
||||
cleaned = await client.delete(f"/api/v1/embedding-spaces/{old.space_id}")
|
||||
assert cleaned.status_code == 200
|
||||
assert cleaned.json()["space"]["state"] == "dropped"
|
||||
live_cleanup = await client.delete(
|
||||
f"/api/v1/embedding-spaces/{spaces['live_space_id']}"
|
||||
)
|
||||
assert live_cleanup.json()["detail"] == "embedding_space_not_retired"
|
||||
|
||||
# Same model and width behind another endpoint: no re-embed.
|
||||
proxied = remote_profile(app, name="proxy", endpoint="https://proxy.example/v1")
|
||||
moved = await client.post(f"/api/v1/model-profiles/{proxied.profile_id}/activate")
|
||||
assert moved.status_code == 200, moved.text
|
||||
assert app.state.embedding_spaces.require_live().space_id == spaces["live_space_id"]
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_cancel_over_the_api(tmp_path, backend):
|
||||
backend.store.put(builtin_base(), 512, {f"m{i}": (f"m {i}", "local") for i in range(4)})
|
||||
|
||||
async def scenario():
|
||||
app = application(tmp_path)
|
||||
async with app.router.lifespan_context(app), client_for(app) as client:
|
||||
await recovered(client)
|
||||
profile = remote_profile(app)
|
||||
plan = (
|
||||
await client.post(
|
||||
"/api/v1/embedding-switches/plan", json={"profile_id": profile.profile_id}
|
||||
)
|
||||
).json()
|
||||
backend.pause = asyncio.Event()
|
||||
job_id = (
|
||||
await client.post(
|
||||
"/api/v1/embedding-switches",
|
||||
json={
|
||||
"profile_id": profile.profile_id,
|
||||
"expected_revision": profile.revision,
|
||||
"plan_token": plan["plan_token"],
|
||||
},
|
||||
)
|
||||
).json()["job_id"]
|
||||
await until(client, lambda body: (body["job"] or {}).get("state") == "reembedding")
|
||||
|
||||
cancelled = await client.post(f"/api/v1/embedding-switches/{job_id}/cancel")
|
||||
assert cancelled.status_code == 200 and cancelled.json()["cancel_requested"]
|
||||
backend.pause.set()
|
||||
done = await until(client, lambda body: body["last_finished"] is not None)
|
||||
|
||||
assert done["last_finished"]["state"] == "rolled_back"
|
||||
assert done["last_finished"]["error_code"] == "embedding_switch_cancelled"
|
||||
missing = await client.post(f"/api/v1/embedding-switches/{uuid4()}/cancel")
|
||||
assert missing.status_code == 404
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_startup_merges_a_stray_space_before_writes_resume(tmp_path, backend):
|
||||
backend.store.put(builtin_base(), 512, {"live-1": ("live", "local")})
|
||||
stray = "memhub_" + "a" * 16
|
||||
backend.store.put(stray, 1024, {"stray-1": ("left by an older build", "hermes-1")})
|
||||
|
||||
async def scenario():
|
||||
app = application(tmp_path)
|
||||
async with app.router.lifespan_context(app), client_for(app) as client:
|
||||
# Core answers at once; the merge runs behind the loading screen.
|
||||
assert (await client.get("/api/v1/health")).status_code == 200
|
||||
await recovered(client)
|
||||
done = await until(client, lambda body: body["last_finished"] is not None)
|
||||
|
||||
assert done["last_finished"]["kind"] == "merge"
|
||||
assert done["last_finished"]["state"] == "done"
|
||||
assert backend.store.ids(builtin_base(), 512) == {"live-1", "stray-1"}
|
||||
spaces = (await client.get("/api/v1/embedding-spaces")).json()
|
||||
stray_row = next(row for row in spaces["spaces"] if row["collection_base"] == stray)
|
||||
assert stray_row["state"] == "retired" and stray_row["keep_until"]
|
||||
assert app.state.embedding_switch.writes_open()
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
class _EmptySearchClient:
|
||||
def search(self, *_args, **_kwargs):
|
||||
return {"request_id": "r", "memories": {"schema": [], "proactive": [], "normal": []}}
|
||||
|
||||
def get(self, *_args, **_kwargs):
|
||||
return None
|
||||
|
||||
|
||||
def test_a_dead_remote_embedding_answers_embedding_unavailable(tmp_path, backend, monkeypatch):
|
||||
status = {"code": 503}
|
||||
|
||||
def transport():
|
||||
return httpx.MockTransport(
|
||||
lambda request: httpx.Response(
|
||||
status["code"], json={"data": [{"embedding": [0.1] * 4}]}
|
||||
)
|
||||
)
|
||||
|
||||
monkeypatch.setattr(embedding_health, "transport_factory", transport)
|
||||
monkeypatch.setattr(
|
||||
"one_memory.hy_sdk_runtime.acquire_client",
|
||||
lambda *args, **kwargs: SimpleNamespace(client=_EmptySearchClient(), close=lambda: None),
|
||||
)
|
||||
|
||||
async def ready(_app):
|
||||
return True
|
||||
|
||||
monkeypatch.setattr("one_memory.api._recall_model_ready", ready)
|
||||
|
||||
async def scenario():
|
||||
app = application(tmp_path)
|
||||
async with app.router.lifespan_context(app), client_for(app) as client:
|
||||
await recovered(client)
|
||||
profile = remote_profile(app)
|
||||
assert (
|
||||
await client.post(f"/api/v1/model-profiles/{profile.profile_id}/activate")
|
||||
).status_code == 200
|
||||
|
||||
recall = await client.post("/api/v1/recall", json={"query": "数据库"})
|
||||
assert recall.status_code == 503
|
||||
assert recall.json()["detail"] == "embedding_unavailable"
|
||||
context = await client.post(
|
||||
"/api/v2/context/search",
|
||||
json={"query": "数据库", "user_id": "local", "kinds": ["memory"]},
|
||||
)
|
||||
assert context.status_code == 503
|
||||
assert context.json()["detail"] == "embedding_unavailable"
|
||||
|
||||
# The service answers again: an empty result is just empty.
|
||||
status["code"] = 200
|
||||
app.state.sdk_client_handle.health.forget()
|
||||
recall = await client.post("/api/v1/recall", json={"query": "数据库"})
|
||||
assert recall.status_code == 200
|
||||
|
||||
asyncio.run(scenario())
|
||||
@@ -0,0 +1,604 @@
|
||||
"""The re-embed adapter against the SDK's real re-embed module and Chroma.
|
||||
|
||||
Skipped on an SDK without `hy_memory.reembed` (before 2.0.10): there the
|
||||
switch is unavailable and the plan says `sdk_reembed_unavailable`, which
|
||||
`test_embedding_switch.py` covers.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from pathlib import Path
|
||||
|
||||
import hy_memory
|
||||
import pytest
|
||||
|
||||
from one_memory.embedding_backend import EmbeddingTarget, PhysicalSpace
|
||||
from one_memory.embedding_spaces import EmbeddingSpace, space_collection_names
|
||||
from one_memory.hy_sdk_migration import SdkReembedBackend
|
||||
from one_memory.hy_sdk_runtime import builtin_retrieval_snapshot
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
SdkReembedBackend(Path("/nonexistent")).unavailable_reason() is not None,
|
||||
reason="the installed SDK has no re-embed module (hy-memory < 2.0.10)",
|
||||
)
|
||||
|
||||
SOURCE = "memhub_" + "1" * 16
|
||||
TARGET = "memhub_e_" + "2" * 16
|
||||
|
||||
|
||||
class _Embedder:
|
||||
"""A deterministic 8-wide stand-in for the target model."""
|
||||
|
||||
def __init__(self, width: int = 8) -> None:
|
||||
self.width = width
|
||||
self.texts: list[str] = []
|
||||
|
||||
async def embed_batch(self, texts: list[str], use_cache: bool = True) -> list[list[float]]:
|
||||
self.texts.extend(texts)
|
||||
return [[float(len(text) % 7 + 1)] + [0.5] * (self.width - 1) for text in texts]
|
||||
|
||||
async def embed(self, text: str, use_cache: bool = True) -> list[float]:
|
||||
return (await self.embed_batch([text]))[0]
|
||||
|
||||
def get_embedding_dims(self) -> int:
|
||||
return self.width
|
||||
|
||||
|
||||
def _seed(data_dir: Path) -> None:
|
||||
import chromadb
|
||||
|
||||
vector = data_dir / "hy-memory" / "governed-v2" / "vector"
|
||||
vector.mkdir(parents=True)
|
||||
client = chromadb.PersistentClient(path=str(vector))
|
||||
names = space_collection_names(SOURCE, 4)
|
||||
memory = client.get_or_create_collection(names["memory"], metadata={"hnsw:space": "cosine"})
|
||||
memory.add(
|
||||
ids=["m1", "m2", "m3"],
|
||||
embeddings=[[0.1, 0.2, 0.3, 0.4]] * 3,
|
||||
documents=["one", "two", "three"],
|
||||
metadatas=[{"user_id": "local"}, {"user_id": "hermes-a"}, {"user_id": "hermes-b"}],
|
||||
)
|
||||
keys = client.get_or_create_collection(names["coding_keys"], metadata={"hnsw:space": "cosine"})
|
||||
keys.add(
|
||||
ids=["e1:0"],
|
||||
embeddings=[[0.4, 0.3, 0.2, 0.1]],
|
||||
documents=["task"],
|
||||
metadatas=[{"user_id": "local"}],
|
||||
)
|
||||
|
||||
|
||||
def _target() -> EmbeddingTarget:
|
||||
space = EmbeddingSpace(
|
||||
space_id="target",
|
||||
state="building",
|
||||
model="stand-in",
|
||||
dims=8,
|
||||
name_dims=8,
|
||||
collection_base=TARGET,
|
||||
naming="v2",
|
||||
created_at="2026-11-01T00:00:00+00:00",
|
||||
)
|
||||
return EmbeddingTarget(profile=builtin_retrieval_snapshot(), space=space, api_key=None)
|
||||
|
||||
|
||||
def test_core_names_a_space_exactly_as_the_sdk_does() -> None:
|
||||
import hy_memory
|
||||
|
||||
for dims in (None, 512, 1024):
|
||||
sdk = hy_memory.vector_space_names(hy_memory.VectorSpace(base_name=SOURCE, dims=dims))
|
||||
assert sdk == space_collection_names(SOURCE, dims)
|
||||
|
||||
|
||||
def test_scan_reembed_verify_and_drop_through_the_sdk(tmp_path: Path, monkeypatch) -> None:
|
||||
_seed(tmp_path)
|
||||
backend = SdkReembedBackend(tmp_path)
|
||||
embedder = _Embedder()
|
||||
monkeypatch.setattr(backend, "_embedding_service", lambda _target: embedder)
|
||||
source = PhysicalSpace(collection_base=SOURCE, name_dims=4, dims=4)
|
||||
target = PhysicalSpace(collection_base=TARGET, name_dims=8, dims=8, model="stand-in")
|
||||
steps = []
|
||||
|
||||
async def scenario() -> None:
|
||||
scanned = await backend.scan(per_user=True)
|
||||
found = {item.collection_base: item for item in scanned}
|
||||
assert found[SOURCE].vectors == 4
|
||||
assert found[SOURCE].per_user == {"local": 2, "hermes-a": 1, "hermes-b": 1}
|
||||
|
||||
await backend.reembed(
|
||||
source, target, embedding=_target(), cancel=lambda: False, on_step=steps.append
|
||||
)
|
||||
assert not any((await backend.diff(source, target)).values())
|
||||
check = await backend.verify(source, target, embedding=_target())
|
||||
assert check.problems == ()
|
||||
assert check.per_user == {"local": 2, "hermes-a": 1, "hermes-b": 1}
|
||||
|
||||
# Resumable: a second run embeds nothing more.
|
||||
before = len(embedder.texts)
|
||||
await backend.reembed(
|
||||
source, target, embedding=_target(), cancel=lambda: False, on_step=steps.append
|
||||
)
|
||||
assert len(embedder.texts) == before
|
||||
|
||||
await backend.drop(target)
|
||||
remaining = {item.collection_base for item in await backend.scan()}
|
||||
assert TARGET not in remaining and SOURCE in remaining
|
||||
|
||||
asyncio.run(scenario())
|
||||
assert {step.role for step in steps} >= {"memory", "coding_keys"}
|
||||
|
||||
|
||||
def test_a_wrong_width_fails_before_anything_is_written(tmp_path: Path, monkeypatch) -> None:
|
||||
from one_memory.embedding_backend import EmbeddingWidthError
|
||||
|
||||
_seed(tmp_path)
|
||||
backend = SdkReembedBackend(tmp_path)
|
||||
monkeypatch.setattr(backend, "_embedding_service", lambda _target: _Embedder(width=6))
|
||||
source = PhysicalSpace(collection_base=SOURCE, name_dims=4, dims=4)
|
||||
target = PhysicalSpace(collection_base=TARGET, name_dims=8, dims=8)
|
||||
|
||||
async def scenario() -> None:
|
||||
with pytest.raises(EmbeddingWidthError):
|
||||
await backend.reembed(
|
||||
source, target, embedding=_target(), cancel=lambda: False, on_step=lambda _: None
|
||||
)
|
||||
found = {item.collection_base: item for item in await backend.scan()}
|
||||
assert TARGET not in found or found[TARGET].vectors == 0
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_cancel_stops_at_a_batch_boundary(tmp_path: Path, monkeypatch) -> None:
|
||||
from one_memory.embedding_backend import ReembedCancelledError
|
||||
|
||||
_seed(tmp_path)
|
||||
backend = SdkReembedBackend(tmp_path)
|
||||
monkeypatch.setattr(backend, "_embedding_service", lambda _target: _Embedder())
|
||||
source = PhysicalSpace(collection_base=SOURCE, name_dims=4, dims=4)
|
||||
target = PhysicalSpace(collection_base=TARGET, name_dims=8, dims=8)
|
||||
|
||||
async def scenario() -> None:
|
||||
with pytest.raises(ReembedCancelledError):
|
||||
await backend.reembed(
|
||||
source, target, embedding=_target(), cancel=lambda: True, on_step=lambda _: None
|
||||
)
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_a_whole_switch_through_the_real_sdk(tmp_path: Path, monkeypatch) -> None:
|
||||
"""Registry, worker and commit over real Chroma collections and the SDK."""
|
||||
import chromadb
|
||||
from pydantic import SecretStr
|
||||
|
||||
from one_memory.embedding_spaces import EmbeddingSpaceStore
|
||||
from one_memory.embedding_switch import EmbeddingSwitchService
|
||||
from one_memory.model_profiles import (
|
||||
ModelCapabilities,
|
||||
ModelProfileInput,
|
||||
ModelProfileStore,
|
||||
ProviderKind,
|
||||
)
|
||||
from one_memory.secrets import InMemorySecretStore
|
||||
|
||||
database = tmp_path / "one-memory.sqlite3"
|
||||
secrets = InMemorySecretStore()
|
||||
profiles = ModelProfileStore(database)
|
||||
spaces = EmbeddingSpaceStore(database)
|
||||
live = spaces.bootstrap(None)
|
||||
vector = tmp_path / "hy-memory" / "governed-v2" / "vector"
|
||||
vector.mkdir(parents=True)
|
||||
client = chromadb.PersistentClient(path=str(vector))
|
||||
memory = client.get_or_create_collection(
|
||||
live.collection_names()["memory"], metadata={"hnsw:space": "cosine"}
|
||||
)
|
||||
memory.add(
|
||||
ids=["a", "b", "c"],
|
||||
embeddings=[[0.1] * 512, [0.2] * 512, [0.3] * 512],
|
||||
documents=["one", "two", "three"],
|
||||
metadatas=[{"user_id": "local"}, {"user_id": "hermes-a"}, {"user_id": "local"}],
|
||||
)
|
||||
stray_base = "memhub_" + "9" * 16
|
||||
stray = client.get_or_create_collection(
|
||||
space_collection_names(stray_base, 768)["memory"], metadata={"hnsw:space": "cosine"}
|
||||
)
|
||||
stray.add(
|
||||
ids=["a", "z"],
|
||||
embeddings=[[0.5] * 768, [0.6] * 768],
|
||||
documents=["stale", "only in the stray"],
|
||||
metadatas=[{"user_id": "local"}, {"user_id": "hermes-b"}],
|
||||
)
|
||||
created = profiles.create(
|
||||
ModelProfileInput(
|
||||
name="remote",
|
||||
provider_kind=ProviderKind.OPENAI_COMPATIBLE,
|
||||
base_url="https://models.example/v1",
|
||||
model="chat",
|
||||
embedding_provider_kind=ProviderKind.OPENAI_COMPATIBLE,
|
||||
embedding_base_url="https://models.example/v1",
|
||||
embedding_model="stand-in",
|
||||
embedding_dimensions=8,
|
||||
api_key=SecretStr("sk-test"),
|
||||
),
|
||||
secrets,
|
||||
)
|
||||
profiles.record_capabilities(
|
||||
created.profile_id,
|
||||
expected_revision=created.revision,
|
||||
capabilities=ModelCapabilities(
|
||||
endpoint_reachable=True,
|
||||
model_available=True,
|
||||
structured_output=True,
|
||||
embedding_available=True,
|
||||
),
|
||||
secrets=secrets,
|
||||
)
|
||||
backend = SdkReembedBackend(tmp_path)
|
||||
monkeypatch.setattr(backend, "_embedding_service", lambda _target: _Embedder())
|
||||
service = EmbeddingSwitchService(
|
||||
database_path=database,
|
||||
data_dir=tmp_path,
|
||||
profiles=profiles,
|
||||
spaces=spaces,
|
||||
secrets=secrets,
|
||||
backend=backend,
|
||||
worker_enabled=False,
|
||||
)
|
||||
service.release_scan_hold()
|
||||
|
||||
async def scenario() -> None:
|
||||
plan = await service.plan(created.profile_id)
|
||||
assert plan.needed and plan.total_vectors == 5 and plan.users == 3
|
||||
job = await service.start_switch(
|
||||
created.profile_id,
|
||||
expected_revision=created.revision,
|
||||
plan_token=plan.plan_token,
|
||||
requested_by="cli",
|
||||
)
|
||||
await service.run(job.job_id)
|
||||
finished = service.get_job(job.job_id)
|
||||
assert finished.state == "done", service.jobs.get(job.job_id)["verification_json"]
|
||||
|
||||
asyncio.run(scenario())
|
||||
new = spaces.require_live()
|
||||
target = client.get_collection(new.collection_names()["memory"])
|
||||
got = target.get(include=["documents", "embeddings"])
|
||||
documents = dict(zip(got["ids"], got["documents"], strict=True))
|
||||
assert documents == {"a": "one", "b": "two", "c": "three", "z": "only in the stray"}
|
||||
assert {len(vector) for vector in got["embeddings"]} == {8}
|
||||
assert profiles.active_identity() == (created.profile_id, created.revision)
|
||||
states = {space.collection_base: space.state for space in spaces.list_spaces()}
|
||||
assert states[live.collection_base] == "retired"
|
||||
assert states[stray_base] == "retired"
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# Nothing stood in: the SDK's EmbedService over HTTP, real Chroma, real Kuzu
|
||||
# ----------------------------------------------------------------------
|
||||
|
||||
|
||||
class _EmbeddingServer:
|
||||
"""A local OpenAI-compatible `/embeddings` endpoint; counts what it was sent."""
|
||||
|
||||
def __init__(self, width: int) -> None:
|
||||
import json
|
||||
import threading
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
|
||||
self.width = width
|
||||
self.inputs: list[str] = []
|
||||
self.dimensions: list[object] = []
|
||||
server = self
|
||||
|
||||
class Handler(BaseHTTPRequestHandler):
|
||||
def do_POST(self) -> None: # noqa: N802 - http.server's name
|
||||
body = json.loads(self.rfile.read(int(self.headers["Content-Length"])))
|
||||
texts = body["input"] if isinstance(body["input"], list) else [body["input"]]
|
||||
server.inputs.extend(texts)
|
||||
server.dimensions.append(body.get("dimensions"))
|
||||
payload = json.dumps(
|
||||
{
|
||||
"object": "list",
|
||||
"model": body["model"],
|
||||
"data": [
|
||||
{
|
||||
"object": "embedding",
|
||||
"index": index,
|
||||
"embedding": [(len(text) % 7 + 1) / 8.0] + [0.5] * (width - 1),
|
||||
}
|
||||
for index, text in enumerate(texts)
|
||||
],
|
||||
"usage": {"prompt_tokens": 1, "total_tokens": 1},
|
||||
}
|
||||
).encode()
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", "application/json")
|
||||
self.send_header("Content-Length", str(len(payload)))
|
||||
self.end_headers()
|
||||
self.wfile.write(payload)
|
||||
|
||||
def log_message(self, *_args: object) -> None:
|
||||
return
|
||||
|
||||
self._server = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
|
||||
self.url = f"http://127.0.0.1:{self._server.server_port}/v1"
|
||||
threading.Thread(target=self._server.serve_forever, daemon=True).start()
|
||||
|
||||
def close(self) -> None:
|
||||
self._server.shutdown()
|
||||
self._server.server_close()
|
||||
|
||||
|
||||
def _home(tmp_path: Path, *, embedding_provider, embedding_url: str, api_key: str | None):
|
||||
from pydantic import SecretStr
|
||||
|
||||
from one_memory.embedding_spaces import EmbeddingSpaceStore
|
||||
from one_memory.model_profiles import (
|
||||
ModelCapabilities,
|
||||
ModelProfileInput,
|
||||
ModelProfileStore,
|
||||
ProviderKind,
|
||||
)
|
||||
from one_memory.secrets import InMemorySecretStore
|
||||
|
||||
database = tmp_path / "one-memory.sqlite3"
|
||||
secrets = InMemorySecretStore()
|
||||
profiles = ModelProfileStore(database)
|
||||
spaces = EmbeddingSpaceStore(database)
|
||||
created = profiles.create(
|
||||
ModelProfileInput(
|
||||
name="target",
|
||||
provider_kind=ProviderKind.OPENAI_COMPATIBLE,
|
||||
base_url=embedding_url,
|
||||
model="chat",
|
||||
embedding_provider_kind=embedding_provider,
|
||||
embedding_base_url=embedding_url,
|
||||
embedding_model="stand-in",
|
||||
embedding_dimensions=8,
|
||||
api_key=SecretStr(api_key) if api_key else None,
|
||||
),
|
||||
secrets,
|
||||
)
|
||||
profiles.record_capabilities(
|
||||
created.profile_id,
|
||||
expected_revision=created.revision,
|
||||
capabilities=ModelCapabilities(
|
||||
endpoint_reachable=True,
|
||||
model_available=True,
|
||||
structured_output=True,
|
||||
embedding_available=True,
|
||||
),
|
||||
secrets=secrets,
|
||||
)
|
||||
return database, secrets, profiles, spaces, created
|
||||
|
||||
|
||||
def _service(database, data_dir, profiles, spaces, secrets, **options):
|
||||
from one_memory.embedding_switch import EmbeddingSwitchService
|
||||
|
||||
service = EmbeddingSwitchService(
|
||||
database_path=database,
|
||||
data_dir=data_dir,
|
||||
profiles=profiles,
|
||||
spaces=spaces,
|
||||
secrets=secrets,
|
||||
backend=SdkReembedBackend(data_dir),
|
||||
worker_enabled=False,
|
||||
**options,
|
||||
)
|
||||
service.release_scan_hold()
|
||||
return service
|
||||
|
||||
|
||||
async def _switch(service, created) -> object:
|
||||
plan = await service.plan(created.profile_id)
|
||||
assert plan.needed and not plan.blockers, plan
|
||||
job = await service.start_switch(
|
||||
created.profile_id,
|
||||
expected_revision=created.revision,
|
||||
plan_token=plan.plan_token,
|
||||
requested_by="cli",
|
||||
)
|
||||
await service.run(job.job_id)
|
||||
finished = service.get_job(job.job_id)
|
||||
assert finished.state == "done", (
|
||||
finished.error_code,
|
||||
service.jobs.get(job.job_id)["verification_json"],
|
||||
service.jobs.get(job.job_id)["error_detail"],
|
||||
)
|
||||
return finished
|
||||
|
||||
|
||||
def _seed_vectors(data_dir: Path, space) -> None:
|
||||
import chromadb
|
||||
|
||||
vector = data_dir / "hy-memory" / "governed-v2" / "vector"
|
||||
vector.mkdir(parents=True, exist_ok=True)
|
||||
client = chromadb.PersistentClient(path=str(vector))
|
||||
memory = client.get_or_create_collection(
|
||||
space.collection_names()["memory"], metadata={"hnsw:space": "cosine"}
|
||||
)
|
||||
memory.add(
|
||||
ids=["a", "b"],
|
||||
embeddings=[[0.1] * 512, [0.2] * 512],
|
||||
documents=["one", "two"],
|
||||
metadatas=[{"user_id": "local"}, {"user_id": "hermes-a"}],
|
||||
)
|
||||
|
||||
|
||||
def test_a_switch_embeds_through_the_sdk_service_over_http(tmp_path: Path) -> None:
|
||||
from one_memory.model_profiles import ProviderKind
|
||||
|
||||
server = _EmbeddingServer(8)
|
||||
try:
|
||||
database, secrets, profiles, spaces, created = _home(
|
||||
tmp_path,
|
||||
embedding_provider=ProviderKind.OPENAI_COMPATIBLE,
|
||||
embedding_url=server.url,
|
||||
api_key="sk-test",
|
||||
)
|
||||
_seed_vectors(tmp_path, spaces.bootstrap(None))
|
||||
service = _service(database, tmp_path, profiles, spaces, secrets)
|
||||
|
||||
asyncio.run(_switch(service, created))
|
||||
|
||||
assert {"one", "two"} <= set(server.inputs)
|
||||
assert set(server.dimensions) == {8}
|
||||
finally:
|
||||
server.close()
|
||||
|
||||
|
||||
def test_an_ollama_embedding_needs_no_generation_key_to_re_embed(tmp_path: Path) -> None:
|
||||
"""A switch only embeds: the generation model's key is not its business."""
|
||||
from one_memory.embedding_backend import EmbeddingTarget
|
||||
from one_memory.embedding_spaces import embedding_identity
|
||||
from one_memory.model_profiles import ProviderKind
|
||||
|
||||
server = _EmbeddingServer(8)
|
||||
try:
|
||||
_database, secrets, profiles, spaces, created = _home(
|
||||
tmp_path,
|
||||
embedding_provider=ProviderKind.OLLAMA,
|
||||
embedding_url=server.url,
|
||||
api_key="sk-test",
|
||||
)
|
||||
profile = profiles.snapshot(created.profile_id, created.revision)
|
||||
identity = embedding_identity(profile)
|
||||
target = EmbeddingSpace(
|
||||
space_id="t",
|
||||
state="building",
|
||||
model=identity.model,
|
||||
dims=8,
|
||||
name_dims=8,
|
||||
collection_base=TARGET,
|
||||
naming="v2",
|
||||
created_at="2026-11-01T00:00:00+00:00",
|
||||
)
|
||||
backend = SdkReembedBackend(tmp_path)
|
||||
|
||||
# `api_key=None` is what `_embedding_target` resolves for Ollama.
|
||||
service = backend._embedding_service(
|
||||
EmbeddingTarget(profile=profile, space=target, api_key=None)
|
||||
)
|
||||
|
||||
assert asyncio.run(service.embed_batch(["hello"], use_cache=False))[0][1] == 0.5
|
||||
finally:
|
||||
server.close()
|
||||
|
||||
|
||||
async def _seed_graph(path: Path, dims: int) -> None:
|
||||
from hy_memory.data.graph_store_kuzu import KuzuGraphStore
|
||||
from hy_memory.models.memory import MemoryLayer, MemoryNode
|
||||
|
||||
config = hy_memory.MemoryConfig()
|
||||
config.graph_store.provider = "kuzu"
|
||||
config.graph_store.db_path = str(path)
|
||||
config.vector_store.embedding_dims = dims
|
||||
store = KuzuGraphStore(config)
|
||||
await store.initialize()
|
||||
assert store._available
|
||||
node = MemoryNode(
|
||||
user_id="local",
|
||||
agent_id="a",
|
||||
session_id="s",
|
||||
layer=MemoryLayer.L6_SCHEMA,
|
||||
content="plans ahead",
|
||||
custom={"retrieval_text": "plans trips weeks ahead"},
|
||||
)
|
||||
node._graph_embedding = [0.1] * dims
|
||||
await store.upsert_memory_node(node)
|
||||
await store.close()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", ["ultra", "pro"])
|
||||
def test_a_real_graph_moves_to_the_target_width(tmp_path: Path, mode: str) -> None:
|
||||
"""In any mode: a pro home's old graph must not be lost with the old space."""
|
||||
from one_memory.hy_graph_admin import forget_graph_widths
|
||||
from one_memory.model_profiles import ProviderKind
|
||||
|
||||
server = _EmbeddingServer(8)
|
||||
try:
|
||||
database, secrets, profiles, spaces, created = _home(
|
||||
tmp_path,
|
||||
embedding_provider=ProviderKind.OPENAI_COMPATIBLE,
|
||||
embedding_url=server.url,
|
||||
api_key="sk-test",
|
||||
)
|
||||
live = spaces.bootstrap(None)
|
||||
_seed_vectors(tmp_path, live)
|
||||
root = tmp_path / "hy-memory" / "governed-v2"
|
||||
asyncio.run(_seed_graph(root / "graph", 512))
|
||||
service = _service(
|
||||
database, tmp_path, profiles, spaces, secrets, graph_enabled=lambda: mode == "ultra"
|
||||
)
|
||||
|
||||
asyncio.run(_switch(service, created))
|
||||
|
||||
new = spaces.require_live()
|
||||
assert new.graph_dir != "graph" and (root / new.graph_dir).exists()
|
||||
assert (root / "graph").exists()
|
||||
assert "plans trips weeks ahead" in server.inputs
|
||||
forget_graph_widths()
|
||||
finally:
|
||||
server.close()
|
||||
|
||||
|
||||
def test_verifying_the_graph_closes_what_reads_reopened(tmp_path: Path) -> None:
|
||||
"""Reads are served during a switch, so a client may hold the graph again."""
|
||||
from one_memory.model_profiles import ProviderKind
|
||||
|
||||
server = _EmbeddingServer(8)
|
||||
quiesced: list[str] = []
|
||||
|
||||
async def quiesce():
|
||||
quiesced.append("closed")
|
||||
return lambda: quiesced.append("released")
|
||||
|
||||
try:
|
||||
database, secrets, profiles, spaces, created = _home(
|
||||
tmp_path,
|
||||
embedding_provider=ProviderKind.OPENAI_COMPATIBLE,
|
||||
embedding_url=server.url,
|
||||
api_key="sk-test",
|
||||
)
|
||||
live = spaces.bootstrap(None)
|
||||
_seed_vectors(tmp_path, live)
|
||||
asyncio.run(_seed_graph(tmp_path / "hy-memory" / "governed-v2" / "graph", 512))
|
||||
service = _service(
|
||||
database,
|
||||
tmp_path,
|
||||
profiles,
|
||||
spaces,
|
||||
secrets,
|
||||
graph_enabled=lambda: True,
|
||||
quiesce_clients=quiesce,
|
||||
)
|
||||
|
||||
asyncio.run(_switch(service, created))
|
||||
|
||||
# Once to copy the graph, once to verify it.
|
||||
assert quiesced == ["closed", "released", "closed", "released"]
|
||||
finally:
|
||||
server.close()
|
||||
|
||||
|
||||
def test_the_first_chroma_open_waits_for_a_client_under_construction(tmp_path: Path) -> None:
|
||||
"""Two first opens on one directory race over Chroma's unlocked System cache."""
|
||||
import threading
|
||||
|
||||
from one_memory.hy_sdk_runtime import first_store_open
|
||||
|
||||
_seed(tmp_path)
|
||||
backend = SdkReembedBackend(tmp_path)
|
||||
finished = threading.Event()
|
||||
|
||||
def scan() -> None:
|
||||
asyncio.run(backend.scan())
|
||||
finished.set()
|
||||
|
||||
with first_store_open(tmp_path):
|
||||
thread = threading.Thread(target=scan)
|
||||
thread.start()
|
||||
assert not finished.wait(0.3)
|
||||
thread.join(10)
|
||||
assert finished.is_set()
|
||||
@@ -48,6 +48,7 @@ def test_health_reports_the_single_core() -> None:
|
||||
# No lifespan ran, so recovery has not begun.
|
||||
"recovery": "pending",
|
||||
"recovery_failed_steps": [],
|
||||
"embedding_switch": None,
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -111,6 +111,11 @@ ALLOWED = {
|
||||
"the offline migration opens Chroma",
|
||||
opened_with="?mode=ro",
|
||||
),
|
||||
("sdk_vector_reattach.py", "stored_collection_width"): Allowed(
|
||||
"Chroma's own chroma.sqlite3, read-only, for the width it fixed on a "
|
||||
"collection: a probe on the live space's profile may only record that one",
|
||||
opened_with="?mode=ro",
|
||||
),
|
||||
("imports/sqlite_access.py", "open_sqlite_readonly"): Allowed(
|
||||
"a host app's own SQLite (Cursor state.vscdb / store.db): one direct "
|
||||
"read-only open when there is no WAL, and one read-only open of the "
|
||||
|
||||
@@ -26,6 +26,7 @@ def snapshot() -> ModelProfileSnapshot:
|
||||
base_url="http://127.0.0.1:11434",
|
||||
model="memory-model",
|
||||
embedding_model="embedding-model",
|
||||
embedding_base_url="http://127.0.0.1:11434",
|
||||
embedding_dimensions=768,
|
||||
timeout_seconds=30,
|
||||
max_concurrency=1,
|
||||
|
||||
@@ -0,0 +1,499 @@
|
||||
"""The SDK retrieval and extraction features MemHub lets a person turn on.
|
||||
|
||||
Graph rerank (usage A), episode recall (usage B), the episode write side and
|
||||
the extraction switches are `MemoryConfig` fields. MemHub hard-codes none of
|
||||
them: one global `memory_features` document sets them per client in
|
||||
`build_sdk_config`, an unset one keeps the SDK default (so nothing set is the
|
||||
config of before, byte for byte), and a change rebuilds the handles the way a
|
||||
mode change does. A feature never touches the vector collection identity.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import dataclasses
|
||||
import json
|
||||
from contextlib import asynccontextmanager
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
from httpx import ASGITransport, AsyncClient
|
||||
from pydantic import SecretStr, ValidationError
|
||||
|
||||
from one_memory import hy_sdk_runtime, sdk_write_processor
|
||||
from one_memory.api import create_app
|
||||
from one_memory.config import Settings
|
||||
from one_memory.memory_features import (
|
||||
FEATURE_NAMES,
|
||||
FEATURES,
|
||||
FEATURES_BY_NAME,
|
||||
MemoryFeatures,
|
||||
effective_features,
|
||||
feature_activity,
|
||||
)
|
||||
from one_memory.memory_server.config_state import ConfigStore
|
||||
from one_memory.model_profiles import ModelProfileSnapshot, ProviderKind
|
||||
from one_memory.secrets import InMemorySecretStore
|
||||
|
||||
TOKEN = "memory-features-test"
|
||||
|
||||
#: What MemHub ran with before the setting existed, on a clean environment.
|
||||
TODAY = {
|
||||
"graph_rerank": False,
|
||||
"same_series_timeline": False,
|
||||
"episode_recall": False,
|
||||
"session_episodes": False,
|
||||
"episode_summary": False,
|
||||
"intention": True,
|
||||
"few_shot": False,
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _clean_sdk_env(monkeypatch):
|
||||
"""The SDK's own env vars feed its defaults; tests start without them."""
|
||||
for spec in FEATURES:
|
||||
monkeypatch.delenv(spec.env, raising=False)
|
||||
monkeypatch.delenv("ONE_MEMORY_GRAPH_STORE", raising=False)
|
||||
|
||||
|
||||
def _remote_profile() -> ModelProfileSnapshot:
|
||||
return ModelProfileSnapshot(
|
||||
profile_id="remote",
|
||||
revision=3,
|
||||
provider_kind=ProviderKind.OPENAI_COMPATIBLE,
|
||||
base_url="https://models.example.test/v1",
|
||||
model="chat-model",
|
||||
embedding_provider_kind=ProviderKind.OPENAI_COMPATIBLE,
|
||||
embedding_model="embed-model",
|
||||
embedding_dimensions=1024,
|
||||
timeout_seconds=60,
|
||||
max_concurrency=1,
|
||||
context_window=None,
|
||||
has_api_key=True,
|
||||
# Since #255 the embedding endpoint is its own field, never inherited.
|
||||
embedding_base_url="https://models.example.test/v1",
|
||||
)
|
||||
|
||||
|
||||
PROFILES = (hy_sdk_runtime.builtin_retrieval_snapshot(), _remote_profile())
|
||||
|
||||
|
||||
def _config(tmp_path: Path, profile: ModelProfileSnapshot, features=()):
|
||||
return hy_sdk_runtime.build_sdk_config(
|
||||
profile, api_key="secret-key", data_dir=tmp_path, features=features
|
||||
)
|
||||
|
||||
|
||||
def _field(config, name: str) -> bool:
|
||||
spec = FEATURES_BY_NAME[name]
|
||||
return getattr(getattr(config, spec.section), spec.field)
|
||||
|
||||
|
||||
# --- the catalog and the document ------------------------------------------
|
||||
|
||||
|
||||
def test_the_document_names_exactly_the_catalog() -> None:
|
||||
assert tuple(MemoryFeatures.model_fields) == FEATURE_NAMES
|
||||
assert set(TODAY) == set(FEATURE_NAMES)
|
||||
# Each feature is a MemoryConfig field the SDK has; no env var is written.
|
||||
defaults = hy_sdk_runtime.sdk_feature_defaults()
|
||||
assert set(defaults) == set(FEATURE_NAMES)
|
||||
|
||||
|
||||
def test_the_defaults_are_todays_behaviour() -> None:
|
||||
assert hy_sdk_runtime.sdk_feature_defaults() == TODAY
|
||||
assert MemoryFeatures().overrides() == ()
|
||||
assert effective_features((), TODAY) == TODAY
|
||||
|
||||
|
||||
def test_the_sdk_env_var_still_moves_the_default_of_an_unset_feature(monkeypatch) -> None:
|
||||
monkeypatch.setenv("MEMORY_EPISODE_RECALL_ENABLED", "true")
|
||||
assert hy_sdk_runtime.sdk_feature_defaults()["episode_recall"] is True
|
||||
|
||||
|
||||
def test_unknown_or_non_boolean_values_are_refused() -> None:
|
||||
with pytest.raises(ValidationError):
|
||||
MemoryFeatures.model_validate({"graph_rerank": True, "local_rerank": True})
|
||||
with pytest.raises(ValidationError):
|
||||
MemoryFeatures.model_validate({"graph_rerank": "true"})
|
||||
with pytest.raises(ValidationError):
|
||||
MemoryFeatures.model_validate({"graph_rerank": 1})
|
||||
with pytest.raises(ValueError, match="unknown_memory_feature"):
|
||||
hy_sdk_runtime.build_sdk_config(
|
||||
PROFILES[0], api_key=None, data_dir=Path("/tmp/unused"), features={"nope": True}
|
||||
)
|
||||
|
||||
|
||||
def test_merge_sets_resets_and_leaves_absent_keys_alone() -> None:
|
||||
saved = MemoryFeatures(graph_rerank=True, few_shot=True)
|
||||
update = MemoryFeatures.model_validate({"few_shot": None, "episode_recall": True})
|
||||
merged = saved.merged(update)
|
||||
assert merged.overrides() == (("episode_recall", True), ("graph_rerank", True))
|
||||
assert saved.merged(MemoryFeatures()) == saved
|
||||
|
||||
|
||||
def test_the_setting_persists_and_an_untouched_install_writes_none(tmp_path) -> None:
|
||||
store = ConfigStore(tmp_path)
|
||||
store.set_sync("cursor", True)
|
||||
assert "memory_features" not in json.loads((tmp_path / "config.json").read_text())
|
||||
store.set_memory_features(MemoryFeatures(graph_rerank=True, intention=False))
|
||||
saved = json.loads((tmp_path / "config.json").read_text())
|
||||
# Only what a person set; an absent key is the SDK default.
|
||||
assert saved["memory_features"] == {"graph_rerank": True, "intention": False}
|
||||
assert ConfigStore(tmp_path).memory_features().overrides() == (
|
||||
("graph_rerank", True),
|
||||
("intention", False),
|
||||
)
|
||||
(tmp_path / "config.json").write_text(json.dumps({"memory_features": {"bogus": True}}))
|
||||
with pytest.raises(Exception, match="config_invalid"):
|
||||
ConfigStore(tmp_path)
|
||||
|
||||
|
||||
# --- how settings reach the SDK ----------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize("profile", PROFILES, ids=("builtin", "remote"))
|
||||
def test_each_feature_reaches_the_built_memory_config(tmp_path, profile) -> None:
|
||||
for spec in FEATURES:
|
||||
flipped = not TODAY[spec.name]
|
||||
config = _config(tmp_path, profile, {spec.name: flipped})
|
||||
assert _field(config, spec.name) is flipped, spec.name
|
||||
for other in FEATURES:
|
||||
if other.name != spec.name:
|
||||
assert _field(config, other.name) is TODAY[other.name], (spec.name, other.name)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("profile", PROFILES, ids=("builtin", "remote"))
|
||||
def test_nothing_set_is_the_config_of_before_byte_for_byte(tmp_path, profile) -> None:
|
||||
before = repr(dataclasses.asdict(_config(tmp_path, profile)))
|
||||
assert repr(dataclasses.asdict(_config(tmp_path, profile, ()))) == before
|
||||
assert repr(dataclasses.asdict(_config(tmp_path, profile, {}))) == before
|
||||
# Setting every feature to its default value changes nothing either.
|
||||
assert repr(dataclasses.asdict(_config(tmp_path, profile, TODAY))) == before
|
||||
|
||||
|
||||
@pytest.mark.parametrize("profile", PROFILES, ids=("builtin", "remote"))
|
||||
def test_features_never_change_the_collection_identity(tmp_path, profile) -> None:
|
||||
def identity(config):
|
||||
return (
|
||||
dataclasses.asdict(config.vector_store),
|
||||
config.embedder.model,
|
||||
config.embedder.embedding_dims,
|
||||
config.graph_store.db_path,
|
||||
config.cache.db_path,
|
||||
)
|
||||
|
||||
baseline = identity(_config(tmp_path, profile))
|
||||
for values in (
|
||||
{name: True for name in FEATURE_NAMES},
|
||||
{name: False for name in FEATURE_NAMES},
|
||||
{"graph_rerank": True, "episode_recall": True},
|
||||
):
|
||||
assert identity(_config(tmp_path, profile, values)) == baseline
|
||||
|
||||
|
||||
def test_acquire_client_builds_the_client_with_the_features(tmp_path, monkeypatch) -> None:
|
||||
built: list[tuple[str, bool, bool]] = []
|
||||
|
||||
class Client:
|
||||
def __init__(self, *, config, mode, runtime, embedding_service):
|
||||
built.append(
|
||||
(
|
||||
mode,
|
||||
config.graph_store.graph_rerank_enabled,
|
||||
config.recall.episode_recall_enabled,
|
||||
)
|
||||
)
|
||||
|
||||
def close(self):
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(hy_sdk_runtime, "HyMemoryClient", Client)
|
||||
monkeypatch.setattr(hy_sdk_runtime, "_acquire_runtime", lambda config, key: object())
|
||||
monkeypatch.setattr(hy_sdk_runtime, "_release_runtime", lambda key: None)
|
||||
monkeypatch.setattr(hy_sdk_runtime, "inspect_graph_width_mismatch", lambda config: None)
|
||||
profile = hy_sdk_runtime.builtin_retrieval_snapshot()
|
||||
hy_sdk_runtime.acquire_client(profile, api_key=None, data_dir=tmp_path, mode="pro").close()
|
||||
hy_sdk_runtime.acquire_client(
|
||||
profile,
|
||||
api_key=None,
|
||||
data_dir=tmp_path,
|
||||
mode="ultra",
|
||||
features=(("episode_recall", True), ("graph_rerank", True)),
|
||||
).close()
|
||||
assert built == [("pro", False, False), ("ultra", True, True)]
|
||||
|
||||
|
||||
# --- active in the current mode ---------------------------------------------
|
||||
|
||||
|
||||
def test_a_feature_needing_a_higher_mode_is_reported_inactive() -> None:
|
||||
on = effective_features((("graph_rerank", True), ("episode_recall", True)), TODAY)
|
||||
assert feature_activity("graph_rerank", on, "pro") == "requires_ultra"
|
||||
assert feature_activity("graph_rerank", on, "ultra") is None
|
||||
# Usage B needs no graph: it acts in pro, not in lite (nothing extracted).
|
||||
assert feature_activity("episode_recall", on, "pro") is None
|
||||
assert feature_activity("episode_recall", on, "lite") == "requires_pro"
|
||||
assert feature_activity("intention", TODAY, "lite") == "requires_pro"
|
||||
assert feature_activity("few_shot", TODAY, "ultra") == "disabled"
|
||||
series = effective_features((("same_series_timeline", True),), TODAY)
|
||||
assert feature_activity("same_series_timeline", series, "ultra") == "requires_graph_rerank"
|
||||
series = effective_features((("same_series_timeline", True), ("graph_rerank", True)), TODAY)
|
||||
assert feature_activity("same_series_timeline", series, "pro") == "requires_ultra"
|
||||
assert feature_activity("same_series_timeline", series, "ultra") is None
|
||||
|
||||
|
||||
# --- rebuilding handles -------------------------------------------------------
|
||||
|
||||
|
||||
class _FakeHandle:
|
||||
def __init__(self, mode: str, features=(), profile_id: str = "p") -> None:
|
||||
self.mode = mode
|
||||
self.features = features
|
||||
self.profile_id = profile_id
|
||||
self.closed = False
|
||||
|
||||
async def client(self):
|
||||
return self
|
||||
|
||||
async def close(self):
|
||||
self.closed = True
|
||||
|
||||
|
||||
def test_the_read_handle_rebuilds_its_client_when_the_features_change(
|
||||
monkeypatch, tmp_path
|
||||
) -> None:
|
||||
"""Reads follow the active profile (#255), the mode and the features."""
|
||||
from one_memory import sdk_read_handle
|
||||
|
||||
features = [()]
|
||||
made: list[_FakeHandle] = []
|
||||
|
||||
def handle_for(profile, *, api_key, data_dir, mode=None, features=()):
|
||||
handle = _FakeHandle(mode or "pro", features, profile.profile_id)
|
||||
made.append(handle)
|
||||
return handle
|
||||
|
||||
monkeypatch.setattr(sdk_read_handle, "SdkClientHandle", handle_for)
|
||||
|
||||
class _NoProfile:
|
||||
def active_identity(self):
|
||||
return None
|
||||
|
||||
reads = sdk_read_handle.ActiveProfileReadHandle(
|
||||
profiles=_NoProfile(), # type: ignore[arg-type]
|
||||
secrets=None, # type: ignore[arg-type]
|
||||
data_dir=tmp_path,
|
||||
memory_mode=lambda: "pro",
|
||||
features=lambda: features[0],
|
||||
)
|
||||
hold = asyncio.Event()
|
||||
holding = asyncio.Event()
|
||||
|
||||
async def read():
|
||||
return await reads.client()
|
||||
|
||||
async def slow_read():
|
||||
client = await reads.client()
|
||||
holding.set()
|
||||
await hold.wait()
|
||||
return client
|
||||
|
||||
async def scenario():
|
||||
in_flight = asyncio.create_task(slow_read())
|
||||
await asyncio.wait_for(holding.wait(), 5)
|
||||
old = made[0]
|
||||
assert (await asyncio.create_task(read())) is old
|
||||
features[0] = (("episode_recall", True),)
|
||||
try:
|
||||
fresh = await asyncio.create_task(read())
|
||||
assert fresh.features == (("episode_recall", True),) and fresh.mode == "pro"
|
||||
# A read that took the old client is still using it.
|
||||
assert not old.closed
|
||||
finally:
|
||||
hold.set()
|
||||
await in_flight
|
||||
for _ in range(10):
|
||||
await asyncio.sleep(0)
|
||||
assert old.closed and not fresh.closed
|
||||
await reads.close()
|
||||
assert fresh.closed
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
class _Profiles:
|
||||
def active_identity(self):
|
||||
return ("profile-a", 1)
|
||||
|
||||
def snapshot(self, profile_id, revision):
|
||||
return SimpleNamespace(profile_id=profile_id, has_api_key=False)
|
||||
|
||||
def resolve_snapshot_secret(self, profile_id, revision, secrets):
|
||||
return None
|
||||
|
||||
|
||||
def test_a_feature_change_keeps_a_busy_write_handle_until_its_write_ends(monkeypatch) -> None:
|
||||
features = [()]
|
||||
opened: list[_FakeHandle] = []
|
||||
|
||||
def handle(profile, **kwargs):
|
||||
value = _FakeHandle(kwargs["mode"], kwargs["features"], profile.profile_id)
|
||||
opened.append(value)
|
||||
return value
|
||||
|
||||
monkeypatch.setattr(sdk_write_processor, "SdkClientHandle", handle)
|
||||
coordinator = sdk_write_processor.SdkWriteCoordinator(
|
||||
profiles=_Profiles(), # type: ignore[arg-type]
|
||||
secrets=None, # type: ignore[arg-type]
|
||||
data_dir=None,
|
||||
transport_factory=lambda _profile: None,
|
||||
sedimentation_enabled=lambda: True,
|
||||
memory_mode=lambda: "pro",
|
||||
memory_features=lambda: features[0],
|
||||
)
|
||||
|
||||
async def scenario():
|
||||
async with coordinator.use_client() as client:
|
||||
assert client.features == ()
|
||||
features[0] = (("few_shot", True),)
|
||||
await coordinator.retire_stale_handles()
|
||||
# In flight: it finishes with the features it started with (#218).
|
||||
assert not opened[0].closed
|
||||
assert opened[0].closed
|
||||
async with coordinator.use_client() as client:
|
||||
assert client.features == (("few_shot", True),)
|
||||
features[0] = ()
|
||||
await coordinator.retire_stale_handles()
|
||||
assert opened[1].closed
|
||||
await coordinator.close()
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
# --- the HTTP surface -----------------------------------------------------------
|
||||
|
||||
|
||||
def application(tmp_path: Path):
|
||||
settings = Settings(
|
||||
data_dir=tmp_path / "data",
|
||||
source_home=tmp_path,
|
||||
api_token=SecretStr(TOKEN),
|
||||
imports_enabled=False,
|
||||
automatic_memory_processing=False,
|
||||
live_sync_enabled=False,
|
||||
)
|
||||
return create_app(settings, secret_store=InMemorySecretStore())
|
||||
|
||||
|
||||
def run_app(tmp_path, monkeypatch, scenario) -> None:
|
||||
async def main() -> None:
|
||||
app = application(tmp_path)
|
||||
|
||||
@asynccontextmanager
|
||||
async def write_client():
|
||||
yield SimpleNamespace()
|
||||
|
||||
async with app.router.lifespan_context(app):
|
||||
monkeypatch.setattr(app.state.sdk_write_coordinator, "use_client", write_client)
|
||||
async with AsyncClient(
|
||||
transport=ASGITransport(app=app),
|
||||
base_url="http://test",
|
||||
headers={"Authorization": f"Bearer {TOKEN}"},
|
||||
) as client:
|
||||
await scenario(app, client)
|
||||
|
||||
asyncio.run(main())
|
||||
|
||||
|
||||
def _by_name(body) -> dict[str, dict]:
|
||||
return {item["name"]: item for item in body["features"]}
|
||||
|
||||
|
||||
@pytest.mark.fresh_sdk
|
||||
def test_features_api_get_put_reset_and_status(tmp_path, monkeypatch) -> None:
|
||||
async def scenario(app, client):
|
||||
current = (await client.get("/api/v1/settings/memory-features")).json()
|
||||
assert current["mode"] == "pro" and current["non_default"] == {}
|
||||
listed = _by_name(current)
|
||||
assert list(listed) == list(FEATURE_NAMES)
|
||||
assert {name: item["enabled"] for name, item in listed.items()} == TODAY
|
||||
assert all(item["configured"] is None for item in listed.values())
|
||||
rerank = listed["graph_rerank"]
|
||||
assert rerank["default"] is False and rerank["active"] is False
|
||||
assert rerank["reason"] == "disabled" and rerank["requires_mode"] == "ultra"
|
||||
assert rerank["sdk_field"] == "graph_store.graph_rerank_enabled"
|
||||
assert rerank["env"] == "MEMORY_GRAPH_RERANK_ENABLED"
|
||||
assert listed["intention"]["active"] is True
|
||||
assert listed["episode_summary"]["cost"] == "llm_call"
|
||||
reads = app.state.sdk_client_handle
|
||||
# The read handle is keyed by (profile, mode, features): what it
|
||||
# builds next is what these return.
|
||||
assert reads.current_features() == ()
|
||||
|
||||
changed = await client.put(
|
||||
"/api/v1/settings/memory-features",
|
||||
json={"graph_rerank": True, "episode_recall": True},
|
||||
)
|
||||
assert changed.status_code == 200, changed.text
|
||||
body = changed.json()
|
||||
assert body["changed"] == ["graph_rerank", "episode_recall"]
|
||||
listed = _by_name(body)
|
||||
# Enabled but the graph store is ultra only: reported, not hidden.
|
||||
assert listed["graph_rerank"]["enabled"] is True
|
||||
assert listed["graph_rerank"]["active"] is False
|
||||
assert listed["graph_rerank"]["reason"] == "requires_ultra"
|
||||
assert listed["episode_recall"]["active"] is True
|
||||
assert body["non_default"] == {"graph_rerank": True, "episode_recall": True}
|
||||
# The handles were rebuilt with the new features.
|
||||
expected = (("episode_recall", True), ("graph_rerank", True))
|
||||
assert app.state.sdk_write_coordinator.memory_features() == expected
|
||||
assert reads.current_features() == expected
|
||||
assert reads.memory_mode() == "pro"
|
||||
|
||||
processing = (await client.get("/api/v1/memory-processing")).json()
|
||||
assert processing["pipeline"]["features"] == {
|
||||
"graph_rerank": True,
|
||||
"episode_recall": True,
|
||||
}
|
||||
|
||||
await client.put("/api/v1/settings/memory-mode", json={"mode": "ultra"})
|
||||
await app.state.memory_graph_backfill.wait()
|
||||
ultra = _by_name((await client.get("/api/v1/settings/memory-features")).json())
|
||||
assert ultra["graph_rerank"]["active"] is True
|
||||
assert ultra["graph_rerank"]["reason"] is None
|
||||
|
||||
again = (
|
||||
await client.put("/api/v1/settings/memory-features", json={"graph_rerank": True})
|
||||
).json()
|
||||
assert again["changed"] == []
|
||||
assert reads.current_features() == expected
|
||||
|
||||
reset = (
|
||||
await client.put(
|
||||
"/api/v1/settings/memory-features",
|
||||
json={"graph_rerank": None, "episode_recall": None},
|
||||
)
|
||||
).json()
|
||||
assert reset["changed"] == ["graph_rerank", "episode_recall"]
|
||||
assert reset["non_default"] == {}
|
||||
assert reads.current_features() == ()
|
||||
saved = json.loads((tmp_path / "data" / "config.json").read_text())
|
||||
assert saved["memory_features"] == {}
|
||||
|
||||
for bad in (
|
||||
{"local_rerank": True},
|
||||
{"graph_rerank": "true"},
|
||||
{"graph_rerank": 1},
|
||||
{"features": {"graph_rerank": True}},
|
||||
):
|
||||
response = await client.put("/api/v1/settings/memory-features", json=bad)
|
||||
assert response.status_code == 422, bad
|
||||
assert (
|
||||
await client.get("/api/v1/settings/memory-features", params={"user_id": "x"})
|
||||
).status_code == 422
|
||||
assert (await client.get("/api/v1/settings/memory-features")).json()["non_default"] == {}
|
||||
|
||||
run_app(tmp_path, monkeypatch, scenario)
|
||||
@@ -318,6 +318,12 @@ def run_with_sdk(tmp_path, monkeypatch, scenario):
|
||||
async with app.router.lifespan_context(app):
|
||||
app.state.sdk_client_handle = Handle()
|
||||
monkeypatch.setattr(app.state.sdk_write_coordinator, "use_client", write_client)
|
||||
# Edits and deletes wait for the startup scan of embedding spaces,
|
||||
# as they would behind the App's loading screen.
|
||||
for _ in range(500):
|
||||
if app.state.embedding_switch.writes_open():
|
||||
break
|
||||
await asyncio.sleep(0.01)
|
||||
async with client_for(app) as client:
|
||||
await scenario(app, sdk, client)
|
||||
|
||||
|
||||
@@ -0,0 +1,867 @@
|
||||
"""The memory mode (lite / pro / ultra), evolve and its nightly schedule.
|
||||
|
||||
Item 5: the mode is a saved setting (default pro, `ONE_MEMORY_GRAPH_STORE`
|
||||
still forces ultra), it decides the mode every SDK client is built with, and
|
||||
a change rebuilds the handles. lite needs no generation model, never calls
|
||||
`add_extracted`, and keeps every raw. Item 6 (P1-3 / R7): `POST
|
||||
/api/v2/evolve` is ultra only, one run per user at a time, never beside that
|
||||
user's write, and loops while the SDK says `should_continue`; the schedule
|
||||
runs nightly and catches up after sleep. R10: the status carries
|
||||
`pipeline.mode` and the last evolve.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from contextlib import asynccontextmanager
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
from httpx import ASGITransport, AsyncClient
|
||||
from pydantic import SecretStr
|
||||
|
||||
from one_memory import hy_sdk_runtime, sdk_write_processor
|
||||
from one_memory.api import _sdk_client_handle, create_app, memory_mode_effects
|
||||
from one_memory.config import Settings
|
||||
from one_memory.hy_graph_admin import effective_memory_mode
|
||||
from one_memory.imports.models import (
|
||||
EvidenceContentKind,
|
||||
EvidenceInput,
|
||||
EvidenceOrigin,
|
||||
OriginMetadata,
|
||||
WrittenBy,
|
||||
materialize_write_input,
|
||||
)
|
||||
from one_memory.memory_evolution import (
|
||||
EvolveModeError,
|
||||
EvolveScheduler,
|
||||
EvolveService,
|
||||
EvolveStateStore,
|
||||
MemoryGraphBackfill,
|
||||
evolve_due,
|
||||
)
|
||||
from one_memory.memory_server.config_state import ConfigStore, EvolveSchedule
|
||||
from one_memory.sdk_memory_runtime import SdkMemoryRuntime
|
||||
from one_memory.secrets import InMemorySecretStore
|
||||
|
||||
TOKEN = "memory-mode-test"
|
||||
|
||||
|
||||
# --- the setting and the client mode -----------------------------------------
|
||||
|
||||
|
||||
def test_mode_defaults_to_pro_and_the_env_still_forces_ultra(tmp_path, monkeypatch) -> None:
|
||||
monkeypatch.delenv("ONE_MEMORY_GRAPH_STORE", raising=False)
|
||||
store = ConfigStore(tmp_path)
|
||||
assert store.memory_mode() == "pro"
|
||||
assert effective_memory_mode(store.memory_mode()) == "pro"
|
||||
store.set_memory_mode("lite")
|
||||
assert ConfigStore(tmp_path).memory_mode() == "lite"
|
||||
assert json.loads((tmp_path / "config.json").read_text())["memory_mode"] == "lite"
|
||||
monkeypatch.setenv("ONE_MEMORY_GRAPH_STORE", "kuzu")
|
||||
assert effective_memory_mode("lite") == "ultra"
|
||||
monkeypatch.setenv("ONE_MEMORY_GRAPH_STORE", "none")
|
||||
assert effective_memory_mode("lite") == "lite"
|
||||
with pytest.raises(ValueError):
|
||||
store.set_memory_mode("max") # type: ignore[arg-type]
|
||||
|
||||
|
||||
def test_an_untouched_install_writes_no_new_settings(tmp_path) -> None:
|
||||
store = ConfigStore(tmp_path)
|
||||
store.set_sync("cursor", True)
|
||||
saved = json.loads((tmp_path / "config.json").read_text())
|
||||
assert "memory_mode" not in saved and "evolve_schedule" not in saved
|
||||
|
||||
|
||||
def test_acquire_client_builds_the_sdk_client_in_the_given_mode(tmp_path, monkeypatch) -> None:
|
||||
built: list[str] = []
|
||||
|
||||
class Client:
|
||||
def __init__(self, *, config, mode, runtime, embedding_service):
|
||||
built.append(mode)
|
||||
|
||||
def close(self):
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(hy_sdk_runtime, "HyMemoryClient", Client)
|
||||
monkeypatch.setattr(hy_sdk_runtime, "_acquire_runtime", lambda config, key: object())
|
||||
monkeypatch.setattr(hy_sdk_runtime, "_release_runtime", lambda key: None)
|
||||
monkeypatch.setattr(hy_sdk_runtime, "inspect_graph_width_mismatch", lambda config: None)
|
||||
monkeypatch.delenv("ONE_MEMORY_GRAPH_STORE", raising=False)
|
||||
profile = hy_sdk_runtime.builtin_retrieval_snapshot()
|
||||
for mode in ("lite", "pro", "ultra"):
|
||||
hy_sdk_runtime.acquire_client(profile, api_key=None, data_dir=tmp_path, mode=mode).close()
|
||||
hy_sdk_runtime.acquire_client(profile, api_key=None, data_dir=tmp_path).close()
|
||||
monkeypatch.setenv("ONE_MEMORY_GRAPH_STORE", "kuzu")
|
||||
hy_sdk_runtime.acquire_client(profile, api_key=None, data_dir=tmp_path).close()
|
||||
assert built == ["lite", "pro", "ultra", "pro", "ultra"]
|
||||
|
||||
|
||||
def test_the_read_handle_rebuilds_on_a_mode_or_profile_change(monkeypatch, tmp_path) -> None:
|
||||
"""Reads follow the active profile's embedding (#255) and the memory mode."""
|
||||
from one_memory import sdk_read_handle
|
||||
|
||||
modes = ["pro"]
|
||||
made: list[_FakeHandle] = []
|
||||
|
||||
def handle_for(profile, *, api_key, data_dir, mode=None, features=()):
|
||||
handle = _FakeHandle(mode or "pro", profile.profile_id)
|
||||
made.append(handle)
|
||||
return handle
|
||||
|
||||
monkeypatch.setattr(sdk_read_handle, "SdkClientHandle", handle_for)
|
||||
profiles = _Profiles(identity=None)
|
||||
reads = sdk_read_handle.ActiveProfileReadHandle(
|
||||
profiles=profiles, # type: ignore[arg-type]
|
||||
secrets=None, # type: ignore[arg-type]
|
||||
data_dir=tmp_path,
|
||||
memory_mode=lambda: modes[0],
|
||||
)
|
||||
|
||||
async def read():
|
||||
return await reads.client()
|
||||
|
||||
async def scenario():
|
||||
first = await asyncio.create_task(read())
|
||||
assert (first.mode, first.profile_id) == ("pro", "builtin-retrieval")
|
||||
modes[0] = "lite"
|
||||
# A read still in flight on the old client keeps it open.
|
||||
hold = asyncio.Event()
|
||||
holding_client = asyncio.Event()
|
||||
|
||||
async def holding():
|
||||
await reads.client()
|
||||
holding_client.set()
|
||||
await hold.wait()
|
||||
|
||||
modes[0] = "pro"
|
||||
holder = asyncio.create_task(holding())
|
||||
await asyncio.wait_for(holding_client.wait(), 5)
|
||||
modes[0] = "lite"
|
||||
try:
|
||||
assert (await asyncio.create_task(read())).mode == "lite"
|
||||
assert not made[0].closed
|
||||
finally:
|
||||
hold.set()
|
||||
await holder
|
||||
for _ in range(10):
|
||||
await asyncio.sleep(0)
|
||||
assert made[0].closed
|
||||
assert [item.mode for item in made] == ["pro", "lite"]
|
||||
await reads.close()
|
||||
assert made[1].closed
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
class _FakeHandle:
|
||||
def __init__(self, mode: str, profile_id: str = "p") -> None:
|
||||
self.mode = mode
|
||||
self.profile_id = profile_id
|
||||
self.closed = False
|
||||
|
||||
async def client(self):
|
||||
return self
|
||||
|
||||
async def close(self):
|
||||
self.closed = True
|
||||
|
||||
|
||||
class _Profiles:
|
||||
def __init__(self, identity=("profile-a", 1)) -> None:
|
||||
self.identity = identity
|
||||
|
||||
def active_identity(self):
|
||||
return self.identity
|
||||
|
||||
def snapshot(self, profile_id, revision):
|
||||
return SimpleNamespace(profile_id=profile_id, has_api_key=False)
|
||||
|
||||
def resolve_snapshot_secret(self, profile_id, revision, secrets):
|
||||
return None
|
||||
|
||||
|
||||
def _coordinator(monkeypatch, modes, profiles=None):
|
||||
opened: list[_FakeHandle] = []
|
||||
|
||||
def handle(profile, **kwargs):
|
||||
value = _FakeHandle(kwargs.get("mode") or "pro", profile.profile_id)
|
||||
opened.append(value)
|
||||
return value
|
||||
|
||||
monkeypatch.setattr(sdk_write_processor, "SdkClientHandle", handle)
|
||||
coordinator = sdk_write_processor.SdkWriteCoordinator(
|
||||
profiles=profiles or _Profiles(), # type: ignore[arg-type]
|
||||
secrets=None, # type: ignore[arg-type]
|
||||
data_dir=None,
|
||||
transport_factory=lambda _profile: None,
|
||||
sedimentation_enabled=lambda: True,
|
||||
memory_mode=lambda: modes[0],
|
||||
)
|
||||
return coordinator, opened
|
||||
|
||||
|
||||
def test_a_mode_change_retires_idle_write_handles_and_keeps_a_busy_one(monkeypatch) -> None:
|
||||
modes = ["pro"]
|
||||
coordinator, opened = _coordinator(monkeypatch, modes)
|
||||
|
||||
async def scenario():
|
||||
async with coordinator.use_client() as client:
|
||||
assert client.mode == "pro"
|
||||
modes[0] = "ultra"
|
||||
await coordinator.retire_stale_modes()
|
||||
# Still in use: closed only when this use ends (#218).
|
||||
assert not opened[0].closed
|
||||
assert opened[0].closed
|
||||
async with coordinator.use_client() as client:
|
||||
assert client.mode == "ultra"
|
||||
modes[0] = "lite"
|
||||
await coordinator.retire_stale_modes()
|
||||
assert opened[1].closed
|
||||
await coordinator.close()
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_lite_needs_no_generation_model(monkeypatch) -> None:
|
||||
modes = ["lite"]
|
||||
coordinator, opened = _coordinator(monkeypatch, modes, _Profiles(identity=None))
|
||||
|
||||
async def scenario():
|
||||
async with coordinator.use_client() as client:
|
||||
assert client.mode == "lite"
|
||||
assert client.profile_id == hy_sdk_runtime.BUILTIN_RETRIEVAL_PROFILE_ID
|
||||
modes[0] = "pro"
|
||||
with pytest.raises(RuntimeError, match="model_profile_not_active"):
|
||||
async with coordinator.use_client():
|
||||
pass
|
||||
await coordinator.close()
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def _record(text: str, *, extracted: bool, user_id: str = "local", role: str = "user"):
|
||||
record = materialize_write_input(
|
||||
EvidenceInput(
|
||||
user_id=user_id,
|
||||
source_id="one-memory-manual" if extracted else "hermes",
|
||||
conversation_id="session-1",
|
||||
message_id=str(uuid4()),
|
||||
role=role,
|
||||
written_by=WrittenBy.PERSON,
|
||||
created_at=datetime(2026, 10, 1, tzinfo=UTC),
|
||||
verbatim=text,
|
||||
origin_meta={EvidenceOrigin.FILE: OriginMetadata()},
|
||||
)
|
||||
)
|
||||
if extracted:
|
||||
record = record.model_copy(update={"content_kind": EvidenceContentKind.MEMORY_CONTENT})
|
||||
return record
|
||||
|
||||
|
||||
class _WritingSdk:
|
||||
def __init__(self) -> None:
|
||||
self.calls: list[tuple[str, object]] = []
|
||||
|
||||
def add(self, messages, **kwargs):
|
||||
self.calls.append(("add", messages))
|
||||
return {"success": True, "raw_ids": ["raw-1"]}
|
||||
|
||||
def add_extracted(self, contents, **kwargs):
|
||||
self.calls.append(("add_extracted", contents))
|
||||
return {"success": True, "memory_ids": ["m-1"]}
|
||||
|
||||
|
||||
class _SdkHandle:
|
||||
def __init__(self, sdk, mode="pro") -> None:
|
||||
self.sdk = sdk
|
||||
self.mode = mode
|
||||
|
||||
async def client(self):
|
||||
return self.sdk
|
||||
|
||||
async def close(self):
|
||||
return None
|
||||
|
||||
|
||||
async def _noop() -> None:
|
||||
return None
|
||||
|
||||
|
||||
def test_a_note_in_lite_is_stored_as_raw_and_never_calls_add_extracted() -> None:
|
||||
sdk = _WritingSdk()
|
||||
note = (_record("I prefer green tea", extracted=True),)
|
||||
|
||||
async def scenario():
|
||||
await sdk_write_processor.SdkWriteProcessor(_SdkHandle(sdk), lite=True).process_sdk_job(
|
||||
note, request_id="r-lite", before_hy_invocation=_noop
|
||||
)
|
||||
await sdk_write_processor.SdkWriteProcessor(_SdkHandle(sdk)).process_sdk_job(
|
||||
note, request_id="r-pro", before_hy_invocation=_noop
|
||||
)
|
||||
|
||||
asyncio.run(scenario())
|
||||
assert sdk.calls == [
|
||||
("add", [{"role": "user", "content": "I prefer green tea"}]),
|
||||
("add_extracted", ["I prefer green tea"]),
|
||||
]
|
||||
|
||||
|
||||
def test_a_lite_write_runs_through_the_coordinator_without_a_model(monkeypatch) -> None:
|
||||
sdk = _WritingSdk()
|
||||
coordinator = sdk_write_processor.SdkWriteCoordinator(
|
||||
profiles=_Profiles(identity=None), # type: ignore[arg-type]
|
||||
secrets=None, # type: ignore[arg-type]
|
||||
data_dir=None,
|
||||
transport_factory=lambda _profile: None,
|
||||
sedimentation_enabled=lambda: True,
|
||||
memory_mode=lambda: "lite",
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
sdk_write_processor, "SdkClientHandle", lambda profile, **kwargs: _SdkHandle(sdk, "lite")
|
||||
)
|
||||
turn = (
|
||||
_record("hello", extracted=False),
|
||||
_record("hi there", extracted=False, role="assistant"),
|
||||
)
|
||||
receipt = asyncio.run(
|
||||
coordinator.process_sdk_job(turn, request_id="r-1", before_hy_invocation=_noop)
|
||||
)
|
||||
assert receipt.raw_ids == ("raw-1",)
|
||||
assert [name for name, _ in sdk.calls] == ["add"]
|
||||
|
||||
|
||||
def test_no_raw_is_ever_deleted_in_any_mode() -> None:
|
||||
"""Owner decision: MemHub keeps every raw; there is no retention window."""
|
||||
from one_memory import sdk_memory_runtime
|
||||
|
||||
for name in ("prune_expired_raw", "start_raw_retention", "_raw_retention_loop"):
|
||||
assert not hasattr(SdkMemoryRuntime, name), name
|
||||
assert not hasattr(sdk_memory_runtime, "RAW_RETENTION_DAYS")
|
||||
|
||||
|
||||
def test_mode_switch_effects() -> None:
|
||||
assert memory_mode_effects("lite", "pro") == (
|
||||
"earlier_raw_not_extracted",
|
||||
"lite_search_stops",
|
||||
)
|
||||
assert memory_mode_effects("ultra", "pro") == ("higher_mode_memories_hidden", "evolve_stopped")
|
||||
assert memory_mode_effects("pro", "ultra", graph_backfill=True) == (
|
||||
"existing_memories_added_to_graph",
|
||||
)
|
||||
assert memory_mode_effects("pro", "ultra", graph_backfill=False) == (
|
||||
"graph_backfill_already_running",
|
||||
)
|
||||
assert memory_mode_effects("pro", "lite") == ("higher_mode_memories_hidden",)
|
||||
assert memory_mode_effects("pro", "pro") == ()
|
||||
|
||||
|
||||
# --- evolve ---------------------------------------------------------------------
|
||||
|
||||
|
||||
class _EvolveSdk:
|
||||
"""`evolve` answers `should_continue` for the first rounds, then stops."""
|
||||
|
||||
def __init__(self, continuing: int = 2, *, gate: asyncio.Event | None = None) -> None:
|
||||
self.calls: list[dict] = []
|
||||
self.continuing = continuing
|
||||
self.gate = gate
|
||||
self.graph_builds: list[str] = []
|
||||
self.active = 0
|
||||
self.peak = 0
|
||||
|
||||
def evolve(self, user_id, agent_id="default_agent", max_batches=None):
|
||||
self.active += 1
|
||||
self.peak = max(self.peak, self.active)
|
||||
try:
|
||||
self.calls.append(
|
||||
{"user_id": user_id, "agent_id": agent_id, "max_batches": max_batches}
|
||||
)
|
||||
more = len([call for call in self.calls if call["user_id"] == user_id]) <= (
|
||||
self.continuing
|
||||
)
|
||||
return {
|
||||
"success": True,
|
||||
"request_id": f"evolve-{len(self.calls)}",
|
||||
"tasks_processed": 2,
|
||||
"results": {"schemas": ["private text"]},
|
||||
"cross_domain_sweeper": {"merged": 1},
|
||||
"should_continue": more,
|
||||
"continuation_reason": "remaining_clusters" if more else None,
|
||||
"remaining_clusters": 3 if more else 0,
|
||||
"remaining_actionable_clusters": 1 if more else 0,
|
||||
"remaining_fresh_facts": 0,
|
||||
"cycle_status": "partial" if more else "complete",
|
||||
"batch_failures": [],
|
||||
"elapsed_ms": 5.0,
|
||||
}
|
||||
finally:
|
||||
self.active -= 1
|
||||
|
||||
def build_memory_graph(self, user_id, *, agent_id=None, limit=10000, clear_first=False):
|
||||
self.graph_builds.append(user_id)
|
||||
return {"success": True, "memories_scanned": 4, "nodes": 4}
|
||||
|
||||
|
||||
class _Coordinator:
|
||||
"""The real per-user exclusivity over a stand-in client."""
|
||||
|
||||
def __init__(self, sdk, monkeypatch) -> None:
|
||||
self.inner, _ = _coordinator(monkeypatch, ["ultra"])
|
||||
self.sdk = sdk
|
||||
|
||||
def exclusive_for_evolve(self, user_id):
|
||||
return self.inner.exclusive_for_evolve(user_id)
|
||||
|
||||
@asynccontextmanager
|
||||
async def use_client(self):
|
||||
yield self.sdk
|
||||
|
||||
|
||||
def test_evolve_refuses_outside_ultra(tmp_path, monkeypatch) -> None:
|
||||
service = EvolveService(
|
||||
coordinator=_Coordinator(_EvolveSdk(), monkeypatch),
|
||||
state=EvolveStateStore(tmp_path),
|
||||
memory_mode=lambda: "pro",
|
||||
)
|
||||
|
||||
async def scenario():
|
||||
with pytest.raises(EvolveModeError):
|
||||
service.start("local")
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_evolve_loops_while_should_continue_and_records_counts_only(tmp_path, monkeypatch) -> None:
|
||||
sdk = _EvolveSdk(continuing=2)
|
||||
state = EvolveStateStore(tmp_path)
|
||||
service = EvolveService(
|
||||
coordinator=_Coordinator(sdk, monkeypatch), state=state, memory_mode=lambda: "ultra"
|
||||
)
|
||||
|
||||
async def scenario():
|
||||
started = service.start("wife", max_batches=2)
|
||||
assert started.state == "running"
|
||||
done = await service.wait(started.run_id)
|
||||
assert done is not None
|
||||
return done
|
||||
|
||||
done = asyncio.run(scenario())
|
||||
assert done.state == "succeeded" and done.rounds == 3 and done.tasks_processed == 6
|
||||
assert done.cycle_status == "complete" and done.should_continue is False
|
||||
assert {call["agent_id"] for call in sdk.calls} == {"all-apps"}
|
||||
assert {call["max_batches"] for call in sdk.calls} == {2}
|
||||
saved = (tmp_path / "evolve-state.json").read_text()
|
||||
assert "private text" not in saved
|
||||
assert json.loads(saved)["users"]["wife"]["last_result"]["rounds"] == 3
|
||||
|
||||
|
||||
def test_evolve_stops_at_the_round_cap(tmp_path, monkeypatch) -> None:
|
||||
sdk = _EvolveSdk(continuing=100)
|
||||
service = EvolveService(
|
||||
coordinator=_Coordinator(sdk, monkeypatch),
|
||||
state=EvolveStateStore(tmp_path),
|
||||
memory_mode=lambda: "ultra",
|
||||
max_rounds=4,
|
||||
)
|
||||
|
||||
async def scenario():
|
||||
return await service.wait(service.start("local").run_id)
|
||||
|
||||
done = asyncio.run(scenario())
|
||||
assert done is not None and done.rounds == 4 and done.capped is True
|
||||
assert done.state == "succeeded"
|
||||
|
||||
|
||||
def test_one_run_per_user_and_never_beside_that_users_write(tmp_path, monkeypatch) -> None:
|
||||
sdk = _EvolveSdk(continuing=0)
|
||||
coordinator = _Coordinator(sdk, monkeypatch)
|
||||
service = EvolveService(
|
||||
coordinator=coordinator, state=EvolveStateStore(tmp_path), memory_mode=lambda: "ultra"
|
||||
)
|
||||
writing = asyncio.Event()
|
||||
release = asyncio.Event()
|
||||
|
||||
async def fake_write(self, records, *, request_id, before_hy_invocation):
|
||||
writing.set()
|
||||
await release.wait()
|
||||
return sdk_write_processor.SdkWriteReceipt()
|
||||
|
||||
monkeypatch.setattr(sdk_write_processor.SdkWriteProcessor, "process_sdk_job", fake_write)
|
||||
record = (_record("hello", extracted=False, user_id="wife"),)
|
||||
|
||||
async def scenario():
|
||||
inner = coordinator.inner
|
||||
write = asyncio.create_task(
|
||||
inner.process_sdk_job(record, request_id="w", before_hy_invocation=_noop)
|
||||
)
|
||||
await writing.wait()
|
||||
first = service.start("wife")
|
||||
again = service.start("wife")
|
||||
assert again.run_id == first.run_id
|
||||
other = service.start("local")
|
||||
await service.wait(other.run_id)
|
||||
# The other user is not held; wife's run waits for wife's write.
|
||||
assert sdk.calls and {call["user_id"] for call in sdk.calls} == {"local"}
|
||||
assert inner.evolving("wife")
|
||||
with pytest.raises(RuntimeError, match="memory_evolve_in_progress"):
|
||||
await inner.process_sdk_job(record, request_id="w2", before_hy_invocation=_noop)
|
||||
release.set()
|
||||
await write
|
||||
done = await service.wait(first.run_id)
|
||||
assert done is not None and done.state == "succeeded"
|
||||
assert not inner.evolving("wife")
|
||||
assert sdk.peak == 1
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_a_failed_round_records_its_code(tmp_path, monkeypatch) -> None:
|
||||
class Failing(_EvolveSdk):
|
||||
def evolve(self, user_id, agent_id="default_agent", max_batches=None):
|
||||
return {
|
||||
"success": False,
|
||||
"error_code": "schema_llm_failed",
|
||||
"failure_kind": "llm_error",
|
||||
"error_message": "the provider said something private",
|
||||
}
|
||||
|
||||
state = EvolveStateStore(tmp_path)
|
||||
service = EvolveService(
|
||||
coordinator=_Coordinator(Failing(), monkeypatch), state=state, memory_mode=lambda: "ultra"
|
||||
)
|
||||
done = asyncio.run(service.wait(_start(service, "local")))
|
||||
assert done is not None and done.state == "failed"
|
||||
assert done.error_code == "schema_llm_failed" and done.failure_kind == "llm_error"
|
||||
assert state.user("local").last_error == "schema_llm_failed"
|
||||
assert "private" not in (tmp_path / "evolve-state.json").read_text()
|
||||
|
||||
|
||||
def _start(service: EvolveService, user_id: str) -> str:
|
||||
async def begin() -> str:
|
||||
return service.start(user_id).run_id
|
||||
|
||||
loop = asyncio.new_event_loop()
|
||||
try:
|
||||
run_id = loop.run_until_complete(begin())
|
||||
loop.run_until_complete(service.wait(run_id))
|
||||
finally:
|
||||
loop.close()
|
||||
return run_id
|
||||
|
||||
|
||||
# --- the nightly schedule -------------------------------------------------------
|
||||
|
||||
LOCAL = datetime.now().astimezone().tzinfo
|
||||
|
||||
|
||||
def _at(day: int, hour: int, minute: int = 0) -> datetime:
|
||||
return datetime(2026, 11, day, hour, minute, tzinfo=LOCAL)
|
||||
|
||||
|
||||
def test_evolve_is_due_after_the_slot_once_a_day_and_catches_up_after_sleep() -> None:
|
||||
schedule = EvolveSchedule()
|
||||
anchor = _at(1, 12)
|
||||
# Never ran: not at noon on the day ultra was enabled, then at 03:00 next day.
|
||||
assert not evolve_due(_at(1, 13), schedule, None, anchor)
|
||||
assert not evolve_due(_at(2, 2, 59), schedule, None, anchor)
|
||||
assert evolve_due(_at(2, 3, 0), schedule, None, anchor)
|
||||
# Ran at 03:00: not again that day, due again the next night.
|
||||
assert not evolve_due(_at(2, 22), schedule, _at(2, 3, 1), anchor)
|
||||
assert evolve_due(_at(3, 3, 5), schedule, _at(2, 3, 1), anchor)
|
||||
# Asleep at 03:00: a run older than the slot is due on wake, the same day.
|
||||
assert evolve_due(_at(3, 9), schedule, _at(2, 3, 1), anchor)
|
||||
# Last run more than 24h ago is due even before today's slot.
|
||||
assert evolve_due(_at(4, 1), schedule, _at(2, 3, 1), anchor)
|
||||
assert not evolve_due(_at(3, 1), schedule, _at(2, 3, 1), anchor)
|
||||
# Moved and turned off.
|
||||
assert evolve_due(_at(2, 23, 30), EvolveSchedule(time="23:30"), _at(2, 3), anchor)
|
||||
assert not evolve_due(_at(3, 9), EvolveSchedule(enabled=False), _at(1, 3), anchor)
|
||||
|
||||
|
||||
def test_the_scheduler_runs_each_due_user_once_in_ultra_only(tmp_path, monkeypatch) -> None:
|
||||
sdk = _EvolveSdk(continuing=0)
|
||||
state = EvolveStateStore(tmp_path)
|
||||
modes = ["pro"]
|
||||
clock = [_at(1, 12)]
|
||||
service = EvolveService(
|
||||
coordinator=_Coordinator(sdk, monkeypatch),
|
||||
state=state,
|
||||
memory_mode=lambda: modes[0],
|
||||
now=lambda: clock[0],
|
||||
)
|
||||
schedule = [EvolveSchedule()]
|
||||
|
||||
async def users():
|
||||
return ("local", "wife")
|
||||
|
||||
scheduler = EvolveScheduler(
|
||||
service=service,
|
||||
schedule=lambda: schedule[0],
|
||||
users_with_memories=users,
|
||||
now=lambda: clock[0],
|
||||
)
|
||||
|
||||
async def scenario():
|
||||
assert await scheduler.tick() == () # pro: nothing
|
||||
modes[0] = "ultra"
|
||||
assert await scheduler.tick() == () # anchored at noon
|
||||
clock[0] = _at(2, 3, 2)
|
||||
assert await scheduler.tick() == ("local", "wife")
|
||||
assert await scheduler.tick() == () # once a night
|
||||
clock[0] = _at(3, 8) # slept through 03:00
|
||||
assert await scheduler.tick() == ("local", "wife")
|
||||
schedule[0] = EvolveSchedule(enabled=False)
|
||||
clock[0] = _at(5, 8)
|
||||
assert await scheduler.tick() == ()
|
||||
|
||||
asyncio.run(scenario())
|
||||
assert [call["user_id"] for call in sdk.calls] == ["local", "wife", "local", "wife"]
|
||||
assert state.user("wife").last_result["trigger"] == "scheduled"
|
||||
|
||||
|
||||
def test_the_graph_backfill_runs_once_per_user(tmp_path, monkeypatch) -> None:
|
||||
sdk = _EvolveSdk()
|
||||
state = EvolveStateStore(tmp_path)
|
||||
backfill = MemoryGraphBackfill(coordinator=_Coordinator(sdk, monkeypatch), state=state)
|
||||
|
||||
async def scenario():
|
||||
backfill.start(["local", "wife", "local"])
|
||||
await backfill.wait()
|
||||
|
||||
asyncio.run(scenario())
|
||||
assert sdk.graph_builds == ["local", "wife"]
|
||||
built = state.graph_build()
|
||||
assert built.state == "succeeded"
|
||||
assert {user: value.state for user, value in built.users.items()} == {
|
||||
"local": "succeeded",
|
||||
"wife": "succeeded",
|
||||
}
|
||||
|
||||
|
||||
# --- the HTTP surface -----------------------------------------------------------
|
||||
|
||||
|
||||
def application(tmp_path: Path):
|
||||
settings = Settings(
|
||||
data_dir=tmp_path / "data",
|
||||
source_home=tmp_path,
|
||||
api_token=SecretStr(TOKEN),
|
||||
imports_enabled=False,
|
||||
automatic_memory_processing=False,
|
||||
live_sync_enabled=False,
|
||||
)
|
||||
return create_app(settings, secret_store=InMemorySecretStore())
|
||||
|
||||
|
||||
def client_for(app) -> AsyncClient:
|
||||
return AsyncClient(
|
||||
transport=ASGITransport(app=app),
|
||||
base_url="http://test",
|
||||
headers={"Authorization": f"Bearer {TOKEN}"},
|
||||
)
|
||||
|
||||
|
||||
class _ReadSdk(_EvolveSdk):
|
||||
def memory_stats(self, *, user_id, statuses=None):
|
||||
return {"total": 1, "scenes": {}}
|
||||
|
||||
|
||||
def run_app(tmp_path, monkeypatch, scenario) -> None:
|
||||
monkeypatch.delenv("ONE_MEMORY_GRAPH_STORE", raising=False)
|
||||
|
||||
async def main() -> None:
|
||||
app = application(tmp_path)
|
||||
sdk = _ReadSdk(continuing=1)
|
||||
|
||||
@asynccontextmanager
|
||||
async def write_client():
|
||||
yield sdk
|
||||
|
||||
async with app.router.lifespan_context(app):
|
||||
monkeypatch.setattr(app.state.sdk_write_coordinator, "use_client", write_client)
|
||||
# Model writes wait for the startup scan of embedding spaces, as
|
||||
# they would behind the App's loading screen.
|
||||
for _ in range(500):
|
||||
if app.state.embedding_switch.writes_open():
|
||||
break
|
||||
await asyncio.sleep(0.01)
|
||||
async with client_for(app) as client:
|
||||
await scenario(app, sdk, client)
|
||||
|
||||
asyncio.run(main())
|
||||
|
||||
|
||||
@pytest.mark.fresh_sdk
|
||||
def test_mode_api_get_put_and_its_effects(tmp_path, monkeypatch) -> None:
|
||||
async def scenario(app, sdk, client):
|
||||
current = (await client.get("/api/v1/settings/memory-mode")).json()
|
||||
assert current["mode"] == current["configured_mode"] == "pro"
|
||||
assert current["forced_by_env"] is False
|
||||
assert app.state.sdk_write_coordinator.memory_mode() == "pro"
|
||||
|
||||
lite = await client.put("/api/v1/settings/memory-mode", json={"mode": "lite"})
|
||||
assert lite.status_code == 200, lite.text
|
||||
body = lite.json()
|
||||
assert body["mode"] == "lite" and body["previous_mode"] == "pro"
|
||||
assert body["effects"] == ["higher_mode_memories_hidden"]
|
||||
assert app.state.sdk_write_coordinator.memory_mode() == "lite"
|
||||
assert _sdk_client_handle(app).memory_mode() == "lite"
|
||||
# lite claims queued work without a generation model.
|
||||
from one_memory.api import _sedimentation_claim_enabled
|
||||
|
||||
assert app.state.model_profile_store.active_identity() is None
|
||||
app.state.history_import_gate_released = True
|
||||
assert _sedimentation_claim_enabled(app, "one-memory-manual")
|
||||
|
||||
ultra = await client.put("/api/v1/settings/memory-mode", json={"mode": "ultra"})
|
||||
assert ultra.status_code == 200, ultra.text
|
||||
assert "existing_memories_added_to_graph" in ultra.json()["effects"]
|
||||
await app.state.memory_graph_backfill.wait()
|
||||
assert sdk.graph_builds == ["local"]
|
||||
reread = (await client.get("/api/v1/settings/memory-mode")).json()
|
||||
assert reread["graph_build"]["state"] == "succeeded"
|
||||
|
||||
for bad in ({"mode": "max"}, {"mode": "pro", "x": 1}):
|
||||
assert (await client.put("/api/v1/settings/memory-mode", json=bad)).status_code == 422
|
||||
assert (
|
||||
await client.get("/api/v1/settings/memory-mode", params={"user_id": "x"})
|
||||
).status_code == 422
|
||||
|
||||
monkeypatch.setenv("ONE_MEMORY_GRAPH_STORE", "kuzu")
|
||||
await client.put("/api/v1/settings/memory-mode", json={"mode": "pro"})
|
||||
forced = (await client.get("/api/v1/settings/memory-mode")).json()
|
||||
assert forced["mode"] == "ultra" and forced["configured_mode"] == "pro"
|
||||
assert forced["forced_by_env"] is True
|
||||
|
||||
run_app(tmp_path, monkeypatch, scenario)
|
||||
|
||||
|
||||
@pytest.mark.fresh_sdk
|
||||
def test_evolve_api_is_ultra_only_async_and_scoped_per_user(tmp_path, monkeypatch) -> None:
|
||||
async def scenario(app, sdk, client):
|
||||
refused = await client.post("/api/v2/evolve", json={})
|
||||
assert refused.status_code == 409
|
||||
assert refused.json()["detail"] == "memory_mode_not_ultra"
|
||||
|
||||
await client.put("/api/v1/settings/memory-mode", json={"mode": "ultra"})
|
||||
await app.state.memory_graph_backfill.wait()
|
||||
started = await client.post(
|
||||
"/api/v2/evolve", json={"user_id": "wife", "agent_id": "all-apps", "max_batches": 2}
|
||||
)
|
||||
assert started.status_code == 202, started.text
|
||||
run = started.json()
|
||||
assert run["state"] == "running" and run["user_id"] == "wife"
|
||||
await app.state.evolve_service.wait(run["run_id"])
|
||||
path = f"/api/v2/evolve/{run['run_id']}"
|
||||
assert (await client.get(path)).status_code == 404
|
||||
done = (await client.get(path, params={"user_id": "wife"})).json()
|
||||
assert done["state"] == "succeeded" and done["rounds"] == 2
|
||||
assert sdk.calls[0] == {"user_id": "wife", "agent_id": "all-apps", "max_batches": 2}
|
||||
|
||||
overview = (await client.get("/api/v2/evolve", params={"user_id": "wife"})).json()
|
||||
assert set(overview["users"]) == {"wife"}
|
||||
assert overview["schedule"] == {"enabled": True, "time": "03:00"}
|
||||
|
||||
processing = (await client.get("/api/v1/memory-processing")).json()
|
||||
pipeline = processing["pipeline"]
|
||||
assert pipeline["mode"] == "ultra"
|
||||
assert pipeline["last_evolve_at"] is not None
|
||||
assert pipeline["evolve"]["users"]["wife"]["last_result"]["state"] == "succeeded"
|
||||
|
||||
for bad in (
|
||||
{"user_id": "a:b"},
|
||||
{"max_batches": 0},
|
||||
{"from_beginning": True},
|
||||
):
|
||||
assert (await client.post("/api/v2/evolve", json=bad)).status_code == 422
|
||||
|
||||
moved = await client.put(
|
||||
"/api/v1/settings/evolve-schedule", json={"time": "04:30", "enabled": False}
|
||||
)
|
||||
assert moved.json() == {"enabled": False, "time": "04:30"}
|
||||
assert (
|
||||
await client.put("/api/v1/settings/evolve-schedule", json={"time": "25:00"})
|
||||
).status_code == 422
|
||||
|
||||
run_app(tmp_path, monkeypatch, scenario)
|
||||
|
||||
|
||||
@pytest.mark.fresh_sdk
|
||||
def test_a_job_written_in_lite_stays_done_and_a_pending_job_takes_the_new_mode(
|
||||
tmp_path, monkeypatch
|
||||
) -> None:
|
||||
"""lite to pro reprocesses nothing: the mode is read when a job runs."""
|
||||
from one_memory.memory_server.durable_queue import DurableWriteDrain, execute_job
|
||||
|
||||
modes = ["lite"]
|
||||
seen: list[str] = []
|
||||
|
||||
class Sdk:
|
||||
def __init__(self, mode):
|
||||
self.mode = mode
|
||||
|
||||
def add(self, messages, **kwargs):
|
||||
seen.append(self.mode)
|
||||
return {"success": True, "raw_ids": [f"raw-{len(seen)}"]}
|
||||
|
||||
monkeypatch.setattr(
|
||||
sdk_write_processor,
|
||||
"SdkClientHandle",
|
||||
lambda profile, **kwargs: _SdkHandle(Sdk(kwargs["mode"]), kwargs["mode"]),
|
||||
)
|
||||
coordinator = sdk_write_processor.SdkWriteCoordinator(
|
||||
profiles=_Profiles(), # type: ignore[arg-type]
|
||||
secrets=None, # type: ignore[arg-type]
|
||||
data_dir=None,
|
||||
transport_factory=lambda _profile: None,
|
||||
sedimentation_enabled=lambda: True,
|
||||
memory_mode=lambda: modes[0],
|
||||
)
|
||||
|
||||
async def never(_job_id):
|
||||
return None
|
||||
|
||||
async def drain_once(queue) -> int:
|
||||
drain = DurableWriteDrain(queue, execute=never, lease=timedelta(seconds=30))
|
||||
claimed = drain.claim(limit=8)
|
||||
for job_id in claimed:
|
||||
await execute_job(
|
||||
job_id,
|
||||
queue=queue,
|
||||
supervisor=coordinator,
|
||||
evidence=None, # type: ignore[arg-type]
|
||||
)
|
||||
return len(claimed)
|
||||
|
||||
async def scenario(app, _sdk, client):
|
||||
queue = app.state.memory_write_queue
|
||||
for request_id in ("turn-lite", "turn-pro"):
|
||||
written = await client.post(
|
||||
"/api/v2/memory-writes",
|
||||
json={
|
||||
"request_id": request_id,
|
||||
"session_id": f"s-{request_id}",
|
||||
"source_id": "one-memory-manual",
|
||||
"messages": [
|
||||
{"role": "user", "content": request_id},
|
||||
{"role": "assistant", "content": "ok"},
|
||||
],
|
||||
},
|
||||
)
|
||||
assert written.status_code == 202, written.text
|
||||
if request_id == "turn-lite":
|
||||
assert await drain_once(queue) == 1
|
||||
modes[0] = "pro"
|
||||
await coordinator.retire_stale_modes()
|
||||
assert await drain_once(queue) == 1
|
||||
assert await drain_once(queue) == 0
|
||||
assert queue.get("turn-lite").state == queue.get("turn-pro").state == "completed"
|
||||
assert seen == ["lite", "pro"]
|
||||
await coordinator.close()
|
||||
|
||||
run_app(tmp_path, monkeypatch, scenario)
|
||||
@@ -66,6 +66,7 @@ def _grounded_setup(tmp_path, *, memory_id: str = "memory-real-data"):
|
||||
base_url="http://127.0.0.1:11434",
|
||||
model="test-model",
|
||||
embedding_model="test-embedding",
|
||||
embedding_base_url="http://127.0.0.1:11434",
|
||||
),
|
||||
secrets,
|
||||
)
|
||||
@@ -263,6 +264,7 @@ def test_answer_uses_only_grounded_confirmed_memory_citations(tmp_path) -> None:
|
||||
base_url="http://127.0.0.1:11434",
|
||||
model="test-model",
|
||||
embedding_model="test-embedding",
|
||||
embedding_base_url="http://127.0.0.1:11434",
|
||||
),
|
||||
secrets,
|
||||
)
|
||||
@@ -468,6 +470,7 @@ def test_answer_rejects_model_citation_outside_retrieved_scope(tmp_path) -> None
|
||||
base_url="http://127.0.0.1:11434",
|
||||
model="test-model",
|
||||
embedding_model="test-embedding",
|
||||
embedding_base_url="http://127.0.0.1:11434",
|
||||
),
|
||||
secrets,
|
||||
)
|
||||
@@ -530,6 +533,7 @@ def _compatible_gateway_setup(tmp_path):
|
||||
base_url="https://api.example.invalid/standard/v1",
|
||||
model="deepseek-v4-pro",
|
||||
embedding_model="text-embedding",
|
||||
embedding_base_url="https://api.example.invalid/standard/v1",
|
||||
api_key="probe-key",
|
||||
),
|
||||
secrets,
|
||||
|
||||
@@ -186,6 +186,7 @@ def _setup(tmp_path: Path, *, model: str = "available") -> _Setup:
|
||||
base_url="https://memory-qa.example.invalid/v1",
|
||||
model="synthetic-generation",
|
||||
embedding_model="synthetic-embedding",
|
||||
embedding_base_url="https://memory-qa.example.invalid/v1",
|
||||
api_key=SecretStr(_SECRET),
|
||||
),
|
||||
secrets,
|
||||
|
||||
@@ -12,10 +12,12 @@ import time
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from hashlib import sha256
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from one_memory.database_access import main_database
|
||||
from one_memory.embedding_spaces import EmbeddingIdentity
|
||||
from one_memory.memory_server.contracts import BatchIdentity, WriteRequest
|
||||
from one_memory.memory_server.durable_queue import (
|
||||
CANCELLED,
|
||||
@@ -1226,6 +1228,10 @@ def _coordinator(monkeypatch, opened: list[_TrackedHandle]):
|
||||
return value
|
||||
|
||||
monkeypatch.setattr(sdk_write_processor, "SdkClientHandle", handle)
|
||||
# One live embedding space the stand-in profiles all fill.
|
||||
identity = EmbeddingIdentity(model="stand-in")
|
||||
monkeypatch.setattr(sdk_write_processor, "embedding_identity", lambda _profile: identity)
|
||||
space = SimpleNamespace(space_id="space", identity=identity)
|
||||
profiles = _SwitchableProfiles()
|
||||
coordinator = sdk_write_processor.SdkWriteCoordinator(
|
||||
profiles=profiles, # type: ignore[arg-type]
|
||||
@@ -1233,6 +1239,7 @@ def _coordinator(monkeypatch, opened: list[_TrackedHandle]):
|
||||
data_dir=None,
|
||||
transport_factory=lambda _profile: None,
|
||||
sedimentation_enabled=lambda: True,
|
||||
spaces=SimpleNamespace(live=lambda: space), # type: ignore[arg-type]
|
||||
)
|
||||
return coordinator, profiles
|
||||
|
||||
|
||||
@@ -25,6 +25,7 @@ def _stores(tmp_path):
|
||||
provider_kind=ProviderKind.OLLAMA,
|
||||
model="hy-bot-phi4",
|
||||
embedding_model="nomic-embed-text",
|
||||
embedding_base_url="http://127.0.0.1:11434",
|
||||
),
|
||||
secrets,
|
||||
)
|
||||
|
||||
@@ -83,6 +83,7 @@ def test_task_freezes_active_profile_revision_and_is_idempotent(tmp_path) -> Non
|
||||
provider_kind=ProviderKind.OLLAMA,
|
||||
model="qwen",
|
||||
embedding_model="nomic-embed-text",
|
||||
embedding_base_url="http://127.0.0.1:11434",
|
||||
),
|
||||
secrets,
|
||||
)
|
||||
@@ -240,6 +241,7 @@ def test_task_extraction_checkpoint_is_idempotent_and_recoverable(tmp_path) -> N
|
||||
provider_kind=ProviderKind.OLLAMA,
|
||||
model="qwen",
|
||||
embedding_model="nomic-embed-text",
|
||||
embedding_base_url="http://127.0.0.1:11434",
|
||||
),
|
||||
secrets,
|
||||
),
|
||||
@@ -293,6 +295,7 @@ def test_task_waits_for_model_and_retries_with_same_snapshot(tmp_path) -> None:
|
||||
provider_kind=ProviderKind.OLLAMA,
|
||||
model="qwen",
|
||||
embedding_model="nomic-embed-text",
|
||||
embedding_base_url="http://127.0.0.1:11434",
|
||||
),
|
||||
secrets,
|
||||
)
|
||||
@@ -337,6 +340,7 @@ def test_blocked_stages_name_the_stuck_stage_and_reason(tmp_path) -> None:
|
||||
provider_kind=ProviderKind.OLLAMA,
|
||||
model="qwen",
|
||||
embedding_model="nomic-embed-text",
|
||||
embedding_base_url="http://127.0.0.1:11434",
|
||||
),
|
||||
secrets,
|
||||
),
|
||||
@@ -433,6 +437,7 @@ def test_task_duration_estimate_ignores_time_the_app_was_closed(tmp_path) -> Non
|
||||
provider_kind=ProviderKind.OLLAMA,
|
||||
model="qwen",
|
||||
embedding_model="nomic-embed-text",
|
||||
embedding_base_url="http://127.0.0.1:11434",
|
||||
),
|
||||
secrets,
|
||||
),
|
||||
@@ -537,6 +542,7 @@ def test_running_task_is_recovered_after_worker_restart(tmp_path) -> None:
|
||||
provider_kind=ProviderKind.OLLAMA,
|
||||
model="qwen",
|
||||
embedding_model="nomic-embed-text",
|
||||
embedding_base_url="http://127.0.0.1:11434",
|
||||
),
|
||||
secrets,
|
||||
),
|
||||
@@ -572,6 +578,7 @@ def test_superseded_profile_cancels_only_claimable_tasks(tmp_path) -> None:
|
||||
provider_kind=ProviderKind.OLLAMA,
|
||||
model="phi",
|
||||
embedding_model="nomic-embed-text",
|
||||
embedding_base_url="http://127.0.0.1:11434",
|
||||
),
|
||||
secrets,
|
||||
),
|
||||
@@ -601,6 +608,7 @@ def test_superseded_profile_cancels_only_claimable_tasks(tmp_path) -> None:
|
||||
provider_kind=ProviderKind.OLLAMA,
|
||||
model="qwen",
|
||||
embedding_model="nomic-embed-text",
|
||||
embedding_base_url="http://127.0.0.1:11434",
|
||||
),
|
||||
secrets,
|
||||
),
|
||||
@@ -630,6 +638,7 @@ def test_superseded_ingest_policy_cancels_only_old_conversation_tasks(tmp_path)
|
||||
provider_kind=ProviderKind.OLLAMA,
|
||||
model="qwen",
|
||||
embedding_model="nomic-embed-text",
|
||||
embedding_base_url="http://127.0.0.1:11434",
|
||||
),
|
||||
secrets,
|
||||
),
|
||||
@@ -676,6 +685,7 @@ def test_task_rejects_inactive_profile_and_request_id_reuse(tmp_path) -> None:
|
||||
provider_kind=ProviderKind.OLLAMA,
|
||||
model="qwen",
|
||||
embedding_model="nomic-embed-text",
|
||||
embedding_base_url="http://127.0.0.1:11434",
|
||||
),
|
||||
secrets,
|
||||
)
|
||||
@@ -1700,6 +1710,7 @@ def _profile_and_tasks(tmp_path) -> tuple[MemoryTaskStore, ModelProfile]:
|
||||
provider_kind=ProviderKind.OLLAMA,
|
||||
model="qwen",
|
||||
embedding_model="nomic-embed-text",
|
||||
embedding_base_url="http://127.0.0.1:11434",
|
||||
),
|
||||
secrets,
|
||||
),
|
||||
@@ -2058,6 +2069,7 @@ def test_a_failed_task_below_the_ceiling_is_still_claimable(tmp_path) -> None:
|
||||
provider_kind=ProviderKind.OLLAMA,
|
||||
model="qwen",
|
||||
embedding_model="nomic-embed-text",
|
||||
embedding_base_url="http://127.0.0.1:11434",
|
||||
),
|
||||
secrets,
|
||||
),
|
||||
|
||||
@@ -55,6 +55,7 @@ def library(tmp_path):
|
||||
provider_kind=ProviderKind.OLLAMA,
|
||||
model="hy-bot-phi4",
|
||||
embedding_model="nomic-embed-text",
|
||||
embedding_base_url="http://127.0.0.1:11434",
|
||||
),
|
||||
secrets,
|
||||
)
|
||||
@@ -611,3 +612,58 @@ def test_granting_a_source_repeatedly_does_not_raise(library):
|
||||
|
||||
with main_database_read(database) as db:
|
||||
assert memory.memory_id in pending(db)
|
||||
|
||||
|
||||
def test_grant_writes_still_work_after_the_cutover_drops_the_link_table(library):
|
||||
"""#233: the grant triggers read `governed_memory_evidence`.
|
||||
|
||||
The SDK cutover drops it, and SQLite resolves a trigger body only when it
|
||||
fires, so every grant, revoke and connect then failed with "no such
|
||||
table". A trigger whose body names a missing table is not installed, and
|
||||
the copy an earlier startup installed is dropped.
|
||||
"""
|
||||
database, *_ = library
|
||||
from one_memory.imports.candidates import RootCandidateKind
|
||||
from one_memory.imports.grants import SourceGrantStore
|
||||
from one_memory.memory_visibility import drop_stale_triggers
|
||||
|
||||
startup(database)
|
||||
with sqlite3.connect(database) as db:
|
||||
db.execute("DROP TABLE governed_memory_evidence")
|
||||
|
||||
def grant_triggers() -> set[str]:
|
||||
with sqlite3.connect(database) as db:
|
||||
return {
|
||||
str(row[0])
|
||||
for row in db.execute(
|
||||
"SELECT name FROM sqlite_master "
|
||||
"WHERE type='trigger' AND tbl_name='source_grants'"
|
||||
)
|
||||
}
|
||||
|
||||
assert grant_triggers() == {
|
||||
"memory_visibility_dirty_on_grants",
|
||||
"memory_visibility_dirty_on_grants_update",
|
||||
}
|
||||
# The `evidence` triggers read the dropped table too; the cutover drops
|
||||
# `evidence` as well, which takes them with it.
|
||||
assert {
|
||||
"memory_visibility_dirty_on_grants",
|
||||
"memory_visibility_dirty_on_grants_update",
|
||||
} <= set(drop_stale_triggers(database))
|
||||
assert grant_triggers() == set()
|
||||
assert drop_stale_triggers(database) == ()
|
||||
# Installing again on this shape skips them instead of bringing them back.
|
||||
startup(database)
|
||||
assert grant_triggers() == set()
|
||||
|
||||
folder = database.parent / "cursor"
|
||||
folder.mkdir(exist_ok=True)
|
||||
store = SourceGrantStore(database, home=database.parent)
|
||||
store.grant(
|
||||
user_id="local",
|
||||
source_id="cursor",
|
||||
root=folder,
|
||||
candidate_kind=RootCandidateKind.DEFAULT,
|
||||
)
|
||||
assert store.revoke(user_id="local", source_id="cursor").revoked_at is not None
|
||||
@@ -51,6 +51,7 @@ def library(tmp_path):
|
||||
provider_kind=ProviderKind.OLLAMA,
|
||||
model="hy-bot-phi4",
|
||||
embedding_model="nomic-embed-text",
|
||||
embedding_base_url="http://127.0.0.1:11434",
|
||||
),
|
||||
secrets,
|
||||
)
|
||||
|
||||
@@ -28,6 +28,7 @@ def test_probe_still_measures_embedding_when_structured_output_fails(tmp_path) -
|
||||
base_url="http://127.0.0.1:11434",
|
||||
model="plain-model",
|
||||
embedding_model="nomic-embed-text",
|
||||
embedding_base_url="http://127.0.0.1:11434",
|
||||
),
|
||||
secrets,
|
||||
)
|
||||
@@ -63,6 +64,7 @@ def test_probe_records_the_real_embedding_width_on_the_profile(tmp_path) -> None
|
||||
base_url="http://127.0.0.1:11434",
|
||||
model="qwen3:8b",
|
||||
embedding_model="nomic-embed-text",
|
||||
embedding_base_url="http://127.0.0.1:11434",
|
||||
),
|
||||
secrets,
|
||||
)
|
||||
@@ -119,6 +121,7 @@ def test_openai_compatible_probe_checks_model_and_structured_output(tmp_path) ->
|
||||
base_url="https://models.example.com/v1",
|
||||
model="memory-model",
|
||||
embedding_model="embedding-model",
|
||||
embedding_base_url="https://models.example.com/v1",
|
||||
api_key=SecretStr(API_KEY),
|
||||
),
|
||||
secrets,
|
||||
@@ -260,6 +263,7 @@ def test_probe_reports_missing_model_without_sending_generation(tmp_path) -> Non
|
||||
provider_kind=ProviderKind.OLLAMA,
|
||||
model="memory-model",
|
||||
embedding_model="embedding-model",
|
||||
embedding_base_url="http://127.0.0.1:11434",
|
||||
),
|
||||
secrets,
|
||||
)
|
||||
|
||||
@@ -50,6 +50,7 @@ def test_model_profile_api_never_returns_key_and_enforces_revision(tmp_path) ->
|
||||
"base_url": "https://models.example.com/v1",
|
||||
"model": "memory-model",
|
||||
"embedding_model": "embedding-model",
|
||||
"embedding_base_url": "https://models.example.com/v1",
|
||||
"api_key": API_KEY,
|
||||
},
|
||||
)
|
||||
@@ -134,6 +135,7 @@ def test_model_profile_probe_route_returns_only_machine_results(tmp_path) -> Non
|
||||
"base_url": "https://models.example.com/v1",
|
||||
"model": "memory-model",
|
||||
"embedding_model": "embedding-model",
|
||||
"embedding_base_url": "https://models.example.com/v1",
|
||||
"api_key": API_KEY,
|
||||
},
|
||||
)
|
||||
@@ -244,6 +246,7 @@ def test_activation_survives_a_corpus_already_imported(tmp_path) -> None:
|
||||
"base_url": "https://models.example.com/v1",
|
||||
"model": "memory-model",
|
||||
"embedding_model": "embedding-model",
|
||||
"embedding_base_url": "https://models.example.com/v1",
|
||||
"api_key": API_KEY,
|
||||
},
|
||||
)
|
||||
|
||||
@@ -198,6 +198,7 @@ def test_active_profile_can_be_deleted_without_changing_other_profiles(tmp_path)
|
||||
provider_kind=ProviderKind.OLLAMA,
|
||||
model="other",
|
||||
embedding_model="nomic-embed-text",
|
||||
embedding_base_url="http://127.0.0.1:11434",
|
||||
),
|
||||
secrets,
|
||||
)
|
||||
@@ -208,6 +209,7 @@ def test_active_profile_can_be_deleted_without_changing_other_profiles(tmp_path)
|
||||
base_url="https://models.example.com/v1",
|
||||
model="qwen",
|
||||
embedding_model="nomic-embed-text",
|
||||
embedding_base_url="https://models.example.com/v1",
|
||||
api_key=SecretStr(API_KEY),
|
||||
),
|
||||
secrets,
|
||||
@@ -349,6 +351,7 @@ def _probed(store, secrets, *, structured: bool, generates: bool = True):
|
||||
provider_kind=ProviderKind.OLLAMA,
|
||||
model="qwen",
|
||||
embedding_model="nomic-embed-text",
|
||||
embedding_base_url="http://127.0.0.1:11434",
|
||||
),
|
||||
secrets,
|
||||
)
|
||||
@@ -394,3 +397,102 @@ def test_a_model_that_cannot_generate_still_cannot_be_turned_on(tmp_path) -> Non
|
||||
|
||||
with pytest.raises(ModelProfileUnavailableError):
|
||||
store.activate(profile.profile_id, secrets)
|
||||
|
||||
|
||||
def _space(store: ModelProfileStore, profile_id: str) -> str:
|
||||
"""The embedding space's collection base, as the Runtime names it."""
|
||||
from one_memory.hy_sdk_runtime import vector_collection_bases
|
||||
|
||||
identity = (profile_id, store.require(profile_id, InMemorySecretStore()).revision)
|
||||
return vector_collection_bases(store.snapshot(*identity))[0]
|
||||
|
||||
|
||||
def test_a_remote_embedding_must_name_its_endpoint(tmp_path) -> None:
|
||||
"""The endpoint no longer follows the generation `base_url`."""
|
||||
from one_memory.model_profiles import EmbeddingBaseUrlRequiredError
|
||||
|
||||
store = ModelProfileStore(tmp_path / "one-memory.sqlite3")
|
||||
secrets = InMemorySecretStore()
|
||||
unnamed = ModelProfileInput(
|
||||
name="Remote",
|
||||
provider_kind=ProviderKind.OPENAI_COMPATIBLE,
|
||||
base_url="https://models.example.com/v1",
|
||||
model="qwen",
|
||||
embedding_model="text-embedding-3-small",
|
||||
api_key=SecretStr(API_KEY),
|
||||
)
|
||||
with pytest.raises(EmbeddingBaseUrlRequiredError):
|
||||
store.create(unnamed, secrets)
|
||||
assert store.list(secrets) == []
|
||||
|
||||
named = store.create(
|
||||
unnamed.model_copy(update={"embedding_base_url": "https://models.example.com/v1"}),
|
||||
secrets,
|
||||
)
|
||||
with pytest.raises(EmbeddingBaseUrlRequiredError):
|
||||
store.update(
|
||||
named.profile_id,
|
||||
ModelProfilePatch(
|
||||
expected_revision=named.revision,
|
||||
embedding_provider_kind=ProviderKind.OLLAMA,
|
||||
embedding_base_url="",
|
||||
),
|
||||
secrets,
|
||||
)
|
||||
|
||||
|
||||
def test_stored_profiles_get_the_endpoint_they_resolved_to_and_keep_their_space(
|
||||
tmp_path,
|
||||
) -> None:
|
||||
"""An older build stored remote embeddings with a NULL endpoint and resolved
|
||||
it to the generation `base_url` at runtime. Opening the store writes that
|
||||
endpoint in once, so the collection, and the memories in it, stay put; a
|
||||
later edit of the generation endpoint no longer moves them.
|
||||
"""
|
||||
from one_memory.hy_sdk_runtime import vector_collection_bases
|
||||
|
||||
database = tmp_path / "one-memory.sqlite3"
|
||||
store = ModelProfileStore(database)
|
||||
secrets = InMemorySecretStore()
|
||||
profile = store.create(
|
||||
ModelProfileInput(
|
||||
name="Remote",
|
||||
provider_kind=ProviderKind.OPENAI_COMPATIBLE,
|
||||
base_url="https://gateway.example.com/v1",
|
||||
model="qwen",
|
||||
embedding_model="text-embedding-3-small",
|
||||
embedding_base_url="https://gateway.example.com/v1",
|
||||
embedding_dimensions=1536,
|
||||
api_key=SecretStr(API_KEY),
|
||||
),
|
||||
secrets,
|
||||
)
|
||||
identity = (profile.profile_id, profile.revision)
|
||||
# The space the older build's Runtime resolved: `embedding_base_url or base_url`.
|
||||
before = vector_collection_bases(store.snapshot(*identity))[0]
|
||||
with sqlite3.connect(database) as db:
|
||||
# The shape an older build stored.
|
||||
db.execute("UPDATE model_profiles SET embedding_base_url=NULL")
|
||||
db.execute("UPDATE model_profile_versions SET embedding_base_url=NULL")
|
||||
|
||||
migrated = ModelProfileStore(database)
|
||||
assert migrated.snapshot(*identity).embedding_base_url == "https://gateway.example.com/v1"
|
||||
assert _space(migrated, profile.profile_id) == before
|
||||
# Idempotent: opening again changes nothing.
|
||||
assert _space(ModelProfileStore(database), profile.profile_id) == before
|
||||
with sqlite3.connect(database) as db:
|
||||
assert db.execute(
|
||||
"SELECT COUNT(*) FROM model_profile_versions WHERE embedding_base_url IS NULL"
|
||||
).fetchone() == (0,)
|
||||
|
||||
edited = migrated.update(
|
||||
profile.profile_id,
|
||||
ModelProfilePatch(
|
||||
expected_revision=profile.revision,
|
||||
base_url="https://gateway.example.com/eu/v1",
|
||||
),
|
||||
secrets,
|
||||
)
|
||||
assert edited.base_url == "https://gateway.example.com/eu/v1"
|
||||
assert edited.embedding_base_url == "https://gateway.example.com/v1"
|
||||
assert _space(migrated, profile.profile_id) == before
|
||||
@@ -650,6 +650,75 @@ def test_two_hermes_grants_each_read_their_own_root_into_separate_batches(tmp_pa
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
@pytest.mark.fresh_sdk
|
||||
def test_a_migrated_home_with_the_old_grant_triggers_can_connect_and_revoke(tmp_path) -> None:
|
||||
"""#233: a home an earlier build migrated kept two `source_grants` triggers
|
||||
whose body reads the dropped `governed_memory_evidence`. Every grant write
|
||||
then failed, so connecting and revoking answered 503. Startup drops them.
|
||||
"""
|
||||
stale = """
|
||||
CREATE TRIGGER memory_visibility_dirty_on_grants{suffix}
|
||||
AFTER {event} ON source_grants
|
||||
BEGIN
|
||||
INSERT INTO memory_visibility_dirty (memory_id)
|
||||
SELECT DISTINCT link.memory_id
|
||||
FROM governed_memory_evidence link
|
||||
JOIN evidence e ON e.evidence_id = link.evidence_id
|
||||
WHERE e.source_id = NEW.source_id
|
||||
ON CONFLICT(memory_id) DO NOTHING;
|
||||
END;"""
|
||||
|
||||
async def scenario():
|
||||
home = tmp_path / "home"
|
||||
_hermes, wife = _hermes_homes(home)
|
||||
app = application(tmp_path, home=home)
|
||||
async with app.router.lifespan_context(app):
|
||||
database = app.state.settings.data_dir / "one-memory.sqlite3"
|
||||
with sqlite3.connect(database) as db:
|
||||
db.execute("CREATE TABLE IF NOT EXISTS memory_visibility_dirty (memory_id TEXT)")
|
||||
db.execute(stale.format(suffix="", event="INSERT"))
|
||||
db.execute(stale.format(suffix="_update", event="UPDATE"))
|
||||
|
||||
app = application(tmp_path, home=home)
|
||||
async with app.router.lifespan_context(app), client_for(app) as client:
|
||||
location = (await client.post("/api/v1/sources/hermes/location")).json()
|
||||
local = await client.post(
|
||||
"/api/v1/sources/connect",
|
||||
json={
|
||||
"source_ids": ["hermes"],
|
||||
"request_id": str(uuid4()),
|
||||
"source_roots": {"hermes": location["root_id"]},
|
||||
},
|
||||
)
|
||||
assert local.status_code == 200, local.text
|
||||
profile = await client.post(
|
||||
"/api/v1/sources/connect",
|
||||
json={
|
||||
"source_ids": ["hermes"],
|
||||
"request_id": str(uuid4()),
|
||||
"user_id": "wife",
|
||||
"root": str(wife),
|
||||
},
|
||||
)
|
||||
assert profile.status_code == 200, profile.text
|
||||
for user_id in ("wife", "local"):
|
||||
revoked = await client.post(
|
||||
"/api/v1/sources/hermes/revoke", params={"user_id": user_id}
|
||||
)
|
||||
assert revoked.status_code == 200, revoked.text
|
||||
assert revoked.json()["revoked_at"] is not None
|
||||
with sqlite3.connect(database) as db:
|
||||
assert (
|
||||
db.execute(
|
||||
"SELECT COUNT(*) FROM sqlite_master "
|
||||
"WHERE type='trigger' AND tbl_name='source_grants'"
|
||||
).fetchone()[0]
|
||||
== 0
|
||||
)
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
@pytest.mark.fresh_sdk
|
||||
def test_only_hermes_takes_another_user_or_a_custom_root(tmp_path) -> None:
|
||||
async def scenario():
|
||||
|
||||
Loaded 100 of 126 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user