mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-10-27 06:15:20 +00:00
Merge branch 'main' into feat/support_hologres
This commit is contained in:
commit
1cdce924df
33 changed files with 3749 additions and 437 deletions
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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}
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
573
reme/core/file_store/zvec_file_store.py
Normal file
573
reme/core/file_store/zvec_file_store.py
Normal 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}")
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
809
reme/core/vector_store/zvec_vector_store.py
Normal file
809
reme/core/vector_store/zvec_vector_store.py
Normal 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")
|
||||
|
|
@ -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(
|
||||
[
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -0,0 +1 @@
|
|||
"""Profile memory tools."""
|
||||
|
|
@ -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}")
|
||||
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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}"
|
||||
|
|
|
|||
234
reme/memory/vector_tools/profiles/file_profile_backend.py
Normal file
234
reme/memory/vector_tools/profiles/file_profile_backend.py
Normal 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)
|
||||
51
reme/memory/vector_tools/profiles/profile_backend.py
Normal file
51
reme/memory/vector_tools/profiles/profile_backend.py
Normal 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."""
|
||||
|
|
@ -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))
|
||||
|
|
|
|||
245
reme/memory/vector_tools/profiles/profile_vector_handler.py
Normal file
245
reme/memory/vector_tools/profiles/profile_vector_handler.py
Normal 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)
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
102
reme/memory/vector_tools/profiles/retrieve_profile.py
Normal file
102
reme/memory/vector_tools/profiles/retrieve_profile.py
Normal 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
|
||||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
55
reme/memory/vector_tools/profiles/vector_profile_backend.py
Normal file
55
reme/memory/vector_tools/profiles/vector_profile_backend.py
Normal 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)
|
||||
164
reme/reme.py
164
reme/reme.py
|
|
@ -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():
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
906
tests/test_zvec_vector_store.py
Normal file
906
tests/test_zvec_vector_store.py
Normal 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())
|
||||
|
|
@ -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()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue