themis-eval 0.2.2__py3-none-any.whl → 1.0.0__py3-none-any.whl

themis/generation/plan.py

@@ -88,9 +88,20 @@ class GenerationPlan:
         }
         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]
+        
+        # If metadata_fields is explicitly specified, use it as a filter (existing behavior)
+        # Otherwise, include all fields by default (new behavior for custom metrics)
+        if self.metadata_fields:
+            # Explicit filter - only include specified fields
+            for field_name in self.metadata_fields:
+                if field_name in row:
+                    metadata[field_name] = row[field_name]
+        else:
+            # No filter - include all fields except those used for other purposes
+            for field_name, field_value in row.items():
+                if field_name not in (self.dataset_id_field, self.reference_field):
+                    metadata[field_name] = field_value
+        
         return metadata
 
     def _build_reference(
@@ -209,9 +220,20 @@ class CartesianExpansionStrategy:
         }
         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]
+        
+        # If metadata_fields is explicitly specified, use it as a filter (existing behavior)
+        # Otherwise, include all fields by default (new behavior for custom metrics)
+        if context.metadata_fields:
+            # Explicit filter - only include specified fields
+            for field_name in context.metadata_fields:
+                if field_name in row:
+                    metadata[field_name] = row[field_name]
+        else:
+            # No filter - include all fields except those used for other purposes
+            for field_name, field_value in row.items():
+                if field_name not in (context.dataset_id_field, context.reference_field):
+                    metadata[field_name] = field_value
+        
         return metadata
 
     def _build_reference(

themis/generation/router.py

@@ -2,22 +2,40 @@
 
 from __future__ import annotations
 
-from typing import Mapping
+from typing import Mapping, Tuple, Union
 
 from themis.core import entities as core_entities
 from themis.interfaces import ModelProvider
 
 
+ProviderKey = Union[str, Tuple[str, str]]
+
+
+def _model_key(provider: str, identifier: str) -> str:
+    return f"{provider}:{identifier}"
+
+
 class ProviderRouter(ModelProvider):
     """Dispatches generation tasks to concrete providers by model identifier."""
 
-    def __init__(self, providers: Mapping[str, ModelProvider]):
-        self._providers = dict(providers)
+    def __init__(self, providers: Mapping[ProviderKey, ModelProvider]):
+        normalized: dict[str, ModelProvider] = {}
+
+        for key, provider in providers.items():
+            if isinstance(key, tuple):
+                provider_name, model_id = key
+                normalized[_model_key(provider_name, model_id)] = provider
+            else:
+                normalized[str(key)] = provider
+
+        self._providers = normalized
 
     def generate(
         self, task: core_entities.GenerationTask
     ) -> core_entities.GenerationRecord:  # type: ignore[override]
-        provider = self._providers.get(task.model.identifier)
+        provider = self._providers.get(_model_key(task.model.provider, task.model.identifier))
+        if provider is None:
+            provider = self._providers.get(task.model.identifier)
         if provider is None:
             known = ", ".join(sorted(self._providers)) or "<none>"
             raise RuntimeError(

themis/generation/runner.py

@@ -26,6 +26,7 @@ class GenerationRunner:
             [core_entities.GenerationTask], strategies.GenerationStrategy
         ]
         | None = None,
+        execution_backend: object | None = None,
         max_parallel: int = 1,
         max_retries: int = 3,
         retry_initial_delay: float = 0.5,
@@ -36,6 +37,7 @@ class GenerationRunner:
         self._strategy_resolver = strategy_resolver or (
             lambda task: strategies.SingleAttemptStrategy()
         )
+        self._execution_backend = execution_backend
         self._max_parallel = max(1, max_parallel)
         self._max_retries = max(1, int(max_retries))
         self._retry_initial_delay = max(0.0, retry_initial_delay)
@@ -53,7 +55,20 @@ class GenerationRunner:
             return
         
         logger.info(f"Runner: Starting execution of {len(task_list)} tasks with {self._max_parallel} workers")
-        
+
+        if self._execution_backend is not None:
+            logger.info("Runner: Using custom execution backend")
+            backend = self._execution_backend
+            try:
+                for result in backend.map(
+                    self._execute_task, task_list, max_workers=self._max_parallel
+                ):
+                    yield result
+            except Exception as e:
+                logger.error(f"Runner: Execution backend failed: {e}")
+                raise
+            return
+
         if self._max_parallel <= 1:
             logger.info("Runner: Using sequential execution (1 worker)")
             for i, task in enumerate(task_list, 1):