arize-phoenix 12.7.1__py3-none-any.whl → 12.9.0__py3-none-any.whl

phoenix/server/api/mutations/prompt_mutations.py

@@ -1,4 +1,4 @@
-from typing import Any, Optional, Union, cast
+from typing import Optional
 
 import strawberry
 from fastapi import Request
@@ -12,18 +12,11 @@ from strawberry.types import Info
 
 from phoenix.db import models
 from phoenix.db.types.identifier import Identifier as IdentifierModel
-from phoenix.db.types.model_provider import ModelProvider
 from phoenix.server.api.auth import IsLocked, IsNotReadOnly, IsNotViewer
 from phoenix.server.api.context import Context
 from phoenix.server.api.exceptions import BadRequest, Conflict, NotFound
-from phoenix.server.api.helpers.prompts.models import (
-    normalize_response_format,
-    normalize_tools,
-    validate_invocation_parameters,
-)
 from phoenix.server.api.input_types.PromptVersionInput import (
     ChatPromptVersionInput,
-    to_pydantic_prompt_chat_template_v1,
 )
 from phoenix.server.api.mutations.prompt_version_tag_mutations import (
     SetPromptVersionTagInput,
@@ -32,7 +25,7 @@ from phoenix.server.api.mutations.prompt_version_tag_mutations import (
 from phoenix.server.api.queries import Query
 from phoenix.server.api.types.Identifier import Identifier
 from phoenix.server.api.types.node import from_global_id_with_expected_type
-from phoenix.server.api.types.Prompt import Prompt, to_gql_prompt_from_orm
+from phoenix.server.api.types.Prompt import Prompt
 from phoenix.server.bearer_auth import PhoenixUser
 
 
@@ -84,63 +77,23 @@ class PromptMutationMixin:
         if "user" in request.scope:
             assert isinstance(user := request.user, PhoenixUser)
             user_id = int(user.identity)
-
-        input_prompt_version = input.prompt_version
-        tool_definitions = [tool.definition for tool in input_prompt_version.tools]
-        tool_choice = cast(
-            Optional[Union[str, dict[str, Any]]],
-            cast(dict[str, Any], input.prompt_version.invocation_parameters).pop(
-                "tool_choice", None
-            ),
-        )
-        model_provider = ModelProvider(input_prompt_version.model_provider)
         try:
-            tools = (
-                normalize_tools(tool_definitions, model_provider, tool_choice)
-                if tool_definitions
-                else None
-            )
-            template = to_pydantic_prompt_chat_template_v1(input_prompt_version.template)
-            response_format = (
-                normalize_response_format(
-                    input_prompt_version.response_format.definition,
-                    model_provider,
-                )
-                if input_prompt_version.response_format
-                else None
-            )
-            invocation_parameters = validate_invocation_parameters(
-                input_prompt_version.invocation_parameters,
-                model_provider,
-            )
+            prompt_version = input.prompt_version.to_orm_prompt_version(user_id)
         except ValidationError as error:
             raise BadRequest(str(error))
-
+        name = IdentifierModel.model_validate(str(input.name))
+        prompt = models.Prompt(
+            name=name,
+            description=input.description,
+            prompt_versions=[prompt_version],
+        )
         async with info.context.db() as session:
-            prompt_version = models.PromptVersion(
-                description=input_prompt_version.description,
-                user_id=user_id,
-                template_type="CHAT",
-                template_format=input_prompt_version.template_format,
-                template=template,
-                invocation_parameters=invocation_parameters,
-                tools=tools,
-                response_format=response_format,
-                model_provider=input_prompt_version.model_provider,
-                model_name=input_prompt_version.model_name,
-            )
-            name = IdentifierModel.model_validate(str(input.name))
-            prompt = models.Prompt(
-                name=name,
-                description=input.description,
-                prompt_versions=[prompt_version],
-            )
             session.add(prompt)
             try:
                 await session.commit()
             except (PostgreSQLIntegrityError, SQLiteIntegrityError):
                 raise Conflict(f"A prompt named '{input.name}' already exists")
-        return to_gql_prompt_from_orm(prompt)
+        return Prompt(id=prompt.id, db_record=prompt)
 
     @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer, IsLocked])  # type: ignore
     async def create_chat_prompt_version(
@@ -153,72 +106,26 @@ class PromptMutationMixin:
         if "user" in request.scope:
             assert isinstance(user := request.user, PhoenixUser)
             user_id = int(user.identity)
-
-        input_prompt_version = input.prompt_version
-        tool_definitions = [tool.definition for tool in input.prompt_version.tools]
-        tool_choice = cast(
-            Optional[Union[str, dict[str, Any]]],
-            cast(dict[str, Any], input.prompt_version.invocation_parameters).pop(
-                "tool_choice", None
-            ),
-        )
-        model_provider = ModelProvider(input_prompt_version.model_provider)
         try:
-            tools = (
-                normalize_tools(tool_definitions, model_provider, tool_choice)
-                if tool_definitions
-                else None
-            )
-            template = to_pydantic_prompt_chat_template_v1(input_prompt_version.template)
-            response_format = (
-                normalize_response_format(
-                    input_prompt_version.response_format.definition,
-                    model_provider,
-                )
-                if input_prompt_version.response_format
-                else None
-            )
-            invocation_parameters = validate_invocation_parameters(
-                input_prompt_version.invocation_parameters,
-                model_provider,
-            )
+            prompt_version = input.prompt_version.to_orm_prompt_version(user_id)
         except ValidationError as error:
             raise BadRequest(str(error))
-
         prompt_id = from_global_id_with_expected_type(
             global_id=input.prompt_id, expected_type_name=Prompt.__name__
         )
+        prompt_version.prompt_id = prompt_id
         async with info.context.db() as session:
-            prompt = await session.get(models.Prompt, prompt_id)
-            if not prompt:
-                raise NotFound(f"Prompt with ID '{input.prompt_id}' not found")
-
-            prompt_version = models.PromptVersion(
-                prompt_id=prompt_id,
-                description=input.prompt_version.description,
-                user_id=user_id,
-                template_type="CHAT",
-                template_format=input.prompt_version.template_format,
-                template=template,
-                invocation_parameters=invocation_parameters,
-                tools=tools,
-                response_format=response_format,
-                model_provider=input.prompt_version.model_provider,
-                model_name=input.prompt_version.model_name,
-            )
             session.add(prompt_version)
-
-        # ensure prompt_version is flushed to the database before creating tags against the
-        # prompt_version id
-        await session.flush()
-
-        if input.tags:
-            for tag in input.tags:
-                await upsert_prompt_version_tag(
-                    session, prompt_id, prompt_version.id, tag.name, tag.description
-                )
-
-        return to_gql_prompt_from_orm(prompt)
+            try:
+                await session.flush()
+            except (PostgreSQLIntegrityError, SQLiteIntegrityError):
+                raise NotFound(f"Prompt with ID '{input.prompt_id}' not found")
+            if input.tags:
+                for tag in input.tags:
+                    await upsert_prompt_version_tag(
+                        session, prompt_id, prompt_version.id, tag.name, tag.description
+                    )
+        return Prompt(id=prompt_id)
 
     @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer])  # type: ignore
     async def delete_prompt(
@@ -288,7 +195,7 @@ class PromptMutationMixin:
                 await session.commit()
             except (PostgreSQLIntegrityError, SQLiteIntegrityError):
                 raise Conflict(f"A prompt named '{input.name}' already exists")
-        return to_gql_prompt_from_orm(new_prompt)
+        return Prompt(id=new_prompt.id, db_record=new_prompt)
 
     @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer, IsLocked])  # type: ignore
     async def patch_prompt(self, info: Info[Context, None], input: PatchPromptInput) -> Prompt:
@@ -310,4 +217,4 @@ class PromptMutationMixin:
             if prompt is None:
                 raise NotFound(f"Prompt with ID '{input.prompt_id}' not found")
 
-        return to_gql_prompt_from_orm(prompt)
+        return Prompt(id=prompt.id, db_record=prompt)

phoenix/server/api/mutations/prompt_version_tag_mutations.py

@@ -16,9 +16,9 @@ from phoenix.server.api.exceptions import BadRequest, Conflict, NotFound
 from phoenix.server.api.queries import Query
 from phoenix.server.api.types.Identifier import Identifier
 from phoenix.server.api.types.node import from_global_id_with_expected_type
-from phoenix.server.api.types.Prompt import Prompt, to_gql_prompt_from_orm
+from phoenix.server.api.types.Prompt import Prompt
 from phoenix.server.api.types.PromptVersion import PromptVersion
-from phoenix.server.api.types.PromptVersionTag import PromptVersionTag, to_gql_prompt_version_tag
+from phoenix.server.api.types.PromptVersionTag import PromptVersionTag
 
 
 @strawberry.input
@@ -75,7 +75,9 @@ class PromptVersionTagMutationMixin:
             await session.delete(prompt_version_tag)
             await session.commit()
             return PromptVersionTagMutationPayload(
-                prompt_version_tag=None, query=Query(), prompt=to_gql_prompt_from_orm(prompt)
+                prompt_version_tag=None,
+                query=Query(),
+                prompt=Prompt(id=prompt.id, db_record=prompt),
             )
 
     @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer, IsLocked])  # type: ignore
@@ -111,9 +113,10 @@ class PromptVersionTagMutationMixin:
             except (PostgreSQLIntegrityError, SQLiteIntegrityError):
                 raise Conflict("Failed to update PromptVersionTag.")
 
-            version_tag = to_gql_prompt_version_tag(updated_tag)
             return PromptVersionTagMutationPayload(
-                prompt_version_tag=version_tag, prompt=to_gql_prompt_from_orm(prompt), query=Query()
+                prompt_version_tag=PromptVersionTag(id=updated_tag.id, db_record=updated_tag),
+                prompt=Prompt(id=prompt.id, db_record=prompt),
+                query=Query(),
             )
 
 

phoenix/server/api/mutations/span_annotations_mutations.py

@@ -21,7 +21,7 @@ from phoenix.server.api.queries import Query
 from phoenix.server.api.types.AnnotationSource import AnnotationSource
 from phoenix.server.api.types.AnnotatorKind import AnnotatorKind
 from phoenix.server.api.types.node import from_global_id_with_expected_type
-from phoenix.server.api.types.SpanAnnotation import SpanAnnotation, to_gql_span_annotation
+from phoenix.server.api.types.SpanAnnotation import SpanAnnotation
 from phoenix.server.bearer_auth import PhoenixUser
 from phoenix.server.dml_event import SpanAnnotationDeleteEvent, SpanAnnotationInsertEvent
 
@@ -138,7 +138,7 @@ class SpanAnnotationMutationMixin:
 
             # Convert the fully loaded annotations to GQL types
             returned_annotations = [
-                to_gql_span_annotation(anno) for anno in ordered_final_annotations
+                SpanAnnotation(id=anno.id, db_record=anno) for anno in ordered_final_annotations
             ]
 
             await session.commit()
@@ -184,7 +184,9 @@ class SpanAnnotationMutationMixin:
             processed_annotation = result.one()
 
             info.context.event_queue.put(SpanAnnotationInsertEvent((processed_annotation.id,)))
-            returned_annotation = to_gql_span_annotation(processed_annotation)
+            returned_annotation = SpanAnnotation(
+                id=processed_annotation.id, db_record=processed_annotation
+            )
             await session.commit()
         return SpanAnnotationMutationPayload(
             span_annotations=[returned_annotation],
@@ -256,7 +258,7 @@ class SpanAnnotationMutationMixin:
                 session.add(span_annotation)
 
             patched_annotations = [
-                to_gql_span_annotation(span_annotation)
+                SpanAnnotation(id=span_annotation.id, db_record=span_annotation)
                 for span_annotation in span_annotations_by_id.values()
             ]
 
@@ -320,7 +322,10 @@ class SpanAnnotationMutationMixin:
                 )
 
         deleted_annotations_gql = [
-            to_gql_span_annotation(deleted_annotations_by_id[id]) for id in span_annotation_ids
+            SpanAnnotation(
+                id=deleted_annotations_by_id[id].id, db_record=deleted_annotations_by_id[id]
+            )
+            for id in span_annotation_ids
         ]
         info.context.event_queue.put(
             SpanAnnotationDeleteEvent(tuple(deleted_annotations_by_id.keys()))

phoenix/server/api/mutations/trace_annotations_mutations.py

@@ -16,7 +16,7 @@ from phoenix.server.api.input_types.PatchAnnotationInput import PatchAnnotationI
 from phoenix.server.api.queries import Query
 from phoenix.server.api.types.AnnotationSource import AnnotationSource
 from phoenix.server.api.types.node import from_global_id_with_expected_type
-from phoenix.server.api.types.TraceAnnotation import TraceAnnotation, to_gql_trace_annotation
+from phoenix.server.api.types.TraceAnnotation import TraceAnnotation
 from phoenix.server.bearer_auth import PhoenixUser
 from phoenix.server.dml_event import TraceAnnotationDeleteEvent, TraceAnnotationInsertEvent
 
@@ -111,7 +111,9 @@ class TraceAnnotationMutationMixin:
             info.context.event_queue.put(TraceAnnotationInsertEvent(inserted_annotation_ids))
 
         returned_annotations = [
-            to_gql_trace_annotation(processed_annotations_map[i])
+            TraceAnnotation(
+                id=processed_annotations_map[i].id, db_record=processed_annotations_map[i]
+            )
             for i in sorted(processed_annotations_map.keys())
         ]
 
@@ -186,7 +188,7 @@ class TraceAnnotationMutationMixin:
             await session.commit()
 
         patched_annotations = [
-            to_gql_trace_annotation(trace_annotation)
+            TraceAnnotation(id=trace_annotation.id, db_record=trace_annotation)
             for trace_annotation in trace_annotations_by_id.values()
         ]
         info.context.event_queue.put(TraceAnnotationInsertEvent(tuple(patch_by_id.keys())))
@@ -245,7 +247,10 @@ class TraceAnnotationMutationMixin:
                 )
 
         deleted_gql_annotations = [
-            to_gql_trace_annotation(deleted_annotations_by_id[id]) for id in trace_annotation_ids
+            TraceAnnotation(
+                id=deleted_annotations_by_id[id].id, db_record=deleted_annotations_by_id[id]
+            )
+            for id in trace_annotation_ids
         ]
         info.context.event_queue.put(
             TraceAnnotationDeleteEvent(tuple(deleted_annotations_by_id.keys()))

phoenix/server/api/mutations/user_mutations.py

@@ -33,7 +33,7 @@ from phoenix.server.api.exceptions import BadRequest, Conflict, NotFound, Unauth
 from phoenix.server.api.input_types.UserRoleInput import UserRoleInput
 from phoenix.server.api.types.AuthMethod import AuthMethod
 from phoenix.server.api.types.node import from_global_id_with_expected_type
-from phoenix.server.api.types.User import User, to_gql_user
+from phoenix.server.api.types.User import User
 from phoenix.server.bearer_auth import PhoenixUser
 from phoenix.server.types import AccessTokenId, ApiKeyId, PasswordResetTokenId, RefreshTokenId
 
@@ -155,7 +155,7 @@ class UserMutationMixin:
             except Exception as error:
                 # Log the error but do not raise it
                 logger.error(f"Failed to send welcome email: {error}")
-        return UserMutationPayload(user=to_gql_user(user))
+        return UserMutationPayload(user=User(id=user.id, db_record=user))
 
     @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer, IsAdmin])  # type: ignore
     async def patch_user(
@@ -204,7 +204,7 @@ class UserMutationMixin:
         assert user
         if should_log_out:
             await info.context.log_out(user.id)
-        return UserMutationPayload(user=to_gql_user(user))
+        return UserMutationPayload(user=User(id=user.id, db_record=user))
 
     @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore
     async def patch_viewer(
@@ -246,7 +246,7 @@ class UserMutationMixin:
             response = info.context.get_response()
             response.delete_cookie(PHOENIX_REFRESH_TOKEN_COOKIE_NAME)
             response.delete_cookie(PHOENIX_ACCESS_TOKEN_COOKIE_NAME)
-        return UserMutationPayload(user=to_gql_user(user))
+        return UserMutationPayload(user=User(id=user.id, db_record=user))
 
     @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer, IsAdmin, IsLocked])  # type: ignore
     async def delete_users(