Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: add default limit to /v1/spans and corresponding client methods #3026

Merged
merged 6 commits into from
Apr 29, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions src/phoenix/server/api/routers/v1/evaluations.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Comment on lines +53 to -54
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fix to OpenAPI schema.

requestBody:
required: true
content:
Expand Down Expand Up @@ -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
Comment on lines +114 to -115
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fix to OpenAPI schema

responses:
200:
description: Success
Expand Down
17 changes: 11 additions & 6 deletions src/phoenix/server/api/routers/v1/spans.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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
Comment on lines +26 to -25
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fix to OpenAPI schema

requestBody:
required: true
content:
Expand Down Expand Up @@ -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
axiomofjoy marked this conversation as resolved.
Show resolved Hide resolved
default: 1000
root_spans_only:
type: boolean
nullable: true
responses:
200:
description: Success
Expand Down Expand Up @@ -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"),
)
)
Expand All @@ -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
"""
Comment on lines -107 to -111
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Remove this route from the OpenAPI schema since it is being deprecated.

return await query_spans_handler(request)
4 changes: 3 additions & 1 deletion src/phoenix/session/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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]]]:
Expand Down Expand Up @@ -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,
},
)
Expand Down
5 changes: 5 additions & 0 deletions src/phoenix/session/data_extractor.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,8 @@
from phoenix.trace.dsl import SpanQuery
from phoenix.trace.trace_dataset import TraceDataset

DEFAULT_SPAN_LIMIT = 1000


class TraceDataExtractor(ABC):
"""
Expand All @@ -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]]]: ...
Expand All @@ -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]:
Expand All @@ -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,
),
Expand Down
4 changes: 3 additions & 1 deletion src/phoenix/session/session.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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]]]:
Expand Down Expand Up @@ -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,
)
Expand Down
19 changes: 15 additions & 4 deletions src/phoenix/trace/dsl/query.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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:
Expand All @@ -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)
Expand All @@ -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(
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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(
Expand Down
22 changes: 22 additions & 0 deletions tests/trace/dsl/test_query.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
Loading