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

refactor(obvec): make it cleaner

Browse source
This commit is contained in:
Chojan Shang 2026-04-03 11:20:34 +08:00
parent 1a72c80aba
commit a1c1be3c57
1 changed files with 191 additions and 262 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

453
reme/core/vector_store/obvec_vector_store.py
View file

@ -1,14 +1,18 @@
"""OceanBase vector store implementation for ReMe.
"""OceanBase / SeekDB vector store for ReMe (pyobvector).
This module provides an OceanBase-based vector store that implements the BaseVectorStore
interface for high-performance dense vector storage and retrieval using pyobvector.
Dense vectors, kNN via ``ObVecClient.ann_search``, JSON metadata filters, and
helpers for coercion / metrics live in this module (pyobvector-aligned).
"""
from __future__ import annotations
import json
from collections.abc import Callable
from pathlib import Path
from typing import Any
from loguru import logger
from sqlalchemy import text as sa_text
from .base_vector_store import BaseVectorStore
from ..embedding import BaseEmbeddingModel
@ -25,39 +29,117 @@ except Exception as e:
VecIndexType = None
VECTOR = None
# HNSW index metric strings for pyobvector ``IndexParam`` (metric_type / distance).
_METRIC_TO_INDEX_DISTANCE: dict[str, str] = {
"cosine": "cosine",
"l2": "l2_distance",
"ip": "inner_product",
}
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 _get_distance_function(metric: str) -> Callable[..., Any]:
from pyobvector import cosine_distance, inner_product, l2_distance
registry: dict[str, Callable[..., Any]] = {
"cosine": cosine_distance,
"l2": l2_distance,
"ip": inner_product,
}
return registry.get(metric.lower(), cosine_distance)
def _similarity_from_distance(metric: str, distance: float) -> float:
m = metric.lower()
if m in ("cosine", "l2"):
return max(0.0, 1.0 - distance / 2.0)
return max(0.0, float(distance))
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]),
)
class ObVecVectorStore(BaseVectorStore):
"""OceanBase-based vector store for dense vector storage and kNN search."""
@staticmethod
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
@staticmethod
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 {}
"""OceanBase or SeekDB vector store for dense vectors and kNN search."""
def __init__(
self,
@ -73,24 +155,9 @@ class ObVecVectorStore(BaseVectorStore):
index_ef_search: int = 100,
**kwargs,
):
"""Initialize the OceanBase vector store with connection parameters.
Args:
collection_name: Name of the collection (table).
db_path: Database path (not used for remote OceanBase, kept for API consistency).
embedding_model: Model instance used to generate vector embeddings.
uri: Connection URI for OceanBase server.
user: Username for authentication.
password: Password for authentication.
database: Database name to use.
index_type: Type of vector index (HNSW).
index_metric: Distance metric (cosine, l2, ip).
index_ef_search: HNSW ef_search parameter.
**kwargs: Additional configuration passed to the base class.
"""
if _OBVECTOR_IMPORT_ERROR is not None:
raise ImportError(
"ObVecVectorStore requires extra dependencies. Install with `pip install pyobvector`",
"ObVecVectorStore requires pyobvector. Install with `pip install pyobvector`",
) from _OBVECTOR_IMPORT_ERROR
super().__init__(
@ -112,9 +179,6 @@ class ObVecVectorStore(BaseVectorStore):
self.embedding_model_dims = embedding_model.dimensions
async def list_collections(self) -> list[str]:
"""List all available table names in the current database."""
# OceanBase doesn't have a direct list collections API
# We'll return the current collection name if it exists
if self.client is None:
return []
try:
@ -124,26 +188,23 @@ class ObVecVectorStore(BaseVectorStore):
rows = result.fetchall()
return [row[0] for row in rows if row]
except Exception as e:
logger.warning(f"Failed to list collections: {e}")
logger.warning("Failed to list collections: {}", e)
return []
async def create_collection(self, collection_name: str, **kwargs):
"""Create a new table with vector support and appropriate indexing."""
dimensions = kwargs.get("dimensions", self.embedding_model_dims)
if self.client is None:
logger.warning("Client not initialized, skipping collection creation")
return
# Check if table already exists
if self.client.check_table_exists(collection_name):
logger.info(f"Collection {collection_name} already exists")
logger.info("Collection {} already exists", collection_name)
return
from sqlalchemy import Column, String, Text, JSON, create_engine
from sqlalchemy import Column, JSON, String
from sqlalchemy.dialects.mysql import LONGTEXT
# Create columns
columns = [
Column("id", String(255), primary_key=True),
Column("content", LONGTEXT),
@ -153,12 +214,7 @@ class ObVecVectorStore(BaseVectorStore):
vidxs: IndexParams | None = None
if self.index_type == "HNSW":
metric_map = {
"cosine": "cosine",
"l2": "l2_distance",
"ip": "inner_product",
}
metric = metric_map.get(self.index_metric, "cosine")
metric = _METRIC_TO_INDEX_DISTANCE.get(self.index_metric, "cosine")
vidxs = IndexParams()
vidxs.add_index(
"vector",
@ -169,63 +225,52 @@ class ObVecVectorStore(BaseVectorStore):
)
try:
# Create table with index
self.client.create_table_with_index_params(
table_name=collection_name,
columns=columns,
vidxs=vidxs,
)
logger.info(f"Created collection {collection_name} with dimensions={dimensions}")
logger.info("Created collection {} with dimensions={}", collection_name, dimensions)
except Exception as e:
logger.error(f"Failed to create collection {collection_name}: {e}")
logger.error("Failed to create collection {}: {}", collection_name, e)
raise
async def delete_collection(self, collection_name: str, **kwargs):
"""Permanently delete a collection table from the database."""
if self.client is None:
logger.warning("Client not initialized, skipping collection deletion")
return
try:
self.client.drop_table_if_exist(collection_name)
logger.info(f"Deleted collection {collection_name}")
logger.info("Deleted collection {}", collection_name)
except Exception as e:
logger.error(f"Failed to delete collection {collection_name}: {e}")
logger.error("Failed to delete collection {}: {}", collection_name, e)
raise
async def copy_collection(self, collection_name: str, **kwargs):
"""Duplicate the current collection to a new one with the given name."""
if self.client is None:
logger.warning("Client not initialized, skipping collection copy")
return
# Check if source collection exists
if not self.client.check_table_exists(self.collection_name):
raise ValueError(f"Source collection {self.collection_name} does not exist")
# Create new collection with same structure
await self.create_collection(collection_name)
# Copy data from source to target
try:
# Get all data from source
source_data = await self.list(limit=None)
# Insert into new collection
if source_data:
await self.insert(source_data, collection_name=collection_name)
logger.info(f"Copied collection {self.collection_name} to {collection_name}")
except Exception as e:
# Clean up the created collection if copy fails
logger.info("Copied collection {} to {}", self.collection_name, collection_name)
except Exception:
try:
self.client.drop_table_if_exist(collection_name)
except:
pass
except Exception as cleanup_err:
logger.warning("Cleanup after failed copy failed: {}", cleanup_err)
raise
async def insert(self, nodes: VectorNode | list[VectorNode], **kwargs):
"""Insert or upsert vector nodes into the OceanBase collection."""
if isinstance(nodes, VectorNode):
nodes = [nodes]
@ -236,16 +281,15 @@ class ObVecVectorStore(BaseVectorStore):
if nodes_without_vectors:
nodes_with_vectors = await self.get_node_embeddings(nodes_without_vectors)
vector_map = {n.vector_id: n for n in nodes_with_vectors}
nodes_to_insert = [vector_map.get(n.vector_id, n) if n.vector is None else n for n in nodes]
nodes_to_insert = [
vector_map.get(n.vector_id, n) if n.vector is None else n for n in nodes
]
else:
nodes_to_insert = nodes
# Prepare data for insertion
data = []
for node in nodes_to_insert:
# Convert vector to list if it's not already
vector_list = node.vector if node.vector is not None else []
data.append({
"id": node.vector_id,
"content": node.content,
@ -253,7 +297,6 @@ class ObVecVectorStore(BaseVectorStore):
"metadata": node.metadata if node.metadata else {},
})
# Determine which collection to use
target_collection = kwargs.get("collection_name", self.collection_name)
try:
@ -261,48 +304,11 @@ class ObVecVectorStore(BaseVectorStore):
table_name=target_collection,
data=data,
)
logger.info(f"Inserted {len(nodes_to_insert)} documents into {target_collection}")
logger.info("Inserted {} documents into {}", len(nodes_to_insert), target_collection)
except Exception as e:
logger.error(f"Failed to insert documents: {e}")
logger.error("Failed to insert documents: {}", e)
raise
def _build_filter_clause(self, filters: dict | None) -> str:
"""Generate a WHERE clause from filter dictionary.
Supports two filter formats:
1. Range query: {"field": [start_value, end_value]} - filters for field >= start_value AND field <= end_value
2. Exact match: {"field": value} - filters for field == value
"""
if not filters:
return ""
conditions = []
for key, value in filters.items():
# Sanitize key to prevent SQL injection
if not key.replace("_", "").replace(".", "").isalnum():
continue
# New syntax: [start, end] represents a range query
if isinstance(value, list) and len(value) == 2:
lo, hi = value[0], value[1]
if isinstance(lo, (int, float)) and isinstance(hi, (int, float)):
conditions.append(
f"(JSON_EXTRACT(metadata, '$.{key}') >= {lo} AND "
f"JSON_EXTRACT(metadata, '$.{key}') <= {hi})",
)
else:
conditions.append(
f"(JSON_EXTRACT(metadata, '$.{key}') >= '{lo}' AND "
f"JSON_EXTRACT(metadata, '$.{key}') <= '{hi}')",
)
else:
if isinstance(value, (int, float)):
conditions.append(f"JSON_EXTRACT(metadata, '$.{key}') = {value}")
else:
conditions.append(f"JSON_EXTRACT(metadata, '$.{key}') = '{value}'")
return " AND ".join(conditions)
async def search(
self,
query: str,
@ -310,32 +316,14 @@ class ObVecVectorStore(BaseVectorStore):
filters: dict | None = None,
**kwargs,
) -> list[VectorNode]:
"""Perform a kNN similarity search based on a text query."""
query_vector = await self.get_embedding(query)
if hasattr(query_vector, "tolist"):
query_vector = query_vector.tolist()
else:
query_vector = [float(x) for x in query_vector]
raw_vec = await self.get_embedding(query)
query_vector = _normalize_embedding_for_ann(raw_vec)
distance_func = _get_distance_function(self.index_metric)
# Determine distance function based on metric
if self.index_metric == "cosine":
from pyobvector import cosine_distance
distance_func = cosine_distance
elif self.index_metric == "l2":
from pyobvector import l2_distance
distance_func = l2_distance
else: # ip (inner product)
from pyobvector import inner_product
distance_func = inner_product
# ann_search passes where_clause to SQLAlchemy where(*...); use text(), not raw str.
from sqlalchemy import text as sa_text
filter_sql = self._build_filter_clause(filters)
filter_sql = _build_metadata_filter_sql(filters)
where_parts = [sa_text(filter_sql)] if filter_sql else None
try:
# pyobvector expects vec_data as a flat list of floats, not [vector].
results = self.client.ann_search(
table_name=self.collection_name,
vec_data=query_vector,
@ -347,54 +335,39 @@ class ObVecVectorStore(BaseVectorStore):
where_clause=where_parts,
)
search_results = []
search_results: list[VectorNode] = []
score_threshold = kwargs.get("score_threshold")
for row in results:
# row format: (id, content, metadata, distance)
if len(row) >= 4:
vector_id = row[0]
content = row[1]
metadata_str = row[2]
distance = row[3]
if len(row) < 4:
continue
vector_id, content, metadata_raw, distance = row[0], row[1], row[2], row[3]
score = _similarity_from_distance(self.index_metric, float(distance))
# Convert distance to similarity score
if self.index_metric == "cosine":
score = max(0.0, 1.0 - distance / 2.0)
elif self.index_metric == "l2":
# For L2, we need to normalize - this is a rough approximation
score = max(0.0, 1.0 - distance / 2.0)
else: # inner product
# For IP, higher is better, no conversion needed
score = max(0.0, distance)
if score_threshold is not None and score < score_threshold:
continue
if score_threshold is not None and score < score_threshold:
continue
metadata: dict[str, Any] = {}
if metadata_raw:
metadata = _coerce_db_metadata(metadata_raw)
metadata["score"] = score
metadata["_distance"] = distance
# Parse metadata
metadata = {}
if metadata_str:
metadata = self._coerce_db_metadata(metadata_str)
metadata["score"] = score
metadata["_distance"] = distance
search_results.append(
VectorNode(
vector_id=vector_id,
content=content or "",
vector=None, # Don't return vector in search results
metadata=metadata,
),
)
search_results.append(
VectorNode(
vector_id=vector_id,
content=content or "",
vector=None,
metadata=metadata,
),
)
return search_results
except Exception as e:
logger.error(f"Search failed: {e}")
logger.error("Search failed: {}", e)
return []
async def delete(self, vector_ids: str | list[str], **kwargs):
"""Remove specific vector records from the collection by their IDs."""
if isinstance(vector_ids, str):
vector_ids = [vector_ids]
@ -403,22 +376,20 @@ class ObVecVectorStore(BaseVectorStore):
try:
self.client.delete(self.collection_name, ids=vector_ids)
logger.info(f"Deleted {len(vector_ids)} documents from {self.collection_name}")
logger.info("Deleted {} documents from {}", len(vector_ids), self.collection_name)
except Exception as e:
logger.error(f"Failed to delete documents: {e}")
logger.error("Failed to delete documents: {}", e)
raise
async def delete_all(self, **kwargs):
"""Remove all vectors from the collection."""
try:
self.client.delete(self.collection_name)
logger.info(f"Deleted all documents from {self.collection_name}")
logger.info("Deleted all documents from {}", self.collection_name)
except Exception as e:
logger.error(f"Failed to delete all documents: {e}")
logger.error("Failed to delete all documents: {}", e)
raise
async def update(self, nodes: VectorNode | list[VectorNode], **kwargs):
"""Update existing vector nodes with new content, embeddings, or metadata."""
if isinstance(nodes, VectorNode):
nodes = [nodes]
@ -429,7 +400,9 @@ class ObVecVectorStore(BaseVectorStore):
if nodes_without_vectors:
nodes_with_vectors = await self.get_node_embeddings(nodes_without_vectors)
vector_map = {n.vector_id: n for n in nodes_with_vectors}
nodes_to_update = [vector_map.get(n.vector_id, n) if n.vector is None and n.content else n for n in nodes]
nodes_to_update = [
vector_map.get(n.vector_id, n) if n.vector is None and n.content else n for n in nodes
]
else:
nodes_to_update = nodes
@ -446,29 +419,29 @@ class ObVecVectorStore(BaseVectorStore):
if node.vector is not None:
updates.append("vector = :vector")
# VECTOR column expects a single literal, not a Python list bound as multi-column
params["vector"] = "[" + ",".join(str(float(v)) for v in node.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 updates:
params["vid"] = node.vector_id
update_sql = (
f"UPDATE `{self.collection_name}` SET {', '.join(updates)} WHERE id = :vid"
)
with self.client.engine.connect() as conn:
with conn.begin():
conn.execute(text(update_sql), params)
if not updates:
continue
logger.info(f"Updated {len(nodes_to_update)} documents in {self.collection_name}")
params["vid"] = node.vector_id
update_sql = (
f"UPDATE `{self.collection_name}` SET {', '.join(updates)} WHERE id = :vid"
)
with self.client.engine.connect() as conn:
with conn.begin():
conn.execute(text(update_sql), params)
logger.info("Updated {} documents in {}", len(nodes_to_update), self.collection_name)
except Exception as e:
logger.error(f"Failed to update documents: {e}")
logger.error("Failed to update documents: {}", e)
raise
async def get(self, vector_ids: str | list[str]) -> VectorNode | list[VectorNode] | None:
"""Retrieve vector nodes by their unique identifiers."""
single_result = isinstance(vector_ids, str)
if single_result:
vector_ids = [vector_ids]
@ -486,23 +459,13 @@ class ObVecVectorStore(BaseVectorStore):
result = self.client.perform_raw_text_sql(select_sql)
rows = result.fetchall()
results = []
for row in rows:
if row:
results.append(
VectorNode(
vector_id=row[0],
content=row[1] or "",
vector=self._coerce_db_vector(row[2]),
metadata=self._coerce_db_metadata(row[3]),
),
)
results = [_vector_node_from_db_row(row) for row in rows if row]
if single_result:
return results[0] if results else None
return results
except Exception as e:
logger.error(f"Failed to get documents: {e}")
logger.error("Failed to get documents: {}", e)
return [] if not single_result else None
async def list(
@ -512,53 +475,29 @@ class ObVecVectorStore(BaseVectorStore):
sort_key: str | None = None,
reverse: bool = False,
) -> list[VectorNode]:
"""Retrieve a list of vector nodes matching the provided filters and limit.
Args:
filters: Dictionary of filter conditions to match vectors
limit: Maximum number of vectors to return
sort_key: Key to sort the results by (e.g., field name in metadata). None for no sorting
reverse: If True, sort in descending order; if False, sort in ascending order
"""
try:
# Build select SQL
select_sql = f"SELECT id, content, vector, metadata FROM `{self.collection_name}`"
where_clause = self._build_filter_clause(filters)
where_clause = _build_metadata_filter_sql(filters)
if where_clause:
select_sql += f" WHERE {where_clause}"
# Add sorting
if sort_key:
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}"
# Add limit
if limit:
if limit is not None:
select_sql += f" LIMIT {limit}"
result = self.client.perform_raw_text_sql(select_sql)
rows = result.fetchall()
results = []
for row in rows:
if row:
results.append(
VectorNode(
vector_id=row[0],
content=row[1] or "",
vector=self._coerce_db_vector(row[2]),
metadata=self._coerce_db_metadata(row[3]),
),
)
return results
return [_vector_node_from_db_row(row) for row in rows if row]
except Exception as e:
logger.error(f"Failed to list documents: {e}")
logger.error("Failed to list documents: {}", e)
return []
async def collection_info(self) -> dict[str, Any]:
"""Fetch metadata including record count for the collection."""
try:
count_sql = f"SELECT COUNT(*) FROM `{self.collection_name}`"
result = self.client.perform_raw_text_sql(count_sql)
@ -570,27 +509,20 @@ class ObVecVectorStore(BaseVectorStore):
"count": count,
}
except Exception as e:
logger.error(f"Failed to get collection info: {e}")
logger.error("Failed to get collection info: {}", e)
return {"name": self.collection_name, "count": 0}
async def reset(self):
"""Purge all data by dropping and recreating the collection table."""
logger.warning(f"Resetting collection {self.collection_name}...")
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):
"""Reset collection with the given name."""
self.collection_name = collection_name
await self.create_collection(collection_name)
logger.info(f"Collection reset to {collection_name}")
logger.info("Collection reset to {}", collection_name)
async def start(self) -> None:
"""Initialize the OceanBase client and ensure the collection exists.
Creates the collection table if it doesn't exist.
"""
# Initialize the client
self.client = ObVecClient(
uri=self.uri,
user=self.user,
@ -599,11 +531,8 @@ class ObVecVectorStore(BaseVectorStore):
)
await super().start()
logger.info(f"OceanBase collection {self.collection_name} initialized")
logger.info("OceanBase collection {} initialized", self.collection_name)
async def close(self):
"""Terminate the OceanBase client connection and release resources."""
# Note: ObVecClient may not have a close method
# We'll just log the closure
self.client = None
logger.info("OceanBase client connection closed")
logger.info("OceanBase client connection closed")