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]>

src/praisonai/praisonai/persistence/knowledge/valkey_vector.py

@@ -115,8 +115,12 @@ class ValkeyVectorKnowledgeStore(KnowledgeStore):
         index_name = self._index_name(name)
         try:
             ft.dropindex(client, index_name)
-        except Exception:
-            return False
+        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}:*"
@@ -139,8 +143,12 @@ class ValkeyVectorKnowledgeStore(KnowledgeStore):
         try:
             ft.info(client, index_name)
             return True
-        except Exception:
-            return False
+        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."""
@@ -158,8 +166,10 @@ class ValkeyVectorKnowledgeStore(KnowledgeStore):
                         name = name[:-4]
                     result.append(name)
             return result
-        except Exception:
-            return []
+        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,
@@ -182,6 +192,7 @@ class ValkeyVectorKnowledgeStore(KnowledgeStore):
                 "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)
@@ -218,6 +229,8 @@ class ValkeyVectorKnowledgeStore(KnowledgeStore):
             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)
@@ -246,6 +259,7 @@ class ValkeyVectorKnowledgeStore(KnowledgeStore):
                     ReturnField("content"),
                     ReturnField("content_hash"),
                     ReturnField("created_at"),
+                    ReturnField("metadata"),
                     ReturnField("score"),
                 ],
                 params={"vec": packed},
@@ -280,7 +294,7 @@ class ValkeyVectorKnowledgeStore(KnowledgeStore):
                 id=bare_id,
                 content=field_map.get("content", ""),
                 embedding=None,
-                metadata={},
+                metadata=json.loads(field_map.get("metadata", "{}") or "{}"),
                 content_hash=field_map.get("content_hash", ""),
                 created_at=float(field_map.get("created_at", 0)),
             ))
@@ -319,12 +333,14 @@ class ValkeyVectorKnowledgeStore(KnowledgeStore):
                 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=metadata,
                     content_hash=content_hash,
                     created_at=created_at,
                 )
@@ -365,8 +381,12 @@ class ValkeyVectorKnowledgeStore(KnowledgeStore):
                 if k in info:
                     return int(info[k])
             return 0
-        except Exception:
-            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."""

src/praisonai/praisonai/persistence/registry.py

@@ -201,8 +201,11 @@ def _register_builtin_knowledge_stores(registry: StoreRegistry):
 
     def _valkey_vector(url=None, **kwargs):
         from .knowledge.valkey_vector import ValkeyVectorKnowledgeStore
-        # url not used; host/port come from kwargs
-        kwargs.pop("url", None)
+        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):
@@ -296,7 +299,11 @@ def _register_builtin_state_stores(registry: StoreRegistry):
 
     def _valkey(url=None, **kwargs):
         from .state.valkey import ValkeyStateStore
-        kwargs.pop("url", None)
+        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):

src/praisonai/praisonai/persistence/state/valkey.py

@@ -79,14 +79,18 @@ class ValkeyStateStore(StateStore):
     def set(self, key: str, value: Any, ttl: Optional[int] = None) -> None:
         """Set a value with optional TTL."""
         client = self._get_client()
-        if not isinstance(value, str):
-            value = json.dumps(value)
+        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
 
@@ -168,8 +172,7 @@ class ValkeyStateStore(StateStore):
     def hset(self, key: str, field: str, value: Any) -> None:
         """Set a field in a hash."""
         client = self._get_client()
-        if not isinstance(value, str):
-            value = json.dumps(value)
+        value = json.dumps(value)
         try:
             client.hset(self._key(key), {field: value})
         except Exception as e:

src/praisonai/praisonai/storage/valkey_adapter.py

@@ -108,9 +108,14 @@ class ValkeyStorageAdapter:
 
         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
 
@@ -283,6 +288,10 @@ class ValkeySearchBackend:
 
     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)
@@ -294,6 +303,12 @@ class ValkeySearchBackend:
 
     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)