microsoft--agent-lightning/agentlightning/client.py

# 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})"