jaf-py 2.4.1__py3-none-any.whl → 2.4.2__py3-none-any.whl

jaf/core/state.py

@@ -0,0 +1,156 @@
+"""
+State management functions for approval handling in HITL scenarios.
+
+This module provides functions to manage approval state transitions
+and integrate with approval storage systems.
+"""
+
+from typing import Dict, Any, Optional
+from dataclasses import replace
+
+from .types import RunState, RunConfig, Interruption, ApprovalValue
+
+
+async def approve(
+    state: RunState[Any],
+    interruption: Interruption,
+    additional_context: Optional[Dict[str, Any]] = None,
+    config: Optional[RunConfig[Any]] = None
+) -> RunState[Any]:
+    """
+    Approve a tool call interruption and update the run state.
+    
+    Args:
+        state: Current run state
+        interruption: The interruption to approve
+        additional_context: Optional additional context for the approval
+        config: Optional run configuration for approval storage
+        
+    Returns:
+        Updated run state with approval recorded
+    """
+    if interruption.type == 'tool_approval':
+        approval_value = ApprovalValue(
+            status='approved',
+            approved=True,
+            additional_context={
+                **(additional_context or {}), 
+                'status': 'approved'
+            }
+        )
+        
+        # Store in approval storage if available
+        if config and config.approval_storage:
+            try:
+                print(f"[JAF:APPROVAL] Storing approval for tool_call_id {interruption.tool_call.id}: {approval_value}")
+                result = await config.approval_storage.store_approval(
+                    state.run_id,
+                    interruption.tool_call.id,
+                    approval_value
+                )
+                if not result.success:
+                    print(f"[JAF:APPROVAL] Failed to store approval: {result.error}")
+                    # Continue with in-memory fallback
+                else:
+                    print(f"[JAF:APPROVAL] Successfully stored approval in storage")
+            except Exception as e:
+                print(f"[JAF:APPROVAL] Approval storage error: {e}")
+                # Continue with in-memory fallback
+        
+        # Update in-memory state
+        new_approvals = {**state.approvals}
+        new_approvals[interruption.tool_call.id] = approval_value
+        
+        return replace(state, approvals=new_approvals)
+    
+    return state
+
+
+async def reject(
+    state: RunState[Any],
+    interruption: Interruption,
+    additional_context: Optional[Dict[str, Any]] = None,
+    config: Optional[RunConfig[Any]] = None
+) -> RunState[Any]:
+    """
+    Reject a tool call interruption and update the run state.
+    
+    Args:
+        state: Current run state
+        interruption: The interruption to reject
+        additional_context: Optional additional context for the rejection
+        config: Optional run configuration for approval storage
+        
+    Returns:
+        Updated run state with rejection recorded
+    """
+    if interruption.type == 'tool_approval':
+        approval_value = ApprovalValue(
+            status='rejected',
+            approved=False,
+            additional_context={
+                **(additional_context or {}), 
+                'status': 'rejected'
+            }
+        )
+        
+        # Store in approval storage if available
+        if config and config.approval_storage:
+            try:
+                print(f"[JAF:APPROVAL] Storing approval for tool_call_id {interruption.tool_call.id}: {approval_value}")
+                result = await config.approval_storage.store_approval(
+                    state.run_id,
+                    interruption.tool_call.id,
+                    approval_value
+                )
+                if not result.success:
+                    print(f"[JAF:APPROVAL] Failed to store approval: {result.error}")
+                    # Continue with in-memory fallback
+                else:
+                    print(f"[JAF:APPROVAL] Successfully stored approval in storage")
+            except Exception as e:
+                print(f"[JAF:APPROVAL] Approval storage error: {e}")
+                # Continue with in-memory fallback
+        
+        # Update in-memory state
+        new_approvals = {**state.approvals}
+        new_approvals[interruption.tool_call.id] = approval_value
+        
+        return replace(state, approvals=new_approvals)
+    
+    return state
+
+
+async def load_approvals_into_state(
+    state: RunState[Any],
+    config: Optional[RunConfig[Any]] = None
+) -> RunState[Any]:
+    """
+    Load approvals from storage into the run state.
+    
+    Args:
+        state: Current run state
+        config: Optional run configuration with approval storage
+        
+    Returns:
+        Updated run state with loaded approvals
+    """
+    if not config or not config.approval_storage:
+        print(f"[JAF:APPROVAL] No approval storage configured, using existing approvals: {state.approvals}")
+        return state
+    
+    try:
+        print(f"[JAF:APPROVAL] Loading approvals from storage for run_id: {state.run_id}")
+        result = await config.approval_storage.get_run_approvals(state.run_id)
+        if result.success and result.data:
+            print(f"[JAF:APPROVAL] Loaded {len(result.data)} approvals from storage: {result.data}")
+            return replace(state, approvals=result.data)
+        else:
+            if not result.success:
+                print(f"[JAF:APPROVAL] Failed to load approvals: {result.error}")
+            else:
+                print(f"[JAF:APPROVAL] No approvals found in storage for run_id: {state.run_id}")
+            return state
+    except Exception as e:
+        print(f"[JAF:APPROVAL] Approval loading error: {e}")
+        return state

jaf/core/tracing.py

@@ -344,6 +344,9 @@ class LangfuseTraceCollector:
         )
         self.active_spans: Dict[str, Any] = {}
         self.trace_spans: Dict[TraceId, Any] = {}
+        # Track tool calls and results for each trace
+        self.trace_tool_calls: Dict[TraceId, List[Dict[str, Any]]] = {}
+        self.trace_tool_results: Dict[TraceId, List[Dict[str, Any]]] = {}
 
     def collect(self, event: TraceEvent) -> None:
         """Collect a trace event and send it to Langfuse."""
@@ -359,9 +362,14 @@ class LangfuseTraceCollector:
                 # Start a new trace for the entire run
                 print(f"[LANGFUSE] Starting trace for run: {trace_id}")
                 
+                # Initialize tracking for this trace
+                self.trace_tool_calls[trace_id] = []
+                self.trace_tool_results[trace_id] = []
+                
                 # Extract user query from the run_start data
                 user_query = None
                 user_id = None
+                conversation_history = []
                 
                 # Debug: Print the event data structure to understand what we're working with
                 if event.data.get("context"):
@@ -401,24 +409,51 @@ class LangfuseTraceCollector:
                             user_id = token_response.email
                             print(f"[LANGFUSE DEBUG] Extracted user_id from attr: {user_id}")
                 
-                # Fallback: try to extract from messages if context didn't work
-                if not user_query and event.data.get("messages"):
-                    print(f"[LANGFUSE DEBUG] Trying fallback from messages")
-                    messages = event.data["messages"]
-                    print(f"[LANGFUSE DEBUG] Found {len(messages)} messages")
-                    # Find the last user message which should be the current query
-                    for i, msg in enumerate(reversed(messages)):
-                        print(f"[LANGFUSE DEBUG] Message {i}: {msg}")
-                        if isinstance(msg, dict) and msg.get("role") == "user":
-                            user_query = msg.get("content", "")
-                            print(f"[LANGFUSE DEBUG] Found user_query from messages: {user_query}")
-                            break
-                        elif hasattr(msg, 'role') and msg.role == 'user':
-                            user_query = msg.content
-                            print(f"[LANGFUSE DEBUG] Found user_query from message attr: {user_query}")
-                            break
+                # Extract conversation history and current user query from messages
+                messages = event.data.get("messages", [])
+                if messages:
+                    print(f"[LANGFUSE DEBUG] Processing {len(messages)} messages")
+                    
+                    # Find the last user message (current query) and extract conversation history (excluding current)
+                    current_user_message_found = False
+                    for i in range(len(messages) - 1, -1, -1):
+                        msg = messages[i]
+                        
+                        if isinstance(msg, dict):
+                            role = msg.get("role")
+                            content = msg.get("content", "")
+                        elif hasattr(msg, 'role'):
+                            role = msg.role
+                            content = getattr(msg, 'content', "")
+                            # Handle both string content and complex content structures
+                            if not isinstance(content, str):
+                                # Try to extract text from complex content
+                                if hasattr(content, '__iter__') and not isinstance(content, str):
+                                    try:
+                                        # If it's a list, try to join text parts
+                                        content = " ".join(str(item) for item in content if item)
+                                    except:
+                                        content = str(content)
+                                else:
+                                    content = str(content)
+                        else:
+                            continue
+                        
+                        # If we haven't found the current user message yet and this is a user message
+                        if not current_user_message_found and (role == "user" or role == 'user'):
+                            user_query = content
+                            current_user_message_found = True
+                            print(f"[LANGFUSE DEBUG] Found current user query: {user_query}")
+                        elif current_user_message_found:
+                            # Add to conversation history (excluding the current user message)
+                            conversation_history.insert(0, {
+                                "role": role,
+                                "content": content,
+                                "timestamp": datetime.now().isoformat() if not hasattr(msg, 'timestamp') else getattr(msg, 'timestamp', datetime.now().isoformat())
+                            })
                 
                 print(f"[LANGFUSE DEBUG] Final extracted - user_query: {user_query}, user_id: {user_id}")
+                print(f"[LANGFUSE DEBUG] Conversation history length: {len(conversation_history)}")
                 
                 # Create comprehensive input data for the trace
                 trace_input = {
@@ -437,31 +472,58 @@ class LangfuseTraceCollector:
                     session_id=event.data.get("session_id"),
                     input=trace_input,
                     metadata={
-                        "framework": "jaf", 
-                        "event_type": "run_start", 
+                        "framework": "jaf",
+                        "event_type": "run_start",
                         "trace_id": str(trace_id),
                         "user_query": user_query,
                         "user_id": user_id or event.data.get("user_id"),
-                        "agent_name": event.data.get("agent_name", "analytics_agent_jaf")
+                        "agent_name": event.data.get("agent_name", "analytics_agent_jaf"),
+                        "conversation_history": conversation_history,
+                        "tool_calls": [],
+                        "tool_results": []
                     }
                 )
                 self.trace_spans[trace_id] = trace
-                # Store user_id and user_query for later use in generations
+                # Store user_id, user_query, and conversation_history for later use
                 trace._user_id = user_id or event.data.get("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'}...")
                 
             elif event.type == "run_end":
                 if trace_id in self.trace_spans:
                     print(f"[LANGFUSE] Ending trace for run: {trace_id}")
-                    # End the trace 
-                    self.trace_spans[trace_id].update(output=event.data)
+                    
+                    # Update the trace metadata with final tool calls and results
+                    final_metadata = {
+                        "framework": "jaf",
+                        "event_type": "run_end",
+                        "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"),
+                        "conversation_history": getattr(self.trace_spans[trace_id], '_conversation_history', []),
+                        "tool_calls": self.trace_tool_calls.get(trace_id, []),
+                        "tool_results": self.trace_tool_results.get(trace_id, [])
+                    }
+                    
+                    # End the trace with updated metadata
+                    self.trace_spans[trace_id].update(
+                        output=event.data,
+                        metadata=final_metadata
+                    )
+                    
                     # Flush to ensure data is sent
                     print(f"[LANGFUSE] Flushing data to Langfuse...")
                     self.langfuse.flush()
                     print(f"[LANGFUSE] Flush completed")
+                    
                     # Clean up
                     del self.trace_spans[trace_id]
+                    if trace_id in self.trace_tool_calls:
+                        del self.trace_tool_calls[trace_id]
+                    if trace_id in self.trace_tool_results:
+                        del self.trace_tool_results[trace_id]
                 else:
                     print(f"[LANGFUSE] No trace found for run_end: {trace_id}")
                     
@@ -549,6 +611,20 @@ class LangfuseTraceCollector:
                 
                 print(f"[LANGFUSE] Starting span for tool call: {tool_name}")
                 
+                # Track this tool call for the trace
+                tool_call_data = {
+                    "tool_name": tool_name,
+                    "arguments": tool_args,
+                    "call_id": event.data.get("call_id"),
+                    "timestamp": datetime.now().isoformat()
+                }
+                
+                # Ensure trace_id exists in tracking
+                if trace_id not in self.trace_tool_calls:
+                    self.trace_tool_calls[trace_id] = []
+                
+                self.trace_tool_calls[trace_id].append(tool_call_data)
+                
                 # Create comprehensive input data for the tool call
                 tool_input = {
                     "tool_name": tool_name,
@@ -579,13 +655,28 @@ class LangfuseTraceCollector:
                     
                     print(f"[LANGFUSE] Ending span for tool call: {tool_name}")
                     
+                    # Track this tool result for the trace
+                    tool_result_data = {
+                        "tool_name": tool_name,
+                        "result": tool_result,
+                        "call_id": event.data.get("call_id"),
+                        "timestamp": datetime.now().isoformat(),
+                        "status": event.data.get("status", "completed"),
+                        "tool_result": event.data.get("tool_result")
+                    }
+                    
+                    if trace_id not in self.trace_tool_results:
+                        self.trace_tool_results[trace_id] = []
+                    
+                    self.trace_tool_results[trace_id].append(tool_result_data)
+                    
                     # Create comprehensive output data for the tool call
                     tool_output = {
                         "tool_name": tool_name,
                         "result": tool_result,
                         "call_id": event.data.get("call_id"),
                         "timestamp": datetime.now().isoformat(),
-                        "status": "completed"
+                        "status": event.data.get("status", "completed")
                     }
                     
                     # End the span with detailed output

jaf/core/types.py

@@ -11,6 +11,7 @@ from collections.abc import Awaitable, AsyncIterator
 from dataclasses import dataclass, field
 from typing import (
     Any,
+    Awaitable,
     Callable,
     Dict,
     Generic,
@@ -28,6 +29,8 @@ from enum import Enum
 
 if TYPE_CHECKING:
     from .tool_results import ToolResult
+    from ..memory.approval_storage import ApprovalStorage
+    from ..memory.types import MemoryConfig
 
 
 # Comprehensive enums for type safety and improved developer experience
@@ -143,13 +146,88 @@ class ToolCallFunction:
     name: str
     arguments: str
 
+@dataclass(frozen=True)
+class Attachment:
+    """Represents an attachment with various content types."""
+    kind: Literal['image', 'document', 'file']
+    mime_type: Optional[str] = None  # e.g. image/png, application/pdf
+    name: Optional[str] = None       # Optional filename
+    url: Optional[str] = None        # Remote URL or data URL
+    data: Optional[str] = None       # Base64 without data: prefix
+    format: Optional[str] = None     # Optional short format like 'pdf', 'txt'
+    use_litellm_format: Optional[bool] = None  # Use LiteLLM native file format
+    
+    def __post_init__(self):
+        """Validate that at least one of url or data is provided."""
+        if self.url is None and self.data is None:
+            raise ValueError("At least one of 'url' or 'data' must be provided for an Attachment.")
+
+@dataclass(frozen=True)
+class MessageContentPart:
+    """Part of multi-part message content."""
+    type: Literal['text', 'image_url', 'file']
+    text: Optional[str] = None
+    image_url: Optional[Dict[str, Any]] = None  # Contains url and optional detail
+    file: Optional[Dict[str, Any]] = None       # Contains file_id and optional format
+
 @dataclass(frozen=True)
 class Message:
-    """A message in the conversation."""
+    """
+    A message in the conversation.
+    
+    BACKWARDS COMPATIBILITY:
+    - Messages created with string content remain fully backwards compatible
+    - 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
+    
+    Examples:
+        # Original usage - still works exactly the same
+        msg = Message(role='user', content='Hello')
+        text = msg.content  # Returns 'Hello' as string
+        
+        # Guaranteed string access (recommended for new code)
+        text = msg.text_content  # Always returns string
+        
+        # Universal text extraction
+        text = get_text_content(msg.content)  # Works with any content type
+    """
     role: ContentRole
-    content: str
+    content: Union[str, List[MessageContentPart]]
+    attachments: Optional[List[Attachment]] = None
     tool_call_id: Optional[str] = None
     tool_calls: Optional[List[ToolCall]] = None
+    
+    @property
+    def text_content(self) -> str:
+        """Get text content as string for backwards compatibility."""
+        return get_text_content(self.content)
+    
+    @classmethod
+    def create(
+        cls,
+        role: ContentRole,
+        content: str,
+        attachments: Optional[List[Attachment]] = None,
+        tool_call_id: Optional[str] = None,
+        tool_calls: Optional[List[ToolCall]] = None
+    ) -> 'Message':
+        """Create a message with string content and optional attachments."""
+        return cls(
+            role=role,
+            content=content,
+            attachments=attachments,
+            tool_call_id=tool_call_id,
+            tool_calls=tool_calls
+        )
+
+def get_text_content(content: Union[str, List[MessageContentPart]]) -> str:
+    """Extract text content from message content."""
+    if isinstance(content, str):
+        return content
+    
+    text_parts = [part.text for part in content if part.type == 'text' and part.text]
+    return ' '.join(text_parts)
 
 @dataclass(frozen=True)
 class ModelConfig:
@@ -179,6 +257,11 @@ class Tool(Protocol[Args, Ctx]):
         """Execute the tool with given arguments and context."""
         ...
 
+    @property
+    def needs_approval(self) -> Union[bool, Callable[[Ctx, Args], Union[bool, Awaitable[bool]]]]:
+        """Whether this tool requires approval before execution."""
+        return False
+
 
 # Function tool configuration for improved DX
 class FunctionToolConfig(TypedDict):
@@ -248,6 +331,13 @@ class Agent(Generic[Ctx, Out]):
 # Guardrail type
 Guardrail = Callable[[Any], Union[ValidationResult, Awaitable[ValidationResult]]]
 
+@dataclass(frozen=True)
+class ApprovalValue:
+    """Represents an approval decision with context."""
+    status: str  # 'pending', 'approved', 'rejected'
+    approved: bool
+    additional_context: Optional[Dict[str, Any]] = None
+
 @dataclass(frozen=True)
 class RunState(Generic[Ctx]):
     """Immutable state of a run."""
@@ -257,6 +347,7 @@ class RunState(Generic[Ctx]):
     current_agent_name: str
     context: Ctx
     turn_count: int
+    approvals: Dict[str, ApprovalValue] = field(default_factory=dict)
 
 # Error types using dataclasses for immutability
 @dataclass(frozen=True)
@@ -335,6 +426,18 @@ class NetworkError:
     is_retryable: bool = True
     endpoint: Optional[str] = None
 
+# Interruption types for HITL
+@dataclass(frozen=True)
+class ToolApprovalInterruption(Generic[Ctx]):
+    """Interruption for tool approval."""
+    type: Literal['tool_approval'] = 'tool_approval'
+    tool_call: ToolCall = field(default_factory=lambda: ToolCall("", "function", ToolCallFunction("", "")))
+    agent: 'Agent[Ctx, Any]' = None
+    session_id: Optional[str] = None
+
+# Union type for all interruptions
+Interruption = Union[ToolApprovalInterruption[Any]]
+
 # Union type for all possible errors
 JAFError = Union[
     MaxTurnsExceeded,
@@ -363,8 +466,14 @@ class ErrorOutcome:
     status: Literal['error'] = 'error'
     error: JAFError = field(default=None)
 
+@dataclass(frozen=True)
+class InterruptedOutcome:
+    """Interrupted outcome for HITL."""
+    status: Literal['interrupted'] = 'interrupted'
+    interruptions: List[Interruption] = field(default_factory=list)
+
 # Union type for outcomes
-RunOutcome = Union[CompletedOutcome[Out], ErrorOutcome]
+RunOutcome = Union[CompletedOutcome[Out], ErrorOutcome, InterruptedOutcome]
 
 @dataclass(frozen=True)
 class RunResult(Generic[Out]):
@@ -601,3 +710,4 @@ class RunConfig(Generic[Ctx]):
     memory: Optional['MemoryConfig'] = None
     conversation_id: Optional[str] = None
     default_tool_timeout: Optional[float] = 30.0  # Default timeout for tool execution in seconds
+    approval_storage: Optional['ApprovalStorage'] = None  # Storage for approval decisions