Compare commits

...
2 Commits
Author SHA1 Message Date
praisonai-triage-agent[bot]andMervin Praison 251470489f fix: address valkey integration issues from code review
- Fix JSON encoding consistency in ValkeyStateStore to prevent type confusion
- Add TTL guards for ExpirySet/ExpiryType to prevent cryptic RuntimeError
- Add warning for score threshold limitation (cosine distance only)
- Fix metadata persistence in KnowledgeDocument (now properly round-tripped)
- Improve exception handling - only catch specific 'index not found' errors
- Fix URL parameter handling in registry factories (reject with clear error)
- Add input validation for vector dimensions and k parameter

Co-authored-by: Mervin Praison <[email protected]>
2026-05-20 17:45:00 +00:00
Edward Liang dbf8b0d6e1 feat: Add Valkey Integration
Signed-off-by: Edward Liang <[email protected]>
2026-05-20 10:23:20 -07:00
14 changed files with 2051 additions and 13 deletions

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()
+41 -2
View File
@@ -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
+3
View File
@@ -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()