synkro 0.4.5__py3-none-any.whl

synkro/generation/tool_simulator.py

@@ -0,0 +1,114 @@
+"""Tool response simulator for training data generation."""
+
+import json
+import uuid
+from typing import TYPE_CHECKING
+
+from synkro.prompts.tool_templates import TOOL_SIMULATION_PROMPT
+
+if TYPE_CHECKING:
+    from synkro.llm.client import LLM
+    from synkro.types.tool import ToolDefinition, ToolCall
+
+
+class ToolSimulator:
+    """
+    Simulates tool responses for training data generation.
+    
+    Uses an LLM to generate realistic, contextual tool responses
+    based on tool definitions and call arguments.
+    
+    Example:
+        >>> from synkro.types.tool import ToolDefinition, ToolCall, ToolFunction
+        >>> simulator = ToolSimulator(tools=[web_search_tool], llm=llm)
+        >>> call = ToolCall(
+        ...     id="call_1",
+        ...     function=ToolFunction(name="web_search", arguments='{"query": "weather NYC"}')
+        ... )
+        >>> response = await simulator.simulate(call)
+        >>> print(response)
+        "NYC: 72°F, sunny with a high of 75°F expected"
+    """
+    
+    def __init__(self, tools: list["ToolDefinition"], llm: "LLM"):
+        """
+        Initialize the simulator.
+        
+        Args:
+            tools: List of available tool definitions
+            llm: LLM client for generating responses
+        """
+        self.tools = {t.name: t for t in tools}
+        self.llm = llm
+    
+    async def simulate(self, tool_call: "ToolCall") -> str:
+        """
+        Simulate a tool response for the given call.
+        
+        Args:
+            tool_call: The tool call to simulate
+            
+        Returns:
+            Simulated tool response content
+        """
+        tool_name = tool_call.function.name
+        
+        if tool_name not in self.tools:
+            return json.dumps({"error": f"Unknown tool: {tool_name}"})
+        
+        tool = self.tools[tool_name]
+        
+        # Format mock responses for the prompt
+        mock_responses = "\n".join(
+            f"- {r}" for r in tool.mock_responses
+        ) if tool.mock_responses else "No example responses provided"
+        
+        prompt = TOOL_SIMULATION_PROMPT.format(
+            TOOL_NAME=tool.name,
+            TOOL_DESCRIPTION=tool.description,
+            TOOL_PARAMETERS=json.dumps(tool.parameters, indent=2),
+            ARGUMENTS=tool_call.function.arguments,
+            MOCK_RESPONSES=mock_responses,
+        )
+        
+        response = await self.llm.generate(prompt)
+        return response.strip()
+    
+    async def simulate_batch(self, tool_calls: list["ToolCall"]) -> list[str]:
+        """
+        Simulate responses for multiple tool calls.
+        
+        Args:
+            tool_calls: List of tool calls to simulate
+            
+        Returns:
+            List of simulated responses in order
+        """
+        import asyncio
+        return await asyncio.gather(*[self.simulate(tc) for tc in tool_calls])
+    
+    def generate_call_id(self) -> str:
+        """Generate a unique tool call ID."""
+        return f"call_{uuid.uuid4().hex[:12]}"
+    
+    def get_tools_description(self) -> str:
+        """
+        Get a formatted description of all available tools.
+        
+        Returns:
+            Formatted string describing all tools
+        """
+        descriptions = []
+        for tool in self.tools.values():
+            descriptions.append(tool.to_system_prompt())
+        return "\n\n".join(descriptions)
+    
+    def get_tools_json(self) -> list[dict]:
+        """
+        Get tools in OpenAI function format.
+        
+        Returns:
+            List of tool definitions in OpenAI format
+        """
+        return [tool.to_openai_format() for tool in self.tools.values()]
+

synkro/llm/__init__.py

@@ -0,0 +1,7 @@
+"""LLM client wrapper for multiple providers via LiteLLM."""
+
+from synkro.llm.client import LLM
+from synkro.llm.rate_limits import auto_workers, get_provider
+
+__all__ = ["LLM", "auto_workers", "get_provider"]
+

synkro/llm/client.py

@@ -0,0 +1,235 @@
+"""Type-safe LLM wrapper using LiteLLM."""
+
+from typing import TypeVar, Type, overload
+
+import litellm
+from litellm import acompletion, supports_response_schema
+from pydantic import BaseModel
+
+# Configure litellm
+litellm.suppress_debug_info = True
+litellm.enable_json_schema_validation=True
+
+from synkro.models import OpenAI, Model, get_model_string
+
+
+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)
+
+    Examples:
+        >>> from synkro import LLM, OpenAI, Anthropic, Google
+
+        # 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)
+
+        # 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,
+    ):
+        """
+        Initialize the LLM client.
+
+        Args:
+            model: Model to use (enum 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
+        """
+        self.model = get_model_string(model)
+        self.temperature = temperature
+        self.max_tokens = max_tokens
+        self._api_key = api_key
+
+    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
+
+        response = await acompletion(**kwargs)
+        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
+
+        response = await acompletion(**kwargs)
+        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
+
+            response = await acompletion(**kwargs)
+            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
+
+        response = await acompletion(**kwargs)
+        return response.choices[0].message.content
+

synkro/llm/rate_limits.py

@@ -0,0 +1,95 @@
+"""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)
+}
+
+# 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
+}
+
+
+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,43 @@
+"""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)
+
+Usage:
+    # Per-provider import (recommended)
+    from synkro.models.openai import OpenAI
+    from synkro.models.anthropic import Anthropic
+    from synkro.models.google import Google
+
+    # Convenience import (all at once)
+    from synkro.models import OpenAI, Anthropic, Google
+"""
+
+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
+
+# Union type for any model
+Model = Union[OpenAI, Anthropic, Google, 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
+    return model
+
+
+__all__ = [
+    "OpenAI",
+    "Anthropic",
+    "Google",
+    "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/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,15 @@
+"""Mode configurations for different dataset types."""
+
+from synkro.modes.config import ModeConfig, get_mode_config
+from synkro.modes.qa import QA_CONFIG
+from synkro.modes.sft import SFT_CONFIG
+from synkro.modes.tool_call import TOOL_CALL_CONFIG
+
+__all__ = [
+    "ModeConfig",
+    "get_mode_config",
+    "QA_CONFIG",
+    "SFT_CONFIG",
+    "TOOL_CALL_CONFIG",
+]
+

synkro/modes/config.py

@@ -0,0 +1,66 @@
+"""Mode configuration that bundles prompts, schema, and formatter per dataset type."""
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Type
+
+if TYPE_CHECKING:
+    from synkro.types.dataset_type import DatasetType
+
+
+@dataclass
+class ModeConfig:
+    """
+    Configuration bundle for a dataset type.
+
+    Defines all the prompts, schemas, and formatters needed
+    for generating a specific type of dataset.
+    """
+
+    # Prompts
+    scenario_prompt: str
+    """Prompt for generating scenarios/questions"""
+
+    response_prompt: str
+    """Prompt for generating responses/answers"""
+
+    grade_prompt: str
+    """Prompt for grading quality"""
+
+    refine_prompt: str
+    """Prompt for refining failed responses"""
+
+    # Output configuration
+    output_description: str
+    """Human-readable description of output format"""
+
+
+def get_mode_config(dataset_type: "DatasetType") -> ModeConfig:
+    """
+    Get the mode configuration for a dataset type.
+
+    Args:
+        dataset_type: The type of dataset to generate
+
+    Returns:
+        ModeConfig with appropriate prompts and settings
+
+    Example:
+        >>> from synkro import DatasetType
+        >>> config = get_mode_config(DatasetType.QA)
+    """
+    from synkro.types.dataset_type import DatasetType
+    from synkro.modes.qa import QA_CONFIG
+    from synkro.modes.sft import SFT_CONFIG
+    from synkro.modes.tool_call import TOOL_CALL_CONFIG
+
+    configs = {
+        DatasetType.QA: QA_CONFIG,
+        DatasetType.SFT: SFT_CONFIG,
+        DatasetType.TOOL_CALL: TOOL_CALL_CONFIG,
+    }
+
+    if dataset_type not in configs:
+        raise ValueError(f"Unknown dataset type: {dataset_type}")
+
+    return configs[dataset_type]
+

synkro/modes/qa.py

@@ -0,0 +1,18 @@
+"""QA mode configuration."""
+
+from synkro.modes.config import ModeConfig
+from synkro.prompts.qa_templates import (
+    QA_SCENARIO_PROMPT,
+    QA_RESPONSE_PROMPT,
+    QA_GRADE_PROMPT,
+    QA_REFINE_PROMPT,
+)
+
+QA_CONFIG = ModeConfig(
+    scenario_prompt=QA_SCENARIO_PROMPT,
+    response_prompt=QA_RESPONSE_PROMPT,
+    grade_prompt=QA_GRADE_PROMPT,
+    refine_prompt=QA_REFINE_PROMPT,
+    output_description="Question-Answer pairs: {question, answer, context}",
+)
+

synkro/modes/sft.py

@@ -0,0 +1,18 @@
+"""SFT mode configuration."""
+
+from synkro.modes.config import ModeConfig
+from synkro.prompts.templates import (
+    SCENARIO_GENERATOR_PROMPT,
+    SINGLE_RESPONSE_PROMPT,
+    SINGLE_GRADE_PROMPT,
+    BATCHED_REFINER_PROMPT,
+)
+
+SFT_CONFIG = ModeConfig(
+    scenario_prompt=SCENARIO_GENERATOR_PROMPT,
+    response_prompt=SINGLE_RESPONSE_PROMPT,
+    grade_prompt=SINGLE_GRADE_PROMPT,
+    refine_prompt=BATCHED_REFINER_PROMPT,
+    output_description="Chat messages: {messages: [system, user, assistant]}",
+)
+