goose/clients/python/goose_client/streaming.py

"""SSE streaming support - essential feature"""

import json
import logging
from typing import Any, AsyncIterator, Dict, Iterator, List, Optional

import httpx
from httpx_sse import aconnect_sse, connect_sse

logger = logging.getLogger(__name__)


class StreamingMixin:
    """Streaming support for Goose client"""

    def stream_reply(
        self, session_id: str, messages: List[Dict[str, Any]], timeout: float = 60.0
    ) -> Iterator[Dict[str, Any]]:
        """Stream responses from /reply endpoint.

        Args:
            session_id: The session ID to send messages to
            messages: List of message dictionaries
            timeout: Request timeout in seconds

        Yields:
            Parsed SSE event data as dictionaries
        """
        request_body = {"session_id": session_id, "messages": messages}

        logger.debug(f"Starting SSE stream for session {session_id}")

        try:
            with httpx.Client(timeout=timeout) as client:
                with connect_sse(
                    client,
                    "POST",
                    f"{self.base_url}/reply",
                    headers={"X-Secret-Key": self.api_key},
                    json=request_body,
                ) as event_source:
                    for sse in event_source.iter_sse():
                        if sse.data:
                            try:
                                event_data = json.loads(sse.data)
                                logger.debug(
                                    f"Received SSE event: {event_data.get('type', 'unknown')}"
                                )
                                yield event_data
                            except json.JSONDecodeError as e:
                                logger.warning(f"Failed to parse SSE data: {e}")
                                continue
        except httpx.TimeoutException:
            logger.error(f"Stream timeout after {timeout} seconds")
            yield {"type": "Error", "error": f"Request timeout after {timeout} seconds"}
        except httpx.HTTPError as e:
            logger.error(f"HTTP error during streaming: {e}")
            yield {"type": "Error", "error": str(e)}
        except Exception as e:
            logger.error(f"Unexpected error during streaming: {e}")
            yield {"type": "Error", "error": str(e)}

    async def astream_reply(
        self, session_id: str, messages: List[Dict[str, Any]], timeout: float = 60.0
    ) -> AsyncIterator[Dict[str, Any]]:
        """Async streaming from /reply endpoint.

        Args:
            session_id: The session ID to send messages to
            messages: List of message dictionaries
            timeout: Request timeout in seconds

        Yields:
            Parsed SSE event data as dictionaries
        """
        request_body = {"session_id": session_id, "messages": messages}

        logger.debug(f"Starting async SSE stream for session {session_id}")

        try:
            async with httpx.AsyncClient(timeout=timeout) as client:
                async with aconnect_sse(
                    client,
                    "POST",
                    f"{self.base_url}/reply",
                    headers={"X-Secret-Key": self.api_key},
                    json=request_body,
                ) as event_source:
                    async for sse in event_source.aiter_sse():
                        if sse.data:
                            try:
                                event_data = json.loads(sse.data)
                                logger.debug(
                                    f"Received async SSE event: {event_data.get('type', 'unknown')}"
                                )
                                yield event_data
                            except json.JSONDecodeError as e:
                                logger.warning(f"Failed to parse SSE data: {e}")
                                continue
        except httpx.TimeoutException:
            logger.error(f"Async stream timeout after {timeout} seconds")
            yield {"type": "Error", "error": f"Request timeout after {timeout} seconds"}
        except httpx.HTTPError as e:
            logger.error(f"Async HTTP error during streaming: {e}")
            yield {"type": "Error", "error": str(e)}
        except Exception as e:
            logger.error(f"Unexpected async error during streaming: {e}")
            yield {"type": "Error", "error": str(e)}

    def parse_stream_content(self, event: Dict[str, Any]) -> Optional[str]:
        """Helper to extract text content from stream events.

        Args:
            event: SSE event dictionary

        Returns:
            Extracted text content or None
        """
        if event.get("type") == "Message":
            message = event.get("message", {})
            text_parts = []
            for content in message.get("content", []):
                if content.get("type") == "text":
                    text_parts.append(content.get("text", ""))
            return "".join(text_parts) if text_parts else None
        return None

    def has_tool_confirmation(self, event: Dict[str, Any]) -> bool:
        """Check if event contains a tool confirmation request.

        Args:
            event: SSE event dictionary

        Returns:
            True if event has tool confirmation request
        """
        if event.get("type") == "Message":
            message = event.get("message", {})
            for content in message.get("content", []):
                if content.get("type") == "toolConfirmationRequest":
                    return True
        return False

    def get_tool_confirmations(self, event: Dict[str, Any]) -> List[Dict[str, Any]]:
        """Extract tool confirmation requests from event.

        Args:
            event: SSE event dictionary

        Returns:
            List of tool confirmation request dictionaries
        """
        confirmations = []
        if event.get("type") == "Message":
            message = event.get("message", {})
            for content in message.get("content", []):
                if content.get("type") == "toolConfirmationRequest":
                    confirmations.append(content)
        return confirmations