arize-phoenix 3.24.0__py3-none-any.whl → 4.0.0__py3-none-any.whl

phoenix/server/api/input_types/SpanSort.py

@@ -1,30 +1,88 @@
-from enum import Enum
-from functools import partial
-from typing import Any, Iterable, Iterator, Optional, Protocol
+from dataclasses import dataclass
+from enum import Enum, auto
+from typing import Any, Optional, Protocol
 
-import pandas as pd
 import strawberry
 from openinference.semconv.trace import SpanAttributes
+from sqlalchemy import and_, desc, nulls_last
+from sqlalchemy.orm import InstrumentedAttribute
+from sqlalchemy.sql.expression import Select
 from strawberry import UNSET
 from typing_extensions import assert_never
 
 import phoenix.trace.v1 as pb
-from phoenix.core.project import WrappedSpan
+from phoenix.db import models
+from phoenix.server.api.types.pagination import CursorSortColumnDataType
 from phoenix.server.api.types.SortDir import SortDir
-from phoenix.trace.schemas import ComputedAttributes, SpanID
+from phoenix.trace.schemas import SpanID
+
+LLM_TOKEN_COUNT_PROMPT = SpanAttributes.LLM_TOKEN_COUNT_PROMPT.split(".")
+LLM_TOKEN_COUNT_COMPLETION = SpanAttributes.LLM_TOKEN_COUNT_COMPLETION.split(".")
+LLM_TOKEN_COUNT_TOTAL = SpanAttributes.LLM_TOKEN_COUNT_TOTAL.split(".")
 
 
 @strawberry.enum
 class SpanColumn(Enum):
-    startTime = "start_time"
-    endTime = "end_time"
-    latencyMs = ComputedAttributes.LATENCY_MS
-    tokenCountTotal = SpanAttributes.LLM_TOKEN_COUNT_TOTAL
-    tokenCountPrompt = SpanAttributes.LLM_TOKEN_COUNT_PROMPT
-    tokenCountCompletion = SpanAttributes.LLM_TOKEN_COUNT_COMPLETION
-    cumulativeTokenCountTotal = ComputedAttributes.CUMULATIVE_LLM_TOKEN_COUNT_TOTAL
-    cumulativeTokenCountPrompt = ComputedAttributes.CUMULATIVE_LLM_TOKEN_COUNT_PROMPT
-    cumulativeTokenCountCompletion = ComputedAttributes.CUMULATIVE_LLM_TOKEN_COUNT_COMPLETION
+    startTime = auto()
+    endTime = auto()
+    latencyMs = auto()
+    tokenCountTotal = auto()
+    tokenCountPrompt = auto()
+    tokenCountCompletion = auto()
+    cumulativeTokenCountTotal = auto()
+    cumulativeTokenCountPrompt = auto()
+    cumulativeTokenCountCompletion = auto()
+
+    @property
+    def column_name(self) -> str:
+        return "f{self.name}_span_sort_column"
+
+    @property
+    def orm_expression(self) -> Any:
+        expr: Any
+        if self is SpanColumn.startTime:
+            expr = models.Span.start_time
+        elif self is SpanColumn.endTime:
+            expr = models.Span.end_time
+        elif self is SpanColumn.latencyMs:
+            expr = models.Span.latency_ms
+        elif self is SpanColumn.tokenCountTotal:
+            expr = models.Span.attributes[LLM_TOKEN_COUNT_TOTAL].as_float()
+        elif self is SpanColumn.tokenCountPrompt:
+            expr = models.Span.attributes[LLM_TOKEN_COUNT_PROMPT].as_float()
+        elif self is SpanColumn.tokenCountCompletion:
+            expr = models.Span.attributes[LLM_TOKEN_COUNT_COMPLETION].as_float()
+        elif self is SpanColumn.cumulativeTokenCountTotal:
+            expr = (
+                models.Span.cumulative_llm_token_count_prompt
+                + models.Span.cumulative_llm_token_count_completion
+            )
+        elif self is SpanColumn.cumulativeTokenCountPrompt:
+            expr = models.Span.cumulative_llm_token_count_prompt
+        elif self is SpanColumn.cumulativeTokenCountCompletion:
+            expr = models.Span.cumulative_llm_token_count_completion
+        else:
+            assert_never(self)
+        return expr.label(self.column_name)
+
+    @property
+    def data_type(self) -> CursorSortColumnDataType:
+        if (
+            self is SpanColumn.cumulativeTokenCountTotal
+            or self is SpanColumn.cumulativeTokenCountPrompt
+            or self is SpanColumn.cumulativeTokenCountCompletion
+        ):
+            return CursorSortColumnDataType.INT
+        if (
+            self is SpanColumn.latencyMs
+            or self is SpanColumn.tokenCountTotal
+            or self is SpanColumn.tokenCountPrompt
+            or self is SpanColumn.tokenCountCompletion
+        ):
+            return CursorSortColumnDataType.FLOAT
+        if self is SpanColumn.startTime or self is SpanColumn.endTime:
+            return CursorSortColumnDataType.DATETIME
+        assert_never(self)
 
 
 @strawberry.enum
@@ -32,6 +90,29 @@ class EvalAttr(Enum):
     score = "score"
     label = "label"
 
+    @property
+    def column_name(self) -> str:
+        return f"{self.value}_eval_sort_column"
+
+    @property
+    def orm_expression(self) -> Any:
+        expr: InstrumentedAttribute[Any]
+        if self is EvalAttr.score:
+            expr = models.SpanAnnotation.score
+        elif self is EvalAttr.label:
+            expr = models.SpanAnnotation.label
+        else:
+            assert_never(self)
+        return expr.label(self.column_name)
+
+    @property
+    def data_type(self) -> CursorSortColumnDataType:
+        if self is EvalAttr.label:
+            return CursorSortColumnDataType.STRING
+        if self is EvalAttr.score:
+            return CursorSortColumnDataType.FLOAT
+        assert_never(self)
+
 
 @strawberry.input
 class EvalResultKey:
@@ -43,6 +124,15 @@ class SupportsGetSpanEvaluation(Protocol):
     def get_span_evaluation(self, span_id: SpanID, name: str) -> Optional[pb.Evaluation]: ...
 
 
+@dataclass(frozen=True)
+class SpanSortConfig:
+    stmt: Select[Any]
+    orm_expression: Any
+    dir: SortDir
+    column_name: str
+    column_data_type: CursorSortColumnDataType
+
+
 @strawberry.input(
     description="The sort key and direction for span connections. Must "
     "specify one and only one of either `col` or `evalResultKey`."
@@ -52,58 +142,38 @@ class SpanSort:
     eval_result_key: Optional[EvalResultKey] = UNSET
     dir: SortDir
 
-    def __call__(
-        self,
-        spans: Iterable[WrappedSpan],
-        evals: Optional[SupportsGetSpanEvaluation] = None,
-    ) -> Iterator[WrappedSpan]:
-        """
-        Sorts the spans by the given key (column or eval) and direction
-        """
-        if self.eval_result_key:
-            get_sort_key_value = partial(
-                _get_eval_result_value,
-                eval_name=self.eval_result_key.name,
-                eval_attr=self.eval_result_key.attr,
-                evals=evals,
+    def update_orm_expr(self, stmt: Select[Any]) -> SpanSortConfig:
+        if (col := self.col) and not self.eval_result_key:
+            expr = col.orm_expression
+            stmt = stmt.add_columns(expr)
+            if self.dir == SortDir.desc:
+                expr = desc(expr)
+            return SpanSortConfig(
+                stmt=stmt.order_by(nulls_last(expr)),
+                orm_expression=col.orm_expression,
+                dir=self.dir,
+                column_name=col.column_name,
+                column_data_type=col.data_type,
             )
-        else:
-            get_sort_key_value = partial(
-                _get_column_value,
-                span_column=self.col or SpanColumn.startTime,
+        if (eval_result_key := self.eval_result_key) and not col:
+            eval_name = eval_result_key.name
+            eval_attr = eval_result_key.attr
+            expr = eval_result_key.attr.orm_expression
+            stmt = stmt.add_columns(expr)
+            if self.dir == SortDir.desc:
+                expr = desc(expr)
+            stmt = stmt.join(
+                models.SpanAnnotation,
+                onclause=and_(
+                    models.SpanAnnotation.span_rowid == models.Span.id,
+                    models.SpanAnnotation.name == eval_name,
+                ),
+            ).order_by(expr)
+            return SpanSortConfig(
+                stmt=stmt,
+                orm_expression=eval_result_key.attr.orm_expression,
+                dir=self.dir,
+                column_name=eval_attr.column_name,
+                column_data_type=eval_attr.data_type,
             )
-        yield from pd.Series(spans, dtype=object).sort_values(
-            key=lambda series: series.apply(get_sort_key_value),
-            ascending=self.dir.value == SortDir.asc.value,
-        )
-
-
-def _get_column_value(span: WrappedSpan, span_column: SpanColumn) -> Any:
-    if span_column is SpanColumn.startTime:
-        return span.start_time
-    if span_column is SpanColumn.endTime:
-        return span.end_time
-    return span[span_column.value]
-
-
-def _get_eval_result_value(
-    span: WrappedSpan,
-    eval_name: str,
-    eval_attr: EvalAttr,
-    evals: Optional[SupportsGetSpanEvaluation] = None,
-) -> Any:
-    """
-    Returns the evaluation result for the given span
-    """
-    if evals is None:
-        return None
-    span_id = span.context.span_id
-    evaluation = evals.get_span_evaluation(span_id, eval_name)
-    if evaluation is None:
-        return None
-    result = evaluation.result
-    if eval_attr is EvalAttr.score:
-        return result.score.value if result.HasField("score") else None
-    if eval_attr is EvalAttr.label:
-        return result.label.value if result.HasField("label") else None
-    assert_never(eval_attr)
+        raise ValueError("Exactly one of `col` or `evalResultKey` must be specified on `SpanSort`.")

phoenix/server/api/routers/v1/__init__.py

@@ -0,0 +1,11 @@
+from starlette.routing import Route
+
+from . import evaluations, spans, traces
+
+V1_ROUTES = [
+    Route("/v1/evaluations", evaluations.post_evaluations, methods=["POST"]),
+    Route("/v1/evaluations", evaluations.get_evaluations, methods=["GET"]),
+    Route("/v1/traces", traces.post_traces, methods=["POST"]),
+    Route("/v1/spans", spans.query_spans_handler, methods=["POST"]),
+    Route("/v1/spans", spans.get_spans_handler, methods=["GET"]),
+]

phoenix/server/api/routers/v1/evaluations.py

@@ -0,0 +1,275 @@
+import gzip
+from itertools import chain
+from typing import AsyncContextManager, Callable, Iterator, Tuple
+
+import pandas as pd
+import pyarrow as pa
+from google.protobuf.message import DecodeError
+from pandas import DataFrame
+from sqlalchemy import select
+from sqlalchemy.engine import Connectable
+from sqlalchemy.ext.asyncio import (
+    AsyncSession,
+)
+from starlette.background import BackgroundTask
+from starlette.datastructures import State
+from starlette.requests import Request
+from starlette.responses import Response, StreamingResponse
+from starlette.status import (
+    HTTP_403_FORBIDDEN,
+    HTTP_404_NOT_FOUND,
+    HTTP_415_UNSUPPORTED_MEDIA_TYPE,
+    HTTP_422_UNPROCESSABLE_ENTITY,
+)
+from typing_extensions import TypeAlias
+
+import phoenix.trace.v1 as pb
+from phoenix.config import DEFAULT_PROJECT_NAME
+from phoenix.db import models
+from phoenix.exceptions import PhoenixEvaluationNameIsMissing
+from phoenix.server.api.routers.utils import table_to_bytes
+from phoenix.session.evaluation import encode_evaluations
+from phoenix.trace.span_evaluations import (
+    DocumentEvaluations,
+    Evaluations,
+    SpanEvaluations,
+    TraceEvaluations,
+)
+
+EvaluationName: TypeAlias = str
+
+
+async def post_evaluations(request: Request) -> Response:
+    """
+    summary: Add evaluations to a span, trace, or document
+    operationId: addEvaluations
+    tags:
+      - evaluations
+    parameters:
+      - name: project-name
+        in: query
+        schema:
+          type: string
+          default: default
+        description: The project name to add the evaluation to
+    requestBody:
+      required: true
+      content:
+        application/x-protobuf:
+          schema:
+            type: string
+            format: binary
+        application/x-pandas-arrow:
+          schema:
+            type: string
+            format: binary
+    responses:
+      200:
+        description: Success
+      403:
+        description: Forbidden
+      415:
+        description: Unsupported content type, only gzipped protobuf and pandas-arrow are supported
+      422:
+        description: Request body is invalid
+    """
+    if request.app.state.read_only:
+        return Response(status_code=HTTP_403_FORBIDDEN)
+    content_type = request.headers.get("content-type")
+    if content_type == "application/x-pandas-arrow":
+        return await _process_pyarrow(request)
+    if content_type != "application/x-protobuf":
+        return Response("Unsupported content type", status_code=HTTP_415_UNSUPPORTED_MEDIA_TYPE)
+    body = await request.body()
+    content_encoding = request.headers.get("content-encoding")
+    if content_encoding == "gzip":
+        body = gzip.decompress(body)
+    elif content_encoding:
+        return Response("Unsupported content encoding", status_code=HTTP_415_UNSUPPORTED_MEDIA_TYPE)
+    evaluation = pb.Evaluation()
+    try:
+        evaluation.ParseFromString(body)
+    except DecodeError:
+        return Response("Request body is invalid", status_code=HTTP_422_UNPROCESSABLE_ENTITY)
+    if not evaluation.name.strip():
+        return Response(
+            "Evaluation name must not be blank/empty",
+            status_code=HTTP_422_UNPROCESSABLE_ENTITY,
+        )
+    await request.state.queue_evaluation_for_bulk_insert(evaluation)
+    return Response()
+
+
+async def get_evaluations(request: Request) -> Response:
+    """
+    summary: Get evaluations from Phoenix
+    operationId: getEvaluation
+    tags:
+      - evaluations
+    parameters:
+      - name: project-name
+        in: query
+        schema:
+          type: string
+          default: default
+        description: The project name to get evaluations from
+    responses:
+      200:
+        description: Success
+      404:
+        description: Not found
+    """
+    project_name = (
+        request.query_params.get("project-name")
+        # read from headers for backwards compatibility
+        or request.headers.get("project-name")
+        or DEFAULT_PROJECT_NAME
+    )
+
+    db: Callable[[], AsyncContextManager[AsyncSession]] = request.app.state.db
+    async with db() as session:
+        connection = await session.connection()
+        trace_evals_dataframe = await connection.run_sync(
+            _read_sql_trace_evaluations_into_dataframe,
+            project_name,
+        )
+        span_evals_dataframe = await connection.run_sync(
+            _read_sql_span_evaluations_into_dataframe,
+            project_name,
+        )
+        document_evals_dataframe = await connection.run_sync(
+            _read_sql_document_evaluations_into_dataframe,
+            project_name,
+        )
+    if (
+        trace_evals_dataframe.empty
+        and span_evals_dataframe.empty
+        and document_evals_dataframe.empty
+    ):
+        return Response(status_code=HTTP_404_NOT_FOUND)
+
+    evals = chain(
+        map(
+            lambda args: TraceEvaluations(*args),
+            _groupby_eval_name(trace_evals_dataframe),
+        ),
+        map(
+            lambda args: SpanEvaluations(*args),
+            _groupby_eval_name(span_evals_dataframe),
+        ),
+        map(
+            lambda args: DocumentEvaluations(*args),
+            _groupby_eval_name(document_evals_dataframe),
+        ),
+    )
+    bytestream = map(lambda evals: table_to_bytes(evals.to_pyarrow_table()), evals)
+    return StreamingResponse(
+        content=bytestream,
+        media_type="application/x-pandas-arrow",
+    )
+
+
+async def _process_pyarrow(request: Request) -> Response:
+    body = await request.body()
+    try:
+        reader = pa.ipc.open_stream(body)
+    except pa.ArrowInvalid:
+        return Response(
+            content="Request body is not valid pyarrow",
+            status_code=HTTP_422_UNPROCESSABLE_ENTITY,
+        )
+    try:
+        evaluations = Evaluations.from_pyarrow_reader(reader)
+    except Exception as e:
+        if isinstance(e, PhoenixEvaluationNameIsMissing):
+            return Response(
+                "Evaluation name must not be blank/empty",
+                status_code=HTTP_422_UNPROCESSABLE_ENTITY,
+            )
+        return Response(
+            content="Invalid data in request body",
+            status_code=HTTP_422_UNPROCESSABLE_ENTITY,
+        )
+    return Response(background=BackgroundTask(_add_evaluations, request.state, evaluations))
+
+
+async def _add_evaluations(state: State, evaluations: Evaluations) -> None:
+    for evaluation in encode_evaluations(evaluations):
+        await state.queue_evaluation_for_bulk_insert(evaluation)
+
+
+def _read_sql_trace_evaluations_into_dataframe(
+    connectable: Connectable,
+    project_name: str,
+) -> DataFrame:
+    """
+    Reads a project's trace evaluations into a pandas dataframe.
+
+    Inputs a synchronous connectable to pandas.read_sql since it does not
+    support async connectables. For more information, see:
+
+    https://stackoverflow.com/questions/70848256/how-can-i-use-pandas-read-sql-on-an-async-connection
+    """
+    return pd.read_sql(
+        select(models.TraceAnnotation, models.Trace.trace_id)
+        .join_from(models.TraceAnnotation, models.Trace)
+        .join_from(models.Trace, models.Project)
+        .where(models.Project.name == project_name)
+        .where(models.TraceAnnotation.annotator_kind == "LLM"),
+        connectable,
+        index_col="trace_id",
+    )
+
+
+def _read_sql_span_evaluations_into_dataframe(
+    connectable: Connectable,
+    project_name: str,
+) -> DataFrame:
+    """
+    Reads a project's span evaluations into a pandas dataframe.
+
+    Inputs a synchronous connectable to pandas.read_sql since it does not
+    support async connectables. For more information, see:
+
+    https://stackoverflow.com/questions/70848256/how-can-i-use-pandas-read-sql-on-an-async-connection
+    """
+    return pd.read_sql_query(
+        select(models.SpanAnnotation, models.Span.span_id)
+        .join_from(models.SpanAnnotation, models.Span)
+        .join_from(models.Span, models.Trace)
+        .join_from(models.Trace, models.Project)
+        .where(models.Project.name == project_name)
+        .where(models.SpanAnnotation.annotator_kind == "LLM"),
+        connectable,
+        index_col="span_id",
+    )
+
+
+def _read_sql_document_evaluations_into_dataframe(
+    connectable: Connectable,
+    project_name: str,
+) -> DataFrame:
+    """
+    Reads a project's document evaluations into a pandas dataframe.
+
+    Inputs a synchronous connectable to pandas.read_sql since it does not
+    support async connectables. For more information, see:
+
+    https://stackoverflow.com/questions/70848256/how-can-i-use-pandas-read-sql-on-an-async-connection
+    """
+    return pd.read_sql(
+        select(models.DocumentAnnotation, models.Span.span_id)
+        .join_from(models.DocumentAnnotation, models.Span)
+        .join_from(models.Span, models.Trace)
+        .join_from(models.Trace, models.Project)
+        .where(models.Project.name == project_name)
+        .where(models.DocumentAnnotation.annotator_kind == "LLM"),
+        connectable,
+    ).set_index(["span_id", "document_position"])
+
+
+def _groupby_eval_name(
+    evals_dataframe: DataFrame,
+) -> Iterator[Tuple[EvaluationName, DataFrame]]:
+    for eval_name, evals_dataframe_for_name in evals_dataframe.groupby("name", as_index=False):
+        yield str(eval_name), evals_dataframe_for_name

phoenix/server/api/routers/v1/spans.py

@@ -0,0 +1,126 @@
+from datetime import timezone
+from typing import AsyncIterator
+
+from starlette.requests import Request
+from starlette.responses import Response, StreamingResponse
+from starlette.status import HTTP_404_NOT_FOUND, HTTP_422_UNPROCESSABLE_ENTITY
+
+from phoenix.config import DEFAULT_PROJECT_NAME
+from phoenix.datetime_utils import normalize_datetime
+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:
+    """
+    summary: Query spans using query DSL
+    operationId: querySpans
+    tags:
+      - spans
+    parameters:
+      - name: project-name
+        in: query
+        schema:
+          type: string
+          default: default
+        description: The project name to get evaluations from
+    requestBody:
+      required: true
+      content:
+        application/json:
+          schema:
+            type: object
+            properties:
+              queries:
+                type: array
+                items:
+                  type: object
+                  properties:
+                    select:
+                      type: object
+                    filter:
+                      type: object
+                    explode:
+                      type: object
+                    concat:
+                      type: object
+                    rename:
+                      type: object
+                    index:
+                      type: object
+              start_time:
+                type: string
+                format: date-time
+              end_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
+      404:
+        description: Not found
+      422:
+        description: Request body is invalid
+    """
+    payload = await request.json()
+    queries = payload.pop("queries", [])
+    project_name = (
+        request.query_params.get("project-name")
+        # read from headers/payload for backward-compatibility
+        or request.headers.get("project-name")
+        or payload.get("project_name")
+        or DEFAULT_PROJECT_NAME
+    )
+    end_time = payload.get("end_time") or payload.get("stop_time")
+    try:
+        span_queries = [SpanQuery.from_dict(query) for query in queries]
+    except Exception as e:
+        return Response(
+            status_code=HTTP_422_UNPROCESSABLE_ENTITY,
+            content=f"Invalid query: {e}",
+        )
+    async with request.app.state.db() as session:
+        results = []
+        for query in span_queries:
+            results.append(
+                await session.run_sync(
+                    query,
+                    project_name=project_name,
+                    start_time=normalize_datetime(
+                        from_iso_format(payload.get("start_time")),
+                        timezone.utc,
+                    ),
+                    end_time=normalize_datetime(
+                        from_iso_format(end_time),
+                        timezone.utc,
+                    ),
+                    limit=payload.get("limit", DEFAULT_SPAN_LIMIT),
+                    root_spans_only=payload.get("root_spans_only"),
+                )
+            )
+    if not results:
+        return Response(status_code=HTTP_404_NOT_FOUND)
+
+    async def content() -> AsyncIterator[bytes]:
+        for result in results:
+            yield df_to_bytes(result)
+
+    return StreamingResponse(
+        content=content(),
+        media_type="application/x-pandas-arrow",
+    )
+
+
+async def get_spans_handler(request: Request) -> Response:
+    return await query_spans_handler(request)