arize-phoenix 7.12.3__py3-none-any.whl → 8.0.1__py3-none-any.whl

phoenix/db/models.py

@@ -13,6 +13,7 @@ from sqlalchemy import (
     ForeignKey,
     Index,
     MetaData,
+    Null,
     String,
     TypeDecorator,
     UniqueConstraint,
@@ -38,6 +39,21 @@ from sqlalchemy.sql import expression
 
 from phoenix.config import get_env_database_schema
 from phoenix.datetime_utils import normalize_datetime
+from phoenix.db.types.identifier import Identifier
+from phoenix.db.types.model_provider import ModelProvider
+from phoenix.server.api.helpers.prompts.models import (
+    PromptInvocationParameters,
+    PromptInvocationParametersRootModel,
+    PromptResponseFormat,
+    PromptResponseFormatRootModel,
+    PromptTemplate,
+    PromptTemplateFormat,
+    PromptTemplateRootModel,
+    PromptTemplateType,
+    PromptTools,
+    is_prompt_invocation_parameters,
+    is_prompt_template,
+)
 
 
 class AuthMethod(Enum):
@@ -99,6 +115,139 @@ class UtcTimeStamp(TypeDecorator[datetime]):
         return normalize_datetime(value, timezone.utc)
 
 
+class _Identifier(TypeDecorator[Identifier]):
+    # See # See https://docs.sqlalchemy.org/en/20/core/custom_types.html
+    cache_ok = True
+    impl = String
+
+    def process_bind_param(self, value: Optional[Identifier], _: Dialect) -> Optional[str]:
+        assert isinstance(value, Identifier) or value is None
+        return None if value is None else value.root
+
+    def process_result_value(self, value: Optional[str], _: Dialect) -> Optional[Identifier]:
+        return None if value is None else Identifier.model_validate(value)
+
+
+class _ModelProvider(TypeDecorator[ModelProvider]):
+    # See # See https://docs.sqlalchemy.org/en/20/core/custom_types.html
+    cache_ok = True
+    impl = String
+
+    def process_bind_param(self, value: Optional[ModelProvider], _: Dialect) -> Optional[str]:
+        if isinstance(value, str):
+            return ModelProvider(value).value
+        return None if value is None else value.value
+
+    def process_result_value(self, value: Optional[str], _: Dialect) -> Optional[ModelProvider]:
+        return None if value is None else ModelProvider(value)
+
+
+class _InvocationParameters(TypeDecorator[PromptInvocationParameters]):
+    # See # See https://docs.sqlalchemy.org/en/20/core/custom_types.html
+    cache_ok = True
+    impl = JSON_
+
+    def process_bind_param(
+        self, value: Optional[PromptInvocationParameters], _: Dialect
+    ) -> Optional[dict[str, Any]]:
+        assert is_prompt_invocation_parameters(value)
+        invocation_parameters = value.model_dump()
+        assert isinstance(invocation_parameters, dict)
+        return invocation_parameters
+
+    def process_result_value(
+        self, value: Optional[dict[str, Any]], _: Dialect
+    ) -> Optional[PromptInvocationParameters]:
+        assert isinstance(value, dict)
+        return PromptInvocationParametersRootModel.model_validate(value).root
+
+
+class _PromptTemplate(TypeDecorator[PromptTemplate]):
+    # See # See https://docs.sqlalchemy.org/en/20/core/custom_types.html
+    cache_ok = True
+    impl = JSON_
+
+    def process_bind_param(
+        self, value: Optional[PromptTemplate], _: Dialect
+    ) -> Optional[dict[str, Any]]:
+        assert is_prompt_template(value)
+        return value.model_dump() if value is not None else None
+
+    def process_result_value(
+        self, value: Optional[dict[str, Any]], _: Dialect
+    ) -> Optional[PromptTemplate]:
+        assert isinstance(value, dict)
+        return PromptTemplateRootModel.model_validate(value).root
+
+
+class _Tools(TypeDecorator[PromptTools]):
+    # See # See https://docs.sqlalchemy.org/en/20/core/custom_types.html
+    cache_ok = True
+    impl = JSON_
+
+    def process_bind_param(
+        self, value: Optional[PromptTools], _: Dialect
+    ) -> Optional[dict[str, Any]]:
+        return value.model_dump() if value is not None else None
+
+    def process_result_value(
+        self, value: Optional[dict[str, Any]], _: Dialect
+    ) -> Optional[PromptTools]:
+        return PromptTools.model_validate(value) if value is not None else None
+
+
+class _PromptResponseFormat(TypeDecorator[PromptResponseFormat]):
+    # See https://docs.sqlalchemy.org/en/20/core/custom_types.html
+    cache_ok = True
+    impl = JSON_
+
+    def process_bind_param(
+        self, value: Optional[PromptResponseFormat], _: Dialect
+    ) -> Optional[dict[str, Any]]:
+        return value.model_dump() if value is not None else None
+
+    def process_result_value(
+        self, value: Optional[dict[str, Any]], _: Dialect
+    ) -> Optional[PromptResponseFormat]:
+        return (
+            PromptResponseFormatRootModel.model_validate(value).root if value is not None else None
+        )
+
+
+class _PromptTemplateType(TypeDecorator[PromptTemplateType]):
+    # See # See https://docs.sqlalchemy.org/en/20/core/custom_types.html
+    cache_ok = True
+    impl = String
+
+    def process_bind_param(self, value: Optional[PromptTemplateType], _: Dialect) -> Optional[str]:
+        if isinstance(value, str):
+            return PromptTemplateType(value).value
+        return None if value is None else value.value
+
+    def process_result_value(
+        self, value: Optional[str], _: Dialect
+    ) -> Optional[PromptTemplateType]:
+        return None if value is None else PromptTemplateType(value)
+
+
+class _TemplateFormat(TypeDecorator[PromptTemplateFormat]):
+    # See # See https://docs.sqlalchemy.org/en/20/core/custom_types.html
+    cache_ok = True
+    impl = String
+
+    def process_bind_param(
+        self, value: Optional[PromptTemplateFormat], _: Dialect
+    ) -> Optional[str]:
+        if isinstance(value, str):
+            return PromptTemplateFormat(value).value
+        return None if value is None else value.value
+
+    def process_result_value(
+        self, value: Optional[str], _: Dialect
+    ) -> Optional[PromptTemplateFormat]:
+        return None if value is None else PromptTemplateFormat(value)
+
+
 class ExperimentRunOutput(TypedDict, total=False):
     task_output: Any
 
@@ -805,3 +954,161 @@ class ApiKey(Base):
     created_at: Mapped[datetime] = mapped_column(UtcTimeStamp, server_default=func.now())
     expires_at: Mapped[Optional[datetime]] = mapped_column(UtcTimeStamp, nullable=True, index=True)
     __table_args__ = (dict(sqlite_autoincrement=True),)
+
+
+class PromptLabel(Base):
+    __tablename__ = "prompt_labels"
+
+    id: Mapped[int] = mapped_column(primary_key=True)
+    name: Mapped[str] = mapped_column(String, unique=True, index=True, nullable=False)
+    description: Mapped[Optional[str]]
+    color: Mapped[str] = mapped_column(String, nullable=True)
+
+    prompts_prompt_labels: Mapped[list["PromptPromptLabel"]] = relationship(
+        "PromptPromptLabel",
+        back_populates="prompt_label",
+        cascade="all, delete-orphan",
+        uselist=True,
+    )
+
+
+class Prompt(Base):
+    __tablename__ = "prompts"
+
+    id: Mapped[int] = mapped_column(primary_key=True)
+    source_prompt_id: Mapped[Optional[int]] = mapped_column(
+        ForeignKey("prompts.id", ondelete="SET NULL"),
+        index=True,
+        nullable=True,
+    )
+    name: Mapped[Identifier] = mapped_column(_Identifier, unique=True, index=True, nullable=False)
+    description: Mapped[Optional[str]]
+    metadata_: Mapped[dict[str, Any]] = mapped_column("metadata")
+    created_at: Mapped[datetime] = mapped_column(UtcTimeStamp, server_default=func.now())
+    updated_at: Mapped[datetime] = mapped_column(
+        UtcTimeStamp, server_default=func.now(), onupdate=func.now()
+    )
+
+    prompts_prompt_labels: Mapped[list["PromptPromptLabel"]] = relationship(
+        "PromptPromptLabel",
+        back_populates="prompt",
+        cascade="all, delete-orphan",
+        uselist=True,
+    )
+
+    prompt_versions: Mapped[list["PromptVersion"]] = relationship(
+        "PromptVersion",
+        back_populates="prompt",
+        cascade="all, delete-orphan",
+        uselist=True,
+    )
+
+    prompt_version_tags: Mapped[list["PromptVersionTag"]] = relationship(
+        "PromptVersionTag",
+        back_populates="prompt",
+        cascade="all, delete-orphan",
+        uselist=True,
+    )
+
+
+class PromptPromptLabel(Base):
+    __tablename__ = "prompts_prompt_labels"
+
+    id: Mapped[int] = mapped_column(primary_key=True)
+    prompt_label_id: Mapped[int] = mapped_column(
+        ForeignKey("prompt_labels.id", ondelete="CASCADE"),
+        index=True,
+        nullable=False,
+    )
+    prompt_id: Mapped[int] = mapped_column(
+        ForeignKey("prompts.id", ondelete="CASCADE"),
+        index=True,
+        nullable=False,
+    )
+
+    prompt_label: Mapped["PromptLabel"] = relationship(
+        "PromptLabel", back_populates="prompts_prompt_labels"
+    )
+    prompt: Mapped["Prompt"] = relationship("Prompt", back_populates="prompts_prompt_labels")
+
+    __table_args__ = (UniqueConstraint("prompt_label_id", "prompt_id"),)
+
+
+class PromptVersion(Base):
+    __tablename__ = "prompt_versions"
+
+    id: Mapped[int] = mapped_column(primary_key=True)
+    prompt_id: Mapped[int] = mapped_column(
+        ForeignKey("prompts.id", ondelete="CASCADE"),
+        index=True,
+        nullable=False,
+    )
+    description: Mapped[Optional[str]] = mapped_column(String, nullable=True)
+    user_id: Mapped[Optional[int]] = mapped_column(
+        ForeignKey("users.id", ondelete="SET NULL"),
+        index=True,
+        nullable=True,
+    )
+    template_type: Mapped[PromptTemplateType] = mapped_column(
+        _PromptTemplateType,
+        CheckConstraint("template_type IN ('CHAT', 'STR')", name="template_type"),
+        nullable=False,
+    )
+    template_format: Mapped[PromptTemplateFormat] = mapped_column(
+        _TemplateFormat,
+        CheckConstraint(
+            "template_format IN ('F_STRING', 'MUSTACHE', 'NONE')", name="template_format"
+        ),
+        nullable=False,
+    )
+    template: Mapped[PromptTemplate] = mapped_column(_PromptTemplate, nullable=False)
+    invocation_parameters: Mapped[PromptInvocationParameters] = mapped_column(
+        _InvocationParameters, nullable=False
+    )
+    tools: Mapped[Optional[PromptTools]] = mapped_column(_Tools, default=Null(), nullable=True)
+    response_format: Mapped[Optional[PromptResponseFormat]] = mapped_column(
+        _PromptResponseFormat, default=Null(), nullable=True
+    )
+    model_provider: Mapped[ModelProvider] = mapped_column(_ModelProvider)
+    model_name: Mapped[str]
+    metadata_: Mapped[dict[str, Any]] = mapped_column("metadata")
+    created_at: Mapped[datetime] = mapped_column(UtcTimeStamp, server_default=func.now())
+
+    prompt: Mapped["Prompt"] = relationship("Prompt", back_populates="prompt_versions")
+
+    prompt_version_tags: Mapped[list["PromptVersionTag"]] = relationship(
+        "PromptVersionTag",
+        back_populates="prompt_version",
+        cascade="all, delete-orphan",
+        uselist=True,
+    )
+
+
+class PromptVersionTag(Base):
+    __tablename__ = "prompt_version_tags"
+
+    id: Mapped[int] = mapped_column(primary_key=True)
+    name: Mapped[Identifier] = mapped_column(_Identifier, nullable=False)
+    description: Mapped[Optional[str]] = mapped_column(String, nullable=True)
+    prompt_id: Mapped[int] = mapped_column(
+        ForeignKey("prompts.id", ondelete="CASCADE"),
+        index=True,
+        nullable=False,
+    )
+    prompt_version_id: Mapped[int] = mapped_column(
+        ForeignKey("prompt_versions.id", ondelete="CASCADE"),
+        index=True,
+        nullable=False,
+    )
+    user_id: Mapped[Optional[int]] = mapped_column(
+        ForeignKey("users.id", ondelete="SET NULL"),
+        index=True,
+        nullable=True,
+    )
+
+    prompt: Mapped["Prompt"] = relationship("Prompt", back_populates="prompt_version_tags")
+    prompt_version: Mapped["PromptVersion"] = relationship(
+        "PromptVersion", back_populates="prompt_version_tags"
+    )
+
+    __table_args__ = (UniqueConstraint("name", "prompt_id"),)

phoenix/db/types/identifier.py

@@ -0,0 +1,7 @@
+from typing import Annotated
+
+from pydantic import Field, RootModel
+
+
+class Identifier(RootModel[str]):
+    root: Annotated[str, Field(pattern=r"^[a-z0-9]([_a-z0-9-]*[a-z0-9])?$")]

phoenix/db/types/model_provider.py

@@ -0,0 +1,8 @@
+from enum import Enum
+
+
+class ModelProvider(Enum):
+    OPENAI = "OPENAI"
+    AZURE_OPENAI = "AZURE_OPENAI"
+    ANTHROPIC = "ANTHROPIC"
+    GOOGLE = "GOOGLE"

phoenix/server/api/context.py

@@ -30,6 +30,7 @@ from phoenix.server.api.dataloaders import (
     LatencyMsQuantileDataLoader,
     MinStartOrMaxEndTimeDataLoader,
     ProjectByNameDataLoader,
+    PromptVersionSequenceNumberDataLoader,
     RecordCountDataLoader,
     SessionIODataLoader,
     SessionNumTracesDataLoader,
@@ -73,6 +74,7 @@ class DataLoaders:
     experiment_sequence_number: ExperimentSequenceNumberDataLoader
     latency_ms_quantile: LatencyMsQuantileDataLoader
     min_start_or_max_end_times: MinStartOrMaxEndTimeDataLoader
+    prompt_version_sequence_number: PromptVersionSequenceNumberDataLoader
     record_counts: RecordCountDataLoader
     session_first_inputs: SessionIODataLoader
     session_last_outputs: SessionIODataLoader

phoenix/server/api/dataloaders/__init__.py

@@ -18,6 +18,7 @@ from .experiment_sequence_number import ExperimentSequenceNumberDataLoader
 from .latency_ms_quantile import LatencyMsQuantileCache, LatencyMsQuantileDataLoader
 from .min_start_or_max_end_times import MinStartOrMaxEndTimeCache, MinStartOrMaxEndTimeDataLoader
 from .project_by_name import ProjectByNameDataLoader
+from .prompt_version_sequence_number import PromptVersionSequenceNumberDataLoader
 from .record_counts import RecordCountCache, RecordCountDataLoader
 from .session_io import SessionIODataLoader
 from .session_num_traces import SessionNumTracesDataLoader
@@ -50,6 +51,7 @@ __all__ = [
     "ExperimentSequenceNumberDataLoader",
     "LatencyMsQuantileDataLoader",
     "MinStartOrMaxEndTimeDataLoader",
+    "PromptVersionSequenceNumberDataLoader",
     "RecordCountDataLoader",
     "SessionIODataLoader",
     "SessionNumTracesDataLoader",

phoenix/server/api/dataloaders/prompt_version_sequence_number.py

@@ -0,0 +1,35 @@
+from typing import Optional
+
+from sqlalchemy import func, select
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+from phoenix.server.types import DbSessionFactory
+
+PromptVersionId: TypeAlias = int
+Key: TypeAlias = PromptVersionId
+Result: TypeAlias = Optional[int]
+
+
+class PromptVersionSequenceNumberDataLoader(DataLoader[Key, Result]):
+    def __init__(self, db: DbSessionFactory) -> None:
+        super().__init__(load_fn=self._load_fn)
+        self._db = db
+
+    async def _load_fn(self, keys: list[Key]) -> list[Result]:
+        prompt_version_ids = keys
+        row_number = (
+            func.row_number().over(
+                partition_by=models.PromptVersion.prompt_id,
+                order_by=models.PromptVersion.id,
+            )
+        ).label("sequence_number")
+        subq = select(models.PromptVersion.id.label("prompt_version_id"), row_number).subquery()
+        stmt = select(subq).where(subq.c.prompt_version_id.in_(prompt_version_ids))
+        async with self._db() as session:
+            result = {
+                prompt_version_id: seq_number
+                async for prompt_version_id, seq_number in await session.stream(stmt)
+            }
+        return [result.get(prompt_version_id) for prompt_version_id in keys]

phoenix/server/api/helpers/jsonschema.py

@@ -0,0 +1,135 @@
+from typing import Annotated, Any, Literal, Union
+
+from jsonschema import Draft7Validator, ValidationError
+from pydantic import AfterValidator, BaseModel, Field
+from typing_extensions import TypeAlias
+
+# This meta-schema describes valid JSON schemas according to the JSON Schema Draft 7 specification.
+# It is copied from https://json-schema.org/draft-07/schema#
+JSON_SCHEMA_DRAFT_7_META_SCHEMA = {
+    "$schema": "http://json-schema.org/draft-07/schema#",
+    "$id": "http://json-schema.org/draft-07/schema#",
+    "title": "Core schema meta-schema",
+    "definitions": {
+        "schemaArray": {"type": "array", "minItems": 1, "items": {"$ref": "#"}},
+        "nonNegativeInteger": {"type": "integer", "minimum": 0},
+        "nonNegativeIntegerDefault0": {
+            "allOf": [{"$ref": "#/definitions/nonNegativeInteger"}, {"default": 0}]
+        },
+        "simpleTypes": {
+            "enum": ["array", "boolean", "integer", "null", "number", "object", "string"]
+        },
+        "stringArray": {
+            "type": "array",
+            "items": {"type": "string"},
+            "uniqueItems": True,
+            "default": [],
+        },
+    },
+    "type": ["object", "boolean"],
+    "properties": {
+        "$id": {"type": "string", "format": "uri-reference"},
+        "$schema": {"type": "string", "format": "uri"},
+        "$ref": {"type": "string", "format": "uri-reference"},
+        "$comment": {"type": "string"},
+        "title": {"type": "string"},
+        "description": {"type": "string"},
+        "default": True,
+        "readOnly": {"type": "boolean", "default": False},
+        "writeOnly": {"type": "boolean", "default": False},
+        "examples": {"type": "array", "items": True},
+        "multipleOf": {"type": "number", "exclusiveMinimum": 0},
+        "maximum": {"type": "number"},
+        "exclusiveMaximum": {"type": "number"},
+        "minimum": {"type": "number"},
+        "exclusiveMinimum": {"type": "number"},
+        "maxLength": {"$ref": "#/definitions/nonNegativeInteger"},
+        "minLength": {"$ref": "#/definitions/nonNegativeIntegerDefault0"},
+        "pattern": {"type": "string", "format": "regex"},
+        "additionalItems": {"$ref": "#"},
+        "items": {"anyOf": [{"$ref": "#"}, {"$ref": "#/definitions/schemaArray"}], "default": True},
+        "maxItems": {"$ref": "#/definitions/nonNegativeInteger"},
+        "minItems": {"$ref": "#/definitions/nonNegativeIntegerDefault0"},
+        "uniqueItems": {"type": "boolean", "default": False},
+        "contains": {"$ref": "#"},
+        "maxProperties": {"$ref": "#/definitions/nonNegativeInteger"},
+        "minProperties": {"$ref": "#/definitions/nonNegativeIntegerDefault0"},
+        "required": {"$ref": "#/definitions/stringArray"},
+        "additionalProperties": {"$ref": "#"},
+        "definitions": {"type": "object", "additionalProperties": {"$ref": "#"}, "default": {}},
+        "properties": {"type": "object", "additionalProperties": {"$ref": "#"}, "default": {}},
+        "patternProperties": {
+            "type": "object",
+            "additionalProperties": {"$ref": "#"},
+            "propertyNames": {"format": "regex"},
+            "default": {},
+        },
+        "dependencies": {
+            "type": "object",
+            "additionalProperties": {
+                "anyOf": [{"$ref": "#"}, {"$ref": "#/definitions/stringArray"}]
+            },
+        },
+        "propertyNames": {"$ref": "#"},
+        "const": True,
+        "enum": {"type": "array", "items": True, "minItems": 1, "uniqueItems": True},
+        "type": {
+            "anyOf": [
+                {"$ref": "#/definitions/simpleTypes"},
+                {
+                    "type": "array",
+                    "items": {"$ref": "#/definitions/simpleTypes"},
+                    "minItems": 1,
+                    "uniqueItems": True,
+                },
+            ]
+        },
+        "format": {"type": "string"},
+        "contentMediaType": {"type": "string"},
+        "contentEncoding": {"type": "string"},
+        "if": {"$ref": "#"},
+        "then": {"$ref": "#"},
+        "else": {"$ref": "#"},
+        "allOf": {"$ref": "#/definitions/schemaArray"},
+        "anyOf": {"$ref": "#/definitions/schemaArray"},
+        "oneOf": {"$ref": "#/definitions/schemaArray"},
+        "not": {"$ref": "#"},
+    },
+    "default": True,
+}
+Draft7Validator.check_schema(JSON_SCHEMA_DRAFT_7_META_SCHEMA)  # ensure the schema is valid
+JSON_SCHEMA_DRAFT_7_VALIDATOR = Draft7Validator(JSON_SCHEMA_DRAFT_7_META_SCHEMA)
+
+
+def validate_json_schema_draft_7_object(schema: dict[str, Any]) -> dict[str, Any]:
+    """
+    Validates that a dictionary is a valid JSON schema according to the JSON
+    Schema Draft 7 specification.
+    """
+    try:
+        JSON_SCHEMA_DRAFT_7_VALIDATOR.validate(schema)
+    except ValidationError as error:
+        raise ValueError(str(error))
+    if schema.get("type") != "object":
+        raise ValueError("The 'type' property must be 'object'")
+    return schema
+
+
+JSONSchemaDraft7ObjectSchemaContent: TypeAlias = Annotated[
+    dict[str, Any],
+    AfterValidator(validate_json_schema_draft_7_object),
+]
+
+
+class JSONSchemaDraft7ObjectSchema(BaseModel):
+    type: Literal["json-schema-draft-7-object-schema"]
+    json_: JSONSchemaDraft7ObjectSchemaContent = Field(
+        ...,
+        alias="json",  # avoid conflict with pydantic json method
+    )
+
+
+JSONSchemaObjectSchema: TypeAlias = Annotated[
+    Union[JSONSchemaDraft7ObjectSchema],
+    Field(discriminator="type"),
+]

phoenix/server/api/helpers/playground_clients.py

@@ -856,7 +856,7 @@ class AnthropicStreamingClient(PlaygroundStreamingClient):
 
 
 @register_llm_client(
-    provider_key=GenerativeProviderKey.GEMINI,
+    provider_key=GenerativeProviderKey.GOOGLE,
     model_names=[
         PROVIDER_DEFAULT,
         "gemini-2.0-flash-exp",
@@ -866,7 +866,7 @@ class AnthropicStreamingClient(PlaygroundStreamingClient):
         "gemini-1.0-pro",
     ],
 )
-class GeminiStreamingClient(PlaygroundStreamingClient):
+class GoogleStreamingClient(PlaygroundStreamingClient):
     def __init__(
         self,
         model: GenerativeModelInput,
@@ -941,7 +941,7 @@ class GeminiStreamingClient(PlaygroundStreamingClient):
     ) -> AsyncIterator[ChatCompletionChunk]:
         import google.generativeai as google_genai
 
-        gemini_message_history, current_message, system_prompt = self._build_gemini_messages(
+        google_message_history, current_message, system_prompt = self._build_google_messages(
             messages
         )
 
@@ -950,17 +950,17 @@ class GeminiStreamingClient(PlaygroundStreamingClient):
             model_args["system_instruction"] = system_prompt
         client = google_genai.GenerativeModel(**model_args)
 
-        gemini_config = google_genai.GenerationConfig(
+        google_config = google_genai.GenerationConfig(
             **invocation_parameters,
         )
-        gemini_params = {
+        google_params = {
             "content": current_message,
-            "generation_config": gemini_config,
+            "generation_config": google_config,
             "stream": True,
         }
 
-        chat = client.start_chat(history=gemini_message_history)
-        stream = await chat.send_message_async(**gemini_params)
+        chat = client.start_chat(history=google_message_history)
+        stream = await chat.send_message_async(**google_params)
         async for event in stream:
             self._attributes.update(
                 {
@@ -971,29 +971,29 @@ class GeminiStreamingClient(PlaygroundStreamingClient):
             )
             yield TextChunk(content=event.text)
 
-    def _build_gemini_messages(
+    def _build_google_messages(
         self,
         messages: list[tuple[ChatCompletionMessageRole, str, Optional[str], Optional[list[str]]]],
     ) -> tuple[list["ContentType"], str, str]:
-        gemini_message_history: list["ContentType"] = []
+        google_message_history: list["ContentType"] = []
         system_prompts = []
         for role, content, _tool_call_id, _tool_calls in messages:
             if role == ChatCompletionMessageRole.USER:
-                gemini_message_history.append({"role": "user", "parts": content})
+                google_message_history.append({"role": "user", "parts": content})
             elif role == ChatCompletionMessageRole.AI:
-                gemini_message_history.append({"role": "model", "parts": content})
+                google_message_history.append({"role": "model", "parts": content})
             elif role == ChatCompletionMessageRole.SYSTEM:
                 system_prompts.append(content)
             elif role == ChatCompletionMessageRole.TOOL:
                 raise NotImplementedError
             else:
                 assert_never(role)
-        if gemini_message_history:
-            prompt = gemini_message_history.pop()["parts"]
+        if google_message_history:
+            prompt = google_message_history.pop()["parts"]
         else:
             prompt = ""
 
-        return gemini_message_history, prompt, "\n".join(system_prompts)
+        return google_message_history, prompt, "\n".join(system_prompts)
 
 
 def initialize_playground_clients() -> None:

phoenix/server/api/helpers/playground_spans.py

@@ -41,6 +41,7 @@ from phoenix.server.api.types.ChatCompletionSubscriptionPayload import (
     TextChunk,
     ToolCallChunk,
 )
+from phoenix.server.api.types.Identifier import Identifier
 from phoenix.trace.attributes import get_attribute_value, unflatten
 from phoenix.trace.schemas import (
     SpanEvent,
@@ -70,6 +71,8 @@ class streaming_llm_span:
     ) -> None:
         self._input = input
         self._attributes: dict[str, Any] = attributes if attributes is not None else {}
+        self._attributes.update(dict(prompt_metadata(input.prompt_name)))
+
         self._attributes.update(
             chain(
                 llm_span_kind(),
@@ -264,6 +267,11 @@ def input_value_and_mime_type(
     yield INPUT_VALUE, safe_json_dumps(input_data)
 
 
+def prompt_metadata(prompt_name: Optional[Identifier]) -> Iterator[tuple[str, Any]]:
+    if prompt_name:
+        yield METADATA, {"phoenix_prompt_id": prompt_name}
+
+
 def _merge_tool_call_chunks(
     chunks_by_id: defaultdict[str, list[ToolCallChunk]],
 ) -> list[dict[str, Any]]:
@@ -442,6 +450,7 @@ LLM_INVOCATION_PARAMETERS = SpanAttributes.LLM_INVOCATION_PARAMETERS
 LLM_TOOLS = SpanAttributes.LLM_TOOLS
 LLM_TOKEN_COUNT_PROMPT = SpanAttributes.LLM_TOKEN_COUNT_PROMPT
 LLM_TOKEN_COUNT_COMPLETION = SpanAttributes.LLM_TOKEN_COUNT_COMPLETION
+METADATA = SpanAttributes.METADATA
 
 MESSAGE_CONTENT = MessageAttributes.MESSAGE_CONTENT
 MESSAGE_ROLE = MessageAttributes.MESSAGE_ROLE