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

themis/experiment/builder.py

@@ -0,0 +1,151 @@
+"""Utilities for assembling experiments from reusable components."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any, Callable, Mapping, Sequence, Type
+
+from themis.config import schema as config
+from themis.core import entities as core_entities
+from themis.evaluation import pipeline as evaluation_pipeline
+from themis.evaluation import strategies as evaluation_strategies
+from themis.experiment import orchestrator
+from themis.experiment import storage as experiment_storage
+from themis.experiment.cache_manager import CacheManager
+from themis.experiment.definitions import (
+    BuiltExperiment,
+    ExperimentDefinition,
+    ModelBinding,
+)
+from themis.experiment.integration_manager import IntegrationManager
+from themis.generation import plan as generation_plan
+from themis.generation import router as generation_router
+from themis.generation import runner as generation_runner
+from themis.generation import strategies as generation_strategies
+from themis.interfaces import ModelProvider
+from themis.providers import create_provider
+
+
+class ExperimentBuilder:
+    """Composable builder for constructing experiment components."""
+
+    def __init__(
+        self,
+        *,
+        extractor,
+        metrics,
+        runner_cls: Type[
+            generation_runner.GenerationRunner
+        ] = generation_runner.GenerationRunner,
+        runner_kwargs: Mapping[str, Any] | None = None,
+        pipeline_cls: Type[
+            evaluation_pipeline.EvaluationPipeline
+        ] = evaluation_pipeline.EvaluationPipeline,
+        pipeline_kwargs: Mapping[str, Any] | None = None,
+        router_cls: Type[ModelProvider] = generation_router.ProviderRouter,
+        router_kwargs: Mapping[str, Any] | None = None,
+        strategy_resolver: Callable[
+            [core_entities.GenerationTask], generation_strategies.GenerationStrategy
+        ]
+        | None = None,
+        evaluation_strategy_resolver: Callable[
+            [core_entities.GenerationRecord], evaluation_strategies.EvaluationStrategy
+        ]
+        | None = None,
+    ) -> None:
+        self._extractor = extractor
+        self._metrics = list(metrics)
+        self._runner_cls = runner_cls
+        self._runner_kwargs = dict(runner_kwargs or {})
+        self._pipeline_cls = pipeline_cls
+        self._pipeline_kwargs = dict(pipeline_kwargs or {})
+        self._router_cls = router_cls
+        self._router_kwargs = dict(router_kwargs or {})
+        self._strategy_resolver = strategy_resolver
+        self._evaluation_strategy_resolver = evaluation_strategy_resolver
+
+    def build(
+        self,
+        definition: ExperimentDefinition,
+        *,
+        storage_dir: str | Path | None = None,
+    ) -> BuiltExperiment:
+        plan_obj = self._build_plan(definition)
+        router = self._build_router(definition.model_bindings)
+        runner_kwargs = dict(self._runner_kwargs)
+        if self._strategy_resolver is not None:
+            runner_kwargs.setdefault("strategy_resolver", self._strategy_resolver)
+        runner = self._runner_cls(provider=router, **runner_kwargs)
+        pipeline_kwargs = dict(self._pipeline_kwargs)
+        if self._evaluation_strategy_resolver is not None:
+            pipeline_kwargs.setdefault(
+                "strategy_resolver", self._evaluation_strategy_resolver
+            )
+        pipeline = self._pipeline_cls(
+            extractor=self._extractor,
+            metrics=self._metrics,
+            **pipeline_kwargs,
+        )
+
+        # Create storage backend
+        storage = (
+            experiment_storage.ExperimentStorage(storage_dir)
+            if storage_dir is not None
+            else None
+        )
+
+        # Create managers for better separation of concerns
+        cache_manager = CacheManager(
+            storage=storage,
+            enable_resume=True,
+            enable_cache=True,
+        )
+        integration_manager = IntegrationManager(config=config.IntegrationsConfig())
+
+        # Create orchestrator with managers
+        orchestrator_obj = orchestrator.ExperimentOrchestrator(
+            generation_plan=plan_obj,
+            generation_runner=runner,
+            evaluation_pipeline=pipeline,
+            cache_manager=cache_manager,
+            integration_manager=integration_manager,
+        )
+
+        return BuiltExperiment(
+            orchestrator=orchestrator_obj,
+            plan=plan_obj,
+            runner=runner,
+            pipeline=pipeline,
+            storage=storage,
+            router=router,
+        )
+
+    def _build_plan(
+        self, definition: ExperimentDefinition
+    ) -> generation_plan.GenerationPlan:
+        return generation_plan.GenerationPlan(
+            templates=list(definition.templates),
+            models=[binding.spec for binding in definition.model_bindings],
+            sampling_parameters=list(definition.sampling_parameters),
+            dataset_id_field=definition.dataset_id_field,
+            reference_field=definition.reference_field,
+            metadata_fields=tuple(definition.metadata_fields),
+            context_builder=definition.context_builder,
+        )
+
+    def _build_router(self, bindings: Sequence[ModelBinding]) -> ModelProvider:
+        providers: dict[str, ModelProvider] = {}
+        for binding in bindings:
+            providers[binding.spec.identifier] = create_provider(
+                binding.provider_name,
+                **binding.provider_options,
+            )
+        return self._router_cls(providers, **self._router_kwargs)
+
+
+__all__ = [
+    "ExperimentBuilder",
+    "ExperimentDefinition",
+    "ModelBinding",
+    "BuiltExperiment",
+]

themis/experiment/cache_manager.py

@@ -0,0 +1,134 @@
+"""Cache management for experiment resumability and storage."""
+
+from __future__ import annotations
+
+from typing import Sequence
+
+from themis.core.entities import EvaluationRecord, GenerationRecord
+from themis.experiment import storage as experiment_storage
+
+
+class CacheManager:
+    """Manages experiment caching and resumability.
+
+    This class handles all storage-related operations including:
+    - Loading cached generation records
+    - Loading cached evaluations
+    - Saving datasets for resumability
+    - Saving generation records and evaluations
+
+    Single Responsibility: Cache and storage management
+    """
+
+    def __init__(
+        self,
+        storage: experiment_storage.ExperimentStorage | None,
+        enable_resume: bool = True,
+        enable_cache: bool = True,
+    ) -> None:
+        """Initialize cache manager.
+
+        Args:
+            storage: Storage backend (None disables caching)
+            enable_resume: Whether to load cached results on resume
+            enable_cache: Whether to save new results to cache
+        """
+        self._storage = storage
+        self._enable_resume = enable_resume
+        self._enable_cache = enable_cache
+
+    @property
+    def has_storage(self) -> bool:
+        """Check if storage is available."""
+        return self._storage is not None
+
+    def cache_dataset(self, run_id: str, dataset: Sequence[dict[str, object]]) -> None:
+        """Cache dataset for future resumability.
+
+        Args:
+            run_id: Unique run identifier
+            dataset: Dataset samples to cache
+        """
+        if self._storage is not None and self._enable_cache:
+            self._storage.cache_dataset(run_id, list(dataset))
+
+    def load_cached_records(self, run_id: str) -> dict[str, GenerationRecord]:
+        """Load cached generation records for resuming.
+
+        Args:
+            run_id: Unique run identifier
+
+        Returns:
+            Dictionary mapping cache keys to generation records
+        """
+        if not self._enable_resume or self._storage is None:
+            return {}
+        return self._storage.load_cached_records(run_id)
+
+    def load_cached_evaluations(
+        self, run_id: str, evaluation_config: dict | None = None
+    ) -> dict[str, EvaluationRecord]:
+        """Load cached evaluation records for resuming.
+
+        Args:
+            run_id: Unique run identifier
+            evaluation_config: Evaluation configuration (metrics, extractor) for cache matching
+
+        Returns:
+            Dictionary mapping cache keys to evaluation records
+        """
+        if not self._enable_resume or self._storage is None:
+            return {}
+        return self._storage.load_cached_evaluations(run_id, evaluation_config=evaluation_config)
+
+    def save_generation_record(
+        self,
+        run_id: str,
+        record: GenerationRecord,
+        cache_key: str,
+    ) -> None:
+        """Save a single generation record.
+
+        Args:
+            run_id: Unique run identifier
+            record: Generation record to save
+            cache_key: Cache key for this record
+        """
+        if self._storage is not None and self._enable_cache:
+            self._storage.append_record(run_id, record, cache_key=cache_key)
+
+    def save_evaluation_record(
+        self,
+        run_id: str,
+        generation_record: GenerationRecord,
+        evaluation_record: EvaluationRecord,
+        evaluation_config: dict | None = None,
+    ) -> None:
+        """Save a single evaluation record.
+
+        Args:
+            run_id: Unique run identifier
+            generation_record: Corresponding generation record
+            evaluation_record: Evaluation record to save
+            evaluation_config: Evaluation configuration for cache invalidation
+        """
+        if self._storage is not None and self._enable_cache:
+            self._storage.append_evaluation(
+                run_id, generation_record, evaluation_record, evaluation_config=evaluation_config
+            )
+
+    def get_run_path(self, run_id: str) -> str | None:
+        """Get filesystem path for a run.
+
+        Args:
+            run_id: Unique run identifier
+
+        Returns:
+            Path to run directory, or None if no storage
+        """
+        if self._storage is None:
+            return None
+        return str(self._storage.get_run_path(run_id))
+
+
+__all__ = ["CacheManager"]