arize-phoenix 10.14.0__py3-none-any.whl → 11.0.0__py3-none-any.whl

phoenix/server/api/helpers/prompts/conversions/aws.py

@@ -0,0 +1,83 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional, Union
+
+from typing_extensions import assert_never
+
+if TYPE_CHECKING:
+    from anthropic.types import (
+        ToolChoiceAnyParam,
+        ToolChoiceAutoParam,
+        ToolChoiceParam,
+        ToolChoiceToolParam,
+    )
+
+    from phoenix.server.api.helpers.prompts.models import (
+        PromptToolChoiceNone,
+        PromptToolChoiceOneOrMore,
+        PromptToolChoiceSpecificFunctionTool,
+        PromptToolChoiceZeroOrMore,
+    )
+
+
+class AwsToolChoiceConversion:
+    @staticmethod
+    def to_aws(
+        obj: Union[
+            PromptToolChoiceNone,
+            PromptToolChoiceZeroOrMore,
+            PromptToolChoiceOneOrMore,
+            PromptToolChoiceSpecificFunctionTool,
+        ],
+        disable_parallel_tool_use: Optional[bool] = None,
+    ) -> ToolChoiceParam:
+        if obj.type == "zero_or_more":
+            choice_auto: ToolChoiceAutoParam = {"type": "auto"}
+            if disable_parallel_tool_use is not None:
+                choice_auto["disable_parallel_tool_use"] = disable_parallel_tool_use
+            return choice_auto
+        if obj.type == "one_or_more":
+            choice_any: ToolChoiceAnyParam = {"type": "any"}
+            if disable_parallel_tool_use is not None:
+                choice_any["disable_parallel_tool_use"] = disable_parallel_tool_use
+            return choice_any
+        if obj.type == "specific_function":
+            choice_tool: ToolChoiceToolParam = {"type": "tool", "name": obj.function_name}
+            if disable_parallel_tool_use is not None:
+                choice_tool["disable_parallel_tool_use"] = disable_parallel_tool_use
+            return choice_tool
+        if obj.type == "none":
+            return {"type": "none"}
+        assert_never(obj.type)
+
+    @staticmethod
+    def from_aws(
+        obj: ToolChoiceParam,
+    ) -> Union[
+        PromptToolChoiceNone,
+        PromptToolChoiceZeroOrMore,
+        PromptToolChoiceOneOrMore,
+        PromptToolChoiceSpecificFunctionTool,
+    ]:
+        from phoenix.server.api.helpers.prompts.models import (
+            PromptToolChoiceNone,
+            PromptToolChoiceOneOrMore,
+            PromptToolChoiceSpecificFunctionTool,
+            PromptToolChoiceZeroOrMore,
+        )
+
+        if obj["type"] == "auto":
+            choice_zero_or_more = PromptToolChoiceZeroOrMore(type="zero_or_more")
+            return choice_zero_or_more
+        if obj["type"] == "any":
+            choice_one_or_more = PromptToolChoiceOneOrMore(type="one_or_more")
+            return choice_one_or_more
+        if obj["type"] == "tool":
+            choice_function_tool = PromptToolChoiceSpecificFunctionTool(
+                type="specific_function",
+                function_name=obj["name"],
+            )
+            return choice_function_tool
+        if obj["type"] == "none":
+            return PromptToolChoiceNone(type="none")
+        assert_never(obj)

phoenix/server/api/helpers/prompts/models.py

@@ -9,6 +9,7 @@ from typing_extensions import Annotated, Self, TypeAlias, TypeGuard, assert_neve
 from phoenix.db.types.db_models import UNDEFINED, DBBaseModel
 from phoenix.db.types.model_provider import ModelProvider
 from phoenix.server.api.helpers.prompts.conversions.anthropic import AnthropicToolChoiceConversion
+from phoenix.server.api.helpers.prompts.conversions.aws import AwsToolChoiceConversion
 from phoenix.server.api.helpers.prompts.conversions.openai import OpenAIToolChoiceConversion
 
 JSONSerializable = Union[None, bool, int, float, str, dict[str, Any], list[Any]]
@@ -312,6 +313,14 @@ class AnthropicToolDefinition(DBBaseModel):
     description: str = UNDEFINED
 
 
+class BedrockToolDefinition(DBBaseModel):
+    """
+    Based on https://github.com/aws/amazon-bedrock-sdk-python/blob/main/src/bedrock/types/tool_param.py#L12
+    """
+
+    toolSpec: dict[str, Any]
+
+
 class PromptOpenAIInvocationParametersContent(DBBaseModel):
     temperature: float = UNDEFINED
     max_tokens: int = UNDEFINED
@@ -397,6 +406,17 @@ class PromptAnthropicInvocationParameters(DBBaseModel):
     anthropic: PromptAnthropicInvocationParametersContent
 
 
+class PromptAwsInvocationParametersContent(DBBaseModel):
+    max_tokens: int = UNDEFINED
+    temperature: float = UNDEFINED
+    top_p: float = UNDEFINED
+
+
+class PromptAwsInvocationParameters(DBBaseModel):
+    type: Literal["aws"]
+    aws: PromptAwsInvocationParametersContent
+
+
 class PromptGoogleInvocationParametersContent(DBBaseModel):
     temperature: float = UNDEFINED
     max_output_tokens: int = UNDEFINED
@@ -421,6 +441,7 @@ PromptInvocationParameters: TypeAlias = Annotated[
         PromptDeepSeekInvocationParameters,
         PromptXAIInvocationParameters,
         PromptOllamaInvocationParameters,
+        PromptAwsInvocationParameters,
     ],
     Field(..., discriminator="type"),
 ]
@@ -443,6 +464,8 @@ def get_raw_invocation_parameters(
         return invocation_parameters.xai.model_dump()
     if isinstance(invocation_parameters, PromptOllamaInvocationParameters):
         return invocation_parameters.ollama.model_dump()
+    if isinstance(invocation_parameters, PromptAwsInvocationParameters):
+        return invocation_parameters.aws.model_dump()
     assert_never(invocation_parameters)
 
 
@@ -459,6 +482,7 @@ def is_prompt_invocation_parameters(
             PromptDeepSeekInvocationParameters,
             PromptXAIInvocationParameters,
             PromptOllamaInvocationParameters,
+            PromptAwsInvocationParameters,
         ),
     )
 
@@ -512,6 +536,11 @@ def validate_invocation_parameters(
             type="ollama",
             ollama=PromptOllamaInvocationParametersContent.model_validate(invocation_parameters),
         )
+    elif model_provider is ModelProvider.AWS:
+        return PromptAwsInvocationParameters(
+            type="aws",
+            aws=PromptAwsInvocationParametersContent.model_validate(invocation_parameters),
+        )
     assert_never(model_provider)
 
 
@@ -530,12 +559,16 @@ def normalize_tools(
     ):
         openai_tools = [OpenAIToolDefinition.model_validate(schema) for schema in schemas]
         tools = [_openai_to_prompt_tool(openai_tool) for openai_tool in openai_tools]
+    elif model_provider is ModelProvider.AWS:
+        bedrock_tools = [BedrockToolDefinition.model_validate(schema) for schema in schemas]
+        tools = [_bedrock_to_prompt_tool(bedrock_tool) for bedrock_tool in bedrock_tools]
     elif model_provider is ModelProvider.ANTHROPIC:
         anthropic_tools = [AnthropicToolDefinition.model_validate(schema) for schema in schemas]
         tools = [_anthropic_to_prompt_tool(anthropic_tool) for anthropic_tool in anthropic_tools]
     else:
         raise ValueError(f"Unsupported model provider: {model_provider}")
     ans = PromptTools(type="tools", tools=tools)
+
     if tool_choice is not None:
         if (
             model_provider is ModelProvider.OPENAI
@@ -545,6 +578,8 @@ def normalize_tools(
             or model_provider is ModelProvider.OLLAMA
         ):
             ans.tool_choice = OpenAIToolChoiceConversion.from_openai(tool_choice)  # type: ignore[arg-type]
+        elif model_provider is ModelProvider.AWS:
+            ans.tool_choice = AwsToolChoiceConversion.from_aws(tool_choice)  # type: ignore[arg-type]
         elif model_provider is ModelProvider.ANTHROPIC:
             choice, disable_parallel_tool_calls = AnthropicToolChoiceConversion.from_anthropic(
                 tool_choice  # type: ignore[arg-type]
@@ -571,6 +606,10 @@ def denormalize_tools(
         denormalized_tools = [_prompt_to_openai_tool(tool) for tool in tools.tools]
         if tools.tool_choice:
             tool_choice = OpenAIToolChoiceConversion.to_openai(tools.tool_choice)
+    elif model_provider is ModelProvider.AWS:
+        denormalized_tools = [_prompt_to_bedrock_tool(tool) for tool in tools.tools]
+        if tools.tool_choice:
+            tool_choice = OpenAIToolChoiceConversion.to_openai(tools.tool_choice)
     elif model_provider is ModelProvider.ANTHROPIC:
         denormalized_tools = [_prompt_to_anthropic_tool(tool) for tool in tools.tools]
         if tools.tool_choice and tools.tool_choice.type != "none":
@@ -614,6 +653,19 @@ def _prompt_to_openai_tool(
     )
 
 
+def _bedrock_to_prompt_tool(
+    tool: BedrockToolDefinition,
+) -> PromptToolFunction:
+    return PromptToolFunction(
+        type="function",
+        function=PromptToolFunctionDefinition(
+            name=tool.toolSpec["name"],
+            description=tool.toolSpec["description"],
+            parameters=tool.toolSpec["inputSchema"]["json"],
+        ),
+    )
+
+
 def _anthropic_to_prompt_tool(
     tool: AnthropicToolDefinition,
 ) -> PromptToolFunction:
@@ -636,3 +688,18 @@ def _prompt_to_anthropic_tool(
         name=function.name,
         description=function.description,
     )
+
+
+def _prompt_to_bedrock_tool(
+    tool: PromptToolFunction,
+) -> BedrockToolDefinition:
+    function = tool.function
+    return BedrockToolDefinition(
+        toolSpec={
+            "name": function.name,
+            "description": function.description,
+            "inputSchema": {
+                "json": function.parameters,
+            },
+        }
+    )

phoenix/server/api/input_types/GenerativeModelInput.py

@@ -17,3 +17,5 @@ class GenerativeModelInput:
     """ The endpoint to use for the model. Only required for Azure OpenAI models. """
     api_version: Optional[str] = UNSET
     """ The API version to use for the model. """
+    region: Optional[str] = UNSET
+    """ The region to use for the model. """

phoenix/server/api/input_types/ProjectSessionSort.py

@@ -13,6 +13,7 @@ class ProjectSessionColumn(Enum):
     endTime = auto()
     tokenCountTotal = auto()
     numTraces = auto()
+    costTotal = auto()
 
     @property
     def data_type(self) -> CursorSortColumnDataType:
@@ -20,6 +21,8 @@ class ProjectSessionColumn(Enum):
             return CursorSortColumnDataType.INT
         if self is ProjectSessionColumn.startTime or self is ProjectSessionColumn.endTime:
             return CursorSortColumnDataType.DATETIME
+        if self is ProjectSessionColumn.costTotal:
+            return CursorSortColumnDataType.FLOAT
         assert_never(self)
 
 

phoenix/server/api/input_types/SpanSort.py

@@ -27,6 +27,7 @@ class SpanColumn(Enum):
     cumulativeTokenCountTotal = auto()
     cumulativeTokenCountPrompt = auto()
     cumulativeTokenCountCompletion = auto()
+    tokenCostTotal = auto()
 
     @property
     def column_name(self) -> str:
@@ -56,6 +57,8 @@ 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
         else:
             assert_never(self)
         return expr.label(self.column_name)
@@ -73,12 +76,25 @@ class SpanColumn(Enum):
             or self is SpanColumn.tokenCountTotal
             or self is SpanColumn.tokenCountPrompt
             or self is SpanColumn.tokenCountCompletion
+            or self is SpanColumn.tokenCostTotal
         ):
             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]) -> Select[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,
+            )
+        return stmt
+
 
 @strawberry.enum
 class EvalAttr(Enum):
@@ -140,6 +156,7 @@ 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 = col.join_tables(stmt)
             stmt = stmt.add_columns(expr)
             if self.dir == SortDir.desc:
                 expr = desc(expr)

phoenix/server/api/mutations/__init__.py

@@ -8,6 +8,7 @@ from phoenix.server.api.mutations.chat_mutations import (
 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.model_mutations import ModelMutationMixin
 from phoenix.server.api.mutations.project_mutations import ProjectMutationMixin
 from phoenix.server.api.mutations.project_trace_retention_policy_mutations import (
     ProjectTraceRetentionPolicyMutationMixin,
@@ -29,6 +30,7 @@ class Mutation(
     DatasetMutationMixin,
     ExperimentMutationMixin,
     ExportEventsMutationMixin,
+    ModelMutationMixin,
     ProjectMutationMixin,
     ProjectTraceRetentionPolicyMutationMixin,
     PromptMutationMixin,

phoenix/server/api/mutations/chat_mutations.py

@@ -1,4 +1,5 @@
 import asyncio
+import logging
 from dataclasses import asdict, field
 from datetime import datetime, timezone
 from itertools import chain, islice
@@ -73,6 +74,8 @@ from phoenix.utilities.template_formatters import (
     TemplateFormatter,
 )
 
+logger = logging.getLogger(__name__)
+
 initialize_playground_clients()
 
 ChatCompletionMessage = tuple[ChatCompletionMessageRole, str, Optional[str], Optional[List[Any]]]
@@ -450,6 +453,19 @@ class ChatCompletionMutationMixin:
             session.add(trace)
             session.add(span)
             await session.flush()
+            try:
+                span_cost = info.context.span_cost_calculator.calculate_cost(
+                    start_time=span.start_time,
+                    attributes=span.attributes,
+                )
+            except Exception as e:
+                logger.exception(f"Failed to calculate cost for span {span.id}: {e}")
+                span_cost = None
+            if span_cost:
+                span_cost.span_rowid = span.id
+                span_cost.trace_rowid = trace.id
+                session.add(span_cost)
+                await session.flush()
 
         gql_span = Span(span_rowid=span.id, db_span=span)
 
@@ -605,5 +621,6 @@ TOOL_CALL_FUNCTION_ARGUMENTS_JSON = ToolCallAttributes.TOOL_CALL_FUNCTION_ARGUME
 TOOL_JSON_SCHEMA = ToolAttributes.TOOL_JSON_SCHEMA
 PROMPT_TEMPLATE_VARIABLES = SpanAttributes.LLM_PROMPT_TEMPLATE_VARIABLES
 
+LLM_PROVIDER = SpanAttributes.LLM_PROVIDER
 
 PLAYGROUND_PROJECT_NAME = "playground"

phoenix/server/api/mutations/model_mutations.py

@@ -0,0 +1,208 @@
+import re
+from datetime import datetime, timezone
+from typing import Optional
+
+import sqlalchemy as sa
+import strawberry
+from sqlalchemy import delete
+from sqlalchemy.exc import IntegrityError as PostgreSQLIntegrityError
+from sqlalchemy.orm import joinedload
+from sqlean.dbapi2 import IntegrityError as SQLiteIntegrityError  # type: ignore[import-untyped]
+from strawberry.relay import GlobalID
+from strawberry.types import Info
+
+from phoenix.db import models
+from phoenix.server.api.auth import IsNotReadOnly
+from phoenix.server.api.context import Context
+from phoenix.server.api.exceptions import BadRequest, Conflict, NotFound
+from phoenix.server.api.queries import Query
+from phoenix.server.api.types.GenerativeModel import GenerativeModel, to_gql_generative_model
+from phoenix.server.api.types.node import from_global_id_with_expected_type
+from phoenix.server.api.types.TokenPrice import TokenKind
+
+
+@strawberry.input
+class TokenPriceInput:
+    token_type: str
+    cost_per_million_tokens: float
+    kind: TokenKind
+
+    @property
+    def token_prices(self) -> models.TokenPrice:
+        """Generate TokenPrice instances based on the input."""
+        return models.TokenPrice(
+            token_type=self.token_type,
+            is_prompt=self.kind == TokenKind.PROMPT,
+            base_rate=self.cost_per_million_tokens / 1_000_000,
+        )
+
+
+@strawberry.input
+class CreateModelMutationInput:
+    name: str
+    provider: Optional[str] = None
+    name_pattern: str
+    costs: list[TokenPriceInput]
+    start_time: Optional[datetime] = None
+
+
+@strawberry.type
+class CreateModelMutationPayload:
+    model: GenerativeModel
+    query: Query
+
+
+@strawberry.input
+class UpdateModelMutationInput:
+    id: GlobalID
+    name: str
+    provider: Optional[str]
+    name_pattern: str
+    costs: list[TokenPriceInput]
+    start_time: Optional[datetime] = None
+
+
+@strawberry.type
+class UpdateModelMutationPayload:
+    model: GenerativeModel
+    query: Query
+
+
+@strawberry.input
+class DeleteModelMutationInput:
+    id: GlobalID
+
+
+@strawberry.type
+class DeleteModelMutationPayload:
+    model: GenerativeModel
+    query: Query
+
+
+@strawberry.type
+class ModelMutationMixin:
+    @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore
+    async def create_model(
+        self,
+        info: Info[Context, None],
+        input: CreateModelMutationInput,
+    ) -> CreateModelMutationPayload:
+        cost_types = set(cost.token_type for cost in input.costs)
+        if "input" not in cost_types:
+            raise BadRequest("input cost is required")
+        if "output" not in cost_types:
+            raise BadRequest("output cost is required")
+        name_pattern = _compile_regular_expression(input.name_pattern)
+        token_prices = [cost.token_prices for cost in input.costs]
+        model = models.GenerativeModel(
+            name=input.name,
+            provider=input.provider,
+            name_pattern=name_pattern,
+            is_built_in=False,
+            token_prices=token_prices,
+            start_time=input.start_time,
+        )
+        async with info.context.db() as session:
+            session.add(model)
+            try:
+                await session.flush()
+            except (PostgreSQLIntegrityError, SQLiteIntegrityError):
+                raise Conflict(f"Model with name '{input.name}' already exists")
+
+        return CreateModelMutationPayload(
+            model=to_gql_generative_model(model),
+            query=Query(),
+        )
+
+    @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore
+    async def update_model(
+        self,
+        info: Info[Context, None],
+        input: UpdateModelMutationInput,
+    ) -> UpdateModelMutationPayload:
+        try:
+            model_id = from_global_id_with_expected_type(input.id, GenerativeModel.__name__)
+        except ValueError:
+            raise BadRequest(f'Invalid model id: "{input.id}"')
+
+        cost_types = set(cost.token_type for cost in input.costs)
+        if "input" not in cost_types:
+            raise BadRequest("input cost is required")
+        if "output" not in cost_types:
+            raise BadRequest("output cost is required")
+        name_pattern = _compile_regular_expression(input.name_pattern)
+        token_prices = [cost.token_prices for cost in input.costs]
+        async with info.context.db() as session:
+            model = await session.scalar(
+                sa.select(models.GenerativeModel)
+                .where(models.GenerativeModel.deleted_at.is_(None))
+                .where(models.GenerativeModel.id == model_id)
+                .options(joinedload(models.GenerativeModel.token_prices))
+            )
+            if model is None:
+                raise NotFound(f'Model "{input.id}" not found')
+            if model.is_built_in:
+                raise BadRequest("Cannot update built-in model")
+
+            await session.execute(
+                delete(models.TokenPrice).where(models.TokenPrice.model_id == model.id)
+            )
+
+            await session.refresh(model)
+
+            model.name = input.name
+            model.provider = input.provider or ""
+            model.name_pattern = name_pattern
+            model.token_prices = token_prices
+            model.start_time = input.start_time
+            session.add(model)
+            try:
+                await session.flush()
+            except (PostgreSQLIntegrityError, SQLiteIntegrityError):
+                raise Conflict(f"Model with name '{input.name}' already exists")
+            await session.refresh(model)
+
+        return UpdateModelMutationPayload(
+            model=to_gql_generative_model(model),
+            query=Query(),
+        )
+
+    @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore
+    async def delete_model(
+        self,
+        info: Info[Context, None],
+        input: DeleteModelMutationInput,
+    ) -> DeleteModelMutationPayload:
+        try:
+            model_id = from_global_id_with_expected_type(input.id, GenerativeModel.__name__)
+        except ValueError:
+            raise BadRequest(f'Invalid model id: "{input.id}"')
+
+        async with info.context.db() as session:
+            model = await session.scalar(
+                sa.update(models.GenerativeModel)
+                .values(deleted_at=datetime.now(timezone.utc))
+                .where(models.GenerativeModel.deleted_at.is_(None))
+                .where(models.GenerativeModel.id == model_id)
+                .returning(models.GenerativeModel)
+            )
+            if model is None:
+                raise NotFound(f'Model "{input.id}" not found')
+            if model.is_built_in:
+                await session.rollback()
+                raise BadRequest("Cannot delete built-in model")
+        return DeleteModelMutationPayload(
+            model=to_gql_generative_model(model),
+            query=Query(),
+        )
+
+
+def _compile_regular_expression(maybe_regex: str) -> re.Pattern[str]:
+    """
+    Compile the given string as a regular expression.
+    Raises a BadRequest error if the given string is not a valid regex.
+    """
+    try:
+        return re.compile(maybe_regex)
+    except re.error as error:
+        raise BadRequest(f"Invalid regex: {str(error)}")