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

jaf/core/tracing.py

@@ -443,10 +443,18 @@ class LangfuseTraceCollector:
             public_key=public_key,
             secret_key=secret_key,
             host=host,
-            release="jaf-py-v2.5.2",
+            release="jaf-py-v2.5.4",
             httpx_client=client
         )
         self._httpx_client = client
+
+        # Detect Langfuse version (v2 has trace() method, v3 does not)
+        self._is_langfuse_v3 = not hasattr(self.langfuse, 'trace')
+        if self._is_langfuse_v3:
+            print("[LANGFUSE] Detected Langfuse v3.x - using OpenTelemetry-based API")
+        else:
+            print("[LANGFUSE] Detected Langfuse v2.x - using legacy API")
+
         self.active_spans: Dict[str, Any] = {}
         self.trace_spans: Dict[TraceId, Any] = {}
         # Track tool calls and results for each trace
@@ -465,6 +473,113 @@ class LangfuseTraceCollector:
             except Exception as e:
                 print(f"[LANGFUSE] Warning: Failed to close httpx client: {e}")
 
+    def _get_event_data(self, event: TraceEvent, key: str, default: Any = None) -> Any:
+        """Extract data from event, handling both dict and dataclass."""
+        if not hasattr(event, 'data'):
+            return default
+
+        # Handle dict
+        if isinstance(event.data, dict):
+            return event.data.get(key, default)
+
+        # Handle dataclass/object with attributes
+        return getattr(event.data, key, default)
+
+    def _create_trace(self, trace_id: TraceId, **kwargs) -> Any:
+        """Create a trace using the appropriate API for the Langfuse version."""
+        if self._is_langfuse_v3:
+            # Langfuse v3: Use start_span() to create a root span (creates trace implicitly)
+            # Extract parameters for v3 API
+            name = kwargs.get('name', 'trace')
+            input_data = kwargs.get('input')
+            metadata = kwargs.get('metadata', {})
+            user_id = kwargs.get('user_id')
+            session_id = kwargs.get('session_id')
+            tags = kwargs.get('tags', [])
+
+            # Add user_id, session_id, and tags to metadata for v3
+            if user_id:
+                metadata['user_id'] = user_id
+            if session_id:
+                metadata['session_id'] = session_id
+            if tags:
+                metadata['tags'] = tags
+
+            # Create root span
+            trace = self.langfuse.start_span(
+                name=name,
+                input=input_data,
+                metadata=metadata
+            )
+
+            # Update trace properties using update_trace()
+            update_params = {}
+            if user_id:
+                update_params['user_id'] = user_id
+            if session_id:
+                update_params['session_id'] = session_id
+            if tags:
+                update_params['tags'] = tags
+
+            if update_params:
+                trace.update_trace(**update_params)
+
+            return trace
+        else:
+            # Langfuse v2: Use trace() method
+            return self.langfuse.trace(**kwargs)
+
+    def _create_generation(self, parent_span: Any, **kwargs) -> Any:
+        """Create a generation using the appropriate API for the Langfuse version."""
+        if self._is_langfuse_v3:
+            # Langfuse v3: Use start_generation() method
+            return parent_span.start_generation(**kwargs)
+        else:
+            # Langfuse v2: Use generation() method
+            return parent_span.generation(**kwargs)
+
+    def _create_span(self, parent_span: Any, **kwargs) -> Any:
+        """Create a span using the appropriate API for the Langfuse version."""
+        if self._is_langfuse_v3:
+            # Langfuse v3: Use start_span() method
+            return parent_span.start_span(**kwargs)
+        else:
+            # Langfuse v2: Use span() method
+            return parent_span.span(**kwargs)
+
+    def _create_event(self, parent_span: Any, **kwargs) -> Any:
+        """Create an event using the appropriate API for the Langfuse version."""
+        if self._is_langfuse_v3:
+            # Langfuse v3: Use create_event() method
+            return parent_span.create_event(**kwargs)
+        else:
+            # Langfuse v2: Use event() method
+            return parent_span.event(**kwargs)
+
+    def _end_span(self, span: Any, **kwargs) -> None:
+        """End a span/generation using the appropriate API for the Langfuse version."""
+        if self._is_langfuse_v3:
+            # Langfuse v3: Call update() first with output/metadata, then end()
+            update_params = {}
+            end_params = {}
+
+            # Separate parameters for update() vs end()
+            for key, value in kwargs.items():
+                if key in ['output', 'metadata', 'model', 'usage']:
+                    update_params[key] = value
+                elif key == 'end_time':
+                    end_params[key] = value
+
+            # Update first if there are parameters
+            if update_params:
+                span.update(**update_params)
+
+            # Then end
+            span.end(**end_params)
+        else:
+            # Langfuse v2: Call end() directly with all parameters
+            span.end(**kwargs)
+
     def collect(self, event: TraceEvent) -> None:
         """Collect a trace event and send it to Langfuse."""
         try:
@@ -489,15 +604,15 @@ class LangfuseTraceCollector:
                 conversation_history = []
                 
                 # Debug: Print the event data structure to understand what we're working with
-                if event.data.get("context"):
-                    context = event.data["context"]
+                if self._get_event_data(event, "context"):
+                    context = self._get_event_data(event, "context")
                     print(f"[LANGFUSE DEBUG] Context type: {type(context)}")
                     print(f"[LANGFUSE DEBUG] Context attributes: {dir(context) if hasattr(context, '__dict__') else 'Not an object'}")
                     if hasattr(context, '__dict__'):
                         print(f"[LANGFUSE DEBUG] Context dict: {context.__dict__}")
                 
                 # Try to extract from context first
-                context = event.data.get("context")
+                context = self._get_event_data(event, "context")
                 if context:
                     # Try direct attribute access
                     if hasattr(context, 'query'):
@@ -527,7 +642,7 @@ class LangfuseTraceCollector:
                             print(f"[LANGFUSE DEBUG] Extracted user_id from attr: {user_id}")
                 
                 # Extract conversation history and current user query from messages
-                messages = event.data.get("messages", [])
+                messages = self._get_event_data(event, "messages", [])
                 if messages:
                     print(f"[LANGFUSE DEBUG] Processing {len(messages)} messages")
                     
@@ -619,20 +734,22 @@ class LangfuseTraceCollector:
                 trace_input = {
                     "user_query": user_query,
                     "run_id": str(trace_id),
-                    "agent_name": event.data.get("agent_name", "analytics_agent_jaf"),
+                    "agent_name": self._get_event_data(event, "agent_name", "analytics_agent_jaf"),
                     "session_info": {
-                        "session_id": event.data.get("session_id"),
-                        "user_id": user_id or event.data.get("user_id")
+                        "session_id": self._get_event_data(event, "session_id"),
+                        "user_id": user_id or self._get_event_data(event, "user_id")
                     }
                 }
 
                 # Extract agent_name for tagging
-                agent_name = event.data.get("agent_name") or "analytics_agent_jaf"
+                agent_name = self._get_event_data(event, "agent_name") or "analytics_agent_jaf"
 
-                trace = self.langfuse.trace(
+                # Use compatibility layer to create trace (works with both v2 and v3)
+                trace = self._create_trace(
+                    trace_id=trace_id,
                     name=agent_name,
-                    user_id=user_id or event.data.get("user_id"),
-                    session_id=event.data.get("session_id"),
+                    user_id=user_id or self._get_event_data(event, "user_id"),
+                    session_id=self._get_event_data(event, "session_id"),
                     input=trace_input,
                     tags=[agent_name],  # Add agent_name as a tag for dashboard filtering
                     metadata={
@@ -640,17 +757,17 @@ class LangfuseTraceCollector:
                         "event_type": "run_start",
                         "trace_id": str(trace_id),
                         "user_query": user_query,
-                        "user_id": user_id or event.data.get("user_id"),
+                        "user_id": user_id or self._get_event_data(event, "user_id"),
                         "agent_name": agent_name,
                         "conversation_history": conversation_history,
                         "tool_calls": [],
                         "tool_results": [],
-                        "user_info": event.data.get("context").user_info if event.data.get("context") and hasattr(event.data.get("context"), 'user_info') else None
+                        "user_info": self._get_event_data(event, "context").user_info if self._get_event_data(event, "context") and hasattr(self._get_event_data(event, "context"), 'user_info') else None
                     }
                 )
                 self.trace_spans[trace_id] = trace
                 # Store user_id, user_query, and conversation_history for later use
-                trace._user_id = user_id or event.data.get("user_id")
+                trace._user_id = user_id or self._get_event_data(event, "user_id")
                 trace._user_query = user_query
                 trace._conversation_history = conversation_history
                 print(f"[LANGFUSE] Created trace with user query: {user_query[:100] if user_query else 'None'}...")
@@ -667,7 +784,7 @@ class LangfuseTraceCollector:
                         "trace_id": str(trace_id),
                         "user_query": getattr(self.trace_spans[trace_id], '_user_query', None),
                         "user_id": getattr(self.trace_spans[trace_id], '_user_id', None),
-                        "agent_name": event.data.get("agent_name", "analytics_agent_jaf"),
+                        "agent_name": self._get_event_data(event, "agent_name", "analytics_agent_jaf"),
                         "conversation_history": conversation_history,
                         "tool_calls": self.trace_tool_calls.get(trace_id, []),
                         "tool_results": self.trace_tool_results.get(trace_id, [])
@@ -695,7 +812,7 @@ class LangfuseTraceCollector:
                     
             elif event.type == "llm_call_start":
                 # Start a generation for LLM calls
-                model = event.data.get("model", "unknown")
+                model = self._get_event_data(event, "model", "unknown")
                 print(f"[LANGFUSE] Starting generation for LLM call with model: {model}")
                 
                 # Get stored user information from the trace
@@ -703,11 +820,13 @@ class LangfuseTraceCollector:
                 user_id = getattr(trace, '_user_id', None)
                 user_query = getattr(trace, '_user_query', None)
                 
-                generation = trace.generation(
+                # Use compatibility layer to create generation (works with both v2 and v3)
+                generation = self._create_generation(
+                    parent_span=trace,
                     name=f"llm-call-{model}",
-                    input=event.data.get("messages"),
+                    input=self._get_event_data(event, "messages"),
                     metadata={
-                        "agent_name": event.data.get("agent_name"), 
+                        "agent_name": self._get_event_data(event, "agent_name"),
                         "model": model,
                         "user_id": user_id,
                         "user_query": user_query
@@ -723,10 +842,10 @@ class LangfuseTraceCollector:
                     print(f"[LANGFUSE] Ending generation for LLM call")
                     # End the generation
                     generation = self.active_spans[span_id]
-                    choice = event.data.get("choice", {})
-                    
+                    choice = self._get_event_data(event, "choice", {})
+
                     # Extract usage from the event data
-                    usage = event.data.get("usage", {})
+                    usage = self._get_event_data(event, "usage", {})
                     
                     # Extract model information from choice data or event data
                     model = choice.get("model", "unknown")
@@ -752,8 +871,10 @@ class LangfuseTraceCollector:
                         print(f"[LANGFUSE] Usage data for automatic cost calculation: {langfuse_usage}")
 
                     # Include model information in the generation end - Langfuse will calculate costs automatically
-                    generation.end(
-                        output=choice, 
+                    # Use compatibility wrapper for ending spans/generations
+                    self._end_span(
+                        span=generation,
+                        output=choice,
                         usage=langfuse_usage,
                         model=model,  # Pass model directly for automatic cost calculation
                         metadata={
@@ -772,9 +893,9 @@ class LangfuseTraceCollector:
                     
             elif event.type == "tool_call_start":
                 # Start a span for tool calls with detailed input information
-                tool_name = event.data.get('tool_name', 'unknown')
-                tool_args = event.data.get("args", {})
-                call_id = event.data.get("call_id")
+                tool_name = self._get_event_data(event, 'tool_name', 'unknown')
+                tool_args = self._get_event_data(event, "args", {})
+                call_id = self._get_event_data(event, "call_id")
                 if not call_id:
                     call_id = f"{tool_name}-{uuid.uuid4().hex[:8]}"
                     try:
@@ -807,7 +928,9 @@ class LangfuseTraceCollector:
                     "timestamp": datetime.now().isoformat()
                 }
                 
-                span = self.trace_spans[trace_id].span(
+                # Use compatibility layer to create span (works with both v2 and v3)
+                span = self._create_span(
+                    parent_span=self.trace_spans[trace_id],
                     name=f"tool-{tool_name}",
                     input=tool_input,
                     metadata={
@@ -824,9 +947,9 @@ class LangfuseTraceCollector:
             elif event.type == "tool_call_end":
                 span_id = self._get_span_id(event)
                 if span_id in self.active_spans:
-                    tool_name = event.data.get('tool_name', 'unknown')
-                    tool_result = event.data.get("result")
-                    call_id = event.data.get("call_id")
+                    tool_name = self._get_event_data(event, 'tool_name', 'unknown')
+                    tool_result = self._get_event_data(event, "result")
+                    call_id = self._get_event_data(event, "call_id")
                     
                     print(f"[LANGFUSE] Ending span for tool call: {tool_name} ({call_id})")
                     
@@ -836,9 +959,9 @@ class LangfuseTraceCollector:
                         "result": tool_result,
                         "call_id": call_id,
                         "timestamp": datetime.now().isoformat(),
-                        "execution_status": event.data.get("execution_status", "completed"),
-                        "status": event.data.get("execution_status", "completed"),  # DEPRECATED: backward compatibility
-                        "tool_result": event.data.get("tool_result")
+                        "execution_status": self._get_event_data(event, "execution_status", "completed"),
+                        "status": self._get_event_data(event, "execution_status", "completed"),  # DEPRECATED: backward compatibility
+                        "tool_result": self._get_event_data(event, "tool_result")
                     }
                     
                     if trace_id not in self.trace_tool_results:
@@ -852,13 +975,15 @@ class LangfuseTraceCollector:
                         "result": tool_result,
                         "call_id": call_id,
                         "timestamp": datetime.now().isoformat(),
-                        "execution_status": event.data.get("execution_status", "completed"),
-                        "status": event.data.get("execution_status", "completed")  # DEPRECATED: backward compatibility
+                        "execution_status": self._get_event_data(event, "execution_status", "completed"),
+                        "status": self._get_event_data(event, "execution_status", "completed")  # DEPRECATED: backward compatibility
                     }
                     
                     # End the span with detailed output
+                    # Use compatibility wrapper for ending spans/generations
                     span = self.active_spans[span_id]
-                    span.end(
+                    self._end_span(
+                        span=span,
                         output=tool_output,
                         metadata={
                             "tool_name": tool_name,
@@ -878,17 +1003,21 @@ class LangfuseTraceCollector:
             elif event.type == "handoff":
                 # Create an event for handoffs
                 print(f"[LANGFUSE] Creating event for handoff")
-                self.trace_spans[trace_id].event(
+                # Use compatibility layer to create event (works with both v2 and v3)
+                self._create_event(
+                    parent_span=self.trace_spans[trace_id],
                     name="agent-handoff",
-                    input={"from": event.data.get("from"), "to": event.data.get("to")},
+                    input={"from": self._get_event_data(event, "from"), "to": self._get_event_data(event, "to")},
                     metadata=event.data
                 )
                 print(f"[LANGFUSE] Handoff event created")
-                    
+
             else:
                 # Create a generic event for other event types
                 print(f"[LANGFUSE] Creating generic event for: {event.type}")
-                self.trace_spans[trace_id].event(
+                # Use compatibility layer to create event (works with both v2 and v3)
+                self._create_event(
+                    parent_span=self.trace_spans[trace_id],
                     name=event.type,
                     input=event.data,
                     metadata={"framework": "jaf", "event_type": event.type}
@@ -902,20 +1031,28 @@ class LangfuseTraceCollector:
             traceback.print_exc()
 
     def _get_trace_id(self, event: TraceEvent) -> Optional[TraceId]:
-        """Extract trace ID from event data."""
-        if hasattr(event, 'data') and isinstance(event.data, dict):
-            # Try snake_case first (Python convention)
-            if 'trace_id' in event.data:
-                return event.data['trace_id']
-            elif 'run_id' in event.data:
-                return TraceId(event.data['run_id'])
-            # Fallback to camelCase (for compatibility)
-            elif 'traceId' in event.data:
-                return event.data['traceId']
-            elif 'runId' in event.data:
-                return TraceId(event.data['runId'])
-        
-        # Debug: print what's actually in the event data
+        """Extract trace ID from event data, handling both dict and dataclass."""
+        if not hasattr(event, 'data'):
+            return None
+
+        # Try snake_case first (Python convention)
+        trace_id = self._get_event_data(event, 'trace_id')
+        if trace_id:
+            return trace_id
+
+        run_id = self._get_event_data(event, 'run_id')
+        if run_id:
+            return TraceId(run_id)
+
+        # Fallback to camelCase (for compatibility)
+        trace_id = self._get_event_data(event, 'traceId')
+        if trace_id:
+            return trace_id
+
+        run_id = self._get_event_data(event, 'runId')
+        if run_id:
+            return TraceId(run_id)
+
         return None
 
     def _get_span_id(self, event: TraceEvent) -> str:
@@ -924,15 +1061,15 @@ class LangfuseTraceCollector:
         
         # Use consistent identifiers that don't depend on timestamp
         if event.type.startswith('tool_call'):
-            call_id = event.data.get('call_id') or event.data.get('tool_call_id')
+            call_id = self._get_event_data(event, 'call_id') or self._get_event_data(event, 'tool_call_id')
             if call_id:
                 return f"tool-{trace_id}-{call_id}"
-            tool_name = event.data.get('tool_name') or event.data.get('toolName', 'unknown')
+            tool_name = self._get_event_data(event, 'tool_name') or self._get_event_data(event, 'toolName', 'unknown')
             return f"tool-{tool_name}-{trace_id}"
         elif event.type.startswith('llm_call'):
             # For LLM calls, use a simpler consistent ID that matches between start and end
             # Get run_id for more consistent matching
-            run_id = event.data.get('run_id') or event.data.get('runId', trace_id)
+            run_id = self._get_event_data(event, 'run_id') or self._get_event_data(event, 'runId', trace_id)
             return f"llm-{run_id}"
         else:
             return f"{event.type}-{trace_id}"

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