arize-phoenix 3.16.0__py3-none-any.whl → 7.7.0__py3-none-any.whl

phoenix/server/api/input_types/SpanSort.py

@@ -1,30 +1,83 @@
-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
 
 
 @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.llm_token_count_total
+        elif self is SpanColumn.tokenCountPrompt:
+            expr = models.Span.llm_token_count_prompt
+        elif self is SpanColumn.tokenCountCompletion:
+            expr = models.Span.llm_token_count_completion
+        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 +85,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 +119,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 +137,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/input_types/TemplateOptions.py

@@ -0,0 +1,10 @@
+import strawberry
+from strawberry.scalars import JSON
+
+from phoenix.server.api.types.TemplateLanguage import TemplateLanguage
+
+
+@strawberry.input
+class TemplateOptions:
+    variables: JSON
+    language: TemplateLanguage

phoenix/server/api/input_types/TraceAnnotationSort.py

@@ -0,0 +1,17 @@
+from enum import Enum
+
+import strawberry
+
+from phoenix.server.api.types.SortDir import SortDir
+
+
+@strawberry.enum
+class TraceAnnotationColumn(Enum):
+    createdAt = "created_at"
+    name = "name"
+
+
+@strawberry.input(description="The sort key and direction for TraceAnnotation connections")
+class TraceAnnotationSort:
+    col: TraceAnnotationColumn
+    dir: SortDir

phoenix/server/api/input_types/UserRoleInput.py

@@ -0,0 +1,9 @@
+from enum import Enum
+
+import strawberry
+
+
+@strawberry.enum
+class UserRoleInput(Enum):
+    ADMIN = "ADMIN"
+    MEMBER = "MEMBER"

phoenix/server/api/mutations/__init__.py

@@ -0,0 +1,28 @@
+import strawberry
+
+from phoenix.server.api.mutations.api_key_mutations import ApiKeyMutationMixin
+from phoenix.server.api.mutations.chat_mutations import (
+    ChatCompletionMutationMixin,
+)
+from phoenix.server.api.mutations.dataset_mutations import DatasetMutationMixin
+from phoenix.server.api.mutations.experiment_mutations import ExperimentMutationMixin
+from phoenix.server.api.mutations.export_events_mutations import ExportEventsMutationMixin
+from phoenix.server.api.mutations.project_mutations import ProjectMutationMixin
+from phoenix.server.api.mutations.span_annotations_mutations import SpanAnnotationMutationMixin
+from phoenix.server.api.mutations.trace_annotations_mutations import TraceAnnotationMutationMixin
+from phoenix.server.api.mutations.user_mutations import UserMutationMixin
+
+
+@strawberry.type
+class Mutation(
+    ApiKeyMutationMixin,
+    DatasetMutationMixin,
+    ExperimentMutationMixin,
+    ExportEventsMutationMixin,
+    ProjectMutationMixin,
+    SpanAnnotationMutationMixin,
+    TraceAnnotationMutationMixin,
+    UserMutationMixin,
+    ChatCompletionMutationMixin,
+):
+    pass

phoenix/server/api/mutations/api_key_mutations.py

@@ -0,0 +1,167 @@
+from datetime import datetime, timezone
+from typing import Optional
+
+import strawberry
+from sqlalchemy import select
+from strawberry import UNSET
+from strawberry.relay import GlobalID
+from strawberry.types import Info
+
+from phoenix.db import enums, models
+from phoenix.server.api.auth import IsAdmin, IsLocked, IsNotReadOnly
+from phoenix.server.api.context import Context
+from phoenix.server.api.exceptions import Unauthorized
+from phoenix.server.api.queries import Query
+from phoenix.server.api.types.node import from_global_id_with_expected_type
+from phoenix.server.api.types.SystemApiKey import SystemApiKey
+from phoenix.server.api.types.UserApiKey import UserApiKey
+from phoenix.server.bearer_auth import PhoenixUser
+from phoenix.server.types import ApiKeyAttributes, ApiKeyClaims, ApiKeyId, UserId
+
+
+@strawberry.type
+class CreateSystemApiKeyMutationPayload:
+    jwt: str
+    api_key: SystemApiKey
+    query: Query
+
+
+@strawberry.input
+class CreateApiKeyInput:
+    name: str
+    description: Optional[str] = UNSET
+    expires_at: Optional[datetime] = UNSET
+
+
+@strawberry.type
+class CreateUserApiKeyMutationPayload:
+    jwt: str
+    api_key: UserApiKey
+    query: Query
+
+
+@strawberry.input
+class CreateUserApiKeyInput:
+    name: str
+    description: Optional[str] = UNSET
+    expires_at: Optional[datetime] = UNSET
+
+
+@strawberry.input
+class DeleteApiKeyInput:
+    id: GlobalID
+
+
+@strawberry.type
+class DeleteApiKeyMutationPayload:
+    apiKeyId: GlobalID
+    query: Query
+
+
+@strawberry.type
+class ApiKeyMutationMixin:
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsAdmin, IsLocked])  # type: ignore
+    async def create_system_api_key(
+        self, info: Info[Context, None], input: CreateApiKeyInput
+    ) -> CreateSystemApiKeyMutationPayload:
+        assert (token_store := info.context.token_store) is not None
+        user_role = enums.UserRole.SYSTEM
+        async with info.context.db() as session:
+            # Get the system user - note this could be pushed into a dataloader
+            system_user = await session.scalar(
+                select(models.User)
+                .join(models.UserRole)  # Join User with UserRole
+                .where(models.UserRole.name == user_role.value)  # Filter where role is SYSTEM
+                .order_by(models.User.id)
+                .limit(1)
+            )
+            if system_user is None:
+                raise ValueError("System user not found")
+        issued_at = datetime.now(timezone.utc)
+        claims = ApiKeyClaims(
+            subject=UserId(system_user.id),
+            issued_at=issued_at,
+            expiration_time=input.expires_at or None,
+            attributes=ApiKeyAttributes(
+                user_role=user_role,
+                name=input.name,
+                description=input.description,
+            ),
+        )
+        token, token_id = await token_store.create_api_key(claims)
+        return CreateSystemApiKeyMutationPayload(
+            jwt=token,
+            api_key=SystemApiKey(
+                id_attr=int(token_id),
+                name=input.name,
+                description=input.description or None,
+                created_at=issued_at,
+                expires_at=input.expires_at or None,
+            ),
+            query=Query(),
+        )
+
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsLocked])  # type: ignore
+    async def create_user_api_key(
+        self, info: Info[Context, None], input: CreateUserApiKeyInput
+    ) -> CreateUserApiKeyMutationPayload:
+        assert (token_store := info.context.token_store) is not None
+        try:
+            user = info.context.request.user  # type: ignore
+            assert isinstance(user, PhoenixUser)
+        except AttributeError:
+            raise ValueError("User not found")
+        issued_at = datetime.now(timezone.utc)
+        claims = ApiKeyClaims(
+            subject=user.identity,
+            issued_at=issued_at,
+            expiration_time=input.expires_at or None,
+            attributes=ApiKeyAttributes(
+                user_role=enums.UserRole.MEMBER,
+                name=input.name,
+                description=input.description,
+            ),
+        )
+        token, token_id = await token_store.create_api_key(claims)
+        return CreateUserApiKeyMutationPayload(
+            jwt=token,
+            api_key=UserApiKey(
+                id_attr=int(token_id),
+                name=input.name,
+                description=input.description or None,
+                created_at=issued_at,
+                expires_at=input.expires_at or None,
+                user_id=int(user.identity),
+            ),
+            query=Query(),
+        )
+
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsAdmin])  # type: ignore
+    async def delete_system_api_key(
+        self, info: Info[Context, None], input: DeleteApiKeyInput
+    ) -> DeleteApiKeyMutationPayload:
+        assert (token_store := info.context.token_store) is not None
+        api_key_id = from_global_id_with_expected_type(
+            input.id, expected_type_name=SystemApiKey.__name__
+        )
+        await token_store.revoke(ApiKeyId(api_key_id))
+        return DeleteApiKeyMutationPayload(apiKeyId=input.id, query=Query())
+
+    @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore
+    async def delete_user_api_key(
+        self, info: Info[Context, None], input: DeleteApiKeyInput
+    ) -> DeleteApiKeyMutationPayload:
+        assert (token_store := info.context.token_store) is not None
+        api_key_id = from_global_id_with_expected_type(
+            input.id, expected_type_name=UserApiKey.__name__
+        )
+        async with info.context.db() as session:
+            api_key = await session.scalar(
+                select(models.ApiKey).where(models.ApiKey.id == api_key_id)
+            )
+            if api_key is None:
+                raise ValueError(f"API key with id {input.id} not found")
+        if int((user := info.context.user).identity) != api_key.user_id and not user.is_admin:
+            raise Unauthorized("User not authorized to delete")
+        await token_store.revoke(ApiKeyId(api_key_id))
+        return DeleteApiKeyMutationPayload(apiKeyId=input.id, query=Query())