jaf-py 2.5.2__py3-none-any.whl → 2.5.4__py3-none-any.whl

jaf/memory/providers/in_memory.py

@@ -11,7 +11,7 @@ from collections import OrderedDict
 from datetime import datetime
 from typing import Any, Dict, List, Optional, Union
 
-from ...core.types import Message
+from ...core.types import Message, MessageId, find_message_index
 from ..types import (
     ConversationMemory,
     Failure,
@@ -307,6 +307,179 @@ class InMemoryProvider(MemoryProvider):
                 cause=e
             ))
 
+    async def truncate_conversation_after(
+        self,
+        conversation_id: str,
+        message_id: MessageId
+    ) -> Result[int, Union[MemoryNotFoundError, MemoryStorageError]]:
+        """
+        Truncate conversation after (and including) the specified message ID.
+        Returns the number of messages removed.
+        """
+        async with self._lock:
+            try:
+                conversation = self._conversations.get(conversation_id)
+                if conversation is None:
+                    return Failure(MemoryNotFoundError(
+                        message=f"Conversation {conversation_id} not found",
+                        provider="InMemory",
+                        conversation_id=conversation_id
+                    ))
+
+                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 = {
+                    **conversation.metadata,
+                    "updated_at": now,
+                    "last_activity": now,
+                    "total_messages": len(truncated_messages),
+                    "regeneration_truncated": True,
+                    "truncated_at": now.isoformat(),
+                    "messages_removed": removed_count
+                }
+                
+                updated_conversation = ConversationMemory(
+                    conversation_id=conversation_id,
+                    user_id=conversation.user_id,
+                    messages=truncated_messages,
+                    metadata=updated_metadata
+                )
+                
+                self._conversations[conversation_id] = updated_conversation
+                self._conversations.move_to_end(conversation_id)
+                
+                print(f"[MEMORY:InMemory] 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="InMemory",
+                    operation="truncate_conversation_after",
+                    cause=e
+                ))
+
+    async def get_conversation_until_message(
+        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.
+        Useful for regeneration scenarios.
+        """
+        async with self._lock:
+            try:
+                conversation = self._conversations.get(conversation_id)
+                if conversation is None:
+                    return Success(None)
+                
+                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:InMemory] 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,
+                    user_id=conversation.user_id,
+                    messages=truncated_messages,
+                    metadata={
+                        **conversation.metadata,
+                        "truncated_for_regeneration": True,
+                        "truncated_until_message": str(message_id),
+                        "original_message_count": len(messages),
+                        "truncated_message_count": len(truncated_messages)
+                    }
+                )
+                
+                print(f"[MEMORY:InMemory] Retrieved conversation {conversation_id} until message {message_id}: {len(truncated_messages)} messages (found at index {until_index})")
+                return Success(truncated_conversation)
+                
+            except Exception as e:
+                return Failure(MemoryStorageError(
+                    message=f"Failed to get conversation until message: {e}",
+                    provider="InMemory",
+                    operation="get_conversation_until_message",
+                    cause=e
+                ))
+
+    async def mark_regeneration_point(
+        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.
+        """
+        async with self._lock:
+            try:
+                conversation = self._conversations.get(conversation_id)
+                if conversation is None:
+                    return Failure(MemoryNotFoundError(
+                        message=f"Conversation {conversation_id} not found",
+                        provider="InMemory",
+                        conversation_id=conversation_id
+                    ))
+                
+                # 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_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)
+                }
+                
+                updated_conversation = ConversationMemory(
+                    conversation_id=conversation.conversation_id,
+                    user_id=conversation.user_id,
+                    messages=conversation.messages,
+                    metadata=updated_metadata
+                )
+                
+                self._conversations[conversation_id] = updated_conversation
+                self._conversations.move_to_end(conversation_id)
+                
+                print(f"[MEMORY:InMemory] 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="InMemory",
+                    operation="mark_regeneration_point",
+                    cause=e
+                ))
+
     async def close(self) -> Result[None, MemoryConnectionError]:
         """Close/cleanup the provider."""
         async with self._lock:

jaf/memory/providers/postgres.py

@@ -9,7 +9,7 @@ import json
 from datetime import datetime
 from typing import Any, Dict, List, Optional, Union
 
-from ...core.types import Message
+from ...core.types import Message, MessageId, find_message_index
 from ..types import (
     ConversationMemory,
     Failure,
@@ -239,6 +239,216 @@ class PostgresProvider(MemoryProvider):
         except Exception as e:
             return Failure(MemoryConnectionError(provider="Postgres", message="Postgres health check failed", cause=e))
 
+    async def truncate_conversation_after(
+        self,
+        conversation_id: str,
+        message_id: MessageId
+    ) -> Result[int, Union[MemoryNotFoundError, MemoryStorageError]]:
+        """
+        Truncate conversation after (and including) the specified message ID.
+        Returns the number of messages removed.
+        """
+        try:
+            # Get the conversation
+            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="Postgres",
+                    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()
+            
+            # Convert any datetime objects in existing metadata to ISO strings
+            serializable_metadata = {}
+            for key, value in conversation.metadata.items():
+                if isinstance(value, datetime):
+                    serializable_metadata[key] = value.isoformat()
+                else:
+                    serializable_metadata[key] = value
+            
+            updated_metadata = {
+                **serializable_metadata,
+                "updated_at": now.isoformat(),
+                "last_activity": now.isoformat(),
+                "total_messages": len(truncated_messages),
+                "regeneration_truncated": True,
+                "truncated_at": now.isoformat(),
+                "messages_removed": removed_count
+            }
+            
+            # Update in database
+            query = f"""
+            UPDATE {self.config.table_name}
+            SET messages = $1::jsonb, metadata = $2::jsonb
+            WHERE conversation_id = $3
+            """
+            
+            await self._db_execute(
+                query,
+                prepare_message_list_for_db(truncated_messages),
+                json.dumps(updated_metadata),
+                conversation_id
+            )
+            
+            print(f"[MEMORY:Postgres] Truncated conversation {conversation_id}: removed {removed_count} messages after message {message_id}")
+            return Success(removed_count)
+            
+        except Exception as e:
+            print(f"[MEMORY:Postgres] DEBUG: Exception in truncate_conversation_after: {e}")
+            import traceback
+            traceback.print_exc()
+            return Failure(MemoryStorageError(
+                message=f"Failed to truncate conversation: {e}",
+                provider="Postgres",
+                operation="truncate_conversation_after",
+                cause=e
+            ))
+
+    async def get_conversation_until_message(
+        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.
+        Useful for regeneration scenarios.
+        """
+        try:
+            # Get the conversation
+            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:Postgres] 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,
+                user_id=conversation.user_id,
+                messages=truncated_messages,
+                metadata={
+                    **conversation.metadata,
+                    "truncated_for_regeneration": True,
+                    "truncated_until_message": str(message_id),
+                    "original_message_count": len(messages),
+                    "truncated_message_count": len(truncated_messages)
+                }
+            )
+            
+            print(f"[MEMORY:Postgres] 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="Postgres",
+                operation="get_conversation_until_message",
+                cause=e
+            ))
+
+    async def mark_regeneration_point(
+        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.
+        """
+        try:
+            # Get the conversation
+            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="Postgres",
+                    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_points.append(regeneration_point)
+            
+            # Update conversation metadata
+            updated_metadata = {
+                **conversation.metadata,
+                "regeneration_points": regeneration_points,
+                "last_regeneration": regeneration_point,
+                "updated_at": datetime.now().isoformat(),
+                "regeneration_count": len(regeneration_points)
+            }
+            
+            # Update in database using JSONB merge
+            query = f"""
+            UPDATE {self.config.table_name}
+            SET metadata = metadata || $1::jsonb
+            WHERE conversation_id = $2
+            """
+            
+            await self._db_execute(
+                query,
+                json.dumps({
+                    "regeneration_points": regeneration_points,
+                    "last_regeneration": regeneration_point,
+                    "updated_at": updated_metadata["updated_at"],
+                    "regeneration_count": len(regeneration_points)
+                }),
+                conversation_id
+            )
+            
+            print(f"[MEMORY:Postgres] 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="Postgres",
+                operation="mark_regeneration_point",
+                cause=e
+            ))
+
     async def close(self) -> Result[None, MemoryConnectionError]:
         try:
             if hasattr(self.client, 'close'):

jaf/memory/providers/redis.py

@@ -8,7 +8,7 @@ Best for production environments with shared state and persistence across restar
 from datetime import datetime
 from typing import Any, Dict, List, Optional, Union
 
-from ...core.types import Message
+from ...core.types import Message, MessageId, find_message_index
 from ..types import (
     ConversationMemory,
     Failure,
@@ -205,6 +205,194 @@ class RedisProvider(MemoryProvider):
         except Exception as 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
+    ) -> Result[int, Union[MemoryNotFoundError, MemoryStorageError]]:
+        """
+        Truncate conversation after (and including) the specified message ID.
+        Returns the number of messages removed.
+        """
+        try:
+            # Get the conversation
+            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
+                ))
+            
+            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 = {
+                **conversation.metadata,
+                "updated_at": now,
+                "last_activity": now,
+                "total_messages": len(truncated_messages),
+                "regeneration_truncated": True,
+                "truncated_at": now.isoformat(),
+                "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
+            )
+            
+            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}")
+            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
+            ))
+
+    async def get_conversation_until_message(
+        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.
+        Useful for regeneration scenarios.
+        """
+        try:
+            # Get the conversation
+            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}")
+                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,
+                user_id=conversation.user_id,
+                messages=truncated_messages,
+                metadata={
+                    **conversation.metadata,
+                    "truncated_for_regeneration": True,
+                    "truncated_until_message": str(message_id),
+                    "original_message_count": len(messages),
+                    "truncated_message_count": len(truncated_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
+            ))
+
+    async def mark_regeneration_point(
+        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.
+        """
+        try:
+            # Get the conversation
+            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
+                ))
+            
+            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_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)
+            }
+            
+            # Store updated conversation
+            updated_conversation = ConversationMemory(
+                conversation_id=conversation.conversation_id,
+                user_id=conversation.user_id,
+                messages=conversation.messages,
+                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}")
+            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
+            ))
+
     async def close(self) -> Result[None, MemoryConnectionError]:
         try:
             await self.redis_client.aclose()

jaf/memory/types.py

@@ -11,7 +11,7 @@ from typing import Any, Dict, Generic, List, Optional, Protocol, TypeVar, Union
 
 from pydantic import BaseModel, Field
 
-from ..core.types import Message, TraceId
+from ..core.types import Message, TraceId, MessageId
 
 # Generic Result type for functional error handling
 T = TypeVar('T')
@@ -135,6 +135,40 @@ class MemoryProvider(Protocol):
         """Close/cleanup the provider."""
         ...
 
+    # Regeneration support methods
+    async def truncate_conversation_after(
+        self,
+        conversation_id: str,
+        message_id: MessageId
+    ) -> Result[int, Union['MemoryNotFoundError', 'MemoryStorageError']]:
+        """
+        Truncate conversation after (and including) the specified message ID.
+        Returns the number of messages removed.
+        """
+        ...
+
+    async def get_conversation_until_message(
+        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.
+        Useful for regeneration scenarios.
+        """
+        ...
+
+    async def mark_regeneration_point(
+        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.
+        """
+        ...
+
 # Configuration models using Pydantic for validation
 
 class InMemoryConfig(BaseModel):

jaf/memory/utils.py

@@ -22,6 +22,7 @@ def serialize_message(msg: Message) -> dict:
     return {
         "role": msg.role,
         "content": msg.content,
+        "message_id": msg.message_id,
         "tool_call_id": msg.tool_call_id,
         "tool_calls": [
             {
@@ -58,6 +59,7 @@ def deserialize_message(msg_data: dict) -> Message:
     return Message(
         role=msg_data["role"],
         content=msg_data["content"],
+        message_id=msg_data.get("message_id"),
         tool_call_id=msg_data.get("tool_call_id"),
         tool_calls=tool_calls
     )