add comment

agentlightning/algorithm/apo/apo.py

@@ -7,23 +7,44 @@ APO with textual gradients that read rollout spans and outputs to modify the pro
 - rollout: same pattern as your example, but task is a dict (T_task)
 """
 
+from __future__ import annotations
+
 import asyncio
 import logging
 import random
 import time
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Counter, Dict, Generic, Iterator, List, Optional, Sequence, Set, Tuple, TypedDict, TypeVar, cast
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Counter,
+    Dict,
+    Generic,
+    Iterator,
+    List,
+    Optional,
+    Sequence,
+    Set,
+    Tuple,
+    TypedDict,
+    TypeVar,
+    cast,
+)
 
 import poml
 from openai import AsyncOpenAI
 
 from agentlightning.adapter.messages import TraceToMessages
 from agentlightning.algorithm.base import Algorithm
-from agentlightning.algorithm.utils import batch_iter_over_dataset
+from agentlightning.algorithm.utils import batch_iter_over_dataset, with_llm_proxy, with_store
 from agentlightning.reward import find_final_reward
 from agentlightning.types import Dataset, NamedResources, PromptTemplate, Rollout, RolloutMode, RolloutStatus
 
+if TYPE_CHECKING:
+    from agentlightning.llm_proxy import LLMProxy
+    from agentlightning.store.base import LightningStore
+
 logger = logging.getLogger(__name__)
 
 T_task = TypeVar("T_task")
@@ -360,8 +381,10 @@ class APO(Algorithm, Generic[T_task]):
             )
         return new_prompt
 
+    @with_store
     async def get_rollout_results(
         self,
+        store: LightningStore,
         rollout: List[Rollout],
         *,
         prefix: Optional[str] = None,
@@ -379,7 +402,6 @@ class APO(Algorithm, Generic[T_task]):
             List of rollout results formatted for APO processing.
         """
         rollout_results: List[RolloutResultForAPO] = []
-        store = self.get_store()
         adapter = self.get_adapter()
         for r in rollout:
             spans = await store.query_spans(r.rollout_id)
@@ -776,8 +798,12 @@ class APO(Algorithm, Generic[T_task]):
                 prefix=prefix,
             )
 
+    @with_llm_proxy()
+    @with_store
     async def run(
         self,
+        store: LightningStore,  # Injected by decorator - callers should not provide this parameter
+        llm_proxy: Optional[LLMProxy],  # Injected by decorator - callers should not provide this parameter
         train_dataset: Optional[Dataset[T_task]] = None,
         val_dataset: Optional[Dataset[T_task]] = None,
     ) -> None:

agentlightning/algorithm/fast.py

@@ -5,11 +5,16 @@ from __future__ import annotations
 import asyncio
 import logging
 from datetime import datetime
-from typing import Any, List, Literal, Optional
+from typing import TYPE_CHECKING, Any, List, Literal, Optional
 
 from agentlightning.types import Attempt, Dataset, Rollout, RolloutStatus, Span
 
 from .base import Algorithm
+from .utils import with_llm_proxy, with_store
+
+if TYPE_CHECKING:
+    from agentlightning.llm_proxy import LLMProxy
+    from agentlightning.store.base import LightningStore
 
 logger = logging.getLogger(__name__)
 
@@ -36,6 +41,8 @@ class Baseline(FastAlgorithm):
     finish, and logs every collected span and reward. It is primarily useful as
     a smoke test for the platform plumbing rather than a performant trainer.
 
+    The baseline algorithm will auto-start a LLM proxy if one is provided and not yet started.
+
     Args:
         n_epochs: Number of dataset passes to execute for both the train and val
             splits during developer experiments.
@@ -180,8 +187,12 @@ class Baseline(FastAlgorithm):
 
             await asyncio.sleep(self.polling_interval)
 
+    @with_llm_proxy()
+    @with_store
     async def run(
         self,
+        store: LightningStore,  # Injected by decorator - callers should not provide this parameter
+        llm_proxy: Optional[LLMProxy],  # Injected by decorator - callers should not provide this parameter
         train_dataset: Optional[Dataset[Any]] = None,
         val_dataset: Optional[Dataset[Any]] = None,
     ) -> None:
@@ -202,8 +213,6 @@ class Baseline(FastAlgorithm):
         logger.debug(f"Train indices: {train_indices}")
         logger.debug(f"Val indices: {val_indices}")
 
-        store = self.get_store()
-
         # Currently we only supports a single resource update at the start.
         initial_resources = self.get_initial_resources()
         if initial_resources is not None:

agentlightning/algorithm/utils.py

@@ -1,11 +1,42 @@
 # Copyright (c) Microsoft. All rights reserved.
 
+from __future__ import annotations
+
+import functools
+import logging
 import random
-from typing import Iterator, List, Sequence, TypeVar
+from collections.abc import Coroutine
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Concatenate,
+    Iterator,
+    List,
+    Literal,
+    Optional,
+    ParamSpec,
+    Sequence,
+    TypeVar,
+    overload,
+)
 
 from agentlightning.types import Dataset
 
+if TYPE_CHECKING:
+    from agentlightning.llm_proxy import LLMProxy
+    from agentlightning.store.base import LightningStore
+
+    from .base import Algorithm
+
 T_task = TypeVar("T_task")
+T_algo = TypeVar("T_algo", bound="Algorithm")
+
+P = ParamSpec("P")
+R = TypeVar("R")
+
+
+logger = logging.getLogger(__name__)
 
 
 def batch_iter_over_dataset(dataset: Dataset[T_task], batch_size: int) -> Iterator[Sequence[T_task]]:
@@ -41,3 +72,106 @@ def batch_iter_over_dataset(dataset: Dataset[T_task], batch_size: int) -> Iterat
                 if len(current_batch) == batch_size:
                     yield [dataset[index] for index in current_batch]
                     current_batch = []
+
+
+def with_store(
+    func: Callable[Concatenate[T_algo, LightningStore, P], Coroutine[Any, Any, R]],
+) -> Callable[Concatenate[T_algo, P], Coroutine[Any, Any, R]]:
+    """Inject the algorithm's `LightningStore` into coroutine methods.
+
+    The decorator calls `Algorithm.get_store()` once per invocation and passes the
+    resulting store as an explicit argument to the wrapped coroutine. Decorated
+    methods therefore receive the resolved store even when invoked by helper
+    utilities rather than directly by the algorithm.
+
+    Args:
+        func: The coroutine that expects `(self, store, *args, **kwargs)`.
+
+    Returns:
+        A coroutine wrapper that automatically retrieves the store and forwards it
+        to `func`.
+    """
+
+    @functools.wraps(func)
+    async def wrapper(self: T_algo, *args: P.args, **kwargs: P.kwargs) -> R:
+        store = self.get_store()
+        return await func(self, store, *args, **kwargs)
+
+    return wrapper
+
+
+@overload
+def with_llm_proxy(
+    required: Literal[False] = False,
+    auto_start: bool = True,
+) -> Callable[
+    [Callable[Concatenate[T_algo, Optional[LLMProxy], P], Coroutine[Any, Any, R]]],
+    Callable[Concatenate[T_algo, P], Coroutine[Any, Any, R]],
+]: ...
+
+
+@overload
+def with_llm_proxy(
+    required: Literal[True],
+    auto_start: bool = True,
+) -> Callable[
+    [Callable[Concatenate[T_algo, LLMProxy, P], Coroutine[Any, Any, R]]],
+    Callable[Concatenate[T_algo, P], Coroutine[Any, Any, R]],
+]: ...
+
+
+def with_llm_proxy(
+    required: bool = False,
+    auto_start: bool = True,
+) -> Callable[
+    [Callable[..., Coroutine[Any, Any, Any]]],
+    Callable[..., Coroutine[Any, Any, Any]],
+]:
+    """Resolve and optionally lifecycle-manage the configured LLM proxy.
+
+    Args:
+        required: When True, raises `ValueError` if the algorithm does not have an
+            [`LLMProxy`][agentlightning.LLMProxy] set. When False, the wrapped coroutine receives
+            `None` if no proxy is available.
+        auto_start: When True, [`LLMProxy.start()`][agentlightning.LLMProxy.start] is invoked if the proxy is not
+            already running before calling `func` and [`LLMProxy.stop()`][agentlightning.LLMProxy.stop] is
+            called afterwards.
+
+    Returns:
+        A decorator that injects the [`LLMProxy`][agentlightning.LLMProxy] (or `None`) as the first
+        argument after `self` and manages automatic startup/shutdown when requested.
+    """
+
+    def decorator(
+        func: Callable[..., Coroutine[Any, Any, Any]],
+    ) -> Callable[..., Coroutine[Any, Any, Any]]:
+        @functools.wraps(func)
+        async def wrapper(self: Algorithm, *args: Any, **kwargs: Any) -> Any:
+            llm_proxy = self.get_llm_proxy()
+
+            if required and llm_proxy is None:
+                raise ValueError(
+                    "LLM proxy is required but not configured. Call set_llm_proxy() before using this method."
+                )
+
+            auto_started = False
+            if auto_start and llm_proxy is not None:
+                if llm_proxy.is_running():
+                    logger.info("Proxy is already running, skipping start")
+                else:
+                    logger.info("Starting proxy, managed by the algorithm")
+                    await llm_proxy.start()
+                    auto_started = True
+
+            try:
+                # At type level, overloads guarantee that if `required=True`
+                # then `func` expects a non-optional LLMProxy.
+                return await func(self, llm_proxy, *args, **kwargs)
+            finally:
+                if auto_started and llm_proxy is not None:
+                    logger.info("Stopping proxy, managed by the algorithm")
+                    await llm_proxy.stop()
+
+        return wrapper
+
+    return decorator

tests/algorithm/test_utils.py

@@ -0,0 +1,115 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+from typing import Any, cast
+from unittest.mock import MagicMock
+
+import pytest
+
+from agentlightning.algorithm.base import Algorithm
+from agentlightning.algorithm.utils import with_llm_proxy, with_store
+from agentlightning.llm_proxy import LLMProxy
+from agentlightning.store.base import LightningStore
+
+
+class _BaseAlgorithm(Algorithm):
+    def run(self, *args: Any, **kwargs: Any) -> None:
+        """Satisfy the abstract interface without invoking training logic."""
+        return None
+
+
+class _StubLLMProxy:
+    """Test double that tracks lifecycle calls."""
+
+    def __init__(self) -> None:
+        self.start_calls = 0
+        self.stop_calls = 0
+        self.running = False
+
+    def is_running(self) -> bool:
+        return self.running
+
+    async def start(self) -> None:
+        self.start_calls += 1
+        self.running = True
+
+    async def stop(self) -> None:
+        self.stop_calls += 1
+        self.running = False
+
+
+@pytest.mark.asyncio
+async def test_with_store_injects_store_argument():
+    class StoreAlgorithm(_BaseAlgorithm):
+        @with_store
+        async def record_store(self, store: LightningStore, payload: str) -> None:
+            self.seen_store = store  # type: ignore[attr-defined]
+            self.seen_payload = payload  # type: ignore[attr-defined]
+
+    algorithm = StoreAlgorithm()
+    fake_store = MagicMock(spec=LightningStore)
+    algorithm.set_store(fake_store)
+
+    await algorithm.record_store("batch-1")
+
+    assert algorithm.seen_store is fake_store  # type: ignore[attr-defined]
+    assert algorithm.seen_payload == "batch-1"  # type: ignore[attr-defined]
+
+
+@pytest.mark.asyncio
+async def test_with_llm_proxy_allows_optional_injection():
+    class OptionalProxyAlgorithm(_BaseAlgorithm):
+        @with_llm_proxy()
+        async def record_proxy(self, llm_proxy: LLMProxy | None, marker: str) -> None:
+            self.seen_proxy = llm_proxy  # type: ignore[attr-defined]
+            self.marker = marker  # type: ignore[attr-defined]
+
+    algorithm = OptionalProxyAlgorithm()
+    algorithm.set_llm_proxy(None)
+
+    await algorithm.record_proxy("optional")
+
+    assert algorithm.seen_proxy is None  # type: ignore[attr-defined]
+    assert algorithm.marker == "optional"  # type: ignore[attr-defined]
+
+
+@pytest.mark.asyncio
+async def test_with_llm_proxy_required_raises_when_missing():
+    class RequiredProxyAlgorithm(_BaseAlgorithm):
+        @with_llm_proxy(required=True)
+        async def record_proxy(self, llm_proxy: LLMProxy) -> None:
+            self.seen_proxy = llm_proxy  # type: ignore[attr-defined]
+
+    algorithm = RequiredProxyAlgorithm()
+    algorithm.set_llm_proxy(None)
+
+    with pytest.raises(ValueError):
+        await algorithm.record_proxy()
+
+
+@pytest.mark.asyncio
+async def test_with_llm_proxy_auto_start_and_stop():
+    class AutoProxyAlgorithm(_BaseAlgorithm):
+        @with_llm_proxy()
+        async def use_proxy(self, llm_proxy: LLMProxy | None) -> None:
+            if llm_proxy is None:
+                raise AssertionError("LLM proxy should be injected")
+            self.seen_proxy = llm_proxy  # type: ignore[attr-defined]
+
+    algorithm = AutoProxyAlgorithm()
+    proxy = _StubLLMProxy()
+    algorithm.set_llm_proxy(cast(LLMProxy, proxy))
+
+    await algorithm.use_proxy()
+
+    assert algorithm.seen_proxy is proxy  # type: ignore[attr-defined]
+    assert proxy.start_calls == 1
+    assert proxy.stop_calls == 1
+
+    # When already running, no extra start/stop should be requested.
+    proxy.running = True
+    await algorithm.use_proxy()
+
+    assert proxy.start_calls == 1
+    assert proxy.stop_calls == 1