gatewaysdk 0.3.2__py3-none-any.whl

gatewaysdk/algorithm/decorator.py

@@ -0,0 +1,264 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import functools
+import inspect
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Awaitable,
+    Dict,
+    Generic,
+    Literal,
+    Optional,
+    Protocol,
+    TypeVar,
+    Union,
+    cast,
+    overload,
+)
+
+from gatewaysdk.adapter import TraceAdapter
+from gatewaysdk.store.base import LightningStore
+from gatewaysdk.types import Dataset, NamedResources
+
+if TYPE_CHECKING:
+    from gatewaysdk.llm_proxy import LLMProxy
+
+from .base import Algorithm
+
+# Algorithm function signature types
+# We've missed a lot of combinations here.
+# Let's add them in future.
+
+
+class AlgorithmFuncSyncFull(Protocol):
+    def __call__(
+        self,
+        *,
+        store: LightningStore,
+        train_dataset: Optional[Dataset[Any]],
+        val_dataset: Optional[Dataset[Any]],
+        llm_proxy: Optional[LLMProxy],
+        adapter: Optional[TraceAdapter[Any]],
+        initial_resources: Optional[NamedResources],
+    ) -> None: ...
+
+
+class AlgorithmFuncSyncOnlyStore(Protocol):
+    def __call__(self, *, store: LightningStore) -> None: ...
+
+
+class AlgorithmFuncSyncOnlyDataset(Protocol):
+    def __call__(self, *, train_dataset: Optional[Dataset[Any]], val_dataset: Optional[Dataset[Any]]) -> None: ...
+
+
+class AlgorithmFuncAsyncFull(Protocol):
+    def __call__(
+        self,
+        *,
+        store: LightningStore,
+        train_dataset: Optional[Dataset[Any]],
+        val_dataset: Optional[Dataset[Any]],
+        llm_proxy: Optional[LLMProxy],
+        adapter: Optional[TraceAdapter[Any]],
+        initial_resources: Optional[NamedResources],
+    ) -> Awaitable[None]: ...
+
+
+class AlgorithmFuncAsyncOnlyStore(Protocol):
+    def __call__(self, *, store: LightningStore) -> Awaitable[None]: ...
+
+
+class AlgorithmFuncAsyncOnlyDataset(Protocol):
+    def __call__(
+        self, *, train_dataset: Optional[Dataset[Any]], val_dataset: Optional[Dataset[Any]]
+    ) -> Awaitable[None]: ...
+
+
+AlgorithmFuncAsync = Union[AlgorithmFuncAsyncOnlyStore, AlgorithmFuncAsyncOnlyDataset, AlgorithmFuncAsyncFull]
+
+AlgorithmFuncSync = Union[AlgorithmFuncSyncOnlyStore, AlgorithmFuncSyncOnlyDataset, AlgorithmFuncSyncFull]
+
+
+class AlgorithmFuncSyncFallback(Protocol):
+    def __call__(self, *args: Any, **kwargs: Any) -> Any: ...
+
+
+class AlgorithmFuncAsyncFallback(Protocol):
+    def __call__(self, *args: Any, **kwargs: Any) -> Awaitable[Any]: ...
+
+
+AlgorithmFuncSyncLike = Union[AlgorithmFuncSync, AlgorithmFuncSyncFallback]
+AlgorithmFuncAsyncLike = Union[AlgorithmFuncAsync, AlgorithmFuncAsyncFallback]
+
+AlgorithmFunc = Union[AlgorithmFuncSyncLike, AlgorithmFuncAsyncLike]
+
+
+AsyncFlag = Literal[True, False]
+AF = TypeVar("AF", bound=AsyncFlag)
+
+
+class FunctionalAlgorithm(Algorithm, Generic[AF]):
+    """An algorithm wrapper built from a callable implementation.
+
+    Functional algorithms let you provide an ordinary function instead of
+    subclassing [`Algorithm`][gatewaysdk.Algorithm]. The wrapper inspects
+    the callable signature to supply optional dependencies
+    such as the store, adapter, and LLM proxy.
+    """
+
+    @overload
+    def __init__(self: "FunctionalAlgorithm[Literal[False]]", algorithm_func: AlgorithmFuncSyncLike) -> None: ...
+
+    @overload
+    def __init__(self: "FunctionalAlgorithm[Literal[True]]", algorithm_func: AlgorithmFuncAsyncLike) -> None: ...
+
+    def __init__(self, algorithm_func: Union[AlgorithmFuncSyncLike, AlgorithmFuncAsyncLike]) -> None:
+        """Wrap a function that implements algorithm behaviour.
+
+        Args:
+            algorithm_func: Sync or async callable implementing the algorithm
+                contract. Arguments are detected automatically based on the
+                function signature.
+        """
+        super().__init__()
+        self._algorithm_func = algorithm_func
+        self._sig = inspect.signature(algorithm_func)
+        self._is_async = inspect.iscoroutinefunction(algorithm_func)
+
+        # Copy function metadata to preserve type hints and other attributes
+        functools.update_wrapper(self, algorithm_func)  # type: ignore
+
+    def is_async(self) -> bool:
+        return self._is_async
+
+    @overload
+    def run(
+        self: "FunctionalAlgorithm[Literal[False]]",
+        train_dataset: Optional[Dataset[Any]] = None,
+        val_dataset: Optional[Dataset[Any]] = None,
+    ) -> None: ...
+
+    @overload
+    def run(
+        self: "FunctionalAlgorithm[Literal[True]]",
+        train_dataset: Optional[Dataset[Any]] = None,
+        val_dataset: Optional[Dataset[Any]] = None,
+    ) -> Awaitable[None]: ...
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        return self._algorithm_func(*args, **kwargs)  # type: ignore
+
+    def run(
+        self,
+        train_dataset: Optional[Dataset[Any]] = None,
+        val_dataset: Optional[Dataset[Any]] = None,
+    ) -> Union[None, Awaitable[None]]:
+        """Execute the wrapped function with injected dependencies.
+
+        Args:
+            train_dataset: Optional training dataset passed through when the
+                callable declares a `train_dataset` parameter.
+            val_dataset: Optional validation dataset passed through when the
+                callable declares a `val_dataset` parameter.
+
+        Returns:
+            None for sync callables or an awaitable when the callable is async.
+
+        Raises:
+            TypeError: If a dataset is provided but the function signature does
+                not accept the corresponding argument.
+        """
+        kwargs: Dict[str, Any] = {}
+        if "store" in self._sig.parameters:
+            kwargs["store"] = self.get_store()
+        if "adapter" in self._sig.parameters:
+            kwargs["adapter"] = self.get_adapter()
+        if "llm_proxy" in self._sig.parameters:
+            kwargs["llm_proxy"] = self.get_llm_proxy()
+        if "initial_resources" in self._sig.parameters:
+            kwargs["initial_resources"] = self.get_initial_resources()
+        if "train_dataset" in self._sig.parameters:
+            kwargs["train_dataset"] = train_dataset
+        elif train_dataset is not None:
+            raise TypeError(
+                f"train_dataset is provided but not supported by the algorithm function: {self._algorithm_func}"
+            )
+        if "val_dataset" in self._sig.parameters:
+            kwargs["val_dataset"] = val_dataset
+        elif val_dataset is not None:
+            raise TypeError(
+                f"val_dataset is provided but not supported by the algorithm function: {self._algorithm_func}"
+            )
+        # both sync and async functions can be called with the same signature
+        result = self._algorithm_func(**kwargs)  # type: ignore[misc]
+        if self._is_async:
+            return cast(Awaitable[None], result)
+        return None
+
+
+@overload
+def algo(func: AlgorithmFuncAsync) -> FunctionalAlgorithm[Literal[True]]: ...
+
+
+@overload
+def algo(func: AlgorithmFuncAsyncFallback) -> FunctionalAlgorithm[Any]: ...
+
+
+@overload
+def algo(func: AlgorithmFuncSync) -> FunctionalAlgorithm[Literal[False]]: ...
+
+
+@overload
+def algo(func: AlgorithmFuncSyncFallback) -> FunctionalAlgorithm[Any]: ...
+
+
+def algo(
+    func: Union[
+        AlgorithmFuncSync,
+        AlgorithmFuncAsync,
+        AlgorithmFuncSyncFallback,
+        AlgorithmFuncAsyncFallback,
+    ],
+) -> Union[FunctionalAlgorithm[Literal[False]], FunctionalAlgorithm[Literal[True]]]:
+    """Convert a callable into a [`FunctionalAlgorithm`][gatewaysdk.algorithm.decorator.FunctionalAlgorithm].
+
+    The decorator inspects the callable signature to decide which dependencies
+    to inject at runtime, enabling concise algorithm definitions that still
+    leverage the full training runtime.
+
+    Args:
+        func: Function implementing the algorithm logic. May be synchronous or
+            asynchronous. The function can expect all of, or a subset of the following parameters:
+
+            - `store`: [`LightningStore`][gatewaysdk.store.base.LightningStore],
+            - `train_dataset`: [`Dataset`][gatewaysdk.Dataset],
+            - `val_dataset`: [`Dataset`][gatewaysdk.Dataset],
+            - `llm_proxy`: [`LLMProxy`][gatewaysdk.LLMProxy],
+            - `adapter`: [`TraceAdapter`][gatewaysdk.TraceAdapter],
+            - `initial_resources`: [`NamedResources`][gatewaysdk.NamedResources],
+
+            If the function does not expect a parameter, the wrapper will not inject it into the call.
+            Using `*args` and `**kwargs` will not work and no parameters will be injected.
+
+    Returns:
+        FunctionalAlgorithm that proxies the callable while exposing the
+        `Algorithm` interface.
+
+    Examples:
+        ```python
+        from gatewaysdk.algorithm.decorator import algo
+
+        @algo
+        def batching_algorithm(*, store, train_dataset, val_dataset):
+            for sample in train_dataset:
+                store.enqueue_rollout(input=sample, mode="train")
+
+        @algo
+        async def async_algorithm(*, store, train_dataset=None, val_dataset=None):
+            await store.enqueue_rollout(input={"prompt": "hello"}, mode="train")
+        ```
+    """
+    return FunctionalAlgorithm(func)

gatewaysdk/algorithm/evals/__init__.py

@@ -0,0 +1,7 @@
+# Copyright (c) Gateway. All rights reserved.
+
+"""Evals algorithm — single-pass agent evaluation."""
+
+from .evals import Evals
+
+__all__ = ["Evals"]

gatewaysdk/algorithm/evals/evals.py

@@ -0,0 +1,217 @@
+# Copyright (c) Gateway. All rights reserved.
+
+"""Evals algorithm — single-pass agent evaluation with configurable parallelism.
+
+Runs every dataset item through the agent exactly once, collects reward scores,
+and reports aggregate metrics. No optimization, no reflection — just evaluate.
+
+Usage via the Gateway Platform UI:
+    1. Select algorithm: "Evals"
+    2. Configure parallelism (how many rollouts run concurrently)
+    3. Pick a dataset and agent
+    4. Start the run
+
+The algorithm enqueues rollouts in batches of ``parallelism``, waits for each
+batch to complete, extracts scores, and logs per-batch and aggregate metrics.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import uuid
+from typing import Any, ClassVar, Dict, List, Optional
+
+from gatewaysdk.algorithm.base import Algorithm
+from gatewaysdk.algorithm.utils import with_llm_proxy, with_store
+from gatewaysdk.emitter import find_final_reward
+from gatewaysdk.types import Dataset, NamedResources, PromptTemplate
+
+logger = logging.getLogger(__name__)
+
+
+class Evals(Algorithm):
+    """Single-pass evaluation algorithm.
+
+    Runs each dataset item through the agent once, collects the reward score,
+    and reports aggregate results. Parallelism controls how many rollouts are
+    in-flight at the same time.
+    """
+
+    # Platform metadata — auto-discovered by registry
+    name: ClassVar[str] = "evals"
+    display_name: ClassVar[str] = "Evals"
+    description: ClassVar[str] = "Single-pass agent evaluation"
+    requirements: ClassVar[Dict[str, Any]] = {"gpu": False, "min_runners": 1}
+    config_schema: ClassVar[Optional[Dict[str, Any]]] = {
+        "type": "object",
+        "properties": {
+            "parallelism": {
+                "type": "integer",
+                "default": 5,
+                "minimum": 1,
+                "maximum": 50,
+                "description": "Number of rollouts to run concurrently",
+            },
+            "rollout_timeout": {
+                "type": "number",
+                "default": 600.0,
+                "minimum": 30.0,
+                "description": "Timeout per rollout batch in seconds",
+            },
+            "val_fraction": {
+                "type": "number",
+                "default": 0.0,
+                "minimum": 0.0,
+                "maximum": 0.5,
+                "description": "Fraction of dataset held out for validation (0 = use all items)",
+            },
+        },
+    }
+
+    def __init__(
+        self,
+        *,
+        parallelism: int = 5,
+        rollout_timeout: float = 600.0,
+    ) -> None:
+        super().__init__()
+        self.parallelism = parallelism
+        self.rollout_timeout = rollout_timeout
+
+    @with_store
+    @with_llm_proxy(required=False)
+    async def run(
+        self,
+        llm_proxy,  # injected by @with_llm_proxy (may be None)
+        store,  # injected by @with_store
+        train_dataset: Optional[Dataset[Any]] = None,
+        val_dataset: Optional[Dataset[Any]] = None,
+    ) -> None:
+        """Run single-pass evaluation on the dataset.
+
+        Uses train_dataset as the evaluation set (the Platform UI sends
+        the selected dataset as train_dataset). val_dataset is ignored.
+        """
+        dataset = train_dataset
+        if dataset is None:
+            raise ValueError("A dataset is required for evaluation.")
+
+        store = self.get_store()
+        assert store is not None
+
+        initial_resources = self.get_initial_resources()
+        if not initial_resources:
+            raise ValueError("Initial resources must be set before running Evals.")
+
+        # Build resources: initial resources + LLM proxy resource (if available)
+        # The LLM proxy resource is required by agents using the @gepa.agent
+        # or @rollout decorators that declare an `llm: LLM` parameter.
+        resources: NamedResources = dict(initial_resources)
+        if llm_proxy is not None:
+            resources["llm"] = llm_proxy.as_resource()
+            logger.info(f"Added LLM proxy resource: model={resources['llm'].model}")
+
+        # Publish resources once
+        resources_id = f"evals-{uuid.uuid4().hex[:8]}"
+        await store.update_resources(resources_id, resources)
+        logger.info(f"Published resources: {resources_id}")
+
+        # Prepare dataset items
+        items: List[Dict[str, Any]] = []
+        for i, item in enumerate(dataset):
+            if isinstance(item, dict):
+                items.append(item)
+            else:
+                items.append({"task": item, "id": str(i)})
+
+        total = len(items)
+        logger.info(
+            f"Starting evaluation: {total} items, parallelism={self.parallelism}, " f"timeout={self.rollout_timeout}s"
+        )
+
+        # Process in batches
+        all_scores: List[Optional[float]] = []
+        step = 0
+
+        for batch_start in range(0, total, self.parallelism):
+            batch = items[batch_start : batch_start + self.parallelism]
+            batch_num = batch_start // self.parallelism + 1
+            total_batches = (total + self.parallelism - 1) // self.parallelism
+
+            logger.info(f"Batch {batch_num}/{total_batches}: " f"enqueuing {len(batch)} rollouts")
+
+            # Enqueue rollouts for this batch
+            rollout_ids: List[str] = []
+            for item in batch:
+                rollout = await store.enqueue_rollout(
+                    input=item,
+                    mode="test",
+                    resources_id=resources_id,
+                )
+                rollout_ids.append(rollout.rollout_id)
+
+            # Wait for all rollouts in the batch to complete
+            try:
+                completed = await store.wait_for_rollouts(
+                    rollout_ids=rollout_ids,
+                    timeout=self.rollout_timeout,
+                )
+            except Exception as e:
+                logger.error(f"Batch {batch_num} failed: {e}")
+                all_scores.extend([None] * len(batch))
+                continue
+
+            # Extract scores from completed rollouts
+            batch_scores: List[Optional[float]] = []
+            for rollout in completed:
+                spans = await store.query_spans(rollout.rollout_id)
+                reward = find_final_reward(spans)
+                batch_scores.append(reward)
+
+            all_scores.extend(batch_scores)
+
+            # Log per-batch metrics
+            valid_scores = [s for s in batch_scores if s is not None]
+            batch_avg = sum(valid_scores) / len(valid_scores) if valid_scores else 0.0
+            batch_completed = sum(1 for r in completed if r.status == "succeeded")
+
+            step += 1
+            if hasattr(store, "log"):
+                store.log(
+                    {
+                        "batch_avg_score": batch_avg,
+                        "batch_completed": batch_completed,
+                        "batch_total": len(batch),
+                    },
+                    step=step,
+                )
+
+            logger.info(
+                f"Batch {batch_num}/{total_batches} done: "
+                f"avg_score={batch_avg:.3f}, "
+                f"completed={batch_completed}/{len(batch)}"
+            )
+
+        # Log aggregate metrics
+        valid_scores = [s for s in all_scores if s is not None]
+        mean_score = sum(valid_scores) / len(valid_scores) if valid_scores else 0.0
+        completion_rate = len(valid_scores) / total if total > 0 else 0.0
+
+        step += 1
+        if hasattr(store, "log"):
+            store.log(
+                {
+                    "mean_score": mean_score,
+                    "completion_rate": completion_rate,
+                    "total_evaluated": len(valid_scores),
+                    "total_items": total,
+                },
+                step=step,
+            )
+
+        logger.info(
+            f"Evaluation complete: mean_score={mean_score:.3f}, "
+            f"completion_rate={completion_rate:.1%}, "
+            f"{len(valid_scores)}/{total} items scored"
+        )