feat: Add tags with prompt version create

app/schema.graphql

@@ -276,6 +276,7 @@ input CreateChatPromptInput {
 input CreateChatPromptVersionInput {
   promptId: GlobalID!
   promptVersion: ChatPromptVersionInput!
+  tags: [SetPromptVersionTagInput!] = null
 }
 
 input CreateDatasetInput {

src/phoenix/server/api/mutations/prompt_mutations.py

@@ -21,6 +21,10 @@
     PromptVersion,
 )
 from phoenix.server.api.input_types.PromptVersionInput import ChatPromptVersionInput
+from phoenix.server.api.mutations.prompt_version_tag_mutations import (
+    SetPromptVersionTagInput,
+    upsert_prompt_version_tag,
+)
 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
@@ -39,6 +43,7 @@ class CreateChatPromptInput:
 class CreateChatPromptVersionInput:
     prompt_id: GlobalID
     prompt_version: ChatPromptVersionInput
+    tags: Optional[list[SetPromptVersionTagInput]] = None
 
 
 @strawberry.input
@@ -199,6 +204,16 @@ async def create_chat_prompt_version(
             )
             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)
 
     @strawberry.mutation

src/phoenix/server/api/mutations/prompt_version_tag_mutations.py

@@ -3,6 +3,7 @@
 import strawberry
 from sqlalchemy import select
 from sqlalchemy.exc import IntegrityError as PostgreSQLIntegrityError
+from sqlalchemy.ext.asyncio import AsyncSession
 from sqlean.dbapi2 import IntegrityError as SQLiteIntegrityError  # type: ignore[import-untyped]
 from strawberry.relay import GlobalID
 from strawberry.types import Info
@@ -96,28 +97,10 @@ async def set_prompt_version_tag(
             )
             if not prompt:
                 raise BadRequest("All prompt version tags must belong to a prompt")
-            name = IdentifierModel.model_validate(str(input.name))
-            existing_tag = await session.scalar(
-                select(models.PromptVersionTag).where(
-                    models.PromptVersionTag.prompt_id == prompt_id,
-                    models.PromptVersionTag.name == name,
-                )
-            )
 
-            if existing_tag:
-                existing_tag.prompt_version_id = prompt_version_id
-                if input.description is not None:
-                    existing_tag.description = input.description
-                updated_tag = existing_tag
-            else:
-                new_tag = models.PromptVersionTag(
-                    name=name,
-                    description=input.description,
-                    prompt_id=prompt_id,
-                    prompt_version_id=prompt_version_id,
-                )
-                session.add(new_tag)
-                updated_tag = new_tag
+            updated_tag = await upsert_prompt_version_tag(
+                session, prompt_id, prompt_version_id, input.name, input.description
+            )
 
             if not updated_tag:
                 raise BadRequest("Failed to create or update PromptVersionTag.")
@@ -131,3 +114,35 @@ async def set_prompt_version_tag(
             return PromptVersionTagMutationPayload(
                 prompt_version_tag=version_tag, prompt=to_gql_prompt_from_orm(prompt), query=Query()
             )
+
+
+async def upsert_prompt_version_tag(
+    session: AsyncSession,
+    prompt_id: int,
+    prompt_version_id: int,
+    name_str: str,
+    description: Optional[str] = None,
+) -> models.PromptVersionTag:
+    name = IdentifierModel.model_validate(name_str)
+
+    existing_tag = await session.scalar(
+        select(models.PromptVersionTag).where(
+            models.PromptVersionTag.prompt_id == prompt_id,
+            models.PromptVersionTag.name == name,
+        )
+    )
+
+    if existing_tag:
+        existing_tag.prompt_version_id = prompt_version_id
+        if description is not None:
+            existing_tag.description = description
+        return existing_tag
+    else:
+        new_tag = models.PromptVersionTag(
+            name=name,
+            description=description,
+            prompt_id=prompt_id,
+            prompt_version_id=prompt_version_id,
+        )
+        session.add(new_tag)
+        return new_tag