From 283a8bf832fab26676649326984989e32398f48d Mon Sep 17 00:00:00 2001 From: dengfuping Date: Fri, 10 Apr 2026 17:43:03 +0800 Subject: [PATCH] refactor(seekdb): add pyseekdb_conn and remote-only host/port config --- reme/core/file_store/seekdb_file_store.py | 59 ++++++++---- reme/core/utils/__init__.py | 4 + reme/core/utils/pyseekdb_conn.py | 60 ++++++++++++ reme/core/vector_store/seekdb_vector_store.py | 91 ++++++++++++++----- tests/test_file_store.py | 31 +++++-- tests/test_vector_store.py | 33 +++++-- 6 files changed, 225 insertions(+), 53 deletions(-) create mode 100644 reme/core/utils/pyseekdb_conn.py diff --git a/reme/core/file_store/seekdb_file_store.py b/reme/core/file_store/seekdb_file_store.py index 1fcce094..4dd38178 100644 --- a/reme/core/file_store/seekdb_file_store.py +++ b/reme/core/file_store/seekdb_file_store.py @@ -1,4 +1,9 @@ -"""seekdb storage backend for file store.""" +"""seekdb storage backend for file store. + +``pyseekdb.Client`` supports **embedded** (``path``) and **remote** OceanBase / +seekdb (``host`` / ``port`` / credentials). SQL-table-oriented helpers can still +use **pyobvector** via ``ObVecFileStore`` if needed. +""" import time from pathlib import Path @@ -8,6 +13,11 @@ from loguru import logger from .base_file_store import BaseFileStore from ..enumeration import MemorySource from ..schema import FileMetadata, MemoryChunk, MemorySearchResult +from ..utils.pyseekdb_conn import ( + DEFAULT_SEEKDB_TENANT, + admin_kwargs_from_client_kwargs, + build_pyseekdb_client_kwargs, +) try: import pyseekdb @@ -28,14 +38,22 @@ def _escape_sql(s: str) -> str: class SeekdbFileStore(BaseFileStore): - """seekdb file storage with vector and full-text search (embedded mode only). + """seekdb file storage with vector and full-text search via ``pyseekdb``. - File metadata is stored in a DB table (same as SqliteFileStore); uses pyseekdb - BaseClient execute/_execute for SQL. No env vars, uses db_path and store_name only. + **Embedded** (default): optional ``path`` for the data directory; if omitted, pyseekdb + uses its default (typically ``./seekdb.db``). **Remote**: ``host`` / ``port`` plus auth. + + File metadata is in a SQL table like SqliteFileStore; raw SQL uses ``execute``/``_execute``. """ def __init__( self, + host: str | None = None, + port: int | None = None, + tenant: str = DEFAULT_SEEKDB_TENANT, + user: str | None = None, + password: str = "", + path: str | None = None, **kwargs, ): if not PYSEEKDB_AVAILABLE: @@ -48,6 +66,16 @@ class SeekdbFileStore(BaseFileStore): self.client: "pyseekdb.Client | None" = None self.collection = None + self._is_remote, self._client_kw = build_pyseekdb_client_kwargs( + path=None if (host and host.strip()) else path, + database=self.store_name, + host=host, + port=port, + tenant=tenant, + user=user, + password=password, + ) + @property def collection_name(self) -> str: """Collection name for chunks.""" @@ -59,8 +87,8 @@ class SeekdbFileStore(BaseFileStore): return f"files_{self.store_name}" def _client_kwargs(self) -> dict: - """Build Client kwargs for embedded mode from db_path and store_name (same as Chroma/SQLite).""" - return {"path": str(self.db_path / "seekdb"), "database": self.store_name} + """Kwargs for ``pyseekdb.Client`` (embedded or remote).""" + return self._client_kw def _sql_client(self): """Underlying BaseClient for raw SQL (pyseekdb Client proxy exposes _server).""" @@ -99,16 +127,15 @@ class SeekdbFileStore(BaseFileStore): return kwargs = self._client_kwargs() - if "path" in kwargs: - path = Path(kwargs["path"]) - path.mkdir(parents=True, exist_ok=True) - database = kwargs.get("database", self.store_name) - try: - admin = pyseekdb.AdminClient(path=str(path)) - if not any(db.name == database for db in admin.list_databases()): - admin.create_database(database) - except Exception as e: - logger.debug("seekdb AdminClient create_database: %s", e) + if not self._is_remote and "path" in kwargs: + Path(kwargs["path"]).parent.mkdir(parents=True, exist_ok=True) + database = kwargs.get("database", self.store_name) + try: + admin = pyseekdb.AdminClient(**admin_kwargs_from_client_kwargs(kwargs)) + if not any(db.name == database for db in admin.list_databases()): + admin.create_database(database) + except Exception as e: + logger.debug("seekdb AdminClient create_database: %s", e) self.client = pyseekdb.Client(**kwargs) dim = self.embedding_dim diff --git a/reme/core/utils/__init__.py b/reme/core/utils/__init__.py index 26ca2827..5ff544bd 100644 --- a/reme/core/utils/__init__.py +++ b/reme/core/utils/__init__.py @@ -19,6 +19,7 @@ from .pydantic_utils import create_pydantic_model from .singleton import singleton from .time import timer, get_now_time from .hf_token_counter_utils import get_hf_token_counter +from .pyseekdb_conn import admin_kwargs_from_client_kwargs, build_pyseekdb_client_kwargs, parse_host_port __all__ = [ "convert_dashscope_to_agentscope", @@ -50,4 +51,7 @@ __all__ = [ "timer", "get_now_time", "get_hf_token_counter", + "admin_kwargs_from_client_kwargs", + "build_pyseekdb_client_kwargs", + "parse_host_port", ] diff --git a/reme/core/utils/pyseekdb_conn.py b/reme/core/utils/pyseekdb_conn.py new file mode 100644 index 00000000..350f5809 --- /dev/null +++ b/reme/core/utils/pyseekdb_conn.py @@ -0,0 +1,60 @@ +"""Build ``pyseekdb.Client`` / ``AdminClient`` kwargs for embedded vs remote OceanBase / seekdb.""" + +from __future__ import annotations + +DEFAULT_SEEKDB_PORT = 2881 +DEFAULT_SEEKDB_TENANT = "sys" +DEFAULT_SEEKDB_USER = "root" + + +def parse_host_port(host: str | None, port: int | None) -> tuple[str | None, int | None]: + """Resolve remote ``host`` / ``port`` (``port`` defaults to :data:`DEFAULT_SEEKDB_PORT`).""" + if host and host.strip(): + return host.strip(), port if port is not None else DEFAULT_SEEKDB_PORT + return None, None + + +def build_pyseekdb_client_kwargs( + *, + path: str | None = None, + database: str, + host: str | None = None, + port: int | None = None, + tenant: str = DEFAULT_SEEKDB_TENANT, + user: str | None = None, + password: str = "", +) -> tuple[bool, dict]: + """Return ``(is_remote, kwargs)`` for ``pyseekdb.Client``. + + Remote when ``host`` is set. Embedded: optional ``path`` for the data directory; if + omitted, pyseekdb uses its default (typically ``seekdb.db`` under the CWD). + """ + h, p = parse_host_port(host, port) + if h: + return True, { + "host": h, + "port": p, + "tenant": tenant, + "database": database, + "user": user if user is not None else DEFAULT_SEEKDB_USER, + "password": password, + } + kw: dict = {"database": database} + if path: + kw["path"] = path + return False, kw + + +def admin_kwargs_from_client_kwargs(client_kw: dict) -> dict: + """Strip ``database`` for ``AdminClient`` (admin uses system DB).""" + if "path" in client_kw: + return {"path": client_kw["path"]} + if "host" in client_kw: + return { + "host": client_kw["host"], + "port": client_kw["port"], + "tenant": client_kw["tenant"], + "user": client_kw["user"], + "password": client_kw["password"], + } + return {} diff --git a/reme/core/vector_store/seekdb_vector_store.py b/reme/core/vector_store/seekdb_vector_store.py index 6318e1b4..f134c482 100644 --- a/reme/core/vector_store/seekdb_vector_store.py +++ b/reme/core/vector_store/seekdb_vector_store.py @@ -1,4 +1,9 @@ -"""seekdb vector store implementation for the ReMe framework.""" +"""seekdb vector store implementation for the ReMe framework. + +Uses ``pyseekdb`` (Chroma-like Collection API) for **embedded** local storage or +**remote** OceanBase / seekdb—the same deployment modes as ``pyseekdb.Client``. +For SQL-table-oriented helpers via ``pyobvector``, see ``ObVecVectorStore``. +""" from pathlib import Path from typing import Any @@ -8,6 +13,11 @@ from loguru import logger from .base_vector_store import BaseVectorStore from ..embedding import BaseEmbeddingModel from ..schema import VectorNode +from ..utils.pyseekdb_conn import ( + DEFAULT_SEEKDB_TENANT, + admin_kwargs_from_client_kwargs, + build_pyseekdb_client_kwargs, +) # Optional: preserve original exception for "raise ... from _PYSEEKDB_IMPORT_ERROR" (better diagnostics) _PYSEEKDB_IMPORT_ERROR = None @@ -25,11 +35,13 @@ except ImportError as e: class SeekdbVectorStore(BaseVectorStore): - """Vector store implementation using seekdb (pyseekdb) for embedded vector search. + """Vector store using ``pyseekdb`` and the Chroma-like Collection API. - Uses the same embedded Client + Collection API as SeekdbFileStore; supports - vector similarity search and metadata filtering. No full-text index by default - (vector_store use case is vector search + filter). + **Embedded** (default): optional ``path`` to the embedded data directory; if omitted, + pyseekdb applies its default (typically a ``seekdb.db`` directory name). **Remote**: + ``host`` / ``port`` plus auth (same deployment style as ``ObVecVectorStore``, without ``uri``). + + Vector similarity search and metadata filtering; no full-text index by default. """ def __init__( @@ -39,16 +51,30 @@ class SeekdbVectorStore(BaseVectorStore): embedding_model: BaseEmbeddingModel, database: str = "reme_vector", distance: str = "cosine", + host: str | None = None, + port: int | None = None, + tenant: str = DEFAULT_SEEKDB_TENANT, + user: str | None = None, + password: str = "", + path: str | None = None, **kwargs: Any, ): """Initialize the seekdb vector store. Args: collection_name: Name of the collection. - db_path: Local path for embedded seekdb storage (e.g. working_dir / "vector_store"). + db_path: Working directory for ReMe (metadata sidecar); also used when resolving + a default location alongside **remote** mode (mirrors ``ObVecVectorStore``). embedding_model: Model used for generating vector embeddings. - database: seekdb database name (all vector_store collections live in this DB). + database: Database name on the seekdb / OceanBase instance. distance: Similarity metric: cosine, euclid, dot. + host: Remote server host (embedded mode if unset or empty). + port: Remote port (default ``2881`` when ``host`` is set). + tenant: OceanBase / seekdb tenant (remote). + user: Remote user (``None`` uses library default ``root``). + password: Remote password. + path: Embedded data directory passed to ``pyseekdb.Client``; omit to use the + library default (typically ``./seekdb.db`` as the directory name). **kwargs: Additional options (ignored for compatibility). """ if _PYSEEKDB_IMPORT_ERROR is not None: @@ -67,12 +93,31 @@ class SeekdbVectorStore(BaseVectorStore): self.client: Any = None self.collection: Any = None + self._is_remote, self._client_kw = build_pyseekdb_client_kwargs( + path=None if (host and host.strip()) else path, + database=self.database, + host=host, + port=port, + tenant=tenant, + user=user, + password=password, + ) + def _client_kwargs(self) -> dict: - """Build Client kwargs for embedded mode.""" - return { - "path": str(self.db_path / "seekdb"), - "database": self.database, - } + """Kwargs passed to ``pyseekdb.Client`` (embedded or remote).""" + return self._client_kw + + def _coerce_embedding_for_upsert(self, vec: Any) -> list[float]: + """Normalize vectors before ``collection.upsert`` (pyseekdb SQL rejects empty hex).""" + if vec is None: + raw: list[float] = [] + elif hasattr(vec, "tolist"): + raw = list(vec.tolist()) + elif isinstance(vec, list): + raw = vec + else: + raw = list(vec) + return self.embedding_model._validate_and_adjust_embedding(raw) @staticmethod def _build_where(filters: dict | None) -> dict | None: @@ -201,10 +246,12 @@ class SeekdbVectorStore(BaseVectorStore): embedding_function=None, ) new_coll = self.client.get_collection(name=collection_name, embedding_function=None) + emb_out = data.get("embeddings") or [] + emb_norm = [self._coerce_embedding_for_upsert(e) for e in emb_out] if emb_out else [] new_coll.upsert( ids=ids, documents=data.get("documents", []), - embeddings=data.get("embeddings", []), + embeddings=emb_norm, metadatas=data.get("metadatas", []), ) logger.info(f"Copied {self.collection_name} to {collection_name}") @@ -224,7 +271,7 @@ class SeekdbVectorStore(BaseVectorStore): nodes_to_insert = nodes ids = [n.vector_id for n in nodes_to_insert] documents = [n.content for n in nodes_to_insert] - embeddings = [n.vector for n in nodes_to_insert] + embeddings = [self._coerce_embedding_for_upsert(n.vector) for n in nodes_to_insert] metadatas = [n.metadata for n in nodes_to_insert] self.collection.upsert(ids=ids, documents=documents, embeddings=embeddings, metadatas=metadatas) logger.info(f"Inserted {len(nodes_to_insert)} nodes into {self.collection_name}") @@ -294,7 +341,7 @@ class SeekdbVectorStore(BaseVectorStore): nodes_to_update = nodes ids = [n.vector_id for n in nodes_to_update] documents = [n.content for n in nodes_to_update] - embeddings = [n.vector for n in nodes_to_update] + embeddings = [self._coerce_embedding_for_upsert(n.vector) for n in nodes_to_update] metadatas = [n.metadata for n in nodes_to_update] self.collection.upsert(ids=ids, documents=documents, embeddings=embeddings, metadatas=metadatas) logger.info(f"Updated {len(nodes_to_update)} nodes in {self.collection_name}") @@ -360,17 +407,19 @@ class SeekdbVectorStore(BaseVectorStore): async def start(self) -> None: """Initialize seekdb client and ensure collection exists.""" - path = Path(self._client_kwargs()["path"]) - path.mkdir(parents=True, exist_ok=True) + kw = self._client_kwargs() + if not self._is_remote and "path" in kw: + Path(kw["path"]).parent.mkdir(parents=True, exist_ok=True) try: - admin = pyseekdb.AdminClient(path=str(path)) + admin = pyseekdb.AdminClient(**admin_kwargs_from_client_kwargs(kw)) if not any(db.name == self.database for db in admin.list_databases()): admin.create_database(self.database) except Exception as e: - logger.debug("seekdb AdminClient: %s", e) - self.client = pyseekdb.Client(**self._client_kwargs()) + logger.debug("seekdb AdminClient create_database: %s", e) + self.client = pyseekdb.Client(**kw) await self.create_collection(self.collection_name) - logger.info(f"seekdb vector store {self.collection_name} initialized") + mode = "remote" if self._is_remote else "embedded" + logger.info(f"seekdb vector store ({mode}) {self.collection_name} initialized") async def close(self) -> None: """Release client; no explicit close in pyseekdb, clear references.""" diff --git a/tests/test_file_store.py b/tests/test_file_store.py index 15df1063..e8838bc8 100644 --- a/tests/test_file_store.py +++ b/tests/test_file_store.py @@ -67,9 +67,14 @@ class TestConfig: LOCAL_DB_PATH = "./test_file_store_local" LOCAL_FTS_ENABLED = True - # SeekdbFileStore settings (embedded mode) + # SeekdbFileStore: embedded by default; SEEKDB_HOST (and optional SEEKDB_PORT) => remote SEEKDB_DB_PATH = "./test_file_store_seekdb" SEEKDB_FTS_ENABLED = True + SEEKDB_HOST = os.environ.get("SEEKDB_HOST") + SEEKDB_PORT = int(os.environ["SEEKDB_PORT"]) if os.environ.get("SEEKDB_PORT") else None + SEEKDB_USER = os.environ.get("SEEKDB_USER", "root") + SEEKDB_PASSWORD = os.environ.get("SEEKDB_PASSWORD", "root") + SEEKDB_TENANT = os.environ.get("SEEKDB_TENANT", "sys") # Embedding model settings EMBEDDING_MODEL_NAME = "text-embedding-v4" @@ -269,13 +274,23 @@ def create_file_store(store_type: str) -> BaseFileStore: raise ImportError( "SeekdbFileStore requires pyseekdb. Install with: pip install reme-ai (pyseekdb is included)", ) - return SeekdbFileStore( - store_name=config.NAME, - db_path=config.SEEKDB_DB_PATH, - embedding_model=embedding_model, - fts_enabled=config.SEEKDB_FTS_ENABLED, - vector_enabled=True, - ) + remote = bool(config.SEEKDB_HOST and config.SEEKDB_HOST.strip()) + kw: dict = { + "store_name": config.NAME, + "db_path": config.SEEKDB_DB_PATH, + "embedding_model": embedding_model, + "fts_enabled": config.SEEKDB_FTS_ENABLED, + "vector_enabled": True, + } + if remote: + kw["host"] = config.SEEKDB_HOST.strip() + kw["port"] = config.SEEKDB_PORT + kw["user"] = config.SEEKDB_USER + kw["password"] = config.SEEKDB_PASSWORD + kw["tenant"] = config.SEEKDB_TENANT + else: + kw["path"] = str(Path(config.SEEKDB_DB_PATH) / "seekdb.db") + return SeekdbFileStore(**kw) else: raise ValueError(f"Unknown store type: {store_type}") diff --git a/tests/test_vector_store.py b/tests/test_vector_store.py index 4f3594d3..0beacc04 100644 --- a/tests/test_vector_store.py +++ b/tests/test_vector_store.py @@ -81,8 +81,14 @@ class TestConfig: PG_USE_HNSW = True # Use HNSW index for faster search PG_USE_DISKANN = False # Use DiskANN index (requires vectorscale extension) - # SeekdbVectorStore settings (embedded, temp dir used if not set) + # SeekdbVectorStore: embedded by default; set SEEKDB_HOST (and optional SEEKDB_PORT) for remote SEEKDB_PATH = None # e.g. "./test_vector_store_seekdb"; None => temp dir + SEEKDB_HOST = os.environ.get("SEEKDB_HOST") + SEEKDB_PORT = int(os.environ["SEEKDB_PORT"]) if os.environ.get("SEEKDB_PORT") else None + SEEKDB_USER = os.environ.get("SEEKDB_USER", "root") + SEEKDB_PASSWORD = os.environ.get("SEEKDB_PASSWORD", "root") + SEEKDB_TENANT = os.environ.get("SEEKDB_TENANT", "sys") + SEEKDB_DATABASE = os.environ.get("SEEKDB_DATABASE", "reme_vector") # ChromaVectorStore settings CHROMA_PATH = "./test_vector_store_chroma" # For local persistent mode @@ -314,13 +320,24 @@ def create_vector_store(store_type: str, collection_name: str) -> BaseVectorStor elif store_type == "seekdb": if SeekdbVectorStore is None: raise ImportError("SeekdbVectorStore not available; install pyseekdb") - return SeekdbVectorStore( - collection_name=collection_name, - embedding_model=embedding_model, - db_path=config.SEEKDB_PATH or tempfile.mkdtemp(prefix="test_seekdb_"), - database="reme_vector", - distance="cosine", - ) + remote = bool(config.SEEKDB_HOST and config.SEEKDB_HOST.strip()) + db_path = config.SEEKDB_PATH or tempfile.mkdtemp(prefix="test_seekdb_") + kw: dict = { + "collection_name": collection_name, + "embedding_model": embedding_model, + "db_path": db_path, + "database": config.SEEKDB_DATABASE if remote else "reme_vector", + "distance": "cosine", + } + if remote: + kw["host"] = config.SEEKDB_HOST.strip() + kw["port"] = config.SEEKDB_PORT + kw["user"] = config.SEEKDB_USER + kw["password"] = config.SEEKDB_PASSWORD + kw["tenant"] = config.SEEKDB_TENANT + else: + kw["path"] = str(Path(db_path) / "seekdb.db") + return SeekdbVectorStore(**kw) else: raise ValueError(f"Unknown store type: {store_type}")