kaiyuan/chatgpt-on-wechat
1
0
Fork 0
mirror of https://github.com/zhayujie/chatgpt-on-wechat.git synced 2026-07-19 21:07:28 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat(memory): support multi-vendor embedding fallback

Browse Source 
Add embedding_provider config knob with native support for
openai / dashscope / doubao / zhipu / linkai, plus an in-chat
/memory status and /memory rebuild-index workflow for switching
vendors safely.
This commit is contained in:
zhayujie
2026-05-20 11:00:53 +08:00
parent a0dfdb79df
commit 3ffb563a44
12 changed files with 1572 additions and 449 deletions
Download Patch File Download Diff File Expand all files Collapse all files

167
agent/memory/embedding.py
View File

@@ -1,167 +0,0 @@
"""
Embedding providers for memory
Supports OpenAI and local embedding models
"""
import hashlib
from abc import ABC, abstractmethod
from typing import List, Optional
class EmbeddingProvider(ABC):
"""Base class for embedding providers"""
@abstractmethod
def embed(self, text: str) -> List[float]:
"""Generate embedding for text"""
pass
@abstractmethod
def embed_batch(self, texts: List[str]) -> List[List[float]]:
"""Generate embeddings for multiple texts"""
pass
@property
@abstractmethod
def dimensions(self) -> int:
"""Get embedding dimensions"""
pass
class OpenAIEmbeddingProvider(EmbeddingProvider):
"""OpenAI embedding provider using REST API"""
def __init__(self, model: str = "text-embedding-3-small", api_key: Optional[str] = None,
api_base: Optional[str] = None, extra_headers: Optional[dict] = None):
"""
Initialize OpenAI embedding provider
Args:
model: Model name (text-embedding-3-small or text-embedding-3-large)
api_key: OpenAI API key
api_base: Optional API base URL
extra_headers: Optional extra headers to include in API requests
"""
self.model = model
self.api_key = api_key
self.api_base = api_base or "https://api.openai.com/v1"
self.extra_headers = extra_headers or {}
# Validate API key
if not self.api_key or self.api_key in ["", "YOUR API KEY", "YOUR_API_KEY"]:
raise ValueError("OpenAI API key is not configured. Please set 'open_ai_api_key' in config.json")
# Set dimensions based on model
self._dimensions = 1536 if "small" in model else 3072
def _call_api(self, input_data):
"""Call OpenAI embedding API using requests"""
import requests
url = f"{self.api_base}/embeddings"
headers = {
"Content-Type": "application/json",
"Authorization": f"Bearer {self.api_key}",
**self.extra_headers,
}
data = {
"input": input_data,
"model": self.model
}
try:
response = requests.post(url, headers=headers, json=data, timeout=5)
response.raise_for_status()
return response.json()
except requests.exceptions.ConnectionError as e:
raise ConnectionError(f"Failed to connect to OpenAI API at {url}. Please check your network connection and api_base configuration. Error: {str(e)}")
except requests.exceptions.Timeout as e:
raise TimeoutError(f"OpenAI API request timed out after 10s. Please check your network connection. Error: {str(e)}")
except requests.exceptions.HTTPError as e:
if e.response.status_code == 401:
raise ValueError(f"Invalid OpenAI API key. Please check your 'open_ai_api_key' in config.json")
elif e.response.status_code == 429:
raise ValueError(f"OpenAI API rate limit exceeded. Please try again later.")
else:
raise ValueError(f"OpenAI API request failed: {e.response.status_code} - {e.response.text}")
def embed(self, text: str) -> List[float]:
"""Generate embedding for text"""
result = self._call_api(text)
return result["data"][0]["embedding"]
def embed_batch(self, texts: List[str]) -> List[List[float]]:
"""Generate embeddings for multiple texts"""
if not texts:
return []
result = self._call_api(texts)
return [item["embedding"] for item in result["data"]]
@property
def dimensions(self) -> int:
return self._dimensions
# LocalEmbeddingProvider removed - only use OpenAI embedding or keyword search
class EmbeddingCache:
"""Cache for embeddings to avoid recomputation"""
def __init__(self):
self.cache = {}
def get(self, text: str, provider: str, model: str) -> Optional[List[float]]:
"""Get cached embedding"""
key = self._compute_key(text, provider, model)
return self.cache.get(key)
def put(self, text: str, provider: str, model: str, embedding: List[float]):
"""Cache embedding"""
key = self._compute_key(text, provider, model)
self.cache[key] = embedding
@staticmethod
def _compute_key(text: str, provider: str, model: str) -> str:
"""Compute cache key"""
content = f"{provider}:{model}:{text}"
return hashlib.md5(content.encode('utf-8')).hexdigest()
def clear(self):
"""Clear cache"""
self.cache.clear()
def create_embedding_provider(
provider: str = "openai",
model: Optional[str] = None,
api_key: Optional[str] = None,
api_base: Optional[str] = None,
extra_headers: Optional[dict] = None
) -> EmbeddingProvider:
"""
Factory function to create embedding provider
Supports "openai" and "linkai" providers (both use OpenAI-compatible REST API).
If initialization fails, caller should fall back to keyword-only search.
Args:
provider: Provider name ("openai" or "linkai")
model: Model name (default: text-embedding-3-small)
api_key: API key (required)
api_base: API base URL
extra_headers: Optional extra headers to include in API requests
Returns:
EmbeddingProvider instance
Raises:
ValueError: If provider is unsupported or api_key is missing
"""
if provider not in ("openai", "linkai"):
raise ValueError(f"Unsupported embedding provider: {provider}. Use 'openai' or 'linkai'.")
model = model or "text-embedding-3-small"
return OpenAIEmbeddingProvider(model=model, api_key=api_key, api_base=api_base, extra_headers=extra_headers)

41
agent/memory/embedding/__init__.py Normal file
View File

@@ -0,0 +1,41 @@
"""
Embedding subsystem for memory.
Public API:
create_embedding_provider, EmbeddingProvider, OpenAIEmbeddingProvider,
EMBEDDING_VENDORS, EmbeddingCache
RebuildResult, clear_index, rebuild_in_process
detect_index_dim, cleanup_legacy_state_file
"""
from agent.memory.embedding.provider import (
EMBEDDING_VENDORS,
DoubaoEmbeddingProvider,
EmbeddingCache,
EmbeddingProvider,
OpenAIEmbeddingProvider,
create_embedding_provider,
)
from agent.memory.embedding.rebuild import (
RebuildResult,
clear_index,
rebuild_in_process,
)
from agent.memory.embedding.state import (
cleanup_legacy_state_file,
detect_index_dim,
)
__all__ = [
"EMBEDDING_VENDORS",
"DoubaoEmbeddingProvider",
"EmbeddingCache",
"EmbeddingProvider",
"OpenAIEmbeddingProvider",
"create_embedding_provider",
"RebuildResult",
"clear_index",
"rebuild_in_process",
"cleanup_legacy_state_file",
"detect_index_dim",
]

486
agent/memory/embedding/provider.py Normal file
View File

@@ -0,0 +1,486 @@
"""
Embedding providers for memory
Supports multiple OpenAI-compatible embedding vendors:
- openai (text-embedding-3-small / large)
- linkai (OpenAI-compatible passthrough)
- dashscope (Aliyun Tongyi text-embedding-v4)
- doubao (ByteDance Doubao Seed1.5 / large-text on Volcengine Ark)
- zhipu (ZhipuAI embedding-3)
Vendor keys here intentionally match the project's bot_type constants in
common.const (OPENAI, LINKAI, QWEN_DASHSCOPE, DOUBAO, ZHIPU_AI).
All providers share a single OpenAI-compatible REST client. Vendor-specific
behaviors (truncation, query instruction prefix) are configured via metadata.
"""
import hashlib
import math
from abc import ABC, abstractmethod
from typing import List, Optional
# HTTP read timeout for a single embeddings request (seconds). A batch of
# 64+ chunks can take 30-50s end-to-end from China-side networks, so 30s is
# routinely too tight; 90s gives meaningful headroom without letting bad
# endpoints hang forever.
EMBEDDING_HTTP_TIMEOUT = 90
class EmbeddingProvider(ABC):
"""Base class for embedding providers"""
@abstractmethod
def embed(self, text: str) -> List[float]:
"""Generate embedding for a single text (treated as a query by default)"""
pass
@abstractmethod
def embed_batch(self, texts: List[str]) -> List[List[float]]:
"""Generate embeddings for multiple texts (treated as documents)"""
pass
def embed_query(self, text: str) -> List[float]:
"""Generate embedding for a query string (may apply vendor instruction prefix)"""
return self.embed(text)
@property
@abstractmethod
def dimensions(self) -> int:
"""Effective embedding dimensions"""
pass
# ---------------------------------------------------------------------------
# Vendor metadata table
# ---------------------------------------------------------------------------
#
# Each entry describes how to reach a vendor's embedding endpoint. Most
# vendors expose an OpenAI-compatible /embeddings API; the few that don't
# (currently: doubao) set `provider_class` to pick a dedicated adapter.
# Fields:
# provider_class : optional adapter key ("doubao"); defaults to OpenAI-compat
# default_base_url : default API base when not overridden by user
# default_model : default embedding model name
# default_dimensions : recommended unified dim when explicit path is enabled
# supports_dim_param : whether the API accepts a `dimensions` request param
# needs_client_truncate : whether to slice + L2-normalize on the client side
# needs_client_normalize : whether to L2-normalize on the client (always safe)
# query_instruction : optional prefix for asymmetric retrieval (Doubao Seed)
# max_batch_size : max texts per /embeddings request; embed_batch
# auto-paginates above this. Conservative defaults.
#
EMBEDDING_VENDORS = {
"openai": {
"default_base_url": "https://api.openai.com/v1",
"default_model": "text-embedding-3-small",
# Match the legacy default so users adding `embedding_provider: openai`
# to an existing index don't need to rebuild. Override via
# embedding_dimensions if you want 1024 / 1536 / 3072.
"default_dimensions": 1536,
"supports_dim_param": True,
"needs_client_truncate": False,
"needs_client_normalize": False,
"query_instruction": "",
# OpenAI permits up to 2048 items per request, but a single call
# carrying hundreds of long chunks routinely exceeds the 30s read
# timeout from China-side networks. 64 keeps each call well under
# both the token-per-request budget and a reasonable wall clock.
"max_batch_size": 64,
},
"linkai": {
"default_base_url": "https://api.link-ai.tech/v1",
"default_model": "text-embedding-3-small",
"default_dimensions": 1536,
"supports_dim_param": True,
"needs_client_truncate": False,
"needs_client_normalize": False,
"query_instruction": "",
"max_batch_size": 64,
},
"dashscope": {
"default_base_url": "https://dashscope.aliyuncs.com/compatible-mode/v1",
"default_model": "text-embedding-v4",
"default_dimensions": 1024,
"supports_dim_param": True,
"needs_client_truncate": False,
"needs_client_normalize": False,
"query_instruction": "",
"max_batch_size": 10, # DashScope hard cap (text-embedding-v4)
},
"doubao": {
# Doubao no longer offers an OpenAI-compatible /v1/embeddings endpoint.
# Current models are unified under /api/v3/embeddings/multimodal
# which uses a structured `input` payload — see DoubaoEmbeddingProvider.
"provider_class": "doubao",
"default_base_url": "https://ark.cn-beijing.volces.com/api/v3",
"default_model": "doubao-embedding-vision-251215",
# Native options: 1024 or 2048. We default to 1024 to align with the
# other Chinese vendors (dashscope/zhipu) and keep storage footprint
# consistent across providers; users can still override via
# `embedding_dimensions: 2048` in config.
"default_dimensions": 1024,
"supports_dim_param": True,
"needs_client_truncate": False,
"needs_client_normalize": False,
"query_instruction": "",
# Multimodal endpoint produces ONE embedding per call (input list is
# a single document's parts, not a batch). embed_batch loops.
"max_batch_size": 1,
},
"zhipu": {
"default_base_url": "https://open.bigmodel.cn/api/paas/v4",
"default_model": "embedding-3",
"default_dimensions": 1024,
"supports_dim_param": True,
"needs_client_truncate": False,
"needs_client_normalize": False,
"query_instruction": "",
"max_batch_size": 64,
},
}
def _l2_normalize(vec: List[float]) -> List[float]:
"""Normalize a vector to unit length (L2 norm). Returns input on zero vector."""
norm = math.sqrt(sum(v * v for v in vec))
if norm == 0:
return vec
return [v / norm for v in vec]
class OpenAIEmbeddingProvider(EmbeddingProvider):
"""
OpenAI-compatible embedding provider.
Used for openai/linkai/dashscope/ark/zhipu by configuring the metadata
fields. The legacy two-arg constructor (model, api_key, api_base) keeps
working, so the original OpenAI/LinkAI fallback code path is unchanged.
"""
def __init__(
self,
model: str = "text-embedding-3-small",
api_key: Optional[str] = None,
api_base: Optional[str] = None,
extra_headers: Optional[dict] = None,
dimensions: Optional[int] = None,
supports_dim_param: bool = True,
needs_client_truncate: bool = False,
needs_client_normalize: bool = False,
query_instruction: str = "",
max_batch_size: int = 256,
):
"""
Args:
model: Model name (e.g. text-embedding-3-small, text-embedding-v4, embedding-3)
api_key: API key (required)
api_base: API base URL (defaults to OpenAI)
extra_headers: Optional extra HTTP headers
dimensions: Target output dimension. Required when supports_dim_param
is False and needs_client_truncate is True (used to slice).
supports_dim_param: Whether the vendor accepts a `dimensions` body param
needs_client_truncate: Slice the returned vector to `dimensions`
needs_client_normalize: L2-normalize on the client after slicing
query_instruction: Optional prefix prepended to query texts only
max_batch_size: Max items per /embeddings request; embed_batch
auto-paginates above this.
"""
self.model = model
self.api_key = api_key
self.api_base = api_base or "https://api.openai.com/v1"
self.extra_headers = extra_headers or {}
self.supports_dim_param = supports_dim_param
self.needs_client_truncate = needs_client_truncate
self.needs_client_normalize = needs_client_normalize
self.query_instruction = query_instruction or ""
self.max_batch_size = max(1, int(max_batch_size or 1))
if not self.api_key or self.api_key in ["", "YOUR API KEY", "YOUR_API_KEY"]:
raise ValueError("Embedding API key is not configured")
if dimensions is not None and dimensions > 0:
self._dimensions = dimensions
else:
# Legacy heuristic for OpenAI text-embedding-3-* family
self._dimensions = 1536 if "small" in model else 3072
def _call_api(self, input_data):
"""Call OpenAI-compatible /embeddings endpoint"""
import requests
url = f"{self.api_base}/embeddings"
headers = {
"Content-Type": "application/json",
"Authorization": f"Bearer {self.api_key}",
**self.extra_headers,
}
data = {
"input": input_data,
"model": self.model,
}
if self.supports_dim_param and self._dimensions:
data["dimensions"] = self._dimensions
try:
response = requests.post(url, headers=headers, json=data, timeout=EMBEDDING_HTTP_TIMEOUT)
response.raise_for_status()
return response.json()
except requests.exceptions.ConnectionError as e:
raise ConnectionError(
f"Failed to connect to embedding API at {url}. "
f"Please check network and api_base. Error: {str(e)}"
)
except requests.exceptions.Timeout as e:
raise TimeoutError(f"Embedding API request timed out. Error: {str(e)}")
except requests.exceptions.HTTPError as e:
if e.response.status_code == 401:
raise ValueError("Invalid embedding API key")
elif e.response.status_code == 429:
raise ValueError("Embedding API rate limit exceeded")
else:
raise ValueError(
f"Embedding API request failed: "
f"{e.response.status_code} - {e.response.text}"
)
def _post_process(self, raw: List[float]) -> List[float]:
"""Apply optional client-side truncation + normalization"""
vec = raw
if self.needs_client_truncate and self._dimensions and len(vec) > self._dimensions:
vec = vec[: self._dimensions]
if self.needs_client_normalize:
vec = _l2_normalize(vec)
return vec
def embed(self, text: str) -> List[float]:
"""Generate embedding (treated as document by default)"""
result = self._call_api(text)
return self._post_process(result["data"][0]["embedding"])
def embed_query(self, text: str) -> List[float]:
"""Generate embedding for a query (applies vendor instruction prefix if any)"""
if self.query_instruction:
text = f"{self.query_instruction}{text}"
return self.embed(text)
def embed_batch(self, texts: List[str]) -> List[List[float]]:
"""Generate embeddings for multiple documents.
Automatically paginates by self.max_batch_size so callers can pass any
number of texts. Order of returned vectors matches the input order.
"""
if not texts:
return []
out: List[List[float]] = []
step = self.max_batch_size
for i in range(0, len(texts), step):
chunk = texts[i:i + step]
result = self._call_api(chunk)
out.extend(self._post_process(item["embedding"]) for item in result["data"])
return out
@property
def dimensions(self) -> int:
return self._dimensions
class DoubaoEmbeddingProvider(EmbeddingProvider):
"""
Doubao (Volcengine Ark) multimodal embedding provider.
Doubao deprecated their OpenAI-compatible /v1/embeddings endpoint and
unified everything under /api/v3/embeddings/multimodal, which uses a
structured `input: [{type, text|image_url|video_url}, ...]` payload.
Notes:
* The endpoint produces ONE embedding per call (input list is multiple
modality parts of a single document, not a batch). embed_batch
therefore loops per-text — no native batch support.
* Native dimensions: 1024 or 2048 (default 1024 to align with other
Chinese vendors). No client-side truncation needed.
* Auth: Bearer ARK API key.
"""
def __init__(
self,
model: str,
api_key: Optional[str] = None,
api_base: Optional[str] = None,
extra_headers: Optional[dict] = None,
dimensions: Optional[int] = None,
):
self.model = model
self.api_key = api_key
self.api_base = api_base or "https://ark.cn-beijing.volces.com/api/v3"
self.extra_headers = extra_headers or {}
if not self.api_key or self.api_key in ["", "YOUR API KEY", "YOUR_API_KEY"]:
raise ValueError("Doubao embedding API key (ark_api_key) is not configured")
if dimensions in (1024, 2048):
self._dimensions = dimensions
elif dimensions is None:
self._dimensions = 1024
else:
raise ValueError(
f"Doubao embedding dimensions must be 1024 or 2048, got {dimensions}"
)
def _call_api(self, text: str) -> List[float]:
"""One call → one embedding. multimodal endpoint takes a single
document represented as a list of typed parts; we send a single
text part."""
import requests
url = f"{self.api_base}/embeddings/multimodal"
headers = {
"Content-Type": "application/json",
"Authorization": f"Bearer {self.api_key}",
**self.extra_headers,
}
payload = {
"model": self.model,
"input": [{"type": "text", "text": text}],
"dimensions": self._dimensions,
"encoding_format": "float",
}
try:
response = requests.post(url, headers=headers, json=payload, timeout=EMBEDDING_HTTP_TIMEOUT)
response.raise_for_status()
body = response.json()
except requests.exceptions.ConnectionError as e:
raise ConnectionError(
f"Failed to connect to Doubao embedding API at {url}. "
f"Please check network and api_base. Error: {str(e)}"
)
except requests.exceptions.Timeout as e:
raise TimeoutError(f"Doubao embedding API request timed out. Error: {str(e)}")
except requests.exceptions.HTTPError as e:
if e.response.status_code == 401:
raise ValueError("Invalid Doubao (ark) embedding API key")
elif e.response.status_code == 429:
raise ValueError("Doubao embedding API rate limit exceeded")
else:
raise ValueError(
f"Doubao embedding API request failed: "
f"{e.response.status_code} - {e.response.text}"
)
# Response shape per docs: {"data": {"embedding": [...]}}
data = body.get("data")
if isinstance(data, dict) and "embedding" in data:
return data["embedding"]
# Some providers wrap as a list of one — be defensive
if isinstance(data, list) and data and "embedding" in data[0]:
return data[0]["embedding"]
raise ValueError(f"Unexpected Doubao embedding response shape: {body}")
def embed(self, text: str) -> List[float]:
return self._call_api(text)
def embed_batch(self, texts: List[str]) -> List[List[float]]:
# Endpoint produces one embedding per call; loop. Order preserved.
return [self._call_api(t) for t in texts]
@property
def dimensions(self) -> int:
return self._dimensions
class EmbeddingCache:
"""In-memory cache for embeddings to avoid recomputation"""
def __init__(self):
self.cache = {}
def get(self, text: str, provider: str, model: str) -> Optional[List[float]]:
key = self._compute_key(text, provider, model)
return self.cache.get(key)
def put(self, text: str, provider: str, model: str, embedding: List[float]):
key = self._compute_key(text, provider, model)
self.cache[key] = embedding
@staticmethod
def _compute_key(text: str, provider: str, model: str) -> str:
content = f"{provider}:{model}:{text}"
return hashlib.md5(content.encode("utf-8")).hexdigest()
def clear(self):
self.cache.clear()
def create_embedding_provider(
provider: str = "openai",
model: Optional[str] = None,
api_key: Optional[str] = None,
api_base: Optional[str] = None,
extra_headers: Optional[dict] = None,
dimensions: Optional[int] = None,
) -> EmbeddingProvider:
"""
Factory function to create an embedding provider.
Backward compatible: when called with provider in {"openai", "linkai"}
and no `dimensions` arg, behaves exactly as before (1536-dim OpenAI).
New providers ("dashscope", "doubao", "zhipu") require explicit configuration
and use the unified 1024-dim defaults from EMBEDDING_VENDORS.
Args:
provider: Vendor key (one of EMBEDDING_VENDORS)
model: Model name (uses vendor default if None)
api_key: API key (required)
api_base: API base URL (uses vendor default if None)
extra_headers: Optional extra HTTP headers
dimensions: Target output dimension (uses vendor default if None)
Returns:
EmbeddingProvider instance
"""
meta = EMBEDDING_VENDORS.get(provider)
if meta is None:
raise ValueError(
f"Unsupported embedding provider: {provider}. "
f"Supported: {sorted(EMBEDDING_VENDORS.keys())}"
)
# Doubao uses a non-OpenAI-compatible multimodal endpoint.
if meta.get("provider_class") == "doubao":
final_dim = dimensions if (dimensions and dimensions > 0) else meta["default_dimensions"]
return DoubaoEmbeddingProvider(
model=model or meta["default_model"],
api_key=api_key,
api_base=api_base or meta["default_base_url"],
extra_headers=extra_headers,
dimensions=final_dim,
)
# Legacy two-arg call for openai/linkai keeps 1536-dim default behavior
# so existing data isn't invalidated.
is_legacy_call = (
provider in ("openai", "linkai")
and dimensions is None
)
if is_legacy_call:
return OpenAIEmbeddingProvider(
model=model or "text-embedding-3-small",
api_key=api_key,
api_base=api_base,
extra_headers=extra_headers,
)
final_dim = dimensions if (dimensions and dimensions > 0) else meta["default_dimensions"]
return OpenAIEmbeddingProvider(
model=model or meta["default_model"],
api_key=api_key,
api_base=api_base or meta["default_base_url"],
extra_headers=extra_headers,
dimensions=final_dim,
supports_dim_param=meta["supports_dim_param"],
needs_client_truncate=meta["needs_client_truncate"],
needs_client_normalize=meta["needs_client_normalize"],
query_instruction=meta["query_instruction"],
max_batch_size=meta.get("max_batch_size", 256),
)

191
agent/memory/embedding/rebuild.py Normal file
View File

@@ -0,0 +1,191 @@
"""
Rebuild memory vector index.
Recommended entry point (in-chat, while agent is running):
/memory rebuild-index
Backward-compatible CLI entry (must run from project root):
python -m agent.memory.rebuild_index
What it does:
1. Probes the embedding endpoint with a tiny call to fail fast on
bad provider/model/key — before touching the index.
2. Clears the SQLite chunks/files tables (workspace markdown stays intact).
3. Runs a fresh sync, regenerating embeddings with the currently configured
provider/model/dimensions.
This is the only safe way to switch embedding_provider after the existing
index has been populated by a different-dim model.
"""
from __future__ import annotations
import asyncio
import sys
from dataclasses import dataclass
from typing import Optional
from common.log import logger
from common.utils import expand_path
@dataclass
class RebuildResult:
"""Outcome of a rebuild_in_process() call"""
ok: bool
removed: int = 0
chunks: int = 0
files: int = 0
error: Optional[str] = None
def clear_index(db_path, storage=None) -> int:
"""Wipe chunks/files, reset FTS5, and clean up any legacy state file.
Args:
db_path: Path of the index DB (also used to locate the legacy state
file for migration cleanup, and — when *storage* is None — to
open a fresh connection).
storage: Optional pre-opened MemoryStorage. When provided we reuse it
so the live connection's triggers stay in sync — opening a second
connection would leave the original one's triggers pointing at a
DROP'd chunks_fts table.
We reset (DROP+recreate) chunks_fts because its shadow tables can become
inconsistent across rebuild cycles, causing bm25() / ORDER BY rank to
raise "database disk image is malformed" even when raw MATCH still works.
Returns number of chunks removed.
"""
from agent.memory.embedding.state import cleanup_legacy_state_file
from agent.memory.storage import MemoryStorage
owns_storage = storage is None
if owns_storage:
storage = MemoryStorage(db_path)
try:
before = storage.conn.execute("SELECT COUNT(*) FROM chunks").fetchone()[0]
storage.conn.execute("DELETE FROM chunks")
storage.conn.execute("DELETE FROM files")
storage.conn.commit()
storage.reset_fts5()
finally:
if owns_storage:
storage.close()
cleanup_legacy_state_file(db_path)
return int(before)
def rebuild_in_process(memory_manager) -> RebuildResult:
"""
Rebuild the index using an existing, fully-initialized MemoryManager.
Used by the in-chat /memory rebuild-index command. The caller already has
config loaded, embedding_provider built, and (optionally) the agent
running, so we only need to:
1. Clear chunks/files + state on the manager's storage.
2. Re-sync (force=True).
NOTE: caller must ensure memory_manager.embedding_provider is set, otherwise
sync() will silently skip embedding generation.
"""
if memory_manager is None:
return RebuildResult(ok=False, error="memory_manager is None")
if memory_manager.embedding_provider is None:
return RebuildResult(ok=False, error="embedding_provider is not initialized")
# Probe the embedding endpoint BEFORE clearing the index. A bad
# provider/model/key would otherwise leave the user with an empty index
# that not even keyword search can serve.
try:
memory_manager.embedding_provider.embed_query("ping")
except Exception as e:
logger.error(f"[RebuildIndex] embedding probe failed, aborting rebuild: {e}")
return RebuildResult(ok=False, error=f"embedding endpoint not reachable: {e}")
db_path = memory_manager.config.get_db_path()
try:
removed = clear_index(db_path, storage=memory_manager.storage)
except Exception as e:
logger.exception("[RebuildIndex] clear_index failed")
return RebuildResult(ok=False, error=f"clear failed: {e}")
try:
asyncio.run(memory_manager.sync(force=True))
except RuntimeError:
# Already inside a running event loop (rare in chat handler thread).
loop = asyncio.new_event_loop()
try:
loop.run_until_complete(memory_manager.sync(force=True))
finally:
loop.close()
except Exception as e:
logger.exception("[RebuildIndex] sync failed")
return RebuildResult(ok=False, removed=removed, error=f"re-embed failed: {e}")
stats = memory_manager.storage.get_stats()
chunks = int(stats.get("chunks", 0))
embedded = int(stats.get("embedded", 0))
# sync() degrades to "no embeddings" on batch failure so keyword search
# still works at startup — but in a /rebuild-index request the user
# explicitly asked for vectors. Surface that as a failure.
if chunks > 0 and embedded == 0:
return RebuildResult(
ok=False,
removed=removed,
chunks=chunks,
files=int(stats.get("files", 0)),
error=(
"embedding API failed during sync; index now has chunks but no "
"vectors. Check embedding provider/model/key and retry."
),
)
return RebuildResult(
ok=True,
removed=removed,
chunks=chunks,
files=int(stats.get("files", 0)),
)
def main() -> int:
"""Standalone CLI entry. Must be run from project root (relative config path)."""
from config import conf, load_config
from agent.memory import MemoryConfig, MemoryManager
load_config()
workspace_root = expand_path(conf().get("agent_workspace", "~/cow"))
memory_config = MemoryConfig(workspace_root=workspace_root)
logger.info(f"[RebuildIndex] Workspace: {workspace_root}")
logger.info(f"[RebuildIndex] Index db: {memory_config.get_db_path()}")
from bridge.agent_initializer import AgentInitializer
initializer = AgentInitializer(bridge=None, agent_bridge=None)
embedding_provider = initializer._init_embedding_provider(memory_config, session_id=None)
if embedding_provider is None:
logger.error(
"[RebuildIndex] No embedding provider could be initialized. "
"Check your config.json. Aborting rebuild."
)
return 1
manager = MemoryManager(memory_config, embedding_provider=embedding_provider)
result = rebuild_in_process(manager)
if not result.ok:
logger.error(f"[RebuildIndex] {result.error}")
return 1
logger.info(
f"[RebuildIndex] Done. removed={result.removed}, "
f"chunks={result.chunks}, files={result.files}"
)
return 0
if __name__ == "__main__":
sys.exit(main())

47
agent/memory/embedding/state.py Normal file
View File

@@ -0,0 +1,47 @@
"""
Embedding-related index utilities.
We don't keep a sidecar state file — the SQLite index is the source of truth
and config.json is the source of intent. The two functions below are the
only things needing on-disk awareness:
detect_index_dim : read the dim of stored vectors (display-only)
cleanup_legacy_state_file: remove old embedding_state.json from earlier
versions; safe no-op when absent.
"""
from __future__ import annotations
import json
import os
from pathlib import Path
from typing import Optional, Union
PathLike = Union[str, os.PathLike]
def detect_index_dim(storage) -> Optional[int]:
"""Return the dim of the first stored embedding, or None if the index
has no embeddings. Used by /memory status."""
try:
row = storage.conn.execute(
"SELECT embedding FROM chunks WHERE embedding IS NOT NULL LIMIT 1"
).fetchone()
except Exception:
return None
if not row or not row["embedding"]:
return None
try:
emb = json.loads(row["embedding"])
return len(emb) if isinstance(emb, list) else None
except (json.JSONDecodeError, TypeError):
return None
def cleanup_legacy_state_file(db_path: PathLike) -> None:
"""Remove old embedding_state.json files from earlier versions.
Safe to call repeatedly; no-op if the file is absent."""
legacy = Path(db_path).parent / "embedding_state.json"
try:
legacy.unlink(missing_ok=True)
except Exception:
pass

299
agent/memory/manager.py
View File

@@ -13,7 +13,7 @@ from datetime import datetime, timedelta
from agent.memory.config import MemoryConfig, get_default_memory_config
from agent.memory.storage import MemoryStorage, MemoryChunk, SearchResult
from agent.memory.chunker import TextChunker
from agent.memory.embedding import create_embedding_provider, EmbeddingProvider
from agent.memory.embedding import EmbeddingProvider
from agent.memory.summarizer import MemoryFlushManager, create_memory_files_if_needed
@@ -50,49 +50,17 @@ class MemoryManager:
overlap_tokens=self.config.chunk_overlap_tokens
)
# Initialize embedding provider (optional, prefer OpenAI, fallback to LinkAI)
self.embedding_provider = None
if embedding_provider:
self.embedding_provider = embedding_provider
else:
# Try OpenAI first
try:
api_key = os.environ.get('OPENAI_API_KEY')
api_base = os.environ.get('OPENAI_API_BASE')
if api_key:
self.embedding_provider = create_embedding_provider(
provider="openai",
model=self.config.embedding_model,
api_key=api_key,
api_base=api_base
)
except Exception as e:
from common.log import logger
logger.warning(f"[MemoryManager] OpenAI embedding failed: {e}")
# Fallback to LinkAI
if self.embedding_provider is None:
try:
linkai_key = os.environ.get('LINKAI_API_KEY')
linkai_base = os.environ.get('LINKAI_API_BASE', 'https://api.link-ai.tech')
if linkai_key:
from common.utils import get_cloud_headers
cloud_headers = get_cloud_headers(linkai_key)
cloud_headers.pop("Authorization", None)
self.embedding_provider = create_embedding_provider(
provider="linkai",
model=self.config.embedding_model,
api_key=linkai_key,
api_base=f"{linkai_base}/v1",
extra_headers=cloud_headers,
)
except Exception as e:
from common.log import logger
logger.warning(f"[MemoryManager] LinkAI embedding failed: {e}")
if self.embedding_provider is None:
from common.log import logger
logger.info(f"[MemoryManager] Memory will work with keyword search only (no vector search)")
# Embedding provider is owned by the caller (agent_initializer is the
# canonical entry point and handles legacy/explicit + state validation).
# When None is passed, memory degrades to keyword-only search instead
# of silently re-initializing a vendor here, which would bypass the
# caller's state checks and risk corrupting the index.
self.embedding_provider = embedding_provider
if self.embedding_provider is None:
from common.log import logger
logger.info(
"[MemoryManager] No embedding provider; memory will use keyword search only"
)
# Initialize memory flush manager
workspace_dir = self.config.get_workspace()
@@ -153,12 +121,14 @@ class MemoryManager:
if self.config.sync_on_search and self._dirty:
await self.sync()
# Perform vector search (if embedding provider available)
from common.log import logger
# Perform vector search (if embedding provider available).
# Failures degrade silently to keyword-only — no exception is raised.
vector_results = []
if self.embedding_provider:
try:
from common.log import logger
query_embedding = self.embedding_provider.embed(query)
query_embedding = self.embedding_provider.embed_query(query)
vector_results = self.storage.search_vector(
query_embedding=query_embedding,
user_id=user_id,
@@ -167,19 +137,19 @@ class MemoryManager:
)
logger.info(f"[MemoryManager] Vector search found {len(vector_results)} results for query: {query}")
except Exception as e:
from common.log import logger
logger.warning(f"[MemoryManager] Vector search failed: {e}")
# Perform keyword search
logger.error(
f"[MemoryManager] Vector search failed, falling back to keyword-only: {e}"
)
# Perform keyword search (also runs as fallback when vector failed)
keyword_results = self.storage.search_keyword(
query=query,
user_id=user_id,
scopes=scopes,
limit=max_results * 2
)
from common.log import logger
logger.info(f"[MemoryManager] Keyword search found {len(keyword_results)} results for query: {query}")
# Merge results
merged = self._merge_results(
vector_results,
@@ -187,7 +157,7 @@ class MemoryManager:
self.config.vector_weight,
self.config.keyword_weight
)
# Filter by min score and limit
filtered = [r for r in merged if r.score >= min_score]
return filtered[:max_results]
@@ -269,132 +239,157 @@ class MemoryManager:
async def sync(self, force: bool = False):
"""
Synchronize memory from files
Synchronize memory from files.
Two-pass design to amortize embedding HTTP cost:
1. Walk all files, chunk those whose hash changed, collect pending
chunks across files. No embedding calls yet.
2. Run a single embed_batch over the union of pending chunks (the
provider auto-paginates by vendor cap), then persist per-file.
For workspaces with many small files (101 files / ~1 chunk each), this
cuts ~100 HTTP calls down to ~ceil(total_chunks / vendor_cap).
Args:
force: Force full reindex
"""
memory_dir = self.config.get_memory_dir()
workspace_dir = self.config.get_workspace()
# Scan MEMORY.md (workspace root)
files_to_scan: List[tuple] = [] # (file_path, source, scope, user_id)
memory_file = Path(workspace_dir) / "MEMORY.md"
if memory_file.exists():
await self._sync_file(memory_file, "memory", "shared", None)
# Scan memory directory (including daily summaries)
files_to_scan.append((memory_file, "memory", "shared", None))
if memory_dir.exists():
for file_path in memory_dir.rglob("*.md"):
# Skip hidden directories (e.g. .dreams/)
if any(part.startswith('.') for part in file_path.relative_to(workspace_dir).parts):
continue
# Determine scope and user_id from path
rel_path = file_path.relative_to(workspace_dir)
parts = rel_path.parts
# Check if it's in daily summary directory
if "daily" in parts:
# Daily summary files
if "users" in parts or len(parts) > 3:
# User-scoped daily summary: memory/daily/{user_id}/2024-01-29.md
user_idx = parts.index("daily") + 1
user_id = parts[user_idx] if user_idx < len(parts) else None
rel_parts = file_path.relative_to(workspace_dir).parts
if "daily" in rel_parts:
if "users" in rel_parts or len(rel_parts) > 3:
user_idx = rel_parts.index("daily") + 1
user_id = rel_parts[user_idx] if user_idx < len(rel_parts) else None
scope = "user"
else:
# Shared daily summary: memory/daily/2024-01-29.md
user_id = None
scope = "shared"
elif "users" in parts:
# User-scoped memory
user_idx = parts.index("users") + 1
user_id = parts[user_idx] if user_idx < len(parts) else None
elif "users" in rel_parts:
user_idx = rel_parts.index("users") + 1
user_id = rel_parts[user_idx] if user_idx < len(rel_parts) else None
scope = "user"
else:
# Shared memory
user_id = None
scope = "shared"
await self._sync_file(file_path, "memory", scope, user_id)
files_to_scan.append((file_path, "memory", scope, user_id))
# Scan knowledge directory (structured knowledge wiki)
from config import conf
if conf().get("knowledge", True):
knowledge_dir = Path(workspace_dir) / "knowledge"
if knowledge_dir.exists():
for file_path in knowledge_dir.rglob("*.md"):
await self._sync_file(file_path, "knowledge", "shared", None)
self._dirty = False
async def _sync_file(
self,
file_path: Path,
source: str,
scope: str,
user_id: Optional[str]
):
"""Sync a single file"""
# Compute file hash
content = file_path.read_text(encoding='utf-8')
file_hash = MemoryStorage.compute_hash(content)
# Get relative path
workspace_dir = self.config.get_workspace()
rel_path = str(file_path.relative_to(workspace_dir))
# Check if file changed
stored_hash = self.storage.get_file_hash(rel_path)
if stored_hash == file_hash:
return # No changes
# Delete old chunks
self.storage.delete_by_path(rel_path)
# Chunk and embed
chunks = self.chunker.chunk_text(content)
if not chunks:
files_to_scan.append((file_path, "knowledge", "shared", None))
# Pass 1: inline chunking + change detection. Inlined (instead of
# calling self._prepare_file_for_sync) so this method does not depend
# on any sibling helpers — keeps it robust against partial reloads
# where the class object is older than the method's source.
pending: List[Dict[str, Any]] = []
workspace_dir_path = self.config.get_workspace()
for file_path, source, scope, user_id in files_to_scan:
try:
content = file_path.read_text(encoding='utf-8')
except Exception:
continue
file_hash = MemoryStorage.compute_hash(content)
rel_path = str(file_path.relative_to(workspace_dir_path))
if self.storage.get_file_hash(rel_path) == file_hash:
continue
chunks = self.chunker.chunk_text(content)
if not chunks:
continue
pending.append({
"file_path": file_path,
"rel_path": rel_path,
"source": source,
"scope": scope,
"user_id": user_id,
"file_hash": file_hash,
"chunks": chunks,
"texts": [c.text for c in chunks],
})
if not pending:
self._dirty = False
return
texts = [chunk.text for chunk in chunks]
if self.embedding_provider:
embeddings = self.embedding_provider.embed_batch(texts)
# Pass 2: single batched embed across all pending chunks.
# CRITICAL: never touch the index until we hold valid embeddings.
# If embed_batch fails, leave the existing index intact (chunks +
# file_hash) so the next sync will retry the same files. Writing
# NULL embeddings + updating file_hash here would mark the file as
# "successfully synced" and silently strand it without vectors.
all_texts: List[str] = []
for entry in pending:
all_texts.extend(entry["texts"])
if not self.embedding_provider:
# No provider configured at all (legacy keyword-only). Persist
# chunks without embeddings — this is the user's intent.
all_embeddings: List[Optional[List[float]]] = [None] * len(all_texts)
else:
embeddings = [None] * len(texts)
# Create memory chunks
memory_chunks = []
for chunk, embedding in zip(chunks, embeddings):
chunk_id = self._generate_chunk_id(rel_path, chunk.start_line, chunk.end_line)
chunk_hash = MemoryStorage.compute_hash(chunk.text)
memory_chunks.append(MemoryChunk(
id=chunk_id,
user_id=user_id,
scope=scope,
source=source,
try:
all_embeddings = self.embedding_provider.embed_batch(all_texts)
except Exception as e:
from common.log import logger
logger.error(
f"[MemoryManager] Batch embedding failed for {len(all_texts)} "
f"chunks across {len(pending)} files: {e}. "
f"Index left untouched; will retry on next sync."
)
# Bail before touching storage. self._dirty stays True so
# callers know there is pending work.
return
# Pass 3: inline persist — same self-contained reasoning as Pass 1.
cursor = 0
for entry in pending:
n = len(entry["texts"])
entry_embeddings = all_embeddings[cursor:cursor + n]
cursor += n
rel_path = entry["rel_path"]
self.storage.delete_by_path(rel_path)
memory_chunks = []
for chunk, embedding in zip(entry["chunks"], entry_embeddings):
chunk_id = self._generate_chunk_id(rel_path, chunk.start_line, chunk.end_line)
chunk_hash = MemoryStorage.compute_hash(chunk.text)
memory_chunks.append(MemoryChunk(
id=chunk_id,
user_id=entry["user_id"],
scope=entry["scope"],
source=entry["source"],
path=rel_path,
start_line=chunk.start_line,
end_line=chunk.end_line,
text=chunk.text,
embedding=embedding,
hash=chunk_hash,
metadata=None,
))
self.storage.save_chunks_batch(memory_chunks)
stat = entry["file_path"].stat()
self.storage.update_file_metadata(
path=rel_path,
start_line=chunk.start_line,
end_line=chunk.end_line,
text=chunk.text,
embedding=embedding,
hash=chunk_hash,
metadata=None
))
# Save
self.storage.save_chunks_batch(memory_chunks)
# Update file metadata
stat = file_path.stat()
self.storage.update_file_metadata(
path=rel_path,
source=source,
file_hash=file_hash,
mtime=int(stat.st_mtime),
size=stat.st_size
)
source=entry["source"],
file_hash=entry["file_hash"],
mtime=int(stat.st_mtime),
size=stat.st_size,
)
self._dirty = False
def flush_memory(
self,
messages: list,

14
agent/memory/rebuild_index.py Normal file
View File

@@ -0,0 +1,14 @@
"""
Backward-compatible shim for the legacy entry point:
python -m agent.memory.rebuild_index
The implementation now lives in agent.memory.embedding.rebuild.
Prefer using `/memory rebuild-index` in chat going forward.
"""
from agent.memory.embedding.rebuild import main
if __name__ == "__main__":
import sys
sys.exit(main())

261
agent/memory/storage.py
View File

@@ -144,45 +144,37 @@ class MemoryStorage:
ON chunks(path, hash)
""")
# Create FTS5 virtual table for keyword search (only if supported)
# Create FTS5 virtual table + triggers (only if supported).
# Self-heal: if the previous process crashed mid-rebuild and left
# triggers pointing at a missing chunks_fts (or vice versa), wipe
# both sides and recreate cleanly. Otherwise next chunks INSERT
# will fail with "no such table: chunks_fts".
if self.fts5_available:
# Use default unicode61 tokenizer (stable and compatible)
# For CJK support, we'll use LIKE queries as fallback
self.conn.execute("""
CREATE VIRTUAL TABLE IF NOT EXISTS chunks_fts USING fts5(
text,
id UNINDEXED,
user_id UNINDEXED,
path UNINDEXED,
source UNINDEXED,
scope UNINDEXED,
content='chunks',
content_rowid='rowid'
if self._fts5_state_inconsistent():
from common.log import logger
logger.warning(
"[MemoryStorage] FTS5 state inconsistent (triggers/table mismatch). "
"Resetting chunks_fts to recover."
)
""")
# Create triggers to keep FTS in sync
self.conn.execute("""
CREATE TRIGGER IF NOT EXISTS chunks_ai AFTER INSERT ON chunks BEGIN
INSERT INTO chunks_fts(rowid, text, id, user_id, path, source, scope)
VALUES (new.rowid, new.text, new.id, new.user_id, new.path, new.source, new.scope);
END
""")
self.conn.execute("""
CREATE TRIGGER IF NOT EXISTS chunks_ad AFTER DELETE ON chunks BEGIN
DELETE FROM chunks_fts WHERE rowid = old.rowid;
END
""")
self.conn.execute("""
CREATE TRIGGER IF NOT EXISTS chunks_au AFTER UPDATE ON chunks BEGIN
UPDATE chunks_fts SET text = new.text, id = new.id,
user_id = new.user_id, path = new.path, source = new.source, scope = new.scope
WHERE rowid = new.rowid;
END
""")
self.conn.execute("DROP TRIGGER IF EXISTS chunks_ai")
self.conn.execute("DROP TRIGGER IF EXISTS chunks_ad")
self.conn.execute("DROP TRIGGER IF EXISTS chunks_au")
self.conn.execute("DROP TABLE IF EXISTS chunks_fts")
self.conn.commit()
self._create_fts5_objects()
# Probe FTS5 shadow tables. The schema may be intact but the
# internal _data/_idx/_docsize blob can still be corrupt — that
# surfaces as "database disk image is malformed" on bm25 / MATCH.
# We rebuild from the chunks table when that happens; data isn't
# lost because chunks (the content table) is the source of truth.
if self._fts5_shadow_corrupt():
from common.log import logger
logger.warning(
"[MemoryStorage] FTS5 shadow tables corrupt; rebuilding from chunks."
)
self._rebuild_fts5_from_chunks()
# Create files metadata table
self.conn.execute("""
CREATE TABLE IF NOT EXISTS files (
@@ -196,7 +188,116 @@ class MemoryStorage:
""")
self.conn.commit()
def _fts5_state_inconsistent(self) -> bool:
"""Detect a half-broken FTS5 setup (e.g. trigger exists but table doesn't)."""
try:
row = self.conn.execute(
"SELECT name FROM sqlite_master WHERE type='table' AND name='chunks_fts'"
).fetchone()
table_exists = row is not None
row = self.conn.execute(
"SELECT COUNT(*) FROM sqlite_master WHERE type='trigger' "
"AND name IN ('chunks_ai','chunks_ad','chunks_au')"
).fetchone()
trigger_count = int(row[0]) if row else 0
except Exception:
return False
# Healthy = both present (3 triggers + table) or both absent.
return table_exists != (trigger_count > 0)
def _create_fts5_objects(self):
"""Create chunks_fts virtual table and the 3 sync triggers.
Idempotent: uses IF NOT EXISTS. Caller must hold self.conn.
"""
self.conn.execute("""
CREATE VIRTUAL TABLE IF NOT EXISTS chunks_fts USING fts5(
text,
id UNINDEXED,
user_id UNINDEXED,
path UNINDEXED,
source UNINDEXED,
scope UNINDEXED,
content='chunks',
content_rowid='rowid'
)
""")
self.conn.execute("""
CREATE TRIGGER IF NOT EXISTS chunks_ai AFTER INSERT ON chunks BEGIN
INSERT INTO chunks_fts(rowid, text, id, user_id, path, source, scope)
VALUES (new.rowid, new.text, new.id, new.user_id, new.path, new.source, new.scope);
END
""")
self.conn.execute("""
CREATE TRIGGER IF NOT EXISTS chunks_ad AFTER DELETE ON chunks BEGIN
DELETE FROM chunks_fts WHERE rowid = old.rowid;
END
""")
self.conn.execute("""
CREATE TRIGGER IF NOT EXISTS chunks_au AFTER UPDATE ON chunks BEGIN
UPDATE chunks_fts SET text = new.text, id = new.id,
user_id = new.user_id, path = new.path,
source = new.source, scope = new.scope
WHERE rowid = new.rowid;
END
""")
def reset_fts5(self):
"""Drop and recreate chunks_fts + triggers in one transaction.
Used by rebuild_index to recover from FTS5 shadow-table corruption
(bm25/ORDER BY rank may raise "database disk image is malformed"
even when raw MATCH still works).
Triggers must be dropped first; otherwise the next chunks INSERT/DELETE
on the existing connection will hit "no such table: chunks_fts".
"""
if not self.fts5_available:
return
self.conn.execute("DROP TRIGGER IF EXISTS chunks_ai")
self.conn.execute("DROP TRIGGER IF EXISTS chunks_ad")
self.conn.execute("DROP TRIGGER IF EXISTS chunks_au")
self.conn.execute("DROP TABLE IF EXISTS chunks_fts")
self._create_fts5_objects()
self.conn.commit()
def _fts5_shadow_corrupt(self) -> bool:
"""Probe whether bm25 over chunks_fts errors out at startup.
Schema (table + triggers) can be intact while the underlying
FTS5 shadow blobs are malformed — typically because the previous
process crashed mid-write or wrote with a different SQLite build.
A cheap MATCH probe surfaces it immediately."""
try:
self.conn.execute(
"SELECT bm25(chunks_fts) FROM chunks_fts WHERE chunks_fts MATCH 'a' LIMIT 1"
).fetchone()
return False
except sqlite3.DatabaseError as e:
msg = str(e).lower()
return "malformed" in msg or "corrupt" in msg
except Exception:
# Any other error (e.g. table missing) is handled by the
# state-inconsistent path; treat as healthy here.
return False
def _rebuild_fts5_from_chunks(self):
"""Drop FTS5, recreate it, then INSERT every row from chunks.
Safe data-wise: chunks (the content table) is the source of truth.
Done in one transaction so a crash leaves either fully old or fully
new state, not a partial rebuild.
"""
# Reset schema first; this clears any malformed shadow blobs.
self.reset_fts5()
# Re-feed content. Triggers handle future writes automatically.
self.conn.execute("""
INSERT INTO chunks_fts(rowid, text, id, user_id, path, source, scope)
SELECT rowid, text, id, user_id, path, source, scope FROM chunks
""")
self.conn.commit()
def save_chunk(self, chunk: MemoryChunk):
"""Save a memory chunk"""
self.conn.execute("""
@@ -283,13 +384,26 @@ class MemoryStorage:
"""
rows = self.conn.execute(query, params).fetchall()
# Calculate cosine similarity
# Calculate cosine similarity. We probe the first row's dim to fail
# loudly on a query/index dim mismatch — otherwise every doc would
# score 0 silently, leaving the user wondering why search broke.
results = []
query_dim = len(query_embedding)
if rows:
first = json.loads(rows[0]['embedding'])
if isinstance(first, list) and len(first) != query_dim:
raise ValueError(
f"Embedding dim mismatch: query is {query_dim}-dim but "
f"index stores {len(first)}-dim vectors. The configured "
f"embedding model differs from the one that built the "
f"index — run /memory rebuild-index to re-embed."
)
for row in rows:
embedding = json.loads(row['embedding'])
similarity = self._cosine_similarity(query_embedding, embedding)
if similarity > 0:
results.append((similarity, row))
@@ -319,27 +433,24 @@ class MemoryStorage:
) -> List[SearchResult]:
"""
Keyword search using FTS5 + LIKE fallback
Strategy:
1. If FTS5 available: Try FTS5 search first (good for English and word-based languages)
2. If no FTS5 or no results and query contains CJK: Use LIKE search
1. If FTS5 available and healthy: try FTS5 first
2. Always fall back to LIKE for CJK queries
3. If FTS5 fails OR returns empty for non-CJK, also try LIKE so a
broken FTS5 shadow table doesn't silently kill keyword search.
"""
if scopes is None:
scopes = ["shared"]
if user_id:
scopes.append("user")
# Try FTS5 search first (if available)
if self.fts5_available:
fts_results = self._search_fts5(query, user_id, scopes, limit)
if fts_results:
return fts_results
# Fallback to LIKE search (always for CJK, or if FTS5 not available)
if not self.fts5_available or MemoryStorage._contains_cjk(query):
return self._search_like(query, user_id, scopes, limit)
return []
return self._search_like(query, user_id, scopes, limit)
def _search_fts5(
self,
@@ -394,7 +505,11 @@ class MemoryStorage:
)
for row in rows
]
except Exception:
except Exception as e:
from common.log import logger
logger.error(
f"[MemoryStorage] FTS5 search failed (caller will fall back to LIKE): {e}"
)
return []
def _search_like(
@@ -404,21 +519,28 @@ class MemoryStorage:
scopes: List[str],
limit: int
) -> List[SearchResult]:
"""LIKE-based search for CJK characters"""
"""LIKE-based search.
Used as the keyword-search fallback when FTS5 is unavailable, fails,
or returns empty. Supports both CJK runs and ASCII word tokens so it
can serve as a true safety net for any query.
"""
import re
# Extract CJK words (2+ characters)
# CJK runs (2+ chars) + ASCII word tokens (3+ chars to avoid noise)
cjk_words = re.findall(r'[\u4e00-\u9fff]{2,}', query)
if not cjk_words:
ascii_words = [t for t in re.findall(r'[A-Za-z0-9_]+', query) if len(t) >= 3]
words = cjk_words + ascii_words
if not words:
return []
scope_placeholders = ','.join('?' * len(scopes))
# Build LIKE conditions for each word
# Build LIKE conditions for each word (case-insensitive for ASCII)
like_conditions = []
params = []
for word in cjk_words:
like_conditions.append("text LIKE ?")
params.append(f'%{word}%')
for word in words:
like_conditions.append("LOWER(text) LIKE ?")
params.append(f'%{word.lower()}%')
where_clause = ' OR '.join(like_conditions)
params.extend(scopes)
@@ -455,7 +577,9 @@ class MemoryStorage:
)
for row in rows
]
except Exception:
except Exception as e:
from common.log import logger
logger.error(f"[MemoryStorage] LIKE search failed: {e}")
return []
def delete_by_path(self, path: str):
@@ -485,14 +609,19 @@ class MemoryStorage:
chunks_count = self.conn.execute("""
SELECT COUNT(*) as cnt FROM chunks
""").fetchone()['cnt']
files_count = self.conn.execute("""
SELECT COUNT(*) as cnt FROM files
""").fetchone()['cnt']
embedded_count = self.conn.execute("""
SELECT COUNT(*) as cnt FROM chunks WHERE embedding IS NOT NULL
""").fetchone()['cnt']
return {
'chunks': chunks_count,
'files': files_count
'files': files_count,
'embedded': embedded_count,
}
def close(self):