themis-eval 0.1.0__py3-none-any.whl → 0.1.1__py3-none-any.whl

themis/generation/plan.py

@@ -0,0 +1,456 @@
+"""Generation planning primitives.
+
+This module provides generation planning with flexible expansion strategies:
+
+1. GenerationPlan: Traditional Cartesian product expansion
+2. FlexibleGenerationPlan: Pluggable expansion strategies
+3. Expansion strategies:
+   - CartesianExpansionStrategy: Full Cartesian product (default)
+   - FilteredExpansionStrategy: Filter specific combinations
+   - ConditionalExpansionStrategy: Route based on conditions
+   - ChainedExpansionStrategy: Chain multiple strategies
+
+Example (Traditional):
+    >>> plan = GenerationPlan(
+    ...     templates=[template1, template2],
+    ...     models=[model1, model2],
+    ...     sampling_parameters=[config1]
+    ... )
+    >>> tasks = list(plan.expand(dataset))
+
+Example (Filtered):
+    >>> plan = FlexibleGenerationPlan(
+    ...     templates=[template1, template2],
+    ...     models=[model1, model2],
+    ...     sampling_parameters=[config1],
+    ...     expansion_strategy=FilteredExpansionStrategy(
+    ...         task_filter=lambda row, tpl, mdl, smp: mdl.identifier != "gpt-4" or row.get("difficulty") == "hard"
+    ...     )
+    ... )
+    >>> tasks = list(plan.expand(dataset))  # Only creates GPT-4 tasks for hard problems
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Callable, Dict, Iterator, Protocol, Sequence
+
+from themis.core import entities as core_entities
+from themis.generation import templates
+
+
+@dataclass
+class GenerationPlan:
+    templates: Sequence[templates.PromptTemplate]
+    models: Sequence[core_entities.ModelSpec]
+    sampling_parameters: Sequence[core_entities.SamplingConfig]
+    dataset_id_field: str = "id"
+    reference_field: str | None = "expected"
+    metadata_fields: Sequence[str] = field(default_factory=tuple)
+    context_builder: Callable[[dict[str, Any]], dict[str, Any]] | None = None
+
+    def expand(
+        self, dataset: Sequence[dict[str, object]]
+    ) -> Iterator[core_entities.GenerationTask]:
+        for row in dataset:
+            row_dict = dict(row)
+            context = self._build_context(row_dict)
+            dataset_id = row_dict.get(self.dataset_id_field)
+            reference = (
+                row_dict.get(self.reference_field) if self.reference_field else None
+            )
+            for template in self.templates:
+                rendered_prompt = template.render_prompt(context)
+                base_metadata = self._build_metadata(template, dataset_id, row_dict)
+                for model in self.models:
+                    for sampling in self.sampling_parameters:
+                        yield core_entities.GenerationTask(
+                            prompt=rendered_prompt,
+                            model=model,
+                            sampling=sampling,
+                            metadata=dict(base_metadata),
+                            reference=self._build_reference(reference),
+                        )
+
+    def _build_context(self, row: dict[str, Any]) -> dict[str, Any]:
+        if self.context_builder is None:
+            return dict(row)
+        return self.context_builder(dict(row))
+
+    def _build_metadata(
+        self,
+        template: templates.PromptTemplate,
+        dataset_id: Any,
+        row: dict[str, Any],
+    ) -> Dict[str, Any]:
+        metadata = {
+            f"template_{key}": value for key, value in (template.metadata or {}).items()
+        }
+        if dataset_id is not None:
+            metadata["dataset_id"] = dataset_id
+        for field_name in self.metadata_fields:
+            if field_name in row:
+                metadata[field_name] = row[field_name]
+        return metadata
+
+    def _build_reference(
+        self, raw_reference: Any | None
+    ) -> core_entities.Reference | None:
+        if raw_reference is None:
+            return None
+        return core_entities.Reference(
+            kind=self.reference_field or "reference", value=raw_reference
+        )
+
+
+# ============================================================================
+# Flexible Generation Planning with Expansion Strategies
+# ============================================================================
+
+
+@dataclass
+class PlanContext:
+    """Context passed to expansion strategies.
+
+    Contains all the information needed to expand dataset rows into tasks.
+    """
+
+    templates: Sequence[templates.PromptTemplate]
+    models: Sequence[core_entities.ModelSpec]
+    sampling_parameters: Sequence[core_entities.SamplingConfig]
+    dataset_id_field: str
+    reference_field: str | None
+    metadata_fields: Sequence[str]
+    context_builder: Callable[[dict], dict] | None
+
+
+class ExpansionStrategy(Protocol):
+    """Strategy for expanding dataset into generation tasks.
+
+    Different strategies can control which combinations of
+    (row, template, model, sampling) are generated.
+    """
+
+    def expand(
+        self,
+        dataset: Sequence[dict[str, Any]],
+        context: PlanContext,
+    ) -> Iterator[core_entities.GenerationTask]:
+        """Expand dataset rows into generation tasks.
+
+        Args:
+            dataset: Dataset rows to expand
+            context: Plan context with templates, models, etc.
+
+        Yields:
+            Generation tasks
+        """
+        ...
+
+
+class CartesianExpansionStrategy:
+    """Traditional Cartesian product expansion (default behavior).
+
+    Generates all possible combinations of:
+    - Each row in dataset
+    - Each template
+    - Each model
+    - Each sampling configuration
+
+    This is the default expansion strategy used by GenerationPlan.
+    """
+
+    def expand(
+        self,
+        dataset: Sequence[dict[str, Any]],
+        context: PlanContext,
+    ) -> Iterator[core_entities.GenerationTask]:
+        """Expand using Cartesian product."""
+        for row in dataset:
+            row_dict = dict(row)
+            ctx = (
+                context.context_builder(row_dict)
+                if context.context_builder
+                else row_dict
+            )
+            dataset_id = row_dict.get(context.dataset_id_field)
+            reference = (
+                row_dict.get(context.reference_field)
+                if context.reference_field
+                else None
+            )
+
+            for template in context.templates:
+                rendered = template.render_prompt(ctx)
+                base_metadata = self._build_metadata(
+                    template, dataset_id, row_dict, context
+                )
+
+                for model in context.models:
+                    for sampling in context.sampling_parameters:
+                        yield core_entities.GenerationTask(
+                            prompt=rendered,
+                            model=model,
+                            sampling=sampling,
+                            metadata=dict(base_metadata),
+                            reference=self._build_reference(reference, context),
+                        )
+
+    def _build_metadata(
+        self,
+        template: templates.PromptTemplate,
+        dataset_id: Any,
+        row: dict[str, Any],
+        context: PlanContext,
+    ) -> Dict[str, Any]:
+        """Build metadata dict for task."""
+        metadata = {
+            f"template_{key}": value for key, value in (template.metadata or {}).items()
+        }
+        if dataset_id is not None:
+            metadata["dataset_id"] = dataset_id
+        for field_name in context.metadata_fields:
+            if field_name in row:
+                metadata[field_name] = row[field_name]
+        return metadata
+
+    def _build_reference(
+        self, raw_reference: Any | None, context: PlanContext
+    ) -> core_entities.Reference | None:
+        """Build reference object."""
+        if raw_reference is None:
+            return None
+        return core_entities.Reference(
+            kind=context.reference_field or "reference", value=raw_reference
+        )
+
+
+class FilteredExpansionStrategy:
+    """Expansion strategy that filters specific combinations.
+
+    Only generates tasks that pass the filter function. Useful for:
+    - Expensive models only on hard problems
+    - Specific templates for specific models
+    - Conditional generation based on metadata
+
+    Example:
+        >>> # Only use GPT-4 on hard problems
+        >>> strategy = FilteredExpansionStrategy(
+        ...     task_filter=lambda row, tpl, mdl, smp: (
+        ...         mdl.identifier != "gpt-4" or row.get("difficulty") == "hard"
+        ...     )
+        ... )
+    """
+
+    def __init__(
+        self,
+        task_filter: Callable[
+            [
+                dict[str, Any],  # row
+                templates.PromptTemplate,  # template
+                core_entities.ModelSpec,  # model
+                core_entities.SamplingConfig,  # sampling
+            ],
+            bool,
+        ],
+    ):
+        """Initialize filtered expansion strategy.
+
+        Args:
+            task_filter: Function that returns True if task should be generated
+        """
+        self._filter = task_filter
+        self._base_strategy = CartesianExpansionStrategy()
+
+    def expand(
+        self,
+        dataset: Sequence[dict[str, Any]],
+        context: PlanContext,
+    ) -> Iterator[core_entities.GenerationTask]:
+        """Expand with filtering."""
+        for row in dataset:
+            row_dict = dict(row)
+            ctx = (
+                context.context_builder(row_dict)
+                if context.context_builder
+                else row_dict
+            )
+            dataset_id = row_dict.get(context.dataset_id_field)
+            reference = (
+                row_dict.get(context.reference_field)
+                if context.reference_field
+                else None
+            )
+
+            for template in context.templates:
+                rendered = template.render_prompt(ctx)
+                base_metadata = self._base_strategy._build_metadata(
+                    template, dataset_id, row_dict, context
+                )
+
+                for model in context.models:
+                    for sampling in context.sampling_parameters:
+                        # Check if this combination should be generated
+                        if self._filter(row_dict, template, model, sampling):
+                            yield core_entities.GenerationTask(
+                                prompt=rendered,
+                                model=model,
+                                sampling=sampling,
+                                metadata=dict(base_metadata),
+                                reference=self._base_strategy._build_reference(
+                                    reference, context
+                                ),
+                            )
+
+
+class ConditionalExpansionStrategy:
+    """Expansion strategy that routes to different strategies based on conditions.
+
+    Evaluates conditions in order and uses the first matching strategy.
+    Falls back to default strategy if no conditions match.
+
+    Example:
+        >>> # Use different strategies for math vs code problems
+        >>> strategy = ConditionalExpansionStrategy(
+        ...     rules=[
+        ...         (lambda row: row.get("type") == "math", math_strategy),
+        ...         (lambda row: row.get("type") == "code", code_strategy),
+        ...     ],
+        ...     default=default_strategy
+        ... )
+    """
+
+    def __init__(
+        self,
+        rules: list[tuple[Callable[[dict], bool], ExpansionStrategy]],
+        default: ExpansionStrategy,
+    ):
+        """Initialize conditional expansion strategy.
+
+        Args:
+            rules: List of (condition, strategy) tuples
+            default: Default strategy if no conditions match
+        """
+        self._rules = rules
+        self._default = default
+
+    def expand(
+        self,
+        dataset: Sequence[dict[str, Any]],
+        context: PlanContext,
+    ) -> Iterator[core_entities.GenerationTask]:
+        """Expand using conditional routing."""
+        # Group rows by which strategy applies
+        strategy_groups: dict[int, list[dict]] = {}
+
+        for row in dataset:
+            # Find first matching rule
+            matched = False
+            for rule_idx, (condition, strategy) in enumerate(self._rules):
+                if condition(row):
+                    strategy_groups.setdefault(rule_idx, []).append(row)
+                    matched = True
+                    break
+
+            if not matched:
+                strategy_groups.setdefault(-1, []).append(row)
+
+        # Expand each group with its strategy
+        for rule_idx, group_rows in strategy_groups.items():
+            if rule_idx == -1:
+                strategy = self._default
+            else:
+                strategy = self._rules[rule_idx][1]
+
+            yield from strategy.expand(group_rows, context)
+
+
+class ChainedExpansionStrategy:
+    """Expansion strategy that chains multiple strategies.
+
+    Applies multiple strategies in sequence, yielding tasks from all of them.
+    Useful for combining different expansion approaches.
+
+    Example:
+        >>> # Generate baseline tasks + additional high-temperature samples for hard problems
+        >>> strategy = ChainedExpansionStrategy([
+        ...     CartesianExpansionStrategy(),
+        ...     FilteredExpansionStrategy(
+        ...         task_filter=lambda row, tpl, mdl, smp: (
+        ...             row.get("difficulty") == "hard" and smp.temperature > 0.5
+        ...         )
+        ...     )
+        ... ])
+    """
+
+    def __init__(self, strategies: Sequence[ExpansionStrategy]):
+        """Initialize chained expansion strategy.
+
+        Args:
+            strategies: List of strategies to apply in sequence
+        """
+        self._strategies = strategies
+
+    def expand(
+        self,
+        dataset: Sequence[dict[str, Any]],
+        context: PlanContext,
+    ) -> Iterator[core_entities.GenerationTask]:
+        """Expand using all strategies in sequence."""
+        for strategy in self._strategies:
+            yield from strategy.expand(dataset, context)
+
+
+@dataclass
+class FlexibleGenerationPlan:
+    """Generation plan with pluggable expansion strategy.
+
+    Allows controlling how dataset rows are expanded into generation tasks
+    using different expansion strategies.
+
+    Example:
+        >>> # Filter expensive model to hard problems only
+        >>> plan = FlexibleGenerationPlan(
+        ...     templates=[template1, template2],
+        ...     models=[cheap_model, expensive_model],
+        ...     sampling_parameters=[config],
+        ...     expansion_strategy=FilteredExpansionStrategy(
+        ...         task_filter=lambda row, tpl, mdl, smp: (
+        ...             mdl.identifier != "expensive" or row.get("difficulty") == "hard"
+        ...         )
+        ...     )
+        ... )
+    """
+
+    templates: Sequence[templates.PromptTemplate]
+    models: Sequence[core_entities.ModelSpec]
+    sampling_parameters: Sequence[core_entities.SamplingConfig]
+    expansion_strategy: ExpansionStrategy | None = None
+    dataset_id_field: str = "id"
+    reference_field: str | None = "expected"
+    metadata_fields: Sequence[str] = field(default_factory=tuple)
+    context_builder: Callable[[dict], dict] | None = None
+
+    def expand(
+        self, dataset: Sequence[dict[str, object]]
+    ) -> Iterator[core_entities.GenerationTask]:
+        """Expand dataset into generation tasks using strategy.
+
+        Args:
+            dataset: Dataset rows to expand
+
+        Yields:
+            Generation tasks
+        """
+        context = PlanContext(
+            templates=self.templates,
+            models=self.models,
+            sampling_parameters=self.sampling_parameters,
+            dataset_id_field=self.dataset_id_field,
+            reference_field=self.reference_field,
+            metadata_fields=self.metadata_fields,
+            context_builder=self.context_builder,
+        )
+
+        strategy = self.expansion_strategy or CartesianExpansionStrategy()
+        yield from strategy.expand(dataset, context)

themis/generation/providers/litellm_provider.py

@@ -0,0 +1,221 @@
+"""LiteLLM provider supporting 100+ LLM providers through a unified interface."""
+
+from __future__ import annotations
+
+import threading
+from dataclasses import dataclass
+from typing import Any, Dict
+
+from themis.core import entities as core_entities
+from themis.interfaces import ModelProvider
+from themis.providers import register_provider
+
+
+@dataclass
+class LiteLLMProvider(ModelProvider):
+    """
+    Universal LLM provider using LiteLLM.
+
+    Supports 100+ providers including:
+    - OpenAI (gpt-4, gpt-3.5-turbo, etc.)
+    - Anthropic (claude-3-opus, claude-3-sonnet, etc.)
+    - Azure OpenAI (azure/<deployment-name>)
+    - AWS Bedrock (bedrock/<model-id>)
+    - Google AI (gemini-pro, etc.)
+    - Cohere, Replicate, Hugging Face, and many more
+
+    Configuration options:
+    - api_key: Optional API key (can also use env vars like OPENAI_API_KEY)
+    - api_base: Optional custom API base URL
+    - timeout: Request timeout in seconds (default: 60)
+    - max_retries: Number of retries for failed requests (default: 2)
+    - n_parallel: Maximum number of parallel requests (default: 10)
+    - drop_params: Whether to drop unsupported params (default: False)
+    - custom_llm_provider: Force a specific provider (e.g., "openai", "anthropic")
+    - extra_kwargs: Additional kwargs to pass to litellm.completion()
+    """
+
+    api_key: str | None = None
+    api_base: str | None = None
+    timeout: int = 60
+    max_retries: int = 2
+    n_parallel: int = 10
+    drop_params: bool = False
+    custom_llm_provider: str | None = None
+    extra_kwargs: Dict[str, Any] | None = None
+
+    def __post_init__(self) -> None:
+        self._semaphore = threading.Semaphore(max(1, self.n_parallel))
+        self._extra_kwargs = self.extra_kwargs or {}
+
+        # Lazy import to avoid import errors if litellm not installed
+        try:
+            import litellm
+
+            self._litellm = litellm
+            # Configure litellm settings
+            litellm.drop_params = self.drop_params
+            if self.max_retries > 0:
+                litellm.num_retries = self.max_retries
+        except ImportError as exc:
+            raise RuntimeError(
+                "LiteLLM is not installed. Install via `pip install litellm` or "
+                "`uv add litellm` to use LiteLLMProvider."
+            ) from exc
+
+    def generate(
+        self, task: core_entities.GenerationTask
+    ) -> core_entities.GenerationRecord:  # type: ignore[override]
+        """Generate a response using LiteLLM."""
+
+        messages = self._build_messages(task)
+        completion_kwargs = self._build_completion_kwargs(task, messages)
+
+        try:
+            with self._semaphore:
+                response = self._litellm.completion(**completion_kwargs)
+
+            # Extract the generated text
+            text = response.choices[0].message.content or ""
+
+            # Extract usage information
+            usage = response.usage if hasattr(response, "usage") else None
+            usage_dict = None
+            metrics = {}
+            if usage:
+                prompt_tokens = getattr(usage, "prompt_tokens", None)
+                completion_tokens = getattr(usage, "completion_tokens", None)
+                total_tokens = getattr(usage, "total_tokens", None)
+
+                metrics["prompt_tokens"] = prompt_tokens
+                metrics["completion_tokens"] = completion_tokens
+                metrics["total_tokens"] = total_tokens
+                # Alias for consistency with other providers
+                metrics["response_tokens"] = completion_tokens
+
+                # Create usage dict for cost tracking
+                if prompt_tokens is not None and completion_tokens is not None:
+                    usage_dict = {
+                        "prompt_tokens": prompt_tokens,
+                        "completion_tokens": completion_tokens,
+                        "total_tokens": total_tokens or (prompt_tokens + completion_tokens),
+                    }
+
+            # Extract model information
+            model_used = getattr(response, "model", task.model.identifier)
+            metrics["model_used"] = model_used
+
+            # Convert response to dict for raw storage
+            raw_data = response.model_dump() if hasattr(response, "model_dump") else {}
+
+            return core_entities.GenerationRecord(
+                task=task,
+                output=core_entities.ModelOutput(text=text, raw=raw_data, usage=usage_dict),
+                error=None,
+                metrics=metrics,
+            )
+
+        except Exception as exc:
+            # Capture detailed error information
+            error_type = type(exc).__name__
+            error_message = str(exc)
+
+            # Extract additional context if available
+            details: Dict[str, Any] = {
+                "error_type": error_type,
+                "model": task.model.identifier,
+            }
+
+            # Check for specific litellm exceptions
+            if hasattr(exc, "status_code"):
+                details["status_code"] = exc.status_code  # type: ignore
+            if hasattr(exc, "llm_provider"):
+                details["llm_provider"] = exc.llm_provider  # type: ignore
+
+            return core_entities.GenerationRecord(
+                task=task,
+                output=None,
+                error=core_entities.ModelError(
+                    message=error_message,
+                    kind=error_type,
+                    details=details,
+                ),
+                metrics={},
+            )
+
+    def _build_messages(
+        self, task: core_entities.GenerationTask
+    ) -> list[dict[str, str]]:
+        """Build messages array from the generation task."""
+        messages = []
+
+        # Add system message if provided in metadata
+        system_prompt = task.prompt.metadata.get(
+            "system_prompt"
+        ) or task.prompt.spec.metadata.get("system_prompt")
+        if system_prompt:
+            messages.append({"role": "system", "content": system_prompt})
+
+        # Add the main user prompt
+        messages.append({"role": "user", "content": task.prompt.text})
+
+        # Support for conversation history if provided
+        conversation_history = task.metadata.get("conversation_history")
+        if conversation_history:
+            # If conversation history is provided, prepend it
+            messages = conversation_history + messages
+
+        return messages
+
+    def _build_completion_kwargs(
+        self, task: core_entities.GenerationTask, messages: list[dict[str, str]]
+    ) -> Dict[str, Any]:
+        """Build the kwargs dictionary for litellm.completion()."""
+
+        kwargs: Dict[str, Any] = {
+            "model": task.model.identifier,
+            "messages": messages,
+            "temperature": task.sampling.temperature,
+            "top_p": task.sampling.top_p,
+            "timeout": self.timeout,
+        }
+
+        # Add max_tokens if specified (negative values mean no limit)
+        if task.sampling.max_tokens >= 0:
+            kwargs["max_tokens"] = task.sampling.max_tokens
+
+        # Add API key if provided
+        if self.api_key:
+            kwargs["api_key"] = self.api_key
+
+        # Add custom API base if provided
+        if self.api_base:
+            kwargs["api_base"] = self.api_base
+
+        # Add custom provider if specified
+        if self.custom_llm_provider:
+            kwargs["custom_llm_provider"] = self.custom_llm_provider
+
+        # Merge any extra kwargs provided in configuration
+        if self._extra_kwargs:
+            kwargs.update(self._extra_kwargs)
+
+        # Allow task-level overrides via metadata
+        litellm_kwargs = task.metadata.get("litellm_kwargs", {})
+        if litellm_kwargs:
+            kwargs.update(litellm_kwargs)
+
+        return kwargs
+
+
+# Register the provider with multiple aliases for convenience
+register_provider("litellm", LiteLLMProvider)
+register_provider("openai", LiteLLMProvider)
+register_provider("anthropic", LiteLLMProvider)
+register_provider("azure", LiteLLMProvider)
+register_provider("bedrock", LiteLLMProvider)
+register_provider("gemini", LiteLLMProvider)
+register_provider("cohere", LiteLLMProvider)
+
+
+__all__ = ["LiteLLMProvider"]