From f3d09aaa3899cf48f19644f9f98e0220526056e6 Mon Sep 17 00:00:00 2001 From: Chojan Shang Date: Thu, 9 Apr 2026 09:17:52 +0100 Subject: [PATCH] feat(vector_store): add OceanBase/seekdb vector store implementation (#201) * feat(vector_store): add OceanBase as a VectorStore * refactor(obvec): make it cleaner * docs: add obvec related info * refactor: minor update * refactor: clean code and pass lint * docs: remove unrelated edit * docs: minor update --- README.md | 2 +- README_ZH.md | 2 +- docs/index.md | 2 +- docs/vector_store_api_guide.md | 59 ++- pyproject.toml | 3 + reme/core/vector_store/__init__.py | 3 + reme/core/vector_store/obvec_vector_store.py | 453 +++++++++++++++++++ tests/test_vector_store.py | 72 ++- 8 files changed, 578 insertions(+), 18 deletions(-) create mode 100644 reme/core/vector_store/obvec_vector_store.py diff --git a/README.md b/README.md index 7c93e489..781f5412 100644 --- a/README.md +++ b/README.md @@ -506,7 +506,7 @@ async def main(): "dimensions": 1024, }, default_vector_store_config={ - "backend": "local", # Supports local/chroma/qdrant/elasticsearch + "backend": "local", # Supports local/chroma/qdrant/elasticsearch/obvec }, ) await reme.start() diff --git a/README_ZH.md b/README_ZH.md index 5fa2e4d2..7e6ccd58 100644 --- a/README_ZH.md +++ b/README_ZH.md @@ -486,7 +486,7 @@ async def main(): "dimensions": 1024, }, default_vector_store_config={ - "backend": "local", # 支持 local/chroma/qdrant/elasticsearch + "backend": "local", # 支持 local/chroma/qdrant/elasticsearch/obvec }, ) await reme.start() diff --git a/docs/index.md b/docs/index.md index 5cbbcb56..d7e2e716 100644 --- a/docs/index.md +++ b/docs/index.md @@ -139,7 +139,7 @@ response = requests.post("http://localhost:8002/retrieve_task_memory", json={ ## 📚 Resources - **[Installation Guide](installation.md)**, **[Quick Start](quick_start.md)**: Get started quickly with practical examples -- **[Vector Storage Setup](vector_store_api_guide.md)**: Configure local/vector databases and usage +- **[Vector Storage Setup](vector_store_api_guide.md)**: Configure local, Elasticsearch, Qdrant, ChromaDB, or ObVec (OceanBase / seekdb via pyobvector) storage and usage - **[MCP Guide](mcp_quick_start.md)**: Create MCP services - **[Personal Memory](personal_memory/personal_memory.md)**, **[Task Memory](task_memory/task_memory.md)** & **[Tool Memory](tool_memory/tool_memory.md)**: Operators used in personal memory, task memory and tool memory. You can modify the config to customize the pipelines. - **[Example Collection](./cookbook/appworld/quickstart.md)**: Real use cases and best practices diff --git a/docs/vector_store_api_guide.md b/docs/vector_store_api_guide.md index df0800ef..b0a06eab 100644 --- a/docs/vector_store_api_guide.md +++ b/docs/vector_store_api_guide.md @@ -33,8 +33,9 @@ FlowLLM provides multiple Vector Store implementations tailored to different use - **QdrantVectorStore** ([source code](https://github.com/flowllm-ai/flowllm/blob/main/flowllm/core/vector_store/qdrant_vector_store.py)): Built on the Qdrant vector database, supporting high-performance vector search. Recommended for large-scale production environments. - **ChromaVectorStore** ([source code](https://github.com/flowllm-ai/flowllm/blob/main/flowllm/core/vector_store/chroma_vector_store.py)): Based on ChromaDB, providing persistent storage and metadata filtering capabilities. - **EsVectorStore** ([source code](https://github.com/flowllm-ai/flowllm/blob/main/flowllm/core/vector_store/es_vector_store.py)): Built on Elasticsearch, enabling powerful combined full-text and vector search functionalities. +- **ObVecVectorStore** ([source code](https://github.com/agentscope-ai/ReMe/blob/main/reme/core/vector_store/obvec_vector_store.py)): Uses [pyobvector](https://pypi.org/project/pyobvector/) against **OceanBase** or **seekdb** (MySQL-compatible wire protocol). Suitable when you already run OceanBase/seekdb or need a SQL-native vector table with HNSW-style ANN search and JSON metadata filters. -All Vector Store implementations inherit from **BaseVectorStore** ([source code](https://github.com/flowllm-ai/flowllm/blob/main/flowllm/core/vector_store/base_vector_store.py)), ensuring a consistent interface specification. +All Vector Store implementations inherit from **BaseVectorStore** ([source code](https://github.com/agentscope-ai/ReMe/blob/main/reme/core/vector_store/base_vector_store.py)) in ReMe, ensuring a consistent interface specification. ## Core Features @@ -108,6 +109,28 @@ The asynchronous interface is particularly useful in the following scenarios: - **hosts**: Elasticsearch host address(es), either a string or a list (default: `http://localhost:9200`). - **basic_auth**: Basic authentication credentials (username and password). +### ObVecVectorStore Configuration + +- **uri**: Server address as `host:port` (default: `127.0.0.1:2881`). +- **user**: MySQL-compatible user. seekdb single-tenant images often use `root`; OceanBase multi-tenant setups typically use `root@` (e.g. `root@test`). +- **password**: Database password (seekdb Docker images commonly set this via `ROOT_PASSWORD`). +- **database**: Logical database name (default: `test`). +- **index_metric**: Distance metric for the vector index: `cosine` or `ip` (inner product); default `cosine`. +- **index_ef_search**: HNSW `ef_search` parameter passed to pyobvector (default: `100`). +- **collection_name**: Table name for the collection (from `VectorStoreConfig`, default `reme`). Use lowercase names if your deployment restricts identifiers. + +**Local seekdb via Docker** + +```text +docker run -d --name reme_seekdb -p 2881:2881 -e ROOT_PASSWORD= quay.io/oceanbase/seekdb:latest +``` + +**Integration tests** (requires a running server, embedding API credentials in `.env`, and matching DB password): + +```shell +OBVEC_PASSWORD= python tests/test_vector_store.py --obvec +``` + ## Configuration File Examples Configure Vector Store in `flowllm/config/default.yaml` under the `vector_store` section. The basic structure is as follows: @@ -128,7 +151,7 @@ vector_store.default.params.= ### Configuration Field Descriptions -- **`backend`** (required): Vector store backend type. Options: `local`, `memory`, `chroma`, `qdrant`, `elasticsearch`. +- **`backend`** (required): Vector store backend type. Options: `local`, `memory`, `chroma`, `qdrant`, `elasticsearch`, `obvec`. - **`embedding_model`** (required): Name of the embedding model configuration, referencing the `embedding_model` section. - **`params`** (optional): Dictionary of backend-specific parameters passed to the vector store constructor. @@ -295,6 +318,35 @@ vector_store.default.backend=elasticsearch vector_store.default.params.hosts='["http://es-node1:9200", "http://es-node2:9200", "http://es-node3:9200"]' ``` +#### 6. ObVecVectorStore Configuration (OceanBase / seekdb) + +**Implementation**: [`reme/core/vector_store/obvec_vector_store.py`](https://github.com/agentscope-ai/ReMe/blob/main/reme/core/vector_store/obvec_vector_store.py) + +**Example (seekdb on localhost)**: + +```yaml +vector_stores: + default: + backend: obvec + embedding_model: default + collection_name: reme + uri: "127.0.0.1:2881" + user: "root" + password: "your-root-password" + database: "test" + index_metric: "cosine" + index_ef_search: 100 +``` + +```shell +vector_stores.default.backend=obvec +vector_stores.default.uri=127.0.0.1:2881 +vector_stores.default.user=root +vector_stores.default.password=your-root-password +``` + +ReMe service YAML uses the key `vector_stores` (plural); CLI overrides use the same nested paths. + ### Complete Configuration Example Below is a complete `default.yaml` example including both embedding model and vector store configurations: @@ -352,8 +404,9 @@ Two types of metadata filtering are supported: - **Development & Testing**: Use MemoryVectorStore or LocalVectorStore—no additional services required. - **Small-Scale Applications**: Use LocalVectorStore or ChromaVectorStore for simplicity and ease of use. -- **Production Environments**: Use QdrantVectorStore or EsVectorStore for high performance and scalability. +- **Production Environments**: Use QdrantVectorStore, EsVectorStore, or ObVecVectorStore (OceanBase/seekdb) for high performance and scalability, depending on your existing infrastructure. - **Hybrid Search**: Use EsVectorStore to combine vector search with full-text search capabilities. +- **OceanBase / seekdb**: Use ObVecVectorStore when you standardize on pyobvector and SQL-accessible vector tables. ## Important Notes diff --git a/pyproject.toml b/pyproject.toml index 91956ce2..11177ffd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -49,6 +49,9 @@ dependencies = [ "openai>=2.8.1", "pandas>=2.3.3", "pydantic>=2.12.4", + "pyobvector>=0.1.20", + # pyobvector imports Expression from sqlglot; removed from sqlglot 30+ top-level API + "sqlglot>=25,<30", "qdrant-client>=1.16.0", "tavily-python>=0.7.13", "tiktoken>=0.12.0", diff --git a/reme/core/vector_store/__init__.py b/reme/core/vector_store/__init__.py index 9d411d03..8429b911 100644 --- a/reme/core/vector_store/__init__.py +++ b/reme/core/vector_store/__init__.py @@ -4,6 +4,7 @@ from .base_vector_store import BaseVectorStore from .chroma_vector_store import ChromaVectorStore from .es_vector_store import ESVectorStore from .local_vector_store import LocalVectorStore +from .obvec_vector_store import ObVecVectorStore from .pgvector_store import PGVectorStore from .qdrant_vector_store import QdrantVectorStore from ..registry_factory import R @@ -13,6 +14,7 @@ __all__ = [ "ChromaVectorStore", "ESVectorStore", "LocalVectorStore", + "ObVecVectorStore", "PGVectorStore", "QdrantVectorStore", ] @@ -20,5 +22,6 @@ __all__ = [ R.vector_stores.register("chroma")(ChromaVectorStore) R.vector_stores.register("es")(ESVectorStore) R.vector_stores.register("local")(LocalVectorStore) +R.vector_stores.register("obvec")(ObVecVectorStore) R.vector_stores.register("pgvector")(PGVectorStore) R.vector_stores.register("qdrant")(QdrantVectorStore) diff --git a/reme/core/vector_store/obvec_vector_store.py b/reme/core/vector_store/obvec_vector_store.py new file mode 100644 index 00000000..c8ad17c5 --- /dev/null +++ b/reme/core/vector_store/obvec_vector_store.py @@ -0,0 +1,453 @@ +"""OceanBase / seekdb vector store for ReMe (pyobvector).""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +from loguru import logger +from sqlalchemy import Column, JSON, String, text as sa_text +from sqlalchemy.dialects.mysql import LONGTEXT + +from .base_vector_store import BaseVectorStore +from ..embedding import BaseEmbeddingModel +from ..schema import VectorNode + +_OBVECTOR_IMPORT_ERROR: Exception | None = None + +try: + from pyobvector import IndexParams, ObVecClient, VecIndexType, VECTOR + from pyobvector import cosine_distance, inner_product +except Exception as e: + _OBVECTOR_IMPORT_ERROR = e + IndexParams = None # type: ignore[misc, assignment] + ObVecClient = None # type: ignore[misc, assignment] + VecIndexType = None # type: ignore[misc, assignment] + VECTOR = None # type: ignore[misc, assignment] + +_COL_SELECT = "id, content, vector, metadata" + + +def _is_safe_metadata_key(key: str) -> bool: + return bool(key.replace("_", "").replace(".", "").isalnum()) + + +def _coerce_db_vector(raw: Any) -> list[float] | None: + if raw is None: + return None + if isinstance(raw, list): + return [float(x) for x in raw] + if isinstance(raw, str): + try: + parsed = json.loads(raw) + if isinstance(parsed, list): + return [float(x) for x in parsed] + except (json.JSONDecodeError, TypeError, ValueError): + pass + return None + + +def _coerce_db_metadata(raw: Any) -> dict[str, Any]: + if raw is None: + return {} + if isinstance(raw, dict): + return raw + if isinstance(raw, str): + try: + parsed = json.loads(raw) + if isinstance(parsed, dict): + return parsed + except (json.JSONDecodeError, TypeError): + pass + return {} + + +def _build_metadata_filter_sql(filters: dict[str, Any] | None) -> str: + if not filters: + return "" + + parts: list[str] = [] + for key, value in filters.items(): + if not _is_safe_metadata_key(key): + continue + + path = f"$.{key}" + if isinstance(value, list) and len(value) == 2: + lo, hi = value[0], value[1] + if isinstance(lo, (int, float)) and isinstance(hi, (int, float)): + parts.append( + f"(JSON_EXTRACT(metadata, '{path}') >= {lo} AND " f"JSON_EXTRACT(metadata, '{path}') <= {hi})", + ) + else: + parts.append( + f"(JSON_EXTRACT(metadata, '{path}') >= '{lo}' AND " f"JSON_EXTRACT(metadata, '{path}') <= '{hi}')", + ) + elif isinstance(value, (int, float)): + parts.append(f"JSON_EXTRACT(metadata, '{path}') = {value}") + else: + parts.append(f"JSON_EXTRACT(metadata, '{path}') = '{value}'") + + return " AND ".join(parts) + + +def _format_vector_sql_literal(vector: list[float]) -> str: + return "[" + ",".join(str(float(v)) for v in vector) + "]" + + +def _normalize_embedding_for_ann(raw: Any) -> list[float]: + if hasattr(raw, "tolist"): + raw = raw.tolist() + return [float(x) for x in raw] + + +def _vector_node_from_db_row(row: tuple[Any, ...]) -> VectorNode: + return VectorNode( + vector_id=row[0], + content=row[1] or "", + vector=_coerce_db_vector(row[2]), + metadata=_coerce_db_metadata(row[3]), + ) + + +def _normalize_nodes(nodes: VectorNode | list[VectorNode]) -> list[VectorNode]: + return [nodes] if isinstance(nodes, VectorNode) else list(nodes) + + +def _sql_table(name: str) -> str: + return f"`{name}`" + + +class ObVecVectorStore(BaseVectorStore): + """OceanBase or seekdb vector store for dense vectors and kNN search. + + Args: + index_metric: ``cosine`` or ``ip`` (inner product). Invalid values raise + ``ValueError``; unsupported strings are not mapped to another metric. + """ + + def __init__( + self, + collection_name: str, + db_path: str | Path, + embedding_model: BaseEmbeddingModel, + uri: str = "127.0.0.1:2881", + user: str = "root", + password: str = "", + database: str = "test", + index_metric: str = "cosine", + index_ef_search: int = 100, + **kwargs, + ): + if _OBVECTOR_IMPORT_ERROR is not None: + raise ImportError( + "ObVecVectorStore requires pyobvector. Install with `pip install pyobvector`", + ) from _OBVECTOR_IMPORT_ERROR + + super().__init__( + collection_name=collection_name, + db_path=db_path, + embedding_model=embedding_model, + **kwargs, + ) + + key = index_metric.strip().lower() + if key not in ("cosine", "ip"): + raise ValueError( + f"ObVecVectorStore index_metric must be 'cosine' or 'ip', got {index_metric!r}", + ) + self.uri = uri + self.user = user + self.password = password + self.database = database + self.index_metric = key + self.index_ef_search = index_ef_search + + self.client: ObVecClient | None = None + self.embedding_model_dims = embedding_model.dimensions + + def _require_client(self) -> ObVecClient: + if self.client is None: + raise RuntimeError("ObVecVectorStore.start() must be called before this operation") + return self.client + + async def list_collections(self) -> list[str]: + client = self._require_client() + result = client.perform_raw_text_sql(f"SHOW TABLES FROM {_sql_table(self.database)}") + rows = result.fetchall() + return [row[0] for row in rows if row] + + def _table_columns_for_create(self, dimensions: int) -> list[Any]: + return [ + Column("id", String(255), primary_key=True), + Column("content", LONGTEXT), + Column("vector", VECTOR(dimensions)), + Column("metadata", JSON), + ] + + def _hnsw_index_params(self, collection_name: str) -> IndexParams: + metric = "cosine" if self.index_metric == "cosine" else "inner_product" + vidxs = IndexParams() + vidxs.add_index( + "vector", + VecIndexType.HNSW, + f"{collection_name}_vidx", + metric_type=metric, + params={"efSearch": self.index_ef_search}, + ) + return vidxs + + async def create_collection(self, collection_name: str, **kwargs): + client = self._require_client() + dimensions = kwargs.get("dimensions", self.embedding_model_dims) + + if client.check_table_exists(collection_name): + logger.info("Collection {} already exists", collection_name) + return + + columns = self._table_columns_for_create(dimensions) + vidxs = self._hnsw_index_params(collection_name) + + client.create_table_with_index_params( + table_name=collection_name, + columns=columns, + vidxs=vidxs, + ) + logger.info("Created collection {} with dimensions={}", collection_name, dimensions) + + async def delete_collection(self, collection_name: str, **kwargs): + client = self._require_client() + client.drop_table_if_exist(collection_name) + logger.info("Deleted collection {}", collection_name) + + async def copy_collection(self, collection_name: str, **kwargs): + client = self._require_client() + + if not client.check_table_exists(self.collection_name): + raise ValueError(f"Source collection {self.collection_name} does not exist") + + await self.create_collection(collection_name) + + try: + source_data = await self.list(limit=None) + if source_data: + await self.insert(source_data, collection_name=collection_name) + logger.info("Copied collection {} to {}", self.collection_name, collection_name) + except Exception: + try: + client.drop_table_if_exist(collection_name) + except Exception as cleanup_err: + logger.warning("Cleanup after failed copy failed: {}", cleanup_err) + raise + + async def insert(self, nodes: VectorNode | list[VectorNode], **kwargs): + nodes = _normalize_nodes(nodes) + if not nodes: + return + + client = self._require_client() + + need_emb = [n for n in nodes if n.vector is None] + if need_emb: + filled = await self.get_node_embeddings(need_emb) + by_id = {n.vector_id: n for n in filled} + nodes_to_insert = [by_id.get(n.vector_id, n) for n in nodes] + else: + nodes_to_insert = nodes + + data = [ + { + "id": node.vector_id, + "content": node.content, + "vector": node.vector if node.vector is not None else [], + "metadata": node.metadata if node.metadata else {}, + } + for node in nodes_to_insert + ] + target = kwargs.get("collection_name", self.collection_name) + client.insert(table_name=target, data=data) + logger.info("Inserted {} documents into {}", len(nodes_to_insert), target) + + async def search( + self, + query: str, + limit: int = 5, + filters: dict | None = None, + **kwargs, + ) -> list[VectorNode]: + client = self._require_client() + raw_vec = await self.get_embedding(query) + query_vector = _normalize_embedding_for_ann(raw_vec) + dist_fn = cosine_distance if self.index_metric == "cosine" else inner_product + + filter_sql = _build_metadata_filter_sql(filters) + where_parts = [sa_text(filter_sql)] if filter_sql else None + + results = client.ann_search( + table_name=self.collection_name, + vec_data=query_vector, + vec_column_name="vector", + distance_func=dist_fn, + with_dist=True, + topk=limit, + output_column_names=["id", "content", "metadata"], + where_clause=where_parts, + ) + + score_threshold = kwargs.get("score_threshold") + out: list[VectorNode] = [] + for row in results: + if len(row) < 4: + raise RuntimeError( + "ann_search row must have id, content, metadata, distance " f"(got {len(row)} columns)", + ) + vid, content, metadata_raw, distance = row[0], row[1], row[2], row[3] + dist_f = float(distance) + if self.index_metric == "cosine": + score = max(0.0, 1.0 - dist_f / 2.0) + else: + score = max(0.0, dist_f) + if score_threshold is not None and score < score_threshold: + continue + meta = _coerce_db_metadata(metadata_raw) if metadata_raw is not None else {} + meta["score"] = score + meta["_distance"] = dist_f + out.append( + VectorNode( + vector_id=vid, + content=content or "", + vector=None, + metadata=meta, + ), + ) + return out + + async def delete(self, vector_ids: str | list[str], **kwargs): + if isinstance(vector_ids, str): + vector_ids = [vector_ids] + if not vector_ids: + return + client = self._require_client() + client.delete(self.collection_name, ids=vector_ids) + logger.info("Deleted {} documents from {}", len(vector_ids), self.collection_name) + + async def delete_all(self, **kwargs): + client = self._require_client() + client.delete(self.collection_name) + logger.info("Deleted all documents from {}", self.collection_name) + + async def update(self, nodes: VectorNode | list[VectorNode], **kwargs): + nodes = _normalize_nodes(nodes) + if not nodes: + return + + client = self._require_client() + need_emb = [n for n in nodes if n.vector is None and bool(n.content)] + if need_emb: + filled = await self.get_node_embeddings(need_emb) + by_id = {n.vector_id: n for n in filled} + nodes_to_update = [ + by_id.get(n.vector_id, n) if (n.vector is None and bool(n.content)) else n for n in nodes + ] + else: + nodes_to_update = nodes + + for node in nodes_to_update: + updates: list[str] = [] + params: dict[str, Any] = {} + + if node.content is not None: + updates.append("content = :content") + params["content"] = node.content + + if node.vector is not None: + updates.append("vector = :vector") + params["vector"] = _format_vector_sql_literal(node.vector) + + if node.metadata is not None: + updates.append("metadata = :metadata") + params["metadata"] = json.dumps(node.metadata) + + if not updates: + continue + + params["vid"] = node.vector_id + update_sql = f"UPDATE {_sql_table(self.collection_name)} SET {', '.join(updates)} WHERE id = :vid" + with client.engine.connect() as conn: + with conn.begin(): + conn.execute(sa_text(update_sql), params) + + logger.info("Updated {} documents in {}", len(nodes_to_update), self.collection_name) + + async def get(self, vector_ids: str | list[str]) -> VectorNode | list[VectorNode] | None: + single = isinstance(vector_ids, str) + if single: + vector_ids = [vector_ids] + if not vector_ids: + return [] if not single else None + + client = self._require_client() + ids_str = "', '".join(vector_ids) + select_sql = f"SELECT {_COL_SELECT} FROM {_sql_table(self.collection_name)} " f"WHERE id IN ('{ids_str}')" + result = client.perform_raw_text_sql(select_sql) + rows = result.fetchall() + parsed = [_vector_node_from_db_row(row) for row in rows if row] + if single: + return parsed[0] if parsed else None + return parsed + + async def list( + self, + filters: dict | None = None, + limit: int | None = None, + sort_key: str | None = None, + reverse: bool = False, + ) -> list[VectorNode]: + client = self._require_client() + select_sql = f"SELECT {_COL_SELECT} FROM {_sql_table(self.collection_name)}" + where_clause = _build_metadata_filter_sql(filters) + if where_clause: + select_sql += f" WHERE {where_clause}" + if sort_key and _is_safe_metadata_key(sort_key): + order = "DESC" if reverse else "ASC" + select_sql += f" ORDER BY JSON_EXTRACT(metadata, '$.{sort_key}') {order}" + if limit is not None: + select_sql += f" LIMIT {limit}" + result = client.perform_raw_text_sql(select_sql) + rows = result.fetchall() + return [_vector_node_from_db_row(row) for row in rows if row] + + async def collection_info(self) -> dict[str, Any]: + """Return collection name and row count.""" + client = self._require_client() + count_sql = f"SELECT COUNT(*) FROM {_sql_table(self.collection_name)}" + result = client.perform_raw_text_sql(count_sql) + row = result.fetchone() + count = row[0] if row else 0 + return {"name": self.collection_name, "count": count} + + async def reset(self): + """Drop and recreate the current collection table.""" + logger.warning("Resetting collection {}...", self.collection_name) + await self.delete_collection(self.collection_name) + await self.create_collection(self.collection_name) + + async def reset_collection(self, collection_name: str): + self.collection_name = collection_name + await self.create_collection(collection_name) + logger.info("Collection reset to {}", collection_name) + + async def start(self) -> None: + self.client = ObVecClient( + uri=self.uri, + user=self.user, + password=self.password, + db_name=self.database, + ) + + await super().start() + logger.info("seekdb / OceanBase vector table {} ready", self.collection_name) + + async def close(self): + self.client = None + logger.info("ObVec client connection closed") diff --git a/tests/test_vector_store.py b/tests/test_vector_store.py index ccdd9c28..7f7264be 100644 --- a/tests/test_vector_store.py +++ b/tests/test_vector_store.py @@ -2,8 +2,8 @@ """Unified test suite for vector store implementations. This module provides comprehensive test coverage for LocalVectorStore, ESVectorStore, -PGVectorStore, QdrantVectorStore, and ChromaVectorStore implementations. Tests can be -run for specific vector stores or all implementations. +PGVectorStore, QdrantVectorStore, ChromaVectorStore, and ObVecVectorStore implementations. +Tests can be run for specific vector stores or all implementations. Usage: python test_vector_store.py --local # Test LocalVectorStore only @@ -11,12 +11,13 @@ Usage: python test_vector_store.py --pgvector # Test PGVectorStore only python test_vector_store.py --qdrant # Test QdrantVectorStore only python test_vector_store.py --chroma # Test ChromaVectorStore only + python test_vector_store.py --obvec # Test ObVecVectorStore only (needs seekdb / OceanBase) python test_vector_store.py --all # Test all vector stores - """ import argparse import asyncio +import os import shutil import tempfile from pathlib import Path @@ -32,12 +33,19 @@ from reme.core.vector_store import ( ChromaVectorStore, LocalVectorStore, ESVectorStore, + ObVecVectorStore, PGVectorStore, QdrantVectorStore, ) load_env() + +def _search_score_for_log(metadata: dict) -> object: + """Similarity score for log lines (implementations use ``metadata['score']``).""" + return metadata.get("score", metadata.get("_score", "N/A")) + + # ==================== Configuration ==================== @@ -73,6 +81,15 @@ class TestConfig: CHROMA_TENANT = None # Set for ChromaDB Cloud tenant CHROMA_DATABASE = None # Set for ChromaDB Cloud database + # ObVecVectorStore: seekdb docker often uses user `root` + ROOT_PASSWORD; OceanBase + # multi-tenant commonly uses `root@` (see pyobvector defaults). + # OBVEC_PASSWORD default `root` matches docker-compose.obvec.yml only—override if your + # seekdb uses another ROOT_PASSWORD (e.g. another compose stack on the same port). + OBVEC_URI = os.environ.get("OBVEC_URI", "127.0.0.1:2881") + OBVEC_USER = os.environ.get("OBVEC_USER", "root") + OBVEC_PASSWORD = os.environ.get("OBVEC_PASSWORD", "root") + OBVEC_DATABASE = os.environ.get("OBVEC_DATABASE", "test") + # Embedding model settings EMBEDDING_MODEL_NAME = "text-embedding-v4" EMBEDDING_DIMENSIONS = 64 @@ -182,7 +199,7 @@ def get_store_type(store: BaseVectorStore) -> str: store: Vector store instance Returns: - str: Type identifier ("local", "es", "pgvector", "qdrant", or "chroma") + str: Type identifier ("local", "es", "pgvector", "qdrant", "chroma", or "obvec") """ if isinstance(store, LocalVectorStore): return "local" @@ -194,6 +211,8 @@ def get_store_type(store: BaseVectorStore) -> str: return "pgvector" elif isinstance(store, ChromaVectorStore): return "chroma" + elif isinstance(store, ObVecVectorStore): + return "obvec" else: raise ValueError(f"Unknown vector store type: {type(store)}") @@ -202,7 +221,7 @@ 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", or "chroma") + store_type: Type of vector store ("local", "es", "pgvector", "qdrant", "chroma", or "obvec") collection_name: Name of the collection Returns: @@ -264,6 +283,18 @@ def create_vector_store(store_type: str, collection_name: str) -> BaseVectorStor tenant=config.CHROMA_TENANT, database=config.CHROMA_DATABASE, ) + elif store_type == "obvec": + return ObVecVectorStore( + collection_name=collection_name, + embedding_model=embedding_model, + db_path=tempfile.mkdtemp(prefix="test_obvec_"), + uri=config.OBVEC_URI, + user=config.OBVEC_USER, + password=config.OBVEC_PASSWORD, + database=config.OBVEC_DATABASE, + index_metric="cosine", + index_ef_search=100, + ) else: raise ValueError(f"Unknown store type: {store_type}") @@ -327,7 +358,7 @@ async def test_search(store: BaseVectorStore, _store_name: str): logger.info(f"Search returned {len(results)} results") for i, r in enumerate(results, 1): - score = r.metadata.get("_score", "N/A") + score = _search_score_for_log(r.metadata) logger.info(f" Result {i}: {r.content[:60]}... (score: {score})") assert len(results) > 0, "Search should return results" @@ -581,9 +612,9 @@ async def test_copy_collection(store: BaseVectorStore, store_name: str): config = TestConfig() copy_collection_name = f"{config.TEST_COLLECTION_PREFIX}_{store_name}_copy" - # Elasticsearch and PostgreSQL require lowercase table/index names + # Elasticsearch, PostgreSQL and OceanBase require lowercase table/index names store_type = get_store_type(store) - if store_type in ("es", "pgvector"): + if store_type in ("es", "pgvector", "obvec"): copy_collection_name = copy_collection_name.lower() # Clean up if exists @@ -1010,7 +1041,7 @@ async def test_search_relevance_ranking(store: BaseVectorStore, _store_name: str logger.info(f"Search results for: '{query}'") for i, result in enumerate(results, 1): - score = result.metadata.get("_score", "N/A") + score = _search_score_for_log(result.metadata) relevance = result.metadata.get("relevance", "unknown") logger.info(f" {i}. [{relevance}] score={score}: {result.content[:60]}...") @@ -1028,7 +1059,7 @@ async def test_search_relevance_ranking(store: BaseVectorStore, _store_name: str results2 = await store.search(query=query2, limit=5) logger.info(f"\nSearch results for: '{query2}'") for i, result in enumerate(results2, 1): - score = result.metadata.get("_score", "N/A") + score = _search_score_for_log(result.metadata) logger.info(f" {i}. score={score}: {result.content[:60]}...") logger.info("✓ Search relevance ranking test passed") @@ -1718,7 +1749,7 @@ async def cleanup_store(store: BaseVectorStore, store_type: str): Args: store: Vector store instance - store_type: Type of vector store ("local" or "es") + store_type: Backend key (e.g. ``"local"``, ``"obvec"``) """ logger.info("=" * 20 + " CLEANUP " + "=" * 20) @@ -1752,6 +1783,13 @@ async def cleanup_store(store: BaseVectorStore, store_type: str): shutil.rmtree(test_dir) logger.info(f"Cleaned up chroma directory: {config.CHROMA_PATH}") + # ObVecVectorStore uses a temp db_path per run (reserved for local sidecar files). + if store_type == "obvec": + obvec_dir = getattr(store, "db_path", None) + if obvec_dir and Path(obvec_dir).exists(): + shutil.rmtree(obvec_dir, ignore_errors=True) + logger.info(f"Cleaned up obvec temp directory: {obvec_dir}") + logger.info("✓ Cleanup completed") except Exception as e: logger.error(f"Cleanup error: {e}") @@ -1772,6 +1810,7 @@ Examples: python test_vector_store.py --pgvector # Test PGVectorStore only python test_vector_store.py --qdrant # Test QdrantVectorStore only python test_vector_store.py --chroma # Test ChromaVectorStore only + python test_vector_store.py --obvec # Test ObVecVectorStore (seekdb / OceanBase) python test_vector_store.py --all # Test all vector stores """, ) @@ -1800,6 +1839,11 @@ Examples: action="store_true", help="Test ChromaVectorStore", ) + parser.add_argument( + "--obvec", + action="store_true", + help="Test ObVecVectorStore", + ) parser.add_argument( "--all", action="store_true", @@ -1818,6 +1862,7 @@ Examples: ("pgvector", "PGVectorStore"), ("qdrant", "QdrantVectorStore"), ("chroma", "ChromaVectorStore"), + ("obvec", "ObVecVectorStore"), ] else: # Build list based on individual flags @@ -1831,6 +1876,8 @@ Examples: stores_to_test.append(("qdrant", "QdrantVectorStore")) if args.chroma: stores_to_test.append(("chroma", "ChromaVectorStore")) + if args.obvec: + stores_to_test.append(("obvec", "ObVecVectorStore")) if not stores_to_test: # Default to all vector stores if no argument provided @@ -1840,10 +1887,11 @@ Examples: ("pgvector", "PGVectorStore"), ("qdrant", "QdrantVectorStore"), ("chroma", "ChromaVectorStore"), + ("obvec", "ObVecVectorStore"), ] print("No vector store specified, defaulting to test all vector stores") print( - "Use --local/--es/--pgvector/--qdrant/--chroma to test specific ones\n", + "Use --local/--es/--pgvector/--qdrant/--chroma/--obvec to test specific ones\n", ) # Run tests for each vector store