spaik-sdk 0.6.2__py3-none-any.whl

spaik_sdk/llm/message_handler.py

@@ -0,0 +1,188 @@
+import time
+import uuid
+from typing import Any, AsyncGenerator, Dict, List, Optional
+
+from spaik_sdk.attachments.models import Attachment
+from spaik_sdk.llm.streaming.streaming_event_handler import EventType, StreamingEventHandler
+from spaik_sdk.recording.base_recorder import BaseRecorder
+from spaik_sdk.thread.models import MessageBlock, MessageBlockType, ThreadMessage
+from spaik_sdk.thread.thread_container import ThreadContainer
+from spaik_sdk.utils.init_logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class MessageHandler:
+    """Manages conversation message history using ThreadContainer."""
+
+    def __init__(
+        self,
+        thread_container: ThreadContainer,
+        assistant_name: str,
+        assistant_id: str,
+        recorder: Optional[BaseRecorder] = None,
+    ):
+        self.thread_container = thread_container
+        self.streaming_handler = StreamingEventHandler(recorder)
+        self.assistant_name = assistant_name
+        self.assistant_id = assistant_id
+        self._update_previous_message_count()
+
+    def _update_previous_message_count(self) -> None:
+        self._previous_message_count = self.thread_container.get_nof_messages_including_system() + 1
+
+    def add_user_message(
+        self,
+        user_input: str,
+        author_id: str,
+        author_name: str,
+        attachments: Optional[List[Attachment]] = None,
+    ) -> None:
+        """Add a user message to both thread container and LangChain messages."""
+
+        # Add to thread container
+        block_id = str(uuid.uuid4())
+        user_message = ThreadMessage(
+            id=str(uuid.uuid4()),
+            ai=False,
+            author_id=author_id,
+            author_name=author_name,
+            timestamp=int(time.time() * 1000),
+            blocks=[
+                MessageBlock(
+                    id=block_id,
+                    streaming=False,
+                    type=MessageBlockType.PLAIN,
+                    content=user_input,
+                )
+            ],
+            attachments=attachments,
+        )
+        self.thread_container.add_message(user_message)
+
+        # Add to LangChain messages for compatibility
+        self._update_previous_message_count()
+
+    def add_error(self, error_text: str, author_id: str = "system") -> str:
+        """Add an error message to the thread container and return message ID"""
+        return self.thread_container.add_error_message(error_text, author_id)
+
+    def handle_cancellation(self) -> None:
+        """Handle cancellation of the agent."""
+        self.thread_container.cancel_generation()
+
+    async def process_agent_token_stream(
+        self,
+        agent_stream,
+    ) -> AsyncGenerator[Dict[str, Any], None]:
+        """Process agent event stream for individual tokens and yield them in real-time."""
+        async for streaming_event in self.streaming_handler.process_stream(agent_stream):
+            if streaming_event.event_type == EventType.MESSAGE_START:
+                logger.debug(f"🔍 Processing MESSAGE_START for message: {streaming_event.message_id}")
+                # Create AI message in thread container
+                assert streaming_event.message_id is not None
+                ai_message = ThreadMessage(
+                    id=streaming_event.message_id,
+                    ai=True,
+                    author_id=self.assistant_id,
+                    author_name=self.assistant_name,
+                    timestamp=int(time.time() * 1000),
+                    blocks=[],
+                )
+                self.thread_container.add_message(ai_message)
+
+            elif streaming_event.event_type == EventType.BLOCK_START:
+                logger.debug(f"🔍 Processing BLOCK_START for block: {streaming_event.block_id}")
+                # Add new block to the message
+                assert streaming_event.block_id is not None
+                assert streaming_event.block_type is not None
+                assert streaming_event.message_id is not None
+                new_block = MessageBlock(
+                    id=streaming_event.block_id,
+                    streaming=True,
+                    type=streaming_event.block_type,
+                    tool_call_id=streaming_event.tool_call_id,
+                    tool_call_args=streaming_event.tool_args,
+                    tool_name=streaming_event.tool_name,
+                )
+                self.thread_container.add_message_block(streaming_event.message_id, new_block)
+
+            elif streaming_event.event_type == EventType.BLOCK_END:
+                logger.debug(f"🔚 Processing BLOCK_END for block: {streaming_event.block_id}")
+                # Mark individual block as non-streaming (completed)
+                assert streaming_event.block_id is not None
+                assert streaming_event.message_id is not None
+                logger.debug(f"🔚 Processing BLOCK_END for block: {streaming_event.block_id}")
+                self.thread_container.finalize_streaming_blocks(streaming_event.message_id, [streaming_event.block_id])
+                logger.debug(f"✅ Block {streaming_event.block_id} marked as completed")
+
+            elif streaming_event.event_type in [EventType.REASONING, EventType.REASONING_SUMMARY, EventType.TOKEN]:
+                logger.debug(f"🔍 Processing {streaming_event.event_type.value} for block: {streaming_event.block_id}")
+                # Add streaming content
+                assert streaming_event.block_id is not None
+                assert streaming_event.content is not None
+                self.thread_container.add_streaming_message_chunk(streaming_event.block_id, streaming_event.content)
+
+                # Yield for external consumption
+                yield {
+                    "type": streaming_event.event_type.value,
+                    "content": streaming_event.content,
+                    "block_id": streaming_event.block_id,
+                    "message_id": streaming_event.message_id,
+                }
+
+            elif streaming_event.event_type == EventType.USAGE_METADATA:
+                logger.info(f"📊 Processing USAGE_METADATA for message: {streaming_event.message_id}")
+                # Store consumption metadata in ThreadContainer
+                if streaming_event.message_id and streaming_event.usage_metadata:
+                    self.thread_container.add_consumption_metadata(streaming_event.message_id, streaming_event.usage_metadata)
+                # Yield usage metadata for external consumption
+                yield {
+                    "type": "usage_metadata",
+                    "message_id": streaming_event.message_id,
+                    "usage_metadata": streaming_event.usage_metadata,
+                }
+
+            elif streaming_event.event_type == EventType.COMPLETE:
+                logger.debug(f"🔍 Processing COMPLETE for message: {streaming_event.message_id}")
+                # Mark all blocks as non-streaming
+                if streaming_event.message_id and streaming_event.blocks:
+                    self.thread_container.finalize_streaming_blocks(streaming_event.message_id, streaming_event.blocks)
+
+                yield {
+                    "type": "complete",
+                    "message": streaming_event.message,
+                    "blocks": streaming_event.blocks,
+                    "message_id": streaming_event.message_id,
+                }
+
+            elif streaming_event.event_type == EventType.TOOL_USE:
+                logger.debug(f"🔍 Processing TOOL_USE for block: {streaming_event.block_id}")
+                # Handle tool use event - the block is already created, just yield for external consumption
+                yield {
+                    "type": "tool_use",
+                    "tool_call_id": streaming_event.tool_call_id,
+                    "tool_name": streaming_event.tool_name,
+                    "tool_args": streaming_event.tool_args,
+                    "block_id": streaming_event.block_id,
+                    "message_id": streaming_event.message_id,
+                }
+
+            elif streaming_event.event_type == EventType.TOOL_RESPONSE:
+                logger.debug(f"🔍 Processing TOOL_RESPONSE for block: {streaming_event.block_id}")
+                # Handle tool response - update the thread container with the response
+                assert streaming_event.tool_call_id is not None
+                assert streaming_event.content is not None
+                self.thread_container.update_tool_use_block_with_response(
+                    streaming_event.tool_call_id, streaming_event.content, streaming_event.error
+                )
+
+                # Yield for external consumption
+                yield {
+                    "type": "tool_response",
+                    "tool_call_id": streaming_event.tool_call_id,
+                    "response": streaming_event.content,
+                    "error": streaming_event.error,
+                    "block_id": streaming_event.block_id,
+                    "message_id": streaming_event.message_id,
+                }

spaik_sdk/llm/streaming/__init__.py

@@ -0,0 +1 @@
+

spaik_sdk/llm/streaming/block_manager.py

@@ -0,0 +1,152 @@
+import time
+import uuid
+from typing import Any, AsyncGenerator, Dict, List, Optional
+
+from spaik_sdk.llm.streaming.models import EventType, StreamingEvent
+from spaik_sdk.thread.models import MessageBlockType
+
+
+class BlockManager:
+    """Manages different types of content blocks during streaming."""
+
+    def __init__(self):
+        self.current_blocks: Dict[str, str] = {}  # block_id -> block_type
+        self.block_timestamps: Dict[str, float] = {}  # block_id -> creation_timestamp
+        self.reasoning_block_id: Optional[str] = None
+        self.regular_block_id: Optional[str] = None
+        self.summary_block_id: Optional[str] = None
+        self.tool_use_blocks: Dict[str, str] = {}  # tool_call_id -> block_id
+        self.reasoning_detected = False
+        self.last_block_type: Optional[str] = None  # Track the last type of block created
+
+    def reset(self):
+        """Reset block manager state."""
+        self.current_blocks = {}
+        self.block_timestamps = {}
+        self.reasoning_block_id = None
+        self.regular_block_id = None
+        self.summary_block_id = None
+        self.tool_use_blocks = {}
+        self.reasoning_detected = False
+        self.last_block_type = None
+
+    def mark_reasoning_detected(self):
+        """Mark that reasoning activity has been detected."""
+        self.reasoning_detected = True
+
+    def reset_reasoning_block(self):
+        """Reset reasoning block for mid-response thinking (creates new reasoning block)."""
+        if self.reasoning_block_id:
+            # Remove the current reasoning block from tracking
+            self.current_blocks.pop(self.reasoning_block_id, None)
+            self.block_timestamps.pop(self.reasoning_block_id, None)  # Remove timestamp too
+        self.reasoning_block_id = None
+
+    def get_block_ids(self) -> List[str]:
+        """Get list of all current block IDs."""
+        return list(self.current_blocks.keys())
+
+    def should_create_new_reasoning_block(self) -> bool:
+        """Check if we need a new reasoning block based on timestamps.
+
+        Rule: If there's any tool call newer than the current reasoning block, create new reasoning block.
+        This ensures reasoning gets properly segmented around tool calls.
+        """
+        if self.reasoning_block_id is None:
+            return False
+
+        current_reasoning_timestamp = self.block_timestamps.get(self.reasoning_block_id, 0)
+
+        # Check if any tool block is newer than our current reasoning block
+        for tool_block_id in self.tool_use_blocks.values():
+            tool_timestamp = self.block_timestamps.get(tool_block_id, 0)
+            if tool_timestamp > current_reasoning_timestamp:
+                return True
+
+        return False
+
+    async def ensure_tool_use_block(
+        self, message_id: str, tool_call_id: str, tool_name: str, tool_args: Dict[str, Any]
+    ) -> AsyncGenerator[StreamingEvent, None]:
+        """Ensure tool use block exists for the given tool call, create if needed."""
+        if tool_call_id not in self.tool_use_blocks:
+            block_id = f"tool_{uuid.uuid4()}"
+            self.tool_use_blocks[tool_call_id] = block_id
+            self.current_blocks[block_id] = "tool_use"
+            self.block_timestamps[block_id] = time.time()
+            self.last_block_type = "tool_use"  # Track that we created a tool block
+
+            yield StreamingEvent(
+                event_type=EventType.BLOCK_START,
+                block_id=block_id,
+                block_type=MessageBlockType.TOOL_USE,
+                message_id=message_id,
+                tool_call_id=tool_call_id,
+                tool_name=tool_name,
+                tool_args=tool_args,
+            )
+
+    def get_tool_use_block_id(self, tool_call_id: str) -> Optional[str]:
+        """Get the block ID for a specific tool call."""
+        return self.tool_use_blocks.get(tool_call_id)
+
+    async def ensure_reasoning_block(self, message_id: str) -> AsyncGenerator[StreamingEvent, None]:
+        """Ensure reasoning block exists, create if needed."""
+        if self.reasoning_block_id is None:
+            self.reasoning_block_id = f"reasoning_{uuid.uuid4()}"
+            self.current_blocks[self.reasoning_block_id] = "reasoning"
+            self.block_timestamps[self.reasoning_block_id] = time.time()
+            self.last_block_type = "reasoning"  # Track that we created a reasoning block
+
+            yield StreamingEvent(
+                event_type=EventType.BLOCK_START,
+                block_id=self.reasoning_block_id,
+                block_type=MessageBlockType.REASONING,
+                message_id=message_id,
+            )
+
+    async def ensure_regular_block(self, message_id: str) -> AsyncGenerator[StreamingEvent, None]:
+        """Ensure regular content block exists, create if needed."""
+        # Create a new regular block if:
+        # 1. No regular block exists yet, OR
+        # 2. The last block created was not a regular block (meaning there was an interruption)
+        should_create_new_block = self.regular_block_id is None or self.last_block_type not in [None, "plain"]
+
+        if should_create_new_block:
+            self.regular_block_id = f"plain_{uuid.uuid4()}"
+            self.current_blocks[self.regular_block_id] = "plain"
+            self.block_timestamps[self.regular_block_id] = time.time()
+            self.last_block_type = "plain"  # Track that we created a regular block
+
+            yield StreamingEvent(
+                event_type=EventType.BLOCK_START, block_id=self.regular_block_id, block_type=MessageBlockType.PLAIN, message_id=message_id
+            )
+
+    async def ensure_summary_block(self, message_id: str) -> AsyncGenerator[StreamingEvent, None]:
+        """Ensure summary block exists, create if needed."""
+        if self.summary_block_id is None:
+            self.summary_block_id = f"summary_{uuid.uuid4()}"
+            self.current_blocks[self.summary_block_id] = "summary"
+            self.block_timestamps[self.summary_block_id] = time.time()
+            self.last_block_type = "summary"  # Track that we created a summary block
+
+        # Note: We don't yield BLOCK_START for summary blocks as they're handled differently
+        # This is an async generator so we need at least one yield or return to make it work
+        return
+        yield  # This line will never be reached but satisfies the type checker
+
+    def get_reasoning_block_id(self) -> Optional[str]:
+        """Get the reasoning block ID."""
+        return self.reasoning_block_id
+
+    def get_regular_block_id(self) -> Optional[str]:
+        """Get the regular content block ID."""
+        return self.regular_block_id
+
+    def get_summary_block_id(self) -> Optional[str]:
+        """Get the summary block ID."""
+        return self.summary_block_id
+
+    def has_reasoning_activity(self) -> bool:
+        """Check if reasoning activity has been detected."""
+        return self.reasoning_detected

spaik_sdk/llm/streaming/models.py

@@ -0,0 +1,42 @@
+from dataclasses import dataclass
+from enum import Enum
+from typing import TYPE_CHECKING, Any, Dict, List, Optional
+
+from spaik_sdk.thread.models import MessageBlockType
+
+if TYPE_CHECKING:
+    from spaik_sdk.llm.consumption.token_usage import TokenUsage
+
+
+class EventType(str, Enum):
+    """Types of streaming events."""
+
+    MESSAGE_START = "message_start"
+    BLOCK_START = "block_start"
+    BLOCK_END = "block_end"
+    REASONING = "reasoning"
+    REASONING_SUMMARY = "reasoning_summary"
+    TOKEN = "token"
+    TOOL_USE = "tool_use"
+    TOOL_RESPONSE = "tool_response"
+    USAGE_METADATA = "usage_metadata"
+    COMPLETE = "complete"
+    ERROR = "error"
+
+
+@dataclass
+class StreamingEvent:
+    """Represents a processed streaming event."""
+
+    event_type: EventType
+    content: Optional[str] = None
+    block_id: Optional[str] = None
+    message_id: Optional[str] = None
+    block_type: Optional[MessageBlockType] = None
+    blocks: Optional[List[str]] = None
+    message: Optional[Any] = None
+    error: Optional[str] = None
+    tool_call_id: Optional[str] = None
+    tool_name: Optional[str] = None
+    tool_args: Optional[Dict[str, Any]] = None
+    usage_metadata: Optional["TokenUsage"] = None

spaik_sdk/llm/streaming/streaming_content_handler.py

@@ -0,0 +1,157 @@
+import uuid
+from typing import Any, AsyncGenerator, Dict, Optional
+
+from spaik_sdk.llm.streaming.block_manager import BlockManager
+from spaik_sdk.llm.streaming.models import EventType, StreamingEvent
+from spaik_sdk.llm.streaming.streaming_state_manager import StreamingStateManager
+from spaik_sdk.utils.init_logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class StreamingContentHandler:
+    """Handles processing of reasoning and regular content during streaming."""
+
+    def __init__(self, block_manager: BlockManager, state_manager: StreamingStateManager):
+        self.block_manager = block_manager
+        self.state_manager = state_manager
+
+    async def handle_reasoning_content(self, reasoning_content: str) -> AsyncGenerator[StreamingEvent, None]:
+        """Handle reasoning content and yield events."""
+        async for event in self._ensure_streaming_started():
+            yield event
+
+        # Type guard to ensure message_id is not None
+        if self.state_manager.current_message_id is None:
+            return
+
+        # Check if we need a new reasoning block based on timestamps
+        if self.block_manager.should_create_new_reasoning_block():
+            # Reset current reasoning block to force creation of a new one
+            self.block_manager.reset_reasoning_block()
+
+        # Check if this is creating a new reasoning block
+        creating_new_block = self.block_manager.get_reasoning_block_id() is None
+
+        # Ensure reasoning block exists
+        async for streaming_event in self.block_manager.ensure_reasoning_block(self.state_manager.current_message_id):
+            yield streaming_event
+
+        # Track that we've created a reasoning block
+        if creating_new_block:
+            self.state_manager.increment_reasoning_blocks()
+
+        # For Google models with thinking_budget, we might not have actual reasoning content
+        # but we still want to show that reasoning is happening
+        if not reasoning_content and self.block_manager.has_reasoning_activity():
+            reasoning_content = "[Thinking process active - reasoning tokens being used internally]"
+
+        # Always yield reasoning content (even if empty) to ensure thread container tracks it
+        yield StreamingEvent(
+            event_type=EventType.REASONING,
+            content=reasoning_content,  # This could be empty string or placeholder text
+            block_id=self.block_manager.get_reasoning_block_id(),
+            message_id=self.state_manager.current_message_id,
+        )
+
+    async def handle_regular_content(self, regular_content: str) -> AsyncGenerator[StreamingEvent, None]:
+        """Handle regular content and yield events."""
+        async for event in self._ensure_streaming_started():
+            yield event
+
+        # Type guard to ensure message_id is not None
+        if self.state_manager.current_message_id is None:
+            return
+
+        # Ensure regular block exists
+        async for streaming_event in self.block_manager.ensure_regular_block(self.state_manager.current_message_id):
+            yield streaming_event
+
+        yield StreamingEvent(
+            event_type=EventType.TOKEN,
+            content=regular_content,
+            block_id=self.block_manager.get_regular_block_id(),
+            message_id=self.state_manager.current_message_id,
+        )
+
+    async def end_thinking_session_if_needed(self) -> AsyncGenerator[StreamingEvent, None]:
+        """End thinking session and emit BLOCK_END event if needed."""
+        if self.state_manager.in_thinking_session:
+            self.state_manager.end_thinking_session()
+            # Emit BLOCK_END event for the current reasoning block
+            current_reasoning_block_id = self.block_manager.get_reasoning_block_id()
+            if current_reasoning_block_id and self.state_manager.current_message_id:
+                logger.debug(f"🔚 Emitting BLOCK_END for reasoning block: {current_reasoning_block_id}")
+                yield StreamingEvent(
+                    event_type=EventType.BLOCK_END, block_id=current_reasoning_block_id, message_id=self.state_manager.current_message_id
+                )
+            else:
+                logger.debug(
+                    "❌ Cannot emit BLOCK_END - block_id: %s, message_id: %s",
+                    current_reasoning_block_id,
+                    self.state_manager.current_message_id,
+                )
+
+    async def end_final_thinking_session_if_needed(self) -> AsyncGenerator[StreamingEvent, None]:
+        """End thinking session at stream end if still active."""
+        if self.state_manager.in_thinking_session:
+            current_reasoning_block_id = self.block_manager.get_reasoning_block_id()
+            if current_reasoning_block_id and self.state_manager.current_message_id:
+                logger.debug(f"🔚 Stream ending - emitting BLOCK_END for final reasoning block: {current_reasoning_block_id}")
+                yield StreamingEvent(
+                    event_type=EventType.BLOCK_END, block_id=current_reasoning_block_id, message_id=self.state_manager.current_message_id
+                )
+            self.state_manager.end_thinking_session()
+
+    async def _ensure_streaming_started(self) -> AsyncGenerator[StreamingEvent, None]:
+        """Ensure streaming has been initialized and yield MESSAGE_START if needed."""
+        if not self.state_manager.streaming_started:
+            self.state_manager.current_message_id = str(uuid.uuid4())
+            self.state_manager.streaming_started = True
+
+            yield StreamingEvent(event_type=EventType.MESSAGE_START, message_id=self.state_manager.current_message_id)
+
+    async def handle_tool_use(self, tool_call_id: str, tool_name: str, tool_args: Dict[str, Any]) -> AsyncGenerator[StreamingEvent, None]:
+        """Handle tool use and yield events."""
+        async for event in self._ensure_streaming_started():
+            yield event
+
+        # Type guard to ensure message_id is not None
+        if self.state_manager.current_message_id is None:
+            return
+
+        # Ensure tool use block exists
+        async for streaming_event in self.block_manager.ensure_tool_use_block(
+            self.state_manager.current_message_id, tool_call_id, tool_name, tool_args
+        ):
+            yield streaming_event
+
+        # Emit tool use event
+        yield StreamingEvent(
+            event_type=EventType.TOOL_USE,
+            block_id=self.block_manager.get_tool_use_block_id(tool_call_id),
+            message_id=self.state_manager.current_message_id,
+            tool_call_id=tool_call_id,
+            tool_name=tool_name,
+            tool_args=tool_args,
+        )
+
+    async def handle_tool_response(
+        self, tool_call_id: str, response: str, error: Optional[str] = None
+    ) -> AsyncGenerator[StreamingEvent, None]:
+        """Handle tool response and yield events."""
+        if self.state_manager.current_message_id is None:
+            return
+
+        # Get the tool use block for this tool call
+        block_id = self.block_manager.get_tool_use_block_id(tool_call_id)
+        if block_id:
+            # Emit tool response event
+            yield StreamingEvent(
+                event_type=EventType.TOOL_RESPONSE,
+                content=response,
+                block_id=block_id,
+                message_id=self.state_manager.current_message_id,
+                tool_call_id=tool_call_id,
+                error=error,
+            )