arize-phoenix 11.29.0__py3-none-any.whl → 11.31.0__py3-none-any.whl

phoenix/server/api/routers/v1/documents.py

@@ -0,0 +1,143 @@
+from typing import Optional
+
+from fastapi import APIRouter, Depends, HTTPException, Query
+from pydantic import Field
+from sqlalchemy import select
+from starlette.requests import Request
+from starlette.status import HTTP_404_NOT_FOUND
+from strawberry.relay import GlobalID
+
+from phoenix.db import models
+from phoenix.db.helpers import SupportedSQLDialect
+from phoenix.db.insertion.helpers import as_kv, insert_on_conflict
+from phoenix.server.api.routers.v1.annotations import SpanDocumentAnnotationData
+from phoenix.server.api.types.DocumentAnnotation import DocumentAnnotation
+from phoenix.server.authorization import is_not_locked
+from phoenix.server.bearer_auth import PhoenixUser
+from phoenix.server.dml_event import DocumentAnnotationInsertEvent
+
+from .models import V1RoutesBaseModel
+from .utils import RequestBody, ResponseBody, add_errors_to_responses
+
+# Since the document annotations are spans related, we place it under spans
+router = APIRouter(tags=["spans"])
+
+
+class AnnotateSpanDocumentsRequestBody(RequestBody[list[SpanDocumentAnnotationData]]):
+    pass
+
+
+class InsertedSpanDocumentAnnotation(V1RoutesBaseModel):
+    id: str = Field(description="The ID of the inserted span document annotation")
+
+
+class AnnotateSpanDocumentsResponseBody(ResponseBody[list[InsertedSpanDocumentAnnotation]]):
+    pass
+
+
+@router.post(
+    "/document_annotations",
+    dependencies=[Depends(is_not_locked)],
+    operation_id="annotateSpanDocuments",
+    responses=add_errors_to_responses(
+        [
+            {
+                "status_code": HTTP_404_NOT_FOUND,
+                "description": "Span not found",
+            },
+            {
+                "status_code": 422,
+                "description": "Invalid request - non-empty identifier not supported",
+            },
+        ]
+    ),
+    response_description="Span document annotation inserted successfully",
+    include_in_schema=True,
+)
+async def annotate_span_documents(
+    request: Request,
+    request_body: AnnotateSpanDocumentsRequestBody,
+    sync: bool = Query(
+        default=False, description="If set to true, the annotations are inserted synchronously."
+    ),
+) -> AnnotateSpanDocumentsResponseBody:
+    if not request_body.data:
+        return AnnotateSpanDocumentsResponseBody(data=[])
+
+    # Validate that identifiers are empty or only whitespace
+    for annotation in request_body.data:
+        if annotation.identifier.strip():
+            raise HTTPException(
+                detail=f"Non-empty identifier '{annotation.identifier}' is not supported",
+                status_code=422,  # Unprocessable Entity
+            )
+
+    user_id: Optional[int] = None
+    if request.app.state.authentication_enabled and isinstance(request.user, PhoenixUser):
+        user_id = int(request.user.identity)
+
+    span_document_annotations = request_body.data
+
+    precursors = [
+        annotation.as_precursor(user_id=user_id) for annotation in span_document_annotations
+    ]
+    if not sync:
+        await request.state.enqueue(*precursors)
+        return AnnotateSpanDocumentsResponseBody(data=[])
+
+    span_ids = {p.span_id for p in precursors}
+    # Account for the fact that the spans could arrive after the annotation
+    async with request.app.state.db() as session:
+        existing_spans = {
+            span_id: (id_, num_docs)
+            async for span_id, id_, num_docs in await session.stream(
+                select(models.Span.span_id, models.Span.id, models.Span.num_documents).filter(
+                    models.Span.span_id.in_(span_ids)
+                )
+            )
+        }
+
+        missing_span_ids = span_ids - set(existing_spans.keys())
+        # We prefer to fail the entire operation if there are missing spans in sync mode
+        if missing_span_ids:
+            raise HTTPException(
+                detail=f"Spans with IDs {', '.join(missing_span_ids)} do not exist.",
+                status_code=HTTP_404_NOT_FOUND,
+            )
+
+        # Validate that document positions are within bounds
+        for annotation in span_document_annotations:
+            _, num_docs = existing_spans[annotation.span_id]
+            if annotation.document_position not in range(num_docs):
+                raise HTTPException(
+                    detail=f"Document position {annotation.document_position} is out of bounds for "
+                    f"span {annotation.span_id} (max: {num_docs - 1})",
+                    status_code=422,  # Unprocessable Entity
+                )
+
+        inserted_document_annotation_ids = []
+        dialect = SupportedSQLDialect(session.bind.dialect.name)
+        for anno in precursors:
+            span_rowid, _ = existing_spans[anno.span_id]
+            values = dict(as_kv(anno.as_insertable(span_rowid).row))
+            span_document_annotation_id = await session.scalar(
+                insert_on_conflict(
+                    values,
+                    dialect=dialect,
+                    table=models.DocumentAnnotation,
+                    unique_by=("name", "span_rowid", "identifier", "document_position"),
+                    constraint_name="uq_document_annotations_name_span_rowid_document_pos_identifier",
+                ).returning(models.DocumentAnnotation.id)
+            )
+            inserted_document_annotation_ids.append(span_document_annotation_id)
+
+    # We queue an event to let the application know that annotations have changed
+    request.state.event_queue.put(
+        DocumentAnnotationInsertEvent(tuple(inserted_document_annotation_ids))
+    )
+    return AnnotateSpanDocumentsResponseBody(
+        data=[
+            InsertedSpanDocumentAnnotation(id=str(GlobalID(DocumentAnnotation.__name__, str(id_))))
+            for id_ in inserted_document_annotation_ids
+        ]
+    )

phoenix/server/api/routers/v1/spans.py

@@ -5,7 +5,7 @@ from collections.abc import AsyncIterator
 from datetime import datetime, timezone
 from enum import Enum
 from secrets import token_urlsafe
-from typing import Annotated, Any, Literal, Optional, Union
+from typing import Annotated, Any, Optional, Union
 
 import pandas as pd
 import sqlalchemy as sa
@@ -27,8 +27,8 @@ from phoenix.datetime_utils import normalize_datetime
 from phoenix.db import models
 from phoenix.db.helpers import SupportedSQLDialect, get_ancestor_span_rowids
 from phoenix.db.insertion.helpers import as_kv, insert_on_conflict
-from phoenix.db.insertion.types import Precursors
 from phoenix.server.api.routers.utils import df_to_bytes
+from phoenix.server.api.routers.v1.annotations import SpanAnnotationData
 from phoenix.server.api.types.node import from_global_id_with_expected_type
 from phoenix.server.authorization import is_not_locked
 from phoenix.server.bearer_auth import PhoenixUser
@@ -850,52 +850,6 @@ async def get_spans_handler(
     return await query_spans_handler(request, request_body, project_name)
 
 
-class SpanAnnotationResult(V1RoutesBaseModel):
-    label: Optional[str] = Field(default=None, description="The label assigned by the annotation")
-    score: Optional[float] = Field(default=None, description="The score assigned by the annotation")
-    explanation: Optional[str] = Field(
-        default=None, description="Explanation of the annotation result"
-    )
-
-
-class SpanAnnotationData(V1RoutesBaseModel):
-    span_id: str = Field(description="OpenTelemetry Span ID (hex format w/o 0x prefix)")
-    name: str = Field(description="The name of the annotation")
-    annotator_kind: Literal["LLM", "CODE", "HUMAN"] = Field(
-        description="The kind of annotator used for the annotation"
-    )
-    result: Optional[SpanAnnotationResult] = Field(
-        default=None, description="The result of the annotation"
-    )
-    metadata: Optional[dict[str, Any]] = Field(
-        default=None, description="Metadata for the annotation"
-    )
-    identifier: str = Field(
-        default="",
-        description=(
-            "The identifier of the annotation. "
-            "If provided, the annotation will be updated if it already exists."
-        ),
-    )
-
-    def as_precursor(self, *, user_id: Optional[int] = None) -> Precursors.SpanAnnotation:
-        return Precursors.SpanAnnotation(
-            datetime.now(timezone.utc),
-            self.span_id,
-            models.SpanAnnotation(
-                name=self.name,
-                annotator_kind=self.annotator_kind,
-                score=self.result.score if self.result else None,
-                label=self.result.label if self.result else None,
-                explanation=self.result.explanation if self.result else None,
-                metadata_=self.metadata or {},
-                identifier=self.identifier,
-                source="API",
-                user_id=user_id,
-            ),
-        )
-
-
 class AnnotateSpansRequestBody(RequestBody[list[SpanAnnotationData]]):
     data: list[SpanAnnotationData]
 
@@ -949,9 +903,11 @@ async def annotate_spans(
     span_ids = {p.span_id for p in precursors}
     async with request.app.state.db() as session:
         existing_spans = {
-            span.span_id: span.id
-            async for span in await session.stream_scalars(
-                select(models.Span).filter(models.Span.span_id.in_(span_ids))
+            span_id: id_
+            async for span_id, id_ in await session.stream(
+                select(models.Span.span_id, models.Span.id).filter(
+                    models.Span.span_id.in_(span_ids)
+                )
             )
         }
 

phoenix/server/api/routers/v1/traces.py

@@ -1,7 +1,6 @@
 import gzip
 import zlib
-from datetime import datetime, timezone
-from typing import Any, Literal, Optional
+from typing import Optional
 
 from fastapi import APIRouter, BackgroundTasks, Depends, Header, HTTPException, Path, Query
 from google.protobuf.message import DecodeError
@@ -10,7 +9,7 @@ from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import (
     ExportTraceServiceResponse,
 )
 from pydantic import Field
-from sqlalchemy import delete, insert, select
+from sqlalchemy import delete, select
 from starlette.concurrency import run_in_threadpool
 from starlette.datastructures import State
 from starlette.requests import Request
@@ -23,8 +22,9 @@ from starlette.status import (
 from strawberry.relay import GlobalID
 
 from phoenix.db import models
-from phoenix.db.insertion.helpers import as_kv
-from phoenix.db.insertion.types import Precursors
+from phoenix.db.helpers import SupportedSQLDialect
+from phoenix.db.insertion.helpers import as_kv, insert_on_conflict
+from phoenix.server.api.routers.v1.annotations import TraceAnnotationData
 from phoenix.server.api.types.node import from_global_id_with_expected_type
 from phoenix.server.authorization import is_not_locked
 from phoenix.server.bearer_auth import PhoenixUser
@@ -33,7 +33,11 @@ from phoenix.trace.otel import decode_otlp_span
 from phoenix.utilities.project import get_project_name
 
 from .models import V1RoutesBaseModel
-from .utils import RequestBody, ResponseBody, add_errors_to_responses
+from .utils import (
+    RequestBody,
+    ResponseBody,
+    add_errors_to_responses,
+)
 
 router = APIRouter(tags=["traces"])
 
@@ -105,54 +109,8 @@ async def post_traces(
     )
 
 
-class TraceAnnotationResult(V1RoutesBaseModel):
-    label: Optional[str] = Field(default=None, description="The label assigned by the annotation")
-    score: Optional[float] = Field(default=None, description="The score assigned by the annotation")
-    explanation: Optional[str] = Field(
-        default=None, description="Explanation of the annotation result"
-    )
-
-
-class TraceAnnotation(V1RoutesBaseModel):
-    trace_id: str = Field(description="OpenTelemetry Trace ID (hex format w/o 0x prefix)")
-    name: str = Field(description="The name of the annotation")
-    annotator_kind: Literal["LLM", "HUMAN"] = Field(
-        description="The kind of annotator used for the annotation"
-    )
-    result: Optional[TraceAnnotationResult] = Field(
-        default=None, description="The result of the annotation"
-    )
-    metadata: Optional[dict[str, Any]] = Field(
-        default=None, description="Metadata for the annotation"
-    )
-    identifier: str = Field(
-        default="",
-        description=(
-            "The identifier of the annotation. "
-            "If provided, the annotation will be updated if it already exists."
-        ),
-    )
-
-    def as_precursor(self, *, user_id: Optional[int] = None) -> Precursors.TraceAnnotation:
-        return Precursors.TraceAnnotation(
-            datetime.now(timezone.utc),
-            self.trace_id,
-            models.TraceAnnotation(
-                name=self.name,
-                annotator_kind=self.annotator_kind,
-                score=self.result.score if self.result else None,
-                label=self.result.label if self.result else None,
-                explanation=self.result.explanation if self.result else None,
-                metadata_=self.metadata or {},
-                identifier=self.identifier,
-                source="APP",
-                user_id=user_id,
-            ),
-        )
-
-
-class AnnotateTracesRequestBody(RequestBody[list[TraceAnnotation]]):
-    data: list[TraceAnnotation] = Field(description="The trace annotations to be upserted")
+class AnnotateTracesRequestBody(RequestBody[list[TraceAnnotationData]]):
+    data: list[TraceAnnotationData] = Field(description="The trace annotations to be upserted")
 
 
 class InsertedTraceAnnotation(V1RoutesBaseModel):
@@ -193,9 +151,11 @@ async def annotate_traces(
     trace_ids = {p.trace_id for p in precursors}
     async with request.app.state.db() as session:
         existing_traces = {
-            trace.trace_id: trace.id
-            async for trace in await session.stream_scalars(
-                select(models.Trace).filter(models.Trace.trace_id.in_(trace_ids))
+            trace_id: id_
+            async for trace_id, id_ in await session.stream(
+                select(models.Trace.trace_id, models.Trace.id).filter(
+                    models.Trace.trace_id.in_(trace_ids)
+                )
             )
         }
 
@@ -206,10 +166,16 @@ async def annotate_traces(
                 status_code=HTTP_404_NOT_FOUND,
             )
         inserted_ids = []
+        dialect = SupportedSQLDialect(session.bind.dialect.name)
         for p in precursors:
             values = dict(as_kv(p.as_insertable(existing_traces[p.trace_id]).row))
             trace_annotation_id = await session.scalar(
-                insert(models.TraceAnnotation).values(**values).returning(models.TraceAnnotation.id)
+                insert_on_conflict(
+                    values,
+                    dialect=dialect,
+                    table=models.TraceAnnotation,
+                    unique_by=("name", "trace_rowid", "identifier"),
+                ).returning(models.TraceAnnotation.id)
             )
             inserted_ids.append(trace_annotation_id)
     request.state.event_queue.put(TraceAnnotationInsertEvent(tuple(inserted_ids)))

phoenix/server/api/types/Dataset.py

@@ -15,9 +15,11 @@ from phoenix.server.api.context import Context
 from phoenix.server.api.exceptions import BadRequest
 from phoenix.server.api.input_types.DatasetVersionSort import DatasetVersionSort
 from phoenix.server.api.types.DatasetExample import DatasetExample
+from phoenix.server.api.types.DatasetExperimentAnnotationSummary import (
+    DatasetExperimentAnnotationSummary,
+)
 from phoenix.server.api.types.DatasetVersion import DatasetVersion
 from phoenix.server.api.types.Experiment import Experiment, to_gql_experiment
-from phoenix.server.api.types.ExperimentAnnotationSummary import ExperimentAnnotationSummary
 from phoenix.server.api.types.node import from_global_id_with_expected_type
 from phoenix.server.api.types.pagination import (
     ConnectionArgs,
@@ -270,53 +272,13 @@ class Dataset(Node):
     @strawberry.field
     async def experiment_annotation_summaries(
         self, info: Info[Context, None]
-    ) -> list[ExperimentAnnotationSummary]:
+    ) -> list[DatasetExperimentAnnotationSummary]:
         dataset_id = self.id_attr
-        repetition_mean_scores_by_example_subquery = (
-            select(
-                models.ExperimentRunAnnotation.name.label("annotation_name"),
-                func.avg(models.ExperimentRunAnnotation.score).label("mean_repetition_score"),
-            )
-            .select_from(models.ExperimentRunAnnotation)
-            .join(
-                models.ExperimentRun,
-                models.ExperimentRunAnnotation.experiment_run_id == models.ExperimentRun.id,
-            )
-            .join(
-                models.Experiment,
-                models.ExperimentRun.experiment_id == models.Experiment.id,
-            )
-            .where(models.Experiment.dataset_id == dataset_id)
-            .group_by(
-                models.ExperimentRun.dataset_example_id,
-                models.ExperimentRunAnnotation.name,
-            )
-            .subquery()
-            .alias("repetition_mean_scores_by_example")
-        )
-        repetition_mean_scores_subquery = (
-            select(
-                repetition_mean_scores_by_example_subquery.c.annotation_name.label(
-                    "annotation_name"
-                ),
-                func.avg(repetition_mean_scores_by_example_subquery.c.mean_repetition_score).label(
-                    "mean_score"
-                ),
-            )
-            .select_from(repetition_mean_scores_by_example_subquery)
-            .group_by(
-                repetition_mean_scores_by_example_subquery.c.annotation_name,
-            )
-            .subquery()
-            .alias("repetition_mean_scores")
-        )
-        repetitions_subquery = (
+        query = (
             select(
                 models.ExperimentRunAnnotation.name.label("annotation_name"),
                 func.min(models.ExperimentRunAnnotation.score).label("min_score"),
                 func.max(models.ExperimentRunAnnotation.score).label("max_score"),
-                func.count().label("count"),
-                func.count(models.ExperimentRunAnnotation.error).label("error_count"),
             )
             .select_from(models.ExperimentRunAnnotation)
             .join(
@@ -329,36 +291,16 @@ class Dataset(Node):
             )
             .where(models.Experiment.dataset_id == dataset_id)
             .group_by(models.ExperimentRunAnnotation.name)
-            .subquery()
-        )
-        run_scores_query = (
-            select(
-                repetition_mean_scores_subquery.c.annotation_name.label("annotation_name"),
-                repetition_mean_scores_subquery.c.mean_score.label("mean_score"),
-                repetitions_subquery.c.min_score.label("min_score"),
-                repetitions_subquery.c.max_score.label("max_score"),
-                repetitions_subquery.c.count.label("count_"),
-                repetitions_subquery.c.error_count.label("error_count"),
-            )
-            .select_from(repetition_mean_scores_subquery)
-            .join(
-                repetitions_subquery,
-                repetitions_subquery.c.annotation_name
-                == repetition_mean_scores_subquery.c.annotation_name,
-            )
-            .order_by(repetition_mean_scores_subquery.c.annotation_name)
+            .order_by(models.ExperimentRunAnnotation.name)
         )
         async with info.context.db() as session:
             return [
-                ExperimentAnnotationSummary(
+                DatasetExperimentAnnotationSummary(
                     annotation_name=scores_tuple.annotation_name,
                     min_score=scores_tuple.min_score,
                     max_score=scores_tuple.max_score,
-                    mean_score=scores_tuple.mean_score,
-                    count=scores_tuple.count_,
-                    error_count=scores_tuple.error_count,
                 )
-                async for scores_tuple in await session.stream(run_scores_query)
+                async for scores_tuple in await session.stream(query)
             ]
 
     @strawberry.field

phoenix/server/api/types/DatasetExperimentAnnotationSummary.py

@@ -0,0 +1,10 @@
+from typing import Optional
+
+import strawberry
+
+
+@strawberry.type
+class DatasetExperimentAnnotationSummary:
+    annotation_name: str
+    min_score: Optional[float]
+    max_score: Optional[float]

phoenix/server/api/types/DocumentAnnotation.py

@@ -0,0 +1,92 @@
+from datetime import datetime
+from typing import TYPE_CHECKING, Annotated, Optional
+
+import strawberry
+from strawberry import Private
+from strawberry.relay import Node, NodeID
+from strawberry.scalars import JSON
+from strawberry.types import Info
+
+from phoenix.db import models
+from phoenix.server.api.context import Context
+from phoenix.server.api.interceptor import GqlValueMediator
+
+from .Annotation import Annotation
+from .AnnotationSource import AnnotationSource
+from .AnnotatorKind import AnnotatorKind
+from .User import User, to_gql_user
+
+if TYPE_CHECKING:
+    from .Span import Span
+
+
+@strawberry.type
+class DocumentAnnotation(Node, Annotation):
+    id_attr: NodeID[int]
+    user_id: Private[Optional[int]]
+    name: str = strawberry.field(
+        description="Name of the annotation, e.g. 'helpfulness' or 'relevance'."
+    )
+    annotator_kind: AnnotatorKind
+    label: Optional[str] = strawberry.field(
+        description="Value of the annotation in the form of a string, e.g. "
+        "'helpful' or 'not helpful'. Note that the label is not necessarily binary."
+    )
+    score: Optional[float] = strawberry.field(
+        description="Value of the annotation in the form of a numeric score.",
+        default=GqlValueMediator(),
+    )
+    explanation: Optional[str] = strawberry.field(
+        description="The annotator's explanation for the annotation result (i.e. "
+        "score or label, or both) given to the subject."
+    )
+    metadata: JSON
+    document_position: int
+    span_rowid: Private[int]
+    identifier: str
+    source: AnnotationSource
+    created_at: datetime = strawberry.field(
+        description="The date and time when the annotation was created."
+    )
+    updated_at: datetime = strawberry.field(
+        description="The date and time when the annotation was last updated."
+    )
+
+    @strawberry.field
+    async def span(self) -> Annotated["Span", strawberry.lazy(".Span")]:
+        from phoenix.server.api.types.Span import Span
+
+        return Span(span_rowid=self.span_rowid)
+
+    @strawberry.field
+    async def user(
+        self,
+        info: Info[Context, None],
+    ) -> Optional[User]:
+        if self.user_id is None:
+            return None
+        user = await info.context.data_loaders.users.load(self.user_id)
+        if user is None:
+            return None
+        return to_gql_user(user)
+
+
+def to_gql_document_annotation(
+    annotation: models.DocumentAnnotation,
+) -> DocumentAnnotation:
+    return DocumentAnnotation(
+        id_attr=annotation.id,
+        user_id=annotation.user_id,
+        name=annotation.name,
+        annotator_kind=AnnotatorKind(annotation.annotator_kind),
+        label=annotation.label,
+        score=annotation.score,
+        explanation=annotation.explanation,
+        metadata=annotation.metadata_,
+        span_rowid=annotation.span_rowid,
+        source=AnnotationSource(annotation.source),
+        identifier=annotation.identifier,
+        document_position=annotation.document_position,
+        created_at=annotation.created_at,
+        updated_at=annotation.updated_at,
+    )

phoenix/server/api/types/Experiment.py

@@ -109,10 +109,10 @@ class Experiment(Node):
 
     @strawberry.field
     async def average_run_latency_ms(self, info: Info[Context, None]) -> Optional[float]:
-        latency_seconds = await info.context.data_loaders.average_experiment_run_latency.load(
+        latency_ms = await info.context.data_loaders.average_experiment_run_latency.load(
             self.id_attr
         )
-        return latency_seconds * 1000 if latency_seconds is not None else None
+        return latency_ms
 
     @strawberry.field
     async def project(self, info: Info[Context, None]) -> Optional[Project]:

phoenix/server/api/types/Span.py

@@ -34,8 +34,11 @@ from phoenix.server.api.input_types.SpanAnnotationSort import (
 )
 from phoenix.server.api.types.AnnotationSummary import AnnotationSummary
 from phoenix.server.api.types.CostBreakdown import CostBreakdown
+from phoenix.server.api.types.DocumentAnnotation import (
+    DocumentAnnotation,
+    to_gql_document_annotation,
+)
 from phoenix.server.api.types.DocumentRetrievalMetrics import DocumentRetrievalMetrics
-from phoenix.server.api.types.Evaluation import DocumentEvaluation
 from phoenix.server.api.types.ExampleRevisionInterface import ExampleRevision
 from phoenix.server.api.types.GenerativeProvider import GenerativeProvider
 from phoenix.server.api.types.MimeType import MimeType
@@ -638,8 +641,11 @@ class Span(Node):
     async def document_evaluations(
         self,
         info: Info[Context, None],
-    ) -> list[DocumentEvaluation]:
-        return await info.context.data_loaders.document_evaluations.load(self.span_rowid)
+    ) -> list[DocumentAnnotation]:
+        return [
+            to_gql_document_annotation(anno)
+            for anno in await info.context.data_loaders.document_evaluations.load(self.span_rowid)
+        ]
 
     @strawberry.field(
         description="Retrieval metrics: NDCG@K, Precision@K, Reciprocal Rank, etc.",

phoenix/server/api/types/TraceAnnotation.py

@@ -1,8 +1,8 @@
-from typing import Optional
+from typing import TYPE_CHECKING, Annotated, Optional
 
 import strawberry
 from strawberry import Private
-from strawberry.relay import GlobalID, Node, NodeID
+from strawberry.relay import Node, NodeID
 from strawberry.scalars import JSON
 from strawberry.types import Info
 
@@ -13,6 +13,9 @@ from phoenix.server.api.types.AnnotatorKind import AnnotatorKind
 from .AnnotationSource import AnnotationSource
 from .User import User, to_gql_user
 
+if TYPE_CHECKING:
+    from .Trace import Trace
+
 
 @strawberry.type
 class TraceAnnotation(Node):
@@ -24,15 +27,15 @@ class TraceAnnotation(Node):
     score: Optional[float]
     explanation: Optional[str]
     metadata: JSON
-    trace_rowid: Private[Optional[int]]
+    trace_rowid: Private[int]
     identifier: str
     source: AnnotationSource
 
     @strawberry.field
-    async def trace_id(self) -> GlobalID:
+    async def trace(self) -> Annotated["Trace", strawberry.lazy(".Trace")]:
         from phoenix.server.api.types.Trace import Trace
 
-        return GlobalID(type_name=Trace.__name__, node_id=str(self.trace_rowid))
+        return Trace(trace_rowid=self.trace_rowid)
 
     @strawberry.field
     async def user(