synkro 0.4.36__py3-none-any.whl

synkro/llm/client.py

@@ -0,0 +1,309 @@
+"""Type-safe LLM wrapper using LiteLLM."""
+
+import warnings
+from typing import TypeVar, Type, overload
+
+import litellm
+from litellm import acompletion, supports_response_schema, completion_cost
+from pydantic import BaseModel
+
+# Configure litellm
+litellm.suppress_debug_info = True
+litellm.enable_json_schema_validation = True
+litellm.drop_params = True  # Drop unsupported params (e.g., temperature for gpt-5)
+
+# Suppress Pydantic serialization warnings from litellm response types
+warnings.filterwarnings(
+    "ignore",
+    message=".*Pydantic serializer warnings.*",
+    category=UserWarning,
+)
+warnings.filterwarnings(
+    "ignore",
+    message=".*Expected `Message`.*",
+    category=UserWarning,
+)
+warnings.filterwarnings(
+    "ignore",
+    message=".*Expected `StreamingChoices`.*",
+    category=UserWarning,
+)
+
+from synkro.models import OpenAI, Model, get_model_string, LocalModel
+
+
+T = TypeVar("T", bound=BaseModel)
+
+
+class LLM:
+    """
+    Type-safe LLM wrapper using LiteLLM for universal provider support.
+
+    Supports structured outputs via native JSON mode for reliable responses.
+
+    Supported providers: OpenAI, Anthropic, Google (Gemini), Local (Ollama, vLLM)
+
+    Examples:
+        >>> from synkro import LLM, OpenAI, Anthropic, Google, Local
+
+        # Use OpenAI
+        >>> llm = LLM(model=OpenAI.GPT_4O_MINI)
+        >>> response = await llm.generate("Hello!")
+
+        # Use Anthropic
+        >>> llm = LLM(model=Anthropic.CLAUDE_35_SONNET)
+
+        # Use Google Gemini
+        >>> llm = LLM(model=Google.GEMINI_25_FLASH)
+
+        # Use local Ollama
+        >>> llm = LLM(model=Local.OLLAMA("llama3.1"))
+
+        # Use local vLLM
+        >>> llm = LLM(model=Local.VLLM("mistral"))
+
+        # Structured output
+        >>> class Output(BaseModel):
+        ...     answer: str
+        ...     confidence: float
+        >>> result = await llm.generate_structured("What is 2+2?", Output)
+        >>> result.answer
+        '4'
+    """
+
+    def __init__(
+        self,
+        model: Model = OpenAI.GPT_4O_MINI,
+        temperature: float = 0.7,
+        max_tokens: int | None = None,
+        api_key: str | None = None,
+        base_url: str | None = None,
+    ):
+        """
+        Initialize the LLM client.
+
+        Args:
+            model: Model to use (enum, LocalModel, or string)
+            temperature: Sampling temperature (0.0-2.0)
+            max_tokens: Maximum tokens to generate (default: None = model's max)
+            api_key: Optional API key override
+            base_url: Optional API base URL (auto-set when using Local models)
+        """
+        # Handle LocalModel - extract endpoint automatically
+        if isinstance(model, LocalModel):
+            self.model = f"{model.provider}/{model.model}"
+            self._base_url = model.endpoint
+        else:
+            self.model = get_model_string(model)
+            self._base_url = base_url
+
+        self.temperature = temperature
+        self.max_tokens = max_tokens
+        self._api_key = api_key
+
+        # Cost and usage tracking
+        self._total_cost = 0.0
+        self._call_count = 0
+
+    async def generate(self, prompt: str, system: str | None = None) -> str:
+        """
+        Generate a text response.
+
+        Args:
+            prompt: The user prompt
+            system: Optional system prompt
+
+        Returns:
+            Generated text response
+        """
+        messages = []
+        if system:
+            messages.append({"role": "system", "content": system})
+        messages.append({"role": "user", "content": prompt})
+
+        kwargs = {
+            "model": self.model,
+            "messages": messages,
+            "temperature": self.temperature,
+            "api_key": self._api_key,
+        }
+        if self.max_tokens is not None:
+            kwargs["max_tokens"] = self.max_tokens
+        if self._base_url:
+            kwargs["api_base"] = self._base_url
+
+        response = await acompletion(**kwargs)
+        self._track_cost(response)
+        return response.choices[0].message.content
+
+    async def generate_batch(
+        self, prompts: list[str], system: str | None = None
+    ) -> list[str]:
+        """
+        Generate responses for multiple prompts in parallel.
+
+        Args:
+            prompts: List of user prompts
+            system: Optional system prompt for all
+
+        Returns:
+            List of generated responses
+        """
+        import asyncio
+
+        tasks = [self.generate(p, system) for p in prompts]
+        return await asyncio.gather(*tasks)
+
+    @overload
+    async def generate_structured(
+        self,
+        prompt: str,
+        response_model: Type[T],
+        system: str | None = None,
+    ) -> T: ...
+
+    @overload
+    async def generate_structured(
+        self,
+        prompt: str,
+        response_model: Type[list[T]],
+        system: str | None = None,
+    ) -> list[T]: ...
+
+    async def generate_structured(
+        self,
+        prompt: str,
+        response_model: Type[T] | Type[list[T]],
+        system: str | None = None,
+    ) -> T | list[T]:
+        """
+        Generate a structured response matching a Pydantic model.
+
+        Uses LiteLLM's native JSON mode with response_format for
+        reliable structured outputs.
+
+        Args:
+            prompt: The user prompt
+            response_model: Pydantic model class for the response
+            system: Optional system prompt
+
+        Returns:
+            Parsed response matching the model
+
+        Example:
+            >>> class Analysis(BaseModel):
+            ...     sentiment: str
+            ...     score: float
+            >>> result = await llm.generate_structured(
+            ...     "Analyze: I love this product!",
+            ...     Analysis
+            ... )
+            >>> result.sentiment
+            'positive'
+        """
+        # Check if model supports structured outputs
+        if not supports_response_schema(model=self.model, custom_llm_provider=None):
+            raise ValueError(
+                f"Model '{self.model}' does not support structured outputs (response_format). "
+                f"Use a model that supports JSON schema like GPT-4o, Gemini 1.5+, or Claude 3.5+."
+            )
+
+        messages = []
+        if system:
+            messages.append({"role": "system", "content": system})
+        messages.append({"role": "user", "content": prompt})
+
+        # Use LiteLLM's native response_format with Pydantic model
+        kwargs = {
+            "model": self.model,
+            "messages": messages,
+            "response_format": response_model,
+            "temperature": self.temperature,
+            "api_key": self._api_key,
+        }
+        if self.max_tokens is not None:
+            kwargs["max_tokens"] = self.max_tokens
+        if self._base_url:
+            kwargs["api_base"] = self._base_url
+
+        response = await acompletion(**kwargs)
+        self._track_cost(response)
+        return response_model.model_validate_json(response.choices[0].message.content)
+
+    async def generate_chat(
+        self, messages: list[dict], response_model: Type[T] | None = None
+    ) -> str | T:
+        """
+        Generate a response for a full conversation.
+
+        Args:
+            messages: List of message dicts with 'role' and 'content'
+            response_model: Optional Pydantic model for structured output
+
+        Returns:
+            Generated response (string or structured)
+        """
+        if response_model:
+            # Check if model supports structured outputs
+            if not supports_response_schema(model=self.model, custom_llm_provider=None):
+                raise ValueError(
+                    f"Model '{self.model}' does not support structured outputs (response_format). "
+                    f"Use a model that supports JSON schema like GPT-4o, Gemini 1.5+, or Claude 3.5+."
+                )
+
+            # Use LiteLLM's native response_format with Pydantic model
+            kwargs = {
+                "model": self.model,
+                "messages": messages,
+                "response_format": response_model,
+                "temperature": self.temperature,
+                "api_key": self._api_key,
+            }
+            if self.max_tokens is not None:
+                kwargs["max_tokens"] = self.max_tokens
+            if self._base_url:
+                kwargs["api_base"] = self._base_url
+
+            response = await acompletion(**kwargs)
+            self._track_cost(response)
+            return response_model.model_validate_json(response.choices[0].message.content)
+
+        kwargs = {
+            "model": self.model,
+            "messages": messages,
+            "temperature": self.temperature,
+            "api_key": self._api_key,
+        }
+        if self.max_tokens is not None:
+            kwargs["max_tokens"] = self.max_tokens
+        if self._base_url:
+            kwargs["api_base"] = self._base_url
+
+        response = await acompletion(**kwargs)
+        self._track_cost(response)
+        return response.choices[0].message.content
+
+    def _track_cost(self, response) -> None:
+        """Track cost and call count from a response."""
+        self._call_count += 1
+        try:
+            cost = completion_cost(completion_response=response)
+            self._total_cost += cost
+        except Exception:
+            # Some models may not have pricing info
+            pass
+
+    @property
+    def total_cost(self) -> float:
+        """Get total cost of all LLM calls made by this client."""
+        return self._total_cost
+
+    @property
+    def call_count(self) -> int:
+        """Get total number of LLM calls made by this client."""
+        return self._call_count
+
+    def reset_tracking(self) -> None:
+        """Reset cost and call tracking."""
+        self._total_cost = 0.0
+        self._call_count = 0

synkro/llm/rate_limits.py

@@ -0,0 +1,99 @@
+"""Automatic worker scaling based on provider rate limits."""
+
+# Known rate limits per provider (requests per minute)
+PROVIDER_RATE_LIMITS = {
+    "openai": 60,  # Tier 1 default, scales with tier
+    "anthropic": 60,  # Standard limit
+    "google": 60,  # Gemini API
+    "gemini": 60,  # Gemini API (alternative prefix)
+    "ollama": 1000,  # Local - no real limit
+    "vllm": 1000,  # Local - no real limit
+}
+
+# Target 80% of rate limit to avoid hitting caps
+UTILIZATION_TARGET = 0.8
+
+# Default workers per provider (pre-computed for convenience)
+DEFAULT_WORKERS = {
+    "openai": 15,  # ~60 RPM / 3 calls = 20, use 15 to be safe
+    "anthropic": 10,  # ~60 RPM, more conservative
+    "google": 15,  # Gemini
+    "gemini": 15,  # Gemini
+    "ollama": 50,  # Local - high parallelism
+    "vllm": 50,  # Local - high parallelism
+}
+
+
+def get_provider(model: str) -> str:
+    """
+    Extract provider name from model string.
+
+    Args:
+        model: Model string like "gpt-4o" or "ollama/llama3.1:8b"
+
+    Returns:
+        Provider name
+    """
+    # Check for explicit prefix
+    if "/" in model:
+        return model.split("/")[0]
+
+    # Infer from model name
+    if model.startswith("gpt") or model.startswith("o1"):
+        return "openai"
+    if model.startswith("claude"):
+        return "anthropic"
+    if model.startswith("gemini"):
+        return "google"
+
+    return "openai"  # Default
+
+
+def auto_workers(model: str) -> int:
+    """
+    Determine optimal worker count based on model's provider.
+
+    This calculates a safe default that won't hit rate limits,
+    accounting for the fact that each trace needs ~3 LLM calls
+    (generate, grade, maybe refine).
+
+    Args:
+        model: Model string
+
+    Returns:
+        Recommended worker count
+
+    Example:
+        >>> auto_workers("gpt-4o")
+        15
+        >>> auto_workers("gemini/gemini-2.5-flash")
+        15
+    """
+    provider = get_provider(model)
+    rpm = PROVIDER_RATE_LIMITS.get(provider, 60)
+
+    # Workers = RPM * utilization / avg_calls_per_trace
+    # Each trace needs ~3 calls (generate, grade, maybe refine)
+    avg_calls_per_trace = 3
+
+    workers = int((rpm * UTILIZATION_TARGET) / avg_calls_per_trace)
+
+    # Clamp to reasonable bounds
+    return max(5, min(workers, 100))
+
+
+def get_default_workers(model: str) -> int:
+    """
+    Quick lookup for worker count.
+
+    Uses pre-computed defaults for common providers.
+
+    Args:
+        model: Model string
+
+    Returns:
+        Default worker count for the provider
+    """
+    provider = get_provider(model)
+    return DEFAULT_WORKERS.get(provider, 10)
+

synkro/models/__init__.py

@@ -0,0 +1,50 @@
+"""Model enums for supported LLM providers.
+
+Supported providers:
+- OpenAI (GPT-4o, GPT-4o-mini)
+- Anthropic (Claude 3.5 Sonnet/Haiku)
+- Google (Gemini 2.5 Flash/Pro)
+- Local (Ollama, vLLM, custom)
+
+Usage:
+    # Per-provider import (recommended)
+    from synkro.models.openai import OpenAI
+    from synkro.models.anthropic import Anthropic
+    from synkro.models.google import Google
+    from synkro.models.local import Local
+
+    # Convenience import (all at once)
+    from synkro.models import OpenAI, Anthropic, Google, Local
+"""
+
+from enum import Enum
+from typing import Union
+
+from synkro.models.openai import OpenAI
+from synkro.models.anthropic import Anthropic
+from synkro.models.google import Google
+from synkro.models.local import Local, LocalModel
+
+# Union type for any model
+Model = Union[OpenAI, Anthropic, Google, LocalModel, str]
+
+
+def get_model_string(model: Model) -> str:
+    """Convert a model enum or string to its string value."""
+    if isinstance(model, Enum):
+        return model.value
+    if isinstance(model, LocalModel):
+        return str(model)
+    return model
+
+
+__all__ = [
+    "OpenAI",
+    "Anthropic",
+    "Google",
+    "Local",
+    "LocalModel",
+    "Model",
+    "get_model_string",
+]
+

synkro/models/anthropic.py

@@ -0,0 +1,26 @@
+"""Anthropic Claude models."""
+
+from enum import Enum
+
+
+class Anthropic(str, Enum):
+    """Anthropic Claude models."""
+
+    # Claude 4.5 (latest)
+    CLAUDE_45_OPUS = "claude-opus-4-5-20250601"
+    """Premium: State-of-the-art for coding and autonomous agents"""
+
+    CLAUDE_45_SONNET = "claude-sonnet-4-5-20250601"
+    """Standard: Default model for most users, faster and more context-aware"""
+
+    CLAUDE_45_HAIKU = "claude-haiku-4-5-20250601"
+    """Light: High-speed, cost-effective, matches Claude 3 Opus intelligence"""
+
+    # Claude 4 (previous gen)
+    CLAUDE_4_SONNET = "claude-sonnet-4-20250514"
+    CLAUDE_4_OPUS = "claude-opus-4-20250514"
+
+    # Claude 3.5 (legacy)
+    CLAUDE_35_SONNET = "claude-3-5-sonnet-20241022"
+    CLAUDE_35_HAIKU = "claude-3-5-haiku-20241022"
+

synkro/models/google.py

@@ -0,0 +1,19 @@
+"""Google Gemini models.
+
+Updated based on: https://ai.google.dev/gemini-api/docs/models#model-versions
+"""
+
+from enum import Enum
+
+
+class Google(str, Enum):
+    """Google Gemini models."""
+
+    GEMINI_3_PRO = "gemini/gemini-3-pro"
+    GEMINI_3_FLASH = "gemini/gemini-3-flash"
+
+    GEMINI_25_FLASH = "gemini/gemini-2.5-flash"
+    GEMINI_25_PRO = "gemini/gemini-2.5-pro"
+
+    GEMINI_2_FLASH = "gemini/gemini-2.0-flash"
+    GEMINI_2_FLASH_LITE = "gemini/gemini-2.0-flash-lite"

synkro/models/local.py

@@ -0,0 +1,104 @@
+"""Local LLM providers (Ollama, vLLM, etc.)."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass
+class LocalModel:
+    """Represents a local model configuration.
+
+    Attributes:
+        provider: The provider name (ollama, vllm, openai)
+        model: The model name
+        endpoint: The API endpoint URL
+    """
+
+    provider: str
+    model: str
+    endpoint: str
+
+    def __str__(self) -> str:
+        return f"{self.provider}/{self.model}"
+
+
+class Local:
+    """Factory for local LLM configurations.
+
+    Provides a clean API for configuring local LLM providers like Ollama and vLLM.
+    Returns LocalModel instances that the LLM client can use to configure the
+    connection automatically.
+
+    Examples:
+        >>> from synkro import LLM, Local
+
+        # Ollama (default localhost:11434)
+        >>> llm = LLM(model=Local.OLLAMA("llama3.1"))
+
+        # vLLM (default localhost:8000)
+        >>> llm = LLM(model=Local.VLLM("mistral"))
+
+        # Custom endpoint
+        >>> llm = LLM(model=Local.OLLAMA("llama3.1", endpoint="http://server:11434"))
+
+        # Any OpenAI-compatible server
+        >>> llm = LLM(model=Local.CUSTOM("my-model", endpoint="http://localhost:8080/v1"))
+    """
+
+    DEFAULT_ENDPOINTS = {
+        "ollama": "http://localhost:11434",
+        "vllm": "http://localhost:8000",
+        "openai": "http://localhost:8000/v1",
+    }
+
+    @classmethod
+    def OLLAMA(cls, model: str, endpoint: str | None = None) -> LocalModel:
+        """Create Ollama model config.
+
+        Args:
+            model: Model name (e.g., "llama3.1", "mistral", "codellama")
+            endpoint: Optional custom endpoint (default: http://localhost:11434)
+
+        Returns:
+            LocalModel configured for Ollama
+        """
+        return LocalModel(
+            provider="ollama",
+            model=model,
+            endpoint=endpoint or cls.DEFAULT_ENDPOINTS["ollama"],
+        )
+
+    @classmethod
+    def VLLM(cls, model: str, endpoint: str | None = None) -> LocalModel:
+        """Create vLLM model config.
+
+        Args:
+            model: Model name (e.g., "mistral-7b", "llama-2-13b")
+            endpoint: Optional custom endpoint (default: http://localhost:8000)
+
+        Returns:
+            LocalModel configured for vLLM
+        """
+        return LocalModel(
+            provider="vllm",
+            model=model,
+            endpoint=endpoint or cls.DEFAULT_ENDPOINTS["vllm"],
+        )
+
+    @classmethod
+    def CUSTOM(cls, model: str, endpoint: str) -> LocalModel:
+        """Create custom local model config for any OpenAI-compatible server.
+
+        Args:
+            model: Model name
+            endpoint: API endpoint URL (required)
+
+        Returns:
+            LocalModel configured for OpenAI-compatible API
+        """
+        return LocalModel(
+            provider="openai",
+            model=model,
+            endpoint=endpoint,
+        )

synkro/models/openai.py

@@ -0,0 +1,31 @@
+"""OpenAI models."""
+
+from enum import Enum
+
+
+class OpenAI(str, Enum):
+    """OpenAI models."""
+
+    # GPT-5 series (latest)
+    GPT_52 = "gpt-5.2"
+    """Flagship: High-speed, human-like dialogue, agentic tool-calling"""
+
+    GPT_5_MINI = "gpt-5-mini"
+    """Mid-tier: Balanced cost and intelligence, primary workhorse"""
+
+    GPT_5_NANO = "gpt-5-nano"
+    """Edge: Extremely low latency, high-volume basic tasks"""
+
+    # GPT-4 series (legacy)
+    GPT_41 = "gpt-4.1"
+    """Legacy flagship: Smartest non-reasoning model from previous gen"""
+
+    GPT_4O = "gpt-4o"
+    GPT_4O_MINI = "gpt-4o-mini"
+
+    # Reasoning models
+    O3 = "o3"
+    O3_MINI = "o3-mini"
+    O1 = "o1"
+    O1_MINI = "o1-mini"
+

synkro/modes/__init__.py

@@ -0,0 +1,13 @@
+"""Mode configurations for different dataset types."""
+
+from synkro.modes.config import ModeConfig, get_mode_config
+from synkro.modes.conversation import CONVERSATION_CONFIG, INSTRUCTION_CONFIG
+from synkro.modes.tool_call import TOOL_CALL_CONFIG
+
+__all__ = [
+    "ModeConfig",
+    "get_mode_config",
+    "CONVERSATION_CONFIG",
+    "INSTRUCTION_CONFIG",
+    "TOOL_CALL_CONFIG",
+]