feat(vector_store): add OceanBase/seekdb vector store implementation (#201)
Some checks are pending
Pre-commit / run (ubuntu-latest) (push) Waiting to run

* 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
This commit is contained in:
Chojan Shang 2026-04-09 09:17:52 +01:00 • committed by GitHub
parent 935e886af3
commit f3d09aaa38
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 578 additions and 18 deletions

View file

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

View file

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

View file

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

View file

@ -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@<tenant>` (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=<your_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=<your_root_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.<param_name>=<param_value>
### 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

View file

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

View file

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

View file

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

View file

@ -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@<tenant>` (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