mirror of
https://github.com/MervinPraison/PraisonAI.git
synced 2026-11-11 11:41:38 -10:00
Compare commits
2
Commits
v4.6.44
...
add-valkey
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
251470489f | ||
|
|
dbf8b0d6e1 |
No files matched your search
@@ -42,6 +42,7 @@ class MemoryBackend(str, Enum):
|
||||
FILE = "file"
|
||||
SQLITE = "sqlite"
|
||||
REDIS = "redis"
|
||||
VALKEY = "valkey"
|
||||
POSTGRES = "postgres"
|
||||
MEM0 = "mem0"
|
||||
MONGODB = "mongodb"
|
||||
|
||||
@@ -346,13 +346,15 @@ class _LazyStorageModule:
|
||||
|
||||
_BACKENDS = {
|
||||
"RedisBackend": "praisonai.storage",
|
||||
"MongoDBBackend": "praisonai.storage",
|
||||
"MongoDBBackend": "praisonai.storage",
|
||||
"PostgreSQLBackend": "praisonai.storage",
|
||||
"DynamoDBBackend": "praisonai.storage",
|
||||
"ValkeyBackend": "praisonai.storage",
|
||||
"ValkeySearchBackend": "praisonai.storage",
|
||||
# Legacy names for compatibility
|
||||
"RedisStorageAdapter": "praisonai.storage",
|
||||
"MongoDBStorageAdapter": "praisonai.storage",
|
||||
"PostgreSQLStorageAdapter": "praisonai.storage",
|
||||
"PostgreSQLStorageAdapter": "praisonai.storage",
|
||||
"DynamoDBStorageAdapter": "praisonai.storage",
|
||||
}
|
||||
|
||||
@@ -390,7 +392,7 @@ def get_backend(
|
||||
Factory function to get a storage backend.
|
||||
|
||||
Args:
|
||||
backend_type: Type of backend ("file", "sqlite", "redis", "mongodb", "postgresql", "dynamodb")
|
||||
backend_type: Type of backend ("file", "sqlite", "redis", "mongodb", "postgresql", "dynamodb", "valkey")
|
||||
**kwargs: Backend-specific arguments
|
||||
|
||||
Returns:
|
||||
@@ -431,9 +433,13 @@ def get_backend(
|
||||
# Lazy load from wrapper
|
||||
backend_class = getattr(_heavy_backends, "DynamoDBBackend")
|
||||
return backend_class(**kwargs)
|
||||
elif backend_type == "valkey":
|
||||
# Lazy load from wrapper
|
||||
backend_class = getattr(_heavy_backends, "ValkeyBackend")
|
||||
return backend_class(**kwargs)
|
||||
else:
|
||||
raise ValueError(f"Unknown backend type: {backend_type}. "
|
||||
f"Supported: file, sqlite, redis, mongodb, postgresql, dynamodb")
|
||||
f"Supported: file, sqlite, redis, mongodb, postgresql, dynamodb, valkey")
|
||||
|
||||
def __getattr__(name: str):
|
||||
"""
|
||||
@@ -442,7 +448,8 @@ def __getattr__(name: str):
|
||||
This allows `from praisonaiagents.storage.backends import RedisBackend` to work
|
||||
while maintaining lazy loading.
|
||||
"""
|
||||
if name in {"RedisBackend", "MongoDBBackend", "PostgreSQLBackend", "DynamoDBBackend"}:
|
||||
if name in {"RedisBackend", "MongoDBBackend", "PostgreSQLBackend", "DynamoDBBackend",
|
||||
"ValkeyBackend", "ValkeySearchBackend"}:
|
||||
return getattr(_heavy_backends, name)
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
|
||||
|
||||
@@ -6,6 +6,7 @@ Databases tested (all via local Docker):
|
||||
2. PostgreSQL — ConversationStore + ManagedAgent roundtrip
|
||||
3. MySQL — ConversationStore + ManagedAgent roundtrip
|
||||
4. Redis — StateStore + metadata persistence
|
||||
4a. Valkey — StateStore + metadata persistence
|
||||
5. MongoDB — StateStore + metadata persistence
|
||||
6. ClickHouse — Raw connectivity + data write/read
|
||||
7. JSON file — DefaultSessionStore + ManagedAgent roundtrip
|
||||
@@ -27,7 +28,6 @@ from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -562,6 +562,206 @@ class TestRedisStateStore:
|
||||
r.close()
|
||||
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# 4a. Valkey — StateStore operations + metadata persistence
|
||||
# ===========================================================================
|
||||
class TestValkeyStateStore:
|
||||
"""Valkey StateStore: real operations against a live Valkey instance."""
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def check_valkey(self):
|
||||
"""Skip if Valkey not available."""
|
||||
try:
|
||||
from praisonai.persistence._valkey_client import create_valkey_client
|
||||
client = create_valkey_client(host="localhost", port=6379)
|
||||
client.ping()
|
||||
client.close()
|
||||
except Exception:
|
||||
pytest.skip("Valkey not available on localhost:6379")
|
||||
|
||||
def test_valkey_state_operations(self):
|
||||
"""StateStore get/set/delete/exists/hash with real Valkey."""
|
||||
from praisonai.persistence.state.valkey import ValkeyStateStore
|
||||
prefix = f"test_{uuid.uuid4().hex[:6]}:"
|
||||
store = ValkeyStateStore(host="localhost", port=6379, prefix=prefix)
|
||||
|
||||
try:
|
||||
# Basic set/get
|
||||
store.set("agent_id", "agent_valkey_001")
|
||||
assert store.get("agent_id") == "agent_valkey_001"
|
||||
|
||||
# JSON roundtrip
|
||||
state = {
|
||||
"agent_id": "agent_valkey_001",
|
||||
"agent_version": 3,
|
||||
"total_input_tokens": 500,
|
||||
"total_output_tokens": 200,
|
||||
"compute_instance_id": "docker_valkey_test",
|
||||
"session_history": [{"id": "s1", "status": "idle"}],
|
||||
}
|
||||
store.set_json("managed_state", state)
|
||||
recovered = store.get_json("managed_state")
|
||||
assert recovered["agent_id"] == "agent_valkey_001"
|
||||
assert recovered["total_input_tokens"] == 500
|
||||
assert recovered["compute_instance_id"] == "docker_valkey_test"
|
||||
assert len(recovered["session_history"]) == 1
|
||||
|
||||
# Hash operations
|
||||
store.hset("agent_meta", "version", "5")
|
||||
store.hset("agent_meta", "env_id", "env_abc")
|
||||
assert str(store.hget("agent_meta", "version")) == "5"
|
||||
all_meta = store.hgetall("agent_meta")
|
||||
assert str(all_meta["env_id"]) == "env_abc"
|
||||
|
||||
# Exists and delete
|
||||
assert store.exists("agent_id")
|
||||
store.delete("agent_id")
|
||||
assert not store.exists("agent_id")
|
||||
|
||||
finally:
|
||||
for k in store.keys():
|
||||
store.delete(k)
|
||||
store.close()
|
||||
|
||||
def test_valkey_managed_state_persist_restore(self):
|
||||
"""Simulate managed agent state persist/restore across store instances."""
|
||||
from praisonai.persistence.state.valkey import ValkeyStateStore
|
||||
prefix = f"mgd_{uuid.uuid4().hex[:6]}:"
|
||||
store = ValkeyStateStore(host="localhost", port=6379, prefix=prefix)
|
||||
|
||||
try:
|
||||
session_id = f"valkey_session_{uuid.uuid4().hex[:8]}"
|
||||
state = {
|
||||
"agent_id": "agent_valkey_mgd",
|
||||
"agent_version": 2,
|
||||
"environment_id": "env_valkey_001",
|
||||
"total_input_tokens": 1000,
|
||||
"total_output_tokens": 500,
|
||||
"compute_instance_id": "modal_valkey_xyz",
|
||||
"session_history": [
|
||||
{"id": session_id, "status": "idle", "title": "Valkey test"},
|
||||
],
|
||||
}
|
||||
store.set_json(f"managed:{session_id}", state)
|
||||
|
||||
# Simulate restart — new store instance, same prefix
|
||||
store2 = ValkeyStateStore(host="localhost", port=6379, prefix=prefix)
|
||||
recovered = store2.get_json(f"managed:{session_id}")
|
||||
assert recovered is not None
|
||||
assert recovered["agent_id"] == "agent_valkey_mgd"
|
||||
assert recovered["total_input_tokens"] == 1000
|
||||
assert recovered["compute_instance_id"] == "modal_valkey_xyz"
|
||||
assert len(recovered["session_history"]) == 1
|
||||
|
||||
store.close()
|
||||
store2.close()
|
||||
finally:
|
||||
store3 = ValkeyStateStore(host="localhost", port=6379, prefix=prefix)
|
||||
for k in store3.keys():
|
||||
store3.delete(k)
|
||||
store3.close()
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# 4b. Valkey Vector Search — ValkeyVectorKnowledgeStore against a live instance
|
||||
# ===========================================================================
|
||||
class TestValkeyVectorSearch:
|
||||
"""Valkey vector search: real operations against a live ValkeySearch instance."""
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def check_valkey_search(self):
|
||||
"""Skip if Valkey is unreachable or ValkeySearch (FT.*) is not loaded."""
|
||||
try:
|
||||
from glide_sync import ft
|
||||
except ImportError:
|
||||
pytest.skip("valkey-glide-sync not installed")
|
||||
|
||||
try:
|
||||
from praisonai.persistence._valkey_client import create_valkey_client
|
||||
except ImportError as e:
|
||||
pytest.skip(f"praisonai not importable: {e}")
|
||||
|
||||
client = None
|
||||
try:
|
||||
client = create_valkey_client(host="localhost", port=6379)
|
||||
client.ping()
|
||||
except Exception as e:
|
||||
if client is not None:
|
||||
client.close()
|
||||
pytest.skip(f"Valkey not available on localhost:6379: {e}")
|
||||
|
||||
try:
|
||||
# Probe FT.* support; raises if module not loaded
|
||||
ft.list(client)
|
||||
except Exception as e:
|
||||
pytest.skip(f"ValkeySearch (FT.*) module not loaded on localhost:6379: {e}")
|
||||
finally:
|
||||
if client is not None:
|
||||
client.close()
|
||||
|
||||
def test_valkey_vector_search_roundtrip(self):
|
||||
"""Full create → insert → search → delete_collection roundtrip."""
|
||||
from praisonai.persistence.knowledge.valkey_vector import ValkeyVectorKnowledgeStore
|
||||
from praisonai.persistence.knowledge.base import KnowledgeDocument
|
||||
|
||||
prefix = f"praisontest_{uuid.uuid4().hex[:6]}_"
|
||||
collection = f"vec_{uuid.uuid4().hex[:6]}"
|
||||
dimension = 4
|
||||
store = ValkeyVectorKnowledgeStore(host="localhost", port=6379, prefix=prefix)
|
||||
|
||||
try:
|
||||
# Create index
|
||||
store.create_collection(collection, dimension=dimension, distance="cosine")
|
||||
assert store.collection_exists(collection)
|
||||
|
||||
# Insert documents with simple, distinguishable embeddings
|
||||
docs = [
|
||||
KnowledgeDocument(
|
||||
id="doc_a",
|
||||
content="alpha document",
|
||||
embedding=[1.0, 0.0, 0.0, 0.0],
|
||||
),
|
||||
KnowledgeDocument(
|
||||
id="doc_b",
|
||||
content="bravo document",
|
||||
embedding=[0.0, 1.0, 0.0, 0.0],
|
||||
),
|
||||
KnowledgeDocument(
|
||||
id="doc_c",
|
||||
content="charlie document",
|
||||
embedding=[0.0, 0.0, 1.0, 0.0],
|
||||
),
|
||||
]
|
||||
inserted = store.insert(collection, docs)
|
||||
assert set(inserted) == {"doc_a", "doc_b", "doc_c"}
|
||||
|
||||
# Vector search: query closest to doc_a
|
||||
results = store.search(collection, query_embedding=[1.0, 0.0, 0.0, 0.0], limit=2)
|
||||
assert len(results) >= 1, "Expected at least one search result"
|
||||
assert results[0].id == "doc_a", f"Expected doc_a as nearest neighbor, got {results[0].id}"
|
||||
assert results[0].content == "alpha document"
|
||||
|
||||
# Verify the bare ID is returned (no prefix leakage)
|
||||
for r in results:
|
||||
assert ":" not in r.id or r.id in {"doc_a", "doc_b", "doc_c"}
|
||||
|
||||
# get() roundtrip
|
||||
got = store.get(collection, ["doc_b"])
|
||||
assert len(got) == 1
|
||||
assert got[0].id == "doc_b"
|
||||
assert got[0].content == "bravo document"
|
||||
|
||||
# Delete one doc
|
||||
deleted = store.delete(collection, ids=["doc_c"])
|
||||
assert deleted == 1
|
||||
finally:
|
||||
try:
|
||||
store.delete_collection(collection)
|
||||
finally:
|
||||
store.close()
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# 5. MongoDB — StateStore operations + metadata persistence
|
||||
# ===========================================================================
|
||||
@@ -673,7 +873,6 @@ class TestClickHouseConnectivity:
|
||||
def test_clickhouse_write_read(self):
|
||||
"""Write and read data from ClickHouse."""
|
||||
import clickhouse_connect
|
||||
|
||||
table = f"praison_test_{uuid.uuid4().hex[:8]}"
|
||||
c = clickhouse_connect.get_client(
|
||||
host="localhost", port=8123, username="clickhouse", password="clickhouse",
|
||||
|
||||
@@ -48,6 +48,12 @@ class TestResolveMemory:
|
||||
result = resolve_memory("redis")
|
||||
assert isinstance(result, MemoryConfig)
|
||||
assert result.backend == MemoryBackend.REDIS
|
||||
|
||||
def test_string_valkey_backend(self):
|
||||
"""'valkey' string returns MemoryConfig with VALKEY backend."""
|
||||
result = resolve_memory("valkey")
|
||||
assert isinstance(result, MemoryConfig)
|
||||
assert result.backend == MemoryBackend.VALKEY
|
||||
|
||||
def test_string_custom_backend(self):
|
||||
"""Custom string backend is preserved."""
|
||||
|
||||
@@ -0,0 +1,44 @@
|
||||
"""
|
||||
Shared Valkey client factory.
|
||||
|
||||
Centralises GlideClientSync creation so that ValkeyStateStore,
|
||||
ValkeyVectorKnowledgeStore, and ValkeyStorageAdapter all use the
|
||||
same connection logic.
|
||||
"""
|
||||
|
||||
from typing import Optional
|
||||
|
||||
try:
|
||||
from glide_sync import (
|
||||
GlideClient as GlideClientSync,
|
||||
GlideClientConfiguration,
|
||||
NodeAddress,
|
||||
ServerCredentials,
|
||||
)
|
||||
except ImportError:
|
||||
GlideClientSync = None
|
||||
GlideClientConfiguration = NodeAddress = ServerCredentials = None
|
||||
|
||||
_MISSING_MSG = (
|
||||
"valkey-glide-sync is required for Valkey support. "
|
||||
"Install with: pip install 'praisonai[valkey]'"
|
||||
)
|
||||
|
||||
|
||||
def create_valkey_client(
|
||||
host: str = "localhost",
|
||||
port: int = 6379,
|
||||
password: Optional[str] = None,
|
||||
db: int = 0,
|
||||
):
|
||||
"""Create and return a GlideClientSync instance."""
|
||||
if GlideClientSync is None:
|
||||
raise ImportError(_MISSING_MSG)
|
||||
addresses = [NodeAddress(host, port)]
|
||||
creds = ServerCredentials(password=password) if password else None
|
||||
config = GlideClientConfiguration(
|
||||
addresses=addresses,
|
||||
credentials=creds,
|
||||
database_id=db,
|
||||
)
|
||||
return GlideClientSync.create(config)
|
||||
@@ -0,0 +1,397 @@
|
||||
"""
|
||||
Valkey Vector implementation of KnowledgeStore.
|
||||
|
||||
Requires: valkey-glide-sync
|
||||
Install: pip install 'praisonai[valkey]'
|
||||
"""
|
||||
|
||||
import re
|
||||
import struct
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from .base import KnowledgeStore, KnowledgeDocument, validate_identifier
|
||||
from .._valkey_client import create_valkey_client
|
||||
|
||||
try:
|
||||
from glide_sync import ft
|
||||
from glide_sync import (
|
||||
TextField, NumericField, VectorField, VectorFieldAttributesHnsw,
|
||||
VectorAlgorithm, DistanceMetricType, VectorType,
|
||||
FtCreateOptions, DataType, FtSearchOptions, ReturnField,
|
||||
)
|
||||
from glide_sync import Batch
|
||||
except ImportError:
|
||||
ft = None
|
||||
Batch = None
|
||||
TextField = NumericField = VectorField = VectorFieldAttributesHnsw = None
|
||||
VectorAlgorithm = DistanceMetricType = VectorType = None
|
||||
FtCreateOptions = DataType = FtSearchOptions = ReturnField = None
|
||||
|
||||
|
||||
_SEARCH_SPECIAL_RE = re.compile(r'([,.<>{}\[\]"\'\\:;!@#$%^&*()\-+=~?/| \t])')
|
||||
_FIELD_NAME_RE = re.compile(r'^[a-zA-Z0-9_]+$')
|
||||
|
||||
|
||||
def _escape_search_value(value: str) -> str:
|
||||
return _SEARCH_SPECIAL_RE.sub(r'\\\1', str(value))
|
||||
|
||||
|
||||
def _decode(v) -> str:
|
||||
return v.decode() if isinstance(v, bytes) else v
|
||||
|
||||
|
||||
class ValkeyVectorKnowledgeStore(KnowledgeStore):
|
||||
"""
|
||||
Valkey-based knowledge store using ValkeySearch FT commands for vector search.
|
||||
|
||||
Requires Valkey with ValkeySearch module.
|
||||
|
||||
Example:
|
||||
store = ValkeyVectorKnowledgeStore(
|
||||
host="localhost",
|
||||
port=6379,
|
||||
)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
host: str = "localhost",
|
||||
port: int = 6379,
|
||||
password: Optional[str] = None,
|
||||
prefix: str = "praison:vec:",
|
||||
):
|
||||
self.host = host
|
||||
self.port = port
|
||||
self.password = password
|
||||
self.prefix = prefix
|
||||
self._client = None
|
||||
|
||||
def _get_client(self):
|
||||
"""Lazy initialize Valkey client."""
|
||||
if self._client is None:
|
||||
self._client = create_valkey_client(
|
||||
host=self.host,
|
||||
port=self.port,
|
||||
password=self.password,
|
||||
)
|
||||
return self._client
|
||||
|
||||
def _index_name(self, collection: str) -> str:
|
||||
return f"{self.prefix}{collection}:idx"
|
||||
|
||||
def _doc_key(self, collection: str, doc_id: str) -> str:
|
||||
return f"{self.prefix}{collection}:{doc_id}"
|
||||
|
||||
def create_collection(
|
||||
self,
|
||||
name: str,
|
||||
dimension: int,
|
||||
distance: str = "cosine",
|
||||
metadata: Optional[Dict[str, Any]] = None,
|
||||
) -> None:
|
||||
"""Create a new index."""
|
||||
validate_identifier(name, "collection name")
|
||||
client = self._get_client()
|
||||
index_name = self._index_name(name)
|
||||
|
||||
try:
|
||||
dist_enum = {"cosine": DistanceMetricType.COSINE, "euclidean": DistanceMetricType.L2, "dot": DistanceMetricType.IP}.get(distance, DistanceMetricType.COSINE)
|
||||
schema = [
|
||||
TextField("content"),
|
||||
TextField("content_hash"),
|
||||
NumericField("created_at"),
|
||||
VectorField("embedding", VectorAlgorithm.HNSW, VectorFieldAttributesHnsw(dimension, dist_enum, VectorType.FLOAT32)),
|
||||
]
|
||||
options = FtCreateOptions(DataType.HASH, prefixes=[f"{self.prefix}{name}:"])
|
||||
ft.create(client, index_name, schema, options)
|
||||
except Exception as e:
|
||||
err = str(e)
|
||||
if "already exists" not in err.lower():
|
||||
raise RuntimeError(f"Failed to create collection in Valkey: {e}") from e
|
||||
|
||||
def delete_collection(self, name: str) -> bool:
|
||||
"""Delete an index and all documents."""
|
||||
client = self._get_client()
|
||||
index_name = self._index_name(name)
|
||||
try:
|
||||
ft.dropindex(client, index_name)
|
||||
except Exception as e:
|
||||
# Only return False for index not found, re-raise other errors
|
||||
err_msg = str(e).lower()
|
||||
if "index not found" in err_msg or "unknown index" in err_msg:
|
||||
return False
|
||||
raise RuntimeError(f"Failed to delete collection: {e}") from e
|
||||
|
||||
# Manually remove orphaned document hashes left behind (no DD flag in ValkeySearch)
|
||||
pattern = f"{self.prefix}{name}:*"
|
||||
cursor: bytes = b"0"
|
||||
while True:
|
||||
result = client.scan(cursor, match=pattern, count=100)
|
||||
cursor = result[0]
|
||||
keys = result[1] or []
|
||||
if keys:
|
||||
client.delete(keys)
|
||||
if not cursor or cursor in (b"0", "0", 0):
|
||||
break
|
||||
|
||||
return True
|
||||
|
||||
def collection_exists(self, name: str) -> bool:
|
||||
"""Check if an index exists."""
|
||||
client = self._get_client()
|
||||
index_name = self._index_name(name)
|
||||
try:
|
||||
ft.info(client, index_name)
|
||||
return True
|
||||
except Exception as e:
|
||||
# Only return False for index not found, re-raise other errors
|
||||
err_msg = str(e).lower()
|
||||
if "index not found" in err_msg or "unknown index" in err_msg:
|
||||
return False
|
||||
raise RuntimeError(f"Failed to check collection existence: {e}") from e
|
||||
|
||||
def list_collections(self) -> List[str]:
|
||||
"""List all collections."""
|
||||
client = self._get_client()
|
||||
try:
|
||||
indexes = ft.list(client)
|
||||
if not indexes:
|
||||
return []
|
||||
result = []
|
||||
for idx in indexes:
|
||||
idx_str = _decode(idx)
|
||||
if idx_str.startswith(self.prefix):
|
||||
name = idx_str[len(self.prefix):]
|
||||
if name.endswith(":idx"):
|
||||
name = name[:-4]
|
||||
result.append(name)
|
||||
return result
|
||||
except Exception as e:
|
||||
# Only return empty list for connection/auth issues in some contexts
|
||||
# but generally we should propagate errors
|
||||
raise RuntimeError(f"Failed to list collections: {e}") from e
|
||||
|
||||
def insert(
|
||||
self,
|
||||
collection: str,
|
||||
documents: List[KnowledgeDocument],
|
||||
) -> List[str]:
|
||||
"""Insert documents using a non-atomic batch (pipeline)."""
|
||||
client = self._get_client()
|
||||
ids = []
|
||||
batch = Batch(is_atomic=False)
|
||||
|
||||
for doc in documents:
|
||||
if doc.embedding is None:
|
||||
raise ValueError(f"Document {doc.id} has no embedding")
|
||||
|
||||
key = self._doc_key(collection, doc.id)
|
||||
embedding_bytes = struct.pack(f"{len(doc.embedding)}f", *doc.embedding)
|
||||
|
||||
batch.hset(key, {
|
||||
"content": doc.content,
|
||||
"content_hash": doc.content_hash or "",
|
||||
"created_at": str(doc.created_at),
|
||||
"metadata": json.dumps(doc.metadata or {}),
|
||||
"embedding": embedding_bytes,
|
||||
})
|
||||
ids.append(doc.id)
|
||||
|
||||
try:
|
||||
client.exec(batch, raise_on_error=True)
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Failed to insert documents into Valkey: {e}") from e
|
||||
|
||||
return ids
|
||||
|
||||
def upsert(
|
||||
self,
|
||||
collection: str,
|
||||
documents: List[KnowledgeDocument],
|
||||
) -> List[str]:
|
||||
"""Insert or update documents."""
|
||||
return self.insert(collection, documents)
|
||||
|
||||
def search(
|
||||
self,
|
||||
collection: str,
|
||||
query_embedding: List[float],
|
||||
limit: int = 5,
|
||||
filters: Optional[Dict[str, Any]] = None,
|
||||
score_threshold: Optional[float] = None,
|
||||
) -> List[KnowledgeDocument]:
|
||||
"""Search for similar documents.
|
||||
|
||||
Args:
|
||||
collection: Collection name.
|
||||
query_embedding: Query vector.
|
||||
limit: Max results.
|
||||
filters: Optional dict of field-value equality filters applied as
|
||||
a ValkeySearch pre-filter (e.g. {"content_hash": "abc"}).
|
||||
score_threshold: Minimum similarity (1 - distance) to include.
|
||||
NOTE: Currently only accurate for cosine distance. For euclidean/L2
|
||||
and inner-product distances, the threshold behavior may be incorrect.
|
||||
"""
|
||||
client = self._get_client()
|
||||
index_name = self._index_name(collection)
|
||||
packed = struct.pack(f"{len(query_embedding)}f", *query_embedding)
|
||||
|
||||
# Build pre-filter expression
|
||||
if filters:
|
||||
filter_parts = []
|
||||
for field, value in filters.items():
|
||||
if not _FIELD_NAME_RE.match(field):
|
||||
raise ValueError(f"Invalid filter field name: {field!r}")
|
||||
if not isinstance(value, str):
|
||||
raise TypeError(
|
||||
f"Filter value for {field!r} must be a string, got {type(value).__name__}"
|
||||
)
|
||||
escaped = _escape_search_value(value)
|
||||
filter_parts.append(f"@{field}:({escaped})")
|
||||
pre_filter = " ".join(filter_parts)
|
||||
query = f"({pre_filter})=>[KNN {limit} @embedding $vec AS score]"
|
||||
else:
|
||||
query = f"*=>[KNN {limit} @embedding $vec AS score]"
|
||||
|
||||
try:
|
||||
options = FtSearchOptions(
|
||||
return_fields=[
|
||||
ReturnField("content"),
|
||||
ReturnField("content_hash"),
|
||||
ReturnField("created_at"),
|
||||
ReturnField("metadata"),
|
||||
ReturnField("score"),
|
||||
],
|
||||
params={"vec": packed},
|
||||
)
|
||||
raw = ft.search(client, index_name, query, options)
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Failed to search Valkey: {e}") from e
|
||||
|
||||
# ft.search returns [count, {doc_id: {field: value, ...}, ...}]
|
||||
if not raw or len(raw) < 2:
|
||||
return []
|
||||
|
||||
documents = []
|
||||
key_prefix = f"{self.prefix}{collection}:"
|
||||
for doc_id_raw, fields in raw[1].items():
|
||||
doc_id_str = _decode(doc_id_raw)
|
||||
bare_id = doc_id_str.removeprefix(key_prefix)
|
||||
|
||||
field_map: Dict[str, str] = {}
|
||||
for fname, fval in fields.items():
|
||||
fname = _decode(fname)
|
||||
fval = _decode(fval)
|
||||
field_map[fname] = fval
|
||||
|
||||
score = float(field_map.get("score", 0))
|
||||
similarity = 1 - score
|
||||
|
||||
if score_threshold is not None and similarity < score_threshold:
|
||||
continue
|
||||
|
||||
documents.append(KnowledgeDocument(
|
||||
id=bare_id,
|
||||
content=field_map.get("content", ""),
|
||||
embedding=None,
|
||||
metadata=json.loads(field_map.get("metadata", "{}") or "{}"),
|
||||
content_hash=field_map.get("content_hash", ""),
|
||||
created_at=float(field_map.get("created_at", 0)),
|
||||
))
|
||||
|
||||
return documents
|
||||
|
||||
def get(
|
||||
self,
|
||||
collection: str,
|
||||
ids: List[str],
|
||||
) -> List[KnowledgeDocument]:
|
||||
"""Get documents by IDs using a non-atomic batch (pipeline)."""
|
||||
client = self._get_client()
|
||||
keys = [self._doc_key(collection, doc_id) for doc_id in ids]
|
||||
|
||||
batch = Batch(is_atomic=False)
|
||||
for key in keys:
|
||||
batch.hgetall(key)
|
||||
|
||||
try:
|
||||
results = client.exec(batch, raise_on_error=True)
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Failed to get documents from Valkey: {e}") from e
|
||||
|
||||
documents = []
|
||||
for doc_id, data in zip(ids, results):
|
||||
if data:
|
||||
embedding = None
|
||||
emb_key = b"embedding" if b"embedding" in data else "embedding"
|
||||
if emb_key in data:
|
||||
raw_bytes = data[emb_key]
|
||||
count = len(raw_bytes) // 4
|
||||
embedding = list(struct.unpack(f"{count}f", raw_bytes))
|
||||
|
||||
content = _decode(data.get(b"content", data.get("content", b"")))
|
||||
content_hash = _decode(data.get(b"content_hash", data.get("content_hash", b"")))
|
||||
created_at_raw = data.get(b"created_at", data.get("created_at", b"0"))
|
||||
created_at = float(_decode(created_at_raw))
|
||||
metadata_raw = _decode(data.get(b"metadata", data.get("metadata", b"{}")))
|
||||
metadata = json.loads(metadata_raw or "{}")
|
||||
|
||||
doc = KnowledgeDocument(
|
||||
id=doc_id,
|
||||
content=content,
|
||||
embedding=embedding,
|
||||
metadata=metadata,
|
||||
content_hash=content_hash,
|
||||
created_at=created_at,
|
||||
)
|
||||
documents.append(doc)
|
||||
|
||||
return documents
|
||||
|
||||
def delete(
|
||||
self,
|
||||
collection: str,
|
||||
ids: Optional[List[str]] = None,
|
||||
filters: Optional[Dict[str, Any]] = None,
|
||||
) -> int:
|
||||
"""Delete documents by IDs. Filter-based deletion is not supported."""
|
||||
if filters and not ids:
|
||||
raise NotImplementedError("Filter-based deletion is not supported by the Valkey backend")
|
||||
if not ids:
|
||||
return 0
|
||||
|
||||
client = self._get_client()
|
||||
keys = [self._doc_key(collection, doc_id) for doc_id in ids]
|
||||
|
||||
try:
|
||||
deleted = client.delete(keys)
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Failed to delete documents from Valkey: {e}") from e
|
||||
|
||||
return deleted
|
||||
|
||||
def count(self, collection: str) -> int:
|
||||
"""Count documents in a collection."""
|
||||
client = self._get_client()
|
||||
index_name = self._index_name(collection)
|
||||
try:
|
||||
info = ft.info(client, index_name)
|
||||
# ft.info returns a Mapping; num_docs may be bytes key
|
||||
for k in (b"num_docs", "num_docs"):
|
||||
if k in info:
|
||||
return int(info[k])
|
||||
return 0
|
||||
except Exception as e:
|
||||
# Only return 0 for index not found, re-raise other errors
|
||||
err_msg = str(e).lower()
|
||||
if "index not found" in err_msg or "unknown index" in err_msg:
|
||||
return 0
|
||||
raise RuntimeError(f"Failed to count documents: {e}") from e
|
||||
|
||||
def close(self) -> None:
|
||||
"""Close the store."""
|
||||
if self._client:
|
||||
try:
|
||||
self._client.close()
|
||||
finally:
|
||||
self._client = None
|
||||
@@ -198,7 +198,16 @@ def _register_builtin_knowledge_stores(registry: StoreRegistry):
|
||||
def _redis_vector(url=None, **kwargs):
|
||||
from .knowledge.redis_vector import RedisVectorKnowledgeStore
|
||||
return RedisVectorKnowledgeStore(url=url, **kwargs)
|
||||
|
||||
|
||||
def _valkey_vector(url=None, **kwargs):
|
||||
from .knowledge.valkey_vector import ValkeyVectorKnowledgeStore
|
||||
if url is not None:
|
||||
raise ValueError(
|
||||
"Valkey vector store does not accept 'url' parameter. "
|
||||
"Use host/port/password kwargs instead."
|
||||
)
|
||||
return ValkeyVectorKnowledgeStore(**kwargs)
|
||||
|
||||
def _cassandra(url=None, **kwargs):
|
||||
kwargs.pop("url", None) # Cassandra doesn't use url parameter
|
||||
from .knowledge.cassandra import CassandraKnowledgeStore
|
||||
@@ -258,9 +267,10 @@ def _register_builtin_knowledge_stores(registry: StoreRegistry):
|
||||
registry.register("milvus", _milvus)
|
||||
registry.register("pgvector", _pgvector)
|
||||
registry.register("redis", _redis_vector)
|
||||
registry.register("valkey", _valkey_vector)
|
||||
registry.register("cassandra", _cassandra)
|
||||
registry.register("clickhouse", _clickhouse)
|
||||
registry.register("mongodb_vector", _mongodb_vector,
|
||||
registry.register("mongodb_vector", _mongodb_vector,
|
||||
aliases=("mongodb_atlas", "mongo_vector"))
|
||||
|
||||
# Register missing backends that were supported in old factory
|
||||
@@ -286,7 +296,16 @@ def _register_builtin_state_stores(registry: StoreRegistry):
|
||||
def _redis(url=None, **kwargs):
|
||||
from .state.redis import RedisStateStore
|
||||
return RedisStateStore(url=url, **kwargs)
|
||||
|
||||
|
||||
def _valkey(url=None, **kwargs):
|
||||
from .state.valkey import ValkeyStateStore
|
||||
if url is not None:
|
||||
raise ValueError(
|
||||
"Valkey state store does not accept 'url' parameter. "
|
||||
"Use host/port/password kwargs instead."
|
||||
)
|
||||
return ValkeyStateStore(**kwargs)
|
||||
|
||||
def _dynamodb(url=None, **kwargs):
|
||||
kwargs.pop("url", None) # DynamoDB doesn't use url parameter
|
||||
from .state.dynamodb import DynamoDBStateStore
|
||||
@@ -322,10 +341,11 @@ def _register_builtin_state_stores(registry: StoreRegistry):
|
||||
return GCSStateStore(bucket_name=bucket, **kwargs)
|
||||
|
||||
registry.register("redis", _redis)
|
||||
registry.register("valkey", _valkey)
|
||||
registry.register("dynamodb", _dynamodb)
|
||||
registry.register("firestore", _firestore)
|
||||
registry.register("mongodb", _mongodb)
|
||||
registry.register("async_mongodb", _async_mongodb,
|
||||
registry.register("async_mongodb", _async_mongodb,
|
||||
aliases=("motor", "mongodb_async"))
|
||||
registry.register("upstash", _upstash)
|
||||
registry.register("memory", _memory)
|
||||
|
||||
@@ -0,0 +1,233 @@
|
||||
"""
|
||||
Valkey implementation of StateStore.
|
||||
|
||||
Requires: valkey-glide-sync
|
||||
Install: pip install 'praisonai[valkey]'
|
||||
"""
|
||||
|
||||
import json
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from .base import StateStore
|
||||
from .._valkey_client import create_valkey_client
|
||||
|
||||
try:
|
||||
from glide_sync import ExpirySet, ExpiryType
|
||||
except ImportError:
|
||||
ExpirySet = None
|
||||
ExpiryType = None
|
||||
|
||||
|
||||
class ValkeyStateStore(StateStore):
|
||||
"""
|
||||
Valkey-based state store for fast key-value operations.
|
||||
|
||||
Example:
|
||||
store = ValkeyStateStore(
|
||||
host="localhost",
|
||||
port=6379,
|
||||
)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
host: str = "localhost",
|
||||
port: int = 6379,
|
||||
password: Optional[str] = None,
|
||||
prefix: str = "praison:",
|
||||
db: int = 0,
|
||||
):
|
||||
self.host = host
|
||||
self.port = port
|
||||
self.password = password
|
||||
self.prefix = prefix
|
||||
self.db = db
|
||||
self._client = None
|
||||
|
||||
def _get_client(self):
|
||||
"""Lazy initialize Valkey client."""
|
||||
if self._client is None:
|
||||
self._client = create_valkey_client(
|
||||
host=self.host,
|
||||
port=self.port,
|
||||
password=self.password,
|
||||
db=self.db,
|
||||
)
|
||||
return self._client
|
||||
|
||||
def _key(self, key: str) -> str:
|
||||
"""Add prefix to key."""
|
||||
return f"{self.prefix}{key}"
|
||||
|
||||
def get(self, key: str) -> Optional[Any]:
|
||||
"""Get a value by key."""
|
||||
client = self._get_client()
|
||||
try:
|
||||
value = client.get(self._key(key))
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Failed to get key from Valkey: {e}") from e
|
||||
|
||||
if value is None:
|
||||
return None
|
||||
if isinstance(value, bytes):
|
||||
value = value.decode()
|
||||
try:
|
||||
return json.loads(value)
|
||||
except json.JSONDecodeError:
|
||||
return value
|
||||
|
||||
def set(self, key: str, value: Any, ttl: Optional[int] = None) -> None:
|
||||
"""Set a value with optional TTL."""
|
||||
client = self._get_client()
|
||||
value = json.dumps(value)
|
||||
|
||||
try:
|
||||
if ttl is not None:
|
||||
if ExpirySet is None or ExpiryType is None:
|
||||
from .._valkey_client import _MISSING_MSG
|
||||
raise ImportError(_MISSING_MSG)
|
||||
client.set(self._key(key), value, expiry=ExpirySet(ExpiryType.SEC, ttl))
|
||||
else:
|
||||
client.set(self._key(key), value)
|
||||
except ImportError:
|
||||
raise
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Failed to set key in Valkey: {e}") from e
|
||||
|
||||
def delete(self, key: str) -> bool:
|
||||
"""Delete a key."""
|
||||
client = self._get_client()
|
||||
try:
|
||||
return client.delete([self._key(key)]) > 0
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Failed to delete key from Valkey: {e}") from e
|
||||
|
||||
def exists(self, key: str) -> bool:
|
||||
"""Check if a key exists."""
|
||||
client = self._get_client()
|
||||
try:
|
||||
return client.exists([self._key(key)]) > 0
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Failed to check key existence in Valkey: {e}") from e
|
||||
|
||||
def keys(self, pattern: str = "*") -> List[str]:
|
||||
"""List keys matching pattern."""
|
||||
client = self._get_client()
|
||||
full_pattern = self._key(pattern)
|
||||
try:
|
||||
cursor: bytes = b"0"
|
||||
all_keys = []
|
||||
while True:
|
||||
result = client.scan(cursor, match=full_pattern, count=100)
|
||||
cursor = result[0]
|
||||
all_keys.extend(result[1] or [])
|
||||
if not cursor or cursor in (b"0", "0"):
|
||||
break
|
||||
prefix_len = len(self.prefix)
|
||||
result_keys = []
|
||||
for k in all_keys:
|
||||
k_str = k.decode() if isinstance(k, bytes) else k
|
||||
result_keys.append(k_str[prefix_len:] if k_str.startswith(self.prefix) else k_str)
|
||||
return result_keys
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Failed to list keys from Valkey: {e}") from e
|
||||
|
||||
def ttl(self, key: str) -> Optional[int]:
|
||||
"""Get remaining TTL in seconds."""
|
||||
client = self._get_client()
|
||||
try:
|
||||
result = client.ttl(self._key(key))
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Failed to get TTL from Valkey: {e}") from e
|
||||
|
||||
if result is None or result < 0:
|
||||
return None
|
||||
return int(result)
|
||||
|
||||
def expire(self, key: str, ttl: int) -> bool:
|
||||
"""Set TTL on existing key."""
|
||||
client = self._get_client()
|
||||
try:
|
||||
return bool(client.expire(self._key(key), ttl))
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Failed to set expire in Valkey: {e}") from e
|
||||
|
||||
def hget(self, key: str, field: str) -> Optional[Any]:
|
||||
"""Get a field from a hash."""
|
||||
client = self._get_client()
|
||||
try:
|
||||
value = client.hget(self._key(key), field)
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Failed to hget from Valkey: {e}") from e
|
||||
|
||||
if value is None:
|
||||
return None
|
||||
if isinstance(value, bytes):
|
||||
value = value.decode()
|
||||
try:
|
||||
return json.loads(value)
|
||||
except json.JSONDecodeError:
|
||||
return value
|
||||
|
||||
def hset(self, key: str, field: str, value: Any) -> None:
|
||||
"""Set a field in a hash."""
|
||||
client = self._get_client()
|
||||
value = json.dumps(value)
|
||||
try:
|
||||
client.hset(self._key(key), {field: value})
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Failed to hset in Valkey: {e}") from e
|
||||
|
||||
def hgetall(self, key: str) -> Dict[str, Any]:
|
||||
"""Get all fields from a hash."""
|
||||
client = self._get_client()
|
||||
try:
|
||||
data = client.hgetall(self._key(key))
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Failed to hgetall from Valkey: {e}") from e
|
||||
|
||||
result = {}
|
||||
if not data:
|
||||
return result
|
||||
for k, v in data.items():
|
||||
k_str = k.decode() if isinstance(k, bytes) else k
|
||||
v_str = v.decode() if isinstance(v, bytes) else v
|
||||
try:
|
||||
result[k_str] = json.loads(v_str)
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
result[k_str] = v_str
|
||||
return result
|
||||
|
||||
def hdel(self, key: str, *fields: str) -> int:
|
||||
"""Delete fields from a hash."""
|
||||
if not fields:
|
||||
return 0
|
||||
client = self._get_client()
|
||||
try:
|
||||
return client.hdel(self._key(key), list(fields))
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Failed to hdel from Valkey: {e}") from e
|
||||
|
||||
def incr(self, key: str, amount: int = 1) -> int:
|
||||
"""Increment a counter."""
|
||||
client = self._get_client()
|
||||
try:
|
||||
return client.incrby(self._key(key), amount)
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Failed to increment key in Valkey: {e}") from e
|
||||
|
||||
def decr(self, key: str, amount: int = 1) -> int:
|
||||
"""Decrement a counter."""
|
||||
client = self._get_client()
|
||||
try:
|
||||
return client.decrby(self._key(key), amount)
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Failed to decrement key in Valkey: {e}") from e
|
||||
|
||||
def close(self) -> None:
|
||||
"""Close the store."""
|
||||
if self._client:
|
||||
try:
|
||||
self._client.close()
|
||||
finally:
|
||||
self._client = None
|
||||
@@ -0,0 +1,492 @@
|
||||
"""
|
||||
Tests for ValkeyVectorKnowledgeStore, ValkeyStateStore, and registry integration.
|
||||
|
||||
Tests cover:
|
||||
- ValkeyVectorKnowledgeStore collection and document operations (mocked glide client)
|
||||
- ValkeyStateStore key-value operations (mocked glide client)
|
||||
- KNOWLEDGE_STORES and STATE_STORES registry entries for "valkey"
|
||||
"""
|
||||
|
||||
import json
|
||||
import struct
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from praisonai.persistence.knowledge.base import KnowledgeDocument
|
||||
from praisonai.persistence.registry import KNOWLEDGE_STORES, STATE_STORES
|
||||
|
||||
|
||||
def _make_knowledge_store():
|
||||
"""Return a ValkeyVectorKnowledgeStore with a pre-injected MagicMock client."""
|
||||
from praisonai.persistence.knowledge.valkey_vector import ValkeyVectorKnowledgeStore
|
||||
|
||||
store = ValkeyVectorKnowledgeStore.__new__(ValkeyVectorKnowledgeStore)
|
||||
store.host = "localhost"
|
||||
store.port = 6379
|
||||
store.password = None
|
||||
store.prefix = "praison:vec:"
|
||||
mock_client = MagicMock()
|
||||
store._client = mock_client
|
||||
store._get_client = lambda: mock_client
|
||||
|
||||
return store, mock_client
|
||||
|
||||
|
||||
def _make_state_store():
|
||||
"""Return a ValkeyStateStore with a pre-injected MagicMock client."""
|
||||
from praisonai.persistence.state.valkey import ValkeyStateStore
|
||||
|
||||
store = ValkeyStateStore.__new__(ValkeyStateStore)
|
||||
store.host = "localhost"
|
||||
store.port = 6379
|
||||
store.password = None
|
||||
store.prefix = "praison:"
|
||||
store.db = 0
|
||||
mock_client = MagicMock()
|
||||
store._client = mock_client
|
||||
store._get_client = lambda: mock_client
|
||||
|
||||
return store, mock_client
|
||||
|
||||
|
||||
class TestValkeyVectorKnowledgeStore:
|
||||
"""Tests for ValkeyVectorKnowledgeStore."""
|
||||
|
||||
_FT_MODULE = "praisonai.persistence.knowledge.valkey_vector.ft"
|
||||
|
||||
def _mock_ft(self, **method_returns):
|
||||
"""Return a mock ft module with configured method returns."""
|
||||
mock_ft = MagicMock()
|
||||
for name, val in method_returns.items():
|
||||
if isinstance(val, Exception):
|
||||
getattr(mock_ft, name).side_effect = val
|
||||
else:
|
||||
getattr(mock_ft, name).return_value = val
|
||||
return mock_ft
|
||||
|
||||
def test_create_collection(self):
|
||||
"""Test create_collection calls ft.create with HNSW schema."""
|
||||
store, _ = _make_knowledge_store()
|
||||
mock_ft = self._mock_ft(create=b"OK")
|
||||
|
||||
with patch(self._FT_MODULE, mock_ft):
|
||||
store.create_collection("docs", dimension=4)
|
||||
|
||||
mock_ft.create.assert_called_once()
|
||||
assert mock_ft.create.call_args[0][1] == "praison:vec:docs:idx"
|
||||
|
||||
def test_create_collection_already_exists(self):
|
||||
"""Test create_collection silently ignores 'already exists' error."""
|
||||
store, _ = _make_knowledge_store()
|
||||
mock_ft = self._mock_ft(create=Exception("already exists"))
|
||||
|
||||
with patch(self._FT_MODULE, mock_ft):
|
||||
store.create_collection("docs", dimension=4) # should not raise
|
||||
|
||||
def test_delete_collection(self):
|
||||
"""Test delete_collection calls ft.dropindex and returns True on success."""
|
||||
store, mock_client = _make_knowledge_store()
|
||||
mock_ft = self._mock_ft(dropindex=b"OK")
|
||||
mock_client.scan.return_value = [b"0", []]
|
||||
|
||||
with patch(self._FT_MODULE, mock_ft):
|
||||
result = store.delete_collection("docs")
|
||||
|
||||
assert result is True
|
||||
mock_ft.dropindex.assert_called_once()
|
||||
|
||||
def test_collection_exists_true(self):
|
||||
"""Test collection_exists returns True when ft.info succeeds."""
|
||||
store, _ = _make_knowledge_store()
|
||||
mock_ft = self._mock_ft(info={b"index_name": b"praison:vec:docs:idx"})
|
||||
|
||||
with patch(self._FT_MODULE, mock_ft):
|
||||
result = store.collection_exists("docs")
|
||||
|
||||
assert result is True
|
||||
|
||||
def test_collection_exists_false(self):
|
||||
"""Test collection_exists returns False when ft.info raises an exception."""
|
||||
store, _ = _make_knowledge_store()
|
||||
mock_ft = self._mock_ft(info=Exception("Unknown index name"))
|
||||
|
||||
with patch(self._FT_MODULE, mock_ft):
|
||||
result = store.collection_exists("docs")
|
||||
|
||||
assert result is False
|
||||
|
||||
def test_list_collections(self):
|
||||
"""Test list_collections returns only names matching the prefix, stripped."""
|
||||
store, _ = _make_knowledge_store()
|
||||
mock_ft = self._mock_ft(list=[b"praison:vec:docs:idx", b"other:idx"])
|
||||
|
||||
with patch(self._FT_MODULE, mock_ft):
|
||||
result = store.list_collections()
|
||||
|
||||
assert "docs" in result
|
||||
assert all("other" not in name for name in result)
|
||||
|
||||
def test_insert(self):
|
||||
"""Test insert uses Batch pipeline and calls client.exec."""
|
||||
store, mock_client = _make_knowledge_store()
|
||||
mock_client.exec.return_value = [4]
|
||||
|
||||
embedding = [0.1, 0.2, 0.3, 0.4]
|
||||
doc = KnowledgeDocument(id="doc1", content="hello", embedding=embedding)
|
||||
ids = store.insert("docs", [doc])
|
||||
|
||||
mock_client.exec.assert_called_once()
|
||||
assert ids == ["doc1"]
|
||||
|
||||
def test_upsert_delegates_to_insert(self):
|
||||
"""Test upsert calls insert (same path)."""
|
||||
store, mock_client = _make_knowledge_store()
|
||||
mock_client.exec.return_value = [4]
|
||||
|
||||
doc = KnowledgeDocument(id="doc1", content="hello", embedding=[0.1, 0.2, 0.3, 0.4])
|
||||
store.upsert("docs", [doc])
|
||||
|
||||
mock_client.exec.assert_called_once()
|
||||
|
||||
def test_search(self):
|
||||
"""Test search calls ft.search and returns a list of KnowledgeDocument."""
|
||||
store, _ = _make_knowledge_store()
|
||||
ft_result = [1, {
|
||||
b"praison:vec:docs:doc1": {
|
||||
b"content": b"hello",
|
||||
b"content_hash": b"abc",
|
||||
b"created_at": b"1234.0",
|
||||
b"score": b"0.1",
|
||||
}
|
||||
}]
|
||||
mock_ft = self._mock_ft(search=ft_result)
|
||||
|
||||
with patch(self._FT_MODULE, mock_ft):
|
||||
results = store.search("docs", query_embedding=[0.1, 0.2, 0.3, 0.4], limit=5)
|
||||
|
||||
assert isinstance(results, list)
|
||||
assert len(results) == 1
|
||||
assert isinstance(results[0], KnowledgeDocument)
|
||||
assert results[0].id == "doc1"
|
||||
mock_ft.search.assert_called_once()
|
||||
|
||||
def test_search_with_filters(self):
|
||||
"""Test search with filters builds a pre-filter KNN query."""
|
||||
store, _ = _make_knowledge_store()
|
||||
ft_result = [1, {
|
||||
b"praison:vec:docs:doc1": {
|
||||
b"content": b"hello",
|
||||
b"content_hash": b"abc",
|
||||
b"created_at": b"1234.0",
|
||||
b"score": b"0.1",
|
||||
}
|
||||
}]
|
||||
mock_ft = self._mock_ft(search=ft_result)
|
||||
|
||||
with patch(self._FT_MODULE, mock_ft):
|
||||
results = store.search(
|
||||
"docs",
|
||||
query_embedding=[0.1, 0.2, 0.3, 0.4],
|
||||
limit=5,
|
||||
filters={"content_hash": "abc"},
|
||||
)
|
||||
|
||||
assert len(results) == 1
|
||||
query_arg = mock_ft.search.call_args[0][2]
|
||||
assert "KNN" in query_arg
|
||||
assert query_arg.startswith("(")
|
||||
|
||||
def test_search_score_threshold_excludes_low_similarity(self):
|
||||
"""Test search excludes documents whose similarity is below score_threshold."""
|
||||
store, _ = _make_knowledge_store()
|
||||
ft_result = [2, {
|
||||
b"praison:vec:docs:doc1": {
|
||||
b"content": b"close",
|
||||
b"content_hash": b"",
|
||||
b"created_at": b"1.0",
|
||||
b"score": b"0.05", # similarity = 0.95 — above threshold
|
||||
},
|
||||
b"praison:vec:docs:doc2": {
|
||||
b"content": b"far",
|
||||
b"content_hash": b"",
|
||||
b"created_at": b"2.0",
|
||||
b"score": b"0.8", # similarity = 0.2 — below threshold
|
||||
},
|
||||
}]
|
||||
mock_ft = self._mock_ft(search=ft_result)
|
||||
|
||||
with patch(self._FT_MODULE, mock_ft):
|
||||
results = store.search(
|
||||
"docs",
|
||||
query_embedding=[0.1, 0.2, 0.3, 0.4],
|
||||
limit=5,
|
||||
score_threshold=0.5,
|
||||
)
|
||||
|
||||
assert len(results) == 1
|
||||
assert results[0].id == "doc1"
|
||||
|
||||
def test_search_empty_result(self):
|
||||
"""Test search returns an empty list when ft.search returns no documents."""
|
||||
store, _ = _make_knowledge_store()
|
||||
mock_ft = self._mock_ft(search=[])
|
||||
|
||||
with patch(self._FT_MODULE, mock_ft):
|
||||
results = store.search("docs", query_embedding=[0.1, 0.2, 0.3, 0.4])
|
||||
|
||||
assert results == []
|
||||
|
||||
def test_get(self):
|
||||
"""Test get uses Batch pipeline and returns a KnowledgeDocument."""
|
||||
store, mock_client = _make_knowledge_store()
|
||||
embedding = [0.1, 0.2, 0.3, 0.4]
|
||||
embedding_bytes = struct.pack(f"{len(embedding)}f", *embedding)
|
||||
mock_client.exec.return_value = [{
|
||||
b"content": b"hello",
|
||||
b"content_hash": b"abc",
|
||||
b"created_at": b"1234.0",
|
||||
b"embedding": embedding_bytes,
|
||||
}]
|
||||
|
||||
results = store.get("docs", ["doc1"])
|
||||
|
||||
assert len(results) == 1
|
||||
assert isinstance(results[0], KnowledgeDocument)
|
||||
assert results[0].id == "doc1"
|
||||
mock_client.exec.assert_called_once()
|
||||
|
||||
def test_delete(self):
|
||||
"""Test delete calls client.delete and returns count deleted."""
|
||||
store, mock_client = _make_knowledge_store()
|
||||
mock_client.delete.return_value = 2
|
||||
|
||||
result = store.delete("docs", ids=["doc1", "doc2"])
|
||||
|
||||
assert result == 2
|
||||
|
||||
def test_count(self):
|
||||
"""Test count calls ft.info and returns integer document count."""
|
||||
store, _ = _make_knowledge_store()
|
||||
mock_ft = self._mock_ft(info={b"num_docs": 7, b"index_name": b"praison:vec:docs:idx"})
|
||||
|
||||
with patch(self._FT_MODULE, mock_ft):
|
||||
result = store.count("docs")
|
||||
|
||||
assert isinstance(result, int)
|
||||
assert result == 7
|
||||
|
||||
def test_close_closes_client(self):
|
||||
"""Test close() calls client.close() and clears the reference."""
|
||||
store, mock_client = _make_knowledge_store()
|
||||
|
||||
store.close()
|
||||
|
||||
mock_client.close.assert_called_once()
|
||||
assert store._client is None
|
||||
|
||||
|
||||
class TestValkeyStateStore:
|
||||
"""Tests for ValkeyStateStore."""
|
||||
|
||||
def test_get_existing(self):
|
||||
"""Test get returns a dict when client.get returns JSON bytes."""
|
||||
store, mock_client = _make_state_store()
|
||||
data = {"result": "ok"}
|
||||
mock_client.get.return_value = json.dumps(data).encode("utf-8")
|
||||
|
||||
result = store.get("mykey")
|
||||
|
||||
mock_client.get.assert_called_once_with("praison:mykey")
|
||||
assert result == data
|
||||
|
||||
def test_get_missing(self):
|
||||
"""Test get returns None when client.get returns None."""
|
||||
store, mock_client = _make_state_store()
|
||||
mock_client.get.return_value = None
|
||||
|
||||
result = store.get("missing")
|
||||
|
||||
assert result is None
|
||||
|
||||
def test_set_no_ttl(self):
|
||||
"""Test set calls client.set with prefixed key when no TTL is given."""
|
||||
store, mock_client = _make_state_store()
|
||||
mock_client.set.return_value = b"OK"
|
||||
|
||||
store.set("mykey", {"x": 1})
|
||||
|
||||
mock_client.set.assert_called_once()
|
||||
call_args = mock_client.set.call_args[0]
|
||||
assert call_args[0] == "praison:mykey"
|
||||
|
||||
def test_set_with_ttl(self):
|
||||
"""Test set with TTL passes SetOptions to client.set (ImportError/RuntimeError tolerated)."""
|
||||
store, mock_client = _make_state_store()
|
||||
mock_client.set.return_value = b"OK"
|
||||
|
||||
# glide's ExpirySet/ExpiryType are lazy-imported inside set(); without a live
|
||||
# glide install this raises ImportError — both paths are valid in CI.
|
||||
try:
|
||||
store.set("ttlkey", {"x": 1}, ttl=300)
|
||||
except (ImportError, RuntimeError):
|
||||
pass
|
||||
|
||||
def test_delete_existing(self):
|
||||
"""Test delete returns True when client.delete returns 1."""
|
||||
store, mock_client = _make_state_store()
|
||||
mock_client.delete.return_value = 1
|
||||
|
||||
result = store.delete("mykey")
|
||||
|
||||
mock_client.delete.assert_called_once_with(["praison:mykey"])
|
||||
assert result is True
|
||||
|
||||
def test_delete_missing(self):
|
||||
"""Test delete returns False when client.delete returns 0."""
|
||||
store, mock_client = _make_state_store()
|
||||
mock_client.delete.return_value = 0
|
||||
|
||||
result = store.delete("absent")
|
||||
|
||||
assert result is False
|
||||
|
||||
def test_exists_true(self):
|
||||
"""Test exists returns True when client.exists returns 1."""
|
||||
store, mock_client = _make_state_store()
|
||||
mock_client.exists.return_value = 1
|
||||
|
||||
assert store.exists("mykey") is True
|
||||
mock_client.exists.assert_called_once_with(["praison:mykey"])
|
||||
|
||||
def test_exists_false(self):
|
||||
"""Test exists returns False when client.exists returns 0."""
|
||||
store, mock_client = _make_state_store()
|
||||
mock_client.exists.return_value = 0
|
||||
|
||||
assert store.exists("absent") is False
|
||||
|
||||
def test_keys(self):
|
||||
"""Test keys returns stripped key names via client.scan."""
|
||||
store, mock_client = _make_state_store()
|
||||
# scan() returns [cursor, [keys]]; cursor b"0" signals end of iteration
|
||||
mock_client.scan.return_value = [b"0", [b"praison:key1", b"praison:key2"]]
|
||||
|
||||
result = store.keys()
|
||||
|
||||
mock_client.scan.assert_called_once()
|
||||
assert result == ["key1", "key2"]
|
||||
|
||||
def test_ttl_with_value(self):
|
||||
"""Test ttl returns seconds remaining when client.ttl returns a positive value."""
|
||||
store, mock_client = _make_state_store()
|
||||
mock_client.ttl.return_value = 30
|
||||
|
||||
result = store.ttl("mykey")
|
||||
|
||||
assert result == 30
|
||||
|
||||
def test_ttl_no_ttl(self):
|
||||
"""Test ttl returns None when client.ttl returns -1 (no expiry)."""
|
||||
store, mock_client = _make_state_store()
|
||||
mock_client.ttl.return_value = -1
|
||||
|
||||
result = store.ttl("mykey")
|
||||
|
||||
assert result is None
|
||||
|
||||
def test_expire(self):
|
||||
"""Test expire calls client.expire and returns its result."""
|
||||
store, mock_client = _make_state_store()
|
||||
mock_client.expire.return_value = True
|
||||
|
||||
result = store.expire("mykey", 600)
|
||||
|
||||
mock_client.expire.assert_called_once_with("praison:mykey", 600)
|
||||
assert result is True
|
||||
|
||||
def test_hget(self):
|
||||
"""Test hget returns a parsed dict when client.hget returns JSON bytes."""
|
||||
store, mock_client = _make_state_store()
|
||||
data = {"field_val": 42}
|
||||
mock_client.hget.return_value = json.dumps(data).encode("utf-8")
|
||||
|
||||
result = store.hget("myhash", "field1")
|
||||
|
||||
mock_client.hget.assert_called_once_with("praison:myhash", "field1")
|
||||
assert result == data
|
||||
|
||||
def test_hset(self):
|
||||
"""Test hset calls client.hset with the prefixed key and a mapping dict."""
|
||||
store, mock_client = _make_state_store()
|
||||
mock_client.hset.return_value = 1
|
||||
|
||||
store.hset("myhash", "field1", "value1")
|
||||
|
||||
mock_client.hset.assert_called_once()
|
||||
call_args = mock_client.hset.call_args[0]
|
||||
assert call_args[0] == "praison:myhash"
|
||||
assert isinstance(call_args[1], dict)
|
||||
assert "field1" in call_args[1]
|
||||
|
||||
def test_hgetall(self):
|
||||
"""Test hgetall returns a decoded dict when client.hgetall returns bytes."""
|
||||
store, mock_client = _make_state_store()
|
||||
mock_client.hgetall.return_value = {
|
||||
b"field1": b'"value1"',
|
||||
b"field2": b"42",
|
||||
}
|
||||
|
||||
result = store.hgetall("myhash")
|
||||
|
||||
mock_client.hgetall.assert_called_once_with("praison:myhash")
|
||||
assert "field1" in result
|
||||
assert "field2" in result
|
||||
|
||||
def test_hdel(self):
|
||||
"""Test hdel calls client.hdel and returns the count of deleted fields."""
|
||||
store, mock_client = _make_state_store()
|
||||
mock_client.hdel.return_value = 2
|
||||
|
||||
result = store.hdel("myhash", "field1", "field2")
|
||||
|
||||
assert result == 2
|
||||
|
||||
def test_incr(self):
|
||||
"""Test incr calls client.incrby and returns the new value."""
|
||||
store, mock_client = _make_state_store()
|
||||
mock_client.incrby.return_value = 5
|
||||
|
||||
result = store.incr("counter")
|
||||
|
||||
assert result == 5
|
||||
mock_client.incrby.assert_called_once_with("praison:counter", 1)
|
||||
|
||||
def test_decr(self):
|
||||
"""Test decr calls client.decrby and returns the new value."""
|
||||
store, mock_client = _make_state_store()
|
||||
mock_client.decrby.return_value = 3
|
||||
|
||||
result = store.decr("counter")
|
||||
|
||||
assert result == 3
|
||||
mock_client.decrby.assert_called_once_with("praison:counter", 1)
|
||||
|
||||
def test_close_closes_client(self):
|
||||
"""Test close() calls client.close() and clears the reference."""
|
||||
store, mock_client = _make_state_store()
|
||||
|
||||
store.close()
|
||||
|
||||
mock_client.close.assert_called_once()
|
||||
assert store._client is None
|
||||
|
||||
|
||||
class TestRegistryValkey:
|
||||
"""Tests for Valkey backend registration in store registries."""
|
||||
|
||||
def test_knowledge_store_valkey_registered(self):
|
||||
"""Test KNOWLEDGE_STORES registry contains 'valkey'."""
|
||||
assert "valkey" in KNOWLEDGE_STORES.list_registered()
|
||||
|
||||
def test_state_store_valkey_registered(self):
|
||||
"""Test STATE_STORES registry contains 'valkey'."""
|
||||
assert "valkey" in STATE_STORES.list_registered()
|
||||
@@ -6,6 +6,7 @@ Heavy implementations that follow StorageBackendProtocol for:
|
||||
- MongoDB (praisonai[mongodb])
|
||||
- PostgreSQL (praisonai[postgresql])
|
||||
- DynamoDB (praisonai[dynamodb])
|
||||
- Valkey (praisonai[valkey])
|
||||
|
||||
These implementations are kept in the wrapper to avoid bloating the core SDK.
|
||||
"""
|
||||
@@ -13,9 +14,12 @@ These implementations are kept in the wrapper to avoid bloating the core SDK.
|
||||
# Lazy imports - only import when needed
|
||||
__all__ = [
|
||||
"RedisStorageAdapter",
|
||||
"MongoDBStorageAdapter",
|
||||
"MongoDBStorageAdapter",
|
||||
"PostgreSQLStorageAdapter",
|
||||
"DynamoDBStorageAdapter",
|
||||
"ValkeyStorageAdapter",
|
||||
"ValkeySearchBackend",
|
||||
"ValkeyBackend",
|
||||
]
|
||||
|
||||
def __getattr__(name: str):
|
||||
@@ -32,5 +36,40 @@ def __getattr__(name: str):
|
||||
elif name == "DynamoDBStorageAdapter":
|
||||
from .dynamodb_adapter import DynamoDBStorageAdapter
|
||||
return DynamoDBStorageAdapter
|
||||
elif name == "ValkeyStorageAdapter":
|
||||
from .valkey_adapter import ValkeyStorageAdapter
|
||||
return ValkeyStorageAdapter
|
||||
elif name == "ValkeySearchBackend":
|
||||
from .valkey_adapter import ValkeySearchBackend
|
||||
return ValkeySearchBackend
|
||||
elif name == "ValkeyBackend":
|
||||
cls = _make_valkey_backend_class()
|
||||
globals()["ValkeyBackend"] = cls
|
||||
return cls
|
||||
else:
|
||||
raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
|
||||
raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
|
||||
|
||||
|
||||
def _make_valkey_backend_class():
|
||||
"""Return ValkeyBackend — ValkeyStorageAdapter pre-configured from env vars."""
|
||||
import os
|
||||
from .valkey_adapter import ValkeyStorageAdapter
|
||||
|
||||
class ValkeyBackend(ValkeyStorageAdapter):
|
||||
def __init__(
|
||||
self,
|
||||
host: str = None,
|
||||
port: int = None,
|
||||
prefix: str = None,
|
||||
ttl=None,
|
||||
password: str = None,
|
||||
):
|
||||
super().__init__(
|
||||
host=host if host is not None else os.environ.get("VALKEY_HOST", "localhost"),
|
||||
port=int(port if port is not None else os.environ.get("VALKEY_PORT", 6379)),
|
||||
prefix=prefix if prefix is not None else os.environ.get("VALKEY_PREFIX", "praisonai:"),
|
||||
ttl=ttl,
|
||||
password=password if password is not None else os.environ.get("VALKEY_PASSWORD"),
|
||||
)
|
||||
|
||||
return ValkeyBackend
|
||||
@@ -0,0 +1,348 @@
|
||||
"""
|
||||
Valkey Storage Adapter for PraisonAI.
|
||||
|
||||
Implements StorageBackendProtocol using Valkey for high-speed storage.
|
||||
This is the wrapper implementation that contains the heavy valkey-glide-sync dependency.
|
||||
"""
|
||||
|
||||
import json
|
||||
import struct
|
||||
from typing import Dict, Any, List, Optional
|
||||
|
||||
from praisonai.persistence._valkey_client import create_valkey_client
|
||||
|
||||
try:
|
||||
from glide_sync import ExpirySet, ExpiryType
|
||||
except ImportError:
|
||||
ExpirySet = None
|
||||
ExpiryType = None
|
||||
|
||||
try:
|
||||
from glide_sync import ft
|
||||
from glide_sync import (
|
||||
VectorField, VectorFieldAttributesHnsw, VectorAlgorithm,
|
||||
DistanceMetricType, VectorType, FtCreateOptions, DataType,
|
||||
FtSearchOptions, ReturnField,
|
||||
)
|
||||
except ImportError:
|
||||
ft = None
|
||||
VectorField = VectorFieldAttributesHnsw = VectorAlgorithm = None
|
||||
DistanceMetricType = VectorType = FtCreateOptions = DataType = None
|
||||
FtSearchOptions = ReturnField = None
|
||||
|
||||
|
||||
class ValkeyStorageAdapter:
|
||||
"""
|
||||
Valkey-based storage backend adapter.
|
||||
|
||||
Uses Valkey for high-speed caching and ephemeral data storage.
|
||||
Implements StorageBackendProtocol from praisonaiagents.storage.protocols.
|
||||
|
||||
Example:
|
||||
```python
|
||||
from praisonai.storage import ValkeyStorageAdapter
|
||||
|
||||
adapter = ValkeyStorageAdapter(host="localhost", port=6379)
|
||||
adapter.save("session_123", {"messages": []})
|
||||
data = adapter.load("session_123")
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
host: str = "localhost",
|
||||
port: int = 6379,
|
||||
prefix: str = "praisonai:",
|
||||
ttl: Optional[int] = None,
|
||||
password: Optional[str] = None,
|
||||
):
|
||||
"""
|
||||
Initialize the Valkey storage adapter.
|
||||
|
||||
Args:
|
||||
host: Valkey server host
|
||||
port: Valkey server port
|
||||
prefix: Key prefix for all stored data
|
||||
ttl: Optional TTL in seconds for all keys
|
||||
password: Optional Valkey password
|
||||
"""
|
||||
self.host = host
|
||||
self.port = port
|
||||
self.prefix = prefix
|
||||
self.ttl = ttl
|
||||
self.password = password
|
||||
self._client = None
|
||||
|
||||
def _get_client(self):
|
||||
"""Lazy initialize Valkey client."""
|
||||
if self._client is None:
|
||||
self._client = create_valkey_client(
|
||||
host=self.host,
|
||||
port=self.port,
|
||||
password=self.password,
|
||||
)
|
||||
return self._client
|
||||
|
||||
def _make_key(self, key: str) -> str:
|
||||
"""Create prefixed key."""
|
||||
return f"{self.prefix}{key}"
|
||||
|
||||
def _scan_keys(self, pattern: str) -> List[bytes]:
|
||||
"""Collect all keys matching pattern via SCAN."""
|
||||
client = self._get_client()
|
||||
cursor: bytes = b"0"
|
||||
all_keys: List[bytes] = []
|
||||
while True:
|
||||
result = client.scan(cursor, match=pattern, count=100)
|
||||
cursor = result[0]
|
||||
all_keys.extend(result[1] or [])
|
||||
if not cursor or cursor in (b"0", "0", 0):
|
||||
break
|
||||
return all_keys
|
||||
|
||||
def save(self, key: str, data: Dict[str, Any]) -> None:
|
||||
"""Save data with the given key."""
|
||||
client = self._get_client()
|
||||
full_key = self._make_key(key)
|
||||
json_data = json.dumps(data, default=str, ensure_ascii=False).encode('utf-8')
|
||||
|
||||
try:
|
||||
if self.ttl is not None:
|
||||
if ExpirySet is None or ExpiryType is None:
|
||||
from praisonai.persistence._valkey_client import _MISSING_MSG
|
||||
raise ImportError(_MISSING_MSG)
|
||||
client.set(full_key, json_data, expiry=ExpirySet(ExpiryType.SEC, self.ttl))
|
||||
else:
|
||||
client.set(full_key, json_data)
|
||||
except ImportError:
|
||||
raise
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Failed to save data to Valkey: {e}") from e
|
||||
|
||||
def load(self, key: str) -> Optional[Dict[str, Any]]:
|
||||
"""Load data by key."""
|
||||
client = self._get_client()
|
||||
full_key = self._make_key(key)
|
||||
|
||||
try:
|
||||
value = client.get(full_key)
|
||||
if value is not None:
|
||||
try:
|
||||
if isinstance(value, bytes):
|
||||
value = value.decode('utf-8')
|
||||
return json.loads(value)
|
||||
except json.JSONDecodeError as e:
|
||||
raise ValueError(f"Invalid JSON data for key '{key}': {e}") from e
|
||||
return None
|
||||
except ValueError:
|
||||
raise
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Failed to load data from Valkey: {e}") from e
|
||||
|
||||
def delete(self, key: str) -> bool:
|
||||
"""Delete data by key."""
|
||||
client = self._get_client()
|
||||
full_key = self._make_key(key)
|
||||
|
||||
try:
|
||||
return client.delete([full_key]) > 0
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Failed to delete data from Valkey: {e}") from e
|
||||
|
||||
def list_keys(self, prefix: str = "") -> List[str]:
|
||||
"""List all keys, optionally filtered by prefix."""
|
||||
pattern = self._make_key(f"{prefix}*")
|
||||
|
||||
try:
|
||||
raw_keys = self._scan_keys(pattern)
|
||||
prefix_len = len(self.prefix)
|
||||
return sorted(
|
||||
(k.decode('utf-8') if isinstance(k, bytes) else k)[prefix_len:]
|
||||
for k in raw_keys
|
||||
)
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Failed to list keys from Valkey: {e}") from e
|
||||
|
||||
def exists(self, key: str) -> bool:
|
||||
"""Check if a key exists."""
|
||||
client = self._get_client()
|
||||
full_key = self._make_key(key)
|
||||
|
||||
try:
|
||||
return client.exists([full_key]) > 0
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Failed to check key existence in Valkey: {e}") from e
|
||||
|
||||
def clear(self) -> int:
|
||||
"""Clear all data with our prefix. Returns number of items deleted."""
|
||||
client = self._get_client()
|
||||
pattern = self._make_key("*")
|
||||
batch_size = 500
|
||||
|
||||
try:
|
||||
all_keys = self._scan_keys(pattern)
|
||||
if not all_keys:
|
||||
return 0
|
||||
total = 0
|
||||
for i in range(0, len(all_keys), batch_size):
|
||||
total += client.delete(all_keys[i:i + batch_size])
|
||||
return total
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Failed to clear data from Valkey: {e}") from e
|
||||
|
||||
def set_ttl(self, key: str, ttl: int) -> bool:
|
||||
"""Set TTL on a specific key."""
|
||||
client = self._get_client()
|
||||
full_key = self._make_key(key)
|
||||
|
||||
try:
|
||||
return client.expire(full_key, ttl)
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Failed to set TTL in Valkey: {e}") from e
|
||||
|
||||
def ping(self) -> bool:
|
||||
"""Test connection to Valkey."""
|
||||
try:
|
||||
client = self._get_client()
|
||||
result = client.ping()
|
||||
return result is not None
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
def close(self) -> None:
|
||||
"""Close the Valkey connection."""
|
||||
if self._client:
|
||||
try:
|
||||
self._client.close()
|
||||
finally:
|
||||
self._client = None
|
||||
|
||||
|
||||
class ValkeySearchBackend:
|
||||
"""
|
||||
Valkey vector search backend using RediSearch / ValkeySearch FT commands.
|
||||
|
||||
Example:
|
||||
```python
|
||||
from praisonai.storage import ValkeySearchBackend
|
||||
|
||||
backend = ValkeySearchBackend(host="localhost", port=6379, vector_dim=1536)
|
||||
backend.create_index()
|
||||
backend.add_document("doc1", "hello world", embedding)
|
||||
results = backend.search(query_embedding, k=5)
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
host: str = "localhost",
|
||||
port: int = 6379,
|
||||
index_name: str = "praisonai_vectors",
|
||||
password: Optional[str] = None,
|
||||
vector_dim: int = 1536,
|
||||
):
|
||||
"""
|
||||
Initialize the Valkey search backend.
|
||||
|
||||
Args:
|
||||
host: Valkey server host
|
||||
port: Valkey server port
|
||||
index_name: Name of the FT index
|
||||
password: Optional Valkey password
|
||||
vector_dim: Dimension of embedding vectors
|
||||
"""
|
||||
self.host = host
|
||||
self.port = port
|
||||
self.index_name = index_name
|
||||
self.password = password
|
||||
self.vector_dim = vector_dim
|
||||
self._client = None
|
||||
|
||||
def _get_client(self):
|
||||
"""Lazy initialize Valkey client."""
|
||||
if self._client is None:
|
||||
self._client = create_valkey_client(
|
||||
host=self.host,
|
||||
port=self.port,
|
||||
password=self.password,
|
||||
)
|
||||
return self._client
|
||||
|
||||
def create_index(self, vector_dim: int = None) -> None:
|
||||
"""Create the vector search index (no-op if already exists)."""
|
||||
client = self._get_client()
|
||||
dim = vector_dim or self.vector_dim
|
||||
|
||||
try:
|
||||
schema = [VectorField(
|
||||
"embedding",
|
||||
VectorAlgorithm.HNSW,
|
||||
VectorFieldAttributesHnsw(dim, DistanceMetricType.COSINE, VectorType.FLOAT32),
|
||||
)]
|
||||
options = FtCreateOptions(DataType.HASH, prefixes=[f"{self.index_name}:"])
|
||||
ft.create(client, self.index_name, schema, options)
|
||||
except Exception as e:
|
||||
err = str(e)
|
||||
if "already exists" not in err.lower() and "index already exists" not in err.lower():
|
||||
raise RuntimeError(f"Failed to create Valkey search index: {e}") from e
|
||||
|
||||
def add_document(self, doc_id: str, text: str, embedding: List[float]) -> None:
|
||||
"""Add a document with its embedding to the index."""
|
||||
if len(embedding) != self.vector_dim:
|
||||
raise ValueError(
|
||||
f"Embedding dimension mismatch: expected {self.vector_dim}, got {len(embedding)}"
|
||||
)
|
||||
client = self._get_client()
|
||||
full_id = f"{self.index_name}:{doc_id}"
|
||||
packed = struct.pack(f"{len(embedding)}f", *embedding)
|
||||
|
||||
try:
|
||||
client.hset(full_id, {"text": text, "embedding": packed})
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Failed to add document to Valkey: {e}") from e
|
||||
|
||||
def search(self, query_embedding: List[float], k: int = 5) -> List[Dict[str, Any]]:
|
||||
"""Search for nearest neighbors by embedding."""
|
||||
if k <= 0:
|
||||
raise ValueError("k must be > 0")
|
||||
if len(query_embedding) != self.vector_dim:
|
||||
raise ValueError(
|
||||
f"Query embedding dimension mismatch: expected {self.vector_dim}, got {len(query_embedding)}"
|
||||
)
|
||||
client = self._get_client()
|
||||
packed = struct.pack(f"{len(query_embedding)}f", *query_embedding)
|
||||
|
||||
try:
|
||||
options = FtSearchOptions(
|
||||
return_fields=[ReturnField("text"), ReturnField("score")],
|
||||
params={"vec": packed},
|
||||
)
|
||||
raw = ft.search(client, self.index_name, f"*=>[KNN {k} @embedding $vec AS score]", options)
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Failed to search Valkey: {e}") from e
|
||||
|
||||
# ft.search returns [count, {doc_id: {field: value, ...}, ...}]
|
||||
if not raw or len(raw) < 2:
|
||||
return []
|
||||
|
||||
results = []
|
||||
key_prefix = f"{self.index_name}:"
|
||||
for doc_id, fields in raw[1].items():
|
||||
doc: Dict[str, Any] = {}
|
||||
raw_id = doc_id.decode('utf-8') if isinstance(doc_id, bytes) else doc_id
|
||||
doc["id"] = raw_id.removeprefix(key_prefix)
|
||||
for fname, fval in fields.items():
|
||||
fname = fname.decode('utf-8') if isinstance(fname, bytes) else fname
|
||||
fval = fval.decode('utf-8') if isinstance(fval, bytes) else fval
|
||||
doc[fname] = fval
|
||||
results.append(doc)
|
||||
|
||||
return results
|
||||
|
||||
def close(self) -> None:
|
||||
"""Close the Valkey connection."""
|
||||
if self._client:
|
||||
try:
|
||||
self._client.close()
|
||||
finally:
|
||||
self._client = None
|
||||
@@ -180,6 +180,9 @@ postgresql = [
|
||||
dynamodb = [
|
||||
"boto3>=1.35.0",
|
||||
]
|
||||
valkey = [
|
||||
"valkey-glide-sync>=2.3.1",
|
||||
]
|
||||
# Cloud-native serverless databases
|
||||
neon = [
|
||||
"psycopg2-binary>=2.9.0",
|
||||
|
||||
Whitespace-only changes.
@@ -0,0 +1,249 @@
|
||||
"""
|
||||
Tests for ValkeyStorageAdapter and ValkeySearchBackend.
|
||||
|
||||
Tests cover:
|
||||
- ValkeyStorageAdapter CRUD operations (mocked glide client)
|
||||
- ValkeySearchBackend index and search operations (mocked)
|
||||
"""
|
||||
|
||||
import json
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
|
||||
def _make_adapter(prefix="praisonai:", ttl=None):
|
||||
"""Return a ValkeyStorageAdapter with a pre-injected MagicMock client."""
|
||||
from praisonai.storage.valkey_adapter import ValkeyStorageAdapter
|
||||
|
||||
adapter = ValkeyStorageAdapter.__new__(ValkeyStorageAdapter)
|
||||
adapter.host = "localhost"
|
||||
adapter.port = 6379
|
||||
adapter.prefix = prefix
|
||||
adapter.ttl = ttl
|
||||
adapter.password = None
|
||||
|
||||
mock_client = MagicMock()
|
||||
adapter._client = mock_client
|
||||
adapter._get_client = lambda: mock_client
|
||||
|
||||
return adapter, mock_client
|
||||
|
||||
|
||||
def _make_search_backend(index_name="praisonai_vectors", vector_dim=4):
|
||||
"""Return a ValkeySearchBackend with a pre-injected MagicMock client."""
|
||||
from praisonai.storage.valkey_adapter import ValkeySearchBackend
|
||||
|
||||
backend = ValkeySearchBackend.__new__(ValkeySearchBackend)
|
||||
backend.host = "localhost"
|
||||
backend.port = 6379
|
||||
backend.index_name = index_name
|
||||
backend.vector_dim = vector_dim
|
||||
backend.password = None
|
||||
|
||||
mock_client = MagicMock()
|
||||
backend._client = mock_client
|
||||
backend._get_client = lambda: mock_client
|
||||
|
||||
return backend, mock_client
|
||||
|
||||
|
||||
class TestValkeyStorageAdapter:
|
||||
"""Tests for ValkeyStorageAdapter."""
|
||||
|
||||
def test_valkey_adapter_save(self):
|
||||
"""Test save calls client.set with prefixed key and JSON bytes."""
|
||||
adapter, mock_client = _make_adapter()
|
||||
|
||||
data = {"name": "test", "value": 42}
|
||||
adapter.save("mykey", data)
|
||||
|
||||
mock_client.set.assert_called_once()
|
||||
call_args = mock_client.set.call_args[0]
|
||||
assert call_args[0] == "praisonai:mykey"
|
||||
assert json.loads(call_args[1]) == data
|
||||
|
||||
def test_valkey_adapter_save_with_ttl(self):
|
||||
"""Test save with TTL calls client.set with expiry kwarg (ImportError tolerated without glide)."""
|
||||
adapter, mock_client = _make_adapter(ttl=300)
|
||||
|
||||
# ExpirySet/ExpiryType are lazy-imported inside save(); without a live
|
||||
# glide install this raises RuntimeError — both paths are valid in CI.
|
||||
try:
|
||||
adapter.save("ttlkey", {"key": "val"})
|
||||
mock_client.set.assert_called_once()
|
||||
_, kwargs = mock_client.set.call_args
|
||||
assert "expiry" in kwargs
|
||||
except (ImportError, RuntimeError):
|
||||
pass
|
||||
|
||||
def test_valkey_adapter_load_existing(self):
|
||||
"""Test load returns a dict when client.get returns JSON bytes."""
|
||||
adapter, mock_client = _make_adapter()
|
||||
|
||||
data = {"result": "ok"}
|
||||
mock_client.get.return_value = json.dumps(data).encode("utf-8")
|
||||
|
||||
loaded = adapter.load("mykey")
|
||||
|
||||
mock_client.get.assert_called_once_with("praisonai:mykey")
|
||||
assert loaded == data
|
||||
|
||||
def test_valkey_adapter_load_missing(self):
|
||||
"""Test load returns None when client.get returns None."""
|
||||
adapter, mock_client = _make_adapter()
|
||||
mock_client.get.return_value = None
|
||||
|
||||
loaded = adapter.load("missing")
|
||||
|
||||
assert loaded is None
|
||||
|
||||
def test_valkey_adapter_delete_existing(self):
|
||||
"""Test delete returns True when client.delete returns 1."""
|
||||
adapter, mock_client = _make_adapter()
|
||||
mock_client.delete.return_value = 1
|
||||
|
||||
result = adapter.delete("mykey")
|
||||
|
||||
mock_client.delete.assert_called_once_with(["praisonai:mykey"])
|
||||
assert result is True
|
||||
|
||||
def test_valkey_adapter_delete_missing(self):
|
||||
"""Test delete returns False when client.delete returns 0."""
|
||||
adapter, mock_client = _make_adapter()
|
||||
mock_client.delete.return_value = 0
|
||||
|
||||
result = adapter.delete("absent")
|
||||
|
||||
assert result is False
|
||||
|
||||
def test_valkey_adapter_list_keys(self):
|
||||
"""Test list_keys returns stripped, sorted key names via client.scan."""
|
||||
adapter, mock_client = _make_adapter()
|
||||
raw_keys = [b"praisonai:beta", b"praisonai:alpha", b"praisonai:gamma"]
|
||||
# scan() returns [cursor, [keys]]; cursor b"0" signals end of iteration
|
||||
mock_client.scan.return_value = [b"0", raw_keys]
|
||||
|
||||
keys = adapter.list_keys()
|
||||
|
||||
mock_client.scan.assert_called_once()
|
||||
assert keys == ["alpha", "beta", "gamma"]
|
||||
|
||||
def test_valkey_adapter_exists_true(self):
|
||||
"""Test exists returns True when client.exists returns 1."""
|
||||
adapter, mock_client = _make_adapter()
|
||||
mock_client.exists.return_value = 1
|
||||
|
||||
assert adapter.exists("mykey") is True
|
||||
mock_client.exists.assert_called_once_with(["praisonai:mykey"])
|
||||
|
||||
def test_valkey_adapter_exists_false(self):
|
||||
"""Test exists returns False when client.exists returns 0."""
|
||||
adapter, mock_client = _make_adapter()
|
||||
mock_client.exists.return_value = 0
|
||||
|
||||
assert adapter.exists("absent") is False
|
||||
|
||||
def test_valkey_adapter_clear(self):
|
||||
"""Test clear deletes all prefixed keys and returns count."""
|
||||
adapter, mock_client = _make_adapter()
|
||||
raw_keys = [b"praisonai:a", b"praisonai:b", b"praisonai:c"]
|
||||
# scan() returns [cursor, [keys]]; cursor b"0" signals end of iteration
|
||||
mock_client.scan.return_value = [b"0", raw_keys]
|
||||
mock_client.delete.return_value = 3
|
||||
|
||||
count = adapter.clear()
|
||||
|
||||
assert count == 3
|
||||
mock_client.delete.assert_called_once()
|
||||
|
||||
def test_valkey_adapter_set_ttl(self):
|
||||
"""Test set_ttl calls client.expire with prefixed key and returns result."""
|
||||
adapter, mock_client = _make_adapter()
|
||||
mock_client.expire.return_value = True
|
||||
|
||||
result = adapter.set_ttl("mykey", 600)
|
||||
|
||||
mock_client.expire.assert_called_once_with("praisonai:mykey", 600)
|
||||
assert result is True
|
||||
|
||||
def test_valkey_adapter_import_error(self):
|
||||
"""Test ImportError with 'valkey' hint when glide package is not available."""
|
||||
import praisonai.persistence._valkey_client as client_mod
|
||||
from praisonai.storage.valkey_adapter import ValkeyStorageAdapter
|
||||
|
||||
adapter = ValkeyStorageAdapter.__new__(ValkeyStorageAdapter)
|
||||
adapter.host = "localhost"
|
||||
adapter.port = 6379
|
||||
adapter.prefix = "praisonai:"
|
||||
adapter.ttl = None
|
||||
adapter.password = None
|
||||
adapter._client = None
|
||||
|
||||
original = client_mod.GlideClientSync
|
||||
try:
|
||||
client_mod.GlideClientSync = None
|
||||
try:
|
||||
adapter._get_client()
|
||||
assert False, "Should have raised ImportError"
|
||||
except ImportError as e:
|
||||
assert "valkey" in str(e).lower()
|
||||
finally:
|
||||
client_mod.GlideClientSync = original
|
||||
|
||||
def test_valkey_adapter_operation_error(self):
|
||||
"""Test RuntimeError raised when a client operation fails."""
|
||||
adapter, mock_client = _make_adapter()
|
||||
mock_client.get.side_effect = Exception("connection refused")
|
||||
|
||||
try:
|
||||
adapter.load("anykey")
|
||||
assert False, "Should have raised RuntimeError"
|
||||
except RuntimeError as e:
|
||||
assert "Valkey" in str(e) or "valkey" in str(e).lower()
|
||||
|
||||
|
||||
class TestValkeySearchBackend:
|
||||
"""Tests for ValkeySearchBackend."""
|
||||
|
||||
def test_create_index(self):
|
||||
"""Test create_index uses ft.create (no-op if already exists)."""
|
||||
backend, _ = _make_search_backend()
|
||||
mock_ft = MagicMock()
|
||||
mock_ft.create.return_value = b"OK"
|
||||
|
||||
with patch("praisonai.storage.valkey_adapter.ft", mock_ft):
|
||||
backend.create_index()
|
||||
|
||||
mock_ft.create.assert_called_once()
|
||||
assert backend.index_name == mock_ft.create.call_args[0][1]
|
||||
|
||||
def test_add_document(self):
|
||||
"""Test add_document calls hset with correct key and fields."""
|
||||
backend, mock_client = _make_search_backend()
|
||||
mock_client.hset.return_value = 2
|
||||
|
||||
backend.add_document("doc1", "hello world", [0.1, 0.2, 0.3, 0.4])
|
||||
|
||||
mock_client.hset.assert_called_once()
|
||||
call_args = mock_client.hset.call_args[0]
|
||||
assert call_args[0] == "praisonai_vectors:doc1"
|
||||
mapping = call_args[1]
|
||||
assert "text" in mapping
|
||||
assert mapping["text"] == "hello world"
|
||||
assert "embedding" in mapping
|
||||
|
||||
def test_search(self):
|
||||
"""Test search uses ft.search and returns a list of dicts."""
|
||||
backend, _ = _make_search_backend()
|
||||
ft_result = [1, {b"praisonai_vectors:doc1": {b"text": b"hello world", b"score": b"0.05"}}]
|
||||
mock_ft = MagicMock()
|
||||
mock_ft.search.return_value = ft_result
|
||||
|
||||
with patch("praisonai.storage.valkey_adapter.ft", mock_ft):
|
||||
results = backend.search([0.1, 0.2, 0.3, 0.4], k=5)
|
||||
|
||||
assert isinstance(results, list)
|
||||
assert len(results) == 1
|
||||
assert results[0]["id"] == "doc1"
|
||||
assert results[0]["text"] == "hello world"
|
||||
mock_ft.search.assert_called_once()
|
||||
|
||||
Reference in new issue
Block a user