Merge branch 'main' into feat/support_hologres

This commit is contained in:
yangtiancheng-ali 2026-05-09 10:21:43 +08:00 • committed by GitHub
commit 1cdce924df
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
33 changed files with 3749 additions and 437 deletions

View file

@ -506,7 +506,7 @@ async def main():
"dimensions": 1024,
},
default_vector_store_config={
"backend": "local", # Supports local/chroma/qdrant/elasticsearch/obvec/hologres
"backend": "local", # Supports local/chroma/qdrant/elasticsearch/obvec/zvec/hologres
},
)
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

@ -643,7 +643,7 @@ class LocomoEvaluator:
},
"personal_retriever": {
"prompt_dict": {
"user_message": self.retriever_prompt,
"user_message_s2": self.retriever_prompt,
},
"params": {
"return_memory_nodes": True,

View file

@ -132,7 +132,7 @@ user_message_retrieve: |
You are a Memory Retrieval Agent specialized in retrieving {memory_type} memories about {memory_target}.
## User Profile
{user_profile}
{profiles}
## User Question
{context}

View file

@ -35,6 +35,7 @@ FlowLLM provides multiple Vector Store implementations tailored to different use
- **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.
- **HologresVectorStore** ([source code](https://github.com/agentscope-ai/ReMe/blob/main/reme/core/vector_store/hologres_store.py)): Uses [asyncpg](https://pypi.org/project/asyncpg/) against **Hologres** (PostgreSQL-compatible). Leverages native `float4[]` vector storage with built-in HGraph index for approximate nearest neighbor search. Suitable when you already run Hologres or need high-performance vector search with JSONB metadata filtering in a PostgreSQL-compatible environment.
- **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.
@ -131,6 +132,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`).
### HologresVectorStore Configuration
@ -166,7 +172,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`, `hologres`.
- **`backend`** (required): Vector store backend type. Options: `local`, `memory`, `chroma`, `qdrant`, `elasticsearch`, `obvec`, `zvec`, `hologres`.
- **`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.
@ -392,6 +398,30 @@ vector_stores.default.password=your-password
vector_stores.default.database=postgres
```
#### 8. 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:
@ -450,6 +480,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, ObVecVectorStore (OceanBase/seekdb), or HologresVectorStore 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.
- **Hologres**: Use HologresVectorStore when you run Hologres and need native HGraph-indexed vector search with PostgreSQL-compatible SQL and JSONB metadata filtering.

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

@ -8,6 +8,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__ = [
@ -19,6 +20,7 @@ __all__ = [
"ObVecVectorStore",
"PGVectorStore",
"QdrantVectorStore",
"ZvecVectorStore",
]
R.vector_stores.register("chroma")(ChromaVectorStore)
@ -28,3 +30,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

@ -1,39 +1,19 @@
"""Personal memory retriever agent for retrieving personal memories through vector search."""
from loguru import logger
from ..base_memory_agent import BaseMemoryAgent
from ....core.enumeration import Role, MemoryType
from ....core.enumeration import MemoryType, Role
from ....core.op import BaseTool
from ....core.schema import Message
from ....core.utils import format_messages
_PROFILE_TOOL_NAMES: tuple[str, ...] = ("retrieve_profile", "read_all_profiles")
_EMPTY_PROFILE_RESULTS: tuple[str, ...] = ("", "No profiles found.", "No new profiles found.")
class PersonalRetriever(BaseMemoryAgent):
"""Retrieve personal memories through vector search and history reading.
clear && python benchmark/halumem/eval_reme.py \
--data_path /Users/yuli/workspace/HaluMem/data/HaluMem-Medium.jsonl \
--reme_model_name qwen3.5-plus \
--batch_size 10000 \
--algo_version default
📊 Question Answering (with LLM answer):
Correct (all): 0.8537
Hallucination (all): 0.1159
Omission (all): 0.0305
Correct (valid): 0.8537
Hallucination (valid): 0.1159
Omission (valid): 0.0305
Valid/Total: 164/164
📊 Question Answering (with original memories):
Correct (all): 0.9085
Hallucination (all): 0.0671
Omission (all): 0.0244
Correct (valid): 0.9085
Hallucination (valid): 0.0671
Omission (valid): 0.0244
Valid/Total: 164/164
"""
"""Retrieve personal memories through vector search and history reading."""
memory_type: MemoryType = MemoryType.PERSONAL
@ -41,36 +21,73 @@ class PersonalRetriever(BaseMemoryAgent):
super().__init__(**kwargs)
self.return_memory_nodes: bool = return_memory_nodes
async def build_messages(self) -> list[Message]:
def _get_context(self) -> str:
if self.context.get("query"):
context = self.context.query
elif self.context.get("messages"):
context = self.description + "\n" + format_messages(self.context.messages)
else:
raise ValueError("input must have either `query` or `messages`")
read_all_profiles_tool: BaseTool | None = self.pop_tool("read_all_profiles")
if read_all_profiles_tool is not None:
all_profiles = await read_all_profiles_tool.call(
memory_target=self.memory_target,
service_context=self.service_context,
)
else:
all_profiles = ""
return self.context.query.strip()
if self.context.get("messages"):
return (self.description + "\n" + format_messages(self.context.messages)).strip()
raise ValueError("input must have either `query` or `messages`")
async def _build_s1_messages(self, context: str) -> list[Message]:
return [
Message(
role=Role.USER,
content=self.prompt_format(
prompt_name="user_message",
prompt_name="user_message_s1",
memory_type=self.memory_type.value,
memory_target=self.memory_target,
user_profile=all_profiles,
context=context.strip(),
context=context,
),
),
]
async def _build_s2_messages(self, context: str, profiles: str) -> list[Message]:
return [
Message(
role=Role.USER,
content=self.prompt_format(
prompt_name="user_message_s2",
memory_type=self.memory_type.value,
memory_target=self.memory_target,
profiles=profiles,
context=context,
),
),
]
def _partition_tools(self) -> tuple[list[BaseTool], list[BaseTool]]:
profile_tools: list[BaseTool] = []
memory_tools: list[BaseTool] = []
for i, tool in enumerate(self.tools):
name = tool.tool_call.name
if name in _PROFILE_TOOL_NAMES:
profile_tools.append(tool)
else:
memory_tools.append(tool)
logger.info(f"[{self.__class__.__name__}] tool_call[{i}]={tool.tool_call.simple_input_dump(as_dict=False)}")
return profile_tools, memory_tools
@staticmethod
def _extract_profile_context(tools: list[BaseTool]) -> str:
outputs = []
for tool in tools:
response = getattr(tool, "response", None)
answer = getattr(response, "answer", "")
if answer and answer not in _EMPTY_PROFILE_RESULTS:
outputs.append(answer)
return "\n".join(outputs)
async def _run_stage(
self,
stage: str,
messages: list[Message],
tools: list[BaseTool],
) -> tuple[list[BaseTool], list[Message], bool]:
for message in messages:
role = message.name or message.role
logger.info(f"[{self.__class__.__name__} {stage}] role={role} {message.simple_dump(as_dict=False)}")
return await self.react(messages, tools, stage=stage)
async def _acting_step(
self,
assistant_message: Message,
@ -91,7 +108,28 @@ class PersonalRetriever(BaseMemoryAgent):
)
async def execute(self):
result = await super().execute()
context = self._get_context()
profile_tools, memory_tools = self._partition_tools()
tools_s1: list[BaseTool] = []
messages_s1: list[Message] = []
success_s1 = True
profiles = ""
if profile_tools:
messages_s1 = await self._build_s1_messages(context)
tools_s1, messages_s1, success_s1 = await self._run_stage("s1-profile", messages_s1, profile_tools)
profiles = self._extract_profile_context(tools_s1)
messages_s2 = await self._build_s2_messages(context, profiles)
tools_s2, messages_s2, success_s2 = await self._run_stage("s2-memory", messages_s2, memory_tools)
answer = messages_s2[-1].content if success_s2 and messages_s2 else ""
result = {
"answer": answer,
"success": success_s1 and success_s2,
"messages": messages_s1 + messages_s2,
"tools": tools_s1 + tools_s2,
}
if self.return_memory_nodes:
result["answer"] = "\n".join(
[

View file

@ -1,8 +1,26 @@
user_message: |
user_message_s1: |
You are a Profile Retrieval Agent specialized in finding profile information about {memory_target}.
## User Question
{context}
## Task
Use the available profile tool to search for profile content that is relevant to the user question.
## Instructions
- If `retrieve_profile` is available, use it to search with focused profile queries derived from the question
- If `read_all_profiles` is available, use it to inspect the full profile list and identify relevant rows
- Focus on profile attributes such as identity, location, work, education, preferences, relationships, and other long-term facts
- Only retrieve information that is directly relevant to the user question
- If no relevant profile information exists, say so clearly
Output a concise summary of the relevant profile information you found.
user_message_s2: |
You are a Memory Retrieval Agent specialized in retrieving {memory_type} memories about {memory_target}.
## User Profile
{user_profile}
## Profile Search Results
{profiles}
## User Question
{context}
@ -14,11 +32,12 @@ user_message: |
**Tool**: `retrieve_memory` (without time constraints)
**Objective**: Cast a wide net to find potentially relevant memories
**Approach**:
- Use the profile search results above as supporting context when forming retrieval queries
- Execute 3-5 diverse search queries using different formulations:
* Original question verbatim
* Rephrased variations (different wording, synonyms)
* Entity-focused queries (extract and search specific names, places, events)
* Keyword-based searches (core concepts, topics)
* Keyword-based searches (core concepts and profile facts)
* Related context queries (broader themes)
### Phase 2(Optional): Temporal Search
@ -31,7 +50,7 @@ user_message: |
- After date: `20200101,99999999` (from 20200101 onwards)
**Approach**:
- Identify temporal constraints from the user question
- Refine Phase 1 queries with 3-5 diverse appropriate different time filters
- Refine Phase 1 queries with 3-5 diverse appropriate time filters
### Phase 3: Deep Dive into History
**Tool**: `read_history`
@ -48,11 +67,11 @@ user_message: |
- Use this to understand the full conversation surrounding a memory
## Response Guidelines
- Base your answer EXCLUSIVELY on user profile, retrieved memories, and history data
- Base your answer EXCLUSIVELY on the profile search results, retrieved memories, and history data
- Never infer, assume, or hallucinate information
- Always cite sources with timestamps: `[timestamp] Memory content`
- Present conflicting information transparently with respective timestamps
- If you find sufficient information to answer the user's question, you may output directly without exhausting all search phases
- Exhaust all search strategies before concluding information doesn't exist
Output a summary of all retrieved memories, user profile, and history data.
Output a summary of all retrieved memories and relevant history data.

View file

@ -3,13 +3,17 @@
from loguru import logger
from ..base_memory_agent import BaseMemoryAgent
from ....core.enumeration import Role, MemoryType
from ....core.enumeration import MemoryType, Role
from ....core.op import BaseTool
from ....core.schema import Message
# Optional profile tools used to pre-load profile context; consumed by the
# summarizer itself and never exposed to the stage-two ReAct loop.
_PROFILE_CONTEXT_TOOLS: tuple[str, ...] = ("retrieve_profile", "read_all_profiles")
class PersonalSummarizer(BaseMemoryAgent):
"""Two-phase personal memory processor: retrieve/add memories then update profile."""
"""Two-phase personal memory processor: add memories, then update profiles."""
memory_type: MemoryType = MemoryType.PERSONAL
@ -62,62 +66,71 @@ class PersonalSummarizer(BaseMemoryAgent):
**kwargs,
)
async def execute(self):
memory_tools = []
profile_tools = []
read_all_profiles_tool: BaseTool | None = None
def _partition_tools(self) -> tuple[list[BaseTool], list[BaseTool], BaseTool | None]:
"""Split attached tools into memory tools, profile tools, and a profile context tool."""
memory_tools: list[BaseTool] = []
profile_tools: list[BaseTool] = []
profile_context_tool: BaseTool | None = None
for i, tool in enumerate(self.tools):
tool_name = tool.tool_call.name
if tool_name == "read_all_profiles":
read_all_profiles_tool = tool
elif "_memory" in tool_name:
name = tool.tool_call.name
if name in _PROFILE_CONTEXT_TOOLS:
profile_context_tool = tool
elif "_memory" in name:
memory_tools.append(tool)
elif "_profile" in tool_name:
elif "_profile" in name:
profile_tools.append(tool)
else:
raise ValueError(f"[{self.__class__.__name__}] unknown tool_name={tool_name}")
raise ValueError(f"[{self.__class__.__name__}] unknown tool_name={name}")
logger.info(f"[{self.__class__.__name__}] tool_call[{i}]={tool.tool_call.simple_input_dump(as_dict=False)}")
return memory_tools, profile_tools, profile_context_tool
stage = "s1-memory"
messages_s1 = await self._build_s1_messages()
for i, message in enumerate(messages_s1):
async def _preload_user_profile(self, tool: BaseTool | None) -> str:
"""Invoke the profile context tool to obtain inline profile text."""
if tool is None:
return ""
call_kwargs: dict = {
"memory_target": self.memory_target,
"service_context": self.service_context,
"retrieved_nodes": self.retrieved_nodes,
}
if tool.tool_call.name == "retrieve_profile":
call_kwargs["query"] = self.context.history_node.content
return await tool.call(**call_kwargs)
async def _run_stage(
self,
stage: str,
messages: list[Message],
tools: list[BaseTool],
) -> tuple[list[BaseTool], list[Message], bool]:
for message in messages:
role = message.name or message.role
logger.info(f"[{self.__class__.__name__} {stage}] role={role} {message.simple_dump(as_dict=False)}")
tools_s1, messages_s1, success_s1 = await self.react(messages_s1, memory_tools, stage=stage)
return await self.react(messages, tools, stage=stage)
if read_all_profiles_tool is not None:
profiles = await read_all_profiles_tool.call(
memory_target=self.memory_target,
service_context=self.service_context,
)
else:
profiles = ""
async def execute(self):
memory_tools, profile_tools, profile_context_tool = self._partition_tools()
messages_s1 = await self._build_s1_messages()
tools_s1, messages_s1, success_s1 = await self._run_stage("s1-memory", messages_s1, memory_tools)
if profile_tools:
stage = "s2-profile"
profiles = await self._preload_user_profile(profile_context_tool)
messages_s2 = await self._build_s2_messages(profiles)
for i, message in enumerate(messages_s2):
role = message.name or message.role
logger.info(f"[{self.__class__.__name__} {stage}] role={role} {message.simple_dump(as_dict=False)}")
tools_s2, messages_s2, success_s2 = await self.react(messages_s2, profile_tools, stage=stage)
tools_s2, messages_s2, success_s2 = await self._run_stage("s2-profile", messages_s2, profile_tools)
else:
tools_s2, messages_s2, success_s2 = [], [], True
answer = (messages_s1[-1].content if success_s1 and messages_s1 else "") + (
messages_s2[-1].content if success_s2 and messages_s2 else ""
)
success = success_s1 and success_s2
messages = messages_s1 + messages_s2
tools = tools_s1 + tools_s2
memory_nodes = []
for tool in tools:
if tool.memory_nodes:
memory_nodes.extend(tool.memory_nodes)
memory_nodes = [node for tool in tools for node in (tool.memory_nodes or [])]
return {
"answer": answer,
"success": success,
"messages": messages,
"success": success_s1 and success_s2,
"messages": messages_s1 + messages_s2,
"tools": tools,
"memory_nodes": memory_nodes,
}

View file

@ -1,5 +1,7 @@
"""memory tools"""
# pylint: disable=no-name-in-module
from .base_memory_tool import BaseMemoryTool
# chunk tools
@ -15,6 +17,7 @@ from .profiles.add_draft_and_read_all_profiles import AddDraftAndReadAllProfiles
from .profiles.add_profile import AddProfile
from .profiles.delete_profile import DeleteProfile
from .profiles.read_all_profiles import ReadAllProfiles
from .profiles.retrieve_profile import RetrieveProfile
from .profiles.update_profile import UpdateProfile
from .profiles.update_profiles_v1 import UpdateProfilesV1
@ -43,6 +46,7 @@ __all__ = [
"AddProfile",
"DeleteProfile",
"ReadAllProfiles",
"RetrieveProfile",
"UpdateProfile",
"UpdateProfilesV1",
# record tools

View file

@ -3,6 +3,7 @@
from abc import ABCMeta
from pathlib import Path
from .profiles.profile_handler import ProfileHandler
from ...core.enumeration import MemoryType
from ...core.op import BaseTool
from ...core.schema import ToolCall, MemoryNode, ToolAttr
@ -16,12 +17,18 @@ class BaseMemoryTool(BaseTool, metaclass=ABCMeta):
enable_multiple: bool = True,
enable_thinking_params: bool = False,
profile_dir: str = "",
profile_backend: str = "filesystem",
profile_store_name: str = "profile",
profile_max_capacity: int = 50,
**kwargs,
):
super().__init__(**kwargs)
self.enable_multiple: bool = enable_multiple
self.enable_thinking_params: bool = enable_thinking_params
self.profile_dir: str = profile_dir
self.profile_backend: str = profile_backend
self.profile_store_name: str = profile_store_name
self.profile_max_capacity: int = profile_max_capacity
def _build_tool_call(self) -> ToolCall:
"""Build and return the tool call schema"""
@ -103,6 +110,19 @@ class BaseMemoryTool(BaseTool, metaclass=ABCMeta):
return self.context.service_context.memory_target_type_mapping
@property
def profile_path(self) -> Path:
def profile_path(self) -> Path | None:
"""Get the path to the profile directory for the current collection."""
if not self.profile_dir:
return None
return Path(self.profile_dir) / self.vector_store.collection_name
def get_profile_handler(self, memory_target: str) -> ProfileHandler:
"""Build a profile handler for the current backend configuration."""
return ProfileHandler(
memory_target=memory_target,
profile_path=self.profile_path,
service_context=self.service_context,
profile_backend=self.profile_backend,
profile_store_name=self.profile_store_name,
max_capacity=self.profile_max_capacity,
)

View file

@ -0,0 +1 @@
"""Profile memory tools."""

View file

@ -1,8 +1,7 @@
"""Add draft profile and read all profiles from local storage"""
"""Add draft profile and read all profiles from the configured backend."""
from loguru import logger
from .profile_handler import ProfileHandler
from ..base_memory_tool import BaseMemoryTool
from ....core.schema import ToolCall
@ -92,9 +91,8 @@ class AddDraftAndReadAllProfiles(BaseMemoryTool):
continue
targets_processed.add(target)
profile_handler = ProfileHandler(profile_path=self.profile_path, memory_target=target)
profiles_str = profile_handler.read_all(add_profile_id=True)
profile_handler = self.get_profile_handler(target)
profiles_str = await profile_handler.aread_all(add_profile_id=True)
if profiles_str:
all_profiles.append(f"## Profiles for {target}:\n{profiles_str}")

View file

@ -1,8 +1,7 @@
"""Add user profile tool"""
"""Add user profile tool."""
from loguru import logger
from .profile_handler import ProfileHandler
from ..base_memory_tool import BaseMemoryTool
from ....core.schema import ToolCall
@ -40,7 +39,7 @@ class AddProfile(BaseMemoryTool):
)
async def execute(self):
profile_handler = ProfileHandler(profile_path=self.profile_path, memory_target=self.memory_target)
profile_handler = self.get_profile_handler(self.memory_target)
# Get parameters
message_time = self.context.get("message_time", "")
@ -58,7 +57,7 @@ class AddProfile(BaseMemoryTool):
}
# Add profile using ProfileHandler
new_nodes = profile_handler.add_batch(profiles=[profile], ref_memory_id=self.history_id)
new_nodes = await profile_handler.aadd_batch(profiles=[profile], ref_memory_id=self.history_id)
self.memory_nodes.extend(new_nodes)
if new_nodes:

View file

@ -1,8 +1,7 @@
"""Delete user profile tool"""
"""Delete user profile tool."""
from loguru import logger
from .profile_handler import ProfileHandler
from ..base_memory_tool import BaseMemoryTool
from ....core.schema import ToolCall
@ -32,7 +31,7 @@ class DeleteProfile(BaseMemoryTool):
)
async def execute(self):
profile_handler = ProfileHandler(profile_path=self.profile_path, memory_target=self.memory_target)
profile_handler = self.get_profile_handler(self.memory_target)
# Get profile_id parameter
profile_id = self.context.get("profile_id", "")
@ -41,7 +40,7 @@ class DeleteProfile(BaseMemoryTool):
return "No profile_id provided, operation cancelled."
# Delete profile using ProfileHandler
success = profile_handler.delete(profile_id)
success = await profile_handler.adelete(profile_id)
if success:
output = f"Successfully deleted profile with ID: {profile_id}"

View file

@ -0,0 +1,234 @@
"""Filesystem-backed profile storage."""
from pathlib import Path
from loguru import logger
from .profile_backend import BaseProfileBackend
from ....core.enumeration import MemoryType
from ....core.schema import MemoryNode
from ....core.utils import CacheHandler, deduplicate_memories
class FileProfileBackend(BaseProfileBackend):
"""Persist user profiles in local JSONL cache files."""
def __init__(self, profile_path: str | Path, memory_target: str, max_capacity: int = 50):
super().__init__(memory_target=memory_target, max_capacity=max_capacity)
self.cache_key: str = self.memory_target.replace(" ", "_").lower()
self.cache_handler: CacheHandler = CacheHandler(profile_path)
def _load_nodes(self) -> list[MemoryNode]:
cached_data = self.cache_handler.load(self.cache_key, auto_clean=False)
if not cached_data:
return []
return [MemoryNode(**data) for data in cached_data]
def _save_nodes(self, nodes: list[MemoryNode], apply_limits: bool = True):
if apply_limits:
nodes = deduplicate_memories(nodes)
if len(nodes) > self.max_capacity:
sorted_nodes = sorted(nodes, key=lambda n: n.message_time)
removed_count = len(sorted_nodes) - self.max_capacity
nodes = sorted_nodes[removed_count:]
logger.info(
f"Capacity limit reached: removed {removed_count} oldest profiles "
f"(kept {len(nodes)}/{self.max_capacity})",
)
nodes_data = [node.model_dump(exclude_none=True) for node in nodes]
self.cache_handler.save(self.cache_key, nodes_data)
logger.info(f"Saved {len(nodes)} profiles to {self.cache_key}")
def get_all_sync(self) -> list[MemoryNode]:
"""Load all profile nodes from cache, ordered by ``message_time``."""
nodes = self._load_nodes()
nodes.sort(key=lambda n: n.message_time)
return nodes
def get_by_sync(self, *, profile_id: str | None = None, profile_key: str | None = None) -> MemoryNode | None:
"""Return the first node matching ``profile_id`` or ``profile_key``."""
if not profile_id and not profile_key:
raise ValueError("Must provide either profile_id or profile_key")
for node in self._load_nodes():
if profile_id and node.memory_id == profile_id:
return node
if profile_key and node.when_to_use == profile_key:
return node
return None
def delete_sync(self, profile_id: str | list[str]) -> bool | int:
"""Remove one id, many ids, or none; returns bool, count, or 0/false if nothing removed."""
nodes = self._load_nodes()
original_count = len(nodes)
if isinstance(profile_id, list):
profile_ids_set = set(profile_id)
nodes = [n for n in nodes if n.memory_id not in profile_ids_set]
deleted_count = original_count - len(nodes)
if deleted_count == 0:
logger.warning(f"No profiles found to delete from {len(profile_id)} IDs")
return 0
self._save_nodes(nodes, apply_limits=False)
logger.info(f"Batch deleted {deleted_count} profiles")
return deleted_count
nodes = [n for n in nodes if n.memory_id != profile_id]
if len(nodes) == original_count:
logger.warning(f"Profile {profile_id} not found")
return False
self._save_nodes(nodes, apply_limits=False)
logger.info(f"Deleted profile {profile_id}")
return True
def delete_all_sync(self) -> int:
"""Clear every cached profile for this target; returns how many were stored."""
nodes = self._load_nodes()
count = len(nodes)
self._save_nodes([], apply_limits=False)
logger.info(f"Deleted all {count} profiles")
return count
def add_sync(self, message_time: str, profile_key: str, profile_value: str, ref_memory_id: str = "") -> MemoryNode:
"""Append a profile row, replacing any existing row with the same key."""
nodes = self._load_nodes()
new_node = MemoryNode(
memory_type=MemoryType.PERSONAL,
memory_target=self.memory_target,
when_to_use=profile_key,
content=profile_value,
message_time=message_time,
ref_memory_id=ref_memory_id,
)
original_count = len(nodes)
nodes = [n for n in nodes if n.when_to_use != profile_key]
if len(nodes) < original_count:
logger.info(f"Removed {original_count - len(nodes)} duplicate profile(s) with key: {profile_key}")
nodes.append(new_node)
self._save_nodes(nodes)
logger.info(f"Added profile: {profile_key}={profile_value}")
return new_node
def add_batch_sync(self, profiles: list[dict], ref_memory_id: str = "") -> list[MemoryNode]:
"""Insert many profiles in one write, deduping by key against existing rows."""
if not profiles:
return []
nodes = self._load_nodes()
new_nodes = [
MemoryNode(
memory_type=MemoryType.PERSONAL,
memory_target=self.memory_target,
when_to_use=p.get("profile_key", ""),
content=p.get("profile_value", ""),
message_time=p.get("message_time", ""),
ref_memory_id=ref_memory_id,
)
for p in profiles
]
new_keys = {n.when_to_use for n in new_nodes}
original_count = len(nodes)
nodes = [n for n in nodes if n.when_to_use not in new_keys]
if len(nodes) < original_count:
logger.info(f"Removed {original_count - len(nodes)} duplicate profile(s) with matching keys")
nodes.extend(new_nodes)
self._save_nodes(nodes)
logger.info(f"Batch added {len(new_nodes)} profiles")
return new_nodes
def update_sync(
self,
profile_id: str,
message_time: str,
profile_key: str,
profile_value: str,
) -> MemoryNode | None:
"""Update fields for ``profile_id``; return ``None`` if that id is missing."""
nodes = self._load_nodes()
target_node = None
for node in nodes:
if node.memory_id == profile_id:
node.when_to_use = profile_key
node.content = profile_value
node.message_time = message_time
target_node = node
break
if target_node is None:
logger.warning(f"Profile {profile_id} not found")
return None
self._save_nodes(nodes, apply_limits=False)
logger.info(f"Updated profile {profile_id}: {profile_key}={profile_value}")
return target_node
def search_sync(self, query: str | list[str], limit: int = 5) -> list[MemoryNode]:
"""Simple substring/token match over key and content, best matches first."""
queries = [query] if isinstance(query, str) else query
query_terms = [q.strip().lower() for q in queries if q and q.strip()]
if not query_terms:
return []
scored_nodes = []
for node in self.get_all_sync():
profile_key = str(node.metadata.get("profile_key", node.when_to_use)).lower()
haystack = f"{profile_key}: {node.content}".lower()
score = 0
for term in query_terms:
if term in haystack:
score += len(term) + 10
else:
token_hits = sum(1 for token in term.split() if token and token in haystack)
score += token_hits
if score > 0:
node.score = float(score)
scored_nodes.append(node)
scored_nodes.sort(key=lambda n: (n.score, n.message_time), reverse=True)
return scored_nodes[:limit]
async def get_all(self) -> list[MemoryNode]:
return self.get_all_sync()
async def get_by(self, *, profile_id: str | None = None, profile_key: str | None = None) -> MemoryNode | None:
return self.get_by_sync(profile_id=profile_id, profile_key=profile_key)
async def delete(self, profile_id: str | list[str]) -> bool | int:
return self.delete_sync(profile_id)
async def delete_all(self) -> int:
return self.delete_all_sync()
async def add(
self,
message_time: str,
profile_key: str,
profile_value: str,
ref_memory_id: str = "",
) -> MemoryNode:
return self.add_sync(message_time, profile_key, profile_value, ref_memory_id)
async def add_batch(self, profiles: list[dict], ref_memory_id: str = "") -> list[MemoryNode]:
return self.add_batch_sync(profiles, ref_memory_id)
async def update(
self,
profile_id: str,
message_time: str,
profile_key: str,
profile_value: str,
) -> MemoryNode | None:
return self.update_sync(profile_id, message_time, profile_key, profile_value)
async def search(self, query: str | list[str], limit: int = 5) -> list[MemoryNode]:
return self.search_sync(query, limit)

View file

@ -0,0 +1,51 @@
"""Profile backend abstractions."""
from abc import ABC, abstractmethod
from ....core.schema import MemoryNode
class BaseProfileBackend(ABC):
"""Abstract interface for profile storage backends."""
def __init__(self, memory_target: str, max_capacity: int = 50):
self.memory_target = memory_target
self.max_capacity = max_capacity
@abstractmethod
async def get_all(self) -> list[MemoryNode]:
"""Return all profile rows for the current user."""
@abstractmethod
async def get_by(self, *, profile_id: str | None = None, profile_key: str | None = None) -> MemoryNode | None:
"""Return one profile row by id or key."""
@abstractmethod
async def delete(self, profile_id: str | list[str]) -> bool | int:
"""Delete one or more profile rows."""
@abstractmethod
async def delete_all(self) -> int:
"""Delete all profile rows for the current user."""
@abstractmethod
async def add(self, message_time: str, profile_key: str, profile_value: str, ref_memory_id: str = "") -> MemoryNode:
"""Add a single profile row."""
@abstractmethod
async def add_batch(self, profiles: list[dict], ref_memory_id: str = "") -> list[MemoryNode]:
"""Add multiple profile rows."""
@abstractmethod
async def update(
self,
profile_id: str,
message_time: str,
profile_key: str,
profile_value: str,
) -> MemoryNode | None:
"""Update one profile row."""
@abstractmethod
async def search(self, query: str | list[str], limit: int = 5) -> list[MemoryNode]:
"""Search profile rows relevant to the query."""

View file

@ -1,195 +1,124 @@
"""Profile Handler for managing user profiles in local memory"""
"""Profile handler facade for filesystem and vector backends."""
# pylint: disable=missing-function-docstring
import asyncio
from pathlib import Path
from loguru import logger
from ....core.enumeration import MemoryType
from .file_profile_backend import FileProfileBackend
from .profile_backend import BaseProfileBackend
from .vector_profile_backend import VectorProfileBackend
from ....core import ServiceContext
from ....core.schema import MemoryNode
from ....core.utils import CacheHandler, deduplicate_memories
class ProfileHandler:
"""User profile CRUD handler"""
"""User profile facade with pluggable storage backends."""
def __init__(self, profile_path: str | Path, memory_target: str, max_capacity: int = 50):
"""init"""
self.memory_target: str = memory_target
self.cache_key: str = self.memory_target.replace(" ", "_").lower()
self.cache_handler: CacheHandler = CacheHandler(profile_path)
self.max_capacity: int = max_capacity
def _load_nodes(self) -> list[MemoryNode]:
"""Load profile nodes"""
cached_data = self.cache_handler.load(self.cache_key, auto_clean=False)
if not cached_data:
return []
return [MemoryNode(**data) for data in cached_data]
def _save_nodes(self, nodes: list[MemoryNode], apply_limits: bool = True):
"""Save nodes with optional deduplication and capacity enforcement"""
if apply_limits:
nodes = deduplicate_memories(nodes)
# Enforce capacity limit by removing the oldest profiles
if len(nodes) > self.max_capacity:
sorted_nodes = sorted(nodes, key=lambda n: n.message_time)
removed_count = len(sorted_nodes) - self.max_capacity
nodes = sorted_nodes[removed_count:]
logger.info(
f"Capacity limit reached: removed {removed_count} oldest profiles "
f"(kept {len(nodes)}/{self.max_capacity})",
)
nodes_data = [node.model_dump(exclude_none=True) for node in nodes]
self.cache_handler.save(self.cache_key, nodes_data)
logger.info(f"Saved {len(nodes)} profiles to {self.cache_key}")
def delete(self, profile_id: str | list[str]) -> bool | int:
"""Delete profile by ID(s), returns True/False for single ID or count for batch delete"""
nodes = self._load_nodes()
original_count = len(nodes)
# Batch delete mode
if isinstance(profile_id, list):
profile_ids_set = set(profile_id)
nodes = [n for n in nodes if n.memory_id not in profile_ids_set]
deleted_count = original_count - len(nodes)
if deleted_count == 0:
logger.warning(f"No profiles found to delete from {len(profile_id)} IDs")
return 0
self._save_nodes(nodes, apply_limits=False)
logger.info(f"Batch deleted {deleted_count} profiles")
return deleted_count
# Single delete mode
nodes = [n for n in nodes if n.memory_id != profile_id]
if len(nodes) == original_count:
logger.warning(f"Profile {profile_id} not found")
return False
self._save_nodes(nodes, apply_limits=False)
logger.info(f"Deleted profile {profile_id}")
return True
def delete_all(self) -> int:
"""Delete all profiles, returns count deleted"""
nodes = self._load_nodes()
count = len(nodes)
self._save_nodes([], apply_limits=False)
logger.info(f"Deleted all {count} profiles")
return count
def add(self, message_time: str, profile_key: str, profile_value: str, ref_memory_id: str = "") -> MemoryNode:
"""Add new profile, returns created MemoryNode"""
nodes = self._load_nodes()
new_node = MemoryNode(
memory_type=MemoryType.PERSONAL,
memory_target=self.memory_target,
when_to_use=profile_key,
content=profile_value,
message_time=message_time,
ref_memory_id=ref_memory_id,
def __init__(
self,
memory_target: str,
profile_path: str | Path | None = None,
service_context: ServiceContext | None = None,
profile_backend: str = "filesystem",
profile_store_name: str = "profile",
max_capacity: int = 50,
):
self.memory_target = memory_target
self.profile_backend = profile_backend
self.profile_store_name = profile_store_name
self.max_capacity = max_capacity
self.cache_key = self.memory_target.replace(" ", "_").lower()
self.backend = self._build_backend(
profile_path=profile_path,
service_context=service_context,
)
# Remove existing nodes with the same when_to_use (profile_key)
original_count = len(nodes)
nodes = [n for n in nodes if n.when_to_use != profile_key]
if len(nodes) < original_count:
logger.info(f"Removed {original_count - len(nodes)} duplicate profile(s) with key: {profile_key}")
nodes.append(new_node)
self._save_nodes(nodes)
logger.info(f"Added profile: {profile_key}={profile_value}")
return new_node
def add_batch(self, profiles: list[dict], ref_memory_id: str = "") -> list[MemoryNode]:
"""Add multiple profiles in batch, returns list of created MemoryNodes"""
if not profiles:
return []
nodes = self._load_nodes()
new_nodes = [
MemoryNode(
memory_type=MemoryType.PERSONAL,
def _build_backend(
self,
profile_path: str | Path | None,
service_context: ServiceContext | None,
) -> BaseProfileBackend:
if self.profile_backend == "filesystem":
if profile_path is None:
raise ValueError("profile_path is required for filesystem profile backend")
return FileProfileBackend(
profile_path=profile_path,
memory_target=self.memory_target,
when_to_use=p.get("profile_key", ""),
content=p.get("profile_value", ""),
message_time=p.get("message_time", ""),
ref_memory_id=ref_memory_id,
max_capacity=self.max_capacity,
)
for p in profiles
]
# Remove existing nodes with the same when_to_use (profile_key)
new_keys = {n.when_to_use for n in new_nodes}
original_count = len(nodes)
nodes = [n for n in nodes if n.when_to_use not in new_keys]
if len(nodes) < original_count:
logger.info(f"Removed {original_count - len(nodes)} duplicate profile(s) with matching keys")
if self.profile_backend == "vector":
if service_context is None:
raise ValueError("service_context is required for vector profile backend")
return VectorProfileBackend(
memory_target=self.memory_target,
service_context=service_context,
vector_store_name=self.profile_store_name,
max_capacity=self.max_capacity,
)
nodes.extend(new_nodes)
self._save_nodes(nodes)
logger.info(f"Batch added {len(new_nodes)} profiles")
return new_nodes
def update(self, profile_id: str, message_time: str, profile_key: str, profile_value: str) -> MemoryNode | None:
"""Update profile by ID, returns updated node or None if not found"""
nodes = self._load_nodes()
target_node = None
for node in nodes:
if node.memory_id == profile_id:
node.when_to_use = profile_key
node.content = profile_value
node.message_time = message_time
target_node = node
break
if target_node is None:
logger.warning(f"Profile {profile_id} not found")
return None
self._save_nodes(nodes, apply_limits=False)
logger.info(f"Updated profile {profile_id}: {profile_key}={profile_value}")
return target_node
def get_by(self, *, profile_id: str | None = None, profile_key: str | None = None) -> MemoryNode | None:
"""Get profile by ID or key"""
if not profile_id and not profile_key:
raise ValueError("Must provide either profile_id or profile_key")
nodes = self._load_nodes()
for node in nodes:
if profile_id and node.memory_id == profile_id:
return node
if profile_key and node.when_to_use == profile_key:
return node
return None
def get_by_id(self, profile_id: str) -> MemoryNode | None:
"""Get profile by ID (convenience method)"""
return self.get_by(profile_id=profile_id)
def get_by_key(self, profile_key: str) -> MemoryNode | None:
"""Get profile by key (convenience method)"""
return self.get_by(profile_key=profile_key)
def get_all(self) -> list[MemoryNode]:
"""Get all profiles, sorted by message_time"""
nodes = self._load_nodes()
nodes.sort(key=lambda n: n.message_time)
return nodes
raise ValueError(f"Unsupported profile backend: {self.profile_backend}")
@staticmethod
def _format_node(node: MemoryNode, add_profile_id: bool = False, add_history_id: bool = False) -> str:
"""Format a single node to string"""
def _run_sync(coro):
try:
asyncio.get_running_loop()
except RuntimeError:
return asyncio.run(coro)
raise RuntimeError(
"Synchronous profile access is not available in an active event loop. Use async methods instead.",
)
async def adelete(self, profile_id: str | list[str]) -> bool | int:
return await self.backend.delete(profile_id)
async def adelete_all(self) -> int:
return await self.backend.delete_all()
async def aadd(
self,
message_time: str,
profile_key: str,
profile_value: str,
ref_memory_id: str = "",
) -> MemoryNode:
return await self.backend.add(message_time, profile_key, profile_value, ref_memory_id)
async def aadd_batch(self, profiles: list[dict], ref_memory_id: str = "") -> list[MemoryNode]:
return await self.backend.add_batch(profiles, ref_memory_id)
async def aupdate(
self,
profile_id: str,
message_time: str,
profile_key: str,
profile_value: str,
) -> MemoryNode | None:
return await self.backend.update(profile_id, message_time, profile_key, profile_value)
async def aget_by(self, *, profile_id: str | None = None, profile_key: str | None = None) -> MemoryNode | None:
return await self.backend.get_by(profile_id=profile_id, profile_key=profile_key)
async def aget_by_id(self, profile_id: str) -> MemoryNode | None:
return await self.aget_by(profile_id=profile_id)
async def aget_by_key(self, profile_key: str) -> MemoryNode | None:
return await self.aget_by(profile_key=profile_key)
async def aget_all(self) -> list[MemoryNode]:
return await self.backend.get_all()
async def asearch(self, query: str | list[str], limit: int = 5) -> list[MemoryNode]:
return await self.backend.search(query=query, limit=limit)
@staticmethod
def format_node(node: MemoryNode, add_profile_id: bool = False, add_history_id: bool = False) -> str:
"""Render a profile ``MemoryNode`` as a single-line string for tools/logs."""
parts = []
profile_key = str(node.metadata.get("profile_key", node.when_to_use))
if add_profile_id:
parts.append(f"profile_id={node.memory_id}")
@ -197,16 +126,70 @@ class ProfileHandler:
if node.message_time:
parts.append(f"[{node.message_time}]")
parts.append(f"{node.when_to_use}: {node.content}")
parts.append(f"{profile_key}: {node.content}")
if add_history_id:
if add_history_id and node.ref_memory_id:
parts.append(f"history_id={node.ref_memory_id}")
return " ".join(parts)
def read_all(self, add_profile_id: bool = False, add_history_id: bool = False) -> str:
"""Read all profiles and return formatted string"""
nodes = self.get_all()
formatted_profiles = [self._format_node(node, add_profile_id, add_history_id) for node in nodes]
async def aread_all(self, add_profile_id: bool = False, add_history_id: bool = False) -> str:
nodes = await self.aget_all()
formatted_profiles = [self.format_node(node, add_profile_id, add_history_id) for node in nodes]
logger.info(f"Read {len(formatted_profiles)} profiles from {self.cache_key}")
return "\n".join(formatted_profiles).strip()
async def aretrieve(
self,
query: str | list[str],
limit: int = 5,
add_profile_id: bool = True,
add_history_id: bool = False,
) -> tuple[list[MemoryNode], str]:
nodes = await self.asearch(query=query, limit=limit)
formatted_profiles = [self.format_node(node, add_profile_id, add_history_id) for node in nodes]
return nodes, "\n".join(formatted_profiles).strip()
def delete(self, profile_id: str | list[str]) -> bool | int:
if isinstance(self.backend, FileProfileBackend):
return self.backend.delete_sync(profile_id)
return self._run_sync(self.adelete(profile_id))
def delete_all(self) -> int:
if isinstance(self.backend, FileProfileBackend):
return self.backend.delete_all_sync()
return self._run_sync(self.adelete_all())
def add(self, message_time: str, profile_key: str, profile_value: str, ref_memory_id: str = "") -> MemoryNode:
if isinstance(self.backend, FileProfileBackend):
return self.backend.add_sync(message_time, profile_key, profile_value, ref_memory_id)
return self._run_sync(self.aadd(message_time, profile_key, profile_value, ref_memory_id))
def add_batch(self, profiles: list[dict], ref_memory_id: str = "") -> list[MemoryNode]:
if isinstance(self.backend, FileProfileBackend):
return self.backend.add_batch_sync(profiles, ref_memory_id)
return self._run_sync(self.aadd_batch(profiles, ref_memory_id))
def update(self, profile_id: str, message_time: str, profile_key: str, profile_value: str) -> MemoryNode | None:
if isinstance(self.backend, FileProfileBackend):
return self.backend.update_sync(profile_id, message_time, profile_key, profile_value)
return self._run_sync(self.aupdate(profile_id, message_time, profile_key, profile_value))
def get_by(self, *, profile_id: str | None = None, profile_key: str | None = None) -> MemoryNode | None:
if isinstance(self.backend, FileProfileBackend):
return self.backend.get_by_sync(profile_id=profile_id, profile_key=profile_key)
return self._run_sync(self.aget_by(profile_id=profile_id, profile_key=profile_key))
def get_by_id(self, profile_id: str) -> MemoryNode | None:
return self._run_sync(self.aget_by_id(profile_id))
def get_by_key(self, profile_key: str) -> MemoryNode | None:
return self._run_sync(self.aget_by_key(profile_key))
def get_all(self) -> list[MemoryNode]:
if isinstance(self.backend, FileProfileBackend):
return self.backend.get_all_sync()
return self._run_sync(self.aget_all())
def read_all(self, add_profile_id: bool = False, add_history_id: bool = False) -> str:
return self._run_sync(self.aread_all(add_profile_id, add_history_id))

View file

@ -0,0 +1,245 @@
"""Vector-backed handler for bounded user profiles."""
import hashlib
from loguru import logger
from ....core import ServiceContext
from ....core.enumeration import MemoryType
from ....core.schema import MemoryNode
from ....core.vector_store import BaseVectorStore
class ProfileVectorHandler:
"""Manage profile rows stored in a dedicated vector collection."""
PROFILE_KIND = "profile"
def __init__(
self,
memory_target: str,
service_context: ServiceContext,
vector_store_name: str = "profile",
max_capacity: int = 50,
):
self.memory_target = memory_target
self.service_context = service_context
self.vector_store_name = vector_store_name
self.max_capacity = max_capacity
self.vector_store: BaseVectorStore = service_context.vector_stores[vector_store_name]
@staticmethod
def build_retrieval_text(profile_key: str, profile_value: str) -> str:
"""Build the text that will be embedded for semantic profile retrieval."""
return f"{profile_key}: {profile_value}".strip(": ")
def build_profile_id(self, profile_key: str) -> str:
"""Build a stable id from user and key."""
hash_obj = hashlib.sha256(f"{self.memory_target}\n{profile_key}".encode("utf-8"))
return hash_obj.hexdigest()[:16]
def _base_filters(self) -> dict:
"""Filters shared by all profile rows in the vector collection."""
return {
"memory_type": MemoryType.IDENTITY.value,
"memory_target": self.memory_target,
"profile_kind": self.PROFILE_KIND,
}
def _build_profile_node(self, profile: dict, ref_memory_id: str = "") -> MemoryNode:
"""Turn a profile dict into a ``MemoryNode`` for upsert into the vector store."""
profile_key = profile.get("profile_key", "").strip()
profile_value = profile.get("profile_value", "").strip()
message_time = profile.get("message_time", "")
ref_id = profile.get("ref_memory_id", ref_memory_id)
metadata = dict(profile.get("metadata", {}))
metadata.update(
{
"profile_key": profile_key,
"profile_kind": self.PROFILE_KIND,
"profile_backend": "vector",
},
)
return MemoryNode(
memory_id=self.build_profile_id(profile_key),
memory_type=MemoryType.IDENTITY,
memory_target=self.memory_target,
when_to_use=self.build_retrieval_text(profile_key, profile_value),
content=profile_value,
message_time=message_time,
ref_memory_id=ref_id,
metadata=metadata,
)
def _vector_profile_matches(self, memory_node: MemoryNode) -> bool:
"""True if ``memory_node`` belongs to this handler's target and profile kind."""
if memory_node.memory_target != self.memory_target:
return False
if memory_node.memory_type is not MemoryType.IDENTITY:
return False
if memory_node.metadata.get("profile_kind") != self.PROFILE_KIND:
return False
return True
async def _get_by_profile_id(self, profile_id: str) -> MemoryNode | None:
"""Load by vector id and validate filters."""
try:
vector_node = await self.vector_store.get(profile_id)
except KeyError:
logger.warning(f"Profile {profile_id} not found in vector store")
return None
if vector_node is None:
logger.warning(f"Profile {profile_id} not found in vector store")
return None
memory_node = MemoryNode.from_vector_node(vector_node)
if not self._vector_profile_matches(memory_node):
return None
return memory_node
async def _get_by_profile_key(self, profile_key: str) -> MemoryNode | None:
"""Load the single row matching ``profile_key`` under base filters."""
filters = {**self._base_filters(), "profile_key": profile_key}
vector_nodes = await self.vector_store.list(filters=filters, limit=1)
if not vector_nodes:
return None
return MemoryNode.from_vector_node(vector_nodes[0])
async def get_all(self) -> list[MemoryNode]:
"""List every profile row for this memory target, sorted by store."""
vector_nodes = await self.vector_store.list(
filters=self._base_filters(),
sort_key="message_time",
reverse=False,
)
return [MemoryNode.from_vector_node(node) for node in vector_nodes]
async def get_by(self, *, profile_id: str | None = None, profile_key: str | None = None) -> MemoryNode | None:
"""Return one profile by stable id or by logical profile key."""
if not profile_id and not profile_key:
raise ValueError("Must provide either profile_id or profile_key")
if profile_id:
return await self._get_by_profile_id(profile_id)
return await self._get_by_profile_key(profile_key or "")
async def delete(self, profile_id: str | list[str]) -> bool | int:
"""Delete one id, many ids, or report zero/false when nothing matched."""
if isinstance(profile_id, list):
profile_ids = list(dict.fromkeys(pid for pid in profile_id if pid))
if not profile_ids:
return 0
existing_nodes = []
for pid in profile_ids:
node = await self.get_by(profile_id=pid)
if node is not None:
existing_nodes.append(node)
if not existing_nodes:
return 0
await self.vector_store.delete([node.memory_id for node in existing_nodes])
return len(existing_nodes)
existing_node = await self.get_by(profile_id=profile_id)
if existing_node is None:
return False
await self.vector_store.delete(existing_node.memory_id)
return True
async def delete_all(self) -> int:
"""Remove all profile vectors for this target; returns how many were deleted."""
nodes = await self.get_all()
if not nodes:
return 0
await self.vector_store.delete([node.memory_id for node in nodes])
return len(nodes)
async def add_batch(self, profiles: list[dict], ref_memory_id: str = "") -> list[MemoryNode]:
"""Upsert many profiles at once (last dict wins per key), then enforce capacity."""
if not profiles:
return []
deduped_profiles: dict[str, dict] = {}
for profile in profiles:
profile_key = profile.get("profile_key", "").strip()
if not profile_key:
continue
deduped_profiles[profile_key] = profile
new_nodes = [
self._build_profile_node(profile, ref_memory_id=ref_memory_id) for profile in deduped_profiles.values()
]
if not new_nodes:
return []
await self.vector_store.delete([node.memory_id for node in new_nodes])
await self.vector_store.insert([node.to_vector_node() for node in new_nodes])
await self.enforce_capacity()
return new_nodes
async def add(self, message_time: str, profile_key: str, profile_value: str, ref_memory_id: str = "") -> MemoryNode:
"""Insert or replace a single profile row."""
nodes = await self.add_batch(
[
{
"message_time": message_time,
"profile_key": profile_key,
"profile_value": profile_value,
},
],
ref_memory_id=ref_memory_id,
)
return nodes[0]
async def update(
self,
profile_id: str,
message_time: str,
profile_key: str,
profile_value: str,
) -> MemoryNode | None:
"""Replace content and key for ``profile_id``; return ``None`` if missing."""
existing_node = await self.get_by(profile_id=profile_id)
if existing_node is None:
return None
new_node = self._build_profile_node(
{
"message_time": message_time,
"profile_key": profile_key,
"profile_value": profile_value,
"ref_memory_id": existing_node.ref_memory_id,
"metadata": existing_node.metadata,
},
)
if existing_node.memory_id != new_node.memory_id:
await self.vector_store.delete(existing_node.memory_id)
else:
await self.vector_store.delete(new_node.memory_id)
await self.vector_store.insert(new_node.to_vector_node())
await self.enforce_capacity()
return new_node
async def search(self, query: str | list[str], limit: int = 5) -> list[MemoryNode]:
"""Semantic search with de-duplication across multiple query strings."""
queries = [query] if isinstance(query, str) else query
seen_nodes: dict[str, MemoryNode] = {}
for item in queries:
if not item or not item.strip():
continue
vector_nodes = await self.vector_store.search(item, limit=limit, filters=self._base_filters())
for vector_node in vector_nodes:
memory_node = MemoryNode.from_vector_node(vector_node)
seen_nodes[memory_node.memory_id] = memory_node
nodes = list(seen_nodes.values())
nodes.sort(key=lambda node: (node.score, node.message_time), reverse=True)
return nodes[:limit]
async def enforce_capacity(self):
"""Drop oldest rows when count exceeds ``max_capacity``."""
nodes = await self.get_all()
overflow = len(nodes) - self.max_capacity
if overflow <= 0:
return
to_delete = [node.memory_id for node in nodes[:overflow]]
await self.vector_store.delete(to_delete)

View file

@ -1,8 +1,7 @@
"""Read user profile tool"""
"""Read user profile tool."""
from loguru import logger
from .profile_handler import ProfileHandler
from ..base_memory_tool import BaseMemoryTool
from ....core.schema import ToolCall
@ -44,8 +43,8 @@ class ReadAllProfiles(BaseMemoryTool):
else:
target = self.memory_target
profile_handler = ProfileHandler(profile_path=self.profile_path, memory_target=target)
profiles_str = profile_handler.read_all(add_profile_id=True)
profile_handler = self.get_profile_handler(target)
profiles_str = await profile_handler.aread_all(add_profile_id=True)
if not profiles_str:
output = "No profiles found."
logger.info(output)

View file

@ -0,0 +1,102 @@
"""Retrieve relevant profile rows."""
from loguru import logger
from .profile_handler import ProfileHandler
from ..base_memory_tool import BaseMemoryTool
from ....core.schema import MemoryNode, ToolCall
class RetrieveProfile(BaseMemoryTool):
"""Tool to retrieve relevant profiles using the configured backend."""
def __init__(self, top_k: int = 5, enable_memory_target: bool = False, **kwargs):
super().__init__(**kwargs)
self.top_k = top_k
self.enable_memory_target = enable_memory_target
def _build_query_parameters(self) -> dict:
properties = {
"query": {
"type": "string",
"description": "query",
},
}
required = ["query"]
if self.enable_memory_target:
properties["memory_target"] = {
"type": "string",
"description": "memory_target",
}
required.append("memory_target")
return {
"type": "object",
"properties": properties,
"required": required,
}
def _build_tool_call(self) -> ToolCall:
return ToolCall(
**{
"description": "Retrieve relevant user profiles using semantic matching.",
"parameters": self._build_query_parameters(),
},
)
def _build_multiple_tool_call(self) -> ToolCall:
return ToolCall(
**{
"description": "Retrieve relevant user profiles using semantic matching.",
"parameters": {
"type": "object",
"properties": {
"query_items": {
"type": "array",
"description": "List of query items.",
"items": self._build_query_parameters(),
},
},
"required": ["query_items"],
},
},
)
async def execute(self):
if self.enable_multiple:
query_items = self.context.get("query_items", [])
else:
query_items = [self.context]
queries_by_target: dict[str, list[str]] = {}
for item in query_items:
target = item["memory_target"] if self.enable_memory_target else self.memory_target
queries_by_target.setdefault(target, []).append(item["query"])
profile_nodes: list[MemoryNode] = []
for target, queries in queries_by_target.items():
profile_handler = self.get_profile_handler(target)
nodes, _ = await profile_handler.aretrieve(
query=queries,
limit=self.top_k,
add_profile_id=True,
add_history_id=True,
)
profile_nodes.extend(nodes)
seen_ids = {node.memory_id: node for node in self.retrieved_nodes if node.memory_id}
new_nodes = []
for node in profile_nodes:
if node.memory_id not in seen_ids:
seen_ids[node.memory_id] = node
new_nodes.append(node)
self.retrieved_nodes.extend(new_nodes)
if not new_nodes:
output = "No new profiles found."
else:
output = "\n".join(
[ProfileHandler.format_node(node, add_profile_id=True, add_history_id=True) for node in new_nodes],
)
logger.info(f"Retrieved {len(profile_nodes)} profiles, {len(new_nodes)} new after deduplication")
return output

View file

@ -1,8 +1,7 @@
"""Update user profile tool"""
"""Update user profile tool."""
from loguru import logger
from .profile_handler import ProfileHandler
from ..base_memory_tool import BaseMemoryTool
from ....core.schema import ToolCall
@ -82,8 +81,8 @@ class UpdateProfile(BaseMemoryTool):
# Delete profiles (using self.memory_target)
if profile_ids_to_delete:
profile_handler = ProfileHandler(profile_path=self.profile_path, memory_target=self.memory_target)
removed_count = profile_handler.delete(profile_ids_to_delete)
profile_handler = self.get_profile_handler(self.memory_target)
removed_count = await profile_handler.adelete(profile_ids_to_delete)
# Add new profiles
if profiles_to_add:
@ -98,14 +97,17 @@ class UpdateProfile(BaseMemoryTool):
# Add profiles for each target
for target, target_profiles in profiles_by_target.items():
profile_handler = ProfileHandler(profile_path=self.profile_path, memory_target=target)
new_nodes = profile_handler.add_batch(profiles=target_profiles, ref_memory_id=self.history_id)
profile_handler = self.get_profile_handler(target)
new_nodes = await profile_handler.aadd_batch(
profiles=target_profiles,
ref_memory_id=self.history_id,
)
self.memory_nodes.extend(new_nodes)
added_count += len(new_nodes)
else:
# Use self.memory_target for all profiles
profile_handler = ProfileHandler(profile_path=self.profile_path, memory_target=self.memory_target)
new_nodes = profile_handler.add_batch(profiles=profiles_to_add, ref_memory_id=self.history_id)
profile_handler = self.get_profile_handler(self.memory_target)
new_nodes = await profile_handler.aadd_batch(profiles=profiles_to_add, ref_memory_id=self.history_id)
self.memory_nodes.extend(new_nodes)
added_count = len(new_nodes)

View file

@ -1,8 +1,7 @@
"""Update user profile tool"""
"""Update user profile tool."""
from loguru import logger
from .profile_handler import ProfileHandler
from ..base_memory_tool import BaseMemoryTool
from ....core.schema import ToolCall
@ -113,8 +112,8 @@ class UpdateProfilesV1(BaseMemoryTool):
for target, profile_ids in delete_by_target.items():
if profile_ids:
profile_ids = sorted(set(profile_ids)) # Remove duplicates and sort
profile_handler = ProfileHandler(profile_path=self.profile_path, memory_target=target)
profile_handler.delete(profile_ids)
profile_handler = self.get_profile_handler(target)
await profile_handler.adelete(profile_ids)
# Step 2: Prepare all profiles to add (both updated and new)
all_profiles_to_add = []
@ -158,8 +157,8 @@ class UpdateProfilesV1(BaseMemoryTool):
added_count = len(profiles_to_add)
for target, target_profiles in profiles_by_target.items():
profile_handler = ProfileHandler(profile_path=self.profile_path, memory_target=target)
new_nodes = profile_handler.add_batch(profiles=target_profiles, ref_memory_id=self.history_id)
profile_handler = self.get_profile_handler(target)
new_nodes = await profile_handler.aadd_batch(profiles=target_profiles, ref_memory_id=self.history_id)
all_memory_nodes.extend(new_nodes)
# Extend memory_nodes for tracking

View file

@ -0,0 +1,55 @@
"""Vector-backed profile storage."""
from .profile_backend import BaseProfileBackend
from .profile_vector_handler import ProfileVectorHandler
from ....core import ServiceContext
from ....core.schema import MemoryNode
class VectorProfileBackend(BaseProfileBackend):
"""Persist user profiles in a dedicated vector store."""
def __init__(
self,
memory_target: str,
service_context: ServiceContext,
vector_store_name: str = "profile",
max_capacity: int = 50,
):
super().__init__(memory_target=memory_target, max_capacity=max_capacity)
self.handler = ProfileVectorHandler(
memory_target=memory_target,
service_context=service_context,
vector_store_name=vector_store_name,
max_capacity=max_capacity,
)
async def get_all(self) -> list[MemoryNode]:
return await self.handler.get_all()
async def get_by(self, *, profile_id: str | None = None, profile_key: str | None = None) -> MemoryNode | None:
return await self.handler.get_by(profile_id=profile_id, profile_key=profile_key)
async def delete(self, profile_id: str | list[str]) -> bool | int:
return await self.handler.delete(profile_id)
async def delete_all(self) -> int:
return await self.handler.delete_all()
async def add(self, message_time: str, profile_key: str, profile_value: str, ref_memory_id: str = "") -> MemoryNode:
return await self.handler.add(message_time, profile_key, profile_value, ref_memory_id)
async def add_batch(self, profiles: list[dict], ref_memory_id: str = "") -> list[MemoryNode]:
return await self.handler.add_batch(profiles, ref_memory_id)
async def update(
self,
profile_id: str,
message_time: str,
profile_key: str,
profile_value: str,
) -> MemoryNode | None:
return await self.handler.update(profile_id, message_time, profile_key, profile_value)
async def search(self, query: str | list[str], limit: int = 5) -> list[MemoryNode]:
return await self.handler.search(query, limit)

View file

@ -14,6 +14,7 @@ from .memory.vector_tools import (
DelegateTask,
ReadAllProfiles,
ReadHistory,
RetrieveProfile,
RetrieveMemory,
UpdateProfilesV1,
)
@ -55,6 +56,10 @@ class ReMe(Application):
target_task_names: list[str] | None = None,
target_tool_names: list[str] | None = None,
enable_profile: bool = True,
profile_backend: str = "filesystem",
profile_store_name: str = "profile",
profile_collection_name: str | None = None,
profile_max_capacity: int = 50,
**kwargs,
):
"""Initialize ReMe with config.
@ -69,8 +74,37 @@ class ReMe(Application):
```
Args:
*args: Positional arguments forwarded to the base `Application`.
llm_api_key: API key used by the default LLM backend when provided.
llm_base_url: Base URL used by the default LLM backend when provided.
embedding_api_key: API key used by the default embedding backend when provided.
embedding_base_url: Base URL used by the default embedding backend when provided.
working_dir: Directory for generated config, logs, caches, and local stores.
config_path: Built-in config name or config file path used to initialize services.
enable_logo: Whether to print the ReMe logo during startup.
log_to_console: Whether to emit logs to the console.
log_to_file: Whether to write logs under `working_dir`.
default_llm_config: Overrides for the default LLM configuration.
default_embedding_model_config: Overrides for the default embedding model configuration.
default_vector_store_config: Configuration for the default memory vector store.
Its `collection_name` is used for normal memory storage.
default_token_counter_config: Overrides for the default token counter configuration.
target_user_names: Personal memory targets to register at initialization.
target_task_names: Procedural memory targets to register at initialization.
target_tool_names: Tool memory targets to register at initialization.
enable_profile: Whether to enable profile functionality. Set to False when using
cloud-based vector stores to avoid local file operations. Default is True.
profile-free memory flows.
profile_backend: Profile storage backend. Use "filesystem" for local JSONL profile
files or "vector" for a dedicated profile vector collection.
profile_store_name: Internal vector store key used to register and look up the
profile vector store in `service_context.vector_stores`. This is not the
database collection name.
profile_collection_name: Dedicated database collection/table name for vector
profiles. When unset, vector profiles use the default memory collection name
with a "_profile" suffix.
profile_max_capacity: Maximum number of profile rows to keep per memory target.
When the limit is exceeded, the oldest profile rows are removed.
**kwargs: Additional keyword arguments forwarded to the base `Application`.
"""
super().__init__(
*args,
@ -92,6 +126,10 @@ class ReMe(Application):
)
self.enable_profile = enable_profile
self.profile_backend = profile_backend
self.profile_store_name = profile_store_name
self.profile_collection_name = profile_collection_name
self.profile_max_capacity = profile_max_capacity
memory_target_type_mapping: dict[str, MemoryType] = {}
if target_user_names:
@ -111,13 +149,16 @@ class ReMe(Application):
self.service_context.memory_target_type_mapping = memory_target_type_mapping
if self.enable_profile:
if self.enable_profile and self.profile_backend == "filesystem":
profile_path = Path(self.service_context.service_config.working_dir) / "profile"
profile_path.mkdir(parents=True, exist_ok=True)
self.profile_dir: str = str(profile_path)
else:
self.profile_dir: str = ""
if self.enable_profile and self.profile_backend == "vector":
self._ensure_profile_vector_store_config()
def _add_meta_memory(self, memory_type: str | MemoryType, memory_target: str):
"""Register or validate a memory target with the given memory type."""
if memory_target in self.service_context.memory_target_type_mapping:
@ -186,6 +227,38 @@ class ReMe(Application):
return result
return result["answer"]
def _ensure_profile_vector_store_config(self) -> None:
"""Ensure the dedicated profile vector store exists in service config."""
vector_store_configs = self.service_context.service_config.vector_stores
if "default" not in vector_store_configs:
raise RuntimeError("Vector profile backend requires a default vector store configuration")
default_config = vector_store_configs["default"]
profile_collection_name = self.profile_collection_name or f"{default_config.collection_name}_profile"
if self.profile_store_name in vector_store_configs:
if self.profile_collection_name:
vector_store_configs[self.profile_store_name] = vector_store_configs[
self.profile_store_name
].model_copy(
update={"collection_name": profile_collection_name},
)
return
vector_store_configs[self.profile_store_name] = default_config.model_copy(
update={"collection_name": profile_collection_name},
)
def _get_profile_tool_kwargs(self, raise_exception: bool) -> dict:
"""Shared profile tool configuration."""
return {
"profile_dir": self.profile_dir,
"profile_backend": self.profile_backend,
"profile_store_name": self.profile_store_name,
"profile_max_capacity": self.profile_max_capacity,
"raise_exception": raise_exception,
}
async def summarize_memory(
self,
messages: list[Message | dict],
@ -211,6 +284,7 @@ class ReMe(Application):
format_messages.append(message)
if version == "default":
profile_tool_kwargs = self._get_profile_tool_kwargs(raise_exception)
personal_summarizer_tools: list = [
AddDraftAndRetrieveSimilarMemory(
enable_thinking_params=enable_thinking_params,
@ -229,20 +303,28 @@ class ReMe(Application):
),
]
if self.enable_profile:
if self.profile_backend == "vector":
profile_context_tool = RetrieveProfile(
top_k=min(5, retrieve_top_k),
enable_thinking_params=False,
enable_memory_target=False,
enable_multiple=False,
**profile_tool_kwargs,
)
else:
profile_context_tool = ReadAllProfiles(
enable_thinking_params=False,
enable_memory_target=False,
**profile_tool_kwargs,
)
personal_summarizer_tools.extend(
[
ReadAllProfiles(
enable_thinking_params=False,
enable_memory_target=False,
profile_dir=self.profile_dir,
raise_exception=raise_exception,
),
profile_context_tool,
UpdateProfilesV1(
enable_thinking_params=enable_thinking_params,
enable_memory_target=False,
enable_multiple=True,
profile_dir=self.profile_dir,
raise_exception=raise_exception,
**profile_tool_kwargs,
),
],
)
@ -381,16 +463,24 @@ class ReMe(Application):
self._ensure_started()
if version == "default":
profile_tool_kwargs = self._get_profile_tool_kwargs(raise_exception)
personal_retriever_tools = []
if self.enable_profile:
personal_retriever_tools.append(
ReadAllProfiles(
if self.profile_backend == "vector":
profile_context_tool = RetrieveProfile(
top_k=min(5, retrieve_top_k),
enable_thinking_params=False,
enable_memory_target=False,
profile_dir=self.profile_dir,
raise_exception=raise_exception,
),
)
enable_multiple=False,
**profile_tool_kwargs,
)
else:
profile_context_tool = ReadAllProfiles(
enable_thinking_params=False,
enable_memory_target=False,
**profile_tool_kwargs,
)
personal_retriever_tools.append(profile_context_tool)
personal_retriever_tools.extend(
[
RetrieveMemory(
@ -509,6 +599,34 @@ class ReMe(Application):
return self._unwrap_memory_result(result, "retrieve_memory", return_dict)
async def retrieve_profile(
self,
query: str | list[str],
user_name: str,
top_k: int = 5,
return_dict: bool = False,
) -> str | dict:
"""Retrieve relevant profile rows for a user."""
self._ensure_started()
if not self.enable_profile:
raise RuntimeError("Profile functionality is disabled.")
profile_handler = self.get_profile_handler(user_name)
if profile_handler is None:
raise RuntimeError("Profile functionality is disabled.")
retrieved_nodes, output = await profile_handler.aretrieve(
query=query,
limit=top_k,
add_profile_id=True,
add_history_id=True,
)
result = {
"answer": output or "No matching profiles found.",
"retrieved_nodes": retrieved_nodes,
}
return self._unwrap_memory_result(result, "retrieve_profile", return_dict)
async def add_memory(
self,
memory_content: str,
@ -675,15 +793,23 @@ class ReMe(Application):
@property
def profile_path(self) -> Path | None:
"""Get the path to the profile directory. Returns None if profile is disabled."""
if not self.enable_profile:
if not self.enable_profile or self.profile_backend != "filesystem":
return None
return Path(self.profile_dir) / self.default_vector_store.collection_name
collection_name = self.service_context.service_config.vector_stores["default"].collection_name
return Path(self.profile_dir) / collection_name
def get_profile_handler(self, user_name: str) -> ProfileHandler | None:
"""Get the profile handler for the specified user. Returns None if profile is disabled."""
if not self.enable_profile:
return None
return ProfileHandler(memory_target=user_name, profile_path=self.profile_path)
return ProfileHandler(
memory_target=user_name,
profile_path=self.profile_path,
service_context=self.service_context,
profile_backend=self.profile_backend,
profile_store_name=self.profile_store_name,
max_capacity=self.profile_max_capacity,
)
def main():

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

@ -1,9 +1,14 @@
"""Tests for ReMe memory error handling and raise_exception propagation."""
from types import SimpleNamespace
import pytest
import reme.reme as reme_module
from reme import ReMe
from reme.core.runtime_context import RuntimeContext
from reme.core.schema import MemoryNode
from reme.memory.vector_tools.history.read_history import ReadHistory
from reme.reme import ReMe
class Recorder:
@ -130,3 +135,34 @@ async def test_summarize_memory_raises_runtime_error_for_unstructured_result(mon
messages=[{"role": "user", "content": "hi", "time_created": "2026-03-20 10:00:00"}],
task_name="demo-task",
)
@pytest.mark.asyncio
async def test_read_history_accepts_single_history_id_in_multiple_mode():
"""Verify multiple-mode history lookup accepts a single history_id string."""
class FakeVectorStore:
"""Minimal vector store stub for ReadHistory tests."""
async def get(self, vector_ids):
"""Return the requested history node."""
assert vector_ids == ["history_123"]
node = MemoryNode(
memory_id="history_123",
memory_type="history",
memory_target="alice",
content="Alice said hello.",
)
return [node.to_vector_node()]
tool = ReadHistory(enable_multiple=True)
tool._vector_store = FakeVectorStore() # pylint: disable=protected-access
tool.context = RuntimeContext(
history_id="history_123",
retrieved_nodes=[],
service_context=SimpleNamespace(memory_target_type_mapping={"alice": "personal"}),
)
result = await tool.execute()
assert "Historical Dialogue[history_123]" in result

View file

@ -2,8 +2,8 @@
"""Unified test suite for vector store implementations.
This module provides comprehensive test coverage for LocalVectorStore, ESVectorStore,
PGVectorStore, QdrantVectorStore, ChromaVectorStore, ObVecVectorStore, and HologresVectorStore
implementations. Tests can be run for specific vector stores or all implementations.
PGVectorStore, QdrantVectorStore, ChromaVectorStore, ObVecVectorStore, HologresVectorStore, and
ZvecVectorStore implementations. Tests can be run for specific vector stores or all implementations.
Usage:
python test_vector_store.py --local # Test LocalVectorStore only
@ -13,6 +13,7 @@ Usage:
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 --hologres # Test HologresVectorStore only
python test_vector_store.py --zvec # Test ZvecVectorStore only
python test_vector_store.py --all # Test all vector stores
"""
@ -38,6 +39,7 @@ from reme.core.vector_store import (
ObVecVectorStore,
PGVectorStore,
QdrantVectorStore,
ZvecVectorStore,
)
load_env()
@ -105,6 +107,8 @@ class TestConfig:
HOLOGRES_SCHEMA = os.environ.get("HOLOGRES_SCHEMA", "public")
HOLOGRES_MIN_SIZE = 1
HOLOGRES_MAX_SIZE = 5
# ZvecVectorStore settings
ZVEC_PATH = "./test_vector_store_zvec" # For local persistent mode
# Embedding model settings
EMBEDDING_MODEL_NAME = "text-embedding-v4"
@ -208,17 +212,7 @@ class SampleDataGenerator:
# ==================== Vector Store Factory ====================
_STORE_TYPE_MAP = {
LocalVectorStore: "local",
QdrantVectorStore: "qdrant",
ESVectorStore: "es",
PGVectorStore: "pgvector",
ChromaVectorStore: "chroma",
ObVecVectorStore: "obvec",
HologresVectorStore: "hologres",
}
# pylint: disable=too-many-return-statements
def get_store_type(store: BaseVectorStore) -> str:
"""Get the type identifier of a vector store instance.
@ -226,102 +220,29 @@ def get_store_type(store: BaseVectorStore) -> str:
store: Vector store instance
Returns:
str: Type identifier ("local", "es", "pgvector", "qdrant", "chroma", "obvec", or "hologres")
str: Type identifier ("local", "es", "pgvector", "qdrant", "chroma", "obvec", "zvec", or "hologres")
"""
for cls, name in _STORE_TYPE_MAP.items():
if isinstance(store, cls):
return name
raise ValueError(f"Unknown vector store type: {type(store)}")
def _build_store_kwargs(store_type, config, collection_name, embedding_model):
"""Build keyword arguments for creating a vector store instance."""
store_kwargs_map = {
"local": lambda: {
"collection_name": collection_name,
"embedding_model": embedding_model,
"db_path": config.LOCAL_ROOT_PATH,
},
"es": lambda: {
"collection_name": collection_name,
"embedding_model": embedding_model,
"db_path": tempfile.mkdtemp(prefix="test_es_"),
"hosts": config.ES_HOSTS,
"basic_auth": config.ES_BASIC_AUTH,
},
"qdrant": lambda: {
"collection_name": collection_name,
"embedding_model": embedding_model,
"db_path": config.QDRANT_PATH or tempfile.mkdtemp(prefix="test_qdrant_"),
"host": config.QDRANT_HOST,
"port": config.QDRANT_PORT,
"url": config.QDRANT_URL,
"api_key": config.QDRANT_API_KEY,
"distance": "cosine",
"on_disk": False,
},
"pgvector": lambda: {
"collection_name": collection_name,
"embedding_model": embedding_model,
"db_path": tempfile.mkdtemp(prefix="test_pgvector_"),
"dsn": config.PG_DSN,
"min_size": config.PG_MIN_SIZE,
"max_size": config.PG_MAX_SIZE,
"use_hnsw": config.PG_USE_HNSW,
"use_diskann": config.PG_USE_DISKANN,
},
"chroma": lambda: {
"collection_name": collection_name,
"embedding_model": embedding_model,
"db_path": config.CHROMA_PATH,
"host": config.CHROMA_HOST,
"port": config.CHROMA_PORT,
"api_key": config.CHROMA_API_KEY,
"tenant": config.CHROMA_TENANT,
"database": config.CHROMA_DATABASE,
},
"obvec": lambda: {
"collection_name": collection_name,
"embedding_model": embedding_model,
"db_path": tempfile.mkdtemp(prefix="test_obvec_"),
"uri": config.OBVEC_URI,
"user": config.OBVEC_USER,
"password": config.OBVEC_PASSWORD,
"database": config.OBVEC_DATABASE,
"index_metric": "cosine",
"index_ef_search": 100,
},
"hologres": lambda: {
"collection_name": collection_name,
"embedding_model": embedding_model,
"db_path": tempfile.mkdtemp(prefix="test_hologres_"),
"host": config.HOLOGRES_HOST,
"port": config.HOLOGRES_PORT,
"database": config.HOLOGRES_DATABASE,
"user": config.HOLOGRES_USER,
"password": config.HOLOGRES_PASSWORD,
"schema": config.HOLOGRES_SCHEMA,
"min_size": config.HOLOGRES_MIN_SIZE,
"max_size": config.HOLOGRES_MAX_SIZE,
**({"dsn": config.HOLOGRES_DSN} if config.HOLOGRES_DSN else {}),
},
}
if store_type not in store_kwargs_map:
raise ValueError(f"Unknown store type: {store_type}")
return store_kwargs_map[store_type]()
_STORE_CLASS_MAP = {
"local": LocalVectorStore,
"es": ESVectorStore,
"qdrant": QdrantVectorStore,
"pgvector": PGVectorStore,
"chroma": ChromaVectorStore,
"obvec": ObVecVectorStore,
"hologres": HologresVectorStore,
}
if isinstance(store, LocalVectorStore):
return "local"
elif isinstance(store, QdrantVectorStore):
return "qdrant"
elif isinstance(store, ESVectorStore):
return "es"
elif isinstance(store, PGVectorStore):
return "pgvector"
elif isinstance(store, ChromaVectorStore):
return "chroma"
elif isinstance(store, ObVecVectorStore):
return "obvec"
elif isinstance(store, ZvecVectorStore):
return "zvec"
elif isinstance(store, HologresVectorStore):
return "hologres"
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.
@ -340,9 +261,93 @@ def create_vector_store(store_type: str, collection_name: str) -> BaseVectorStor
dimensions=config.EMBEDDING_DIMENSIONS,
)
kwargs = _build_store_kwargs(store_type, config, collection_name, embedding_model)
store_cls = _STORE_CLASS_MAP[store_type]
return store_cls(**kwargs)
if store_type == "local":
return LocalVectorStore(
collection_name=collection_name,
embedding_model=embedding_model,
db_path=config.LOCAL_ROOT_PATH,
)
elif store_type == "es":
return ESVectorStore(
collection_name=collection_name,
embedding_model=embedding_model,
db_path=tempfile.mkdtemp(prefix="test_es_"),
hosts=config.ES_HOSTS,
basic_auth=config.ES_BASIC_AUTH,
)
elif store_type == "qdrant":
return QdrantVectorStore(
collection_name=collection_name,
embedding_model=embedding_model,
db_path=config.QDRANT_PATH or tempfile.mkdtemp(prefix="test_qdrant_"),
host=config.QDRANT_HOST,
port=config.QDRANT_PORT,
url=config.QDRANT_URL,
api_key=config.QDRANT_API_KEY,
distance="cosine",
on_disk=False,
)
elif store_type == "pgvector":
return PGVectorStore(
collection_name=collection_name,
embedding_model=embedding_model,
db_path=tempfile.mkdtemp(prefix="test_pgvector_"),
dsn=config.PG_DSN,
min_size=config.PG_MIN_SIZE,
max_size=config.PG_MAX_SIZE,
use_hnsw=config.PG_USE_HNSW,
use_diskann=config.PG_USE_DISKANN,
)
elif store_type == "chroma":
return ChromaVectorStore(
collection_name=collection_name,
embedding_model=embedding_model,
db_path=config.CHROMA_PATH,
host=config.CHROMA_HOST,
port=config.CHROMA_PORT,
api_key=config.CHROMA_API_KEY,
tenant=config.CHROMA_TENANT,
database=config.CHROMA_DATABASE,
)
elif store_type == "obvec":
return ObVecVectorStore(
collection_name=collection_name,
embedding_model=embedding_model,
db_path=tempfile.mkdtemp(prefix="test_obvec_"),
uri=config.OBVEC_URI,
user=config.OBVEC_USER,
password=config.OBVEC_PASSWORD,
database=config.OBVEC_DATABASE,
index_metric="cosine",
index_ef_search=100,
)
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",
)
elif store_type == "hologres":
kwargs = {
"collection_name": collection_name,
"embedding_model": embedding_model,
"db_path": tempfile.mkdtemp(prefix="test_hologres_"),
"host": config.HOLOGRES_HOST,
"port": config.HOLOGRES_PORT,
"database": config.HOLOGRES_DATABASE,
"user": config.HOLOGRES_USER,
"password": config.HOLOGRES_PASSWORD,
"schema": config.HOLOGRES_SCHEMA,
"min_size": config.HOLOGRES_MIN_SIZE,
"max_size": config.HOLOGRES_MAX_SIZE,
}
if config.HOLOGRES_DSN:
kwargs["dsn"] = config.HOLOGRES_DSN
return HologresVectorStore(**kwargs)
else:
raise ValueError(f"Unknown store type: {store_type}")
# ==================== Test Functions ====================
@ -1836,6 +1841,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}")
@ -1896,6 +1908,11 @@ Examples:
action="store_true",
help="Test HologresVectorStore",
)
parser.add_argument(
"--zvec",
action="store_true",
help="Test ZvecVectorStore",
)
parser.add_argument(
"--all",
action="store_true",
@ -1916,6 +1933,7 @@ Examples:
("chroma", "ChromaVectorStore"),
("obvec", "ObVecVectorStore"),
("hologres", "HologresVectorStore"),
("zvec", "ZvecVectorStore"),
]
else:
# Build list based on individual flags
@ -1933,6 +1951,8 @@ Examples:
stores_to_test.append(("obvec", "ObVecVectorStore"))
if args.hologres:
stores_to_test.append(("hologres", "HologresVectorStore"))
if args.zvec:
stores_to_test.append(("zvec", "ZvecVectorStore"))
if not stores_to_test:
# Default to all vector stores if no argument provided
@ -1944,10 +1964,11 @@ Examples:
("chroma", "ChromaVectorStore"),
("obvec", "ObVecVectorStore"),
("hologres", "HologresVectorStore"),
("zvec", "ZvecVectorStore"),
]
print("No vector store specified, defaulting to test all vector stores")
print(
"Use --local/--es/--pgvector/--qdrant/--chroma/--obvec/--hologres to test specific ones\n",
"Use --local/--es/--pgvector/--qdrant/--chroma/--obvec/--zvec/--hologres 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()