jaf-py 2.5.10__py3-none-any.whl → 2.5.12__py3-none-any.whl

jaf/memory/providers/redis.py

@@ -24,10 +24,12 @@ from ..types import (
 
 try:
     import redis.asyncio as redis
+
     RedisClient = redis.Redis
 except ImportError:
     RedisClient = Any
 
+
 class RedisProvider(MemoryProvider):
     """
     Redis implementation of MemoryProvider.
@@ -45,18 +47,20 @@ class RedisProvider(MemoryProvider):
     def _serialize(self, conversation: ConversationMemory) -> str:
         """Serialize conversation using shared utilities."""
         from ..utils import serialize_conversation_for_json
+
         return serialize_conversation_for_json(conversation)
 
     def _deserialize(self, data: str) -> ConversationMemory:
         """Deserialize conversation using shared utilities."""
         from ..utils import deserialize_conversation_from_json
+
         return deserialize_conversation_from_json(data)
 
     async def store_messages(
         self,
         conversation_id: str,
         messages: List[Message],
-        metadata: Optional[Dict[str, Any]] = None
+        metadata: Optional[Dict[str, Any]] = None,
     ) -> Result[None, MemoryStorageError]:
         try:
             now = datetime.now()
@@ -69,18 +73,21 @@ class RedisProvider(MemoryProvider):
                     "updated_at": now,
                     "total_messages": len(messages),
                     "last_activity": now,
-                    **(metadata or {})
-                }
+                    **(metadata or {}),
+                },
             )
             key = self._get_key(conversation_id)
             await self.redis_client.set(key, self._serialize(conversation), ex=self.config.ttl)
             return Success(None)
         except Exception as e:
-            return Failure(MemoryStorageError(operation="store_messages", provider="Redis", message=str(e), cause=e))
+            return Failure(
+                MemoryStorageError(
+                    operation="store_messages", provider="Redis", message=str(e), cause=e
+                )
+            )
 
     async def get_conversation(
-        self,
-        conversation_id: str
+        self, conversation_id: str
     ) -> Result[Optional[ConversationMemory], MemoryStorageError]:
         try:
             key = self._get_key(conversation_id)
@@ -94,13 +101,17 @@ class RedisProvider(MemoryProvider):
             await self.redis_client.set(key, self._serialize(conversation), ex=self.config.ttl)
             return Success(conversation)
         except Exception as e:
-            return Failure(MemoryStorageError(operation="get_conversation", provider="Redis", message=str(e), cause=e))
+            return Failure(
+                MemoryStorageError(
+                    operation="get_conversation", provider="Redis", message=str(e), cause=e
+                )
+            )
 
     async def append_messages(
         self,
         conversation_id: str,
         messages: List[Message],
-        metadata: Optional[Dict[str, Any]] = None
+        metadata: Optional[Dict[str, Any]] = None,
     ) -> Result[None, Union[MemoryNotFoundError, MemoryStorageError]]:
         result = await self.get_conversation(conversation_id)
         if isinstance(result, Failure):
@@ -108,7 +119,13 @@ class RedisProvider(MemoryProvider):
 
         existing = result.data
         if not existing:
-            return Failure(MemoryNotFoundError(conversation_id=conversation_id, provider="Redis", message=f"Conversation {conversation_id} not found"))
+            return Failure(
+                MemoryNotFoundError(
+                    conversation_id=conversation_id,
+                    provider="Redis",
+                    message=f"Conversation {conversation_id} not found",
+                )
+            )
 
         # Convert tuple back to list, append new messages, then store
         all_messages = list(existing.messages) + messages
@@ -117,14 +134,13 @@ class RedisProvider(MemoryProvider):
             "updated_at": datetime.now(),
             "last_activity": datetime.now(),
             "total_messages": len(all_messages),
-            **(metadata or {})
+            **(metadata or {}),
         }
 
         return await self.store_messages(conversation_id, all_messages, updated_metadata)
 
     async def find_conversations(
-        self,
-        query: MemoryQuery
+        self, query: MemoryQuery
     ) -> Result[List[ConversationMemory], MemoryStorageError]:
         try:
             keys = await self.redis_client.keys(f"{self.config.key_prefix}*")
@@ -144,12 +160,14 @@ class RedisProvider(MemoryProvider):
                     conversations.append(conv)
             return Success(conversations)
         except Exception as e:
-            return Failure(MemoryStorageError(operation="find_conversations", provider="Redis", message=str(e), cause=e))
+            return Failure(
+                MemoryStorageError(
+                    operation="find_conversations", provider="Redis", message=str(e), cause=e
+                )
+            )
 
     async def get_recent_messages(
-        self,
-        conversation_id: str,
-        limit: int = 50
+        self, conversation_id: str, limit: int = 50
     ) -> Result[List[Message], Union[MemoryNotFoundError, MemoryStorageError]]:
         result = await self.get_conversation(conversation_id)
         if isinstance(result, Failure):
@@ -157,36 +175,47 @@ class RedisProvider(MemoryProvider):
 
         conversation = result.data
         if not conversation:
-            return Failure(MemoryNotFoundError(conversation_id=conversation_id, provider="Redis", message=f"Conversation {conversation_id} not found"))
+            return Failure(
+                MemoryNotFoundError(
+                    conversation_id=conversation_id,
+                    provider="Redis",
+                    message=f"Conversation {conversation_id} not found",
+                )
+            )
 
         return Success(conversation.messages[-limit:])
 
-    async def delete_conversation(
-        self,
-        conversation_id: str
-    ) -> Result[bool, MemoryStorageError]:
+    async def delete_conversation(self, conversation_id: str) -> Result[bool, MemoryStorageError]:
         try:
             deleted = await self.redis_client.delete(self._get_key(conversation_id))
             return Success(deleted > 0)
         except Exception as e:
-            return Failure(MemoryStorageError(operation="delete_conversation", provider="Redis", message=str(e), cause=e))
+            return Failure(
+                MemoryStorageError(
+                    operation="delete_conversation", provider="Redis", message=str(e), cause=e
+                )
+            )
 
-    async def clear_user_conversations(
-        self,
-        user_id: str
-    ) -> Result[int, MemoryStorageError]:
+    async def clear_user_conversations(self, user_id: str) -> Result[int, MemoryStorageError]:
         # This is inefficient in Redis, consider a different approach for production
-        return Failure(MemoryStorageError(operation="clear_user_conversations", provider="Redis", message="clear_user_conversations not efficiently supported"))
+        return Failure(
+            MemoryStorageError(
+                operation="clear_user_conversations",
+                provider="Redis",
+                message="clear_user_conversations not efficiently supported",
+            )
+        )
 
     async def get_stats(
-        self,
-        user_id: Optional[str] = None
+        self, user_id: Optional[str] = None
     ) -> Result[Dict[str, Any], MemoryStorageError]:
         try:
             keys = await self.redis_client.keys(f"{self.config.key_prefix}*")
             return Success({"total_conversations": len(keys)})
         except Exception as e:
-            return Failure(MemoryStorageError(operation="get_stats", provider="Redis", message=str(e), cause=e))
+            return Failure(
+                MemoryStorageError(operation="get_stats", provider="Redis", message=str(e), cause=e)
+            )
 
     async def health_check(self) -> Result[Dict[str, Any], MemoryConnectionError]:
         start_time = datetime.now()
@@ -194,21 +223,23 @@ class RedisProvider(MemoryProvider):
             await self.redis_client.ping()
             latency_ms = (datetime.now() - start_time).total_seconds() * 1000
             db_size = await self.redis_client.dbsize()
-            return Success({
-                "healthy": True,
-                "latency_ms": latency_ms,
-                "provider": "Redis",
-                "details": {
-                    "db_size": db_size
+            return Success(
+                {
+                    "healthy": True,
+                    "latency_ms": latency_ms,
+                    "provider": "Redis",
+                    "details": {"db_size": db_size},
                 }
-            })
+            )
         except Exception as e:
-            return Failure(MemoryConnectionError(provider="Redis", message="Redis health check failed", cause=e))
+            return Failure(
+                MemoryConnectionError(
+                    provider="Redis", message="Redis health check failed", cause=e
+                )
+            )
 
     async def truncate_conversation_after(
-        self,
-        conversation_id: str,
-        message_id: MessageId
+        self, conversation_id: str, message_id: MessageId
     ) -> Result[int, Union[MemoryNotFoundError, MemoryStorageError]]:
         """
         Truncate conversation after (and including) the specified message ID.
@@ -219,27 +250,29 @@ class RedisProvider(MemoryProvider):
             conv_result = await self.get_conversation(conversation_id)
             if isinstance(conv_result, Failure):
                 return conv_result
-            
+
             if not conv_result.data:
-                return Failure(MemoryNotFoundError(
-                    message=f"Conversation {conversation_id} not found",
-                    provider="Redis",
-                    conversation_id=conversation_id
-                ))
-            
+                return Failure(
+                    MemoryNotFoundError(
+                        message=f"Conversation {conversation_id} not found",
+                        provider="Redis",
+                        conversation_id=conversation_id,
+                    )
+                )
+
             conversation = conv_result.data
             messages = list(conversation.messages)
             truncate_index = find_message_index(messages, message_id)
-            
+
             if truncate_index is None:
                 # Message not found, nothing to truncate
                 return Success(0)
-            
+
             # Truncate messages from the found index onwards
             original_count = len(messages)
             truncated_messages = messages[:truncate_index]
             removed_count = original_count - len(truncated_messages)
-            
+
             # Update conversation with truncated messages
             now = datetime.now()
             updated_metadata = {
@@ -249,35 +282,39 @@ class RedisProvider(MemoryProvider):
                 "total_messages": len(truncated_messages),
                 "regeneration_truncated": True,
                 "truncated_at": now.isoformat(),
-                "messages_removed": removed_count
+                "messages_removed": removed_count,
             }
-            
+
             # Store updated conversation
             updated_conversation = ConversationMemory(
                 conversation_id=conversation_id,
                 user_id=conversation.user_id,
                 messages=truncated_messages,
-                metadata=updated_metadata
+                metadata=updated_metadata,
             )
-            
+
             key = self._get_key(conversation_id)
-            await self.redis_client.set(key, self._serialize(updated_conversation), ex=self.config.ttl)
-            
-            print(f"[MEMORY:Redis] Truncated conversation {conversation_id}: removed {removed_count} messages after message {message_id}")
+            await self.redis_client.set(
+                key, self._serialize(updated_conversation), ex=self.config.ttl
+            )
+
+            print(
+                f"[MEMORY:Redis] Truncated conversation {conversation_id}: removed {removed_count} messages after message {message_id}"
+            )
             return Success(removed_count)
-            
+
         except Exception as e:
-            return Failure(MemoryStorageError(
-                message=f"Failed to truncate conversation: {e}",
-                provider="Redis",
-                operation="truncate_conversation_after",
-                cause=e
-            ))
+            return Failure(
+                MemoryStorageError(
+                    message=f"Failed to truncate conversation: {e}",
+                    provider="Redis",
+                    operation="truncate_conversation_after",
+                    cause=e,
+                )
+            )
 
     async def get_conversation_until_message(
-        self,
-        conversation_id: str,
-        message_id: MessageId
+        self, conversation_id: str, message_id: MessageId
     ) -> Result[Optional[ConversationMemory], Union[MemoryNotFoundError, MemoryStorageError]]:
         """
         Get conversation history up to (but not including) the specified message ID.
@@ -288,22 +325,24 @@ class RedisProvider(MemoryProvider):
             conv_result = await self.get_conversation(conversation_id)
             if isinstance(conv_result, Failure):
                 return conv_result
-            
+
             if not conv_result.data:
                 return Success(None)
-            
+
             conversation = conv_result.data
             messages = list(conversation.messages)
             until_index = find_message_index(messages, message_id)
-            
+
             if until_index is None:
                 # Message not found, return None as lightweight indicator
-                print(f"[MEMORY:Redis] Message {message_id} not found in conversation {conversation_id}")
+                print(
+                    f"[MEMORY:Redis] Message {message_id} not found in conversation {conversation_id}"
+                )
                 return Success(None)
-            
+
             # Return conversation up to (but not including) the specified message
             truncated_messages = messages[:until_index]
-            
+
             # Create a copy of the conversation with truncated messages
             truncated_conversation = ConversationMemory(
                 conversation_id=conversation.conversation_id,
@@ -314,26 +353,27 @@ class RedisProvider(MemoryProvider):
                     "truncated_for_regeneration": True,
                     "truncated_until_message": str(message_id),
                     "original_message_count": len(messages),
-                    "truncated_message_count": len(truncated_messages)
-                }
+                    "truncated_message_count": len(truncated_messages),
+                },
+            )
+
+            print(
+                f"[MEMORY:Redis] Retrieved conversation {conversation_id} until message {message_id}: {len(truncated_messages)} messages"
             )
-            
-            print(f"[MEMORY:Redis] Retrieved conversation {conversation_id} until message {message_id}: {len(truncated_messages)} messages")
             return Success(truncated_conversation)
-            
+
         except Exception as e:
-            return Failure(MemoryStorageError(
-                message=f"Failed to get conversation until message: {e}",
-                provider="Redis",
-                operation="get_conversation_until_message",
-                cause=e
-            ))
+            return Failure(
+                MemoryStorageError(
+                    message=f"Failed to get conversation until message: {e}",
+                    provider="Redis",
+                    operation="get_conversation_until_message",
+                    cause=e,
+                )
+            )
 
     async def mark_regeneration_point(
-        self,
-        conversation_id: str,
-        message_id: MessageId,
-        regeneration_metadata: Dict[str, Any]
+        self, conversation_id: str, message_id: MessageId, regeneration_metadata: Dict[str, Any]
     ) -> Result[None, Union[MemoryNotFoundError, MemoryStorageError]]:
         """
         Mark a regeneration point in the conversation for audit purposes.
@@ -343,64 +383,79 @@ class RedisProvider(MemoryProvider):
             conv_result = await self.get_conversation(conversation_id)
             if isinstance(conv_result, Failure):
                 return conv_result
-            
+
             if not conv_result.data:
-                return Failure(MemoryNotFoundError(
-                    message=f"Conversation {conversation_id} not found",
-                    provider="Redis",
-                    conversation_id=conversation_id
-                ))
-            
+                return Failure(
+                    MemoryNotFoundError(
+                        message=f"Conversation {conversation_id} not found",
+                        provider="Redis",
+                        conversation_id=conversation_id,
+                    )
+                )
+
             conversation = conv_result.data
-            
+
             # Add regeneration point to metadata
             regeneration_points = conversation.metadata.get("regeneration_points", [])
             regeneration_point = {
                 "message_id": str(message_id),
                 "timestamp": datetime.now().isoformat(),
-                **regeneration_metadata
+                **regeneration_metadata,
             }
             regeneration_points.append(regeneration_point)
-            
+
             # Update conversation metadata
             updated_metadata = {
                 **conversation.metadata,
                 "regeneration_points": regeneration_points,
                 "last_regeneration": regeneration_point,
                 "updated_at": datetime.now(),
-                "regeneration_count": len(regeneration_points)
+                "regeneration_count": len(regeneration_points),
             }
-            
+
             # Store updated conversation
             updated_conversation = ConversationMemory(
                 conversation_id=conversation.conversation_id,
                 user_id=conversation.user_id,
                 messages=conversation.messages,
-                metadata=updated_metadata
+                metadata=updated_metadata,
             )
-            
+
             key = self._get_key(conversation_id)
-            await self.redis_client.set(key, self._serialize(updated_conversation), ex=self.config.ttl)
-            
-            print(f"[MEMORY:Redis] Marked regeneration point for conversation {conversation_id} at message {message_id}")
+            await self.redis_client.set(
+                key, self._serialize(updated_conversation), ex=self.config.ttl
+            )
+
+            print(
+                f"[MEMORY:Redis] Marked regeneration point for conversation {conversation_id} at message {message_id}"
+            )
             return Success(None)
-            
+
         except Exception as e:
-            return Failure(MemoryStorageError(
-                message=f"Failed to mark regeneration point: {e}",
-                provider="Redis",
-                operation="mark_regeneration_point",
-                cause=e
-            ))
+            return Failure(
+                MemoryStorageError(
+                    message=f"Failed to mark regeneration point: {e}",
+                    provider="Redis",
+                    operation="mark_regeneration_point",
+                    cause=e,
+                )
+            )
 
     async def close(self) -> Result[None, MemoryConnectionError]:
         try:
             await self.redis_client.aclose()
             return Success(None)
         except Exception as e:
-            return Failure(MemoryConnectionError(provider="Redis", message="Failed to close Redis connection", cause=e))
+            return Failure(
+                MemoryConnectionError(
+                    provider="Redis", message="Failed to close Redis connection", cause=e
+                )
+            )
+
 
-async def create_redis_provider(config: RedisConfig) -> Result[RedisProvider, MemoryConnectionError]:
+async def create_redis_provider(
+    config: RedisConfig,
+) -> Result[RedisProvider, MemoryConnectionError]:
     try:
         # These will be passed to the Redis client constructor
         # and will override any values parsed from the URL.
@@ -421,4 +476,6 @@ async def create_redis_provider(config: RedisConfig) -> Result[RedisProvider, Me
         await redis_client.ping()
         return Success(RedisProvider(config, redis_client))
     except Exception as e:
-        return Failure(MemoryConnectionError(provider="Redis", message="Failed to connect to Redis", cause=e))
+        return Failure(
+            MemoryConnectionError(provider="Redis", message="Failed to connect to Redis", cause=e)
+        )