diff --git a/.gitignore b/.gitignore index 264c7372..1e52cf66 100644 --- a/.gitignore +++ b/.gitignore @@ -5,7 +5,10 @@ venv/ .ipynb_checkpoints .__pycache__ -__pycache__ +__pycache__/ +*.pyc +*.pyo +*.pyd *.log tmp* temp* diff --git a/example.env b/example.env index d1a27415..9824e6a4 100644 --- a/example.env +++ b/example.env @@ -1,5 +1,18 @@ +# LLM (required for most flows) LLM_API_KEY=sk-xxxx LLM_BASE_URL=https://xxxx/v1 + +# Embedding (optional; Application uses EMBEDDING_* or falls back to LLM_* when unset) #EMBEDDING_API_KEY=sk-xxxx #EMBEDDING_BASE_URL=https://xxxx/v1 + +# Web search (optional) #TAVILY_API_KEY=xxxx + +# Seekdb / pyseekdb (optional; requires Python >=3.11, installed automatically on 3.11+) +# Embedded: leave SEEKDB_HOST unset. Remote: set host (and port if not 2881). +#SEEKDB_HOST=127.0.0.1 +#SEEKDB_PORT=2881 +#SEEKDB_USER=root +#SEEKDB_PASSWORD= +#SEEKDB_DATABASE=test diff --git a/pyproject.toml b/pyproject.toml index 89d9020e..024ae35d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -36,6 +36,8 @@ dependencies = [ "aiofiles>=24.1.0", "asyncpg>=0.31.0", "chromadb>=1.3.5", + # pyseekdb 1.2+ requires Python >=3.11 (no wheels on 3.10) + "pyseekdb>=1.2.0; python_version >= '3.11'", "dashscope>=1.25.1", "elasticsearch>=9.2.0", "fastapi>=0.121.3", diff --git a/reme/core/embedding/base_embedding_model.py b/reme/core/embedding/base_embedding_model.py index aeea6869..e40526d0 100644 --- a/reme/core/embedding/base_embedding_model.py +++ b/reme/core/embedding/base_embedding_model.py @@ -55,8 +55,8 @@ class BaseEmbeddingModel(ABC): enable_cache: Whether to enable embedding cache **kwargs: Additional model-specific parameters """ - self.api_key: str = api_key - self.base_url: str = base_url + self.api_key: str | None = api_key + self.base_url: str | None = base_url self.model_name = model_name self.dimensions = dimensions self.use_dimensions = use_dimensions diff --git a/reme/core/file_store/__init__.py b/reme/core/file_store/__init__.py index a8457406..f071956d 100644 --- a/reme/core/file_store/__init__.py +++ b/reme/core/file_store/__init__.py @@ -1,8 +1,8 @@ """File store module for persistent memory management. This module provides storage backends for memory chunks and file metadata, -including SQLite-based, ChromaDB-based, and pure-Python local implementations -with vector and full-text search. +including SQLite-based, ChromaDB-based, seekdb-based, and pure-Python local +implementations with vector and full-text search. """ from .base_file_store import BaseFileStore @@ -24,3 +24,11 @@ R.file_stores.register("sqlite")(SqliteFileStore) R.file_stores.register("chroma")(ChromaFileStore) R.file_stores.register("local")(LocalFileStore) R.file_stores.register("zvec")(ZvecFileStore) + +try: + from .seekdb_file_store import SeekdbFileStore + + R.file_stores.register("seekdb")(SeekdbFileStore) + __all__.append("SeekdbFileStore") +except ImportError: + pass diff --git a/reme/core/file_store/seekdb_file_store.py b/reme/core/file_store/seekdb_file_store.py new file mode 100644 index 00000000..3f38400f --- /dev/null +++ b/reme/core/file_store/seekdb_file_store.py @@ -0,0 +1,601 @@ +"""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 + +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 ( + admin_kwargs_from_client_kwargs, + build_pyseekdb_client_kwargs, +) + +try: + import pyseekdb + from pyseekdb import Configuration, HNSWConfiguration, FulltextIndexConfig + + PYSEEKDB_AVAILABLE = True +except ImportError: + PYSEEKDB_AVAILABLE = False + pyseekdb = None + Configuration = None + HNSWConfiguration = None + FulltextIndexConfig = None + + +def _escape_sql(s: str) -> str: + """Escape single quotes for SQL string literals (MySQL/seekdb).""" + return s.replace("\\", "\\\\").replace("'", "''") + + +class SeekdbFileStore(BaseFileStore): + """seekdb file storage with vector and full-text search via ``pyseekdb``. + + **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, + user: str | None = None, + password: str = "", + path: str | None = None, + **kwargs, + ): + if not PYSEEKDB_AVAILABLE: + raise ImportError( + "pyseekdb is required for SeekdbFileStore. " + "Install it with: pip install reme-ai (pyseekdb is included)", + ) + + super().__init__(**kwargs) + 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, + user=user, + password=password, + ) + + @property + def collection_name(self) -> str: + """Collection name for chunks.""" + return f"chunks_{self.store_name}" + + @property + def files_table_name(self) -> str: + """Table name for file metadata (same as SQLite).""" + return f"files_{self.store_name}" + + def _client_kwargs(self) -> dict: + """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).""" + if self.client is None: + return None + return getattr(self.client, "_server", self.client) + + def _execute_sql(self, sql: str): + """Execute SQL via pyseekdb embedded/server client. Returns fetchall() for SELECT/SHOW/DESCRIBE.""" + client = self._sql_client() + if client is None: + raise RuntimeError("seekdb client not initialized") + for attr in ("execute", "_execute"): + run_sql = getattr(client, attr, None) + if run_sql is not None: + return run_sql(sql) + raise RuntimeError("seekdb client has no execute/_execute for raw SQL") + + def _create_files_table(self) -> None: + """Create file metadata table (same schema as SQLite).""" + sql = f""" + CREATE TABLE IF NOT EXISTS `{self.files_table_name}` ( + path VARCHAR(1024), + source VARCHAR(128), + hash VARCHAR(256), + mtime REAL, + size BIGINT, + PRIMARY KEY (path, source) + ) + """ + self._execute_sql(sql) + logger.debug(f"seekdb files table: {self.files_table_name}") + + async def start(self) -> None: + """Initialize seekdb client, collection, and files table.""" + if self.client is not None: + return + + kwargs = self._client_kwargs() + 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 + config = Configuration( + hnsw=HNSWConfiguration(dimension=dim, distance="cosine"), + fulltext_config=FulltextIndexConfig(analyzer="space"), + ) + self.collection = self.client.get_or_create_collection( + name=self.collection_name, + configuration=config, + embedding_function=None, + ) + self._create_files_table() + + logger.info( + f"seekdb initialized with collection: {self.collection_name}, " f"files table: {self.files_table_name}", + ) + + async def upsert_file( + self, + file_meta: FileMetadata, + source: MemorySource, + chunks: list[MemoryChunk], + ) -> None: + """Insert or update file and its chunks.""" + if not chunks: + return + + await self.delete_file(file_meta.path, source) + chunks = await self.get_chunk_embeddings(chunks) + + ids = [] + documents = [] + embeddings = [] + metadatas = [] + + now = int(time.time() * 1000) + for chunk in chunks: + ids.append(chunk.id) + documents.append(chunk.text) + embeddings.append(chunk.embedding) + metadatas.append( + { + "path": file_meta.path, + "source": source.value, + "start_line": chunk.start_line, + "end_line": chunk.end_line, + "hash": chunk.hash, + "updated_at": now, + }, + ) + + self.collection.upsert( + ids=ids, + documents=documents, + embeddings=embeddings, + metadatas=metadatas, + ) + # File metadata in DB table (same as SQLite) + p, s = _escape_sql(file_meta.path), _escape_sql(source.value) + h = _escape_sql(file_meta.hash) + mtime = file_meta.mtime_ms + size = file_meta.size + sql = ( + f"REPLACE INTO `{self.files_table_name}` (path, source, hash, mtime, size) " + f"VALUES ('{p}', '{s}', '{h}', {mtime}, {size})" + ) + self._execute_sql(sql) + + async def delete_file(self, path: str, source: MemorySource) -> None: + """Delete file and all its chunks.""" + results = self.collection.get( + where={"$and": [{"path": path}, {"source": source.value}]}, + include=[], + ) + if results.get("ids"): + self.collection.delete(ids=results["ids"]) + p, s = _escape_sql(path), _escape_sql(source.value) + self._execute_sql(f"DELETE FROM `{self.files_table_name}` WHERE path = '{p}' AND source = '{s}'") + + async def delete_file_chunks(self, path: str, chunk_ids: list[str]) -> None: + """Delete specific chunks for a file (chunk count comes from collection at get_file_metadata).""" + if not chunk_ids: + return + self.collection.delete(ids=chunk_ids) + + async def upsert_chunks( + self, + chunks: list[MemoryChunk], + source: MemorySource, + ) -> None: + """Insert or update specific chunks.""" + if not chunks: + return + chunks = await self.get_chunk_embeddings(chunks) + ids = [] + documents = [] + embeddings = [] + metadatas = [] + now = int(time.time() * 1000) + for chunk in chunks: + ids.append(chunk.id) + documents.append(chunk.text) + embeddings.append(chunk.embedding) + metadatas.append( + { + "path": chunk.path, + "source": source.value, + "start_line": chunk.start_line, + "end_line": chunk.end_line, + "hash": chunk.hash, + "updated_at": now, + }, + ) + self.collection.upsert( + ids=ids, + documents=documents, + embeddings=embeddings, + metadatas=metadatas, + ) + + async def list_files(self, source: MemorySource) -> list[str]: + """List all indexed files for a source (from files table).""" + s = _escape_sql(source.value) + rows = self._execute_sql(f"SELECT path FROM `{self.files_table_name}` WHERE source = '{s}'") + if not rows: + return [] + return [row[0] if isinstance(row, (list, tuple)) else row.get("path") for row in rows] + + async def get_file_metadata( + self, + path: str, + source: MemorySource, + ) -> FileMetadata | None: + """Get file metadata from files table and chunk count from collection.""" + p, s = _escape_sql(path), _escape_sql(source.value) + rows = self._execute_sql( + f"SELECT hash, mtime, size FROM `{self.files_table_name}` WHERE path = '{p}' AND source = '{s}'", + ) + if not rows: + return None + row = rows[0] + if isinstance(row, (list, tuple)): + hash_val, mtime, size = row[0], row[1], row[2] + else: + hash_val, mtime, size = row["hash"], row["mtime"], row["size"] + results = self.collection.get( + where={"$and": [{"path": path}, {"source": source.value}]}, + include=[], + ) + chunk_count = len(results.get("ids") or []) + return FileMetadata( + path=path, + hash=hash_val or "", + mtime_ms=mtime or 0, + size=size or 0, + chunk_count=chunk_count, + ) + + async def update_file_metadata(self, file_meta: FileMetadata, source: MemorySource) -> None: + """Update file metadata in files table without affecting chunks.""" + p = _escape_sql(file_meta.path) + s = _escape_sql(source.value) + h = _escape_sql(file_meta.hash) + mtime = file_meta.mtime_ms + size = file_meta.size + sql = ( + f"REPLACE INTO `{self.files_table_name}` (path, source, hash, mtime, size) " + f"VALUES ('{p}', '{s}', '{h}', {mtime}, {size})" + ) + self._execute_sql(sql) + + async def get_file_chunks( + self, + path: str, + source: MemorySource, + ) -> list[MemoryChunk]: + """Get all chunks for a file.""" + results = self.collection.get( + where={"$and": [{"path": path}, {"source": source.value}]}, + include=["documents", "embeddings", "metadatas"], + ) + chunks = [] + ids = results.get("ids") or [] + documents = results.get("documents") or [] + embeddings = results.get("embeddings") or [] + metadatas = results.get("metadatas") or [] + for i, chunk_id in enumerate(ids): + meta = metadatas[i] if i < len(metadatas) else {} + doc = documents[i] if i < len(documents) else "" + emb = embeddings[i] if embeddings and i < len(embeddings) else None + chunks.append( + MemoryChunk( + id=chunk_id, + path=meta.get("path", path), + source=MemorySource(meta.get("source", source.value)), + start_line=meta.get("start_line", 0), + end_line=meta.get("end_line", 0), + text=doc, + hash=meta.get("hash", ""), + embedding=emb, + ), + ) + chunks.sort(key=lambda c: c.start_line) + return chunks + + async def vector_search( + self, + query: str, + limit: int, + sources: list[MemorySource] | None = None, + ) -> list[MemorySearchResult]: + """Perform vector similarity search.""" + if not self.vector_enabled or not query: + return [] + query_embedding = await self.get_embedding(query) + if not query_embedding: + return [] + + where_filter = None + if sources: + if len(sources) == 1: + where_filter = {"source": sources[0].value} + else: + where_filter = {"source": {"$in": [s.value for s in sources]}} + + results = self.collection.query( + query_embeddings=[query_embedding], + n_results=limit, + where=where_filter, + include=["documents", "metadatas", "distances"], + ) + + search_results = [] + if results.get("ids") and results["ids"][0]: + for i, _ in enumerate(results["ids"][0]): + metadata = results["metadatas"][0][i] + distance = results["distances"][0][i] + score = max(0.0, 1.0 - distance / 2.0) + search_results.append( + MemorySearchResult( + path=metadata["path"], + start_line=metadata["start_line"], + end_line=metadata["end_line"], + score=score, + snippet=results["documents"][0][i], + source=MemorySource(metadata["source"]), + raw_metric=distance, + ), + ) + search_results.sort(key=lambda r: r.score, reverse=True) + return search_results + + async def keyword_search( + self, + query: str, + limit: int, + sources: list[MemorySource] | None = None, + ) -> list[MemorySearchResult]: + """Perform full-text keyword search.""" + if not self.fts_enabled or not query or not query.strip(): + return [] + + where_filter = None + if sources: + if len(sources) == 1: + where_filter = {"source": sources[0].value} + else: + where_filter = {"source": {"$in": [s.value for s in sources]}} + + results = self.collection.get( + where=where_filter, + where_document={"$contains": query.strip()}, + limit=limit, + include=["documents", "metadatas"], + ) + + search_results = [] + ids = results.get("ids") or [] + documents = results.get("documents") or [] + metadatas = results.get("metadatas") or [] + query_lower = query.lower() + words = query_lower.split() + n_words = max(1, len(words)) + + for i, _ in enumerate(ids): + meta = metadatas[i] if i < len(metadatas) else {} + text = documents[i] if i < len(documents) else "" + match_count = sum(1 for w in words if w in text.lower()) + score = min(1.0, match_count / n_words + (0.2 if query_lower in text.lower() else 0.0)) + search_results.append( + MemorySearchResult( + path=meta.get("path", ""), + start_line=meta.get("start_line", 0), + end_line=meta.get("end_line", 0), + score=score, + snippet=text, + source=MemorySource(meta.get("source", "memory")), + ), + ) + search_results.sort(key=lambda r: r.score, reverse=True) + return search_results[:limit] + + async def hybrid_search( + self, + query: str, + limit: int, + sources: list[MemorySource] | None = None, + vector_weight: float = 0.7, + candidate_multiplier: float = 3.0, + ) -> list[MemorySearchResult]: + """Perform hybrid search using seekdb native hybrid_search when possible.""" + if not query or not query.strip(): + return [] + + assert 0.0 <= vector_weight <= 1.0 + candidates = min(200, max(1, int(limit * candidate_multiplier))) + where_filter = None + if sources: + if len(sources) == 1: + where_filter = {"source": sources[0].value} + else: + where_filter = {"source": {"$in": [s.value for s in sources]}} + + if self.vector_enabled and self.fts_enabled: + try: + query_embedding = await self.get_embedding(query) + if not query_embedding: + return await self.keyword_search(query, limit, sources) + + results = self.collection.hybrid_search( + query={ + "where_document": {"$contains": query.strip()}, + "where": where_filter, + "n_results": candidates, + }, + knn={ + "query_embeddings": [query_embedding], + "where": where_filter, + "n_results": candidates, + }, + rank={"rrf": {}}, + n_results=limit, + include=["documents", "metadatas", "distances"], + ) + except Exception as e: + logger.warning(f"seekdb hybrid_search failed, fallback to merge: {e}") + return await self._hybrid_search_merge( + query, + limit, + sources, + vector_weight, + candidate_multiplier, + ) + else: + return await self._hybrid_search_merge( + query, + limit, + sources, + vector_weight, + candidate_multiplier, + ) + + search_results = [] + raw_ids = results.get("ids") or [] + ids = raw_ids[0] if raw_ids and isinstance(raw_ids[0], list) else raw_ids + if not ids: + return [] + documents = results.get("documents") + metadatas = results.get("metadatas") + distances = results.get("distances") + doc_list = (documents[0] if documents and isinstance(documents[0], list) else documents) or [] + meta_list = (metadatas[0] if metadatas and isinstance(metadatas[0], list) else metadatas) or [] + dist_list = (distances[0] if distances and isinstance(distances[0], list) else distances) or [] + for i, _ in enumerate(ids): + meta = meta_list[i] if i < len(meta_list) else {} + doc = doc_list[i] if i < len(doc_list) else "" + dist = dist_list[i] if i < len(dist_list) else 0.0 + score = max(0.0, 1.0 - dist / 2.0) + search_results.append( + MemorySearchResult( + path=meta.get("path", ""), + start_line=meta.get("start_line", 0), + end_line=meta.get("end_line", 0), + score=score, + snippet=doc, + source=MemorySource(meta.get("source", "memory")), + raw_metric=dist, + ), + ) + return search_results[:limit] + + async def _hybrid_search_merge( + self, + query: str, + limit: int, + sources: list[MemorySource] | None, + vector_weight: float, + candidate_multiplier: float, + ) -> list[MemorySearchResult]: + """Fallback: merge vector and keyword results like Chroma.""" + candidates = min(200, max(1, int(limit * candidate_multiplier))) + text_weight = 1.0 - vector_weight + + if self.vector_enabled and self.fts_enabled: + keyword_results = await self.keyword_search(query, candidates, sources) + vector_results = await self.vector_search(query, candidates, sources) + if not keyword_results: + return vector_results[:limit] + if not vector_results: + return keyword_results[:limit] + merged = self._merge_hybrid_results( + vector=vector_results, + keyword=keyword_results, + vector_weight=vector_weight, + text_weight=text_weight, + ) + return merged[:limit] + if self.vector_enabled: + return await self.vector_search(query, limit, sources) + return await self.keyword_search(query, limit, sources) + + @staticmethod + def _merge_hybrid_results( + vector: list[MemorySearchResult], + keyword: list[MemorySearchResult], + vector_weight: float, + text_weight: float, + ) -> list[MemorySearchResult]: + """Merge vector and keyword results with weighted scoring.""" + merged: dict[str, MemorySearchResult] = {} + for result in vector: + result.score = result.score * vector_weight + merged[result.merge_key] = result + for result in keyword: + key = result.merge_key + if key in merged: + merged[key].score += result.score * text_weight + else: + result.score = result.score * text_weight + merged[key] = result + results = list(merged.values()) + results.sort(key=lambda r: r.score, reverse=True) + return results + + async def clear_all(self) -> None: + """Clear all indexed data (collection + files table).""" + self.client.delete_collection(self.collection_name) + dim = self.embedding_dim + config = Configuration( + hnsw=HNSWConfiguration(dimension=dim, distance="cosine"), + fulltext_config=FulltextIndexConfig(analyzer="space"), + ) + self.collection = self.client.get_or_create_collection( + name=self.collection_name, + configuration=config, + embedding_function=None, + ) + self._execute_sql(f"DELETE FROM `{self.files_table_name}`") + logger.info(f"Cleared all data from seekdb collection: {self.collection_name} and files table") + + async def close(self) -> None: + """Close client (file metadata is in DB, no persist needed).""" + self.client = None + self.collection = None 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..f1c76e2c --- /dev/null +++ b/reme/core/utils/pyseekdb_conn.py @@ -0,0 +1,57 @@ +"""Build ``pyseekdb.Client`` / ``AdminClient`` kwargs for embedded vs remote OceanBase / seekdb.""" + +from __future__ import annotations + +DEFAULT_SEEKDB_PORT = 2881 +DEFAULT_SEEKDB_USER = "root" +DEFAULT_SEEKDB_DATABASE = "test" + + +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, + 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, + "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"], + "user": client_kw["user"], + "password": client_kw["password"], + } + return {} diff --git a/reme/core/vector_store/__init__.py b/reme/core/vector_store/__init__.py index 7b5a60bb..8ea58d76 100644 --- a/reme/core/vector_store/__init__.py +++ b/reme/core/vector_store/__init__.py @@ -31,3 +31,11 @@ R.vector_stores.register("obvec")(ObVecVectorStore) R.vector_stores.register("pgvector")(PGVectorStore) R.vector_stores.register("qdrant")(QdrantVectorStore) R.vector_stores.register("zvec")(ZvecVectorStore) + +try: + from .seekdb_vector_store import SeekdbVectorStore + + R.vector_stores.register("seekdb")(SeekdbVectorStore) + __all__.append("SeekdbVectorStore") +except ImportError: + pass diff --git a/reme/core/vector_store/seekdb_vector_store.py b/reme/core/vector_store/seekdb_vector_store.py new file mode 100644 index 00000000..3e0bf24f --- /dev/null +++ b/reme/core/vector_store/seekdb_vector_store.py @@ -0,0 +1,437 @@ +"""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 + +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_DATABASE, + 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 + +try: + import pyseekdb + from pyseekdb import Configuration, HNSWConfiguration + + PYSEEKDB_AVAILABLE = True +except ImportError as e: + _PYSEEKDB_IMPORT_ERROR = e + pyseekdb = None + Configuration = None + HNSWConfiguration = None + + +class SeekdbVectorStore(BaseVectorStore): + """Vector store using ``pyseekdb`` and the Chroma-like Collection API. + + **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__( + self, + collection_name: str, + db_path: str | Path, + embedding_model: BaseEmbeddingModel, + database: str = DEFAULT_SEEKDB_DATABASE, + distance: str = "cosine", + host: str | None = None, + port: int | None = None, + 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: 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: 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). + 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: + raise ImportError( + "seekdb vector store requires pyseekdb. Install with `pip install pyseekdb` or `pip install reme-ai`", + ) from _PYSEEKDB_IMPORT_ERROR + + super().__init__( + collection_name=collection_name, + db_path=db_path, + embedding_model=embedding_model, + **kwargs, + ) + self.database = database + self.distance = distance.lower() + 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, + user=user, + password=password, + ) + + def _client_kwargs(self) -> dict: + """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) + dim = self.embedding_model.dimensions + actual_len = len(raw) + if actual_len == dim: + return raw + if actual_len < dim: + logger.warning( + f"Embedding dimensions {actual_len} < {dim}, padding with zeros", + ) + return raw + [0.0] * (dim - actual_len) + logger.warning(f"Embedding dimensions {actual_len} > {dim}, truncating") + return raw[:dim] + + @staticmethod + def _build_where(filters: dict | None) -> dict | None: + """Build seekdb/Chroma-style where clause from universal filter format. + + Supports exact match and range: {"key": value} or {"key": [start, end]}. + """ + if not filters: + return None + conditions = [] + for key, value in filters.items(): + if value == "*": + continue + if isinstance(value, list) and len(value) == 2: + conditions.append({key: {"$gte": value[0]}}) + conditions.append({key: {"$lte": value[1]}}) + elif isinstance(value, dict) and ("gte" in value or "lte" in value or "gt" in value or "lt" in value): + for op, val in value.items(): + if op in ("gte", "lte", "gt", "lt") and val is not None: + conditions.append({key: {"$" + op: val}}) + else: + conditions.append({key: {"$eq": value}}) + if not conditions: + return None + return conditions[0] if len(conditions) == 1 else {"$and": conditions} + + @staticmethod + def _parse_results( + ids: list, + documents: list | None = None, + metadatas: list | None = None, + embeddings: list | None = None, + distances: list | None = None, + include_score: bool = False, + ) -> list[VectorNode]: + """Convert seekdb get/query result to list of VectorNode.""" + nodes = [] + documents = documents or [] + metadatas = metadatas or [] + embeddings = embeddings or [] + distances = distances or [] + if ids and isinstance(ids, list) and ids and isinstance(ids[0], list): + ids = ids[0] + if documents and isinstance(documents[0], list): + documents = documents[0] + if metadatas and isinstance(metadatas[0], list): + metadatas = metadatas[0] + if embeddings and isinstance(embeddings[0], list): + embeddings = embeddings[0] + if distances and isinstance(distances[0], (list, tuple)): + distances = distances[0] + for i, vector_id in enumerate(ids): + meta = metadatas[i] if i < len(metadatas) and metadatas[i] is not None else {} + if include_score and i < len(distances): + meta = dict(meta) + meta["score"] = 1.0 - (float(distances[i]) / 2.0) if distances[i] is not None else 0.0 + nodes.append( + VectorNode( + vector_id=str(vector_id), + content=documents[i] if i < len(documents) and documents[i] is not None else "", + vector=embeddings[i] if i < len(embeddings) else None, + metadata=meta, + ), + ) + return nodes + + async def list_collections(self) -> list[str]: + """List collection names in the current database.""" + if self.client is None: + return [] + try: + colls = self.client.list_collections() + return [c.name if hasattr(c, "name") else str(c) for c in colls] + except Exception as e: + logger.debug("seekdb list_collections: %s", e) + return [self.collection_name] + + async def create_collection(self, collection_name: str, **kwargs: Any) -> None: + """Create or get collection with HNSW vector index.""" + if self.client is None: + raise RuntimeError("seekdb client not initialized; call start() first") + dimensions = kwargs.get("dimensions", self.embedding_model.dimensions) + config = Configuration( + hnsw=HNSWConfiguration(dimension=dimensions, distance=self.distance), + ) + coll = self.client.get_or_create_collection( + name=collection_name, + configuration=config, + embedding_function=None, + ) + if collection_name == self.collection_name: + self.collection = coll + logger.info(f"seekdb collection {collection_name} ready (dim={dimensions})") + + async def delete_collection(self, collection_name: str, **kwargs: Any) -> None: + """Remove the collection from the database.""" + if self.client is None: + return + try: + self.client.delete_collection(collection_name) + if collection_name == self.collection_name: + self.collection = None + logger.info(f"Deleted seekdb collection {collection_name}") + except Exception as e: + logger.warning("seekdb delete_collection %s: %s", collection_name, e) + + async def copy_collection(self, collection_name: str, **kwargs: Any) -> None: + """Copy current collection to a new one.""" + if self.collection is None: + raise RuntimeError("No current collection") + data = self.collection.get(include=["documents", "metadatas", "embeddings"]) + ids = data.get("ids") or [] + if not ids: + logger.warning("Source collection is empty") + return + dims = self.embedding_model.dimensions + embs = data.get("embeddings") + if embs and (isinstance(embs[0], list) and embs[0]) or (not isinstance(embs[0], list) and embs): + dims = len(embs[0]) if isinstance(embs[0], list) else len(embs) + config = Configuration( + hnsw=HNSWConfiguration(dimension=dims, distance=self.distance), + ) + self.client.get_or_create_collection( + name=collection_name, + configuration=config, + 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=emb_norm, + metadatas=data.get("metadatas", []), + ) + logger.info(f"Copied {self.collection_name} to {collection_name}") + + async def insert(self, nodes: VectorNode | list[VectorNode], **kwargs: Any) -> None: + """Insert vector nodes; generate embeddings for nodes that lack them.""" + if isinstance(nodes, VectorNode): + nodes = [nodes] + if not nodes: + return + nodes_without_vectors = [n for n in nodes if n.vector is None] + if nodes_without_vectors: + nodes_with_vectors = await self.get_node_embeddings(nodes_without_vectors) + vector_map = {n.vector_id: n for n in nodes_with_vectors} + nodes_to_insert = [vector_map.get(n.vector_id, n) if n.vector is None else n for n in nodes] + else: + nodes_to_insert = nodes + ids = [n.vector_id for n in nodes_to_insert] + documents = [n.content 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}") + + async def search( + self, + query: str, + limit: int = 5, + filters: dict | None = None, + **kwargs: Any, + ) -> list[VectorNode]: + """Vector similarity search with optional metadata filter.""" + query_vector = await self.get_embedding(query) + where = self._build_where(filters) + results = self.collection.query( + query_embeddings=[query_vector], + n_results=limit, + where=where, + include=["documents", "metadatas", "distances"], + ) + ids = results.get("ids") or [] + documents = results.get("documents") + metadatas = results.get("metadatas") + distances = results.get("distances") + nodes = self._parse_results( + ids, + documents=documents, + metadatas=metadatas, + embeddings=results.get("embeddings"), + distances=distances, + include_score=True, + ) + score_threshold = kwargs.get("score_threshold") + if score_threshold is not None: + nodes = [n for n in nodes if n.metadata.get("score", 0) >= score_threshold] + return nodes + + async def delete(self, vector_ids: str | list[str], **kwargs: Any) -> None: + """Delete points by id.""" + if isinstance(vector_ids, str): + vector_ids = [vector_ids] + if not vector_ids: + return + self.collection.delete(ids=vector_ids) + logger.info(f"Deleted {len(vector_ids)} nodes from {self.collection_name}") + + async def delete_all(self, **kwargs: Any) -> None: + """Remove all points from the collection.""" + data = self.collection.get(include=[]) + ids = data.get("ids") or [] + if ids: + self.collection.delete(ids=ids) + logger.info(f"Deleted all {len(ids)} nodes from {self.collection_name}") + + async def update(self, nodes: VectorNode | list[VectorNode], **kwargs: Any) -> None: + """Update nodes (upsert by id).""" + if isinstance(nodes, VectorNode): + nodes = [nodes] + if not nodes: + return + nodes_without_vectors = [n for n in nodes if n.vector is None and n.content] + if nodes_without_vectors: + nodes_with_vectors = await self.get_node_embeddings(nodes_without_vectors) + vector_map = {n.vector_id: n for n in nodes_with_vectors} + nodes_to_update = [vector_map.get(n.vector_id, n) if n.vector is None and n.content else n for n in nodes] + else: + nodes_to_update = nodes + ids = [n.vector_id for n in nodes_to_update] + documents = [n.content 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}") + + async def get(self, vector_ids: str | list[str]) -> VectorNode | list[VectorNode]: + """Fetch nodes by id.""" + single = isinstance(vector_ids, str) + ids = [vector_ids] if single else list(vector_ids) + if not ids: + return None if single else [] + results = self.collection.get( + ids=ids, + include=["documents", "metadatas", "embeddings"], + ) + rids = results.get("ids") or [] + nodes = self._parse_results( + rids, + documents=results.get("documents"), + metadatas=results.get("metadatas"), + embeddings=results.get("embeddings"), + ) + if single: + return nodes[0] if nodes else None + return nodes + + async def list( + self, + filters: dict | None = None, + limit: int | None = None, + sort_key: str | None = None, + reverse: bool = True, + ) -> list[VectorNode]: + """List nodes with optional filter, limit, and sort by metadata key.""" + where = self._build_where(filters) + # When sorting in memory, fetch candidates first (cap like default list); do not pass + # user limit to get() or we sort an arbitrary first page only (see ChromaVectorStore.list). + if sort_key is not None: + fetch_limit = 10000 + else: + fetch_limit = limit if limit is not None else 10000 + results = self.collection.get( + where=where, + limit=fetch_limit, + include=["documents", "metadatas", "embeddings"], + ) + ids = results.get("ids") or [] + nodes = self._parse_results( + ids, + documents=results.get("documents"), + metadatas=results.get("metadatas"), + embeddings=results.get("embeddings"), + ) + if sort_key: + + def key_fn(n): + v = n.metadata.get(sort_key) + if v is None: + return float("-inf") if not reverse else float("inf") + return v + + nodes.sort(key=key_fn, reverse=reverse) + if limit is not None: + nodes = nodes[:limit] + return nodes + + async def start(self) -> None: + """Initialize seekdb client and ensure collection exists.""" + 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(**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 create_database: %s", e) + self.client = pyseekdb.Client(**kw) + await self.create_collection(self.collection_name) + 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.""" + self.client = None + self.collection = None + logger.info("seekdb vector store closed") diff --git a/tests/test_file_store.py b/tests/test_file_store.py index 0e34764f..5420e8c9 100644 --- a/tests/test_file_store.py +++ b/tests/test_file_store.py @@ -2,19 +2,22 @@ """Unified test suite for file store implementations. This module provides comprehensive test coverage for SqliteFileStore, ChromaFileStore, -LocalFileStore and future file store implementations. Tests can be run for specific stores -or all implementations. +LocalFileStore, ZvecFileStore, SeekdbFileStore and future file store implementations. +Tests can be run for specific stores or all implementations. Usage: python test_file_store.py --sqlite # Test SqliteFileStore only python test_file_store.py --chroma # Test ChromaFileStore only python test_file_store.py --local # Test LocalFileStore only + python test_file_store.py --zvec # Test ZvecFileStore only + python test_file_store.py --seekdb # Test SeekdbFileStore (pyseekdb; Python >=3.11) python test_file_store.py --all # Test all file stores """ import argparse import asyncio import hashlib +import os import shutil import time from pathlib import Path @@ -29,6 +32,14 @@ from reme.core.file_store.chroma_file_store import ChromaFileStore from reme.core.file_store.local_file_store import LocalFileStore from reme.core.file_store.sqlite_file_store import SqliteFileStore from reme.core.file_store.zvec_file_store import ZvecFileStore + +try: + from reme.core.file_store.seekdb_file_store import SeekdbFileStore + + SEEKDB_AVAILABLE = True +except ImportError: + SeekdbFileStore = None # type: ignore[misc, assignment] + SEEKDB_AVAILABLE = False from reme.core.schema.file_metadata import FileMetadata from reme.core.schema.memory_chunk import MemoryChunk from reme.core.utils import load_env @@ -62,6 +73,14 @@ class TestConfig: LOCAL_DB_PATH = "./test_file_store_local" LOCAL_FTS_ENABLED = True + # 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", "") + # Embedding model settings EMBEDDING_MODEL_NAME = "text-embedding-v4" EMBEDDING_DIMENSIONS = 64 @@ -196,7 +215,7 @@ def get_store_type(store: BaseFileStore) -> str: store: File store instance Returns: - str: Type identifier ("sqlite", "chroma", etc.) + str: Type identifier ("sqlite", "chroma", "local", "seekdb", etc.) """ if isinstance(store, SqliteFileStore): return "sqlite" @@ -206,6 +225,8 @@ def get_store_type(store: BaseFileStore) -> str: return "local" elif isinstance(store, ZvecFileStore): return "zvec" + elif SEEKDB_AVAILABLE and isinstance(store, SeekdbFileStore): + return "seekdb" else: raise ValueError(f"Unknown file store type: {type(store)}") @@ -221,10 +242,16 @@ def create_file_store(store_type: str) -> BaseFileStore: """ config = TestConfig() - # Initialize embedding model + # api_key/base_url from env (same as Application.embedding_api_key in application.py) + embedding_api_key = os.environ.get("EMBEDDING_API_KEY") or os.environ.get("OPENAI_API_KEY") + embedding_base_url = os.environ.get("EMBEDDING_BASE_URL") or os.environ.get("OPENAI_BASE_URL") + # use_dimensions=True so API returns 64-dim, matching collection/table schema for all backends embedding_model = OpenAIEmbeddingModel( model_name=config.EMBEDDING_MODEL_NAME, dimensions=config.EMBEDDING_DIMENSIONS, + use_dimensions=True, + api_key=embedding_api_key or None, + base_url=embedding_base_url or None, ) if store_type == "sqlite": @@ -257,6 +284,27 @@ def create_file_store(store_type: str) -> BaseFileStore: fts_enabled=config.ZVEC_FTS_ENABLED, dimension=config.EMBEDDING_DIMENSIONS, ) + elif store_type == "seekdb": + if not SEEKDB_AVAILABLE: + raise ImportError( + "SeekdbFileStore requires pyseekdb (Python >=3.11). Install: pip install 'pyseekdb>=1.2.0'", + ) + 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 + else: + kw["path"] = str(Path(config.SEEKDB_DB_PATH) / "seekdb.db") + return SeekdbFileStore(**kw) else: raise ValueError(f"Unknown store type: {store_type}") @@ -306,6 +354,12 @@ async def test_start_store(store: BaseFileStore, _store_name: str): assert store._initialized, "Zvec engine should be initialized" logger.info(f"✓ ZvecFileStore ready (collection: {store.collection_name})") + # Verify SeekdbFileStore initialized + if SEEKDB_AVAILABLE and isinstance(store, SeekdbFileStore): + assert store.client is not None, "seekdb client should be initialized" + assert store.collection is not None, "seekdb collection should exist" + logger.info(f"✓ seekdb collection created: {store.collection_name}") + async def test_upsert_file(store: BaseFileStore, _store_name: str) -> tuple[FileMetadata, List[MemoryChunk]]: """Test file and chunks insertion.""" @@ -1047,6 +1101,18 @@ async def cleanup_store(store: BaseFileStore, store_type: str): metadata_file.unlink() logger.info(f"✓ Cleaned up metadata file: {metadata_file}") + # Clean up SeekdbFileStore directory and metadata file + if store_type == "seekdb": + config = TestConfig() + db_dir = Path(config.SEEKDB_DB_PATH) + if db_dir.exists(): + shutil.rmtree(db_dir) + logger.info(f"✓ Cleaned up directory: {db_dir}") + metadata_file = db_dir.parent / f"{config.NAME}_file_metadata.json" + if metadata_file.exists(): + metadata_file.unlink() + logger.info(f"✓ Cleaned up metadata file: {metadata_file}") + logger.info("✓ Cleanup completed") except Exception as e: logger.error(f"Cleanup error: {e}") @@ -1093,6 +1159,11 @@ Examples: action="store_true", help="Run tests for all available file stores", ) + parser.add_argument( + "--seekdb", + action="store_true", + help="Test SeekdbFileStore (requires pyseekdb on Python >=3.11)", + ) args = parser.parse_args() @@ -1106,6 +1177,10 @@ Examples: ("local", "LocalFileStore"), ("zvec", "ZvecFileStore"), ] + if SEEKDB_AVAILABLE: + stores_to_test.append(("seekdb", "SeekdbFileStore")) + else: + logger.warning("SeekdbFileStore skipped (pyseekdb not installed or Python <3.11)") else: # Build list based on individual flags if args.sqlite: @@ -1116,6 +1191,8 @@ Examples: stores_to_test.append(("local", "LocalFileStore")) if args.zvec: stores_to_test.append(("zvec", "ZvecFileStore")) + if args.seekdb: + stores_to_test.append(("seekdb", "SeekdbFileStore")) if not stores_to_test: # Default to all file stores if no argument provided @@ -1125,8 +1202,10 @@ Examples: ("local", "LocalFileStore"), ("zvec", "ZvecFileStore"), ] + if SEEKDB_AVAILABLE: + stores_to_test.append(("seekdb", "SeekdbFileStore")) print("No file store specified, defaulting to test all file stores") - print("Use --sqlite, --chroma, --local, or --zvec to test specific ones\n") + print("Use --sqlite, --chroma, --local, --zvec, or --seekdb to test specific ones\n") # Run tests for each file store for store_type, store_name in stores_to_test: diff --git a/tests/test_vector_store.py b/tests/test_vector_store.py index 346be967..57260a04 100644 --- a/tests/test_vector_store.py +++ b/tests/test_vector_store.py @@ -2,8 +2,9 @@ """Unified test suite for vector store implementations. This module provides comprehensive test coverage for LocalVectorStore, ESVectorStore, -PGVectorStore, QdrantVectorStore, ChromaVectorStore, ObVecVectorStore, HologresVectorStore, and -ZvecVectorStore implementations. Tests can be run for specific vector stores or all implementations. +PGVectorStore, QdrantVectorStore, ChromaVectorStore, ObVecVectorStore, HologresVectorStore, +ZvecVectorStore, and SeekdbVectorStore implementations. Tests can be run for specific vector +stores or all implementations. Usage: python test_vector_store.py --local # Test LocalVectorStore only @@ -14,6 +15,7 @@ Usage: python test_vector_store.py --obvec # Test ObVecVectorStore only (needs seekdb / OceanBase) python test_vector_store.py --hologres # Test HologresVectorStore only python test_vector_store.py --zvec # Test ZvecVectorStore only + python test_vector_store.py --seekdb # Test SeekdbVectorStore (pyseekdb; Python >=3.11) python test_vector_store.py --all # Test all vector stores """ @@ -42,6 +44,14 @@ from reme.core.vector_store import ( ZvecVectorStore, ) +try: + from reme.core.vector_store import SeekdbVectorStore + + SEEKDB_AVAILABLE = True +except ImportError: + SeekdbVectorStore = None + SEEKDB_AVAILABLE = False + load_env() @@ -77,6 +87,14 @@ class TestConfig: PG_USE_HNSW = True # Use HNSW index for faster search PG_USE_DISKANN = False # Use DiskANN index (requires vectorscale extension) + # 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", "") + SEEKDB_DATABASE = os.environ.get("SEEKDB_DATABASE", "test") + # ChromaVectorStore settings CHROMA_PATH = "./test_vector_store_chroma" # For local persistent mode CHROMA_HOST = None # Set to host address for remote mode (e.g., "localhost") @@ -220,7 +238,9 @@ def get_store_type(store: BaseVectorStore) -> str: store: Vector store instance Returns: - str: Type identifier ("local", "es", "pgvector", "qdrant", "chroma", "obvec", "zvec", or "hologres") + str: Type identifier ( + "local", "es", "pgvector", "qdrant", "chroma", "obvec", "zvec", "hologres", or "seekdb" + ) """ if isinstance(store, LocalVectorStore): return "local" @@ -238,6 +258,8 @@ def get_store_type(store: BaseVectorStore) -> str: return "zvec" elif isinstance(store, HologresVectorStore): return "hologres" + elif SeekdbVectorStore is not None and isinstance(store, SeekdbVectorStore): + return "seekdb" else: raise ValueError(f"Unknown vector store type: {type(store)}") @@ -247,7 +269,9 @@ def create_vector_store(store_type: str, collection_name: str) -> BaseVectorStor """Create a vector store instance based on type. Args: - store_type: Type of vector store ("local", "es", "pgvector", "qdrant", "chroma", "obvec", or "hologres") + store_type: Type of vector store ( + "local", "es", "pgvector", "qdrant", "chroma", "obvec", "zvec", "hologres", or "seekdb" + ) collection_name: Name of the collection Returns: @@ -255,10 +279,16 @@ def create_vector_store(store_type: str, collection_name: str) -> BaseVectorStor """ config = TestConfig() - # Initialize embedding model + # api_key/base_url from env (Application layer does the same via embedding_api_key in application.py) + embedding_api_key = os.environ.get("EMBEDDING_API_KEY") or os.environ.get("OPENAI_API_KEY") + embedding_base_url = os.environ.get("EMBEDDING_BASE_URL") or os.environ.get("OPENAI_BASE_URL") + # use_dimensions=True so API output matches HNSW dim / vector column embedding_model = OpenAIEmbeddingModel( model_name=config.EMBEDDING_MODEL_NAME, dimensions=config.EMBEDDING_DIMENSIONS, + use_dimensions=True, + api_key=embedding_api_key or None, + base_url=embedding_base_url or None, ) if store_type == "local": @@ -346,6 +376,26 @@ def create_vector_store(store_type: str, collection_name: str) -> BaseVectorStor if config.HOLOGRES_DSN: kwargs["dsn"] = config.HOLOGRES_DSN return HologresVectorStore(**kwargs) + elif store_type == "seekdb": + if SeekdbVectorStore is None: + raise ImportError("SeekdbVectorStore not available; install pyseekdb (Python >=3.11)") + 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, + "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 + else: + kw["path"] = str(Path(db_path) / "seekdb.db") + return SeekdbVectorStore(**kw) else: raise ValueError(f"Unknown store type: {store_type}") @@ -667,6 +717,7 @@ async def test_copy_collection(store: BaseVectorStore, store_name: str): store_type = get_store_type(store) if store_type in ("es", "pgvector", "obvec", "hologres"): copy_collection_name = copy_collection_name.lower() + # seekdb uses collection names as-is # Clean up if exists collections = await store.list_collections() @@ -1848,6 +1899,13 @@ async def cleanup_store(store: BaseVectorStore, store_type: str): shutil.rmtree(test_dir) logger.info(f"Cleaned up zvec directory: {config.ZVEC_PATH}") + # Clean up temp directory if SeekdbVectorStore used temp dir + if store_type == "seekdb" and config.SEEKDB_PATH is None and getattr(store, "db_path", None): + test_dir = Path(store.db_path) + if test_dir.exists() and "test_seekdb_" in str(test_dir): + shutil.rmtree(test_dir, ignore_errors=True) + logger.info(f"Cleaned up seekdb temp directory: {test_dir}") + logger.info("✓ Cleanup completed") except Exception as e: logger.error(f"Cleanup error: {e}") @@ -1870,6 +1928,8 @@ Examples: python test_vector_store.py --chroma # Test ChromaVectorStore only python test_vector_store.py --obvec # Test ObVecVectorStore (seekdb / OceanBase) python test_vector_store.py --hologres # Test HologresVectorStore + python test_vector_store.py --zvec # Test ZvecVectorStore + python test_vector_store.py --seekdb # Test SeekdbVectorStore only python test_vector_store.py --all # Test all vector stores """, ) @@ -1913,6 +1973,11 @@ Examples: action="store_true", help="Test ZvecVectorStore", ) + parser.add_argument( + "--seekdb", + action="store_true", + help="Test SeekdbVectorStore (requires pyseekdb on Python >=3.11)", + ) parser.add_argument( "--all", action="store_true", @@ -1935,6 +2000,8 @@ Examples: ("hologres", "HologresVectorStore"), ("zvec", "ZvecVectorStore"), ] + if SEEKDB_AVAILABLE: + stores_to_test.append(("seekdb", "SeekdbVectorStore")) else: # Build list based on individual flags if args.local: @@ -1953,6 +2020,12 @@ Examples: stores_to_test.append(("hologres", "HologresVectorStore")) if args.zvec: stores_to_test.append(("zvec", "ZvecVectorStore")) + if args.seekdb: + if not SEEKDB_AVAILABLE: + raise ImportError( + "seekdb tests require pyseekdb (Python >=3.11); install: pip install 'pyseekdb>=1.2.0'", + ) + stores_to_test.append(("seekdb", "SeekdbVectorStore")) if not stores_to_test: # Default to all vector stores if no argument provided @@ -1966,9 +2039,12 @@ Examples: ("hologres", "HologresVectorStore"), ("zvec", "ZvecVectorStore"), ] + if SEEKDB_AVAILABLE: + stores_to_test.append(("seekdb", "SeekdbVectorStore")) print("No vector store specified, defaulting to test all vector stores") print( - "Use --local/--es/--pgvector/--qdrant/--chroma/--obvec/--zvec/--hologres to test specific ones\n", + "Use --local/--es/--pgvector/--qdrant/--chroma/--obvec/--zvec/--hologres/--seekdb " + "to test specific ones\n", ) # Run tests for each vector store