kaiyuan/chatgpt-on-wechat
1
0
Fork 0
mirror of https://github.com/zhayujie/chatgpt-on-wechat.git synced 2026-07-19 12:47:25 +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

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