arize-phoenix 8.32.1__py3-none-any.whl → 9.0.0__py3-none-any.whl

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

@@ -3,9 +3,10 @@ from __future__ import annotations
 from enum import Enum
 from typing import Any, Literal, Mapping, Optional, Union
 
-from pydantic import BaseModel, ConfigDict, Field, RootModel, model_validator
+from pydantic import Field, RootModel, model_validator
 from typing_extensions import Annotated, Self, TypeAlias, TypeGuard, assert_never
 
+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.openai import OpenAIToolChoiceConversion
@@ -13,25 +14,6 @@ from phoenix.server.api.helpers.prompts.conversions.openai import OpenAIToolChoi
 JSONSerializable = Union[None, bool, int, float, str, dict[str, Any], list[Any]]
 
 
-class Undefined:
-    """
-    A singleton class that represents an unset or undefined value. Needed since Pydantic
-    can't natively distinguish between an undefined value and a value that is set to
-    None.
-    """
-
-    def __new__(cls) -> Any:
-        if not hasattr(cls, "_instance"):
-            cls._instance = super().__new__(cls)
-        return cls._instance
-
-    def __bool__(self) -> bool:
-        return False
-
-
-UNDEFINED: Any = Undefined()
-
-
 class PromptTemplateType(str, Enum):
     STRING = "STR"
     CHAT = "CHAT"
@@ -50,33 +32,18 @@ class PromptTemplateFormat(str, Enum):
     NONE = "NONE"
 
 
-class PromptModel(BaseModel):
-    model_config = ConfigDict(
-        extra="forbid",  # disallow extra attributes
-        use_enum_values=True,
-        validate_assignment=True,
-    )
-
-    def __init__(self, *args: Any, **kwargs: Any) -> None:
-        kwargs = {k: v for k, v in kwargs.items() if v is not UNDEFINED}
-        super().__init__(*args, **kwargs)
-
-    def model_dump(self, *args: Any, **kwargs: Any) -> dict[str, Any]:
-        return super().model_dump(*args, exclude_unset=True, by_alias=True, **kwargs)
-
-
-class TextContentPart(PromptModel):
+class TextContentPart(DBBaseModel):
     type: Literal["text"]
     text: str
 
 
-class ToolCallFunction(PromptModel):
+class ToolCallFunction(DBBaseModel):
     type: Literal["function"]
     name: str
     arguments: str
 
 
-class ToolCallContentPart(PromptModel):
+class ToolCallContentPart(DBBaseModel):
     type: Literal["tool_call"]
     tool_call_id: str
     tool_call: Annotated[
@@ -85,7 +52,7 @@ class ToolCallContentPart(PromptModel):
     ]
 
 
-class ToolResultContentPart(PromptModel):
+class ToolResultContentPart(DBBaseModel):
     type: Literal["tool_result"]
     tool_call_id: str
     tool_result: JSONSerializable
@@ -131,17 +98,17 @@ class RoleConversion:
         assert_never(role)
 
 
-class PromptMessage(PromptModel):
+class PromptMessage(DBBaseModel):
     role: Role
     content: Union[str, Annotated[list[ContentPart], Field(..., min_length=1)]]
 
 
-class PromptChatTemplate(PromptModel):
+class PromptChatTemplate(DBBaseModel):
     type: Literal["chat"]
     messages: list[PromptMessage]
 
 
-class PromptStringTemplate(PromptModel):
+class PromptStringTemplate(DBBaseModel):
     type: Literal["string"]
     template: str
 
@@ -159,12 +126,12 @@ class PromptTemplateRootModel(RootModel[PromptTemplate]):
     root: PromptTemplate
 
 
-class PromptToolFunction(PromptModel):
+class PromptToolFunction(DBBaseModel):
     type: Literal["function"]
     function: PromptToolFunctionDefinition
 
 
-class PromptToolFunctionDefinition(PromptModel):
+class PromptToolFunctionDefinition(DBBaseModel):
     name: str
     description: str = UNDEFINED
     parameters: dict[str, Any] = UNDEFINED
@@ -174,26 +141,26 @@ class PromptToolFunctionDefinition(PromptModel):
 PromptTool: TypeAlias = Annotated[Union[PromptToolFunction], Field(..., discriminator="type")]
 
 
-class PromptTools(PromptModel):
+class PromptTools(DBBaseModel):
     type: Literal["tools"]
     tools: Annotated[list[PromptTool], Field(..., min_length=1)]
     tool_choice: PromptToolChoice = UNDEFINED
     disable_parallel_tool_calls: bool = UNDEFINED
 
 
-class PromptToolChoiceNone(PromptModel):
+class PromptToolChoiceNone(DBBaseModel):
     type: Literal["none"]
 
 
-class PromptToolChoiceZeroOrMore(PromptModel):
+class PromptToolChoiceZeroOrMore(DBBaseModel):
     type: Literal["zero_or_more"]
 
 
-class PromptToolChoiceOneOrMore(PromptModel):
+class PromptToolChoiceOneOrMore(DBBaseModel):
     type: Literal["one_or_more"]
 
 
-class PromptToolChoiceSpecificFunctionTool(PromptModel):
+class PromptToolChoiceSpecificFunctionTool(DBBaseModel):
     type: Literal["specific_function"]
     function_name: str
 
@@ -209,7 +176,7 @@ PromptToolChoice: TypeAlias = Annotated[
 ]
 
 
-class PromptOpenAIJSONSchema(PromptModel):
+class PromptOpenAIJSONSchema(DBBaseModel):
     """
     Based on https://github.com/openai/openai-python/blob/d16e6edde5a155626910b5758a0b939bfedb9ced/src/openai/types/shared/response_format_json_schema.py#L13
     """
@@ -223,7 +190,7 @@ class PromptOpenAIJSONSchema(PromptModel):
     strict: Optional[bool] = UNDEFINED
 
 
-class PromptOpenAIResponseFormatJSONSchema(PromptModel):
+class PromptOpenAIResponseFormatJSONSchema(DBBaseModel):
     """
     Based on https://github.com/openai/openai-python/blob/d16e6edde5a155626910b5758a0b939bfedb9ced/src/openai/types/shared/response_format_json_schema.py#L40
     """
@@ -232,12 +199,12 @@ class PromptOpenAIResponseFormatJSONSchema(PromptModel):
     type: Literal["json_schema"]
 
 
-class PromptResponseFormatJSONSchema(PromptModel):
+class PromptResponseFormatJSONSchema(DBBaseModel):
     type: Literal["json_schema"]
     json_schema: PromptResponseFormatJSONSchemaDefinition
 
 
-class PromptResponseFormatJSONSchemaDefinition(PromptModel):
+class PromptResponseFormatJSONSchemaDefinition(DBBaseModel):
     name: str
     description: str = UNDEFINED
     schema_: dict[str, Any] = Field(UNDEFINED, alias="schema")
@@ -305,7 +272,7 @@ def denormalize_response_format(
 
 
 # OpenAI tool definitions
-class OpenAIFunctionDefinition(PromptModel):
+class OpenAIFunctionDefinition(DBBaseModel):
     """
     Based on https://github.com/openai/openai-python/blob/1e07c9d839e7e96f02d0a4b745f379a43086334c/src/openai/types/shared_params/function_definition.py#L13
     """
@@ -316,7 +283,7 @@ class OpenAIFunctionDefinition(PromptModel):
     strict: Optional[bool] = UNDEFINED
 
 
-class OpenAIToolDefinition(PromptModel):
+class OpenAIToolDefinition(DBBaseModel):
     """
     Based on https://github.com/openai/openai-python/blob/1e07c9d839e7e96f02d0a4b745f379a43086334c/src/openai/types/chat/chat_completion_tool_param.py#L12
     """
@@ -326,7 +293,7 @@ class OpenAIToolDefinition(PromptModel):
 
 
 # Anthropic tool definitions
-class AnthropicCacheControlParam(PromptModel):
+class AnthropicCacheControlParam(DBBaseModel):
     """
     Based on https://github.com/anthropics/anthropic-sdk-python/blob/93cbbbde964e244f02bf1bd2b579c5fabce4e267/src/anthropic/types/cache_control_ephemeral_param.py#L10
     """
@@ -334,7 +301,7 @@ class AnthropicCacheControlParam(PromptModel):
     type: Literal["ephemeral"]
 
 
-class AnthropicToolDefinition(PromptModel):
+class AnthropicToolDefinition(DBBaseModel):
     """
     Based on https://github.com/anthropics/anthropic-sdk-python/blob/93cbbbde964e244f02bf1bd2b579c5fabce4e267/src/anthropic/types/tool_param.py#L22
     """
@@ -345,7 +312,7 @@ class AnthropicToolDefinition(PromptModel):
     description: str = UNDEFINED
 
 
-class PromptOpenAIInvocationParametersContent(PromptModel):
+class PromptOpenAIInvocationParametersContent(DBBaseModel):
     temperature: float = UNDEFINED
     max_tokens: int = UNDEFINED
     max_completion_tokens: int = UNDEFINED
@@ -356,7 +323,7 @@ class PromptOpenAIInvocationParametersContent(PromptModel):
     reasoning_effort: Literal["low", "medium", "high"] = UNDEFINED
 
 
-class PromptOpenAIInvocationParameters(PromptModel):
+class PromptOpenAIInvocationParameters(DBBaseModel):
     type: Literal["openai"]
     openai: PromptOpenAIInvocationParametersContent
 
@@ -365,21 +332,21 @@ class PromptAzureOpenAIInvocationParametersContent(PromptOpenAIInvocationParamet
     pass
 
 
-class PromptAzureOpenAIInvocationParameters(PromptModel):
+class PromptAzureOpenAIInvocationParameters(DBBaseModel):
     type: Literal["azure_openai"]
     azure_openai: PromptAzureOpenAIInvocationParametersContent
 
 
-class PromptAnthropicThinkingConfigDisabled(PromptModel):
+class PromptAnthropicThinkingConfigDisabled(DBBaseModel):
     type: Literal["disabled"]
 
 
-class PromptAnthropicThinkingConfigEnabled(PromptModel):
+class PromptAnthropicThinkingConfigEnabled(DBBaseModel):
     type: Literal["enabled"]
     budget_tokens: int = Field(..., ge=1024)
 
 
-class PromptAnthropicInvocationParametersContent(PromptModel):
+class PromptAnthropicInvocationParametersContent(DBBaseModel):
     max_tokens: int
     temperature: float = UNDEFINED
     top_p: float = UNDEFINED
@@ -398,12 +365,12 @@ class PromptAnthropicInvocationParametersContent(PromptModel):
         return self
 
 
-class PromptAnthropicInvocationParameters(PromptModel):
+class PromptAnthropicInvocationParameters(DBBaseModel):
     type: Literal["anthropic"]
     anthropic: PromptAnthropicInvocationParametersContent
 
 
-class PromptGoogleInvocationParametersContent(PromptModel):
+class PromptGoogleInvocationParametersContent(DBBaseModel):
     temperature: float = UNDEFINED
     max_output_tokens: int = UNDEFINED
     stop_sequences: list[str] = UNDEFINED
@@ -413,7 +380,7 @@ class PromptGoogleInvocationParametersContent(PromptModel):
     top_k: int = UNDEFINED
 
 
-class PromptGoogleInvocationParameters(PromptModel):
+class PromptGoogleInvocationParameters(DBBaseModel):
     type: Literal["google"]
     google: PromptGoogleInvocationParametersContent
 
@@ -524,7 +491,7 @@ def denormalize_tools(
     tools: PromptTools, model_provider: ModelProvider
 ) -> tuple[list[dict[str, Any]], Optional[Any]]:
     assert tools.type == "tools"
-    denormalized_tools: list[PromptModel]
+    denormalized_tools: list[DBBaseModel]
     tool_choice: Optional[Any] = None
     if model_provider is ModelProvider.OPENAI or model_provider is ModelProvider.AZURE_OPENAI:
         denormalized_tools = [_prompt_to_openai_tool(tool) for tool in tools.tools]

phoenix/server/api/input_types/CreateSpanAnnotationInput.py

@@ -4,6 +4,7 @@ import strawberry
 from strawberry.relay import GlobalID
 from strawberry.scalars import JSON
 
+from phoenix.server.api.types.AnnotationSource import AnnotationSource
 from phoenix.server.api.types.AnnotatorKind import AnnotatorKind
 
 
@@ -16,3 +17,11 @@ class CreateSpanAnnotationInput:
     score: Optional[float] = None
     explanation: Optional[str] = None
     metadata: JSON = strawberry.field(default_factory=dict)
+    source: AnnotationSource
+    identifier: Optional[str] = strawberry.UNSET
+
+
+@strawberry.input
+class CreateSpanNoteInput:
+    span_id: GlobalID
+    note: str

phoenix/server/api/input_types/CreateTraceAnnotationInput.py

@@ -4,6 +4,7 @@ import strawberry
 from strawberry.relay import GlobalID
 from strawberry.scalars import JSON
 
+from phoenix.server.api.types.AnnotationSource import AnnotationSource
 from phoenix.server.api.types.AnnotatorKind import AnnotatorKind
 
 
@@ -16,3 +17,5 @@ class CreateTraceAnnotationInput:
     score: Optional[float] = None
     explanation: Optional[str] = None
     metadata: JSON = strawberry.field(default_factory=dict)
+    source: AnnotationSource
+    identifier: Optional[str] = strawberry.UNSET

phoenix/server/api/input_types/PatchAnnotationInput.py

@@ -5,6 +5,7 @@ from strawberry import UNSET
 from strawberry.relay import GlobalID
 from strawberry.scalars import JSON
 
+from phoenix.server.api.types.AnnotationSource import AnnotationSource
 from phoenix.server.api.types.AnnotatorKind import AnnotatorKind
 
 
@@ -17,3 +18,5 @@ class PatchAnnotationInput:
     score: Optional[float] = UNSET
     explanation: Optional[str] = UNSET
     metadata: Optional[JSON] = UNSET
+    identifier: Optional[str] = UNSET
+    source: Optional[AnnotationSource] = UNSET

phoenix/server/api/input_types/SpanAnnotationFilter.py

@@ -0,0 +1,67 @@
+from typing import Optional
+
+import strawberry
+from strawberry import UNSET
+from strawberry.relay import GlobalID
+
+from phoenix.db import models
+from phoenix.server.api.exceptions import BadRequest
+from phoenix.server.api.types.AnnotationSource import AnnotationSource
+from phoenix.server.api.types.node import from_global_id_with_expected_type
+
+
+@strawberry.input
+class SpanAnnotationFilterCondition:
+    names: Optional[list[str]] = UNSET
+    sources: Optional[list[AnnotationSource]] = UNSET
+    user_ids: Optional[list[Optional[GlobalID]]] = UNSET
+
+    def __post_init__(self) -> None:
+        if isinstance(self.names, list) and not self.names:
+            raise BadRequest("names must be a non-empty list")
+        if isinstance(self.sources, list) and not self.sources:
+            raise BadRequest("sources must be a non-empty list")
+        if isinstance(self.user_ids, list) and not self.user_ids:
+            raise BadRequest("user ids must be a non-empty list")
+
+
+@strawberry.input
+class SpanAnnotationFilter:
+    include: Optional[SpanAnnotationFilterCondition] = UNSET
+    exclude: Optional[SpanAnnotationFilterCondition] = UNSET
+
+    def __post_init__(self) -> None:
+        if self.include is UNSET and self.exclude is UNSET:
+            raise BadRequest("include and exclude cannot both be unset")
+
+
+def satisfies_filter(span_annotation: models.SpanAnnotation, filter: SpanAnnotationFilter) -> bool:
+    """
+    Returns true if the span annotation satisfies the filter and false otherwise.
+    """
+    span_annotation_source = AnnotationSource(span_annotation.source)
+    if include := filter.include:
+        if include.names and span_annotation.name not in include.names:
+            return False
+        if include.sources and span_annotation_source not in include.sources:
+            return False
+        if include.user_ids:
+            user_rowids = [
+                from_global_id_with_expected_type(user_id, "User") if user_id is not None else None
+                for user_id in include.user_ids
+            ]
+            if span_annotation.user_id not in user_rowids:
+                return False
+    if exclude := filter.exclude:
+        if exclude.names and span_annotation.name in exclude.names:
+            return False
+        if exclude.sources and span_annotation_source in exclude.sources:
+            return False
+        if exclude.user_ids:
+            user_rowids = [
+                from_global_id_with_expected_type(user_id, "User") if user_id is not None else None
+                for user_id in exclude.user_ids
+            ]
+            if span_annotation.user_id in user_rowids:
+                return False
+    return True

phoenix/server/api/mutations/__init__.py

@@ -1,5 +1,6 @@
 import strawberry
 
+from phoenix.server.api.mutations.annotation_config_mutations import AnnotationConfigMutationMixin
 from phoenix.server.api.mutations.api_key_mutations import ApiKeyMutationMixin
 from phoenix.server.api.mutations.chat_mutations import (
     ChatCompletionMutationMixin,
@@ -8,6 +9,9 @@ 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.project_trace_retention_policy_mutations import (
+    ProjectTraceRetentionPolicyMutationMixin,
+)
 from phoenix.server.api.mutations.prompt_label_mutations import PromptLabelMutationMixin
 from phoenix.server.api.mutations.prompt_mutations import PromptMutationMixin
 from phoenix.server.api.mutations.prompt_version_tag_mutations import PromptVersionTagMutationMixin
@@ -19,12 +23,14 @@ from phoenix.server.api.mutations.user_mutations import UserMutationMixin
 
 @strawberry.type
 class Mutation(
+    AnnotationConfigMutationMixin,
     ApiKeyMutationMixin,
     ChatCompletionMutationMixin,
     DatasetMutationMixin,
     ExperimentMutationMixin,
     ExportEventsMutationMixin,
     ProjectMutationMixin,
+    ProjectTraceRetentionPolicyMutationMixin,
     PromptMutationMixin,
     PromptVersionTagMutationMixin,
     PromptLabelMutationMixin,