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

jaf/__init__.py

@@ -191,7 +191,7 @@ def generate_run_id() -> RunId:
     """Generate a new run ID."""
     return create_run_id(str(uuid.uuid4()))
 
-__version__ = "2.5.3"
+__version__ = "2.5.4"
 __all__ = [
     # Core types and functions
     "TraceId", "RunId", "ValidationResult", "Message", "ModelConfig",

jaf/core/engine.py

@@ -293,15 +293,7 @@ async def _load_conversation_history(state: RunState[Ctx], config: RunConfig[Ctx
         # For HITL scenarios, append new messages to memory messages
         # This prevents duplication when resuming from interruptions
         if memory_messages:
-            combined_messages = memory_messages + [
-                msg for msg in state.messages 
-                if not any(
-                    mem_msg.role == msg.role and 
-                    mem_msg.content == msg.content and 
-                    getattr(mem_msg, 'tool_calls', None) == getattr(msg, 'tool_calls', None)
-                    for mem_msg in memory_messages
-                )
-            ]
+            combined_messages = memory_messages + list(state.messages)
         else:
             combined_messages = list(state.messages)
 

jaf/core/regeneration.py

@@ -0,0 +1,392 @@
+"""
+Regeneration functionality for the JAF framework.
+
+This module implements conversation regeneration where a specific message can be
+regenerated, removing all subsequent messages and creating a new conversation path.
+"""
+
+import time
+from dataclasses import replace
+from typing import Any, TypeVar, Optional
+
+from .types import (
+    RunState, RunConfig, RunResult, 
+    RegenerationRequest, RegenerationContext,
+    MessageId, Message, ErrorOutcome, ModelBehaviorError,
+    find_message_index, truncate_messages_after, get_message_by_id,
+    generate_run_id, generate_trace_id
+)
+from .engine import run as engine_run
+from ..memory.types import Success, Failure
+
+Ctx = TypeVar('Ctx')
+Out = TypeVar('Out')
+
+async def regenerate_conversation(
+    regeneration_request: RegenerationRequest,
+    config: RunConfig[Ctx],
+    context: Ctx,
+    agent_name: str
+) -> RunResult[Out]:
+    """
+    Regenerate a conversation from a specific message ID.
+    
+    This function:
+    1. Loads the full conversation from memory
+    2. Finds the message to regenerate from
+    3. Truncates the conversation at that point
+    4. Creates a new RunState with truncated conversation
+    5. Executes the regeneration through the normal engine flow
+    6. Updates memory with the new conversation path
+    
+    Args:
+        regeneration_request: The regeneration request containing conversation_id and message_id
+        config: The run configuration
+        context: The context for the regeneration
+        agent_name: The name of the agent to use for regeneration
+        
+    Returns:
+        RunResult with the regenerated conversation outcome
+    """
+    if not config.memory or not config.memory.provider or not config.conversation_id:
+        return RunResult(
+            final_state=RunState(
+                run_id=generate_run_id(),
+                trace_id=generate_trace_id(),
+                messages=[],
+                current_agent_name=agent_name,
+                context=context,
+                turn_count=0
+            ),
+            outcome=ErrorOutcome(error=ModelBehaviorError(
+                detail="Regeneration requires memory provider and conversation_id to be configured"
+            ))
+        )
+
+    # Load the conversation from memory
+    conversation_result = await config.memory.provider.get_conversation(regeneration_request.conversation_id)
+    if isinstance(conversation_result, Failure):
+        return RunResult(
+            final_state=RunState(
+                run_id=generate_run_id(),
+                trace_id=generate_trace_id(),
+                messages=[],
+                current_agent_name=agent_name,
+                context=context,
+                turn_count=0
+            ),
+            outcome=ErrorOutcome(error=ModelBehaviorError(
+                detail=f"Failed to load conversation: {conversation_result.error}"
+            ))
+        )
+
+    conversation_memory = conversation_result.data
+    if not conversation_memory:
+        return RunResult(
+            final_state=RunState(
+                run_id=generate_run_id(),
+                trace_id=generate_trace_id(),
+                messages=[],
+                current_agent_name=agent_name,
+                context=context,
+                turn_count=0
+            ),
+            outcome=ErrorOutcome(error=ModelBehaviorError(
+                detail=f"Conversation {regeneration_request.conversation_id} not found"
+            ))
+        )
+
+    # Convert tuple back to list for processing
+    original_messages = list(conversation_memory.messages)
+    
+    # Find the message to regenerate from
+    regenerate_message = get_message_by_id(original_messages, regeneration_request.message_id)
+    if not regenerate_message:
+        return RunResult(
+            final_state=RunState(
+                run_id=generate_run_id(),
+                trace_id=generate_trace_id(),
+                messages=original_messages,
+                current_agent_name=agent_name,
+                context=context,
+                turn_count=len([m for m in original_messages if (m.role.value if hasattr(m.role, 'value') else m.role) == 'assistant'])
+            ),
+            outcome=ErrorOutcome(error=ModelBehaviorError(
+                detail=f"Message {regeneration_request.message_id} not found in conversation"
+            ))
+        )
+
+    # Get the index of the message to regenerate
+    regenerate_index = find_message_index(original_messages, regeneration_request.message_id)
+    if regenerate_index is None:
+        return RunResult(
+            final_state=RunState(
+                run_id=generate_run_id(),
+                trace_id=generate_trace_id(),
+                messages=original_messages,
+                current_agent_name=agent_name,
+                context=context,
+                turn_count=len([m for m in original_messages if (m.role.value if hasattr(m.role, 'value') else m.role) == 'assistant'])
+            ),
+            outcome=ErrorOutcome(error=ModelBehaviorError(
+                detail=f"Failed to find index for message {regeneration_request.message_id}"
+            ))
+        )
+
+    def determine_regeneration_type(messages, regenerate_index, context):
+        """Determine if this is pure regeneration or edit scenario."""
+        if context and context.get("replace_user_message"):
+            return "edit"
+        
+        regenerate_message = messages[regenerate_index]
+        if regenerate_message.role in ['assistant', 'ASSISTANT']:
+            for i in range(regenerate_index - 1, -1, -1):
+                if messages[i].role in ['user', 'USER']:
+                    return "pure"
+        return "edit"
+
+    # Determine regeneration type
+    regen_type = determine_regeneration_type(original_messages, regenerate_index, regeneration_request.context or {})
+    print(f"[JAF:REGENERATION] Detected regeneration type: {regen_type}")
+
+    if regen_type == "pure":
+        # For pure regeneration, find the user message that started this conversation turn
+        user_message_index = None
+        for i in range(regenerate_index - 1, -1, -1):
+            if original_messages[i].role in ['user', 'USER']:
+                user_message_index = i
+                break
+        
+        if user_message_index is not None:
+            # Truncate AFTER the user message (keeps user message, removes tool calls/outputs)
+            truncated_messages = original_messages[:user_message_index + 1]
+            print(f"[JAF:REGENERATION] Pure regeneration: truncated to user message at index {user_message_index}")
+        else:
+            truncated_messages = original_messages[:regenerate_index]
+            print(f"[JAF:REGENERATION] Pure regeneration fallback: no user message found")
+    else:
+        # Edit regeneration: truncate at the specified point and add replacement query
+        truncated_messages = original_messages[:regenerate_index]
+        
+        if (regeneration_request.context and 
+            regeneration_request.context.get("replace_user_message")):
+            
+            from .types import ContentRole, Message
+            replacement_user_message = Message(
+                role=ContentRole.USER,
+                content=regeneration_request.context.get("replace_user_message")
+            )
+            truncated_messages.append(replacement_user_message)
+            print(f"[JAF:REGENERATION] Edit regeneration: replaced user query with: {regeneration_request.context.get('replace_user_message')}")
+    
+    print(f"[JAF:REGENERATION] Truncated conversation to {len(truncated_messages)} messages")
+    
+
+    print(f"[JAF:REGENERATION] About to store {len(truncated_messages)} truncated messages to memory")
+    
+    def serialize_metadata(metadata):
+        import json
+        import datetime
+        
+        def json_serializer(obj):
+            if isinstance(obj, datetime.datetime):
+                return obj.isoformat()
+            elif isinstance(obj, datetime.date):
+                return obj.isoformat()
+            elif hasattr(obj, '__dict__'):
+                return obj.__dict__
+            return str(obj)
+        
+        try:
+            json_str = json.dumps(metadata, default=json_serializer)
+            return json.loads(json_str)
+        except Exception as e:
+            print(f"[JAF:REGENERATION] Warning: Metadata serialization failed: {e}")
+            return {
+                "regeneration_truncated": True,
+                "regeneration_point": str(regeneration_request.message_id),
+                "original_message_count": len(original_messages),
+                "truncated_at_index": regenerate_index,
+                "turn_count": len([m for m in truncated_messages if (m.role.value if hasattr(m.role, 'value') else m.role) == 'assistant'])
+            }
+    
+    metadata = serialize_metadata({
+        **conversation_memory.metadata,
+        "regeneration_truncated": True,
+        "regeneration_point": str(regeneration_request.message_id),
+        "original_message_count": len(original_messages),
+        "truncated_at_index": regenerate_index,
+        "turn_count": len([m for m in truncated_messages if (m.role.value if hasattr(m.role, 'value') else m.role) == 'assistant'])
+    })
+    
+    store_result = await config.memory.provider.store_messages(
+        regeneration_request.conversation_id,
+        truncated_messages,
+        metadata
+    )
+    
+    print(f"[JAF:REGENERATION] Store result type: {type(store_result)}")
+    if isinstance(store_result, Failure):
+        print(f"[JAF:REGENERATION] Store failed with error: {store_result.error}")
+        return RunResult(
+            final_state=RunState(
+                run_id=generate_run_id(),
+                trace_id=generate_trace_id(),
+                messages=original_messages,
+                current_agent_name=agent_name,
+                context=context,
+                turn_count=len([m for m in original_messages if (m.role.value if hasattr(m.role, 'value') else m.role) == 'assistant'])
+            ),
+            outcome=ErrorOutcome(error=ModelBehaviorError(
+                detail=f"Failed to store truncated conversation: {store_result.error}"
+            ))
+        )
+    else:
+        print(f"[JAF:REGENERATION] Store successful, proceeding to engine execution")
+
+    # Create regeneration context for later use
+    regeneration_context = RegenerationContext(
+        original_message_count=len(original_messages),
+        truncated_at_index=regenerate_index,
+        regenerated_message_id=regeneration_request.message_id,
+        regeneration_id=f"regen_{int(time.time() * 1000)}_{regeneration_request.message_id}",
+        timestamp=int(time.time() * 1000)
+    )
+
+    # Calculate turn count from truncated messages
+    truncated_turn_count = len([m for m in truncated_messages if (m.role.value if hasattr(m.role, 'value') else m.role) == 'assistant'])
+    
+    final_context = context
+    print(f"[JAF:REGENERATION] Using provided context: {type(context).__name__}")
+
+    # Create initial state for regeneration with truncated conversation
+    initial_state = RunState(
+        run_id=generate_run_id(),
+        trace_id=generate_trace_id(),
+        messages=[],  
+        current_agent_name=agent_name,
+        context=final_context,
+        turn_count=truncated_turn_count,
+        approvals={}  # Reset approvals for regeneration
+    )
+
+    print(f"[JAF:REGENERATION] Starting regeneration from message {regeneration_request.message_id}")
+    print(f"[JAF:REGENERATION] Original messages: {len(original_messages)}, Truncated to: {len(truncated_messages)}")
+    print(f"[JAF:REGENERATION] Regeneration context: {regeneration_context}")
+
+    # Create a modified config for regeneration that ensures memory storage
+    regeneration_config = replace(
+        config,
+        conversation_id=regeneration_request.conversation_id,
+        memory=replace(config.memory, auto_store=True, store_on_completion=True) if config.memory else None
+    )
+
+    # Execute the regeneration through the normal engine flow
+    print(f"[JAF:REGENERATION] About to execute engine with {len(truncated_messages)} messages")
+    print(f"[JAF:REGENERATION] Final message: {truncated_messages[-1] if truncated_messages else 'None'}")
+    
+    result = await engine_run(initial_state, regeneration_config)
+
+    print(f"[JAF:REGENERATION] Regeneration completed with status: {result.outcome.status}")
+    if hasattr(result, 'final_state') and hasattr(result.final_state, 'messages'):
+        print(f"[JAF:REGENERATION] Final state has {len(result.final_state.messages)} messages")
+        assistant_msgs = [m for m in result.final_state.messages if m.role in ['assistant', 'ASSISTANT']]
+        print(f"[JAF:REGENERATION] Found {len(assistant_msgs)} assistant messages in result")
+    
+    # After successful regeneration, mark the regeneration point and preserve metadata
+    if result.outcome.status == 'completed' and config.memory and config.memory.provider:
+        try:
+            print(f"[JAF:REGENERATION] Marking regeneration point after successful regeneration")
+            
+            # Get the current conversation to preserve regeneration metadata
+            current_conv_result = await config.memory.provider.get_conversation(regeneration_request.conversation_id)
+            print(f"[JAF:REGENERATION] Retrieved conversation for preservation: {hasattr(current_conv_result, 'data') and current_conv_result.data is not None}")
+            
+            if hasattr(current_conv_result, 'data') and current_conv_result.data:
+                current_metadata = current_conv_result.data.metadata
+                regeneration_points = current_metadata.get('regeneration_points', [])
+                print(f"[JAF:REGENERATION] Found {len(regeneration_points)} regeneration points in metadata before marking")
+                
+                # Mark the regeneration point by calling the provider method directly
+                mark_result = await config.memory.provider.mark_regeneration_point(
+                    regeneration_request.conversation_id,
+                    regeneration_request.message_id,
+                    {
+                        "regeneration_id": regeneration_context.regeneration_id,
+                        "original_message_count": len(original_messages),
+                        "truncated_at_index": regenerate_index,
+                        "timestamp": regeneration_context.timestamp
+                    }
+                )
+                
+                if isinstance(mark_result, Failure):
+                    print(f"[JAF:REGENERATION] Warning: Failed to mark regeneration point: {mark_result.error}")
+                else:
+                    print(f"[JAF:REGENERATION] Successfully marked regeneration point")
+                
+                # Get the updated conversation with the new regeneration point
+                updated_conv_result = await config.memory.provider.get_conversation(regeneration_request.conversation_id)
+                if hasattr(updated_conv_result, 'data') and updated_conv_result.data:
+                    updated_metadata = updated_conv_result.data.metadata
+                    updated_regeneration_points = updated_metadata.get('regeneration_points', [])
+                    print(f"[JAF:REGENERATION] Found {len(updated_regeneration_points)} regeneration points after marking")
+                    
+                    # Ensure final metadata includes the regeneration points
+                    final_metadata = {
+                        **updated_metadata,
+                        'regeneration_points': updated_regeneration_points,
+                        'regeneration_count': len(updated_regeneration_points),
+                        'last_regeneration': updated_regeneration_points[-1] if updated_regeneration_points else None,
+                        'regeneration_preserved': True,
+                        'final_preservation_timestamp': int(time.time() * 1000)
+                    }
+                    
+                    # Store the final conversation with preserved regeneration metadata
+                    await config.memory.provider.store_messages(
+                        regeneration_request.conversation_id,
+                        result.final_state.messages,
+                        final_metadata
+                    )
+                    print(f"[JAF:REGENERATION] Final preservation completed with {len(updated_regeneration_points)} regeneration points")
+            else:
+                print(f"[JAF:REGENERATION] No conversation data found for preservation")
+                    
+        except Exception as e:
+            print(f"[JAF:REGENERATION] Warning: Failed to preserve regeneration points: {e}")
+            import traceback
+            traceback.print_exc()
+    
+    return result
+
+
+async def get_regeneration_points(
+    conversation_id: str,
+    config: RunConfig[Ctx]
+) -> Optional[list]:
+    """
+    Get all regeneration points for a conversation.
+    
+    Args:
+        conversation_id: The conversation ID
+        config: The run configuration
+        
+    Returns:
+        List of regeneration points or None if not available
+    """
+    if not config.memory or not config.memory.provider:
+        return None
+        
+    try:
+        conversation_result = await config.memory.provider.get_conversation(conversation_id)
+        if hasattr(conversation_result, 'data') and conversation_result.data:
+            metadata = conversation_result.data.metadata
+            regeneration_points = metadata.get('regeneration_points', [])
+            print(f"[JAF:REGENERATION] Retrieved {len(regeneration_points)} regeneration points for {conversation_id}")
+            return regeneration_points
+        else:
+            print(f"[JAF:REGENERATION] No conversation data found for {conversation_id}")
+            return []
+    except Exception as e:
+        print(f"[JAF:REGENERATION] Failed to get regeneration points: {e}")
+    
+    return []

jaf/core/tracing.py

@@ -443,7 +443,7 @@ class LangfuseTraceCollector:
             public_key=public_key,
             secret_key=secret_key,
             host=host,
-            release="jaf-py-v2.5.3",
+            release="jaf-py-v2.5.4",
             httpx_client=client
         )
         self._httpx_client = client

jaf/core/types.py

@@ -94,6 +94,11 @@ class RunId(str):
     def __new__(cls, value: str) -> 'RunId':
         return str.__new__(cls, value)
 
+class MessageId(str):
+    """Branded string type for message IDs."""
+    def __new__(cls, value: str) -> 'MessageId':
+        return str.__new__(cls, value)
+
 def create_trace_id(id_str: str) -> TraceId:
     """Create a TraceId from a string."""
     return TraceId(id_str)
@@ -102,6 +107,36 @@ def create_run_id(id_str: str) -> RunId:
     """Create a RunId from a string."""
     return RunId(id_str)
 
+def create_message_id(id_str: Union[str, MessageId]) -> MessageId:
+    """
+    Create a MessageId from a string or return existing MessageId.
+    
+    Args:
+        id_str: Either a string to convert to MessageId or an existing MessageId
+        
+    Returns:
+        MessageId: A validated MessageId instance
+        
+    Raises:
+        ValueError: If the input is invalid or empty
+    """
+    # Handle None input
+    if id_str is None:
+        raise ValueError("Message ID cannot be None")
+    
+    # If already a MessageId, return as-is
+    if isinstance(id_str, MessageId):
+        return id_str
+    
+    # Convert string to MessageId with validation
+    if isinstance(id_str, str):
+        if not id_str.strip():
+            raise ValueError("Message ID cannot be empty or whitespace")
+        return MessageId(id_str.strip())
+    
+    # Handle any other type
+    raise ValueError(f"Message ID must be a string or MessageId, got {type(id_str)}")
+
 def generate_run_id() -> RunId:
     """Generate a new unique run ID."""
     import time
@@ -114,6 +149,12 @@ def generate_trace_id() -> TraceId:
     import uuid
     return TraceId(f"trace_{int(time.time() * 1000)}_{uuid.uuid4().hex[:8]}")
 
+def generate_message_id() -> MessageId:
+    """Generate a new unique message ID."""
+    import time
+    import uuid
+    return MessageId(f"msg_{int(time.time() * 1000)}_{uuid.uuid4().hex[:8]}")
+
 # Type variables for generic contexts and outputs
 Ctx = TypeVar('Ctx')
 Out = TypeVar('Out')
@@ -180,12 +221,16 @@ class Message:
     - Direct access to .content returns the original string when created with string
     - Use .text_content property for guaranteed string access in all cases
     - Use get_text_content() function to extract text from any content type
+    - message_id is optional for backward compatibility
     
     Examples:
         # Original usage - still works exactly the same
         msg = Message(role='user', content='Hello')
         text = msg.content  # Returns 'Hello' as string
         
+        # New usage with message ID
+        msg = Message(role='user', content='Hello', message_id='msg_123')
+        
         # Guaranteed string access (recommended for new code)
         text = msg.text_content  # Always returns string
         
@@ -197,6 +242,27 @@ class Message:
     attachments: Optional[List[Attachment]] = None
     tool_call_id: Optional[str] = None
     tool_calls: Optional[List[ToolCall]] = None
+    message_id: Optional[MessageId] = None  # Optional for backward compatibility
+    
+    def __post_init__(self):
+        """
+        Auto-generate message ID if not provided.
+        
+        This implementation uses object.__setattr__ to bypass frozen dataclass restrictions,
+        which is a recommended pattern for one-time initialization of computed fields in
+        frozen dataclasses. This ensures:
+        
+        1. Backward compatibility - existing code with message_id=None continues to work
+        2. Immutability - the dataclass remains frozen after initialization
+        3. Guaranteed unique IDs - every message gets a unique identifier
+        4. Clean API - users don't need to manually generate IDs in most cases
+        
+        This pattern is preferred over using field(default_factory=...) because it
+        maintains the Optional[MessageId] type hint for backward compatibility while
+        ensuring the field is never actually None after object creation.
+        """
+        if self.message_id is None:
+            object.__setattr__(self, 'message_id', generate_message_id())
     
     @property
     def text_content(self) -> str:
@@ -210,7 +276,8 @@ class Message:
         content: str,
         attachments: Optional[List[Attachment]] = None,
         tool_call_id: Optional[str] = None,
-        tool_calls: Optional[List[ToolCall]] = None
+        tool_calls: Optional[List[ToolCall]] = None,
+        message_id: Optional[MessageId] = None
     ) -> 'Message':
         """Create a message with string content and optional attachments."""
         return cls(
@@ -218,7 +285,8 @@ class Message:
             content=content,
             attachments=attachments,
             tool_call_id=tool_call_id,
-            tool_calls=tool_calls
+            tool_calls=tool_calls,
+            message_id=message_id
         )
 
 def get_text_content(content: Union[str, List[MessageContentPart]]) -> str:
@@ -824,3 +892,42 @@ class RunConfig(Generic[Ctx]):
     default_fast_model: Optional[str] = None  # Default model for fast operations like guardrails
     default_tool_timeout: Optional[float] = 300.0  # Default timeout for tool execution in seconds
     approval_storage: Optional['ApprovalStorage'] = None  # Storage for approval decisions
+
+# Regeneration types for conversation management
+@dataclass(frozen=True)
+class RegenerationRequest:
+    """Request to regenerate a conversation from a specific message."""
+    conversation_id: str
+    message_id: MessageId  # ID of the message to regenerate from
+    context: Optional[Dict[str, Any]] = None  # Optional context override
+
+@dataclass(frozen=True)
+class RegenerationContext:
+    """Context information for a regeneration operation."""
+    original_message_count: int
+    truncated_at_index: int
+    regenerated_message_id: MessageId
+    regeneration_id: str  # Unique ID for this regeneration operation
+    timestamp: int  # Unix timestamp in milliseconds
+
+# Message utility functions
+def find_message_index(messages: List[Message], message_id: MessageId) -> Optional[int]:
+    """Find the index of a message by its ID."""
+    for i, msg in enumerate(messages):
+        if msg.message_id == message_id:
+            return i
+    return None
+
+def truncate_messages_after(messages: List[Message], message_id: MessageId) -> List[Message]:
+    """Truncate messages after (and including) the specified message ID."""
+    index = find_message_index(messages, message_id)
+    if index is None:
+        return messages  # Message not found, return unchanged
+    return messages[:index]
+
+def get_message_by_id(messages: List[Message], message_id: MessageId) -> Optional[Message]:
+    """Get a message by its ID."""
+    for msg in messages:
+        if msg.message_id == message_id:
+            return msg
+    return None