add zvec vector/file store (#218)

This commit is contained in:
lichen2015 2026-05-08 17:12:11 +08:00 committed by GitHub
parent 72eabfa858
commit f42cf60706
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
11 changed files with 2409 additions and 8 deletions

View file

@ -506,7 +506,7 @@ async def main():
"dimensions": 1024,
},
default_vector_store_config={
"backend": "local", # Supports local/chroma/qdrant/elasticsearch/obvec
"backend": "local", # Supports local/chroma/qdrant/elasticsearch/obvec/zvec
},
)
await reme.start()

View file

@ -486,7 +486,7 @@ async def main():
"dimensions": 1024,
},
default_vector_store_config={
"backend": "local", # 支持 local/chroma/qdrant/elasticsearch/obvec
"backend": "local", # 支持 local/chroma/qdrant/elasticsearch/obvec/zvec
},
)
await reme.start()

View file

@ -34,6 +34,7 @@ FlowLLM provides multiple Vector Store implementations tailored to different use
- **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.
- **ZvecVectorStore** ([source code](https://github.com/agentscope-ai/ReMe/blob/main/reme/core/vector_store/zvec_vector_store.py)): Built on zvec, a high-performance local vector database with strong-schema support and HNSW indexing. Suitable for single-machine deployments requiring fast vector search.
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.
@ -130,6 +131,11 @@ docker run -d --name reme_seekdb -p 2881:2881 -e ROOT_PASSWORD=<your_root_passwo
```shell
OBVEC_PASSWORD=<your_root_password> python tests/test_vector_store.py --obvec
```
### ZvecVectorStore Configuration
- **db_path**: Local storage path for persistent mode (required).
- **dimension**: Dimensionality of the embedding vectors (default: `1024`).
- **distance**: Distance metric — supports `cosine`, `l2`, `ip` (default: `cosine`).
## Configuration File Examples
@ -151,7 +157,7 @@ vector_store.default.params.<param_name>=<param_value>
### Configuration Field Descriptions
- **`backend`** (required): Vector store backend type. Options: `local`, `memory`, `chroma`, `qdrant`, `elasticsearch`, `obvec`.
- **`backend`** (required): Vector store backend type. Options: `local`, `memory`, `chroma`, `qdrant`, `elasticsearch`, `obvec`, `zvec`.
- **`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.
@ -347,6 +353,30 @@ vector_stores.default.password=your-root-password
ReMe service YAML uses the key `vector_stores` (plural); CLI overrides use the same nested paths.
#### 7. ZvecVectorStore Configuration
Persistent local storage based on zvec with HNSW indexing and strong-schema support.
**Implementation**: [`reme/core/vector_store/zvec_vector_store.py`](https://github.com/agentscope-ai/ReMe/blob/main/reme/core/vector_store/zvec_vector_store.py)
```yaml
vector_store:
default:
backend: zvec
embedding_model: default
params:
db_path: "./zvec_vector_store" # Local storage path (required)
dimension: 1024 # Vector dimension (optional; default: 1024)
distance: "cosine" # Distance metric (optional; default: cosine; options: cosine, l2, ip)
```
```shell
vector_store.default.backend=zvec
vector_store.default.params.db_path=./zvec_vector_store
vector_store.default.params.dimension=1024
vector_store.default.params.distance=cosine
```
### Complete Configuration Example
Below is a complete `default.yaml` example including both embedding model and vector store configurations:
@ -405,6 +435,7 @@ 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, EsVectorStore, or ObVecVectorStore (OceanBase/seekdb) for high performance and scalability, depending on your existing infrastructure.
- **High-Performance Local Search**: Use ZvecVectorStore for single-machine deployments requiring fast HNSW-based vector search with local persistence.
- **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.

View file

@ -9,6 +9,7 @@ from .base_file_store import BaseFileStore
from .chroma_file_store import ChromaFileStore
from .local_file_store import LocalFileStore
from .sqlite_file_store import SqliteFileStore
from .zvec_file_store import ZvecFileStore
from ..registry_factory import R
__all__ = [
@ -16,8 +17,10 @@ __all__ = [
"ChromaFileStore",
"LocalFileStore",
"SqliteFileStore",
"ZvecFileStore",
]
R.file_stores.register("sqlite")(SqliteFileStore)
R.file_stores.register("chroma")(ChromaFileStore)
R.file_stores.register("local")(LocalFileStore)
R.file_stores.register("zvec")(ZvecFileStore)

View file

@ -0,0 +1,573 @@
"""Zvec storage backend for file store."""
from __future__ import annotations
import json
import time
from pathlib import Path
from typing import Any
from .base_file_store import BaseFileStore
from ..enumeration import MemorySource
from ..schema import FileMetadata, MemoryChunk, MemorySearchResult
from ..utils import get_logger
logger = get_logger()
_ZVEC_IMPORT_ERROR: Exception | None = None
try:
import zvec # type: ignore[import-untyped]
from zvec import (
CollectionOption,
CollectionSchema,
DataType,
Doc,
FieldSchema,
HnswIndexParam,
InvertIndexParam,
VectorQuery,
VectorSchema,
)
from zvec.typing import MetricType
except Exception as e:
_ZVEC_IMPORT_ERROR = e
zvec = None # type: ignore[assignment]
# zvec max topk (will be lifted to 100,000 in zvec v0.3.2+)
_ZVEC_MAX_TOPK = 1024
# Default vector field name
_DEFAULT_VECTOR_FIELD = "embedding"
def _escape(value: str) -> str:
"""Escape a string value for zvec filter expressions."""
return value.replace("'", "\\'")
def _build_file_store_schema(name: str, dimension: int) -> CollectionSchema:
"""Build a zvec CollectionSchema for file store chunks."""
return CollectionSchema(
name=name,
fields=[
FieldSchema("content", DataType.STRING, nullable=True, index_param=InvertIndexParam()),
FieldSchema("path", DataType.STRING, nullable=True, index_param=InvertIndexParam()),
FieldSchema("source", DataType.STRING, nullable=True, index_param=InvertIndexParam()),
FieldSchema("start_line", DataType.INT64, nullable=True),
FieldSchema("end_line", DataType.INT64, nullable=True),
FieldSchema("hash", DataType.STRING, nullable=True),
FieldSchema("updated_at", DataType.INT64, nullable=True),
FieldSchema("file_metadata", DataType.STRING, nullable=True),
],
vectors=[
VectorSchema(
name=_DEFAULT_VECTOR_FIELD,
data_type=DataType.VECTOR_FP32,
dimension=dimension,
index_param=HnswIndexParam(metric_type=MetricType.COSINE),
),
],
)
def _chunk_to_doc(chunk: MemoryChunk, file_meta_json: str = "{}") -> Doc:
"""Convert a MemoryChunk to a zvec Doc."""
fields: dict[str, Any] = {
"content": chunk.text,
"path": chunk.path,
"source": chunk.source.value if chunk.source else "",
"start_line": chunk.start_line,
"end_line": chunk.end_line,
"hash": chunk.hash,
"updated_at": int(time.time() * 1000),
"file_metadata": file_meta_json,
}
vectors: dict[str, Any] = {}
if chunk.embedding is not None:
vectors[_DEFAULT_VECTOR_FIELD] = chunk.embedding
return Doc(id=chunk.id, fields=fields, vectors=vectors)
def _doc_to_chunk(doc: Doc) -> MemoryChunk:
"""Convert a zvec Doc to a MemoryChunk."""
raw_vector = doc.vector(_DEFAULT_VECTOR_FIELD)
vector = raw_vector if isinstance(raw_vector, list) and len(raw_vector) > 0 else None
return MemoryChunk(
id=str(doc.id),
path=str(doc.field("path") or ""),
source=MemorySource(str(doc.field("source") or "")),
start_line=int(doc.field("start_line") or 0),
end_line=int(doc.field("end_line") or 0),
text=str(doc.field("content") or ""),
hash=str(doc.field("hash") or ""),
embedding=vector,
)
def _build_source_filter(sources: list[MemorySource] | None) -> str | None:
"""Build a zvec filter expression for source filtering."""
if not sources:
return None
if len(sources) == 1:
return f"source='{_escape(sources[0].value)}'"
vals = ", ".join(f"'{_escape(s.value)}'" for s in sources)
return f"source IN ({vals})"
class ZvecFileStore(BaseFileStore):
"""Zvec file storage with vector and keyword search.
Provides zvec-backed persistent storage with:
- Vector similarity search (native zvec HNSW)
- Keyword search (Python substring matching on fetched results)
- Hybrid search (weighted fusion of vector and keyword results)
Note:
Keyword search operates on chunks fetched from zvec, which is subject
to the topk limit (1024 in zvec < v0.3.2, 100,000 in v0.3.2+).
For collections with more chunks than the topk limit, keyword search
may not scan all documents.
"""
def __init__(
self,
store_name: str,
db_path: str | Path,
embedding_model: Any | None = None,
vector_enabled: bool = False,
fts_enabled: bool = True,
dimension: int = 1024,
**kwargs: Any,
):
if _ZVEC_IMPORT_ERROR is not None:
raise ImportError(
"Zvec requires extra dependencies. Install with `pip install zvec`",
) from _ZVEC_IMPORT_ERROR
super().__init__(
store_name=store_name,
db_path=db_path,
embedding_model=embedding_model,
vector_enabled=vector_enabled,
fts_enabled=fts_enabled,
**kwargs,
)
self.dimension = dimension
self._collection = None
self._initialized = False
self._metadata_file: Path = self.db_path / f"{store_name}_file_metadata.json"
self._metadata_cache: dict[str, dict[str, FileMetadata]] = {}
@property
def collection_name(self) -> str:
"""Get the name of the zvec collection for this store."""
return f"chunks_{self.store_name}"
# ------------------------------------------------------------------
# Lifecycle
# ------------------------------------------------------------------
async def start(self) -> None:
"""Initialize zvec engine and open the collection."""
if not self._initialized:
try:
zvec.init()
except RuntimeError:
pass
self._initialized = True
self.db_path.mkdir(parents=True, exist_ok=True)
collection_path = str(self.db_path / self.collection_name)
option = CollectionOption(read_only=False, enable_mmap=True)
try:
self._collection = zvec.open(collection_path, option)
logger.info(f"Opened existing zvec file store collection: {collection_path}")
except Exception:
schema = _build_file_store_schema(self.collection_name, self.dimension)
self._collection = zvec.create_and_open(
path=collection_path,
schema=schema,
option=option,
)
logger.info(f"Created new zvec file store collection: {collection_path}")
self._metadata_cache = await self._load_metadata()
async def close(self) -> None:
"""Close zvec collection and persist metadata."""
if self._metadata_cache:
await self._save_metadata(self._metadata_cache)
if self._collection is not None:
try:
self._collection.flush()
except Exception as e:
logger.warning(f"Failed to flush collection on close: {e}")
self._collection = None
# ------------------------------------------------------------------
# Metadata management
# ------------------------------------------------------------------
async def _load_metadata(self) -> dict[str, dict[str, FileMetadata]]:
"""Load file metadata from JSON file."""
if not self._metadata_file.exists():
return {}
try:
data = json.loads(self._metadata_file.read_text(encoding="utf-8"))
result: dict[str, dict[str, FileMetadata]] = {}
for source, files in data.items():
result[source] = {}
for path, meta in files.items():
result[source][path] = FileMetadata(**meta)
return result
except Exception as e:
logger.warning(f"Failed to load metadata from {self._metadata_file}: {e}")
return {}
async def _save_metadata(self, metadata: dict[str, dict[str, FileMetadata]]) -> None:
"""Save file metadata to JSON file."""
try:
out: dict[str, dict[str, dict]] = {}
for source, files in metadata.items():
out[source] = {}
for path, meta in files.items():
out[source][path] = {
"path": meta.path,
"hash": meta.hash,
"mtime_ms": meta.mtime_ms,
"size": meta.size,
"chunk_count": meta.chunk_count,
}
self._metadata_file.write_text(
json.dumps(out, indent=2, ensure_ascii=False),
encoding="utf-8",
)
except Exception as e:
logger.error(f"Failed to save metadata to {self._metadata_file}: {e}")
# ------------------------------------------------------------------
# CRUD operations
# ------------------------------------------------------------------
async def upsert_file(
self,
file_meta: FileMetadata,
source: MemorySource,
chunks: list[MemoryChunk],
) -> None:
"""Insert or update a file and its chunks."""
if not chunks:
return
# Delete existing chunks for this file first
await self.delete_file(file_meta.path, source)
# Generate embeddings
chunks = await self.get_chunk_embeddings(chunks)
file_meta_json = json.dumps(
{
"path": file_meta.path,
"hash": file_meta.hash,
"mtime_ms": file_meta.mtime_ms,
"size": file_meta.size,
"chunk_count": len(chunks),
},
ensure_ascii=False,
)
docs = [_chunk_to_doc(c, file_meta_json) for c in chunks]
self._collection.insert(docs)
# Update metadata cache
if source.value not in self._metadata_cache:
self._metadata_cache[source.value] = {}
self._metadata_cache[source.value][file_meta.path] = FileMetadata(
hash=file_meta.hash,
mtime_ms=file_meta.mtime_ms,
size=file_meta.size,
path=file_meta.path,
chunk_count=len(chunks),
)
async def delete_file(self, path: str, source: MemorySource) -> None:
"""Delete a file and all its chunks."""
filter_expr = f"path='{_escape(path)}' AND source='{_escape(source.value)}'"
results = self._collection.query(topk=_ZVEC_MAX_TOPK, filter=filter_expr, include_vector=False)
ids_to_delete = [doc.id for doc in results]
if ids_to_delete:
self._collection.delete(ids_to_delete)
if source.value in self._metadata_cache:
self._metadata_cache[source.value].pop(path, None)
async def delete_file_chunks(self, path: str, chunk_ids: list[str]) -> None:
"""Delete specific chunks for a file."""
if not chunk_ids:
return
self._collection.delete(chunk_ids)
async def upsert_chunks(
self,
chunks: list[MemoryChunk],
source: MemorySource,
) -> None:
"""Insert or update specific chunks."""
if not chunks:
return
chunks = await self.get_chunk_embeddings(chunks)
docs = [_chunk_to_doc(c) for c in chunks]
self._collection.upsert(docs)
# ------------------------------------------------------------------
# Listing and metadata
# ------------------------------------------------------------------
async def list_files(self, source: MemorySource) -> list[str]:
"""List all indexed files for a source."""
if source.value not in self._metadata_cache:
return []
return list(self._metadata_cache[source.value].keys())
async def get_file_metadata(
self,
path: str,
source: MemorySource,
) -> FileMetadata | None:
"""Get file metadata."""
if source.value not in self._metadata_cache:
return None
return self._metadata_cache[source.value].get(path)
async def update_file_metadata(self, file_meta: FileMetadata, source: MemorySource) -> None:
"""Update file metadata without affecting chunks."""
if source.value not in self._metadata_cache:
self._metadata_cache[source.value] = {}
self._metadata_cache[source.value][file_meta.path] = FileMetadata(
hash=file_meta.hash,
mtime_ms=file_meta.mtime_ms,
size=file_meta.size,
path=file_meta.path,
chunk_count=file_meta.chunk_count,
)
async def get_file_chunks(
self,
path: str,
source: MemorySource,
) -> list[MemoryChunk]:
"""Get all chunks for a file."""
filter_expr = f"path='{_escape(path)}' AND source='{_escape(source.value)}'"
results = self._collection.query(
topk=_ZVEC_MAX_TOPK,
filter=filter_expr,
include_vector=True,
)
chunks = [_doc_to_chunk(doc) for doc in results]
chunks.sort(key=lambda c: c.start_line)
return chunks
# ------------------------------------------------------------------
# Search
# ------------------------------------------------------------------
async def vector_search(
self,
query: str,
limit: int,
sources: list[MemorySource] | None = None,
) -> list[MemorySearchResult]:
"""Perform vector similarity search."""
if not self.vector_enabled or not query:
return []
query_embedding = await self.get_embedding(query)
if not query_embedding:
return []
filter_expr = _build_source_filter(sources)
vq = VectorQuery(field_name=_DEFAULT_VECTOR_FIELD, vector=query_embedding)
try:
results = self._collection.query(
vectors=vq,
topk=min(limit, _ZVEC_MAX_TOPK),
filter=filter_expr,
include_vector=False,
)
except Exception as e:
logger.error(f"Vector search failed: {e}")
return []
search_results = []
for doc in results:
score = doc.score if doc.score is not None else 0.0
# zvec cosine score might need normalization depending on version
search_results.append(
MemorySearchResult(
path=str(doc.field("path") or ""),
start_line=int(doc.field("start_line") or 0),
end_line=int(doc.field("end_line") or 0),
score=score,
snippet=str(doc.field("content") or ""),
source=MemorySource(str(doc.field("source") or "")),
raw_metric=score,
),
)
search_results.sort(key=lambda r: r.score, reverse=True)
return search_results[:limit]
async def keyword_search(
self,
query: str,
limit: int,
sources: list[MemorySource] | None = None,
) -> list[MemorySearchResult]:
"""Perform keyword search via Python substring matching.
Fetches chunks from zvec (subject to topk limit) then matches
keywords in Python. For collections larger than the topk limit,
not all documents are scanned.
"""
if not self.fts_enabled or not query:
return []
words = query.split()
if not words:
return []
# Fetch candidate chunks from zvec
filter_expr = _build_source_filter(sources)
results = self._collection.query(
topk=_ZVEC_MAX_TOPK,
filter=filter_expr,
include_vector=False,
)
query_lower = query.lower()
words_lower = [w.lower() for w in words]
n_words = len(words)
search_results = []
for doc in results:
text = str(doc.field("content") or "")
text_lower = text.lower()
match_count = sum(1 for w in words_lower if w in text_lower)
if match_count == 0:
continue
base_score = match_count / n_words
phrase_bonus = 0.2 if n_words > 1 and query_lower in text_lower else 0.0
score = min(1.0, base_score + phrase_bonus)
search_results.append(
MemorySearchResult(
path=str(doc.field("path") or ""),
start_line=int(doc.field("start_line") or 0),
end_line=int(doc.field("end_line") or 0),
score=score,
snippet=text,
source=MemorySource(str(doc.field("source") or "")),
),
)
search_results.sort(key=lambda r: r.score, reverse=True)
return search_results[:limit]
async def hybrid_search(
self,
query: str,
limit: int,
sources: list[MemorySource] | None = None,
vector_weight: float = 0.7,
candidate_multiplier: float = 3.0,
) -> list[MemorySearchResult]:
"""Perform hybrid search combining vector and keyword search."""
assert 0.0 <= vector_weight <= 1.0, f"vector_weight must be between 0 and 1, got {vector_weight}"
candidates = min(200, max(1, int(limit * candidate_multiplier)))
text_weight = 1.0 - vector_weight
if self.vector_enabled and self.fts_enabled:
keyword_results = await self.keyword_search(query, candidates, sources)
vector_results = await self.vector_search(query, candidates, sources)
if not keyword_results:
return vector_results[:limit]
elif not vector_results:
return keyword_results[:limit]
else:
return self._merge_hybrid_results(
vector=vector_results,
keyword=keyword_results,
vector_weight=vector_weight,
text_weight=text_weight,
)[:limit]
elif self.vector_enabled:
return await self.vector_search(query, limit, sources)
elif self.fts_enabled:
return await self.keyword_search(query, limit, sources)
else:
return []
@staticmethod
def _merge_hybrid_results(
vector: list[MemorySearchResult],
keyword: list[MemorySearchResult],
vector_weight: float,
text_weight: float,
) -> list[MemorySearchResult]:
"""Merge vector and keyword search results with weighted scoring."""
merged: dict[str, MemorySearchResult] = {}
for result in vector:
result.score = result.score * vector_weight
merged[result.merge_key] = result
for result in keyword:
key = result.merge_key
if key in merged:
merged[key].score += result.score * text_weight
else:
result.score = result.score * text_weight
merged[key] = result
results = list(merged.values())
results.sort(key=lambda r: r.score, reverse=True)
return results
# ------------------------------------------------------------------
# Maintenance
# ------------------------------------------------------------------
async def clear_all(self) -> None:
"""Clear all indexed data."""
# Delete all documents
stats = self._collection.stats
count = stats.doc_count if stats else 0
if count > 0:
try:
self._collection.delete_by_filter("content!=''")
except Exception:
remaining = count
while remaining > 0:
batch = self._collection.query(
topk=min(remaining, _ZVEC_MAX_TOPK),
include_vector=False,
)
if not batch:
break
self._collection.delete([doc.id for doc in batch])
remaining -= len(batch)
self._metadata_cache = {}
await self._save_metadata({})
logger.info(f"Cleared all data from zvec file store: {self.collection_name}")

View file

@ -7,6 +7,7 @@ from .local_vector_store import LocalVectorStore
from .obvec_vector_store import ObVecVectorStore
from .pgvector_store import PGVectorStore
from .qdrant_vector_store import QdrantVectorStore
from .zvec_vector_store import ZvecVectorStore
from ..registry_factory import R
__all__ = [
@ -17,6 +18,7 @@ __all__ = [
"ObVecVectorStore",
"PGVectorStore",
"QdrantVectorStore",
"ZvecVectorStore",
]
R.vector_stores.register("chroma")(ChromaVectorStore)
@ -25,3 +27,4 @@ R.vector_stores.register("local")(LocalVectorStore)
R.vector_stores.register("obvec")(ObVecVectorStore)
R.vector_stores.register("pgvector")(PGVectorStore)
R.vector_stores.register("qdrant")(QdrantVectorStore)
R.vector_stores.register("zvec")(ZvecVectorStore)

View file

@ -0,0 +1,809 @@
"""Zvec vector store implementation for the ReMe framework."""
from __future__ import annotations
import json
from pathlib import Path
from typing import Any
from loguru import logger
from .base_vector_store import BaseVectorStore
from ..embedding import BaseEmbeddingModel
from ..schema import VectorNode
_ZVEC_IMPORT_ERROR: Exception | None = None
try:
import zvec # type: ignore[import-untyped]
from zvec import (
CollectionOption,
CollectionSchema,
DataType,
Doc,
FieldSchema,
HnswIndexParam,
InvertIndexParam,
VectorQuery,
VectorSchema,
)
from zvec.typing import MetricType
except Exception as e:
_ZVEC_IMPORT_ERROR = e
zvec = None # type: ignore[assignment]
# Default vector field name used inside zvec collections
_DEFAULT_VECTOR_FIELD = "embedding"
# Default scalar content field name for storing text
_CONTENT_FIELD = "content"
# Field name for JSON-serialized metadata
_METADATA_FIELD = "metadata"
# Metadata fields promoted to top-level zvec schema columns for native filtering.
# These are the most commonly filtered keys in ReMe's memory system.
# Defining them as independent schema columns allows zvec to perform
# filtering at the database level instead of Python post-filtering.
# Format: {metadata_key: (zvec_data_type_str, has_inverted_index)}
_PROMOTED_FIELD_SPECS: dict[str, tuple[str, bool]] = {
"memory_type": ("STRING", True), # Inverted index for exact match filtering
"memory_target": ("STRING", True), # Inverted index for exact match filtering
"author": ("STRING", False),
"time_int": ("INT64", False), # Numeric for range queries
}
# zvec data-type string → DataType enum mapping (populated after import)
_DATATYPE_MAP: dict[str, Any] = {} # filled in _build_collection_schema
def _escape_zvec_string(value: str) -> str:
"""Escape a string value for use in zvec filter expressions."""
return value.replace("'", "\\'")
def _build_zvec_filter(
filters: dict | None,
promoted_fields: set[str],
) -> tuple[str | None, dict | None]:
"""Split ReMe filter dict into a zvec native filter expression and remaining post-filters.
For filter keys that correspond to promoted schema fields, native
zvec filter expressions are generated. Non-promoted keys are
kept for Python post-filtering.
Args:
filters: ReMe-style filter dictionary.
promoted_fields: Set of metadata keys that exist as top-level schema columns.
Returns:
(native_filter_expr, post_filter_dict) either may be None.
"""
if not filters:
return None, None
native_conditions: list[str] = []
post_filters: dict = {}
for key, value in filters.items():
if key.startswith("$"):
# Compound operators ($or, $and, $not) — keep for post-filtering
post_filters[key] = value
continue
if key not in promoted_fields:
# Not a promoted field — use post-filtering
post_filters[key] = value
continue
# Build native filter condition for promoted fields
field_type = _PROMOTED_FIELD_SPECS.get(key, ("STRING", False))[0]
if isinstance(value, list) and len(value) == 2:
# Range query: [start, end]
if field_type == "INT64":
native_conditions.append(f"{key} >= {value[0]} AND {key} <= {value[1]}")
else:
# STRING range — use >= and <= with string escaping
native_conditions.append(
f"{key} >= '{_escape_zvec_string(str(value[0]))}' "
f"AND {key} <= '{_escape_zvec_string(str(value[1]))}'",
)
elif isinstance(value, bool):
native_conditions.append(f"{key} = {str(value).upper()}")
elif isinstance(value, (int, float)):
native_conditions.append(f"{key} = {value}")
elif isinstance(value, str):
native_conditions.append(f"{key} = '{_escape_zvec_string(value)}'")
else:
# Unsupported type — fall back to post-filtering
post_filters[key] = value
native_filter = " AND ".join(native_conditions) if native_conditions else None
return native_filter, post_filters if post_filters else None
def _metric_type_from_str(metric: str) -> Any:
"""Convert a string metric name to zvec MetricType enum value."""
if zvec is None:
return None
mapping = {
"cosine": MetricType.COSINE,
"l2": MetricType.L2,
"ip": MetricType.IP,
}
return mapping.get(metric.lower(), MetricType.COSINE)
def _build_collection_schema(
name: str,
dimension: int,
metric: str = "cosine",
) -> CollectionSchema:
"""Build a zvec CollectionSchema for ReMe usage.
The schema contains:
- "content" (STRING, inverted index) text content
- "metadata" (STRING) JSON-serialized metadata dictionary
- Promoted metadata fields (STRING / INT64) for native zvec filtering
- "embedding" (VECTOR_FP32, dimension, HNSW index) the vector field
Promoted fields are commonly filtered metadata keys defined as top-level
schema columns so that zvec can perform filtering natively instead of
Python post-filtering. The full metadata is still stored as JSON in the
"metadata" field for complete round-trip serialization.
zvec automatically manages the document ID (string type); we do NOT
define an "id" field in the schema.
"""
# Populate the DataType map on first call
if not _DATATYPE_MAP:
_DATATYPE_MAP.update(
{
"STRING": DataType.STRING,
"INT64": DataType.INT64,
},
)
distance = _metric_type_from_str(metric)
# Base fields
fields = [
FieldSchema("content", DataType.STRING, nullable=True, index_param=InvertIndexParam()),
FieldSchema("metadata", DataType.STRING, nullable=True),
]
# Add promoted metadata fields as top-level schema columns
for field_name, (type_str, has_inv_index) in _PROMOTED_FIELD_SPECS.items():
dt = _DATATYPE_MAP[type_str]
idx_param = InvertIndexParam() if has_inv_index else None
fields.append(FieldSchema(field_name, dt, nullable=True, index_param=idx_param))
return CollectionSchema(
name=name,
fields=fields,
vectors=[
VectorSchema(
name=_DEFAULT_VECTOR_FIELD,
data_type=DataType.VECTOR_FP32,
dimension=dimension,
index_param=HnswIndexParam(metric_type=distance),
),
],
)
def _vector_node_to_doc(node: VectorNode) -> Doc:
"""Convert a ReMe VectorNode to a zvec Doc.
Metadata is serialized as a JSON string into the "metadata" field.
The "score" key is excluded since it is a computed value, not stored data.
Promoted metadata fields are also extracted as top-level Doc fields
for native zvec filtering.
The vector is placed under the default vector field name.
The zvec Doc id must be a string.
"""
# Filter out computed score before serialization
meta_to_store = {k: v for k, v in node.metadata.items() if k != "score"}
fields: dict[str, Any] = {
"content": node.content,
"metadata": json.dumps(meta_to_store) if meta_to_store else "{}",
}
# Extract promoted metadata fields as top-level schema columns
for field_name, (type_str, _) in _PROMOTED_FIELD_SPECS.items():
value = meta_to_store.get(field_name)
if value is not None:
# Ensure correct type: INT64 fields must be int
if type_str == "INT64" and not isinstance(value, int):
try:
value = int(value)
except (ValueError, TypeError):
continue
fields[field_name] = value
vectors: dict[str, Any] = {}
if node.vector is not None:
vectors[_DEFAULT_VECTOR_FIELD] = node.vector
return Doc(id=str(node.vector_id), fields=fields, vectors=vectors)
def _doc_to_vector_node(doc: Doc, include_score: bool = False) -> VectorNode:
"""Convert a zvec Doc back to a ReMe VectorNode.
The "metadata" field is parsed from JSON. The "content" field becomes
the node content. If ``include_score`` is True, the search score is
added to the metadata dictionary.
"""
metadata: dict[str, str | bool | int | float] = {}
# Parse JSON metadata
raw_metadata = doc.field("metadata")
if raw_metadata:
try:
parsed = json.loads(raw_metadata)
if isinstance(parsed, dict):
metadata.update(parsed)
except (json.JSONDecodeError, TypeError):
logger.warning(f"Failed to parse metadata JSON: {raw_metadata}")
if include_score and doc.score is not None:
metadata["score"] = doc.score
# Extract vector — doc.vector() returns list or empty dict
raw_vector = doc.vector(_DEFAULT_VECTOR_FIELD)
vector = raw_vector if isinstance(raw_vector, list) and len(raw_vector) > 0 else None
content = doc.field("content") or ""
return VectorNode(
vector_id=str(doc.id),
content=str(content),
vector=vector,
metadata=metadata,
)
def _apply_filters_post(nodes: list[VectorNode], filters: dict | None) -> list[VectorNode]:
"""Apply ReMe-style filter dict as post-filtering on metadata.
Used as a fallback for metadata keys that are NOT promoted to top-level
schema columns (and thus cannot be filtered natively by zvec). Promoted
fields are handled by zvec's native ``filter`` parameter instead.
Supports:
- Exact match: {"field": value}
- Range query: {"field": [start, end]}
"""
if not filters:
return nodes
filtered = []
for node in nodes:
match = True
for key, value in filters.items():
if key.startswith("$"):
# Skip compound operators for post-filtering
continue
node_value = node.metadata.get(key)
# Range query: [start, end]
if isinstance(value, list) and len(value) == 2:
if node_value is None:
match = False
break
try:
if not value[0] <= node_value <= value[1]:
match = False
break
except TypeError:
match = False
break
else:
# Exact match
if node_value != value:
match = False
break
if match:
filtered.append(node)
return filtered
class ZvecVectorStore(BaseVectorStore):
"""Zvec-based vector store implementation.
Zvec is a high-performance vector database. This adapter bridges the
ReMe ``BaseVectorStore`` interface with zvec's Python API.
Supports local persistent storage via ``db_path``.
Args:
collection_name: Name of the vector collection.
db_path: Local storage path for persistent mode.
embedding_model: Model used for generating vector embeddings.
dimension: Dimensionality of the embedding vectors (default: 1024).
distance: Distance metric cosine / l2 / ip (default: cosine).
**kwargs: Additional zvec-specific configuration.
"""
def __init__(
self,
collection_name: str,
db_path: str | Path,
embedding_model: BaseEmbeddingModel,
dimension: int = 1024,
distance: str = "cosine",
**kwargs: Any,
):
"""Initialize the Zvec vector store."""
if _ZVEC_IMPORT_ERROR is not None:
raise ImportError(
"Zvec requires extra dependencies. Install with `pip install zvec`",
) from _ZVEC_IMPORT_ERROR
super().__init__(
collection_name=collection_name,
db_path=db_path,
embedding_model=embedding_model,
**kwargs,
)
self.dimension = dimension
self.distance = distance
self._collection = None
self._initialized = False
# Set of promoted field names that exist in the current collection's schema.
# Populated during start() by inspecting the schema. Only fields present
# in the schema can use native zvec filtering; the rest fall back to
# Python post-filtering.
self._promoted_fields_in_schema: set[str] = set()
# ------------------------------------------------------------------
# Lifecycle
# ------------------------------------------------------------------
async def start(self) -> None:
"""Initialize the Zvec engine and open the collection.
Calls ``zvec.init()`` once, then tries to ``zvec.open()`` an existing
collection or ``zvec.create_and_open()`` a new one.
After opening, detects which promoted fields exist in the schema
and attempts to add missing numeric fields via ``add_column``.
"""
if not self._initialized:
try:
zvec.init()
except RuntimeError:
# Already initialized — safe to ignore
pass
self._initialized = True
self.db_path.mkdir(parents=True, exist_ok=True)
collection_path = str(self.db_path / self.collection_name)
option = CollectionOption(read_only=False, enable_mmap=True)
try:
# Try opening an existing collection first
self._collection = zvec.open(collection_path, option)
logger.info(f"Opened existing Zvec collection at {collection_path}")
except Exception:
# Collection doesn't exist — create it
schema = _build_collection_schema(
name=self.collection_name,
dimension=self.dimension,
metric=self.distance,
)
self._collection = zvec.create_and_open(
path=collection_path,
schema=schema,
option=option,
)
logger.info(f"Created new Zvec collection at {collection_path}")
# Detect which promoted fields exist in the current schema
self._detect_promoted_fields()
# Try to add missing numeric promoted fields to existing collections
# (zvec's add_column only supports numeric types: INT64, FLOAT, etc.)
self._ensure_numeric_promoted_columns()
async def close(self) -> None:
"""Flush pending writes and release the collection handle."""
if self._collection is not None:
try:
self._collection.flush()
except Exception as e:
logger.warning(f"Failed to flush collection on close: {e}")
self._collection = None
logger.info(f"Zvec vector store for collection {self.collection_name} closed")
# ------------------------------------------------------------------
# Collection management
# ------------------------------------------------------------------
async def list_collections(self) -> list[str]:
"""Retrieve a list of collection names in the db_path directory.
Zvec doesn't have a global ``list_collections`` API; we scan the
db_path directory for zvec collection folders.
"""
if not self.db_path.exists():
return []
collections = []
for child in self.db_path.iterdir():
if child.is_dir():
collections.append(child.name)
return collections
async def create_collection(self, collection_name: str, **kwargs) -> None:
"""Create a new collection with the specified name and distance metric."""
if not self._initialized:
try:
zvec.init()
except RuntimeError:
pass
self._initialized = True
self.db_path.mkdir(parents=True, exist_ok=True)
collection_path = str(self.db_path / collection_name)
dimension = kwargs.get("dimension", self.dimension)
metric = kwargs.get("distance_metric", self.distance)
schema = _build_collection_schema(
name=collection_name,
dimension=dimension,
metric=metric,
)
option = CollectionOption(read_only=False, enable_mmap=True)
collection = zvec.create_and_open(path=collection_path, schema=schema, option=option)
if collection_name == self.collection_name:
self._collection = collection
logger.info(f"Created collection `{collection_name}`")
async def delete_collection(self, collection_name: str, **kwargs) -> None:
"""Permanently remove a collection from disk."""
# If it's the active collection, destroy it via zvec API
if self._collection is not None and collection_name == self.collection_name:
try:
self._collection.destroy()
self._collection = None
deleted = True
except Exception as _e:
logger.warning(f"Failed to destroy collection {collection_name}: {_e}")
deleted = False
else:
# For non-active collections, remove the directory
collection_path = self.db_path / collection_name
if collection_path.exists():
import shutil
shutil.rmtree(collection_path, ignore_errors=True)
deleted = True
else:
deleted = False
logger.info(f"Deleted collection {collection_name}: {deleted}")
async def copy_collection(self, collection_name: str, **kwargs) -> None:
"""Duplicate the current collection to a new one with the given name.
Uses ``shutil.copytree`` to directly copy the collection directory on
disk, which is both faster and complete it avoids the topk limit of
``list()`` (max 1024 docs) that would cause data loss for large
collections.
The source collection is flushed before copying to ensure all
pending writes are persisted to disk.
"""
import shutil
# Flush source collection so all data is on disk
if self._collection is not None:
self._collection.flush()
src_path = self.db_path / self.collection_name
dst_path = self.db_path / collection_name
if not src_path.exists():
logger.warning(f"Source collection directory not found: {src_path}")
return
if dst_path.exists():
logger.warning(f"Target collection already exists: {dst_path}, removing it first")
shutil.rmtree(dst_path, ignore_errors=True)
shutil.copytree(src_path, dst_path)
logger.info(
f"Copied collection {self.collection_name} to {collection_name} "
f"(directory copy: {src_path} -> {dst_path})",
)
# ------------------------------------------------------------------
# CRUD operations
# ------------------------------------------------------------------
async def insert(self, nodes: VectorNode | list[VectorNode], **kwargs) -> None:
"""Add one or more vector nodes into the current collection.
Automatically generates embeddings for nodes that lack vectors.
"""
if isinstance(nodes, VectorNode):
nodes = [nodes]
if not nodes:
return
# Batch generate embeddings for nodes that need them
nodes_without_vectors = [n for n in nodes if n.vector is None]
if nodes_without_vectors:
nodes_with_vectors = await self.get_node_embeddings(nodes_without_vectors)
vector_map = {n.vector_id: n for n in nodes_with_vectors}
nodes_to_insert = [vector_map.get(n.vector_id, n) if n.vector is None else n for n in nodes]
else:
nodes_to_insert = nodes
batch_size = kwargs.get("batch_size", 100)
for i in range(0, len(nodes_to_insert), batch_size):
batch = nodes_to_insert[i : i + batch_size]
docs = [_vector_node_to_doc(n) for n in batch]
self._collection.insert(docs)
logger.info(f"Inserted {len(nodes_to_insert)} nodes into {self.collection_name}")
async def search(
self,
query: str,
limit: int = 5,
filters: dict | None = None,
**kwargs,
) -> list[VectorNode]:
"""Find the most similar vector nodes based on a text query.
Uses zvec's ``query()`` method with a ``VectorQuery`` built from the
embedding of the query text. Promoted metadata fields are filtered
natively via zvec's ``filter`` parameter; remaining filters are
applied as post-filtering in Python.
"""
query_vector = await self.get_embedding(query)
vq = VectorQuery(
field_name=_DEFAULT_VECTOR_FIELD,
vector=query_vector,
)
include_vector = kwargs.get("include_embeddings", False)
# Split filters: native zvec filter vs Python post-filter
native_filter, post_filters = _build_zvec_filter(filters, self._promoted_fields_in_schema)
# Over-fetch to compensate for post-filtering
_ZVEC_MAX_TOPK = 1024
# When post-filters remain, we need to fetch more results because
# many may be filtered out. Use the maximum allowed to minimize misses.
fetch_limit = _ZVEC_MAX_TOPK if post_filters else min(limit, _ZVEC_MAX_TOPK)
results = self._collection.query(
vectors=vq,
topk=fetch_limit,
filter=native_filter,
include_vector=include_vector,
)
nodes = [_doc_to_vector_node(doc, include_score=True) for doc in results]
# Post-filter on non-promoted metadata fields
nodes = _apply_filters_post(nodes, post_filters)
score_threshold = kwargs.get("score_threshold")
if score_threshold is not None:
nodes = [n for n in nodes if n.metadata.get("score", 0) >= score_threshold]
return nodes[:limit]
async def delete(self, vector_ids: str | list[str], **kwargs) -> None:
"""Remove specific vectors from the collection using their identifiers."""
if isinstance(vector_ids, str):
vector_ids = [vector_ids]
if not vector_ids:
return
self._collection.delete(vector_ids)
logger.info(f"Deleted {len(vector_ids)} nodes from {self.collection_name}")
async def delete_all(self, **kwargs) -> None:
"""Remove all vectors from the collection.
Uses zvec's ``delete_by_filter`` with a condition that matches all
documents (content is not empty), or falls back to query + delete
in batches (zvec topk max is 1024).
"""
stats = self._collection.stats
count = stats.doc_count if stats else 0
if count > 0:
try:
# Use delete_by_filter for efficiency
self._collection.delete_by_filter("content!=''")
except Exception:
# Fallback: fetch all IDs in batches then delete
_ZVEC_MAX_TOPK = 1024
remaining = count
while remaining > 0:
all_docs = self._collection.query(topk=min(remaining, _ZVEC_MAX_TOPK), include_vector=False)
if not all_docs:
break
ids = [doc.id for doc in all_docs]
self._collection.delete(ids)
remaining -= len(ids)
logger.info(f"Deleted all {count} nodes from {self.collection_name}")
async def update(self, nodes: VectorNode | list[VectorNode], **kwargs) -> None:
"""Update existing vectors using zvec's ``upsert``.
Automatically regenerates embeddings for nodes whose content changed
but lack an updated vector.
"""
if isinstance(nodes, VectorNode):
nodes = [nodes]
if not nodes:
return
# Batch generate embeddings for nodes that need them
nodes_without_vectors = [n for n in nodes if n.vector is None and n.content]
if nodes_without_vectors:
nodes_with_vectors = await self.get_node_embeddings(nodes_without_vectors)
vector_map = {n.vector_id: n for n in nodes_with_vectors}
nodes_to_update = [vector_map.get(n.vector_id, n) if n.vector is None and n.content else n for n in nodes]
else:
nodes_to_update = nodes
docs = [_vector_node_to_doc(n) for n in nodes_to_update]
self._collection.upsert(docs)
logger.info(f"Updated {len(nodes_to_update)} nodes in {self.collection_name}")
async def get(self, vector_ids: str | list[str]) -> VectorNode | list[VectorNode]:
"""Fetch specific vector nodes from the collection by their IDs."""
is_single = isinstance(vector_ids, str)
ids = [vector_ids] if is_single else vector_ids
result_dict = self._collection.fetch(ids)
nodes = [_doc_to_vector_node(doc) for doc in result_dict.values()]
return nodes[0] if is_single and nodes else (nodes if not is_single else None)
async def list(
self,
filters: dict | None = None,
limit: int | None = None,
sort_key: str | None = None,
reverse: bool = True,
) -> list[VectorNode]:
"""Retrieve vectors matching optional metadata filters.
Uses zvec's ``query()`` without a vector query to list all documents.
Promoted metadata fields are filtered natively via zvec's ``filter``
parameter; remaining filters are applied as post-filtering in Python.
Args:
filters: Dictionary of filter conditions to match vectors.
limit: Maximum number of vectors to return.
sort_key: Key to sort the results by (in metadata).
reverse: If True, sort in descending order; otherwise ascending.
"""
# Split filters: native zvec filter vs Python post-filter
native_filter, post_filters = _build_zvec_filter(filters, self._promoted_fields_in_schema)
# Determine fetch limit — zvec max topk is 1024 (will be lifted to 100,000 in zvec v0.3.2+)
_ZVEC_MAX_TOPK = 1024
fetch_limit = min(limit or _ZVEC_MAX_TOPK, _ZVEC_MAX_TOPK)
if sort_key or post_filters:
fetch_limit = _ZVEC_MAX_TOPK # fetch max and sort/filter in Python
results = self._collection.query(
topk=fetch_limit,
filter=native_filter,
include_vector=True,
)
nodes = [_doc_to_vector_node(doc) for doc in results]
# Post-filter on non-promoted metadata fields
nodes = _apply_filters_post(nodes, post_filters)
# Apply sorting if sort_key is provided
if sort_key:
def _sort_key_func(node: VectorNode):
value = node.metadata.get(sort_key)
if value is None:
return float("-inf") if not reverse else float("inf")
return value
nodes.sort(key=_sort_key_func, reverse=reverse)
if limit is not None:
nodes = nodes[:limit]
return nodes
# ------------------------------------------------------------------
# Helpers
# ------------------------------------------------------------------
def _detect_promoted_fields(self) -> None:
"""Detect which promoted fields exist in the current collection's schema.
Compares the set of promoted field names against the actual schema
and populates ``_promoted_fields_in_schema`` accordingly. Only fields
present in the schema can use native zvec filtering.
"""
if self._collection is None:
return
try:
schema = self._collection.schema
existing_fields = {f.name for f in schema.fields} if schema.fields else set()
except Exception as e:
logger.warning(f"Failed to read collection schema: {e}")
existing_fields = set()
self._promoted_fields_in_schema = set(_PROMOTED_FIELD_SPECS.keys()) & existing_fields
missing = set(_PROMOTED_FIELD_SPECS.keys()) - existing_fields
if missing:
logger.info(
f"Promoted fields not in schema (will use post-filtering): {missing}",
)
def _ensure_numeric_promoted_columns(self) -> None:
"""Add missing numeric promoted fields to existing collections.
zvec's ``add_column`` only supports numeric types (INT64, FLOAT, etc.).
STRING fields cannot be added via ``add_column`` and must be defined
at collection creation time. For those, we fall back to post-filtering.
"""
if self._collection is None:
return
missing = set(_PROMOTED_FIELD_SPECS.keys()) - self._promoted_fields_in_schema
if not missing:
return
# Populate the DataType map if needed
if not _DATATYPE_MAP:
_DATATYPE_MAP.update(
{
"STRING": DataType.STRING,
"INT64": DataType.INT64,
},
)
for field_name in missing:
type_str, _ = _PROMOTED_FIELD_SPECS[field_name]
# Only numeric types can be added via add_column
if type_str not in ("INT64", "INT32", "FLOAT", "DOUBLE"):
continue
try:
dt = _DATATYPE_MAP[type_str]
self._collection.add_column(FieldSchema(field_name, dt, nullable=True))
self._promoted_fields_in_schema.add(field_name)
logger.info(f"Added promoted column '{field_name}' to existing collection")
except Exception as e:
logger.warning(f"Failed to add column '{field_name}': {e}")
async def count(self) -> int:
"""Return the total number of documents in the current collection."""
stats = self._collection.stats
return stats.doc_count if stats else 0
async def reset(self):
"""Reset the current collection by destroying and recreating it."""
logger.warning(f"Resetting collection {self.collection_name}...")
await self.delete_collection(self.collection_name)
await self.create_collection(self.collection_name)
logger.info(f"Collection {self.collection_name} has been reset")

View file

@ -28,6 +28,7 @@ from reme.core.file_store.base_file_store import BaseFileStore
from reme.core.file_store.chroma_file_store import ChromaFileStore
from reme.core.file_store.local_file_store import LocalFileStore
from reme.core.file_store.sqlite_file_store import SqliteFileStore
from reme.core.file_store.zvec_file_store import ZvecFileStore
from reme.core.schema.file_metadata import FileMetadata
from reme.core.schema.memory_chunk import MemoryChunk
from reme.core.utils import load_env
@ -53,6 +54,10 @@ class TestConfig:
CHROMA_DB_PATH = "./test_file_store_chroma"
CHROMA_FTS_ENABLED = True
# ZvecFileStore settings
ZVEC_DB_PATH = "./test_file_store_zvec"
ZVEC_FTS_ENABLED = True
# LocalFileStore settings
LOCAL_DB_PATH = "./test_file_store_local"
LOCAL_FTS_ENABLED = True
@ -199,6 +204,8 @@ def get_store_type(store: BaseFileStore) -> str:
return "chroma"
elif isinstance(store, LocalFileStore):
return "local"
elif isinstance(store, ZvecFileStore):
return "zvec"
else:
raise ValueError(f"Unknown file store type: {type(store)}")
@ -242,6 +249,14 @@ def create_file_store(store_type: str) -> BaseFileStore:
embedding_model=embedding_model,
fts_enabled=config.LOCAL_FTS_ENABLED,
)
elif store_type == "zvec":
return ZvecFileStore(
store_name=config.NAME,
db_path=config.ZVEC_DB_PATH,
embedding_model=embedding_model,
fts_enabled=config.ZVEC_FTS_ENABLED,
dimension=config.EMBEDDING_DIMENSIONS,
)
else:
raise ValueError(f"Unknown store type: {store_type}")
@ -284,6 +299,13 @@ async def test_start_store(store: BaseFileStore, _store_name: str):
assert isinstance(store._files, dict), "Files index should be a dict"
logger.info(f"✓ LocalFileStore ready (chunks file: {store._chunks_file})")
# Verify ZvecFileStore initialized
if isinstance(store, ZvecFileStore):
# pylint: disable=protected-access
assert store._collection is not None, "Zvec collection should be initialized"
assert store._initialized, "Zvec engine should be initialized"
logger.info(f"✓ ZvecFileStore ready (collection: {store.collection_name})")
async def test_upsert_file(store: BaseFileStore, _store_name: str) -> tuple[FileMetadata, List[MemoryChunk]]:
"""Test file and chunks insertion."""
@ -1013,6 +1035,18 @@ async def cleanup_store(store: BaseFileStore, store_type: str):
json_file.unlink()
logger.info(f"✓ Cleaned up file: {json_file}")
# Clean up zvec directory and metadata file
if store_type == "zvec":
config = TestConfig()
db_dir = Path(config.ZVEC_DB_PATH)
if db_dir.exists():
shutil.rmtree(db_dir)
logger.info(f"✓ Cleaned up directory: {db_dir}")
metadata_file = db_dir.parent / f"{config.NAME}_file_metadata.json"
if metadata_file.exists():
metadata_file.unlink()
logger.info(f"✓ Cleaned up metadata file: {metadata_file}")
logger.info("✓ Cleanup completed")
except Exception as e:
logger.error(f"Cleanup error: {e}")
@ -1049,6 +1083,11 @@ Examples:
action="store_true",
help="Test LocalFileStore",
)
parser.add_argument(
"--zvec",
action="store_true",
help="Test ZvecFileStore",
)
parser.add_argument(
"--all",
action="store_true",
@ -1065,6 +1104,7 @@ Examples:
("sqlite", "SqliteFileStore"),
("chroma", "ChromaFileStore"),
("local", "LocalFileStore"),
("zvec", "ZvecFileStore"),
]
else:
# Build list based on individual flags
@ -1074,6 +1114,8 @@ Examples:
stores_to_test.append(("chroma", "ChromaFileStore"))
if args.local:
stores_to_test.append(("local", "LocalFileStore"))
if args.zvec:
stores_to_test.append(("zvec", "ZvecFileStore"))
if not stores_to_test:
# Default to all file stores if no argument provided
@ -1081,9 +1123,10 @@ Examples:
("sqlite", "SqliteFileStore"),
("chroma", "ChromaFileStore"),
("local", "LocalFileStore"),
("zvec", "ZvecFileStore"),
]
print("No file store specified, defaulting to test all file stores")
print("Use --sqlite, --chroma, or --local to test specific ones\n")
print("Use --sqlite, --chroma, --local, or --zvec to test specific ones\n")
# Run tests for each file store
for store_type, store_name in stores_to_test:

View file

@ -2,7 +2,7 @@
"""Unified test suite for vector store implementations.
This module provides comprehensive test coverage for LocalVectorStore, ESVectorStore,
PGVectorStore, QdrantVectorStore, ChromaVectorStore, and ObVecVectorStore implementations.
PGVectorStore, QdrantVectorStore, ChromaVectorStore, ObVecVectorStore and ZvecVectorStore implementations.
Tests can be run for specific vector stores or all implementations.
Usage:
@ -12,6 +12,7 @@ Usage:
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 --zvec # Test ZvecVectorStore only
python test_vector_store.py --all # Test all vector stores
"""
@ -36,6 +37,7 @@ from reme.core.vector_store import (
ObVecVectorStore,
PGVectorStore,
QdrantVectorStore,
ZvecVectorStore,
)
load_env()
@ -90,6 +92,9 @@ class TestConfig:
OBVEC_PASSWORD = os.environ.get("OBVEC_PASSWORD", "root")
OBVEC_DATABASE = os.environ.get("OBVEC_DATABASE", "test")
# ZvecVectorStore settings
ZVEC_PATH = "./test_vector_store_zvec" # For local persistent mode
# Embedding model settings
EMBEDDING_MODEL_NAME = "text-embedding-v4"
EMBEDDING_DIMENSIONS = 64
@ -192,6 +197,7 @@ class SampleDataGenerator:
# ==================== Vector Store Factory ====================
# pylint: disable=too-many-return-statements
def get_store_type(store: BaseVectorStore) -> str:
"""Get the type identifier of a vector store instance.
@ -199,7 +205,7 @@ def get_store_type(store: BaseVectorStore) -> str:
store: Vector store instance
Returns:
str: Type identifier ("local", "es", "pgvector", "qdrant", "chroma", or "obvec")
str: Type identifier ("local", "es", "pgvector", "qdrant", "chroma", "obvec", or "zvec")
"""
if isinstance(store, LocalVectorStore):
return "local"
@ -213,10 +219,13 @@ def get_store_type(store: BaseVectorStore) -> str:
return "chroma"
elif isinstance(store, ObVecVectorStore):
return "obvec"
elif isinstance(store, ZvecVectorStore):
return "zvec"
else:
raise ValueError(f"Unknown vector store type: {type(store)}")
# pylint: disable=too-many-return-statements
def create_vector_store(store_type: str, collection_name: str) -> BaseVectorStore:
"""Create a vector store instance based on type.
@ -295,6 +304,14 @@ def create_vector_store(store_type: str, collection_name: str) -> BaseVectorStor
index_metric="cosine",
index_ef_search=100,
)
elif store_type == "zvec":
return ZvecVectorStore(
collection_name=collection_name,
embedding_model=embedding_model,
db_path=config.ZVEC_PATH or tempfile.mkdtemp(prefix="test_zvec_"),
dimension=config.EMBEDDING_DIMENSIONS,
distance="cosine",
)
else:
raise ValueError(f"Unknown store type: {store_type}")
@ -1790,6 +1807,13 @@ async def cleanup_store(store: BaseVectorStore, store_type: str):
shutil.rmtree(obvec_dir, ignore_errors=True)
logger.info(f"Cleaned up obvec temp directory: {obvec_dir}")
# Clean up local directory if ZvecVectorStore
if store_type == "zvec" and config.ZVEC_PATH:
test_dir = Path(config.ZVEC_PATH)
if test_dir.exists():
shutil.rmtree(test_dir)
logger.info(f"Cleaned up zvec directory: {config.ZVEC_PATH}")
logger.info("✓ Cleanup completed")
except Exception as e:
logger.error(f"Cleanup error: {e}")
@ -1844,6 +1868,11 @@ Examples:
action="store_true",
help="Test ObVecVectorStore",
)
parser.add_argument(
"--zvec",
action="store_true",
help="Test ZvecVectorStore",
)
parser.add_argument(
"--all",
action="store_true",
@ -1863,6 +1892,7 @@ Examples:
("qdrant", "QdrantVectorStore"),
("chroma", "ChromaVectorStore"),
("obvec", "ObVecVectorStore"),
("zvec", "ZvecVectorStore"),
]
else:
# Build list based on individual flags
@ -1878,6 +1908,8 @@ Examples:
stores_to_test.append(("chroma", "ChromaVectorStore"))
if args.obvec:
stores_to_test.append(("obvec", "ObVecVectorStore"))
if args.zvec:
stores_to_test.append(("zvec", "ZvecVectorStore"))
if not stores_to_test:
# Default to all vector stores if no argument provided
@ -1888,10 +1920,11 @@ Examples:
("qdrant", "QdrantVectorStore"),
("chroma", "ChromaVectorStore"),
("obvec", "ObVecVectorStore"),
("zvec", "ZvecVectorStore"),
]
print("No vector store specified, defaulting to test all vector stores")
print(
"Use --local/--es/--pgvector/--qdrant/--chroma/--obvec to test specific ones\n",
"Use --local/--es/--pgvector/--qdrant/--chroma/--obvec/--zvec to test specific ones\n",
)
# Run tests for each vector store

View file

@ -0,0 +1,906 @@
"""Test suite for ZvecVectorStore implementation.
Comprehensive tests covering CRUD operations, search, filtering,
collection management, and edge cases for the zvec vector store adapter.
Usage:
python -m pytest tests/test_zvec_vector_store.py -v
python tests/test_zvec_vector_store.py
"""
# pylint: disable=redefined-outer-name,unused-argument
from __future__ import annotations
import asyncio
import shutil
import tempfile
from pathlib import Path
from typing import List
from uuid import uuid4
import pytest
from loguru import logger
from reme.core.schema import VectorNode
from reme.core.vector_store import ZvecVectorStore
# ---------------------------------------------------------------------------
# Skip entire module if zvec native library is not installed
# ---------------------------------------------------------------------------
try:
import zvec as _zvec # noqa: F401 — just checking availability
except ImportError:
pytest.skip("zvec native library not installed", allow_module_location=True)
# ==================== Configuration ====================
class TestConfig:
"""Configuration for zvec test execution."""
ZVEC_ROOT_PATH = tempfile.mkdtemp(prefix="test_zvec_")
EMBEDDING_DIMENSION = 64 # Small dimension for faster tests
TEST_COLLECTION_PREFIX = "test_zvec_vs"
# ==================== Sample Data ====================
def create_sample_nodes(prefix: str = "") -> List[VectorNode]:
"""Create sample VectorNode instances for testing."""
id_prefix = f"{prefix}_" if prefix else ""
return [
VectorNode(
vector_id=f"{id_prefix}node1",
content="Artificial intelligence is a technology that simulates human intelligence.",
metadata={
"node_type": "tech",
"category": "AI",
"source": "research",
"priority": "high",
"year": "2023",
},
),
VectorNode(
vector_id=f"{id_prefix}node2",
content="Machine learning is a subset of artificial intelligence.",
metadata={
"node_type": "tech",
"category": "ML",
"source": "research",
"priority": "high",
"year": "2022",
},
),
VectorNode(
vector_id=f"{id_prefix}node3",
content="Deep learning uses neural networks with multiple layers.",
metadata={
"node_type": "tech_new",
"category": "DL",
"source": "blog",
"priority": "medium",
"year": "2024",
},
),
VectorNode(
vector_id=f"{id_prefix}node4",
content="I love eating delicious seafood, especially fresh fish.",
metadata={
"node_type": "food",
"category": "preference",
"source": "personal",
"priority": "low",
"year": "2023",
},
),
VectorNode(
vector_id=f"{id_prefix}node5",
content="Natural language processing enables computers to understand human language.",
metadata={
"node_type": "tech",
"category": "NLP",
"source": "research",
"priority": "high",
"year": "2024",
},
),
]
# ==================== Fixtures ====================
class MockEmbeddingModel:
"""A mock embedding model that generates deterministic random vectors.
Avoids external API calls during testing. Produces unit-normalized
vectors so that cosine similarity works correctly.
"""
def __init__(self, dimension: int = 64):
self.dimension = dimension
async def get_embedding(self, query: str) -> list[float]:
"""Generate a deterministic embedding from a query string."""
import hashlib
import struct
h = hashlib.sha256(query.encode()).digest()
# Repeat hash to fill dimension
full_hash = b""
while len(full_hash) < self.dimension * 4:
full_hash += hashlib.sha256(h + full_hash).digest()
vec = list(struct.unpack(f"<{self.dimension}f", full_hash[: self.dimension * 4]))
# Normalize to unit vector
norm = sum(x * x for x in vec) ** 0.5
if norm > 0:
vec = [x / norm for x in vec]
return vec
async def get_embeddings(self, queries: list[str]) -> list[list[float]]:
"""Generate embeddings for multiple queries."""
return [await self.get_embedding(q) for q in queries]
async def get_node_embedding(self, node: VectorNode) -> VectorNode:
"""Assign embedding to a single node."""
if node.content:
node.vector = await self.get_embedding(node.content)
return node
async def get_node_embeddings(self, nodes: list[VectorNode]) -> list[VectorNode]:
"""Assign embeddings to multiple nodes."""
return [await self.get_node_embedding(n) for n in nodes]
@pytest.fixture
def embedding_model():
"""Provide a MockEmbeddingModel for tests."""
return MockEmbeddingModel(dimension=TestConfig.EMBEDDING_DIMENSION)
@pytest.fixture
def zvec_store(embedding_model, tmp_path):
"""Create and start a ZvecVectorStore for testing.
Yields the store and cleans up afterwards.
"""
collection_name = f"{TestConfig.TEST_COLLECTION_PREFIX}_{uuid4().hex[:8]}"
store = ZvecVectorStore(
collection_name=collection_name,
db_path=str(tmp_path / "zvec_db"),
embedding_model=embedding_model,
dimension=TestConfig.EMBEDDING_DIMENSION,
distance="cosine",
)
async def _setup():
await store.start()
return store
store = asyncio.get_event_loop().run_until_complete(_setup())
yield store
async def _teardown():
try:
await store.close()
except Exception:
pass
# Clean up temp directory
db_path = Path(str(tmp_path / "zvec_db"))
if db_path.exists():
shutil.rmtree(db_path, ignore_errors=True)
asyncio.get_event_loop().run_until_complete(_teardown())
# ==================== Helper ====================
def run(coro):
"""Run an async coroutine in the current event loop."""
return asyncio.get_event_loop().run_until_complete(coro)
# ==================== Test: Collection Lifecycle ====================
class TestCollectionLifecycle:
"""Tests for collection creation, listing, deletion, and copy."""
def test_create_collection(self, zvec_store):
"""Test that a collection is created during start()."""
collections = run(zvec_store.list_collections())
assert zvec_store.collection_name in collections
def test_list_collections(self, zvec_store):
"""Test listing collections."""
collections = run(zvec_store.list_collections())
assert isinstance(collections, list)
assert len(collections) >= 1
def test_delete_collection(self, zvec_store, embedding_model, tmp_path):
"""Test deleting a collection."""
# Create a secondary collection
coll_name = f"del_test_{uuid4().hex[:8]}"
store2 = ZvecVectorStore(
collection_name=coll_name,
db_path=str(tmp_path / "zvec_db"),
embedding_model=embedding_model,
dimension=TestConfig.EMBEDDING_DIMENSION,
)
run(store2.start())
collections = run(zvec_store.list_collections())
assert coll_name in collections
run(zvec_store.delete_collection(coll_name))
collections = run(zvec_store.list_collections())
assert coll_name not in collections
def test_copy_collection(self, zvec_store, embedding_model, tmp_path):
"""Test copying a collection."""
# Insert some data first
nodes = create_sample_nodes("copy")
run(zvec_store.insert(nodes))
copy_name = f"copy_test_{uuid4().hex[:8]}"
run(zvec_store.copy_collection(copy_name))
# Verify copy exists
collections = run(zvec_store.list_collections())
assert copy_name in collections
# Clean up
run(zvec_store.delete_collection(copy_name))
# ==================== Test: Insert ====================
class TestInsert:
"""Tests for node insertion (single and batch)."""
def test_insert_single_node(self, zvec_store):
"""Test inserting a single node."""
node = VectorNode(
vector_id="single_1",
content="This is a single node insertion test",
metadata={"test_type": "single_insert"},
)
run(zvec_store.insert(node))
result = run(zvec_store.get("single_1"))
assert result is not None
assert result.vector_id == "single_1"
assert "single node" in result.content
def test_insert_batch_nodes(self, zvec_store):
"""Test inserting multiple nodes in batch."""
nodes = create_sample_nodes("batch")
run(zvec_store.insert(nodes))
all_nodes = run(zvec_store.list(limit=10))
assert len(all_nodes) >= len(nodes)
def test_insert_node_with_vector(self, zvec_store):
"""Test inserting a node that already has a vector."""
node = VectorNode(
vector_id="prevec_1",
content="Node with pre-computed vector",
vector=[0.1] * TestConfig.EMBEDDING_DIMENSION,
metadata={"test_type": "pre_vector"},
)
run(zvec_store.insert(node))
result = run(zvec_store.get("prevec_1"))
assert result is not None
assert result.vector is not None
# ==================== Test: Search ====================
class TestSearch:
"""Tests for vector similarity search."""
@pytest.fixture(autouse=True)
def _insert_sample_data(self, zvec_store):
"""Insert sample data before each search test."""
nodes = create_sample_nodes("search")
run(zvec_store.insert(nodes))
def test_basic_search(self, zvec_store):
"""Test basic vector search."""
results = run(zvec_store.search(query="What is artificial intelligence?", limit=3))
assert len(results) > 0
for r in results:
assert isinstance(r, VectorNode)
assert r.content
def test_search_with_limit(self, zvec_store):
"""Test search with various limits."""
results = run(zvec_store.search(query="technology", limit=2))
assert len(results) <= 2
def test_search_with_filter(self, zvec_store):
"""Test vector search with metadata filter."""
results = run(
zvec_store.search(
query="What is artificial intelligence?",
limit=5,
filters={"node_type": "tech"},
),
)
# All results should have node_type == "tech"
for r in results:
assert r.metadata.get("node_type") == "tech"
def test_search_with_multiple_filters(self, zvec_store):
"""Test search with multiple metadata filters (AND)."""
results = run(
zvec_store.search(
query="What is artificial intelligence?",
limit=5,
filters={"node_type": "tech", "source": "research"},
),
)
for r in results:
assert r.metadata.get("node_type") == "tech"
assert r.metadata.get("source") == "research"
def test_search_relevance_ranking(self, zvec_store):
"""Test that search results have scores and are relevant."""
results = run(zvec_store.search(query="artificial intelligence", limit=5))
assert len(results) > 0
# All results should have a score
for r in results:
assert "score" in r.metadata
assert r.metadata["score"] > 0
# The top result should be highly relevant (AI content matches AI query)
top_content = results[0].content.lower()
assert "artificial intelligence" in top_content or "intelligence" in top_content or "ai" in top_content
# ==================== Test: Get ====================
class TestGet:
"""Tests for retrieving nodes by ID."""
@pytest.fixture(autouse=True)
def _insert_sample_data(self, zvec_store):
"""Insert sample data before each get test."""
nodes = create_sample_nodes("get")
run(zvec_store.insert(nodes))
def test_get_single_id(self, zvec_store):
"""Test retrieving a single node by ID."""
result = run(zvec_store.get("get_node1"))
assert result is not None
assert result.vector_id == "get_node1"
def test_get_multiple_ids(self, zvec_store):
"""Test retrieving multiple nodes by IDs."""
results = run(zvec_store.get(["get_node1", "get_node2"]))
assert isinstance(results, list)
assert len(results) >= 2
result_ids = {r.vector_id for r in results}
assert "get_node1" in result_ids
assert "get_node2" in result_ids
def test_get_nonexistent_id(self, zvec_store):
"""Test retrieving a non-existent ID."""
result = run(zvec_store.get("nonexistent_id_xyz"))
assert result is None or result == []
# ==================== Test: List ====================
class TestList:
"""Tests for listing nodes with optional filters and sorting."""
@pytest.fixture(autouse=True)
def _insert_sample_data(self, zvec_store):
"""Insert sample data before each list test."""
nodes = create_sample_nodes("list")
run(zvec_store.insert(nodes))
def test_list_all(self, zvec_store):
"""Test listing all nodes."""
results = run(zvec_store.list(limit=20))
assert len(results) > 0
def test_list_with_filter(self, zvec_store):
"""Test listing nodes with metadata filter."""
results = run(zvec_store.list(filters={"category": "AI"}, limit=10))
for r in results:
assert r.metadata.get("category") == "AI"
def test_list_with_sorting(self, zvec_store):
"""Test listing with sorting by metadata key."""
# Insert nodes with numeric metadata for sorting
sort_nodes = [
VectorNode(
vector_id=f"sort_{i}",
content=f"Sort test node {i}",
metadata={"rating": str(50 + i * 5), "test_type": "sort_test"},
)
for i in range(10)
]
run(zvec_store.insert(sort_nodes))
results = run(
zvec_store.list(
filters={"test_type": "sort_test"},
sort_key="rating",
reverse=True,
limit=5,
),
)
assert len(results) <= 5
# Verify descending order
ratings = [r.metadata.get("rating") for r in results]
for i in range(len(ratings) - 1):
assert ratings[i] >= ratings[i + 1]
# ==================== Test: Update ====================
class TestUpdate:
"""Tests for updating existing nodes."""
@pytest.fixture(autouse=True)
def _insert_sample_data(self, zvec_store):
"""Insert sample data before each update test."""
nodes = create_sample_nodes("upd")
run(zvec_store.insert(nodes))
def test_update_single_node(self, zvec_store):
"""Test updating a single node's content and metadata."""
updated = VectorNode(
vector_id="upd_node2",
content="Machine learning is a powerful subset of AI that learns from data.",
metadata={
"node_type": "tech",
"category": "ML",
"updated": "true",
},
)
run(zvec_store.update(updated))
result = run(zvec_store.get("upd_node2"))
assert result is not None
assert result.metadata.get("updated") == "true"
def test_update_batch(self, zvec_store):
"""Test batch updating multiple nodes."""
updates = [
VectorNode(
vector_id="upd_node1",
content="Updated content for node 1",
metadata={"node_type": "tech", "batch_updated": "true"},
),
VectorNode(
vector_id="upd_node3",
content="Updated content for node 3",
metadata={"node_type": "tech_new", "batch_updated": "true"},
),
]
run(zvec_store.update(updates))
results = run(zvec_store.get(["upd_node1", "upd_node3"]))
if isinstance(results, list):
for r in results:
assert r.metadata.get("batch_updated") == "true"
# ==================== Test: Delete ====================
class TestDelete:
"""Tests for deleting nodes."""
@pytest.fixture(autouse=True)
def _insert_sample_data(self, zvec_store):
"""Insert sample data before each delete test."""
nodes = create_sample_nodes("del")
run(zvec_store.insert(nodes))
def test_delete_single(self, zvec_store):
"""Test deleting a single node by ID."""
run(zvec_store.delete("del_node4"))
# Verify deletion
result = run(zvec_store.get("del_node4"))
assert result is None or result == []
def test_delete_batch(self, zvec_store):
"""Test batch deleting multiple nodes by IDs."""
# First insert some extra nodes to delete
extra_nodes = [
VectorNode(
vector_id=f"del_extra_{i}",
content=f"Extra node {i} for batch delete test",
metadata={"test_type": "batch_delete"},
)
for i in range(5)
]
run(zvec_store.insert(extra_nodes))
ids = [f"del_extra_{i}" for i in range(5)]
run(zvec_store.delete(ids))
# Verify all deleted
for nid in ids:
result = run(zvec_store.get(nid))
assert result is None or result == []
def test_delete_all(self, zvec_store):
"""Test deleting all nodes from the collection."""
run(zvec_store.delete_all())
# Collection should be empty now
remaining = run(zvec_store.list(limit=100))
assert len(remaining) == 0
# ==================== Test: Edge Cases ====================
class TestEdgeCases:
"""Tests for edge cases and boundary conditions."""
def test_empty_content(self, zvec_store):
"""Test inserting a node with empty content."""
node = VectorNode(
vector_id="edge_empty",
content="",
metadata={"type": "empty"},
)
# Empty content may fail embedding — that's OK, we just want to see it handled
try:
run(zvec_store.insert([node]))
except Exception:
pass # Expected if embedding fails on empty string
def test_long_content(self, zvec_store):
"""Test inserting a node with very long content."""
node = VectorNode(
vector_id="edge_long",
content="A" * 5000,
metadata={"type": "long_content"},
)
run(zvec_store.insert([node]))
result = run(zvec_store.get("edge_long"))
assert result is not None
assert len(result.content) == 5000
def test_special_characters(self, zvec_store):
"""Test content with special characters."""
node = VectorNode(
vector_id="edge_special",
content="Special chars: @#$%^&*()[]{}|;:',.<>?/~`",
metadata={"type": "special_chars"},
)
run(zvec_store.insert([node]))
result = run(zvec_store.get("edge_special"))
assert result is not None
assert "@#$%" in result.content
def test_unicode_content(self, zvec_store):
"""Test content with Unicode characters."""
node = VectorNode(
vector_id="edge_unicode",
content="Unicode test: 你好世界 مرحبا Привет",
metadata={"type": "unicode"},
)
run(zvec_store.insert([node]))
result = run(zvec_store.get("edge_unicode"))
assert result is not None
assert "你好世界" in result.content
def test_nonexistent_id(self, zvec_store):
"""Test getting a non-existent ID."""
result = run(zvec_store.get("nonexistent_xyz_999"))
assert result is None or result == []
def test_metadata_with_empty_string_value(self, zvec_store):
"""Test metadata containing empty string values."""
node = VectorNode(
vector_id="edge_meta_empty",
content="Testing empty metadata values",
metadata={"field1": "value1", "field2": "", "field3": "value3"},
)
run(zvec_store.insert([node]))
result = run(zvec_store.get("edge_meta_empty"))
assert result is not None
def test_search_nonexistent_filter(self, zvec_store):
"""Test search with a filter value that doesn't match anything."""
nodes = create_sample_nodes("edge_filter")
run(zvec_store.insert(nodes))
results = run(
zvec_store.search(
query="test",
limit=10,
filters={"category": "NONEXISTENT_CATEGORY"},
),
)
assert len(results) == 0
# ==================== Test: Batch Operations ====================
class TestBatchOperations:
"""Tests for large-scale batch insert, update, and delete."""
def test_batch_insert_100_nodes(self, zvec_store):
"""Test inserting 100 nodes in batch."""
batch_nodes = [
VectorNode(
vector_id=f"batch_{i}",
content=f"This is batch test content number {i} about technology and science.",
metadata={
"batch_id": str(i // 10),
"index": str(i),
"category": ["tech", "science", "business"][i % 3],
},
)
for i in range(100)
]
run(zvec_store.insert(batch_nodes))
all_nodes = run(zvec_store.list(limit=150))
assert len(all_nodes) >= 100
def test_batch_update_20_nodes(self, zvec_store):
"""Test batch updating 20 nodes."""
# Insert first
nodes = [
VectorNode(
vector_id=f"bupd_{i}",
content=f"Batch update test {i}",
metadata={"index": str(i)},
)
for i in range(30)
]
run(zvec_store.insert(nodes))
# Update first 20
updates = [
VectorNode(
vector_id=f"bupd_{i}",
content=f"UPDATED content {i}",
metadata={"index": str(i), "updated": "true"},
)
for i in range(20)
]
run(zvec_store.update(updates))
# Verify
results = run(zvec_store.list(filters={"updated": "true"}, limit=30))
assert len(results) >= 20
def test_batch_delete_50_nodes(self, zvec_store):
"""Test batch deleting 50 nodes."""
# Insert
nodes = [
VectorNode(
vector_id=f"bdel_{i}",
content=f"Batch delete test {i}",
metadata={"index": str(i)},
)
for i in range(50)
]
run(zvec_store.insert(nodes))
# Delete
ids = [f"bdel_{i}" for i in range(50)]
run(zvec_store.delete(ids))
# Verify
remaining = run(zvec_store.list(limit=200))
batch_remaining = [n for n in remaining if n.vector_id.startswith("bdel_")]
assert len(batch_remaining) == 0
# ==================== Test: Concurrent Operations ====================
class TestConcurrentOperations:
"""Tests for concurrent read/write operations."""
def test_concurrent_inserts_and_searches(self, zvec_store):
"""Test that concurrent inserts and searches work without errors."""
async def _run():
# Concurrent inserts
insert_tasks = []
for i in range(5):
batch = [
VectorNode(
vector_id=f"conc_{i}_{j}",
content=f"Concurrent test content {i}-{j}",
metadata={"thread_id": str(i)},
)
for j in range(10)
]
insert_tasks.append(zvec_store.insert(batch))
await asyncio.gather(*insert_tasks)
# Concurrent searches
search_tasks = [zvec_store.search(query="concurrent test", limit=5) for _ in range(5)]
search_results = await asyncio.gather(*search_tasks)
# All searches should return results
for results in search_results:
assert len(results) > 0
run(_run())
# ==================== Test: Data Model Conversion ====================
class TestDataModelConversion:
"""Tests for VectorNode <-> zvec Doc conversion helpers."""
def test_vector_node_to_doc_roundtrip(self, zvec_store):
"""Test that VectorNode -> Doc -> VectorNode roundtrip preserves data."""
from reme.core.vector_store.zvec_vector_store import (
_vector_node_to_doc,
_doc_to_vector_node,
)
original = VectorNode(
vector_id="roundtrip_1",
content="Roundtrip test content",
vector=[0.1] * TestConfig.EMBEDDING_DIMENSION,
metadata={"key1": "value1", "key2": "42", "key3": "true"},
)
doc = _vector_node_to_doc(original)
assert doc.id == "roundtrip_1"
assert doc.field("content") == "Roundtrip test content"
restored = _doc_to_vector_node(doc, include_score=False)
assert restored.vector_id == "roundtrip_1"
assert restored.content == "Roundtrip test content"
assert restored.metadata.get("key1") == "value1"
def test_post_filter_exact_match(self):
"""Test post-filtering with exact match."""
from reme.core.vector_store.zvec_vector_store import _apply_filters_post
nodes = [
VectorNode(vector_id="1", content="a", metadata={"category": "AI"}),
VectorNode(vector_id="2", content="b", metadata={"category": "ML"}),
VectorNode(vector_id="3", content="c", metadata={"category": "AI"}),
]
filtered = _apply_filters_post(nodes, {"category": "AI"})
assert len(filtered) == 2
assert all(n.metadata["category"] == "AI" for n in filtered)
def test_post_filter_range_query(self):
"""Test post-filtering with range query."""
from reme.core.vector_store.zvec_vector_store import _apply_filters_post
nodes = [
VectorNode(vector_id="1", content="a", metadata={"year": 2022}),
VectorNode(vector_id="2", content="b", metadata={"year": 2023}),
VectorNode(vector_id="3", content="c", metadata={"year": 2024}),
]
filtered = _apply_filters_post(nodes, {"year": [2023, 2024]})
assert len(filtered) == 2
def test_post_filter_none_and_empty(self):
"""Test post-filtering with None and empty filters."""
from reme.core.vector_store.zvec_vector_store import _apply_filters_post
nodes = [VectorNode(vector_id="1", content="a", metadata={})]
# None filter returns all
assert _apply_filters_post(nodes, None) == nodes
# Empty filter returns all
assert _apply_filters_post(nodes, {}) == nodes
def test_score_excluded_from_stored_metadata(self):
"""Test that score is excluded when converting VectorNode to Doc."""
from reme.core.vector_store.zvec_vector_store import _vector_node_to_doc
node = VectorNode(
vector_id="score_test",
content="test",
vector=[0.1] * TestConfig.EMBEDDING_DIMENSION,
metadata={"key1": "val1", "score": 0.95},
)
doc = _vector_node_to_doc(node)
# The metadata JSON should NOT contain the score key
import json
stored_meta = json.loads(doc.field("metadata"))
assert "score" not in stored_meta
assert "key1" in stored_meta
# ==================== Main Entry Point ====================
async def run_standalone_tests():
"""Run tests standalone (without pytest) for quick validation."""
tmp_dir = tempfile.mkdtemp(prefix="test_zvec_standalone_")
embedding_model = MockEmbeddingModel(dimension=TestConfig.EMBEDDING_DIMENSION)
store = ZvecVectorStore(
collection_name="standalone_test",
db_path=tmp_dir,
embedding_model=embedding_model,
dimension=TestConfig.EMBEDDING_DIMENSION,
distance="cosine",
)
try:
await store.start()
logger.info("✓ Store started")
# Insert
nodes = create_sample_nodes("std")
await store.insert(nodes)
logger.info(f"✓ Inserted {len(nodes)} nodes")
# Search
results = await store.search(query="artificial intelligence", limit=3)
logger.info(f"✓ Search returned {len(results)} results")
for r in results:
logger.info(f" - {r.vector_id}: {r.content[:50]}... (score={r.metadata.get('score')})")
# Get
result = await store.get("std_node1")
logger.info(f"✓ Get: {result.vector_id if result else 'None'}")
# List
all_nodes = await store.list(limit=10)
logger.info(f"✓ List: {len(all_nodes)} nodes")
# Update
await store.update(
VectorNode(
vector_id="std_node1",
content="Updated content",
metadata={"updated": "true"},
),
)
result = await store.get("std_node1")
logger.info(f"✓ Update: metadata.updated={result.metadata.get('updated') if result else 'N/A'}")
# Delete
await store.delete("std_node4")
result = await store.get("std_node4")
logger.info(f"✓ Delete: {'gone' if result is None or result == [] else 'still exists'}")
# Count
count = await store.count()
logger.info(f"✓ Count: {count} nodes")
logger.info("✓ All standalone tests passed!")
finally:
await store.close()
shutil.rmtree(tmp_dir, ignore_errors=True)
if __name__ == "__main__":
asyncio.run(run_standalone_tests())

View file

@ -20,7 +20,7 @@ async def main():
"dimensions": 1024,
},
default_vector_store_config={
"backend": "local", # 支持 local/chroma/qdrant/elasticsearch
"backend": "local", # 支持 local/chroma/qdrant/elasticsearch/zvec
},
)
await reme.start()