arize-phoenix 10.0.4__py3-none-any.whl → 12.28.1__py3-none-any.whl

phoenix/server/api/input_types/PromptVersionInput.py

@@ -1,10 +1,11 @@
 import json
-from typing import Optional, cast
+from typing import Any, Optional, Union, cast
 
 import strawberry
 from strawberry import UNSET
 from strawberry.scalars import JSON
 
+from phoenix.db import models
 from phoenix.db.types.model_provider import ModelProvider
 from phoenix.server.api.helpers.prompts.models import (
     ContentPart,
@@ -12,11 +13,15 @@ from phoenix.server.api.helpers.prompts.models import (
     PromptMessage,
     PromptMessageRole,
     PromptTemplateFormat,
+    PromptTemplateType,
     RoleConversion,
     TextContentPart,
     ToolCallContentPart,
     ToolCallFunction,
     ToolResultContentPart,
+    normalize_response_format,
+    normalize_tools,
+    validate_invocation_parameters,
 )
 
 
@@ -83,6 +88,52 @@ class ChatPromptVersionInput:
     model_provider: ModelProvider
     model_name: str
 
+    def __post_init__(self) -> None:
+        self.invocation_parameters = {
+            k: v for k, v in self.invocation_parameters.items() if v is not None
+        }
+
+    def to_orm_prompt_version(
+        self,
+        user_id: Optional[int],
+    ) -> models.PromptVersion:
+        tool_definitions = [tool.definition for tool in self.tools]
+        tool_choice = cast(
+            Optional[Union[str, dict[str, Any]]],
+            cast(dict[str, Any], self.invocation_parameters).pop("tool_choice", None),
+        )
+        model_provider = ModelProvider(self.model_provider)
+        tools = (
+            normalize_tools(tool_definitions, model_provider, tool_choice)
+            if tool_definitions
+            else None
+        )
+        template = to_pydantic_prompt_chat_template_v1(self.template)
+        response_format = (
+            normalize_response_format(
+                self.response_format.definition,
+                model_provider,
+            )
+            if self.response_format
+            else None
+        )
+        invocation_parameters = validate_invocation_parameters(
+            self.invocation_parameters,
+            model_provider,
+        )
+        return models.PromptVersion(
+            description=self.description,
+            user_id=user_id,
+            template_type=PromptTemplateType.CHAT,
+            template_format=self.template_format,
+            template=template,
+            invocation_parameters=invocation_parameters,
+            tools=tools,
+            response_format=response_format,
+            model_provider=ModelProvider(self.model_provider),
+            model_name=self.model_name,
+        )
+
 
 def to_pydantic_prompt_chat_template_v1(
     prompt_chat_template_input: PromptChatTemplateInput,

phoenix/server/api/input_types/SpanSort.py

@@ -3,7 +3,7 @@ from enum import Enum, auto
 from typing import Any, Optional, Protocol
 
 import strawberry
-from sqlalchemy import and_, desc, nulls_last
+from sqlalchemy import and_, desc, func, nulls_last, select
 from sqlalchemy.orm import InstrumentedAttribute
 from sqlalchemy.sql.expression import Select
 from strawberry import UNSET
@@ -11,6 +11,7 @@ from typing_extensions import assert_never
 
 import phoenix.trace.v1 as pb
 from phoenix.db import models
+from phoenix.db.helpers import truncate_name
 from phoenix.server.api.types.pagination import CursorSortColumnDataType
 from phoenix.server.api.types.SortDir import SortDir
 from phoenix.trace.schemas import SpanID
@@ -27,13 +28,14 @@ class SpanColumn(Enum):
     cumulativeTokenCountTotal = auto()
     cumulativeTokenCountPrompt = auto()
     cumulativeTokenCountCompletion = auto()
+    cumulativeTokenCostTotal = auto()
+    tokenCostTotal = auto()
 
     @property
     def column_name(self) -> str:
-        return f"{self.name}_span_sort_column"
+        return truncate_name(f"{self.name}_span_sort_column")
 
-    @property
-    def orm_expression(self) -> Any:
+    def as_orm_expression(self, joined_table: Optional[Any] = None) -> Any:
         expr: Any
         if self is SpanColumn.startTime:
             expr = models.Span.start_time
@@ -56,6 +58,11 @@ class SpanColumn(Enum):
             expr = models.Span.cumulative_llm_token_count_prompt
         elif self is SpanColumn.cumulativeTokenCountCompletion:
             expr = models.Span.cumulative_llm_token_count_completion
+        elif self is SpanColumn.tokenCostTotal:
+            expr = models.SpanCost.total_cost
+        elif self is SpanColumn.cumulativeTokenCostTotal:
+            assert joined_table is not None
+            expr = joined_table.c.cumulative_total_cost
         else:
             assert_never(self)
         return expr.label(self.column_name)
@@ -73,12 +80,41 @@ class SpanColumn(Enum):
             or self is SpanColumn.tokenCountTotal
             or self is SpanColumn.tokenCountPrompt
             or self is SpanColumn.tokenCountCompletion
+            or self is SpanColumn.tokenCostTotal
+            or self is SpanColumn.cumulativeTokenCostTotal
         ):
             return CursorSortColumnDataType.FLOAT
         if self is SpanColumn.startTime or self is SpanColumn.endTime:
             return CursorSortColumnDataType.DATETIME
         assert_never(self)
 
+    def join_tables(self, stmt: Select[Any]) -> tuple[Select[Any], Any]:
+        """
+        If needed, joins tables required for the sort column.
+        """
+        if self is SpanColumn.tokenCostTotal:
+            return stmt.join_from(
+                models.Span,
+                models.SpanCost,
+                onclause=models.SpanCost.span_rowid == models.Span.id,
+            ), models.SpanCost
+        if self is SpanColumn.cumulativeTokenCostTotal:
+            trace_costs = (
+                select(
+                    func.sum(models.SpanCost.total_cost).label("cumulative_total_cost"),
+                    models.SpanCost.trace_rowid,
+                )
+                .select_from(models.SpanCost)
+                .group_by(models.SpanCost.trace_rowid)
+                .subquery()
+            )
+            stmt = stmt.join(
+                trace_costs,
+                onclause=models.Span.trace_rowid == trace_costs.c.trace_rowid,
+            )
+            return stmt, trace_costs
+        return stmt, None
+
 
 @strawberry.enum
 class EvalAttr(Enum):
@@ -139,13 +175,14 @@ class SpanSort:
 
     def update_orm_expr(self, stmt: Select[Any]) -> SpanSortConfig:
         if (col := self.col) and not self.eval_result_key:
-            expr = col.orm_expression
+            stmt, joined_table = col.join_tables(stmt)
+            expr = col.as_orm_expression(joined_table)
             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,
+                orm_expression=col.as_orm_expression(joined_table),
                 dir=self.dir,
                 column_name=col.column_name,
                 column_data_type=col.data_type,
@@ -163,7 +200,7 @@ class SpanSort:
                     models.SpanAnnotation.span_rowid == models.Span.id,
                     models.SpanAnnotation.name == eval_name,
                 ),
-            ).order_by(expr)
+            ).order_by(nulls_last(expr))
             return SpanSortConfig(
                 stmt=stmt,
                 orm_expression=eval_result_key.attr.orm_expression,

phoenix/server/api/input_types/TimeBinConfig.py

@@ -0,0 +1,23 @@
+from enum import Enum
+
+import strawberry
+
+
+@strawberry.enum
+class TimeBinScale(Enum):
+    MINUTE = "minute"
+    HOUR = "hour"
+    DAY = "day"
+    WEEK = "week"
+    MONTH = "month"
+    YEAR = "year"
+
+
+@strawberry.input
+class TimeBinConfig:
+    scale: TimeBinScale = strawberry.field(
+        default=TimeBinScale.HOUR, description="The scale of time bins for aggregation."
+    )
+    utc_offset_minutes: int = strawberry.field(
+        default=0, description="Offset in minutes from UTC for local time binning."
+    )

phoenix/server/api/input_types/UpdateAnnotationInput.py

@@ -0,0 +1,34 @@
+from typing import Optional
+
+import strawberry
+from strawberry.relay import GlobalID
+from strawberry.scalars import JSON
+
+from phoenix.server.api.exceptions import BadRequest
+from phoenix.server.api.types.AnnotationSource import AnnotationSource
+from phoenix.server.api.types.AnnotatorKind import AnnotatorKind
+
+
+@strawberry.input
+class UpdateAnnotationInput:
+    id: GlobalID
+    name: str
+    annotator_kind: AnnotatorKind = AnnotatorKind.HUMAN
+    label: Optional[str] = None
+    score: Optional[float] = None
+    explanation: Optional[str] = None
+    metadata: JSON = strawberry.field(default_factory=dict)
+    source: AnnotationSource = AnnotationSource.APP
+
+    def __post_init__(self) -> None:
+        self.name = self.name.strip()
+        if isinstance(self.label, str):
+            self.label = self.label.strip()
+        if not self.label:
+            self.label = None
+        if isinstance(self.explanation, str):
+            self.explanation = self.explanation.strip()
+        if not self.explanation:
+            self.explanation = None
+        if self.score is None and not self.label and not self.explanation:
+            raise BadRequest("At least one of score, label, or explanation must be not null/empty.")

phoenix/server/api/input_types/UserRoleInput.py

@@ -7,3 +7,4 @@ import strawberry
 class UserRoleInput(Enum):
     ADMIN = "ADMIN"
     MEMBER = "MEMBER"
+    VIEWER = "VIEWER"

phoenix/server/api/mutations/__init__.py

@@ -5,10 +5,16 @@ 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_label_mutations import DatasetLabelMutationMixin
 from phoenix.server.api.mutations.dataset_mutations import DatasetMutationMixin
+from phoenix.server.api.mutations.dataset_split_mutations import DatasetSplitMutationMixin
 from phoenix.server.api.mutations.experiment_mutations import ExperimentMutationMixin
 from phoenix.server.api.mutations.export_events_mutations import ExportEventsMutationMixin
+from phoenix.server.api.mutations.model_mutations import ModelMutationMixin
 from phoenix.server.api.mutations.project_mutations import ProjectMutationMixin
+from phoenix.server.api.mutations.project_session_annotations_mutations import (
+    ProjectSessionAnnotationMutationMixin,
+)
 from phoenix.server.api.mutations.project_trace_retention_policy_mutations import (
     ProjectTraceRetentionPolicyMutationMixin,
 )
@@ -26,15 +32,19 @@ class Mutation(
     AnnotationConfigMutationMixin,
     ApiKeyMutationMixin,
     ChatCompletionMutationMixin,
+    DatasetLabelMutationMixin,
     DatasetMutationMixin,
+    DatasetSplitMutationMixin,
     ExperimentMutationMixin,
     ExportEventsMutationMixin,
+    ModelMutationMixin,
     ProjectMutationMixin,
     ProjectTraceRetentionPolicyMutationMixin,
     PromptMutationMixin,
     PromptVersionTagMutationMixin,
     PromptLabelMutationMixin,
     SpanAnnotationMutationMixin,
+    ProjectSessionAnnotationMutationMixin,
     TraceAnnotationMutationMixin,
     TraceMutationMixin,
     UserMutationMixin,

phoenix/server/api/mutations/annotation_config_mutations.py

@@ -23,7 +23,7 @@ from phoenix.db.types.annotation_configs import (
 from phoenix.db.types.annotation_configs import (
     FreeformAnnotationConfig as FreeformAnnotationConfigModel,
 )
-from phoenix.server.api.auth import IsNotReadOnly
+from phoenix.server.api.auth import IsNotReadOnly, IsNotViewer
 from phoenix.server.api.context import Context
 from phoenix.server.api.exceptions import BadRequest, Conflict, NotFound
 from phoenix.server.api.queries import Query
@@ -197,7 +197,7 @@ def _to_pydantic_freeform_annotation_config(
 
 @strawberry.type
 class AnnotationConfigMutationMixin:
-    @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore[misc]
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer])  # type: ignore[misc]
     async def create_annotation_config(
         self,
         info: Info[Context, None],
@@ -236,7 +236,7 @@ class AnnotationConfigMutationMixin:
             annotation_config=to_gql_annotation_config(annotation_config),
         )
 
-    @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore[misc]
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer])  # type: ignore[misc]
     async def update_annotation_config(
         self,
         info: Info[Context, None],
@@ -285,7 +285,7 @@ class AnnotationConfigMutationMixin:
             annotation_config=to_gql_annotation_config(annotation_config),
         )
 
-    @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore[misc]
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer])  # type: ignore[misc]
     async def delete_annotation_configs(
         self,
         info: Info[Context, None],
@@ -317,7 +317,7 @@ class AnnotationConfigMutationMixin:
             ],
         )
 
-    @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore[misc]
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer])  # type: ignore[misc]
     async def add_annotation_config_to_project(
         self,
         info: Info[Context, None],
@@ -374,10 +374,10 @@ class AnnotationConfigMutationMixin:
                 )
             return AddAnnotationConfigToProjectPayload(
                 query=Query(),
-                project=Project(project_rowid=project_id),
+                project=Project(id=project_id),
             )
 
-    @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore[misc]
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer])  # type: ignore[misc]
     async def remove_annotation_config_from_project(
         self,
         info: Info[Context, None],
@@ -409,5 +409,5 @@ class AnnotationConfigMutationMixin:
                 raise NotFound("Could not find one or more input project annotation configs")
         return RemoveAnnotationConfigFromProjectPayload(
             query=Query(),
-            project=Project(project_rowid=project_id),
+            project=Project(id=project_id),
         )

phoenix/server/api/mutations/api_key_mutations.py

@@ -1,5 +1,5 @@
 from datetime import datetime, timezone
-from typing import Optional
+from typing import Literal, Optional
 
 import strawberry
 from sqlalchemy import select
@@ -7,8 +7,9 @@ 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.db import models
+from phoenix.db.models import UserRoleName
+from phoenix.server.api.auth import IsAdmin, IsLocked, IsNotReadOnly, IsNotViewer
 from phoenix.server.api.context import Context
 from phoenix.server.api.exceptions import Unauthorized
 from phoenix.server.api.queries import Query
@@ -60,18 +61,18 @@ class DeleteApiKeyMutationPayload:
 
 @strawberry.type
 class ApiKeyMutationMixin:
-    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsAdmin, IsLocked])  # type: ignore
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer, 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
+        user_role: UserRoleName = "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
+                .where(models.UserRole.name == user_role)  # Filter where role is SYSTEM
                 .order_by(models.User.id)
                 .limit(1)
             )
@@ -91,13 +92,7 @@ class ApiKeyMutationMixin:
         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,
-            ),
+            api_key=SystemApiKey(id=int(token_id)),
             query=Query(),
         )
 
@@ -112,12 +107,20 @@ class ApiKeyMutationMixin:
         except AttributeError:
             raise ValueError("User not found")
         issued_at = datetime.now(timezone.utc)
+        # Determine user role for API key
+        user_role: Literal["ADMIN", "MEMBER", "VIEWER"]
+        if user.is_admin:
+            user_role = "ADMIN"
+        elif user.is_viewer:
+            user_role = "VIEWER"
+        else:
+            user_role = "MEMBER"
         claims = ApiKeyClaims(
             subject=user.identity,
             issued_at=issued_at,
             expiration_time=input.expires_at or None,
             attributes=ApiKeyAttributes(
-                user_role=enums.UserRole.MEMBER,
+                user_role=user_role,
                 name=input.name,
                 description=input.description,
             ),
@@ -125,18 +128,11 @@ class ApiKeyMutationMixin:
         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),
-            ),
+            api_key=UserApiKey(id=int(token_id)),
             query=Query(),
         )
 
-    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsAdmin])  # type: ignore
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer, IsAdmin])  # type: ignore
     async def delete_system_api_key(
         self, info: Info[Context, None], input: DeleteApiKeyInput
     ) -> DeleteApiKeyMutationPayload: