Support native advanced queries in LightningStore (#318)

Author: ultmaster

agentlightning/adapter/base.py

@@ -1,6 +1,6 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-from typing import Generic, List, TypeVar
+from typing import Generic, Sequence, TypeVar
 
 from opentelemetry.sdk.trace import ReadableSpan
 
@@ -66,7 +66,7 @@ def adapt(self, source: T_from, /) -> T_to:
         raise NotImplementedError("Adapter.adapt() is not implemented")
 
 
-class OtelTraceAdapter(Adapter[List[ReadableSpan], T_to], Generic[T_to]):
+class OtelTraceAdapter(Adapter[Sequence[ReadableSpan], T_to], Generic[T_to]):
     """Base class for adapters that convert OpenTelemetry trace spans into other formats.
 
     This specialization of [`Adapter`][agentlightning.Adapter] expects a list of
@@ -84,7 +84,7 @@ class OtelTraceAdapter(Adapter[List[ReadableSpan], T_to], Generic[T_to]):
     """
 
 
-class TraceAdapter(Adapter[List[Span], T_to], Generic[T_to]):
+class TraceAdapter(Adapter[Sequence[Span], T_to], Generic[T_to]):
     """Base class for adapters that convert trace spans into other formats.
 
     This class specializes [`Adapter`][agentlightning.Adapter] for working with

agentlightning/adapter/messages.py

@@ -4,7 +4,7 @@
 
 import json
 from collections import defaultdict
-from typing import TYPE_CHECKING, Any, Dict, Generator, Iterable, List, Optional, TypedDict, Union, cast
+from typing import TYPE_CHECKING, Any, Dict, Generator, Iterable, List, Optional, Sequence, TypedDict, Union, cast
 
 from pydantic import TypeAdapter
 
@@ -208,7 +208,7 @@ class TraceToMessages(TraceAdapter[List[OpenAIMessages]]):
         children of the associated completion span.
     """
 
-    def get_tool_calls(self, completion: Span, all_spans: List[Span], /) -> Iterable[Dict[str, Any]]:
+    def get_tool_calls(self, completion: Span, all_spans: Sequence[Span], /) -> Iterable[Dict[str, Any]]:
         """Yield tool call payloads for a completion span.
 
         Args:
@@ -231,7 +231,7 @@ def get_tool_calls(self, completion: Span, all_spans: List[Span], /) -> Iterable
             if tool_call:
                 yield tool_call
 
-    def adapt(self, source: List[Span], /) -> List[OpenAIMessages]:
+    def adapt(self, source: Sequence[Span], /) -> List[OpenAIMessages]:
         """Transform trace spans into OpenAI chat payloads.
 
         Args:

agentlightning/adapter/triplet.py

@@ -6,7 +6,7 @@
 import logging
 import re
 from enum import Enum
-from typing import Any, Dict, List, Optional, Tuple, Union, cast
+from typing import Any, Dict, List, Optional, Sequence, Tuple, Union, cast
 
 from opentelemetry.sdk.trace import ReadableSpan
 from pydantic import BaseModel
@@ -670,7 +670,7 @@ def visualize(
         trace_tree.visualize(filename, interested_span_match=interested_span_match)
         return trace_tree
 
-    def adapt(self, source: Union[List[Span], List[ReadableSpan]], /) -> List[Triplet]:  # type: ignore
+    def adapt(self, source: Union[Sequence[Span], Sequence[ReadableSpan]], /) -> List[Triplet]:  # type: ignore
         """Convert tracer spans into [`Triplet`][agentlightning.Triplet] trajectories.
 
         Args:
@@ -800,7 +800,7 @@ def _request_id_from_attrs(self, attrs: Dict[str, Any]) -> Optional[str]:
         rid = attrs.get("gen_ai.response.id") or attrs.get("llm.hosted_vllm.id")
         return str(rid) if isinstance(rid, str) and rid else None
 
-    def adapt(self, source: List[Span], /) -> List[Triplet]:  # type: ignore
+    def adapt(self, source: Sequence[Span], /) -> List[Triplet]:  # type: ignore
         """Convert LLM Proxy spans into [`Triplet`][agentlightning.Triplet] trajectories.
 
         Args:

agentlightning/algorithm/fast.py

@@ -143,7 +143,7 @@ async def _enqueue_rollouts(
         store = self.get_store()
 
         for index in train_indices + val_indices:
-            queuing_rollouts = await store.query_rollouts(status=["queuing", "requeuing"])
+            queuing_rollouts = await store.query_rollouts(status_in=["queuing", "requeuing"])
             if len(queuing_rollouts) <= 1:
                 # Only enqueue a new rollout when there is at most 1 rollout in the queue.
                 sample = dataset[index]
@@ -222,7 +222,7 @@ async def run(
                     f"Processing index {index}. {len(train_indices)} train indices and {len(val_indices)} val indices in total."
                 )
                 while True:
-                    queuing_rollouts = await store.query_rollouts(status=["queuing", "requeuing"])
+                    queuing_rollouts = await store.query_rollouts(status_in=["queuing", "requeuing"])
                     if len(queuing_rollouts) <= self.max_queue_length:
                         # Only enqueue a new rollout when there is at most "max_queue_length" rollout in the queue.
                         sample = concatenated_dataset[index]

agentlightning/store/base.py

@@ -18,6 +18,7 @@
     Span,
     TaskInput,
     Worker,
+    WorkerStatus,
 )
 
 
@@ -292,30 +293,77 @@ async def add_otel_span(
         raise NotImplementedError()
 
     async def query_rollouts(
-        self, *, status: Optional[Sequence[RolloutStatus]] = None, rollout_ids: Optional[Sequence[str]] = None
-    ) -> List[Rollout]:
+        self,
+        *,
+        status_in: Optional[Sequence[RolloutStatus]] = None,
+        rollout_id_in: Optional[Sequence[str]] = None,
+        rollout_id_contains: Optional[str] = None,
+        filter_logic: Literal["and", "or"] = "and",
+        sort_by: Optional[str] = None,
+        sort_order: Literal["asc", "desc"] = "asc",
+        limit: int = -1,
+        offset: int = 0,
+        # Deprecated fields
+        status: Optional[Sequence[RolloutStatus]] = None,
+        rollout_ids: Optional[Sequence[str]] = None,
+    ) -> Sequence[Rollout]:
         """Retrieve rollouts filtered by status and/or explicit identifiers.
 
+        This interface supports structured filtering, sorting, and pagination so
+        callers can build simple dashboards without copying data out of the
+        store. The legacy parameters `status` and `rollout_ids` remain valid and
+        are treated as aliases for `status_in` and `rollout_id_in`
+        respectively—when both the new and deprecated parameters are supplied
+        the new parameters take precedence.
+
         Args:
-            status: Optional whitelist of [`RolloutStatus`][agentlightning.RolloutStatus] values.
-            rollout_ids: Optional whitelist of rollout identifiers to include.
+            status_in: Optional whitelist of [`RolloutStatus`][agentlightning.RolloutStatus] values.
+            rollout_id_in: Optional whitelist of rollout identifiers to include.
+            rollout_id_contains: Optional substring match for rollout identifiers.
+            filter_logic: Logical operator to combine filters.
+            sort_by: Optional field to sort by. Must reference a numeric or string
+                field on [`Rollout`][agentlightning.Rollout].
+            sort_order: Direction to sort when `sort_by` is provided.
+            limit: Maximum number of rows to return. Use `-1` for "no limit".
+            offset: Number of rows to skip before returning results.
+            status: Deprecated field. Use `status_in` instead.
+            rollout_ids: Deprecated field. Use `rollout_id_in` instead.
 
         Returns:
-            A list of matching rollouts. Ordering is backend-defined but must be deterministic.
+            A sequence of matching rollouts (or [`AttemptedRollout`][agentlightning.AttemptedRollout]
+            when attempts exist). Ordering is deterministic when `sort_by` is set.
+            The return value is not guaranteed to be a list.
 
         Raises:
             NotImplementedError: Subclasses must implement the query.
         """
         raise NotImplementedError()
 
-    async def query_attempts(self, rollout_id: str) -> List[Attempt]:
+    async def query_attempts(
+        self,
+        rollout_id: str,
+        *,
+        sort_by: Optional[str] = "sequence_id",
+        sort_order: Literal["asc", "desc"] = "asc",
+        limit: int = -1,
+        offset: int = 0,
+    ) -> Sequence[Attempt]:
         """Return every attempt ever created for `rollout_id` in ascending sequence order.
 
+        The parameters allow callers to re-order or paginate the attempts so that
+        large retry histories can be streamed lazily.
+
         Args:
             rollout_id: Identifier of the rollout being inspected.
+            sort_by: Field to sort by. Must be a numeric or string field of
+                [`Attempt`][agentlightning.Attempt]. Defaults to `sequence_id` (oldest first).
+            sort_order: Order to sort by.
+            limit: Limit on the number of results. `-1` for unlimited.
+            offset: Offset into the results.
 
         Returns:
-            Attempts sorted by `sequence_id` (oldest first). Returns an empty list when none exist.
+            Sequence of Attempts. Returns an empty sequence when none exist.
+            The return value is not guaranteed to be a list.
 
         Raises:
             NotImplementedError: Subclasses must implement the query.
@@ -352,11 +400,35 @@ async def get_latest_attempt(self, rollout_id: str) -> Optional[Attempt]:
         """
         raise NotImplementedError()
 
-    async def query_resources(self) -> List[ResourcesUpdate]:
+    async def query_resources(
+        self,
+        *,
+        resources_id: Optional[str] = None,
+        resources_id_contains: Optional[str] = None,
+        # Filter logic is not supported here because I can't see why it's needed.
+        sort_by: Optional[str] = None,
+        sort_order: Literal["asc", "desc"] = "asc",
+        limit: int = -1,
+        offset: int = 0,
+    ) -> Sequence[ResourcesUpdate]:
         """List every stored resource snapshot in insertion order.
 
+        Supports lightweight filtering, sorting, and pagination for embedding in
+        dashboards.
+
+        Args:
+            resources_id: Optional identifier of the resources to include.
+            resources_id_contains: Optional substring match for resources identifiers.
+            sort_by: Optional field to sort by (must be numeric or string on
+                [`ResourcesUpdate`][agentlightning.ResourcesUpdate]).
+            sort_order: Order to sort by.
+            limit: Limit on the number of results. `-1` for unlimited.
+            offset: Offset into the results.
+
         Returns:
-            A chronological list of [`ResourcesUpdate`][agentlightning.ResourcesUpdate] objects.
+            [`ResourcesUpdate`][agentlightning.ResourcesUpdate] objects.
+            By default, resources are sorted in a deterministic but undefined order.
+            The return value is not guaranteed to be a list.
 
         Raises:
             NotImplementedError: Subclasses must implement retrieval.
@@ -439,19 +511,61 @@ async def wait_for_rollouts(self, *, rollout_ids: List[str], timeout: Optional[f
         """
         raise NotImplementedError()
 
-    async def query_spans(self, rollout_id: str, attempt_id: str | Literal["latest"] | None = None) -> List[Span]:
+    async def query_spans(
+        self,
+        rollout_id: str,
+        attempt_id: str | Literal["latest"] | None = None,
+        *,
+        # Filtering
+        trace_id: Optional[str] = None,
+        trace_id_contains: Optional[str] = None,
+        span_id: Optional[str] = None,
+        span_id_contains: Optional[str] = None,
+        parent_id: Optional[str] = None,
+        parent_id_contains: Optional[str] = None,
+        name: Optional[str] = None,
+        name_contains: Optional[str] = None,
+        filter_logic: Literal["and", "or"] = "and",
+        # Pagination
+        limit: int = -1,
+        offset: int = 0,
+        # Sorting
+        sort_by: Optional[str] = "sequence_id",
+        sort_order: Literal["asc", "desc"] = "asc",
+    ) -> Sequence[Span]:
         """Return the stored spans for a rollout, optionally scoped to one attempt.
 
-        Spans must be returned in ascending `sequence_id` order. Implementations may raise
-        a `RuntimeError` when spans were evicted or expired.
+        Supports a handful of filters that cover the most common debugging
+        scenarios (matching `trace_id`/`span_id`/`parent_id` or substring
+        matches on the span name). `attempt_id="latest"` acts as a convenience
+        that resolves the most recent attempt before evaluating filters. When
+        `attempt_id=None`, spans across every attempt are eligible. By default
+        results are sorted by `sequence_id` (oldest first). Implementations may
+        raise a `RuntimeError` when spans were evicted or expired.
 
         Args:
             rollout_id: Identifier of the rollout being inspected.
             attempt_id: Attempt identifier to filter by. Pass `"latest"` to retrieve only the
                 most recent attempt, or `None` to return all spans across attempts.
+            trace_id: Optional trace ID to filter by.
+            trace_id_contains: Optional substring match for trace IDs.
+            span_id: Optional span ID to filter by.
+            span_id_contains: Optional substring match for span IDs.
+            parent_id: Optional parent span ID to filter by.
+            parent_id_contains: Optional substring match for parent span IDs.
+            name: Optional span name to filter by.
+            name_contains: Optional substring match for span names.
+            filter_logic: Logical operator to combine the optional filters above.
+                The `rollout_id` argument is always applied with AND semantics.
+            limit: Limit on the number of results. `-1` for unlimited.
+            offset: Offset into the results.
+            sort_by: Field to sort by. Must be a numeric or string field of
+                [`Span`][agentlightning.Span].
+            sort_order: Order to sort by.
 
         Returns:
             An ordered list of spans (possibly empty).
+            The return value is not guaranteed to be a list.
 
         Raises:
             NotImplementedError: Subclasses must implement the query.
@@ -578,11 +692,29 @@ async def update_attempt(
 
     async def query_workers(
         self,
-    ) -> List[Worker]:
+        *,
+        status_in: Optional[Sequence[WorkerStatus]] = None,
+        worker_id_contains: Optional[str] = None,
+        filter_logic: Literal["and", "or"] = "and",
+        sort_by: Optional[str] = None,
+        sort_order: Literal["asc", "desc"] = "asc",
+        limit: int = -1,
+        offset: int = 0,
+    ) -> Sequence[Worker]:
         """Query all workers in the system.
 
+        Args:
+            status_in: Optional whitelist of [`WorkerStatus`][agentlightning.WorkerStatus] values.
+            worker_id_contains: Optional substring match for worker identifiers.
+            filter_logic: Logical operator to combine the optional filters above.
+            sort_by: Field to sort by. Must be a numeric or string field of [`Worker`][agentlightning.Worker].
+            sort_order: Order to sort by.
+            limit: Limit on the number of results. `-1` for unlimited.
+            offset: Offset into the results.
+
         Returns:
-            A list of all workers.
+            Sequence of Workers. Returns an empty sequence when none exist.
+            The return value is not guaranteed to be a list.
         """
         raise NotImplementedError()