spaik-sdk 0.6.2__py3-none-any.whl

spaik_sdk/__init__.py

@@ -0,0 +1,21 @@
+"""Spaik SDK - SDK for building various kinds of agentic + AI solutions."""
+
+from .agent.base_agent import BaseAgent
+from .models.llm_config import LLMConfig
+from .models.llm_model import LLMModel
+from .models.providers.provider_type import ProviderType
+from .thread.models import MessageBlock, MessageBlockType, ThreadMessage
+from .thread.thread_container import ThreadContainer
+
+__all__ = [
+    "BaseAgent",
+    "LLMConfig",
+    "LLMModel",
+    "ProviderType",
+    "ThreadContainer",
+    "ThreadMessage",
+    "MessageBlock",
+    "MessageBlockType",
+]
+
+__version__ = "0.0.1"

spaik_sdk/agent/base_agent.py

@@ -0,0 +1,249 @@
+import asyncio
+from abc import ABC
+from typing import Any, AsyncGenerator, Dict, List, Optional, Type, TypeVar
+
+from langchain_core.language_models import BaseChatModel
+from langchain_core.tools import BaseTool
+from langgraph.graph.state import CompiledStateGraph
+from pydantic import BaseModel
+
+from spaik_sdk.attachments.models import Attachment
+from spaik_sdk.config.env import env_config
+from spaik_sdk.llm.cancellation_handle import CancellationHandle
+from spaik_sdk.llm.cost.builtin_cost_provider import BuiltinCostProvider
+from spaik_sdk.llm.cost.cost_estimate import CostEstimate
+from spaik_sdk.llm.cost.cost_provider import CostProvider
+from spaik_sdk.llm.langchain_service import LangChainService
+from spaik_sdk.models.llm_config import LLMConfig
+from spaik_sdk.models.llm_model import LLMModel
+from spaik_sdk.models.providers.provider_type import ProviderType
+from spaik_sdk.prompt.get_prompt_loader import get_prompt_loader
+from spaik_sdk.prompt.prompt_loader import PromptLoader
+from spaik_sdk.prompt.prompt_loader_mode import PromptLoaderMode
+from spaik_sdk.recording.conditional_recorder import ConditionalRecorder
+from spaik_sdk.thread.adapters.cli.live_cli import LiveCLI
+from spaik_sdk.thread.adapters.event_adapter import EventAdapter
+from spaik_sdk.thread.adapters.sync_adapter import SyncAdapter
+from spaik_sdk.thread.models import BlockFullyAddedEvent, ThreadEvent, ThreadMessage
+from spaik_sdk.thread.thread_container import ThreadContainer
+from spaik_sdk.tools.tool_provider import ToolProvider
+from spaik_sdk.tracing.agent_trace import AgentTrace
+from spaik_sdk.tracing.trace_sink import TraceSink
+from spaik_sdk.utils.init_logger import init_logger
+
+logger = init_logger(__name__)
+
+T = TypeVar("T", bound=BaseModel)
+
+
+class BaseAgent(ABC):
+    def __init__(
+        self,
+        system_prompt_args: Dict[str, Any] = {},
+        system_prompt_version: Optional[str] = None,
+        system_prompt: Optional[str] = None,
+        prompt_loader: Optional[PromptLoader] = None,
+        prompt_loader_mode: Optional[PromptLoaderMode] = None,
+        llm_config: Optional[LLMConfig] = None,
+        llm_model: Optional[LLMModel] = None,
+        reasoning: Optional[bool] = None,
+        trace: Optional[AgentTrace] = None,
+        trace_sink: Optional[TraceSink] = None,
+        thread_container: Optional[ThreadContainer] = None,
+        tools: Optional[List[BaseTool]] = None,
+        tool_providers: Optional[List[ToolProvider]] = None,
+        recorder: Optional[ConditionalRecorder] = None,
+        cancellation_handle: Optional[CancellationHandle] = None,
+        cost_provider: Optional[CostProvider] = None,
+    ):
+        logger.debug("Initializing BaseAgent")
+        self.prompt_loader = prompt_loader or get_prompt_loader(prompt_loader_mode)
+        self.system_prompt = system_prompt or self._get_system_prompt(system_prompt_args, system_prompt_version)
+        self.trace = trace or AgentTrace(self.system_prompt, self.__class__.__name__, trace_sink=trace_sink)
+        self.thread_container = thread_container or ThreadContainer(self.system_prompt)
+        self.tools = tools or self._create_tools(tool_providers)
+        self.llm_config = llm_config or self.create_llm_config(llm_model, reasoning)
+        self.recorder = recorder.get_recorder() if recorder is not None else None
+        self.playback = recorder.get_playback() if recorder is not None else None
+        self.thread_container.subscribe(self._on_thread_event)
+        self.cancellation_handle = cancellation_handle
+        self.cost_provider = cost_provider or BuiltinCostProvider()
+
+    def get_response_stream(
+        self,
+        user_input: Optional[str] = None,
+        attachments: Optional[List[Attachment]] = None,
+    ) -> AsyncGenerator[Dict[str, Any], None]:
+        self.trace.add_input(user_input)
+        langchain_service = self._create_langchain_service()
+        return langchain_service.execute_stream_tokens(user_input, self.tools, attachments)
+
+    async def get_event_stream(
+        self,
+        user_input: Optional[str] = None,
+        attachments: Optional[List[Attachment]] = None,
+    ) -> AsyncGenerator[ThreadEvent, None]:
+        event_adapter = EventAdapter(self.thread_container)
+        async for _event in self.get_response_stream(user_input, attachments):
+            new_events = event_adapter.flush()
+            for new_event in new_events:
+                yield new_event
+        event_adapter.cleanup()
+
+    def get_response(
+        self,
+        user_input: Optional[str] = None,
+        attachments: Optional[List[Attachment]] = None,
+    ) -> ThreadMessage:
+        return asyncio.run(self.get_response_async(user_input, attachments))
+
+    async def get_response_async(
+        self,
+        user_input: Optional[str] = None,
+        attachments: Optional[List[Attachment]] = None,
+    ) -> ThreadMessage:
+        sync_adapter = SyncAdapter(self.thread_container)
+        await sync_adapter.run_async(self.get_response_stream(user_input, attachments))
+        ret = await sync_adapter.wait_for_completion_async()
+        if ret is None:
+            raise ValueError("No response received")
+        self.thread_container.complete_generation()
+        return ret
+
+    def get_response_text(
+        self,
+        user_input: Optional[str] = None,
+        attachments: Optional[List[Attachment]] = None,
+    ) -> str:
+        return self.get_response(user_input, attachments).get_text_content()
+
+    async def get_response_text_async(
+        self,
+        user_input: Optional[str] = None,
+        attachments: Optional[List[Attachment]] = None,
+    ) -> str:
+        return (await self.get_response_async(user_input, attachments)).get_text_content()
+
+    def get_structured_response(self, prompt: str, output_schema: Type[T]) -> T:
+        self.trace.add_structured_response_input(prompt, output_schema)
+        llm_config = self.llm_config.as_structured_response_config()
+        langchain_service = self._create_langchain_service(llm_config)
+        ret = langchain_service.get_structured_response(prompt, output_schema)
+        self.trace.add_structured_response_output(ret)
+        return ret
+
+    def run_cli(self):
+        asyncio.run(LiveCLI(self.thread_container).run_interactive(self))
+
+    def create_llm_config(self, llm_model: Optional[LLMModel] = None, reasoning: Optional[bool] = None) -> LLMConfig:
+        if llm_model is None:
+            llm_model = self.get_llm_model()
+
+        provider_type = ProviderType.from_family(llm_model.family)
+
+        return LLMConfig(
+            model=llm_model,
+            provider_type=provider_type,
+            reasoning=reasoning if reasoning is not None else llm_model.reasoning,
+            tool_usage=len(self.tools) > 0,
+        )
+
+    def get_llm_model(self) -> LLMModel:
+        return env_config.get_default_model()
+
+    def is_reasoning(self) -> bool:
+        return True
+
+    def get_tool_providers(self) -> List[ToolProvider]:
+        return []
+
+    def set_thread_container(self, thread_container: ThreadContainer) -> None:
+        self.thread_container = thread_container
+        if self.thread_container.system_prompt is None:
+            self.thread_container.system_prompt = self.system_prompt
+        self.thread_container.subscribe(self._on_thread_event)
+
+    def set_cancellation_handle(self, cancellation_handle: Optional[CancellationHandle]) -> None:
+        self.cancellation_handle = cancellation_handle
+
+    def _get_prompt(self, prompt_name: str, args: Dict[str, Any], version: Optional[str] = None) -> str:
+        return self.prompt_loader.get_agent_prompt(self.__class__, prompt_name, args, version)
+
+    def _get_system_prompt(self, args: Dict[str, Any], version: Optional[str] = None) -> str:
+        return self.prompt_loader.get_system_prompt(self.__class__, args, version)
+
+    def _create_tools(self, tool_providers: Optional[List[ToolProvider]] = None) -> List[BaseTool]:
+        tool_providers = tool_providers or self.get_tool_providers()
+        tools = []
+        for provider in tool_providers:
+            tools.extend(provider.get_tools())
+        return tools
+
+    def _create_langchain_service(self, llm_config: Optional[LLMConfig] = None) -> LangChainService:
+        if self.thread_container is None:
+            raise ValueError("Thread container is not set")
+        return LangChainService(
+            llm_config or self.llm_config,
+            self.thread_container,
+            self.__class__.__name__,
+            self.__class__.__name__,
+            self.recorder,
+            self.playback,
+            self.cancellation_handle,
+        )
+
+    def _on_thread_event(self, event: ThreadEvent) -> None:
+        logger.debug(f"Thread event: {event}")
+        """Handle thread events and forward to trace"""
+        if isinstance(event, BlockFullyAddedEvent):
+            self.trace.add_block(event.block)
+
+    def get_cost(self, latest_only: bool = False) -> CostEstimate:
+        token_usage = self.thread_container.get_token_usage()
+        if latest_only:
+            token_usage = self.thread_container.get_latest_token_usage()
+        if token_usage is None:
+            from spaik_sdk.llm.consumption.token_usage import TokenUsage
+
+            token_usage = TokenUsage(input_tokens=0, output_tokens=0)
+        return self.cost_provider.get_cost_estimate(self.get_llm_model(), token_usage)
+
+    def get_langchain_model(self) -> BaseChatModel:
+        """Get the underlying LangChain chat model.
+
+        Returns the configured LangChain BaseChatModel that can be used directly
+        with LangChain or LangGraph for custom workflows.
+
+        Returns:
+            BaseChatModel: The LangChain chat model instance.
+        """
+        return self.llm_config.get_model_wrapper().get_langchain_model()
+
+    def get_react_agent(self) -> CompiledStateGraph:
+        """Get a LangGraph react agent with the agent's tools and model.
+
+        Returns a compiled LangGraph react agent that can be used directly
+        with LangGraph workflows, custom orchestration, or integrated into
+        larger agent systems.
+
+        The react agent includes all tools configured on this agent and uses
+        the agent's LLM configuration.
+
+        Returns:
+            CompiledStateGraph: A compiled LangGraph react agent.
+
+        Example:
+            ```python
+            agent = MyAgent()
+            react_agent = agent.get_react_agent()
+
+            # Use directly with LangGraph
+            result = await react_agent.ainvoke({"messages": messages})
+
+            # Or stream events
+            async for event in react_agent.astream_events({"messages": messages}, version="v2"):
+                print(event)
+            ```
+        """
+        langchain_service = self._create_langchain_service()
+        return langchain_service.create_executor(self.tools)

spaik_sdk/attachments/__init__.py

@@ -0,0 +1,22 @@
+from spaik_sdk.attachments.builder import AttachmentBuilder
+from spaik_sdk.attachments.file_storage_provider import (
+    get_file_storage,
+    get_or_create_file_storage,
+    reset_file_storage,
+    set_file_storage,
+)
+from spaik_sdk.attachments.models import Attachment, FileMetadata
+from spaik_sdk.attachments.storage.base_file_storage import BaseFileStorage
+from spaik_sdk.attachments.storage.impl.local_file_storage import LocalFileStorage
+
+__all__ = [
+    "Attachment",
+    "AttachmentBuilder",
+    "BaseFileStorage",
+    "FileMetadata",
+    "LocalFileStorage",
+    "get_file_storage",
+    "get_or_create_file_storage",
+    "reset_file_storage",
+    "set_file_storage",
+]

spaik_sdk/attachments/builder.py

@@ -0,0 +1,61 @@
+from pathlib import Path
+from typing import List, Optional, Union
+
+from spaik_sdk.attachments.mime_types import guess_mime_type
+from spaik_sdk.attachments.models import Attachment
+from spaik_sdk.attachments.storage.base_file_storage import BaseFileStorage
+
+
+class AttachmentBuilder:
+    def __init__(self, file_storage: BaseFileStorage):
+        self.file_storage = file_storage
+
+    async def from_path(
+        self,
+        path: Union[str, Path],
+        owner_id: str = "system",
+        mime_type: Optional[str] = None,
+    ) -> Attachment:
+        path = Path(path)
+
+        if not path.exists():
+            raise FileNotFoundError(f"File not found: {path}")
+
+        if mime_type is None:
+            mime_type = guess_mime_type(path)
+            if mime_type is None:
+                raise ValueError(f"Could not determine MIME type for: {path}")
+
+        with open(path, "rb") as f:
+            data = f.read()
+
+        metadata = await self.file_storage.store(
+            data=data,
+            mime_type=mime_type,
+            owner_id=owner_id,
+            filename=path.name,
+        )
+
+        return metadata.to_attachment()
+
+    async def from_paths(
+        self,
+        paths: List[Union[str, Path]],
+        owner_id: str = "system",
+    ) -> List[Attachment]:
+        return [await self.from_path(p, owner_id=owner_id) for p in paths]
+
+    async def from_bytes(
+        self,
+        data: bytes,
+        mime_type: str,
+        owner_id: str = "system",
+        filename: Optional[str] = None,
+    ) -> Attachment:
+        metadata = await self.file_storage.store(
+            data=data,
+            mime_type=mime_type,
+            owner_id=owner_id,
+            filename=filename,
+        )
+        return metadata.to_attachment()

spaik_sdk/attachments/file_storage_provider.py

@@ -0,0 +1,27 @@
+from typing import Optional
+
+from spaik_sdk.attachments.storage.base_file_storage import BaseFileStorage
+from spaik_sdk.attachments.storage.impl.local_file_storage import LocalFileStorage
+
+_file_storage: Optional[BaseFileStorage] = None
+
+
+def get_file_storage() -> Optional[BaseFileStorage]:
+    return _file_storage
+
+
+def set_file_storage(storage: BaseFileStorage) -> None:
+    global _file_storage
+    _file_storage = storage
+
+
+def get_or_create_file_storage() -> BaseFileStorage:
+    global _file_storage
+    if _file_storage is None:
+        _file_storage = LocalFileStorage()
+    return _file_storage
+
+
+def reset_file_storage() -> None:
+    global _file_storage
+    _file_storage = None

spaik_sdk/attachments/mime_types.py

@@ -0,0 +1,118 @@
+from pathlib import Path
+from typing import Optional
+
+
+class MimeTypes:
+    # Images
+    PNG = "image/png"
+    JPEG = "image/jpeg"
+    GIF = "image/gif"
+    WEBP = "image/webp"
+    SVG = "image/svg+xml"
+
+    # Documents
+    PDF = "application/pdf"
+
+    # Audio
+    MP3 = "audio/mpeg"
+    WAV = "audio/wav"
+    OGG = "audio/ogg"
+    WEBM_AUDIO = "audio/webm"
+
+    # Video
+    MP4 = "video/mp4"
+    WEBM_VIDEO = "video/webm"
+    MOV = "video/quicktime"
+
+    # Text
+    PLAIN = "text/plain"
+    HTML = "text/html"
+    CSV = "text/csv"
+
+
+IMAGE_TYPES = frozenset(
+    {
+        MimeTypes.PNG,
+        MimeTypes.JPEG,
+        MimeTypes.GIF,
+        MimeTypes.WEBP,
+        MimeTypes.SVG,
+    }
+)
+
+DOCUMENT_TYPES = frozenset(
+    {
+        MimeTypes.PDF,
+    }
+)
+
+AUDIO_TYPES = frozenset(
+    {
+        MimeTypes.MP3,
+        MimeTypes.WAV,
+        MimeTypes.OGG,
+        MimeTypes.WEBM_AUDIO,
+    }
+)
+
+VIDEO_TYPES = frozenset(
+    {
+        MimeTypes.MP4,
+        MimeTypes.WEBM_VIDEO,
+        MimeTypes.MOV,
+    }
+)
+
+TEXT_TYPES = frozenset(
+    {
+        MimeTypes.PLAIN,
+        MimeTypes.HTML,
+        MimeTypes.CSV,
+    }
+)
+
+ALL_SUPPORTED_TYPES = IMAGE_TYPES | DOCUMENT_TYPES | AUDIO_TYPES | VIDEO_TYPES | TEXT_TYPES
+
+EXTENSION_TO_MIME: dict[str, str] = {
+    ".png": MimeTypes.PNG,
+    ".jpg": MimeTypes.JPEG,
+    ".jpeg": MimeTypes.JPEG,
+    ".gif": MimeTypes.GIF,
+    ".webp": MimeTypes.WEBP,
+    ".svg": MimeTypes.SVG,
+    ".pdf": MimeTypes.PDF,
+    ".mp3": MimeTypes.MP3,
+    ".wav": MimeTypes.WAV,
+    ".ogg": MimeTypes.OGG,
+    ".mp4": MimeTypes.MP4,
+    ".webm": MimeTypes.WEBM_VIDEO,
+    ".mov": MimeTypes.MOV,
+    ".txt": MimeTypes.PLAIN,
+    ".html": MimeTypes.HTML,
+    ".csv": MimeTypes.CSV,
+}
+
+
+def guess_mime_type(path: str | Path) -> Optional[str]:
+    suffix = Path(path).suffix.lower()
+    return EXTENSION_TO_MIME.get(suffix)
+
+
+def is_image(mime_type: str) -> bool:
+    return mime_type in IMAGE_TYPES
+
+
+def is_document(mime_type: str) -> bool:
+    return mime_type in DOCUMENT_TYPES
+
+
+def is_audio(mime_type: str) -> bool:
+    return mime_type in AUDIO_TYPES
+
+
+def is_video(mime_type: str) -> bool:
+    return mime_type in VIDEO_TYPES
+
+
+def is_supported(mime_type: str) -> bool:
+    return mime_type in ALL_SUPPORTED_TYPES

spaik_sdk/attachments/models.py

@@ -0,0 +1,63 @@
+import time
+from dataclasses import dataclass, field
+from typing import Any, Dict, Optional
+
+
+@dataclass
+class Attachment:
+    file_id: str
+    mime_type: str
+    filename: Optional[str] = None
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "file_id": self.file_id,
+            "mime_type": self.mime_type,
+            "filename": self.filename,
+        }
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "Attachment":
+        return cls(
+            file_id=data["file_id"],
+            mime_type=data["mime_type"],
+            filename=data.get("filename"),
+        )
+
+
+@dataclass
+class FileMetadata:
+    file_id: str
+    mime_type: str
+    owner_id: str
+    size_bytes: int
+    created_at: int = field(default_factory=lambda: int(time.time() * 1000))
+    filename: Optional[str] = None
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "file_id": self.file_id,
+            "mime_type": self.mime_type,
+            "filename": self.filename,
+            "owner_id": self.owner_id,
+            "created_at": self.created_at,
+            "size_bytes": self.size_bytes,
+        }
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "FileMetadata":
+        return cls(
+            file_id=data["file_id"],
+            mime_type=data["mime_type"],
+            filename=data.get("filename"),
+            owner_id=data["owner_id"],
+            created_at=data["created_at"],
+            size_bytes=data["size_bytes"],
+        )
+
+    def to_attachment(self) -> Attachment:
+        return Attachment(
+            file_id=self.file_id,
+            mime_type=self.mime_type,
+            filename=self.filename,
+        )

spaik_sdk/attachments/provider_support.py

@@ -0,0 +1,53 @@
+from spaik_sdk.attachments.mime_types import (
+    AUDIO_TYPES,
+    DOCUMENT_TYPES,
+    IMAGE_TYPES,
+    VIDEO_TYPES,
+    MimeTypes,
+)
+
+OPENAI_SUPPORTED = frozenset(
+    {
+        MimeTypes.PNG,
+        MimeTypes.JPEG,
+        MimeTypes.GIF,
+        MimeTypes.WEBP,
+        MimeTypes.PDF,
+    }
+)
+
+ANTHROPIC_SUPPORTED = frozenset(
+    {
+        MimeTypes.PNG,
+        MimeTypes.JPEG,
+        MimeTypes.GIF,
+        MimeTypes.WEBP,
+        MimeTypes.PDF,
+    }
+)
+
+GOOGLE_SUPPORTED = frozenset(
+    {
+        *IMAGE_TYPES,
+        *DOCUMENT_TYPES,
+        *AUDIO_TYPES,
+        *VIDEO_TYPES,
+    }
+)
+
+PROVIDER_FAMILY_SUPPORT: dict[str, frozenset[str]] = {
+    "openai": OPENAI_SUPPORTED,
+    "anthropic": ANTHROPIC_SUPPORTED,
+    "google": GOOGLE_SUPPORTED,
+    "azure": OPENAI_SUPPORTED,
+    "ollama": IMAGE_TYPES,
+}
+
+
+def get_supported_types(provider_family: str) -> frozenset[str]:
+    return PROVIDER_FAMILY_SUPPORT.get(provider_family.lower(), IMAGE_TYPES)
+
+
+def is_supported_by_provider(mime_type: str, provider_family: str) -> bool:
+    supported = get_supported_types(provider_family)
+    return mime_type in supported

spaik_sdk/attachments/storage/base_file_storage.py

@@ -0,0 +1,32 @@
+from abc import ABC, abstractmethod
+from typing import Optional
+
+from spaik_sdk.attachments.models import FileMetadata
+
+
+class BaseFileStorage(ABC):
+    @abstractmethod
+    async def store(
+        self,
+        data: bytes,
+        mime_type: str,
+        owner_id: str,
+        filename: Optional[str] = None,
+    ) -> FileMetadata:
+        pass
+
+    @abstractmethod
+    async def retrieve(self, file_id: str) -> tuple[bytes, FileMetadata]:
+        pass
+
+    @abstractmethod
+    async def get_metadata(self, file_id: str) -> Optional[FileMetadata]:
+        pass
+
+    @abstractmethod
+    async def delete(self, file_id: str) -> bool:
+        pass
+
+    @abstractmethod
+    async def exists(self, file_id: str) -> bool:
+        pass