chore: import upstream snapshot with attribution

agentlightning/client.py

@@ -0,0 +1,408 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Utilities for interacting with legacy Agent Lightning servers.
+
+This module contains compatibility shims that speak the deprecated HTTP
+interface used by older Agent Lightning deployments. Modern code should prefer
+the store-based APIs exposed by `agentlightning.store`, but keeping these
+clients available makes it easier to migrate existing workflows incrementally.
+"""
+
+import asyncio
+import logging
+import time
+import urllib.parse
+import warnings
+from typing import Any, Dict, List, Optional, Union
+
+import aiohttp
+import requests
+
+from .types import NamedResources, ResourcesUpdate, RolloutLegacy, Task, TaskIfAny, TaskInput
+
+logger = logging.getLogger(__name__)
+
+
+class AgentLightningClient:
+    """Client wrapper for the legacy version-aware Agent Lightning server.
+
+    The client exposes synchronous and asynchronous helpers for polling tasks,
+    retrieving resource bundles, and submitting rollouts. It also maintains a
+    simple in-memory cache keyed by the server-provided resource identifier to
+    avoid redundant network requests.
+
+    !!! warning "Deprecated"
+        [`AgentLightningClient`][agentlightning.client.AgentLightningClient] is part of
+        the legacy client/server stack. New code should rely on the store-based APIs
+        implemented in `agentlightning.store`.
+
+    Attributes:
+        endpoint: Base URL of the Agent Lightning server.
+        poll_interval: Delay in seconds between polling attempts when no task is
+            available.
+        timeout: Timeout in seconds applied to HTTP requests.
+        task_count: Number of tasks claimed during the lifetime of this client.
+    """
+
+    _next_task_uri = "/task"
+    _resources_uri = "/resources"
+    _latest_resources_uri = "/resources/latest"
+    _report_rollout_uri = "/rollout"
+
+    def __init__(self, endpoint: str, poll_interval: float = 5.0, timeout: float = 10.0):
+        """Initialize the client.
+
+        Args:
+            endpoint: Root URL of the Agent Lightning server.
+            poll_interval: Seconds to wait between polling attempts.
+            timeout: Seconds before a request to the server is considered timed out.
+        """
+        warnings.warn(
+            "AgentLightningClient is deprecated. Please use LightningStoreClient instead.", DeprecationWarning
+        )
+        self.endpoint = endpoint
+        self.task_count = 0
+        self.poll_interval = poll_interval
+        self.timeout = timeout
+        self._resource_cache: Dict[str, ResourcesUpdate] = {}  # TODO: mechanism to evict cache
+        self._default_headers = {"X-AgentLightning-Client": "true"}
+
+    async def _request_json_async(self, url: str) -> Optional[Dict[str, Any]]:
+        """Perform an asynchronous ``GET`` request and parse the JSON payload.
+
+        Args:
+            url: Fully qualified URL to query.
+
+        Returns:
+            Parsed JSON body as a dictionary if the request succeeds; otherwise ``None``.
+        """
+        timeout = aiohttp.ClientTimeout(total=self.timeout)
+        async with aiohttp.ClientSession(timeout=timeout) as session:
+            try:
+                async with session.get(url, headers=self._default_headers) as resp:
+                    resp.raise_for_status()
+                    return await resp.json()
+            except Exception as e:
+                logger.debug(f"Async GET request failed for {url}: {e}")
+                return None
+
+    async def _post_json_async(self, url: str, payload: Dict[str, Any]) -> Optional[Dict[str, Any]]:
+        """Perform an asynchronous ``POST`` request with a JSON body.
+
+        Args:
+            url: Fully qualified URL that accepts the payload.
+            payload: Dictionary that will be serialized and sent as JSON.
+
+        Returns:
+            Parsed JSON body as a dictionary if the request succeeds; otherwise ``None``.
+        """
+        timeout = aiohttp.ClientTimeout(total=self.timeout)
+        async with aiohttp.ClientSession(timeout=timeout) as session:
+            try:
+                async with session.post(url, json=payload, headers=self._default_headers) as resp:
+                    resp.raise_for_status()
+                    return await resp.json()
+            except Exception as e:
+                logger.debug(f"Async POST request failed for {url}: {e}")
+                return None
+
+    async def poll_next_task_async(self) -> Optional[Task]:
+        """Poll the server asynchronously until a task becomes available.
+
+        Returns:
+            The next [`Task`][agentlightning.Task] exposed by the server,
+            or ``None`` if polling fails.
+        """
+        url = urllib.parse.urljoin(self.endpoint, self._next_task_uri)
+        while True:
+            response = await self._request_json_async(url)
+            if response:
+                task_if_any = TaskIfAny.model_validate(response)
+                if task_if_any.is_available and task_if_any.task:
+                    self.task_count += 1
+                    logger.info(f"[Task {self.task_count} Received] ID: {task_if_any.task.rollout_id}")
+                    return task_if_any.task
+            logger.debug(f"No task available yet. Retrying in {self.poll_interval} seconds...")
+            await asyncio.sleep(self.poll_interval)
+
+    async def get_resources_by_id_async(self, resource_id: str) -> Optional[ResourcesUpdate]:
+        """Fetch a specific resource bundle by identifier.
+
+        Args:
+            resource_id: Identifier sourced from the task metadata.
+
+        Returns:
+            Cached or freshly downloaded
+            [`ResourcesUpdate`][agentlightning.ResourcesUpdate], or
+            ``None`` when the server returns an error.
+        """
+        if resource_id in self._resource_cache:
+            logger.debug(f"Found resources '{resource_id}' in cache.")
+            return self._resource_cache[resource_id]
+
+        url = urllib.parse.urljoin(self.endpoint, f"{self._resources_uri}/{resource_id}")
+        response = await self._request_json_async(url)
+        if response:
+            resources_update = ResourcesUpdate.model_validate(response)
+            self._resource_cache[resource_id] = resources_update
+            logger.info(f"Fetched and cached resources for ID: {resource_id}")
+            return resources_update
+        return None
+
+    async def get_latest_resources_async(self) -> Optional[ResourcesUpdate]:
+        """Fetch the most recent resource bundle advertised by the server.
+
+        Returns:
+            [`ResourcesUpdate`][agentlightning.ResourcesUpdate] for the
+            newest version, or ``None`` when unavailable.
+        """
+        url = urllib.parse.urljoin(self.endpoint, self._latest_resources_uri)
+        response = await self._request_json_async(url)
+        if response:
+            resources_update = ResourcesUpdate.model_validate(response)
+            # Cache this result as well
+            self._resource_cache[resources_update.resources_id] = resources_update
+            return resources_update
+        return None
+
+    async def post_rollout_async(self, rollout: RolloutLegacy) -> Optional[Dict[str, Any]]:
+        """Submit a completed rollout back to the server.
+
+        Args:
+            rollout: Legacy rollout payload produced by the executor.
+
+        Returns:
+            Parsed JSON response returned by the server, or ``None`` when the request fails.
+        """
+        url = urllib.parse.urljoin(self.endpoint, self._report_rollout_uri)
+        payload = rollout.model_dump(mode="json")
+        return await self._post_json_async(url, payload)
+
+    def _request_json(self, url: str) -> Optional[Dict[str, Any]]:
+        """Perform a blocking ``GET`` request and parse the JSON payload.
+
+        Args:
+            url: Fully qualified URL to query.
+
+        Returns:
+            Parsed JSON body as a dictionary if the request succeeds; otherwise ``None``.
+        """
+        try:
+            response = requests.get(url, timeout=self.timeout, headers=self._default_headers)
+            response.raise_for_status()
+            return response.json()
+        except requests.exceptions.RequestException as e:
+            logger.debug(f"Sync GET request failed for {url}: {e}")
+            return None
+
+    def _post_json(self, url: str, payload: Dict[str, Any]) -> Optional[Dict[str, Any]]:
+        """Perform a blocking ``POST`` request with a JSON payload.
+
+        Args:
+            url: Fully qualified URL that accepts the payload.
+            payload: Dictionary that will be serialized and sent as JSON.
+
+        Returns:
+            Parsed JSON body as a dictionary if the request succeeds; otherwise ``None``.
+        """
+        try:
+            response = requests.post(url, json=payload, timeout=self.timeout, headers=self._default_headers)
+            response.raise_for_status()
+            return response.json()
+        except requests.exceptions.RequestException as e:
+            logger.debug(f"Sync POST request failed for {url}: {e}")
+            return None
+
+    def poll_next_task(self) -> Optional[Task]:
+        """Poll the server synchronously until a task becomes available.
+
+        Returns:
+            The next [`Task`][agentlightning.Task] available for execution, or
+            ``None`` if polling fails.
+        """
+        url = urllib.parse.urljoin(self.endpoint, self._next_task_uri)
+        while True:
+            response = self._request_json(url)
+            if response:
+                task_if_any = TaskIfAny.model_validate(response)
+                if task_if_any.is_available and task_if_any.task:
+                    self.task_count += 1
+                    logger.info(f"[Task {self.task_count} Received] ID: {task_if_any.task.rollout_id}")
+                    return task_if_any.task
+            logger.debug(f"No task available yet. Retrying in {self.poll_interval} seconds...")
+            time.sleep(self.poll_interval)
+
+    def get_resources_by_id(self, resource_id: str) -> Optional[ResourcesUpdate]:
+        """Fetch a specific resource bundle by identifier.
+
+        Args:
+            resource_id: Identifier sourced from the task metadata.
+
+        Returns:
+            Cached or freshly downloaded
+            [`ResourcesUpdate`][agentlightning.ResourcesUpdate], or
+            ``None`` when the server returns an error.
+        """
+        if resource_id in self._resource_cache:
+            logger.debug(f"Found resources '{resource_id}' in cache.")
+            return self._resource_cache[resource_id]
+
+        url = urllib.parse.urljoin(self.endpoint, f"{self._resources_uri}/{resource_id}")
+        response = self._request_json(url)
+        if response:
+            resources_update = ResourcesUpdate.model_validate(response)
+            self._resource_cache[resource_id] = resources_update
+            logger.info(f"Fetched and cached resources for ID: {resource_id}")
+            return resources_update
+        return None
+
+    def get_latest_resources(self) -> Optional[ResourcesUpdate]:
+        """Fetch the most recent resource bundle advertised by the server.
+
+        Returns:
+            [`ResourcesUpdate`][agentlightning.ResourcesUpdate] for the
+            newest version, or ``None`` when unavailable.
+        """
+        url = urllib.parse.urljoin(self.endpoint, self._latest_resources_uri)
+        response = self._request_json(url)
+        if response:
+            resources_update = ResourcesUpdate.model_validate(response)
+            self._resource_cache[resources_update.resources_id] = resources_update
+            return resources_update
+        return None
+
+    def post_rollout(self, rollout: RolloutLegacy) -> Optional[Dict[str, Any]]:
+        """Submit a completed rollout back to the server.
+
+        Args:
+            rollout: Legacy rollout payload produced by the executor.
+
+        Returns:
+            Parsed JSON response returned by the server, or ``None`` when the request fails.
+        """
+        url = urllib.parse.urljoin(self.endpoint, self._report_rollout_uri)
+        payload = rollout.model_dump(mode="json")
+        return self._post_json(url, payload)
+
+
+class DevTaskLoader(AgentLightningClient):
+    """In-memory task loader used for development and integration tests.
+
+    The loader mimics the behavior of the legacy HTTP server by storing tasks and
+    resources locally. Polling methods simply iterate over the provided collection,
+    allowing rapid iteration without provisioning any external infrastructure.
+
+    !!! warning "Deprecated"
+
+        [`DevTaskLoader`][agentlightning.client.DevTaskLoader] is a compatibility shim.
+        Prefer [`Trainer.dev`][agentlightning.Trainer.dev] for new code.
+    """
+
+    def __init__(
+        self,
+        tasks: Union[List[TaskInput], List[Task]],
+        resources: Union[NamedResources, ResourcesUpdate],
+        **kwargs: Any,
+    ):
+        """Initialize the loader with predefined tasks and resources.
+
+        Args:
+            tasks: Sequence of task inputs or preconstructed tasks that will be served in
+                order.
+            resources: Static resources returned for any `resources_id` query.
+            **kwargs: Additional keyword arguments forwarded to the parent client.
+
+        Raises:
+            ValueError: If no tasks are provided or both [`Task`][agentlightning.Task]
+                and [`TaskInput`][agentlightning.TaskInput] instances are mixed.
+        """
+        warnings.warn("DevTaskLoader is deprecated. Please use Trainer.dev instead.", DeprecationWarning)
+        super().__init__(endpoint="local://", **kwargs)
+        self._tasks = tasks.copy()
+        if len(self._tasks) == 0:
+            raise ValueError("DevTaskLoader requires at least one task to be provided.")
+
+        # Check if tasks are mixture of TaskInput and Task
+        if any(isinstance(task, Task) for task in self._tasks):
+            if not all(isinstance(task, Task) for task in self._tasks):
+                raise ValueError("All tasks must be either Task or TaskInput objects.")
+
+        self._task_index = 0
+
+        if isinstance(resources, ResourcesUpdate):
+            self._resources_update = resources
+        else:
+            self._resources_update = ResourcesUpdate(
+                resources_id="local", resources=resources, create_time=time.time(), update_time=time.time(), version=1
+            )
+
+        # Store rollouts posted back to the loader for easy debugging of local runs
+        self._rollouts: List[RolloutLegacy] = []
+
+    @property
+    def rollouts(self) -> List[RolloutLegacy]:
+        """Return the rollouts posted back to the loader during development runs."""
+        return self._rollouts
+
+    def poll_next_task(self) -> Optional[Task]:
+        """Return the next task from the local queue.
+
+        If [`TaskInput`][agentlightning.TaskInput] instances were provided,
+        they are converted into [`Task`][agentlightning.Task] objects on the
+        fly. Otherwise, the preconstructed tasks are returned in sequence.
+
+        Returns:
+            Next task to execute.
+        """
+        if self._task_index >= len(self._tasks):
+            self._task_index = 0
+
+        task_or_input = self._tasks[self._task_index]
+
+        if isinstance(task_or_input, Task):
+            task = task_or_input
+        else:
+            rollout_id = f"local_task_{self._task_index + 1:03d}"
+            task = Task(
+                rollout_id=rollout_id,
+                input=task_or_input,
+                resources_id=self._resources_update.resources_id,
+                create_time=time.time(),
+            )
+
+        self._task_index += 1
+        self.task_count += 1
+        logger.info(f"[Task {self.task_count} Received] Task ID: {task.rollout_id}")
+        return task
+
+    def get_resources_by_id(self, resource_id: str) -> Optional[ResourcesUpdate]:
+        logger.debug(f"DevTaskLoader checking resources for ID: {resource_id}")
+        if resource_id != self._resources_update.resources_id:
+            raise ValueError(
+                f"Resource ID '{resource_id}' not found. Only '{self._resources_update.resources_id}' is available."
+            )
+        return self._resources_update
+
+    def get_latest_resources(self) -> Optional[ResourcesUpdate]:
+        logger.debug("DevTaskLoader returning latest resources.")
+        return self._resources_update
+
+    def post_rollout(self, rollout: RolloutLegacy) -> Optional[Dict[str, Any]]:
+        logger.debug(f"DevTaskLoader received rollout for task: {rollout.rollout_id}")
+        self._rollouts.append(rollout)
+        return {"status": "received", "rollout_id": rollout.rollout_id}
+
+    async def poll_next_task_async(self) -> Optional[Task]:
+        return self.poll_next_task()
+
+    async def get_resources_by_id_async(self, resource_id: str) -> Optional[ResourcesUpdate]:
+        return self.get_resources_by_id(resource_id)
+
+    async def get_latest_resources_async(self) -> Optional[ResourcesUpdate]:
+        return self.get_latest_resources()
+
+    async def post_rollout_async(self, rollout: RolloutLegacy) -> Optional[Dict[str, Any]]:
+        return self.post_rollout(rollout)
+
+    def __repr__(self):
+        return f"DevTaskLoader(num_tasks={len(self._tasks)}, resources={self._resources_update.resources})"