spaik-sdk 0.6.2__py3-none-any.whl

spaik_sdk/llm/streaming/streaming_event_handler.py

@@ -0,0 +1,215 @@
+from typing import AsyncGenerator, Optional, Union
+
+from langchain_core.messages import AIMessage, AIMessageChunk
+
+from spaik_sdk.llm.consumption.token_usage import TokenUsage
+from spaik_sdk.llm.streaming.block_manager import BlockManager
+from spaik_sdk.llm.streaming.models import EventType, StreamingEvent
+from spaik_sdk.llm.streaming.streaming_content_handler import StreamingContentHandler
+from spaik_sdk.llm.streaming.streaming_state_manager import StreamingStateManager
+from spaik_sdk.recording.base_recorder import BaseRecorder
+from spaik_sdk.utils.init_logger import init_logger
+
+logger = init_logger(__name__)
+
+AIMessageType = Union[AIMessage, AIMessageChunk]
+
+
+class StreamingEventHandler:
+    """Handles LangChain 1.x streaming events."""
+
+    def __init__(self, recorder: Optional[BaseRecorder] = None):
+        self.recorder = recorder
+        self.block_manager = BlockManager()
+        self.state_manager = StreamingStateManager()
+        self.content_handler = StreamingContentHandler(self.block_manager, self.state_manager)
+        self._processed_message_ids: set[str] = set()
+        self._final_message: Optional[AIMessageType] = None
+        self._got_chat_model_stream: bool = False
+
+    def reset(self) -> None:
+        self.block_manager.reset()
+        self.state_manager.reset()
+        self._processed_message_ids.clear()
+        self._final_message = None
+        self._got_chat_model_stream = False
+
+    async def process_stream(self, agent_stream) -> AsyncGenerator[StreamingEvent, None]:
+        """Process LangChain 1.x agent stream events."""
+        self.reset()
+
+        async for event in agent_stream:
+            if self.recorder is not None:
+                self.recorder.record_token(event)
+
+            event_type = event.get("event", "")
+            data = event.get("data", {})
+            logger.trace(f"Stream event: {event_type}")
+
+            # on_chat_model_stream - real-time token streaming (preferred)
+            if event_type == "on_chat_model_stream":
+                self._got_chat_model_stream = True
+                chunk = data.get("chunk")
+                if isinstance(chunk, AIMessageChunk):
+                    async for streaming_event in self._handle_chunk(chunk):
+                        yield streaming_event
+
+            # on_chain_stream - complete messages (fallback if no chat_model_stream)
+            elif event_type == "on_chain_stream":
+                if not self._got_chat_model_stream:
+                    ai_message = self._extract_ai_message(data.get("chunk", {}))
+                    if ai_message and not self._is_duplicate(ai_message):
+                        async for streaming_event in self._handle_message(ai_message):
+                            yield streaming_event
+                        self._final_message = ai_message
+
+            # on_chat_model_end - usage metadata from the model
+            elif event_type == "on_chat_model_end":
+                output = data.get("output")
+                if isinstance(output, (AIMessage, AIMessageChunk)):
+                    self._final_message = output
+                    async for streaming_event in self._emit_usage_if_available(output):
+                        yield streaming_event
+
+            # on_chain_end - final state
+            elif event_type == "on_chain_end":
+                output = data.get("output", {})
+                if isinstance(output, dict) and "messages" in output:
+                    for msg in output["messages"]:
+                        if isinstance(msg, (AIMessage, AIMessageChunk)):
+                            if self._final_message is None:
+                                self._final_message = msg
+                            async for streaming_event in self._emit_usage_if_available(msg):
+                                yield streaming_event
+                            break
+
+            # on_tool_end - tool execution completed
+            elif event_type == "on_tool_end":
+                output = data.get("output")
+                if output is not None:
+                    tool_call_id = getattr(output, "tool_call_id", None)
+                    content = getattr(output, "content", str(output))
+                    if tool_call_id:
+                        async for streaming_event in self.content_handler.handle_tool_response(
+                            tool_call_id, content if isinstance(content, str) else str(content)
+                        ):
+                            yield streaming_event
+
+        # End any active thinking session
+        async for event in self.content_handler.end_final_thinking_session_if_needed():
+            yield event
+
+        # Emit final COMPLETE event
+        if self._final_message or self.state_manager.current_message_id:
+            yield StreamingEvent(
+                event_type=EventType.COMPLETE,
+                message=self._final_message,
+                blocks=self.block_manager.get_block_ids(),
+                message_id=self.state_manager.current_message_id,
+            )
+
+    def _is_duplicate(self, message: AIMessageType) -> bool:
+        msg_id = getattr(message, "id", None)
+        if not msg_id:
+            return False
+        if msg_id in self._processed_message_ids:
+            return True
+        self._processed_message_ids.add(msg_id)
+        return False
+
+    def _extract_ai_message(self, chunk: dict) -> Optional[AIMessageType]:
+        if "messages" in chunk:
+            for msg in chunk["messages"]:
+                if isinstance(msg, (AIMessage, AIMessageChunk)):
+                    return msg
+        if "model" in chunk and isinstance(chunk["model"], dict):
+            if "messages" in chunk["model"]:
+                for msg in chunk["model"]["messages"]:
+                    if isinstance(msg, (AIMessage, AIMessageChunk)):
+                        return msg
+        return None
+
+    async def _handle_chunk(self, chunk: AIMessageChunk) -> AsyncGenerator[StreamingEvent, None]:
+        """Handle streaming chunk - real-time content."""
+        content = chunk.content
+
+        if isinstance(content, str) and content:
+            async for event in self.content_handler.handle_regular_content(content):
+                yield event
+            self.state_manager.mark_text_content_received()
+
+        elif isinstance(content, list):
+            for block in content:
+                if isinstance(block, dict):
+                    block_type = block.get("type")
+                    if block_type == "text":
+                        text = block.get("text", "")
+                        if text:
+                            async for event in self.content_handler.handle_regular_content(text):
+                                yield event
+                            self.state_manager.mark_text_content_received()
+                    elif block_type in ("reasoning", "thinking"):
+                        reasoning = block.get("reasoning", "") or block.get("thinking", "")
+                        if reasoning:
+                            async for event in self.content_handler.handle_reasoning_content(reasoning):
+                                yield event
+                elif isinstance(block, str) and block:
+                    async for event in self.content_handler.handle_regular_content(block):
+                        yield event
+                    self.state_manager.mark_text_content_received()
+
+        if hasattr(chunk, "tool_calls") and chunk.tool_calls:
+            for tool_call in chunk.tool_calls:
+                tool_id = tool_call.get("id") if isinstance(tool_call, dict) else getattr(tool_call, "id", None)
+                tool_name = tool_call.get("name") if isinstance(tool_call, dict) else getattr(tool_call, "name", None)
+                tool_args = tool_call.get("args", {}) if isinstance(tool_call, dict) else getattr(tool_call, "args", {})
+                if tool_id and tool_name:
+                    async for event in self.content_handler.handle_tool_use(tool_id, tool_name, tool_args):
+                        yield event
+
+    async def _handle_message(self, message: AIMessageType) -> AsyncGenerator[StreamingEvent, None]:
+        """Handle complete message (from on_chain_stream fallback)."""
+        content = message.content
+
+        if isinstance(content, str) and content:
+            async for event in self.content_handler.handle_regular_content(content):
+                yield event
+            self.state_manager.mark_text_content_received()
+
+        elif isinstance(content, list):
+            for block in content:
+                if isinstance(block, dict):
+                    block_type = block.get("type")
+                    if block_type == "text":
+                        async for event in self.content_handler.handle_regular_content(block.get("text", "")):
+                            yield event
+                        self.state_manager.mark_text_content_received()
+                    elif block_type in ("reasoning", "thinking"):
+                        reasoning = block.get("reasoning", "") or block.get("thinking", "")
+                        async for event in self.content_handler.handle_reasoning_content(reasoning):
+                            yield event
+                elif isinstance(block, str) and block:
+                    async for event in self.content_handler.handle_regular_content(block):
+                        yield event
+                    self.state_manager.mark_text_content_received()
+
+        if hasattr(message, "tool_calls") and message.tool_calls:
+            for tool_call in message.tool_calls:
+                tool_id = tool_call.get("id") if isinstance(tool_call, dict) else getattr(tool_call, "id", None)
+                tool_name = tool_call.get("name") if isinstance(tool_call, dict) else getattr(tool_call, "name", None)
+                tool_args = tool_call.get("args", {}) if isinstance(tool_call, dict) else getattr(tool_call, "args", {})
+                if tool_id and tool_name:
+                    async for event in self.content_handler.handle_tool_use(tool_id, tool_name, tool_args):
+                        yield event
+
+    async def _emit_usage_if_available(self, message: AIMessageType) -> AsyncGenerator[StreamingEvent, None]:
+        """Emit usage metadata if available on message."""
+        if hasattr(message, "usage_metadata") and message.usage_metadata:
+            yield StreamingEvent(
+                event_type=EventType.USAGE_METADATA,
+                message_id=self.state_manager.current_message_id,
+                usage_metadata=TokenUsage.from_langchain(message.usage_metadata),
+            )
+
+
+__all__ = ["StreamingEventHandler"]

spaik_sdk/llm/streaming/streaming_state_manager.py

@@ -0,0 +1,58 @@
+from typing import Optional
+
+from spaik_sdk.utils.init_logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class StreamingStateManager:
+    """Manages state for streaming operations."""
+
+    def __init__(self):
+        self.current_message_id: Optional[str] = None
+        self.streaming_started = False
+
+        # Track mid-response thinking state
+        self.last_block_type = None
+        self.has_text_content = False
+        self.reasoning_blocks_created = 0  # Track how many reasoning blocks we've created
+        self.in_thinking_session = False  # Track if we're currently in a thinking session
+
+    def reset(self):
+        """Reset all state for new stream."""
+        self.current_message_id = None
+        self.streaming_started = False
+        self.last_block_type = None
+        self.has_text_content = False
+        self.reasoning_blocks_created = 0
+        self.in_thinking_session = False
+
+    def start_thinking_session(self):
+        """Mark the start of a thinking session."""
+        if not self.in_thinking_session:
+            logger.debug("🧠 Starting thinking session")
+            self.in_thinking_session = True
+
+    def end_thinking_session(self):
+        """Mark the end of a thinking session."""
+        if self.in_thinking_session:
+            logger.debug("Ending thinking session - got text content")
+            self.in_thinking_session = False
+
+    def increment_reasoning_blocks(self):
+        """Increment the count of reasoning blocks created."""
+        self.reasoning_blocks_created += 1
+        logger.debug(f"Created reasoning block #{self.reasoning_blocks_created}")
+
+    def should_create_new_thinking_session(self, reasoning_content: bool, current_block_type: str) -> bool:
+        """Determine if we should create a new thinking session (mid-response thinking)."""
+        return reasoning_content and self.has_text_content and self.last_block_type == "text" and current_block_type == "thinking"
+
+    def update_block_type(self, block_type: str):
+        """Update the last block type."""
+        if block_type:
+            self.last_block_type = block_type
+
+    def mark_text_content_received(self):
+        """Mark that we've received text content."""
+        self.has_text_content = True

spaik_sdk/models/factories/anthropic_factory.py

@@ -0,0 +1,33 @@
+from typing import Any, Dict, Optional
+
+from spaik_sdk.models.factories.base_model_factory import BaseModelFactory
+from spaik_sdk.models.llm_config import LLMConfig
+from spaik_sdk.models.llm_families import LLMFamilies
+from spaik_sdk.models.llm_model import LLMModel
+from spaik_sdk.models.model_registry import ModelRegistry
+
+
+class AnthropicModelFactory(BaseModelFactory):
+    MODELS = ModelRegistry.get_by_family(LLMFamilies.ANTHROPIC)
+
+    def supports_model(self, model: LLMModel) -> bool:
+        return model in AnthropicModelFactory.MODELS
+
+    def get_cache_control(self, config: LLMConfig) -> Optional[Dict[str, Any]]:
+        return {"type": "ephemeral"}
+
+    def get_model_specific_config(self, config: LLMConfig) -> Dict[str, Any]:
+        allow_reasoning = config.reasoning and not config.structured_response
+        model_config: Dict[str, Any] = {
+            "model_name": config.model.name,
+            "streaming": config.streaming,
+            "max_tokens": config.max_output_tokens,
+        }
+
+        # Handle thinking mode via model_kwargs for LangChain compatibility
+        if allow_reasoning:
+            model_config["thinking"] = {"type": "enabled", "budget_tokens": config.reasoning_budget_tokens}
+        else:
+            model_config["temperature"] = config.temperature
+
+        return model_config

spaik_sdk/models/factories/base_model_factory.py

@@ -0,0 +1,71 @@
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any, Dict, Optional
+
+if TYPE_CHECKING:
+    from spaik_sdk.models.llm_config import LLMConfig
+
+from spaik_sdk.models.llm_model import LLMModel
+from spaik_sdk.models.llm_wrapper import LLMWrapper
+from spaik_sdk.models.providers.base_provider import BaseProvider
+
+
+class BaseModelFactory(ABC):
+    def create_model(self, config: "LLMConfig", provider: BaseProvider) -> LLMWrapper:
+        """Create a model wrapper for the given config and provider instance."""
+        # Check if this factory supports the model with this config
+        if not self.supports_model_config(config):
+            raise ValueError(f"Factory doesn't support model config: {config}")
+
+        # Get provider config and cache control
+        provider_config = provider.get_model_config(config)
+        cache_control = self.get_cache_control(config)
+
+        # Get model-specific configuration from subclass
+        model_specific_config = self.get_model_specific_config(config)
+
+        # Build complete model config
+        model_config = {**model_specific_config, **provider_config}
+
+        # Let provider create the langchain model
+        langchain_model = provider.create_langchain_model(config, model_config)
+
+        return LLMWrapper(langchain_model, cache_control, config.model)
+
+    @abstractmethod
+    def supports_model(self, model: LLMModel) -> bool:
+        """Check if this factory supports the given model (basic check)."""
+        pass
+
+    def supports_model_config(self, config: "LLMConfig") -> bool:
+        return self.supports_model(config.model)
+
+    @abstractmethod
+    def get_cache_control(self, config: "LLMConfig") -> Optional[Dict[str, Any]]:
+        """Get cache control settings for this factory's models."""
+        pass
+
+    @abstractmethod
+    def get_model_specific_config(self, config: "LLMConfig") -> Dict[str, Any]:
+        """Get model-specific configuration for the given config."""
+        pass
+
+    @classmethod
+    def create_factory(cls, config: "LLMConfig") -> "BaseModelFactory":
+        """Factory method to create appropriate factory instance."""
+
+        from spaik_sdk.models.factories.anthropic_factory import AnthropicModelFactory
+        from spaik_sdk.models.factories.google_factory import GoogleModelFactory
+        from spaik_sdk.models.factories.ollama_factory import OllamaModelFactory
+        from spaik_sdk.models.factories.openai_factory import OpenAIModelFactory
+
+        factories = [
+            AnthropicModelFactory(),
+            OpenAIModelFactory(),
+            GoogleModelFactory(),
+            OllamaModelFactory(),
+        ]
+        for factory in factories:
+            if factory.supports_model_config(config):
+                return factory
+
+        raise ValueError(f"No factory found that supports model config: {config}")

spaik_sdk/models/factories/google_factory.py

@@ -0,0 +1,30 @@
+from typing import Any, Dict, Optional
+
+from spaik_sdk.models.factories.base_model_factory import BaseModelFactory
+from spaik_sdk.models.llm_config import LLMConfig
+from spaik_sdk.models.llm_families import LLMFamilies
+from spaik_sdk.models.llm_model import LLMModel
+from spaik_sdk.models.model_registry import ModelRegistry
+
+
+class GoogleModelFactory(BaseModelFactory):
+    MODELS = ModelRegistry.get_by_family(LLMFamilies.GOOGLE)
+
+    def supports_model(self, model: LLMModel) -> bool:
+        return model in GoogleModelFactory.MODELS
+
+    def get_cache_control(self, config: LLMConfig) -> Optional[Dict[str, Any]]:
+        return {"type": "permanent"}
+
+    def get_model_specific_config(self, config: LLMConfig) -> Dict[str, Any]:
+        model_config: Dict[str, Any] = {"model": config.model.name, "temperature": config.temperature}
+
+        if config.reasoning:
+            model_config["thinking_budget"] = config.reasoning_budget_tokens
+            model_config["include_thoughts"] = True
+
+        # Handle streaming - Google models use disable_streaming instead of streaming
+        if not config.streaming:
+            model_config["disable_streaming"] = True
+
+        return model_config

spaik_sdk/models/factories/ollama_factory.py

@@ -0,0 +1,41 @@
+from typing import Any, Dict, Optional
+
+from spaik_sdk.models.factories.base_model_factory import BaseModelFactory
+from spaik_sdk.models.llm_config import LLMConfig
+from spaik_sdk.models.llm_families import LLMFamilies
+from spaik_sdk.models.llm_model import LLMModel
+from spaik_sdk.models.model_registry import ModelRegistry
+from spaik_sdk.utils.init_logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class OllamaModelFactory(BaseModelFactory):
+    # Models are dynamically created by users with LLMModel(family="ollama", name="...")
+    # So we start with an empty registry and let users add models as needed
+    MODELS = ModelRegistry.get_by_family(LLMFamilies.OLLAMA)
+
+    def supports_model(self, model: LLMModel) -> bool:
+        return model.family == "ollama"
+
+    def get_cache_control(self, config: LLMConfig) -> Optional[Dict[str, Any]]:
+        # Ollama doesn't support prompt caching in the same way as cloud providers
+        return None
+
+    def get_model_specific_config(self, config: LLMConfig) -> Dict[str, Any]:
+        model_config: Dict[str, Any] = {
+            "model": config.model.name,
+            "temperature": config.temperature,
+        }
+
+        # Enable streaming if requested
+        if config.streaming:
+            model_config["streaming"] = True
+
+        # Handle reasoning configuration for models that support it (like deepseek-r1)
+        # reasoning=True should separate <think> content to additional_kwargs['reasoning_content']
+        # reasoning=None/False leaves <think> tags in main content
+        if config.reasoning is not None:
+            model_config["reasoning"] = config.reasoning
+
+        return model_config

spaik_sdk/models/factories/openai_factory.py

@@ -0,0 +1,50 @@
+from typing import Any, Dict, Optional
+
+from spaik_sdk.models.factories.base_model_factory import BaseModelFactory
+from spaik_sdk.models.llm_config import LLMConfig
+from spaik_sdk.models.llm_families import LLMFamilies
+from spaik_sdk.models.llm_model import LLMModel
+from spaik_sdk.models.model_registry import ModelRegistry
+from spaik_sdk.utils.init_logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class OpenAIModelFactory(BaseModelFactory):
+    MODELS = ModelRegistry.get_by_family(LLMFamilies.OPENAI)
+
+    def supports_model(self, model: LLMModel) -> bool:
+        return model in OpenAIModelFactory.MODELS
+
+    def supports_model_config(self, config: LLMConfig) -> bool:
+        # First check basic model support
+        if not self.supports_model(config.model):
+            return False
+        if config.reasoning and not config.model.reasoning:
+            # let's not fail here, but we should log a warning
+            logger.warning(f"Model {config.model} does not support reasoning")
+        return True
+
+    def get_cache_control(self, config: LLMConfig) -> Optional[Dict[str, Any]]:
+        if config.model.prompt_caching:
+            return {"type": "permanent"}
+        return None
+
+    def get_model_specific_config(self, config: LLMConfig) -> Dict[str, Any]:
+        model_config: Dict[str, Any] = {"model": config.model.name, "streaming": config.streaming}
+        # Add parallel tool calls if tool usage is enabled
+        if config.tool_usage:
+            model_config["model_kwargs"] = {"parallel_tool_calls": True}
+
+        # Add model-specific configurations for reasoning models
+        if config.model.reasoning:
+            # Enable Responses API for reasoning models
+            model_config["use_responses_api"] = True
+
+            # Configure reasoning through model_kwargs as per LangChain docs
+            if config.reasoning_summary:
+                model_config["model_kwargs"] = {"reasoning": {"effort": config.reasoning_effort, "summary": config.reasoning_summary}}
+        else:
+            model_config["temperature"] = config.temperature
+
+        return model_config

spaik_sdk/models/llm_config.py

@@ -0,0 +1,46 @@
+from dataclasses import dataclass, replace
+from typing import Optional
+
+from spaik_sdk.models.llm_model import LLMModel
+from spaik_sdk.models.llm_wrapper import LLMWrapper
+from spaik_sdk.models.providers.base_provider import BaseProvider
+from spaik_sdk.models.providers.provider_type import ProviderType
+
+
+@dataclass
+class LLMConfig:
+    model: LLMModel
+    provider_type: Optional[ProviderType] = None
+    reasoning: bool = True
+    tool_usage: bool = True
+    streaming: bool = True
+    reasoning_summary: str = "detailed"  # Options: "auto", "concise", "detailed", None
+    reasoning_effort: str = "medium"  # Options: "low", "medium", "high"
+    max_output_tokens: int = 8192
+    reasoning_budget_tokens: int = 4096
+    temperature: float = 0.1
+    structured_response: bool = False
+
+    _model_wrapper: Optional[LLMWrapper] = None
+
+    def get_model_wrapper(self) -> LLMWrapper:
+        if self._model_wrapper is None:
+            self._model_wrapper = self.create_model_wrapper()
+        return self._model_wrapper
+
+    def create_model_wrapper(self) -> LLMWrapper:
+        provider = self.get_provider()
+        factory = self.get_factory()
+        return factory.create_model(self, provider)
+
+    def get_provider(self) -> BaseProvider:
+        return BaseProvider.create_provider(self.provider_type)
+
+    def get_factory(self):
+        # Late import to avoid circular dependency
+        from spaik_sdk.models.factories.base_model_factory import BaseModelFactory
+
+        return BaseModelFactory.create_factory(self)
+
+    def as_structured_response_config(self) -> "LLMConfig":
+        return replace(self, structured_response=True)

spaik_sdk/models/llm_families.py

@@ -0,0 +1,7 @@
+class LLMFamilies:
+    """Model identifier with extensibility support."""
+
+    ANTHROPIC = "anthropic"
+    OPENAI = "openai"
+    GOOGLE = "google"
+    OLLAMA = "ollama"

spaik_sdk/models/llm_model.py

@@ -0,0 +1,17 @@
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class LLMModel:
+    """Model identifier with extensibility support."""
+
+    family: str
+    name: str
+    reasoning: bool = True
+    prompt_caching: bool = False
+
+    def __str__(self) -> str:
+        return self.name
+
+    def __repr__(self) -> str:
+        return f"LLMModel('{self.name}')"

spaik_sdk/models/llm_wrapper.py

@@ -0,0 +1,25 @@
+from typing import Any, Dict, Optional
+
+from langchain_core.language_models.chat_models import BaseChatModel
+
+from spaik_sdk.models.llm_model import LLMModel
+
+
+class LLMWrapper:
+    def __init__(self, langchain_model: BaseChatModel, cache_control: Optional[Dict[str, Any]], model_type: LLMModel):
+        """Initialize wrapper with pre-created langchain model and cache control."""
+        self._langchain_model = langchain_model
+        self._cache_control = cache_control
+        self._model_type = model_type
+
+    def get_langchain_model(self) -> BaseChatModel:
+        """Get the underlying langchain model instance."""
+        return self._langchain_model
+
+    def get_cache_control(self) -> Optional[Dict[str, Any]]:
+        """Get cache control settings for this model."""
+        return self._cache_control
+
+    def get_model_type(self) -> LLMModel:
+        """Get the model type enum."""
+        return self._model_type