scriptoriumadmin/ReMe
1
0
Fork 0
mirror of https://github.com/agentscope-ai/ReMe.git synced 2026-09-12 23:01:15 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 9

refactor(seekdb): add pyseekdb_conn and remote-only host/port config

Browse source
This commit is contained in:
dengfuping 2026-04-10 17:43:03 +08:00
parent 279a995eaf
commit 283a8bf832
6 changed files with 225 additions and 53 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

59
reme/core/file_store/seekdb_file_store.py
View file

@ -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

4
reme/core/utils/__init__.py
View file

@ -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",
]

60
reme/core/utils/pyseekdb_conn.py Normal file
View file

@ -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 {}

91
reme/core/vector_store/seekdb_vector_store.py
View file

@ -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."""

31
tests/test_file_store.py
View file

@ -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}")

33
tests/test_vector_store.py
View file

@ -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}")