feat: add default limit to /v1/spans and corresponding client metho…

…ds (#3026) - adds default limit of 1000 to /v1/spans - similarly updates default in phoenix.Client and phoenix.Session methods - fixes openapi schema, which was in a broken state
Arize-ai · Apr 29, 2024 · e5698d7 · e5698d7
1 parent dcac234
commit e5698d7
Show file tree

Hide file tree

Showing 7 changed files with 61 additions and 14 deletions.
diff --git a/src/phoenix/server/api/routers/v1/evaluations.py b/src/phoenix/server/api/routers/v1/evaluations.py
@@ -50,8 +50,8 @@ async def post_evaluations(request: Request) -> Response:
         in: query
         schema:
           type: string
+          default: default
         description: The project name to add the evaluation to
-        default: default
     requestBody:
       required: true
       content:
@@ -111,8 +111,8 @@ async def get_evaluations(request: Request) -> Response:
         in: query
         schema:
           type: string
+          default: default
         description: The project name to get evaluations from
-        default: default
     responses:
       200:
         description: Success

diff --git a/src/phoenix/server/api/routers/v1/spans.py b/src/phoenix/server/api/routers/v1/spans.py
@@ -8,6 +8,8 @@
 from phoenix.server.api.routers.utils import df_to_bytes, from_iso_format
 from phoenix.trace.dsl import SpanQuery
 
+DEFAULT_SPAN_LIMIT = 1000
+
 
 # TODO: Add property details to SpanQuery schema
 async def query_spans_handler(request: Request) -> Response:
@@ -21,8 +23,8 @@ async def query_spans_handler(request: Request) -> Response:
         in: query
         schema:
           type: string
+          default: default
         description: The project name to get evaluations from
-        default: default
     requestBody:
       required: true
       content:
@@ -50,11 +52,18 @@ async def query_spans_handler(request: Request) -> Response:
               start_time:
                 type: string
                 format: date-time
+                nullable: true
               stop_time:
                 type: string
                 format: date-time
+                nullable: true
+              limit:
+                type: integer
+                nullable: true
+                default: 1000
               root_spans_only:
                 type: boolean
+                nullable: true
     responses:
       200:
         description: Success
@@ -87,6 +96,7 @@ async def query_spans_handler(request: Request) -> Response:
                     project_name=project_name,
                     start_time=from_iso_format(payload.get("start_time")),
                     stop_time=from_iso_format(payload.get("stop_time")),
+                    limit=payload.get("limit", DEFAULT_SPAN_LIMIT),
                     root_spans_only=payload.get("root_spans_only"),
                 )
             )
@@ -104,9 +114,4 @@ async def content() -> AsyncIterator[bytes]:
 
 
 async def get_spans_handler(request: Request) -> Response:
-    """
-    summary: Deprecated route for querying for spans, use the POST method instead
-    operationId: legacyQuerySpans
-    deprecated: true
-    """
     return await query_spans_handler(request)
diff --git a/src/phoenix/session/client.py b/src/phoenix/session/client.py
@@ -21,7 +21,7 @@
     get_env_port,
     get_env_project_name,
 )
-from phoenix.session.data_extractor import TraceDataExtractor
+from phoenix.session.data_extractor import DEFAULT_SPAN_LIMIT, TraceDataExtractor
 from phoenix.trace import Evaluations, TraceDataset
 from phoenix.trace.dsl import SpanQuery
 from phoenix.trace.otel import encode_span_to_otlp
@@ -65,6 +65,7 @@ def query_spans(
         *queries: SpanQuery,
         start_time: Optional[datetime] = None,
         stop_time: Optional[datetime] = None,
+        limit: Optional[int] = DEFAULT_SPAN_LIMIT,
         root_spans_only: Optional[bool] = None,
         project_name: Optional[str] = None,
     ) -> Optional[Union[pd.DataFrame, List[pd.DataFrame]]]:
@@ -93,6 +94,7 @@ def query_spans(
                 "queries": [q.to_dict() for q in queries],
                 "start_time": _to_iso_format(start_time),
                 "stop_time": _to_iso_format(stop_time),
+                "limit": limit,
                 "root_spans_only": root_spans_only,
             },
         )

diff --git a/src/phoenix/session/data_extractor.py b/src/phoenix/session/data_extractor.py
@@ -8,6 +8,8 @@
 from phoenix.trace.dsl import SpanQuery
 from phoenix.trace.trace_dataset import TraceDataset
 
+DEFAULT_SPAN_LIMIT = 1000
+
 
 class TraceDataExtractor(ABC):
     """
@@ -21,6 +23,7 @@ def query_spans(
         *queries: SpanQuery,
         start_time: Optional[datetime] = None,
         stop_time: Optional[datetime] = None,
+        limit: Optional[int] = DEFAULT_SPAN_LIMIT,
         root_spans_only: Optional[bool] = None,
         project_name: Optional[str] = None,
     ) -> Optional[Union[pd.DataFrame, List[pd.DataFrame]]]: ...
@@ -31,6 +34,7 @@ def get_spans_dataframe(
         *,
         start_time: Optional[datetime] = None,
         stop_time: Optional[datetime] = None,
+        limit: Optional[int] = DEFAULT_SPAN_LIMIT,
         root_spans_only: Optional[bool] = None,
         project_name: Optional[str] = None,
     ) -> Optional[pd.DataFrame]:
@@ -40,6 +44,7 @@ def get_spans_dataframe(
                 SpanQuery().where(filter_condition or ""),
                 start_time=start_time,
                 stop_time=stop_time,
+                limit=limit,
                 root_spans_only=root_spans_only,
                 project_name=project_name,
             ),

diff --git a/src/phoenix/session/session.py b/src/phoenix/session/session.py
@@ -44,7 +44,7 @@
 from phoenix.server.thread_server import ThreadServer
 from phoenix.services import AppService
 from phoenix.session.client import Client
-from phoenix.session.data_extractor import TraceDataExtractor
+from phoenix.session.data_extractor import DEFAULT_SPAN_LIMIT, TraceDataExtractor
 from phoenix.session.evaluation import encode_evaluations
 from phoenix.trace import Evaluations
 from phoenix.trace.dsl.query import SpanQuery
@@ -141,6 +141,7 @@ def query_spans(
         *queries: SpanQuery,
         start_time: Optional[datetime] = None,
         stop_time: Optional[datetime] = None,
+        limit: Optional[int] = DEFAULT_SPAN_LIMIT,
         root_spans_only: Optional[bool] = None,
         project_name: Optional[str] = None,
     ) -> Optional[Union[pd.DataFrame, List[pd.DataFrame]]]:
@@ -174,6 +175,7 @@ def query_spans(
             *queries,
             start_time=start_time,
             stop_time=stop_time,
+            limit=limit,
             root_spans_only=root_spans_only,
             project_name=project_name,
         )

diff --git a/src/phoenix/trace/dsl/query.py b/src/phoenix/trace/dsl/query.py
@@ -38,6 +38,8 @@
 from phoenix.trace.dsl.filter import Projector
 from phoenix.trace.schemas import ATTRIBUTE_PREFIX
 
+DEFAULT_SPAN_LIMIT = 1000
+
 # supported SQL dialects
 _SQLITE: Literal["sqlite"] = "sqlite"
 _POSTGRESQL: Literal["postgresql"] = "postgresql"
@@ -499,6 +501,7 @@ def __call__(
         project_name: Optional[str] = None,
         start_time: Optional[datetime] = None,
         stop_time: Optional[datetime] = None,
+        limit: Optional[int] = DEFAULT_SPAN_LIMIT,
         root_spans_only: Optional[bool] = None,
     ) -> pd.DataFrame:
         if not project_name:
@@ -507,10 +510,11 @@ def __call__(
             return _get_spans_dataframe(
                 session,
                 project_name,
-                self._filter,
-                start_time,
-                stop_time,
-                root_spans_only,
+                span_filter=self._filter,
+                start_time=start_time,
+                stop_time=stop_time,
+                limit=limit,
+                root_spans_only=root_spans_only,
             )
         assert session.bind is not None
         dialect = cast(Literal["sqlite", "postgresql"], session.bind.dialect.name)
@@ -528,6 +532,8 @@ def __call__(
             stmt = stmt.where(start_time <= models.Span.start_time)
         if stop_time:
             stmt = stmt.where(models.Span.start_time < stop_time)
+        if limit is not None:
+            stmt = stmt.limit(limit)
         if root_spans_only:
             parent = aliased(models.Span)
             stmt = stmt.outerjoin(
@@ -662,9 +668,12 @@ def from_dict(
 def _get_spans_dataframe(
     session: Session,
     project_name: str,
+    /,
+    *,
     span_filter: Optional[SpanFilter] = None,
     start_time: Optional[datetime] = None,
     stop_time: Optional[datetime] = None,
+    limit: Optional[int] = DEFAULT_SPAN_LIMIT,
     root_spans_only: Optional[bool] = None,
 ) -> pd.DataFrame:
     # use legacy labels for backward-compatibility
@@ -694,6 +703,8 @@ def _get_spans_dataframe(
         stmt = stmt.where(start_time <= models.Span.start_time)
     if stop_time:
         stmt = stmt.where(models.Span.start_time < stop_time)
+    if limit is not None:
+        stmt = stmt.limit(limit)
     if root_spans_only:
         parent = aliased(models.Span)
         stmt = stmt.outerjoin(

diff --git a/tests/trace/dsl/test_query.py b/tests/trace/dsl/test_query.py
@@ -250,6 +250,28 @@ async def test_stop_time(session: AsyncSession, default_project: None, abc_proje
     )
 
 
+async def test_limit(session: AsyncSession, default_project: None, abc_project: None) -> None:
+    sq = SpanQuery()
+    actual = await session.run_sync(sq, project_name="abc", limit=2)
+    assert actual.index.tolist() == ["234", "345"]
+
+
+async def test_limit_with_select_statement(
+    session: AsyncSession, default_project: None, abc_project: None
+) -> None:
+    sq = SpanQuery().select("context.span_id")
+    expected = pd.DataFrame(
+        {
+            "context.span_id": ["234", "345"],
+        }
+    ).set_index("context.span_id")
+    actual = await session.run_sync(sq, project_name="abc", limit=2)
+    assert_frame_equal(
+        actual.sort_index().sort_index(axis=1),
+        expected.sort_index().sort_index(axis=1),
+    )
+
+
 async def test_filter_for_none(
     session: AsyncSession, default_project: None, abc_project: None
 ) -> None: