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

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):