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

phoenix/db/engines.py

@@ -5,7 +5,7 @@ import logging
 from collections.abc import Callable
 from enum import Enum
 from sqlite3 import Connection
-from typing import Any
+from typing import Any, Optional
 
 import aiosqlite
 import numpy as np
@@ -168,29 +168,142 @@ def aio_postgresql_engine(
     log_to_stdout: bool = False,
     log_migrations_to_stdout: bool = True,
 ) -> AsyncEngine:
-    asyncpg_url, asyncpg_args = get_pg_config(url, "asyncpg")
-    engine = create_async_engine(
-        url=asyncpg_url,
-        connect_args=asyncpg_args,
-        echo=log_to_stdout,
-        json_serializer=_dumps,
+    from phoenix.config import (
+        get_env_postgres_iam_token_lifetime,
+        get_env_postgres_use_iam_auth,
     )
+
+    use_iam_auth = get_env_postgres_use_iam_auth()
+
+    asyncpg_url, asyncpg_args = get_pg_config(url, "asyncpg", enforce_ssl=use_iam_auth)
+
+    iam_config: Optional[dict[str, Any]] = None
+    token_lifetime: int = 0
+    if use_iam_auth:
+        iam_config = _extract_iam_config_from_url(url)
+        token_lifetime = get_env_postgres_iam_token_lifetime()
+
+        async def iam_async_creator() -> Any:
+            import asyncpg  # type: ignore
+
+            from phoenix.db.iam_auth import generate_aws_rds_token
+
+            assert iam_config is not None
+            token = generate_aws_rds_token(
+                host=iam_config["host"],
+                port=iam_config["port"],
+                user=iam_config["user"],
+            )
+
+            conn_kwargs = {
+                "host": iam_config["host"],
+                "port": iam_config["port"],
+                "user": iam_config["user"],
+                "password": token,
+                "database": iam_config["database"],
+            }
+
+            if asyncpg_args:
+                conn_kwargs.update(asyncpg_args)
+
+            return await asyncpg.connect(**conn_kwargs)
+
+        engine = create_async_engine(
+            url=asyncpg_url,
+            async_creator=iam_async_creator,
+            echo=log_to_stdout,
+            json_serializer=_dumps,
+            pool_recycle=token_lifetime,
+        )
+    else:
+        engine = create_async_engine(
+            url=asyncpg_url,
+            connect_args=asyncpg_args,
+            echo=log_to_stdout,
+            json_serializer=_dumps,
+        )
+
     if not migrate:
         return engine
 
-    psycopg_url, psycopg_args = get_pg_config(url, "psycopg")
-    sync_engine = sqlalchemy.create_engine(
-        url=psycopg_url,
-        connect_args=psycopg_args,
-        echo=log_migrations_to_stdout,
-        json_serializer=_dumps,
-    )
+    psycopg_url, psycopg_args = get_pg_config(url, "psycopg", enforce_ssl=use_iam_auth)
+
+    if use_iam_auth:
+        assert iam_config is not None
+
+        def iam_sync_creator() -> Any:
+            import psycopg
+
+            from phoenix.db.iam_auth import generate_aws_rds_token
+
+            token = generate_aws_rds_token(
+                host=iam_config["host"],
+                port=iam_config["port"],
+                user=iam_config["user"],
+            )
+
+            conn_kwargs = {
+                "host": iam_config["host"],
+                "port": iam_config["port"],
+                "user": iam_config["user"],
+                "password": token,
+                "dbname": iam_config["database"],
+            }
+
+            if psycopg_args:
+                conn_kwargs.update(psycopg_args)
+
+            return psycopg.connect(**conn_kwargs)
+
+        sync_engine = sqlalchemy.create_engine(
+            url=psycopg_url,
+            creator=iam_sync_creator,
+            echo=log_migrations_to_stdout,
+            json_serializer=_dumps,
+            pool_recycle=token_lifetime,
+        )
+    else:
+        sync_engine = sqlalchemy.create_engine(
+            url=psycopg_url,
+            connect_args=psycopg_args,
+            echo=log_migrations_to_stdout,
+            json_serializer=_dumps,
+        )
+
     if schema := get_env_database_schema():
         event.listen(sync_engine, "connect", set_postgresql_search_path(schema))
     migrate_in_thread(sync_engine)
     return engine
 
 
+def _extract_iam_config_from_url(url: URL) -> dict[str, Any]:
+    """Extract connection parameters needed for IAM authentication from a SQLAlchemy URL.
+
+    Args:
+        url: SQLAlchemy database URL
+
+    Returns:
+        Dictionary with host, port, user, and database
+    """
+    host = url.host
+    if not host:
+        raise ValueError("Database host is required for IAM authentication")
+
+    port = url.port or 5432
+    user = url.username
+    if not user:
+        raise ValueError("Database user is required for IAM authentication")
+
+    database = url.database or "postgres"
+
+    return {
+        "host": host,
+        "port": port,
+        "user": user,
+        "database": database,
+    }
+
+
 def _dumps(obj: Any) -> str:
     return orjson.dumps(obj, default=_default).decode()
 

phoenix/db/iam_auth.py

@@ -0,0 +1,64 @@
+from __future__ import annotations
+
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+def generate_aws_rds_token(
+    host: str,
+    port: int,
+    user: str,
+) -> str:
+    """Generate an AWS RDS IAM authentication token.
+
+    This function creates a short-lived (15 minutes) authentication token for connecting
+    to AWS RDS/Aurora PostgreSQL instances using IAM database authentication.
+
+    The AWS region is automatically resolved using boto3.
+
+    Args:
+        host: The database hostname (e.g., 'mydb.abc123.us-west-2.rds.amazonaws.com')
+        port: The database port (typically 5432 for PostgreSQL)
+        user: The database username (must match an IAM-enabled database user)
+
+    Returns:
+        A temporary authentication token string to use as the database password
+
+    Raises:
+        ImportError: If boto3 is not installed
+        Exception: If AWS credentials/region are not configured or token generation fails
+
+    Example:
+        >>> token = generate_aws_rds_token(
+        ...     host='mydb.us-west-2.rds.amazonaws.com',
+        ...     port=5432,
+        ...     user='myuser'
+        ... )
+    """
+    try:
+        import boto3  # type: ignore
+    except ImportError as e:
+        raise ImportError(
+            "boto3 is required for AWS RDS IAM authentication. "
+            "Install it with: pip install 'arize-phoenix[aws]'"
+        ) from e
+
+    try:
+        client = boto3.client("rds")
+
+        logger.debug(f"Generating AWS RDS IAM auth token for user '{user}' at {host}:{port}")
+        token = client.generate_db_auth_token(  # pyright: ignore
+            DBHostname=host,
+            Port=port,
+            DBUsername=user,
+        )
+
+        return str(token)  # pyright: ignore
+
+    except Exception as e:
+        logger.error(
+            f"Failed to generate AWS RDS IAM authentication token: {e}. "
+            "Ensure AWS credentials are configured and have 'rds-db:connect' permission."
+        )
+        raise

phoenix/db/pg_config.py

@@ -10,12 +10,14 @@ from typing_extensions import assert_never
 def get_pg_config(
     url: URL,
     driver: Literal["psycopg", "asyncpg"],
+    enforce_ssl: bool = False,
 ) -> tuple[URL, dict[str, Any]]:
     """Convert SQLAlchemy URL to driver-specific configuration.
 
     Args:
         url: SQLAlchemy URL
         driver: "psycopg" or "asyncpg"
+        enforce_ssl: If True, ensure SSL is enabled (required for AWS RDS IAM auth)
 
     Returns:
         Tuple of (base_url, connect_args):
@@ -26,6 +28,14 @@ def get_pg_config(
     query = url.query
     ssl_args = _get_ssl_args(query)
 
+    if enforce_ssl and not ssl_args:
+        ssl_args = {"sslmode": "require"}
+    elif enforce_ssl and ssl_args.get("sslmode") == "disable":
+        raise ValueError(
+            "SSL cannot be disabled when using AWS RDS IAM authentication. "
+            "Remove 'sslmode=disable' from the connection string."
+        )
+
     # Create base URL without SSL parameters
     base_url = url.set(
         drivername=f"postgresql+{driver}",

phoenix/server/api/context.py

@@ -35,6 +35,7 @@ from phoenix.server.api.dataloaders import (
     ExperimentRepeatedRunGroupsDataLoader,
     ExperimentRunAnnotations,
     ExperimentRunCountsDataLoader,
+    ExperimentRunsByExperimentAndExampleDataLoader,
     ExperimentSequenceNumberDataLoader,
     LastUsedTimesByGenerativeModelIdDataLoader,
     LatencyMsQuantileDataLoader,
@@ -71,6 +72,7 @@ from phoenix.server.api.dataloaders import (
     SpanProjectsDataLoader,
     TableFieldsDataLoader,
     TokenCountDataLoader,
+    TokenPricesByModelDataLoader,
     TraceAnnotationsByTraceDataLoader,
     TraceByTraceIdsDataLoader,
     TraceRetentionPolicyIdByProjectIdDataLoader,
@@ -100,27 +102,38 @@ class DataLoaders:
         AverageExperimentRepeatedRunGroupLatencyDataLoader
     )
     average_experiment_run_latency: AverageExperimentRunLatencyDataLoader
+    dataset_example_fields: TableFieldsDataLoader
     dataset_example_revisions: DatasetExampleRevisionsDataLoader
     dataset_example_spans: DatasetExampleSpansDataLoader
     dataset_labels: DatasetLabelsDataLoader
+    dataset_label_fields: TableFieldsDataLoader
     dataset_dataset_splits: DatasetDatasetSplitsDataLoader
     dataset_examples_and_versions_by_experiment_run: (
         DatasetExamplesAndVersionsByExperimentRunDataLoader
     )
     dataset_example_splits: DatasetExampleSplitsDataLoader
+    dataset_fields: TableFieldsDataLoader
+    dataset_split_fields: TableFieldsDataLoader
+    dataset_version_fields: TableFieldsDataLoader
+    document_annotation_fields: TableFieldsDataLoader
     document_evaluation_summaries: DocumentEvaluationSummaryDataLoader
     document_evaluations: DocumentEvaluationsDataLoader
     document_retrieval_metrics: DocumentRetrievalMetricsDataLoader
     experiment_annotation_summaries: ExperimentAnnotationSummaryDataLoader
     experiment_dataset_splits: ExperimentDatasetSplitsDataLoader
     experiment_error_rates: ExperimentErrorRatesDataLoader
+    experiment_fields: TableFieldsDataLoader
     experiment_repeated_run_group_annotation_summaries: (
         ExperimentRepeatedRunGroupAnnotationSummariesDataLoader
     )
     experiment_repeated_run_groups: ExperimentRepeatedRunGroupsDataLoader
+    experiment_run_annotation_fields: TableFieldsDataLoader
     experiment_run_annotations: ExperimentRunAnnotations
     experiment_run_counts: ExperimentRunCountsDataLoader
+    experiment_run_fields: TableFieldsDataLoader
+    experiment_runs_by_experiment_and_example: ExperimentRunsByExperimentAndExampleDataLoader
     experiment_sequence_number: ExperimentSequenceNumberDataLoader
+    generative_model_fields: TableFieldsDataLoader
     last_used_times_by_generative_model_id: LastUsedTimesByGenerativeModelIdDataLoader
     latency_ms_quantile: LatencyMsQuantileDataLoader
     min_start_or_max_end_times: MinStartOrMaxEndTimeDataLoader
@@ -130,7 +143,12 @@ class DataLoaders:
     project_fields: TableFieldsDataLoader
     project_trace_retention_policy_fields: TableFieldsDataLoader
     projects_by_trace_retention_policy_id: ProjectIdsByTraceRetentionPolicyIdDataLoader
+    prompt_fields: TableFieldsDataLoader
+    prompt_label_fields: TableFieldsDataLoader
     prompt_version_sequence_number: PromptVersionSequenceNumberDataLoader
+    prompt_version_tag_fields: TableFieldsDataLoader
+    project_session_annotation_fields: TableFieldsDataLoader
+    project_session_fields: TableFieldsDataLoader
     record_counts: RecordCountDataLoader
     session_annotations_by_session: SessionAnnotationsBySessionDataLoader
     session_first_inputs: SessionIODataLoader
@@ -139,6 +157,7 @@ class DataLoaders:
     session_num_traces_with_error: SessionNumTracesWithErrorDataLoader
     session_token_usages: SessionTokenUsagesDataLoader
     session_trace_latency_ms_quantile: SessionTraceLatencyMsQuantileDataLoader
+    span_annotation_fields: TableFieldsDataLoader
     span_annotations: SpanAnnotationsDataLoader
     span_by_id: SpanByIdDataLoader
     span_cost_by_span: SpanCostBySpanDataLoader
@@ -167,12 +186,16 @@ class DataLoaders:
     span_fields: TableFieldsDataLoader
     span_projects: SpanProjectsDataLoader
     token_counts: TokenCountDataLoader
+    token_prices_by_model: TokenPricesByModelDataLoader
+    trace_annotation_fields: TableFieldsDataLoader
     trace_annotations_by_trace: TraceAnnotationsByTraceDataLoader
     trace_by_trace_ids: TraceByTraceIdsDataLoader
     trace_fields: TableFieldsDataLoader
     trace_retention_policy_id_by_project_id: TraceRetentionPolicyIdByProjectIdDataLoader
     trace_root_spans: TraceRootSpansDataLoader
     user_roles: UserRolesDataLoader
+    user_api_key_fields: TableFieldsDataLoader
+    user_fields: TableFieldsDataLoader
     users: UsersDataLoader
 
 

phoenix/server/api/dataloaders/__init__.py

@@ -33,6 +33,9 @@ from .experiment_repeated_run_group_annotation_summaries import (
 from .experiment_repeated_run_groups import ExperimentRepeatedRunGroupsDataLoader
 from .experiment_run_annotations import ExperimentRunAnnotations
 from .experiment_run_counts import ExperimentRunCountsDataLoader
+from .experiment_runs_by_experiment_and_example import (
+    ExperimentRunsByExperimentAndExampleDataLoader,
+)
 from .experiment_sequence_number import ExperimentSequenceNumberDataLoader
 from .last_used_times_by_generative_model_id import LastUsedTimesByGenerativeModelIdDataLoader
 from .latency_ms_quantile import LatencyMsQuantileCache, LatencyMsQuantileDataLoader
@@ -73,6 +76,7 @@ from .span_descendants import SpanDescendantsDataLoader
 from .span_projects import SpanProjectsDataLoader
 from .table_fields import TableFieldsDataLoader
 from .token_counts import TokenCountCache, TokenCountDataLoader
+from .token_prices_by_model import TokenPricesByModelDataLoader
 from .trace_annotations_by_trace import TraceAnnotationsByTraceDataLoader
 from .trace_by_trace_ids import TraceByTraceIdsDataLoader
 from .trace_retention_policy_id_by_project_id import TraceRetentionPolicyIdByProjectIdDataLoader
@@ -102,6 +106,7 @@ __all__ = [
     "ExperimentRepeatedRunGroupAnnotationSummariesDataLoader",
     "ExperimentRunAnnotations",
     "ExperimentRunCountsDataLoader",
+    "ExperimentRunsByExperimentAndExampleDataLoader",
     "ExperimentSequenceNumberDataLoader",
     "LastUsedTimesByGenerativeModelIdDataLoader",
     "LatencyMsQuantileDataLoader",
@@ -139,6 +144,7 @@ __all__ = [
     "SpanProjectsDataLoader",
     "TableFieldsDataLoader",
     "TokenCountDataLoader",
+    "TokenPricesByModelDataLoader",
     "TraceAnnotationsByTraceDataLoader",
     "TraceByTraceIdsDataLoader",
     "TraceRetentionPolicyIdByProjectIdDataLoader",

phoenix/server/api/dataloaders/experiment_repeated_run_groups.py

@@ -1,7 +1,6 @@
 from dataclasses import dataclass
 
 from sqlalchemy import select, tuple_
-from sqlalchemy.orm import joinedload
 from strawberry.dataloader import DataLoader
 from typing_extensions import TypeAlias
 
@@ -38,7 +37,6 @@ class ExperimentRepeatedRunGroupsDataLoader(DataLoader[Key, Result]):
                 ).in_(set(keys))
             )
             .order_by(models.ExperimentRun.repetition_number)
-            .options(joinedload(models.ExperimentRun.trace).load_only(models.Trace.trace_id))
         )
 
         async with self._db() as session:

phoenix/server/api/dataloaders/experiment_runs_by_experiment_and_example.py

@@ -0,0 +1,44 @@
+from collections import defaultdict
+
+from sqlalchemy import select, tuple_
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+from phoenix.server.types import DbSessionFactory
+
+ExperimentId: TypeAlias = int
+DatasetExampleId: TypeAlias = int
+Key: TypeAlias = tuple[ExperimentId, DatasetExampleId]
+Result: TypeAlias = list[models.ExperimentRun]
+
+
+class ExperimentRunsByExperimentAndExampleDataLoader(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]:
+        runs_by_key: defaultdict[Key, Result] = defaultdict(list)
+
+        async with self._db() as session:
+            stmt = (
+                select(models.ExperimentRun)
+                .where(
+                    tuple_(
+                        models.ExperimentRun.experiment_id,
+                        models.ExperimentRun.dataset_example_id,
+                    ).in_(keys)
+                )
+                .order_by(
+                    models.ExperimentRun.experiment_id,
+                    models.ExperimentRun.dataset_example_id,
+                    models.ExperimentRun.repetition_number,
+                )
+            )
+            result = await session.stream_scalars(stmt)
+            async for run in result:
+                key = (run.experiment_id, run.dataset_example_id)
+                runs_by_key[key].append(run)
+
+        return [runs_by_key[key] for key in keys]

phoenix/server/api/dataloaders/span_costs.py

@@ -1,7 +1,6 @@
 from typing import Optional
 
 from sqlalchemy import select
-from sqlalchemy.orm import joinedload, load_only
 from strawberry.dataloader import DataLoader
 from typing_extensions import TypeAlias
 
@@ -22,14 +21,9 @@ class SpanCostsDataLoader(DataLoader[Key, Result]):
         span_ids = list(set(keys))
         async with self._db() as session:
             costs = {
-                span.id: span.span_cost
-                async for span in await session.stream_scalars(
-                    select(models.Span)
-                    .where(models.Span.id.in_(span_ids))
-                    .options(
-                        load_only(models.Span.id),
-                        joinedload(models.Span.span_cost),
-                    )
+                span_cost.span_rowid: span_cost
+                async for span_cost in await session.stream_scalars(
+                    select(models.SpanCost).where(models.SpanCost.span_rowid.in_(span_ids))
                 )
             }
         return [costs.get(span_id) for span_id in keys]

phoenix/server/api/dataloaders/token_prices_by_model.py

@@ -0,0 +1,30 @@
+from collections import defaultdict
+
+from sqlalchemy import select
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+from phoenix.server.types import DbSessionFactory
+
+ModelId: TypeAlias = int
+Key: TypeAlias = ModelId
+Result: TypeAlias = list[models.TokenPrice]
+
+
+class TokenPricesByModelDataLoader(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]:
+        model_ids = keys
+        token_prices: defaultdict[Key, Result] = defaultdict(list)
+
+        async with self._db() as session:
+            async for token_price in await session.stream_scalars(
+                select(models.TokenPrice).where(models.TokenPrice.model_id.in_(model_ids))
+            ):
+                token_prices[token_price.model_id].append(token_price)
+
+        return [token_prices[model_id] for model_id in keys]

phoenix/server/api/helpers/playground_clients.py

@@ -699,7 +699,7 @@ class BedrockStreamingClient(PlaygroundStreamingClient):
         self.model_name = model.name
         self.client = boto3.client(
             service_name="bedrock-runtime",
-            region_name="us-east-1",  # match the default region in the UI
+            region_name=self.region,
             aws_access_key_id=self.aws_access_key_id,
             aws_secret_access_key=self.aws_secret_access_key,
             aws_session_token=self.aws_session_token,
@@ -805,7 +805,7 @@ class BedrockStreamingClient(PlaygroundStreamingClient):
 
         # Build the request parameters for Converse API
         converse_params: dict[str, Any] = {
-            "modelId": f"us.{self.model_name}",
+            "modelId": self.model_name,
             "messages": converse_messages,
             "inferenceConfig": {
                 "maxTokens": invocation_parameters["max_tokens"],
@@ -953,7 +953,7 @@ class BedrockStreamingClient(PlaygroundStreamingClient):
         }
 
         response = self.client.invoke_model_with_response_stream(
-            modelId=f"us.{self.model_name}",  # or another Claude model
+            modelId=self.model_name,
             contentType="application/json",
             accept="application/json",
             body=json.dumps(bedrock_params),

phoenix/server/api/input_types/PromptVersionInput.py

@@ -1,10 +1,11 @@
 import json
-from typing import Optional, cast
+from typing import Any, Optional, Union, cast
 
 import strawberry
 from strawberry import UNSET
 from strawberry.scalars import JSON
 
+from phoenix.db import models
 from phoenix.db.types.model_provider import ModelProvider
 from phoenix.server.api.helpers.prompts.models import (
     ContentPart,
@@ -12,11 +13,15 @@ from phoenix.server.api.helpers.prompts.models import (
     PromptMessage,
     PromptMessageRole,
     PromptTemplateFormat,
+    PromptTemplateType,
     RoleConversion,
     TextContentPart,
     ToolCallContentPart,
     ToolCallFunction,
     ToolResultContentPart,
+    normalize_response_format,
+    normalize_tools,
+    validate_invocation_parameters,
 )
 
 
@@ -88,6 +93,47 @@ class ChatPromptVersionInput:
             k: v for k, v in self.invocation_parameters.items() if v is not None
         }
 
+    def to_orm_prompt_version(
+        self,
+        user_id: Optional[int],
+    ) -> models.PromptVersion:
+        tool_definitions = [tool.definition for tool in self.tools]
+        tool_choice = cast(
+            Optional[Union[str, dict[str, Any]]],
+            cast(dict[str, Any], self.invocation_parameters).pop("tool_choice", None),
+        )
+        model_provider = ModelProvider(self.model_provider)
+        tools = (
+            normalize_tools(tool_definitions, model_provider, tool_choice)
+            if tool_definitions
+            else None
+        )
+        template = to_pydantic_prompt_chat_template_v1(self.template)
+        response_format = (
+            normalize_response_format(
+                self.response_format.definition,
+                model_provider,
+            )
+            if self.response_format
+            else None
+        )
+        invocation_parameters = validate_invocation_parameters(
+            self.invocation_parameters,
+            model_provider,
+        )
+        return models.PromptVersion(
+            description=self.description,
+            user_id=user_id,
+            template_type=PromptTemplateType.CHAT,
+            template_format=self.template_format,
+            template=template,
+            invocation_parameters=invocation_parameters,
+            tools=tools,
+            response_format=response_format,
+            model_provider=ModelProvider(self.model_provider),
+            model_name=self.model_name,
+        )
+
 
 def to_pydantic_prompt_chat_template_v1(
     prompt_chat_template_input: PromptChatTemplateInput,

phoenix/server/api/mutations/annotation_config_mutations.py

@@ -374,7 +374,7 @@ class AnnotationConfigMutationMixin:
                 )
             return AddAnnotationConfigToProjectPayload(
                 query=Query(),
-                project=Project(project_rowid=project_id),
+                project=Project(id=project_id),
             )
 
     @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer])  # type: ignore[misc]
@@ -409,5 +409,5 @@ class AnnotationConfigMutationMixin:
                 raise NotFound("Could not find one or more input project annotation configs")
         return RemoveAnnotationConfigFromProjectPayload(
             query=Query(),
-            project=Project(project_rowid=project_id),
+            project=Project(id=project_id),
         )

phoenix/server/api/mutations/api_key_mutations.py

@@ -92,13 +92,7 @@ class ApiKeyMutationMixin:
         token, token_id = await token_store.create_api_key(claims)
         return CreateSystemApiKeyMutationPayload(
             jwt=token,
-            api_key=SystemApiKey(
-                id_attr=int(token_id),
-                name=input.name,
-                description=input.description or None,
-                created_at=issued_at,
-                expires_at=input.expires_at or None,
-            ),
+            api_key=SystemApiKey(id=int(token_id)),
             query=Query(),
         )
 
@@ -134,14 +128,7 @@ class ApiKeyMutationMixin:
         token, token_id = await token_store.create_api_key(claims)
         return CreateUserApiKeyMutationPayload(
             jwt=token,
-            api_key=UserApiKey(
-                id_attr=int(token_id),
-                name=input.name,
-                description=input.description or None,
-                created_at=issued_at,
-                expires_at=input.expires_at or None,
-                user_id=int(user.identity),
-            ),
+            api_key=UserApiKey(id=int(token_id)),
             query=Query(),
         )
 

phoenix/server/api/mutations/chat_mutations.py

@@ -19,6 +19,7 @@ from openinference.semconv.trace import (
 from opentelemetry.sdk.trace.id_generator import RandomIdGenerator as DefaultOTelIDGenerator
 from opentelemetry.trace import StatusCode
 from sqlalchemy import insert, select
+from strawberry import Private
 from strawberry.relay import GlobalID
 from strawberry.types import Info
 from typing_extensions import assert_never
@@ -101,7 +102,7 @@ class ChatCompletionToolCall:
 
 @strawberry.type
 class ChatCompletionMutationPayload:
-    db_span: strawberry.Private[models.Span]
+    db_span: Private[models.Span]
     content: Optional[str]
     tool_calls: List[ChatCompletionToolCall]
     span: Span
@@ -502,7 +503,7 @@ class ChatCompletionMutationMixin:
                 session.add(span_cost)
                 await session.flush()
 
-        gql_span = Span(span_rowid=span.id, db_span=span)
+        gql_span = Span(id=span.id, db_record=span)
 
         info.context.event_queue.put(SpanInsertEvent(ids=(project_id,)))