mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-10-09 03:20:54 +00:00
feat(vector_store): add OceanBase/seekdb vector store implementation (#201)
Some checks are pending
Pre-commit / run (ubuntu-latest) (push) Waiting to run
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:
parent
935e886af3
commit
f3d09aaa38
8 changed files with 578 additions and 18 deletions
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
453
reme/core/vector_store/obvec_vector_store.py
Normal file
453
reme/core/vector_store/obvec_vector_store.py
Normal 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")
|
||||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue