jaf-py 2.5.9__py3-none-any.whl → 2.5.11__py3-none-any.whl

jaf/core/streaming.py

@@ -14,26 +14,38 @@ from typing import AsyncIterator, Dict, List, Optional, Any, Union, Callable
 from enum import Enum
 
 from .types import (
-    RunState, RunConfig, Message, TraceEvent, RunId, TraceId,
-    ContentRole, ToolCall, JAFError, CompletedOutcome, ErrorOutcome, ModelBehaviorError
+    RunState,
+    RunConfig,
+    Message,
+    TraceEvent,
+    RunId,
+    TraceId,
+    ContentRole,
+    ToolCall,
+    JAFError,
+    CompletedOutcome,
+    ErrorOutcome,
+    ModelBehaviorError,
 )
 
 
 class StreamingEventType(str, Enum):
     """Types of streaming events."""
-    START = 'start'
-    CHUNK = 'chunk'
-    TOOL_CALL = 'tool_call'
-    TOOL_RESULT = 'tool_result'
-    AGENT_SWITCH = 'agent_switch'
-    ERROR = 'error'
-    COMPLETE = 'complete'
-    METADATA = 'metadata'
+
+    START = "start"
+    CHUNK = "chunk"
+    TOOL_CALL = "tool_call"
+    TOOL_RESULT = "tool_result"
+    AGENT_SWITCH = "agent_switch"
+    ERROR = "error"
+    COMPLETE = "complete"
+    METADATA = "metadata"
 
 
 @dataclass(frozen=True)
 class StreamingChunk:
     """A chunk of streaming content."""
+
     content: str
     delta: str  # The new content added in this chunk
     is_complete: bool = False
@@ -43,15 +55,17 @@ class StreamingChunk:
 @dataclass(frozen=True)
 class StreamingToolCall:
     """Streaming tool call information."""
+
     tool_name: str
     arguments: Dict[str, Any]
     call_id: str
-    status: str = 'started'  # 'started', 'executing', 'completed', 'failed'
+    status: str = "started"  # 'started', 'executing', 'completed', 'failed'
 
 
 @dataclass(frozen=True)
 class StreamingToolResult:
     """Streaming tool result information."""
+
     tool_name: str
     call_id: str
     result: Any
@@ -62,6 +76,7 @@ class StreamingToolResult:
 @dataclass(frozen=True)
 class StreamingMetadata:
     """Metadata about the streaming session."""
+
     agent_name: str
     model_name: str
     turn_count: int
@@ -73,25 +88,28 @@ class StreamingMetadata:
 @dataclass(frozen=True)
 class StreamingEvent:
     """A streaming event containing progressive updates."""
+
     type: StreamingEventType
     data: Union[StreamingChunk, StreamingToolCall, StreamingToolResult, StreamingMetadata, JAFError]
     timestamp: float = field(default_factory=time.time)
     run_id: Optional[RunId] = None
     trace_id: Optional[TraceId] = None
-    
+
     def to_dict(self) -> Dict[str, Any]:
         """Convert streaming event to dictionary for serialization."""
         return {
-            'type': self.type.value,
-            'data': self._serialize_data(),
-            'timestamp': self.timestamp,
-            'run_id': str(self.run_id) if self.run_id else None,
-            'trace_id': str(self.trace_id) if self.trace_id else None
+            "type": self.type.value,
+            "data": self._serialize_data(),
+            "timestamp": self.timestamp,
+            "run_id": str(self.run_id) if self.run_id else None,
+            "trace_id": str(self.trace_id) if self.trace_id else None,
         }
-    
+
     def _serialize_data(self) -> Dict[str, Any]:
         """Serialize the data field based on its type."""
-        if isinstance(self.data, (StreamingChunk, StreamingToolCall, StreamingToolResult, StreamingMetadata)):
+        if isinstance(
+            self.data, (StreamingChunk, StreamingToolCall, StreamingToolResult, StreamingMetadata)
+        ):
             # Convert dataclass to dict
             result = {}
             for field_name, field_value in self.data.__dict__.items():
@@ -100,12 +118,12 @@ class StreamingEvent:
             return result
         elif isinstance(self.data, JAFError):
             return {
-                'error_type': self.data._tag,
-                'detail': getattr(self.data, 'detail', str(self.data))
+                "error_type": self.data._tag,
+                "detail": getattr(self.data, "detail", str(self.data)),
             }
         else:
-            return {'value': self.data}
-    
+            return {"value": self.data}
+
     def to_json(self) -> str:
         """Convert streaming event to JSON string."""
         return json.dumps(self.to_dict())
@@ -115,7 +133,7 @@ class StreamingBuffer:
     """
     Buffer for accumulating streaming content and managing state.
     """
-    
+
     def __init__(self):
         self.content: str = ""
         self.chunks: List[StreamingChunk] = []
@@ -124,31 +142,31 @@ class StreamingBuffer:
         self.metadata: Optional[StreamingMetadata] = None
         self.is_complete: bool = False
         self.error: Optional[JAFError] = None
-    
+
     def add_chunk(self, chunk: StreamingChunk) -> None:
         """Add a content chunk to the buffer."""
         self.chunks.append(chunk)
         self.content += chunk.delta
         if chunk.is_complete:
             self.is_complete = True
-    
+
     def add_tool_call(self, tool_call: StreamingToolCall) -> None:
         """Add a tool call to the buffer."""
         self.tool_calls.append(tool_call)
-    
+
     def add_tool_result(self, tool_result: StreamingToolResult) -> None:
         """Add a tool result to the buffer."""
         self.tool_results.append(tool_result)
-    
+
     def set_metadata(self, metadata: StreamingMetadata) -> None:
         """Set session metadata."""
         self.metadata = metadata
-    
+
     def set_error(self, error: JAFError) -> None:
         """Set error state."""
         self.error = error
         self.is_complete = True
-    
+
     def get_final_message(self) -> Message:
         """Get the final accumulated message."""
         tool_calls = None
@@ -156,38 +174,31 @@ class StreamingBuffer:
             tool_calls = [
                 ToolCall(
                     id=tc.call_id,
-                    type='function',
-                    function={'name': tc.tool_name, 'arguments': json.dumps(tc.arguments)}
+                    type="function",
+                    function={"name": tc.tool_name, "arguments": json.dumps(tc.arguments)},
                 )
                 for tc in self.tool_calls
             ]
-        
-        return Message(
-            role=ContentRole.ASSISTANT,
-            content=self.content,
-            tool_calls=tool_calls
-        )
+
+        return Message(role=ContentRole.ASSISTANT, content=self.content, tool_calls=tool_calls)
 
 
 async def run_streaming(
-    initial_state: RunState,
-    config: RunConfig,
-    chunk_size: int = 50,
-    include_metadata: bool = True
+    initial_state: RunState, config: RunConfig, chunk_size: int = 50, include_metadata: bool = True
 ) -> AsyncIterator[StreamingEvent]:
     """
     Run an agent with streaming output.
-    
+
     This function provides real-time streaming of agent responses, tool calls,
     and execution metadata. It yields StreamingEvent objects that can be
     consumed by clients for progressive UI updates.
-    
+
     Args:
         initial_state: Initial run state
         config: Run configuration
         chunk_size: Size of content chunks for streaming (characters)
         include_metadata: Whether to include performance metadata
-        
+
     Yields:
         StreamingEvent: Progressive updates during execution
     """
@@ -203,10 +214,10 @@ async def run_streaming(
             model_name="unknown",
             turn_count=initial_state.turn_count,
             total_tokens=0,
-            execution_time_ms=0
+            execution_time_ms=0,
         ),
         run_id=initial_state.run_id,
-        trace_id=initial_state.trace_id
+        trace_id=initial_state.trace_id,
     )
 
     tool_call_ids: Dict[str, str] = {}  # Map call_id -> tool_name for in-flight tool calls
@@ -225,37 +236,36 @@ async def run_streaming(
                     return getattr(payload, key)
             return None
 
-        if event.type == 'tool_call_start':
-            tool_name = _get_event_value(['tool_name', 'toolName']) or 'unknown'
-            args = _get_event_value(['args', 'arguments'])
-            call_id = _get_event_value(['call_id', 'tool_call_id', 'toolCallId'])
+        if event.type == "tool_call_start":
+            tool_name = _get_event_value(["tool_name", "toolName"]) or "unknown"
+            args = _get_event_value(["args", "arguments"])
+            call_id = _get_event_value(["call_id", "tool_call_id", "toolCallId"])
 
             if not call_id:
                 call_id = f"call_{uuid.uuid4().hex[:8]}"
                 if isinstance(payload, dict):
-                    payload['call_id'] = call_id
+                    payload["call_id"] = call_id
 
             tool_call_ids[call_id] = tool_name
 
             tool_call = StreamingToolCall(
-                tool_name=tool_name,
-                arguments=args,
-                call_id=call_id,
-                status='started'
+                tool_name=tool_name, arguments=args, call_id=call_id, status="started"
             )
             streaming_event = StreamingEvent(
                 type=StreamingEventType.TOOL_CALL,
                 data=tool_call,
                 run_id=initial_state.run_id,
-                trace_id=initial_state.trace_id
+                trace_id=initial_state.trace_id,
             )
-        elif event.type == 'tool_call_end':
-            tool_name = _get_event_value(['tool_name', 'toolName']) or 'unknown'
-            call_id = _get_event_value(['call_id', 'tool_call_id', 'toolCallId'])
+        elif event.type == "tool_call_end":
+            tool_name = _get_event_value(["tool_name", "toolName"]) or "unknown"
+            call_id = _get_event_value(["call_id", "tool_call_id", "toolCallId"])
 
             if not call_id:
                 # Fallback to locate a pending tool call with the same tool name
-                matching_call_id = next((cid for cid, name in tool_call_ids.items() if name == tool_name), None)
+                matching_call_id = next(
+                    (cid for cid, name in tool_call_ids.items() if name == tool_name), None
+                )
                 if matching_call_id:
                     call_id = matching_call_id
                 else:
@@ -268,16 +278,16 @@ async def run_streaming(
             tool_result = StreamingToolResult(
                 tool_name=tool_name,
                 call_id=call_id,
-                result=_get_event_value(['result']),
-                status=_get_event_value(['status']) or 'completed'
+                result=_get_event_value(["result"]),
+                status=_get_event_value(["status"]) or "completed",
             )
             streaming_event = StreamingEvent(
                 type=StreamingEventType.TOOL_RESULT,
                 data=tool_result,
                 run_id=initial_state.run_id,
-                trace_id=initial_state.trace_id
+                trace_id=initial_state.trace_id,
             )
-        
+
         if streaming_event:
             try:
                 event_queue.put_nowait(streaming_event)
@@ -296,7 +306,7 @@ async def run_streaming(
         on_event=event_handler,
         memory=config.memory,
         conversation_id=config.conversation_id,
-        prefer_streaming=config.prefer_streaming
+        prefer_streaming=config.prefer_streaming,
     )
 
     from .engine import run
@@ -330,30 +340,30 @@ async def run_streaming(
             type=StreamingEventType.ERROR,
             data=error,
             run_id=initial_state.run_id,
-            trace_id=initial_state.trace_id
+            trace_id=initial_state.trace_id,
         )
         return
 
-    if result.outcome.status == 'completed':
+    if result.outcome.status == "completed":
         final_content = str(result.outcome.output) if result.outcome.output else ""
-        
+
         # Stream content in chunks
         for i in range(0, len(final_content), chunk_size):
-            chunk_content = final_content[i:i + chunk_size]
+            chunk_content = final_content[i : i + chunk_size]
             is_final_chunk = i + chunk_size >= len(final_content)
-            
+
             chunk = StreamingChunk(
-                content=final_content[:i + len(chunk_content)],
+                content=final_content[: i + len(chunk_content)],
                 delta=chunk_content,
                 is_complete=is_final_chunk,
-                token_count=len(final_content.split()) if is_final_chunk else None
+                token_count=len(final_content.split()) if is_final_chunk else None,
             )
-            
+
             yield StreamingEvent(
                 type=StreamingEventType.CHUNK,
                 data=chunk,
                 run_id=initial_state.run_id,
-                trace_id=initial_state.trace_id
+                trace_id=initial_state.trace_id,
             )
             # Remove artificial delay for better performance
 
@@ -364,27 +374,27 @@ async def run_streaming(
                 model_name=config.model_override or "default",
                 turn_count=result.final_state.turn_count,
                 total_tokens=len(final_content.split()),
-                execution_time_ms=execution_time
+                execution_time_ms=execution_time,
             )
             yield StreamingEvent(
                 type=StreamingEventType.METADATA,
                 data=metadata,
                 run_id=initial_state.run_id,
-                trace_id=initial_state.trace_id
+                trace_id=initial_state.trace_id,
             )
-        
+
         yield StreamingEvent(
             type=StreamingEventType.COMPLETE,
             data=StreamingChunk(content=final_content, delta="", is_complete=True),
             run_id=initial_state.run_id,
-            trace_id=initial_state.trace_id
+            trace_id=initial_state.trace_id,
         )
     else:
         yield StreamingEvent(
             type=StreamingEventType.ERROR,
             data=result.outcome.error,
             run_id=initial_state.run_id,
-            trace_id=initial_state.trace_id
+            trace_id=initial_state.trace_id,
         )
 
 
@@ -392,33 +402,31 @@ class StreamingCollector:
     """
     Collects streaming events for analysis and replay.
     """
-    
+
     def __init__(self):
         self.events: List[StreamingEvent] = []
         self.buffers: Dict[str, StreamingBuffer] = {}
-    
+
     async def collect_stream(
-        self,
-        stream: AsyncIterator[StreamingEvent],
-        run_id: Optional[str] = None
+        self, stream: AsyncIterator[StreamingEvent], run_id: Optional[str] = None
     ) -> StreamingBuffer:
         """
         Collect all events from a stream and return the final buffer.
-        
+
         Args:
             stream: Async iterator of streaming events
             run_id: Optional run ID for tracking
-            
+
         Returns:
             StreamingBuffer: Final accumulated buffer
         """
         buffer_key = run_id or "default"
         buffer = StreamingBuffer()
         self.buffers[buffer_key] = buffer
-        
+
         async for event in stream:
             self.events.append(event)
-            
+
             if event.type == StreamingEventType.CHUNK:
                 buffer.add_chunk(event.data)
             elif event.type == StreamingEventType.TOOL_CALL:
@@ -431,46 +439,45 @@ class StreamingCollector:
                 buffer.set_error(event.data)
             elif event.type == StreamingEventType.COMPLETE:
                 buffer.is_complete = True
-        
+
         return buffer
-    
+
     def get_events_for_run(self, run_id: str) -> List[StreamingEvent]:
         """Get all events for a specific run."""
-        return [
-            event for event in self.events
-            if event.run_id and str(event.run_id) == run_id
-        ]
-    
+        return [event for event in self.events if event.run_id and str(event.run_id) == run_id]
+
     def replay_stream(self, run_id: str, delay_ms: int = 50) -> AsyncIterator[StreamingEvent]:
         """
         Replay a collected stream with optional delay.
-        
+
         Args:
             run_id: Run ID to replay
             delay_ms: Delay between events in milliseconds
-            
+
         Yields:
             StreamingEvent: Replayed events
         """
+
         async def _replay():
             events = self.get_events_for_run(run_id)
             for event in events:
                 yield event
                 if delay_ms > 0:
                     await asyncio.sleep(delay_ms / 1000)
-        
+
         return _replay()
 
 
 # Utility functions for streaming integration
 
+
 def create_sse_response(event: StreamingEvent) -> str:
     """
     Create a Server-Sent Events (SSE) formatted response.
-    
+
     Args:
         event: Streaming event to format
-        
+
     Returns:
         str: SSE-formatted string
     """
@@ -478,12 +485,11 @@ def create_sse_response(event: StreamingEvent) -> str:
 
 
 async def stream_to_websocket(
-    stream: AsyncIterator[StreamingEvent],
-    websocket_send: Callable[[str], None]
+    stream: AsyncIterator[StreamingEvent], websocket_send: Callable[[str], None]
 ) -> None:
     """
     Stream events to a WebSocket connection.
-    
+
     Args:
         stream: Stream of events
         websocket_send: WebSocket send function
@@ -493,21 +499,22 @@ async def stream_to_websocket(
 
 
 def create_streaming_middleware(
-    on_event: Optional[Callable[[StreamingEvent], None]] = None
+    on_event: Optional[Callable[[StreamingEvent], None]] = None,
 ) -> Callable[[AsyncIterator[StreamingEvent]], AsyncIterator[StreamingEvent]]:
     """
     Create middleware for processing streaming events.
-    
+
     Args:
         on_event: Optional callback for each event
-        
+
     Returns:
         Middleware function
     """
+
     async def middleware(stream: AsyncIterator[StreamingEvent]) -> AsyncIterator[StreamingEvent]:
         async for event in stream:
             if on_event:
                 on_event(event)
             yield event
-    
+
     return middleware