spaik-sdk 0.6.2__py3-none-any.whl

spaik_sdk/config/get_credentials_provider.py

@@ -0,0 +1,14 @@
+from spaik_sdk.config.credentials_provider import CredentialsProvider
+from spaik_sdk.config.env import env_config
+from spaik_sdk.config.env_credentials_provider import EnvCredentialsProvider
+
+
+def get_credentials_provider() -> CredentialsProvider:
+    provider_type = env_config.get_credentials_provider_type()
+    if provider_type == "env":
+        return EnvCredentialsProvider()
+    else:
+        raise ValueError(f"Unsupported provider type: {provider_type}")
+
+
+credentials_provider = get_credentials_provider()

spaik_sdk/image_gen/__init__.py

@@ -0,0 +1,9 @@
+from spaik_sdk.image_gen.image_generator import ImageGenerator
+from spaik_sdk.image_gen.options import ImageFormat, ImageGenOptions, ImageQuality
+
+__all__ = [
+    "ImageGenerator",
+    "ImageGenOptions",
+    "ImageQuality",
+    "ImageFormat",
+]

spaik_sdk/image_gen/image_generator.py

@@ -0,0 +1,83 @@
+import re
+from pathlib import Path
+from time import time
+
+from spaik_sdk.config.env import env_config
+from spaik_sdk.config.get_credentials_provider import credentials_provider
+from spaik_sdk.image_gen.options import ImageFormat, ImageGenOptions
+from spaik_sdk.image_gen.providers import google as google_provider
+from spaik_sdk.image_gen.providers import openai as openai_provider
+
+
+def _slugify(text: str, max_length: int = 50) -> str:
+    slug = re.sub(r"[^a-zA-Z0-9]+", "_", text.lower())
+    slug = slug.strip("_")
+    return slug[:max_length]
+
+
+def _generate_filename(prompt: str, output_format: ImageFormat) -> str:
+    slug = _slugify(prompt)
+    timestamp = int(time())
+    extension = output_format.value
+    return f"{slug}_{timestamp}.{extension}"
+
+
+class ImageGenerator:
+    def __init__(
+        self,
+        model: str | None = None,
+        output_dir: str = ".",
+        endpoint: str | None = None,
+        headers: dict[str, str] | None = None,
+    ):
+        self.model = model or env_config.get_image_model()
+        self.output_dir = Path(output_dir)
+        self.endpoint = endpoint
+        self.headers = headers
+
+    def _get_provider(self) -> str:
+        if self.model.startswith("gpt-image"):
+            return "openai"
+        elif self.model.startswith("gemini"):
+            return "google"
+        else:
+            raise ValueError(f"Unknown image model provider for: {self.model}")
+
+    async def generate_image(
+        self,
+        prompt: str,
+        options: ImageGenOptions | None = None,
+        output_filename: str | None = None,
+    ) -> Path:
+        opts = options or ImageGenOptions()
+        provider = self._get_provider()
+
+        if provider == "openai":
+            api_key = credentials_provider.get_provider_key("openai")
+            image_bytes = await openai_provider.generate_image(
+                prompt=prompt,
+                model=self.model,
+                api_key=api_key,
+                options=opts,
+                endpoint=self.endpoint,
+                headers=self.headers,
+            )
+        elif provider == "google":
+            api_key = credentials_provider.get_provider_key("google")
+            image_bytes = await google_provider.generate_image(
+                prompt=prompt,
+                model=self.model,
+                api_key=api_key,
+                options=opts,
+                endpoint=self.endpoint,
+                headers=self.headers,
+            )
+        else:
+            raise ValueError(f"Unsupported provider: {provider}")
+
+        filename = output_filename or _generate_filename(prompt, opts.output_format)
+        self.output_dir.mkdir(parents=True, exist_ok=True)
+        output_path = self.output_dir / filename
+
+        output_path.write_bytes(image_bytes)
+        return output_path

spaik_sdk/image_gen/options.py

@@ -0,0 +1,24 @@
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any
+
+
+class ImageQuality(Enum):
+    LOW = "low"
+    MEDIUM = "medium"
+    HIGH = "high"
+
+
+class ImageFormat(Enum):
+    PNG = "png"
+    JPEG = "jpeg"
+    WEBP = "webp"
+
+
+@dataclass
+class ImageGenOptions:
+    width: int = 1024
+    height: int = 1024
+    quality: ImageQuality = ImageQuality.MEDIUM
+    output_format: ImageFormat = ImageFormat.PNG
+    vendor: dict[str, Any] = field(default_factory=dict)

spaik_sdk/image_gen/providers/google.py

@@ -0,0 +1,75 @@
+import base64
+
+import httpx
+
+from spaik_sdk.image_gen.options import ImageGenOptions
+
+GOOGLE_GENERATIVE_ENDPOINT = "https://generativelanguage.googleapis.com/v1/models"
+
+
+def _derive_resolution_and_aspect(width: int, height: int) -> tuple[str, str]:
+    from math import gcd
+
+    divisor = gcd(width, height)
+    aspect_w = width // divisor
+    aspect_h = height // divisor
+    aspect_ratio = f"{aspect_w}:{aspect_h}"
+
+    max_dim = max(width, height)
+    if max_dim <= 1024:
+        resolution = "1K"
+    elif max_dim <= 2048:
+        resolution = "2K"
+    else:
+        resolution = "4K"
+
+    return resolution, aspect_ratio
+
+
+async def generate_image(
+    prompt: str,
+    model: str,
+    api_key: str,
+    options: ImageGenOptions,
+    endpoint: str | None = None,
+    headers: dict[str, str] | None = None,
+) -> bytes:
+    base_url = endpoint or GOOGLE_GENERATIVE_ENDPOINT
+    url = f"{base_url}/{model}:generateContent?key={api_key}"
+
+    request_headers = {
+        "Content-Type": "application/json",
+    }
+    if headers:
+        request_headers.update(headers)
+
+    resolution, aspect_ratio = _derive_resolution_and_aspect(options.width, options.height)
+
+    generation_config: dict = {
+        "responseModalities": ["image", "text"],
+        "resolution": resolution,
+        "aspectRatio": aspect_ratio,
+    }
+    generation_config.update(options.vendor)
+
+    payload = {
+        "contents": [{"parts": [{"text": prompt}]}],
+        "generationConfig": generation_config,
+    }
+
+    async with httpx.AsyncClient(timeout=120.0) as client:
+        response = await client.post(url, headers=request_headers, json=payload)
+        response.raise_for_status()
+        data = response.json()
+
+    candidates = data.get("candidates", [])
+    if not candidates:
+        raise ValueError("No image generated by Gemini API")
+
+    parts = candidates[0].get("content", {}).get("parts", [])
+    for part in parts:
+        if "inlineData" in part:
+            image_b64 = part["inlineData"]["data"]
+            return base64.b64decode(image_b64)
+
+    raise ValueError("No image data found in Gemini API response")

spaik_sdk/image_gen/providers/openai.py

@@ -0,0 +1,60 @@
+import base64
+
+import httpx
+
+from spaik_sdk.image_gen.options import ImageFormat, ImageGenOptions
+
+OPENAI_IMAGES_ENDPOINT = "https://api.openai.com/v1/images/generations"
+
+
+async def generate_image(
+    prompt: str,
+    model: str,
+    api_key: str,
+    options: ImageGenOptions,
+    endpoint: str | None = None,
+    headers: dict[str, str] | None = None,
+) -> bytes:
+    url = endpoint or OPENAI_IMAGES_ENDPOINT
+
+    request_headers = {
+        "Authorization": f"Bearer {api_key}",
+        "Content-Type": "application/json",
+    }
+    if headers:
+        request_headers.update(headers)
+
+    size = f"{options.width}x{options.height}"
+
+    response_format_map = {
+        ImageFormat.PNG: "png",
+        ImageFormat.JPEG: "jpeg",
+        ImageFormat.WEBP: "webp",
+    }
+
+    payload: dict = {
+        "model": model,
+        "prompt": prompt,
+        "size": size,
+        "quality": options.quality.value,
+        "output_format": response_format_map[options.output_format],
+        "n": 1,
+    }
+    payload.update(options.vendor)
+
+    async with httpx.AsyncClient(timeout=120.0) as client:
+        response = await client.post(url, headers=request_headers, json=payload)
+        if response.status_code != 200:
+            raise ValueError(f"OpenAI API error {response.status_code}: {response.text}")
+        data = response.json()
+
+    image_data = data["data"][0]
+    if "b64_json" in image_data:
+        return base64.b64decode(image_data["b64_json"])
+    elif "url" in image_data:
+        async with httpx.AsyncClient(timeout=120.0) as client:
+            img_response = await client.get(image_data["url"])
+            img_response.raise_for_status()
+            return img_response.content
+    else:
+        raise ValueError(f"Unexpected response format: {image_data.keys()}")

spaik_sdk/llm/cancellation_handle.py

@@ -0,0 +1,10 @@
+from abc import ABC, abstractmethod
+
+
+class CancellationHandle(ABC):
+    """Abstract base class for handling cancellation of LLM operations."""
+
+    @abstractmethod
+    async def is_cancelled(self) -> bool:
+        """Check if the operation has been cancelled."""
+        pass

spaik_sdk/llm/consumption/consumption_estimate.py

@@ -0,0 +1,26 @@
+from dataclasses import dataclass, field
+from typing import Any, Dict
+
+from spaik_sdk.llm.consumption.token_usage import TokenUsage
+
+
+@dataclass
+class ConsumptionEstimate:
+    """Consumption estimation for a request."""
+
+    token_usage: TokenUsage
+    metadata: Dict[str, Any] = field(default_factory=dict)
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary for serialization."""
+        return {
+            "token_usage": {
+                "input_tokens": self.token_usage.input_tokens,
+                "output_tokens": self.token_usage.output_tokens,
+                "total_tokens": self.token_usage.total_tokens,
+                "reasoning_tokens": self.token_usage.reasoning_tokens,
+                "cache_creation_tokens": self.token_usage.cache_creation_tokens,
+                "cache_read_tokens": self.token_usage.cache_read_tokens,
+            },
+            "metadata": self.metadata or {},
+        }

spaik_sdk/llm/consumption/consumption_estimate_builder.py

@@ -0,0 +1,113 @@
+from typing import Any, Dict, Optional
+
+from spaik_sdk.llm.consumption.consumption_estimate import ConsumptionEstimate
+from spaik_sdk.llm.consumption.token_usage import TokenUsage
+from spaik_sdk.utils.init_logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class ConsumptionEstimateBuilder:
+    """Builder for creating consumption estimates from various data sources."""
+
+    def __init__(self):
+        self._token_usage: Optional[TokenUsage] = None
+        self._metadata: Dict[str, Any] = {}
+
+    def from_usage_metadata(self, usage_metadata: Dict[str, Any]) -> "ConsumptionEstimateBuilder":
+        """Extract consumption data from LangChain usage_metadata."""
+
+        logger.info(f"usage_metadata: {usage_metadata}")
+        if not usage_metadata:
+            return self
+
+        # Extract basic token counts
+        input_tokens = usage_metadata.get("input_tokens", 0)
+        output_tokens = usage_metadata.get("output_tokens", 0)
+        total_tokens = usage_metadata.get("total_tokens", 0)
+
+        # Extract reasoning tokens (Google/OpenAI)
+        reasoning_tokens = 0
+        output_token_details = usage_metadata.get("output_token_details")
+        if output_token_details:
+            reasoning_tokens = output_token_details.get("reasoning", 0)
+
+        # Extract cache tokens (Anthropic)
+        cache_creation_tokens = 0
+        cache_read_tokens = 0
+        input_token_details = usage_metadata.get("input_token_details")
+        if input_token_details:
+            cache_creation_tokens = input_token_details.get("cache_creation", 0)
+            cache_read_tokens = input_token_details.get("cache_read", 0)
+
+        self._token_usage = TokenUsage(
+            input_tokens=input_tokens,
+            output_tokens=output_tokens,
+            total_tokens=total_tokens,
+            reasoning_tokens=reasoning_tokens,
+            cache_creation_tokens=cache_creation_tokens,
+            cache_read_tokens=cache_read_tokens,
+        )
+
+        logger.info(f"self._token_usage: {self._token_usage}")
+        return self
+
+    def from_response_metadata(self, response_metadata: Any) -> "ConsumptionEstimateBuilder":
+        """Extract consumption data from response_metadata when usage_metadata is unavailable."""
+        if not response_metadata:
+            return self
+
+        # Store metadata for potential estimation
+        self._metadata.update(
+            {
+                "finish_reason": getattr(response_metadata, "finish_reason", None),
+                "stop_reason": getattr(response_metadata, "stop_reason", None),
+                "response_id": getattr(response_metadata, "id", None),
+            }
+        )
+
+        return self
+
+    def from_event_metadata(self, metadata: Dict[str, Any]) -> "ConsumptionEstimateBuilder":
+        """Extract relevant info from LangChain event metadata."""
+        if not metadata:
+            return self
+
+        # Store any potentially useful metadata
+        self._metadata.update(
+            {
+                "ls_provider": metadata.get("ls_provider"),
+                "ls_model_name": metadata.get("ls_model_name"),
+            }
+        )
+
+        return self
+
+    def estimate_from_content(self, content: str) -> "ConsumptionEstimateBuilder":
+        """Rough estimation when no usage metadata is available."""
+        if not content or self._token_usage:
+            return self
+
+        # Very rough token estimation (4 chars per token average)
+        estimated_output_tokens = len(content) // 4
+        estimated_input_tokens = estimated_output_tokens * 5
+
+        self._token_usage = TokenUsage(
+            input_tokens=estimated_input_tokens,
+            output_tokens=estimated_output_tokens,
+            total_tokens=estimated_output_tokens,
+        )
+
+        self._metadata["estimation_method"] = "content_length"
+
+        return self
+
+    def build(self) -> Optional[ConsumptionEstimate]:
+        """Build the final ConsumptionEstimate."""
+        if not self._token_usage:
+            return None
+
+        return ConsumptionEstimate(
+            token_usage=self._token_usage,
+            metadata=self._metadata,
+        )

spaik_sdk/llm/consumption/consumption_extractor.py

@@ -0,0 +1,59 @@
+from typing import Any, Dict, Optional
+
+from spaik_sdk.llm.consumption.consumption_estimate import ConsumptionEstimate
+from spaik_sdk.llm.consumption.consumption_estimate_builder import ConsumptionEstimateBuilder
+from spaik_sdk.utils.init_logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class ConsumptionExtractor:
+    """Extracts consumption information from LangChain streaming events."""
+
+    def extract_from_stream_end(self, data: Dict[str, Any]) -> Optional[ConsumptionEstimate]:
+        """Extract consumption data from on_chat_model_end event data."""
+        builder = ConsumptionEstimateBuilder()
+
+        # Extract from event metadata if available
+        if "metadata" in data:
+            builder.from_event_metadata(data["metadata"])
+
+        # Extract from output data
+        if "output" in data:
+            output = data["output"]
+
+            # Try usage_metadata first (preferred)
+            usage_metadata = getattr(output, "usage_metadata", None)
+            if usage_metadata:
+                usage_metadata = dict(usage_metadata)
+                builder.from_usage_metadata(usage_metadata)
+            else:
+                # Fallback to response_metadata
+                response_metadata = getattr(output, "response_metadata", None)
+                if response_metadata:
+                    builder.from_response_metadata(response_metadata)
+
+                # Last resort: estimate from content
+                content = getattr(output, "content", "")
+                if isinstance(content, str):
+                    builder.estimate_from_content(content)
+
+        consumption_estimate = builder.build()
+
+        if consumption_estimate:
+            cache_info = ""
+            if consumption_estimate.token_usage.cache_creation_tokens > 0 or consumption_estimate.token_usage.cache_read_tokens > 0:
+                cache_info = (
+                    f", cache_create: {consumption_estimate.token_usage.cache_creation_tokens}, "
+                    f"cache_read: {consumption_estimate.token_usage.cache_read_tokens}"
+                )
+
+            logger.info(
+                f"📊 Consumption tracking - "
+                f"Tokens: {consumption_estimate.token_usage.total_tokens} "
+                f"(in: {consumption_estimate.token_usage.input_tokens}, out: {consumption_estimate.token_usage.output_tokens}, "
+                f"reasoning: {consumption_estimate.token_usage.reasoning_tokens}{cache_info})"
+            )
+            logger.debug(f"📊 Full consumption data: {consumption_estimate.to_dict()}")
+
+        return consumption_estimate

spaik_sdk/llm/consumption/token_usage.py

@@ -0,0 +1,31 @@
+from dataclasses import dataclass
+from typing import Any, Mapping
+
+
+@dataclass
+class TokenUsage:
+    input_tokens: int = 0
+    output_tokens: int = 0
+    total_tokens: int = 0
+    reasoning_tokens: int = 0
+    cache_creation_tokens: int = 0
+    cache_read_tokens: int = 0
+
+    def __post_init__(self):
+        if self.total_tokens == 0:
+            self.total_tokens = self.input_tokens + self.output_tokens
+
+    @classmethod
+    def from_langchain(cls, usage: Mapping[str, Any]) -> "TokenUsage":
+        """Create TokenUsage from LangChain usage_metadata."""
+        output_details = usage.get("output_token_details", {})
+        input_details = usage.get("input_token_details", {})
+
+        return cls(
+            input_tokens=usage.get("input_tokens", 0),
+            output_tokens=usage.get("output_tokens", 0),
+            total_tokens=usage.get("total_tokens", 0),
+            reasoning_tokens=output_details.get("reasoning", 0) if isinstance(output_details, dict) else 0,
+            cache_creation_tokens=input_details.get("cache_creation", 0) if isinstance(input_details, dict) else 0,
+            cache_read_tokens=input_details.get("cache_read", 0) if isinstance(input_details, dict) else 0,
+        )

spaik_sdk/llm/converters.py

@@ -0,0 +1,146 @@
+import base64
+from typing import Any, Dict, List, Optional, Union, cast
+
+from langchain_core.messages import AIMessage, BaseMessage, HumanMessage
+
+from spaik_sdk.attachments.models import Attachment
+from spaik_sdk.attachments.provider_support import is_supported_by_provider
+from spaik_sdk.attachments.storage.base_file_storage import BaseFileStorage
+from spaik_sdk.thread.models import MessageBlock, MessageBlockType, ThreadMessage
+from spaik_sdk.utils.init_logger import init_logger
+
+logger = init_logger(__name__)
+
+
+def convert_thread_message_to_langchain(
+    thread_message: ThreadMessage,
+    include_author_name: bool = False,
+) -> BaseMessage:
+    processed_blocks = [_process_message_block(block) for block in thread_message.blocks]
+    content = "\n".join(filter(None, processed_blocks))
+    author_prefix = f"{'[' + thread_message.author_name + ']: ' if include_author_name else ''}"
+    if thread_message.ai:
+        return AIMessage(content=f"{author_prefix}{content}")
+    else:
+        return HumanMessage(content=f"{author_prefix}{content}")
+
+
+async def convert_thread_message_to_langchain_multimodal(
+    thread_message: ThreadMessage,
+    file_storage: BaseFileStorage,
+    provider_family: str = "openai",
+    include_author_name: bool = False,
+) -> BaseMessage:
+    processed_blocks = [_process_message_block(block) for block in thread_message.blocks]
+    text_content = "\n".join(filter(None, processed_blocks))
+    author_prefix = f"{'[' + thread_message.author_name + ']: ' if include_author_name else ''}"
+
+    if thread_message.ai:
+        return AIMessage(content=f"{author_prefix}{text_content}")
+
+    if not thread_message.attachments:
+        return HumanMessage(content=f"{author_prefix}{text_content}")
+
+    content_parts: List[Dict[str, Any]] = []
+
+    if text_content:
+        content_parts.append({"type": "text", "text": f"{author_prefix}{text_content}"})
+
+    for attachment in thread_message.attachments:
+        # Check if MIME type is supported by the provider
+        if not is_supported_by_provider(attachment.mime_type, provider_family):
+            logger.warning(f"Skipping unsupported attachment type {attachment.mime_type} for provider {provider_family}")
+            filename = attachment.filename or "unnamed file"
+            fallback_msg = (
+                f"[SYSTEM: The user attempted to attach a file '{filename}' of type '{attachment.mime_type}', "
+                f"but this file type is not supported by the current AI model. "
+                f"Please inform the user that you cannot process this file type "
+                f"and suggest supported alternatives like images (PNG, JPEG, GIF, WEBP) or PDF documents.]"
+            )
+            content_parts.append({"type": "text", "text": fallback_msg})
+            continue
+
+        attachment_content = await _convert_attachment_to_content_part(attachment, file_storage, provider_family)
+        if attachment_content:
+            content_parts.append(attachment_content)
+
+    if not content_parts:
+        return HumanMessage(content="")
+
+    return HumanMessage(content=cast(List[Union[str, Dict[str, Any]]], content_parts))
+
+
+async def _convert_attachment_to_content_part(
+    attachment: Attachment,
+    file_storage: BaseFileStorage,
+    provider_family: str = "openai",
+) -> Optional[Dict[str, Any]]:
+    try:
+        data, metadata = await file_storage.retrieve(attachment.file_id)
+    except FileNotFoundError:
+        return None
+
+    mime_type = attachment.mime_type
+    b64_data = base64.b64encode(data).decode("utf-8")
+
+    if mime_type.startswith("image/"):
+        if provider_family == "anthropic":
+            return {
+                "type": "image",
+                "source": {"type": "base64", "media_type": mime_type, "data": b64_data},
+            }
+        return {
+            "type": "image_url",
+            "image_url": {"url": f"data:{mime_type};base64,{b64_data}"},
+        }
+    elif mime_type == "application/pdf":
+        if provider_family == "anthropic":
+            return {
+                "type": "document",
+                "source": {"type": "base64", "media_type": mime_type, "data": b64_data},
+            }
+        return {
+            "type": "image_url",
+            "image_url": {"url": f"data:{mime_type};base64,{b64_data}"},
+        }
+    elif mime_type.startswith("audio/"):
+        return {
+            "type": "input_audio",
+            "input_audio": {"data": b64_data, "format": _get_audio_format(mime_type)},
+        }
+    elif mime_type.startswith("video/"):
+        return {
+            "type": "video_url",
+            "video_url": {"url": f"data:{mime_type};base64,{b64_data}"},
+        }
+    else:
+        return {
+            "type": "file",
+            "file": {"data": b64_data, "mime_type": mime_type},
+        }
+
+
+def _get_audio_format(mime_type: str) -> str:
+    format_map = {
+        "audio/mpeg": "mp3",
+        "audio/mp3": "mp3",
+        "audio/wav": "wav",
+        "audio/ogg": "ogg",
+        "audio/webm": "webm",
+    }
+    return format_map.get(mime_type, "wav")
+
+
+def _process_message_block(block: MessageBlock) -> str:
+    if block.type == MessageBlockType.REASONING:
+        return "<thinking/>"
+    elif block.type == MessageBlockType.TOOL_USE:
+        tool_name = block.tool_name or "unknown"
+        return f'<tool_call tool="{tool_name}"/>'
+    elif block.type == MessageBlockType.ERROR:
+        content = block.content or "unknown error"
+        return f'<error msg="{content}"/>'
+    elif block.type == MessageBlockType.PLAIN:
+        return block.content or ""
+    else:
+        return block.content or ""

spaik_sdk/llm/cost/__init__.py

@@ -0,0 +1 @@
+