autobyteus 1.2.0__py3-none-any.whl → 1.2.3__py3-none-any.whl

autobyteus/agent/streaming/handlers/api_tool_call_streaming_response_handler.py

@@ -0,0 +1,303 @@
+# file: autobyteus/autobyteus/agent/streaming/api_tool_call_streaming_response_handler.py
+"""
+ApiToolCallStreamingResponseHandler: Handler for API-provided tool calls.
+
+This handler processes SDK-provided tool calls from providers like OpenAI,
+Anthropic, and Gemini. It emits SegmentEvents and uses an internal
+ToolInvocationAdapter to create ToolInvocations.
+"""
+import json
+import uuid
+import logging
+from dataclasses import dataclass
+from typing import Optional, List, Dict, Callable
+
+from .streaming_response_handler import StreamingResponseHandler
+from ..segments.segment_events import SegmentEvent, SegmentType, SegmentEventType
+from ..adapters.invocation_adapter import ToolInvocationAdapter
+from ..api_tool_call.file_content_streamer import WriteFileContentStreamer, PatchFileContentStreamer
+from autobyteus.agent.tool_invocation import ToolInvocation
+from autobyteus.llm.utils.response_types import ChunkResponse
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ToolCallState:
+    """Tracks the state of an in-progress tool call."""
+    segment_id: str
+    name: str
+    accumulated_args: str = ""
+    segment_type: SegmentType = SegmentType.TOOL_CALL
+    streamer: Optional[object] = None
+    path: Optional[str] = None
+    segment_started: bool = False
+    pending_content: str = ""
+
+
+class ApiToolCallStreamingResponseHandler(StreamingResponseHandler):
+    """
+    Handler for API-provided tool calls (OpenAI, Anthropic, Gemini native tool calling).
+
+    Responsibilities:
+    1. Emit TEXT segments for text content
+    2. Emit TOOL_CALL segments for SDK-provided tool calls
+    3. Use internal ToolInvocationAdapter to create ToolInvocations
+
+    Key Design:
+    - Handler emits SegmentEvents
+    - Internal adapter processes events to create ToolInvocations
+    - get_all_invocations() returns adapter results (consistent interface)
+    """
+
+    def __init__(
+        self,
+        on_segment_event: Optional[Callable[[SegmentEvent], None]] = None,
+        on_tool_invocation: Optional[Callable[[ToolInvocation], None]] = None,
+        segment_id_prefix: str = "",
+    ):
+        self._on_segment_event = on_segment_event
+        self._on_tool_invocation = on_tool_invocation
+        self._segment_id_prefix = segment_id_prefix
+
+        # Internal adapter for creating invocations from events
+        self._adapter = ToolInvocationAdapter()
+
+        # State tracking
+        self._text_segment_id: Optional[str] = None
+        self._active_tools: Dict[int, ToolCallState] = {}  # index -> state
+        self._all_events: List[SegmentEvent] = []
+        self._all_invocations: List[ToolInvocation] = []
+        self._is_finalized = False
+
+    def _generate_id(self) -> str:
+        return f"{self._segment_id_prefix}{uuid.uuid4().hex}"
+
+    @staticmethod
+    def _resolve_segment_type(tool_name: str):
+        if tool_name == "write_file":
+            return SegmentType.WRITE_FILE, WriteFileContentStreamer()
+        if tool_name == "patch_file":
+            return SegmentType.PATCH_FILE, PatchFileContentStreamer()
+        return SegmentType.TOOL_CALL, None
+
+    def _emit(self, event: SegmentEvent) -> None:
+        """Emit event and process through internal adapter."""
+        self._all_events.append(event)
+        
+        # Notify callback
+        if self._on_segment_event:
+            try:
+                self._on_segment_event(event)
+            except Exception as e:
+                logger.error(f"Error in on_segment_event callback: {e}")
+
+        # Process through internal adapter
+        invocation = self._adapter.process_event(event)
+        if invocation:
+            self._all_invocations.append(invocation)
+            if self._on_tool_invocation:
+                try:
+                    self._on_tool_invocation(invocation)
+                except Exception as e:
+                    logger.error(f"Error in on_tool_invocation callback: {e}")
+
+    def feed(self, chunk: ChunkResponse) -> List[SegmentEvent]:
+        if self._is_finalized:
+            raise RuntimeError("Handler has been finalized.")
+
+        events = []
+
+        # 1. Handle text content → TEXT segment
+        if chunk.content:
+            if self._text_segment_id is None:
+                self._text_segment_id = self._generate_id()
+                start_event = SegmentEvent.start(
+                    segment_id=self._text_segment_id,
+                    segment_type=SegmentType.TEXT
+                )
+                self._emit(start_event)
+                events.append(start_event)
+
+            content_event = SegmentEvent.content(
+                segment_id=self._text_segment_id,
+                delta=chunk.content
+            )
+            self._emit(content_event)
+            events.append(content_event)
+
+        # 2. Handle tool calls from SDK
+        if chunk.tool_calls:
+            for delta in chunk.tool_calls:
+                if delta.index not in self._active_tools:
+                    # New tool call - emit SEGMENT_START
+                    seg_id = delta.call_id or self._generate_id()
+                    tool_name = delta.name or ""
+                    segment_type, streamer = self._resolve_segment_type(tool_name)
+                    self._active_tools[delta.index] = ToolCallState(
+                        segment_id=seg_id,
+                        name=tool_name,
+                        accumulated_args="",
+                        segment_type=segment_type,
+                        streamer=streamer,
+                    )
+                    if segment_type == SegmentType.TOOL_CALL:
+                        start_event = SegmentEvent.start(
+                            segment_id=seg_id,
+                            segment_type=segment_type,
+                            tool_name=tool_name,
+                        )
+                        self._active_tools[delta.index].segment_started = True
+                        self._emit(start_event)
+                        events.append(start_event)
+
+                # Accumulate arguments and emit content delta (for UI streaming)
+                if delta.arguments_delta:
+                    state = self._active_tools[delta.index]
+                    state.accumulated_args += delta.arguments_delta
+
+                    if state.segment_type == SegmentType.TOOL_CALL:
+                        if not state.segment_started:
+                            start_event = SegmentEvent.start(
+                                segment_id=state.segment_id,
+                                segment_type=state.segment_type,
+                                tool_name=state.name,
+                            )
+                            state.segment_started = True
+                            self._emit(start_event)
+                            events.append(start_event)
+                        content_event = SegmentEvent.content(
+                            segment_id=state.segment_id,
+                            delta=delta.arguments_delta,
+                        )
+                        self._emit(content_event)
+                        events.append(content_event)
+                    else:
+                        update = state.streamer.feed(delta.arguments_delta) if state.streamer else None
+                        if update and update.path and not state.path:
+                            state.path = update.path
+                        if not state.segment_started and state.path:
+                            start_event = SegmentEvent.start(
+                                segment_id=state.segment_id,
+                                segment_type=state.segment_type,
+                                tool_name=state.name,
+                                path=state.path,
+                            )
+                            state.segment_started = True
+                            self._emit(start_event)
+                            events.append(start_event)
+                            if state.pending_content:
+                                content_event = SegmentEvent.content(
+                                    segment_id=state.segment_id,
+                                    delta=state.pending_content,
+                                )
+                                self._emit(content_event)
+                                events.append(content_event)
+                                state.pending_content = ""
+                        if update and update.content_delta:
+                            if state.segment_started:
+                                content_event = SegmentEvent.content(
+                                    segment_id=state.segment_id,
+                                    delta=update.content_delta,
+                                )
+                                self._emit(content_event)
+                                events.append(content_event)
+                            else:
+                                state.pending_content += update.content_delta
+
+                # Update name if provided later
+                if delta.name and not self._active_tools[delta.index].name:
+                    self._active_tools[delta.index].name = delta.name
+
+        return events
+
+    def finalize(self) -> List[SegmentEvent]:
+        if self._is_finalized:
+            return []
+
+        self._is_finalized = True
+        events = []
+
+        # Close text segment
+        if self._text_segment_id:
+            end_event = SegmentEvent.end(segment_id=self._text_segment_id)
+            self._emit(end_event)
+            events.append(end_event)
+
+        # Close tool segments with pre-parsed arguments in metadata
+        for state in self._active_tools.values():
+            if state.segment_type in {SegmentType.WRITE_FILE, SegmentType.PATCH_FILE}:
+                if not state.segment_started:
+                    start_meta = {"tool_name": state.name}
+                    if state.path:
+                        start_meta["path"] = state.path
+                    start_event = SegmentEvent.start(
+                        segment_id=state.segment_id,
+                        segment_type=state.segment_type,
+                        **start_meta,
+                    )
+                    state.segment_started = True
+                    self._emit(start_event)
+                    events.append(start_event)
+                    if state.pending_content:
+                        content_event = SegmentEvent.content(
+                            segment_id=state.segment_id,
+                            delta=state.pending_content,
+                        )
+                        self._emit(content_event)
+                        events.append(content_event)
+                        state.pending_content = ""
+            if state.segment_type == SegmentType.TOOL_CALL:
+                # Parse accumulated JSON arguments
+                try:
+                    parsed_args = json.loads(state.accumulated_args) if state.accumulated_args else {}
+                except json.JSONDecodeError as e:
+                    logger.error(f"Failed to parse tool arguments for {state.name}: {e}")
+                    parsed_args = {}
+
+                # Emit SEGMENT_END with arguments in metadata
+                # The internal adapter will use this to create ToolInvocation
+                end_event = SegmentEvent(
+                    event_type=SegmentEventType.END,
+                    segment_id=state.segment_id,
+                    payload={
+                        "metadata": {
+                            "tool_name": state.name,
+                            "arguments": parsed_args,  # Pre-parsed for adapter
+                        }
+                    }
+                )
+            else:
+                metadata = {}
+                if state.path:
+                    metadata["path"] = state.path
+                end_event = SegmentEvent(
+                    event_type=SegmentEventType.END,
+                    segment_id=state.segment_id,
+                    payload={"metadata": metadata} if metadata else {},
+                )
+            self._emit(end_event)
+            events.append(end_event)
+
+        if self._all_invocations:
+            logger.info(
+                "ApiToolCallStreamingResponseHandler finalized %d tool invocations.",
+                len(self._all_invocations),
+            )
+
+        return events
+
+    def get_all_events(self) -> List[SegmentEvent]:
+        return self._all_events.copy()
+
+    def get_all_invocations(self) -> List[ToolInvocation]:
+        """Returns invocations created by the internal adapter."""
+        return self._all_invocations.copy()
+
+    def reset(self) -> None:
+        self._text_segment_id = None
+        self._active_tools.clear()
+        self._all_events.clear()
+        self._all_invocations.clear()
+        self._adapter = ToolInvocationAdapter()
+        self._is_finalized = False

autobyteus/agent/streaming/handlers/parsing_streaming_response_handler.py

@@ -0,0 +1,107 @@
+"""
+ParsingStreamingResponseHandler: Concrete implementation of StreamingResponseHandler that uses a parser.
+"""
+from typing import Optional, List, Callable, Any, Union
+import logging
+
+from .streaming_response_handler import StreamingResponseHandler
+from ..parser.parser_factory import create_streaming_parser, resolve_parser_name
+from ..segments.segment_events import SegmentEvent
+from ..adapters.invocation_adapter import ToolInvocationAdapter
+from ..parser.parser_context import ParserConfig
+from autobyteus.agent.tool_invocation import ToolInvocation
+from autobyteus.llm.utils.response_types import ChunkResponse
+
+logger = logging.getLogger(__name__)
+
+
+class ParsingStreamingResponseHandler(StreamingResponseHandler):
+    """
+    Handler that uses a StreamingParser to process the LLM response.
+    """
+    
+    def __init__(
+        self,
+        on_segment_event: Optional[Callable[[SegmentEvent], None]] = None,
+        on_tool_invocation: Optional[Callable[[ToolInvocation], None]] = None,
+        config: Optional[ParserConfig] = None,
+        parser_name: Optional[str] = None,
+    ):
+        """
+        Initialize the parsing response handler.
+        """
+        self._parser_name = resolve_parser_name(parser_name)
+        self._parser_config = config
+        self._parser = create_streaming_parser(
+            config=config,
+            parser_name=self._parser_name,
+        )
+        self._adapter = ToolInvocationAdapter(
+            json_tool_parser=self._parser.config.json_tool_parser
+        )
+        self._on_segment_event = on_segment_event
+        self._on_tool_invocation = on_tool_invocation
+        self._is_finalized = False
+        
+        # Accumulated data
+        self._all_events: List[SegmentEvent] = []
+        self._all_invocations: List[ToolInvocation] = []
+
+    def feed(self, chunk: ChunkResponse) -> List[SegmentEvent]:
+        if self._is_finalized:
+            raise RuntimeError("Handler has been finalized, cannot feed more chunks.")
+        
+        # Extract text content from ChunkResponse (ignore tool_calls - not our concern)
+        text_content = chunk.content if isinstance(chunk, ChunkResponse) else chunk
+        if not text_content:
+            return []
+        
+        events = self._parser.feed(text_content)
+        self._process_events(events)
+        return events
+
+    def finalize(self) -> List[SegmentEvent]:
+        if self._is_finalized:
+            return []
+        
+        self._is_finalized = True
+        events = self._parser.finalize()
+        self._process_events(events)
+        return events
+
+    def _process_events(self, events: List[SegmentEvent]) -> None:
+        for event in events:
+            self._all_events.append(event)
+            
+            if self._on_segment_event:
+                try:
+                    self._on_segment_event(event)
+                except Exception as e:
+                    logger.error(f"Error in on_segment_event callback: {e}")
+            
+            invocation = self._adapter.process_event(event)
+            if invocation:
+                self._all_invocations.append(invocation)
+                if self._on_tool_invocation:
+                    try:
+                        self._on_tool_invocation(invocation)
+                    except Exception as e:
+                        logger.error(f"Error in on_tool_invocation callback: {e}")
+
+    def get_all_events(self) -> List[SegmentEvent]:
+        return self._all_events.copy()
+
+    def get_all_invocations(self) -> List[ToolInvocation]:
+        return self._all_invocations.copy()
+
+    def reset(self) -> None:
+        self._parser = create_streaming_parser(
+            config=self._parser_config,
+            parser_name=self._parser_name,
+        )
+        self._adapter = ToolInvocationAdapter(
+            json_tool_parser=self._parser.config.json_tool_parser
+        )
+        self._all_events.clear()
+        self._all_invocations.clear()
+        self._is_finalized = False

autobyteus/agent/streaming/handlers/pass_through_streaming_response_handler.py

@@ -0,0 +1,107 @@
+"""
+PassThroughStreamingResponseHandler: Implementation of StreamingResponseHandler that bypasses parsing.
+"""
+import uuid
+import logging
+from typing import Optional, List, Callable
+
+from .streaming_response_handler import StreamingResponseHandler
+from ..segments.segment_events import SegmentEvent, SegmentType, SegmentEventType
+from autobyteus.agent.tool_invocation import ToolInvocation
+from autobyteus.llm.utils.response_types import ChunkResponse
+
+logger = logging.getLogger(__name__)
+
+
+class PassThroughStreamingResponseHandler(StreamingResponseHandler):
+    """
+    Handler that passes raw text chunks directly as segment events without parsing.
+    Used when an agent has no tools configured.
+    """
+    
+    def __init__(
+        self,
+        on_segment_event: Optional[Callable[[SegmentEvent], None]] = None,
+        on_tool_invocation: Optional[Callable[[ToolInvocation], None]] = None,
+        segment_id_prefix: Optional[str] = None,
+    ):
+        """
+        Initialize the pass-through handler.
+        
+        Args:
+            on_segment_event: Callback for UI streaming.
+            on_tool_invocation: Callback for tools (unused here, kept for compatibility).
+            segment_id_prefix: Prefix for the single text segment ID.
+        """
+        self._on_segment_event = on_segment_event
+        self._segment_id_prefix = segment_id_prefix or f"pt_{uuid.uuid4().hex}:"
+        self._segment_id = f"{self._segment_id_prefix}text_0"
+        self._is_active = False
+        self._is_finalized = False
+        self._all_events: List[SegmentEvent] = []
+
+    def feed(self, chunk: ChunkResponse) -> List[SegmentEvent]:
+        if self._is_finalized:
+            raise RuntimeError("Handler has been finalized, cannot feed more chunks.")
+        
+        # Extract text content from ChunkResponse
+        text_content = chunk.content if isinstance(chunk, ChunkResponse) else chunk
+        if not text_content:
+            return []
+        
+        events = []
+        
+        # Start segment if not active
+        if not self._is_active:
+            self._is_active = True
+            start_event = SegmentEvent.start(
+                segment_id=self._segment_id,
+                segment_type=SegmentType.TEXT
+            )
+            events.append(start_event)
+        
+        # Content event
+        content_event = SegmentEvent.content(
+            segment_id=self._segment_id,
+            delta=text_content
+        )
+        events.append(content_event)
+        
+        self._process_events(events)
+        return events
+
+    def finalize(self) -> List[SegmentEvent]:
+        if self._is_finalized:
+            return []
+        
+        self._is_finalized = True
+        events = []
+        
+        if self._is_active:
+            end_event = SegmentEvent.end(segment_id=self._segment_id)
+            events.append(end_event)
+            self._is_active = False
+            
+        self._process_events(events)
+        return events
+
+    def _process_events(self, events: List[SegmentEvent]) -> None:
+        for event in events:
+            self._all_events.append(event)
+            if self._on_segment_event:
+                try:
+                    self._on_segment_event(event)
+                except Exception as e:
+                    logger.error(f"Error in on_segment_event callback: {e}")
+
+    def get_all_events(self) -> List[SegmentEvent]:
+        return self._all_events.copy()
+
+    def get_all_invocations(self) -> List[ToolInvocation]:
+        return []
+
+    def reset(self) -> None:
+        self._segment_id = f"{self._segment_id_prefix}text_{uuid.uuid4().hex}"
+        self._is_active = False
+        self._is_finalized = False
+        self._all_events.clear()

autobyteus/agent/streaming/handlers/streaming_handler_factory.py

@@ -0,0 +1,177 @@
+# file: autobyteus/autobyteus/agent/streaming/handlers/streaming_handler_factory.py
+"""
+Factory for selecting the appropriate StreamingResponseHandler implementation.
+
+This factory encapsulates all configuration logic including:
+- Format/mode resolution from environment
+- ParserConfig construction for text parsing modes
+- JSON profile selection for provider-aware parsing
+- Tool schema building for API tool call mode
+"""
+from __future__ import annotations
+
+import logging
+import uuid
+from dataclasses import dataclass, field
+from typing import Optional, Callable, List, Dict
+
+from autobyteus.agent.tool_invocation import ToolInvocation
+from autobyteus.llm.providers import LLMProvider
+from autobyteus.utils.tool_call_format import resolve_tool_call_format
+
+from .streaming_response_handler import StreamingResponseHandler
+from .parsing_streaming_response_handler import ParsingStreamingResponseHandler
+from .pass_through_streaming_response_handler import PassThroughStreamingResponseHandler
+from .api_tool_call_streaming_response_handler import ApiToolCallStreamingResponseHandler
+from ..parser.parser_context import ParserConfig
+from ..parser.json_parsing_strategies.registry import get_json_tool_parsing_profile
+from ..segments.segment_events import SegmentEvent
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class StreamingHandlerResult:
+    """
+    Result of creating a streaming handler.
+    
+    Attributes:
+        handler: The configured streaming response handler.
+        tool_schemas: Pre-built tool schemas for API mode (None for other modes).
+    """
+    handler: StreamingResponseHandler
+    tool_schemas: Optional[List[Dict]] = None
+
+
+class StreamingResponseHandlerFactory:
+    """Factory for building streaming response handlers based on minimal inputs."""
+
+    @staticmethod
+    def create(
+        *,
+        tool_names: List[str],
+        provider: Optional[LLMProvider],
+        segment_id_prefix: Optional[str] = None,
+        on_segment_event: Optional[Callable[[SegmentEvent], None]] = None,
+        on_tool_invocation: Optional[Callable[[ToolInvocation], None]] = None,
+        agent_id: Optional[str] = None,
+    ) -> StreamingHandlerResult:
+        """
+        Create a streaming response handler with all necessary configuration.
+        
+        Args:
+            tool_names: List of tool names the agent has access to.
+            provider: The LLM provider being used (for provider-specific parsing).
+            segment_id_prefix: Optional prefix for segment IDs.
+            on_segment_event: Callback for UI streaming events.
+            on_tool_invocation: Callback when a tool invocation is created.
+            agent_id: Agent identifier for logging.
+            
+        Returns:
+            StreamingHandlerResult containing the handler and optional tool schemas.
+        """
+        # Resolve format from environment
+        format_override = resolve_tool_call_format()
+        parse_tool_calls = bool(tool_names)
+        
+        # Generate segment ID prefix if not provided
+        if segment_id_prefix is None:
+            segment_id_prefix = f"turn_{uuid.uuid4().hex}:"
+        
+        # No tools → PassThrough handler
+        if not parse_tool_calls:
+            logger.debug(
+                "Agent '%s': No tools enabled - using PassThroughStreamingResponseHandler",
+                agent_id or "unknown",
+            )
+            return StreamingHandlerResult(
+                handler=PassThroughStreamingResponseHandler(
+                    on_segment_event=on_segment_event,
+                    on_tool_invocation=on_tool_invocation,
+                    segment_id_prefix=segment_id_prefix,
+                ),
+                tool_schemas=None,
+            )
+        
+        # API tool call mode → ApiToolCall handler + build schemas
+        if format_override == "api_tool_call":
+            logger.debug(
+                "Agent '%s': Using ApiToolCallStreamingResponseHandler",
+                agent_id or "unknown",
+            )
+            # Build tool schemas for API mode
+            tool_schemas = StreamingResponseHandlerFactory._build_tool_schemas(
+                tool_names=tool_names,
+                provider=provider,
+            )
+            return StreamingHandlerResult(
+                handler=ApiToolCallStreamingResponseHandler(
+                    on_segment_event=on_segment_event,
+                    on_tool_invocation=on_tool_invocation,
+                    segment_id_prefix=segment_id_prefix,
+                ),
+                tool_schemas=tool_schemas,
+            )
+        
+        # Text parsing mode (XML/JSON/Sentinel) → Parsing handler
+        parser_name = StreamingResponseHandlerFactory._resolve_parser_name(
+            format_override=format_override,
+            provider=provider,
+        )
+        
+        # Build ParserConfig with provider-aware JSON profile
+        json_profile = get_json_tool_parsing_profile(provider)
+        parser_config = ParserConfig(
+            parse_tool_calls=parse_tool_calls,
+            json_tool_patterns=json_profile.signature_patterns,
+            json_tool_parser=json_profile.parser,
+            segment_id_prefix=segment_id_prefix,
+        )
+        
+        logger.debug(
+            "Agent '%s': Using ParsingStreamingResponseHandler with %s parser",
+            agent_id or "unknown",
+            parser_name,
+        )
+        return StreamingHandlerResult(
+            handler=ParsingStreamingResponseHandler(
+                on_segment_event=on_segment_event,
+                on_tool_invocation=on_tool_invocation,
+                config=parser_config,
+                parser_name=parser_name,
+            ),
+            tool_schemas=None,
+        )
+
+    @staticmethod
+    def _resolve_parser_name(
+        *,
+        format_override: Optional[str],
+        provider: Optional[LLMProvider],
+    ) -> str:
+        """Resolve which parser to use based on format and provider."""
+        if format_override in {"xml", "json", "sentinel"}:
+            return format_override
+        # Default: XML for Anthropic, JSON for others
+        return "xml" if provider == LLMProvider.ANTHROPIC else "json"
+
+    @staticmethod
+    def _build_tool_schemas(
+        tool_names: List[str],
+        provider: Optional[LLMProvider],
+    ) -> Optional[List[Dict]]:
+        """Build tool schemas for API tool call mode."""
+        if not tool_names:
+            return None
+        
+        # Import here to avoid circular dependency
+        from autobyteus.tools.usage.tool_schema_provider import ToolSchemaProvider
+        
+        schemas = ToolSchemaProvider().build_schema(tool_names, provider)
+        if schemas:
+            logger.debug(
+                "Built %d tool schemas for API tool calls (provider: %s)",
+                len(schemas),
+                provider,
+            )
+        return schemas if schemas else None