From c54341960096d33106e51f4f31b62d0e5511d146 Mon Sep 17 00:00:00 2001 From: Dustin Ngo Date: Fri, 9 Aug 2024 12:00:52 -0400 Subject: [PATCH 1/5] Propagate span annotation metadata to examples on all mutations --- .../server/api/mutations/dataset_mutations.py | 48 +++++++++++++++++-- 1 file changed, 43 insertions(+), 5 deletions(-) diff --git a/src/phoenix/server/api/mutations/dataset_mutations.py b/src/phoenix/server/api/mutations/dataset_mutations.py index f1d5e817d7..2981e1189a 100644 --- a/src/phoenix/server/api/mutations/dataset_mutations.py +++ b/src/phoenix/server/api/mutations/dataset_mutations.py @@ -230,7 +230,7 @@ async def add_examples_to_dataset( ) -> DatasetMutationPayload: dataset_id = input.dataset_id # Extract the span rowids from the input examples if they exist - span_ids = span_ids = [example.span_id for example in input.examples if example.span_id] + span_ids = [example.span_id for example in input.examples if example.span_id] span_rowids = { from_global_id_with_expected_type(global_id=span_id, expected_type_name=Span.__name__) for span_id in set(span_ids) @@ -260,6 +260,8 @@ async def add_examples_to_dataset( ) .returning(models.DatasetVersion.id) ) + + # Fetch spans and span annotations spans = ( await session.execute( select(models.Span.id) @@ -267,9 +269,36 @@ async def add_examples_to_dataset( .where(models.Span.id.in_(span_rowids)) ) ).all() - # Just validate that the number of spans matches the number of span_ids - # to ensure that the span_ids are valid - assert len(spans) == len(span_rowids) + + span_annotations = ( + await session.execute( + select( + models.SpanAnnotation.span_rowid, + models.SpanAnnotation.name, + models.SpanAnnotation.label, + models.SpanAnnotation.score, + models.SpanAnnotation.explanation, + models.SpanAnnotation.metadata_, + models.SpanAnnotation.annotator_kind, + ) + .select_from(models.SpanAnnotation) + .where(models.SpanAnnotation.span_rowid.in_(span_rowids)) + ) + ).all() + + span_annotations_by_span: Dict[int, Dict[Any, Any]] = {span.id: {} for span in spans} + for annotation in span_annotations: + span_id = annotation.span_rowid + if span_id not in span_annotations_by_span: + span_annotations_by_span[span_id] = dict() + span_annotations_by_span[span_id][annotation.name] = { + "label": annotation.label, + "score": annotation.score, + "explanation": annotation.explanation, + "metadata": annotation.metadata_, + "annotator_kind": annotation.annotator_kind, + } + DatasetExample = models.DatasetExample dataset_example_rowids = ( await session.scalars( @@ -299,7 +328,16 @@ async def add_examples_to_dataset( DatasetExampleRevision.dataset_version_id.key: dataset_version_rowid, DatasetExampleRevision.input.key: example.input, DatasetExampleRevision.output.key: example.output, - DatasetExampleRevision.metadata_.key: example.metadata, + DatasetExampleRevision.metadata_.key: { + **(example.metadata or {}), + "annotations": span_annotations_by_span.get( + from_global_id_with_expected_type( + global_id=example.span_id, + expected_type_name=Span.__name__, + ), + {} + ), + }, DatasetExampleRevision.revision_kind.key: "CREATE", } for dataset_example_rowid, example in zip( From f747c50910b3248f4ea68535e045fd898fa72aa1 Mon Sep 17 00:00:00 2001 From: Dustin Ngo Date: Fri, 9 Aug 2024 13:53:16 -0400 Subject: [PATCH 2/5] Refactor logic for better type-checker clarity --- .../server/api/mutations/dataset_mutations.py | 30 ++++++++++--------- 1 file changed, 16 insertions(+), 14 deletions(-) diff --git a/src/phoenix/server/api/mutations/dataset_mutations.py b/src/phoenix/server/api/mutations/dataset_mutations.py index 2981e1189a..2226945731 100644 --- a/src/phoenix/server/api/mutations/dataset_mutations.py +++ b/src/phoenix/server/api/mutations/dataset_mutations.py @@ -320,9 +320,17 @@ async def add_examples_to_dataset( assert len(dataset_example_rowids) == len(input.examples) assert all(map(lambda id: isinstance(id, int), dataset_example_rowids)) DatasetExampleRevision = models.DatasetExampleRevision - await session.execute( - insert(DatasetExampleRevision), - [ + + dataset_example_revisions = [] + for dataset_example_rowid, example in zip(dataset_example_rowids, input.examples): + span_annotation = {} + if example.span_id: + span_id = from_global_id_with_expected_type( + global_id=example.span_id, + expected_type_name=Span.__name__, + ) + span_annotation = span_annotations_by_span.get(span_id, {}) + dataset_example_revisions.append( { DatasetExampleRevision.dataset_example_id.key: dataset_example_rowid, DatasetExampleRevision.dataset_version_id.key: dataset_version_rowid, @@ -330,20 +338,14 @@ async def add_examples_to_dataset( DatasetExampleRevision.output.key: example.output, DatasetExampleRevision.metadata_.key: { **(example.metadata or {}), - "annotations": span_annotations_by_span.get( - from_global_id_with_expected_type( - global_id=example.span_id, - expected_type_name=Span.__name__, - ), - {} - ), + "annotations": span_annotation, }, DatasetExampleRevision.revision_kind.key: "CREATE", } - for dataset_example_rowid, example in zip( - dataset_example_rowids, input.examples - ) - ], + ) + await session.execute( + insert(DatasetExampleRevision), + dataset_example_revisions, ) info.context.event_queue.put(DatasetInsertEvent((dataset.id,))) return DatasetMutationPayload(dataset=to_gql_dataset(dataset)) From 4fd7ab7907d4a4d5732b6b16a4fc1c0bcb691810 Mon Sep 17 00:00:00 2001 From: Dustin Ngo Date: Mon, 12 Aug 2024 17:42:50 -0400 Subject: [PATCH 3/5] Add annotation info to preview --- src/phoenix/server/api/types/Span.py | 22 ++++++++++++++++++++-- 1 file changed, 20 insertions(+), 2 deletions(-) diff --git a/src/phoenix/server/api/types/Span.py b/src/phoenix/server/api/types/Span.py index 4e59730561..559af0eb5a 100644 --- a/src/phoenix/server/api/types/Span.py +++ b/src/phoenix/server/api/types/Span.py @@ -240,7 +240,7 @@ async def descendants( @strawberry.field( description="The span's attributes translated into an example revision for a dataset", ) # type: ignore - def as_example_revision(self) -> SpanAsExampleRevision: + def as_example_revision(self, info: Info[Context, None]) -> SpanAsExampleRevision: db_span = self.db_span attributes = db_span.attributes span_io = _SpanIO( @@ -256,10 +256,28 @@ def as_example_revision(self) -> SpanAsExampleRevision: llm_output_messages=get_attribute_value(attributes, LLM_OUTPUT_MESSAGES), retrieval_documents=get_attribute_value(attributes, RETRIEVAL_DOCUMENTS), ) + + # Fetch annotations associated with this span + span_annotations = self.span_annotations(info) + annotations = dict() + for annotation in span_annotations: + annotations[annotation.name] = { + "label": annotation.label, + "score": annotation.score, + "explanation": annotation.explanation, + "metadata": annotation.metadata, + "annotator_kind": annotation.annotator_kind, + } + # Merge annotations into the metadata + metadata = { + **attributes, + "annotations": annotations, + } + return SpanAsExampleRevision( input=get_dataset_example_input(span_io), output=get_dataset_example_output(span_io), - metadata=attributes, + metadata=metadata, ) @strawberry.field(description="The project that this span belongs to.") # type: ignore From 4a2595efd747391cf1c118e35ac1489973e4ed7e Mon Sep 17 00:00:00 2001 From: Dustin Ngo Date: Mon, 12 Aug 2024 18:29:55 -0400 Subject: [PATCH 4/5] Await annotation resolver --- src/phoenix/server/api/types/Span.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/phoenix/server/api/types/Span.py b/src/phoenix/server/api/types/Span.py index 559af0eb5a..33c38ed0d2 100644 --- a/src/phoenix/server/api/types/Span.py +++ b/src/phoenix/server/api/types/Span.py @@ -240,7 +240,7 @@ async def descendants( @strawberry.field( description="The span's attributes translated into an example revision for a dataset", ) # type: ignore - def as_example_revision(self, info: Info[Context, None]) -> SpanAsExampleRevision: + async def as_example_revision(self, info: Info[Context, None]) -> SpanAsExampleRevision: db_span = self.db_span attributes = db_span.attributes span_io = _SpanIO( @@ -258,7 +258,7 @@ def as_example_revision(self, info: Info[Context, None]) -> SpanAsExampleRevisio ) # Fetch annotations associated with this span - span_annotations = self.span_annotations(info) + span_annotations = await self.span_annotations(info) annotations = dict() for annotation in span_annotations: annotations[annotation.name] = { From 8d67a9e8288a4c796d56ad3610774f51a9c69e61 Mon Sep 17 00:00:00 2001 From: Dustin Ngo Date: Mon, 12 Aug 2024 18:35:33 -0400 Subject: [PATCH 5/5] Use AnnotatorKind value --- src/phoenix/server/api/types/Span.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/phoenix/server/api/types/Span.py b/src/phoenix/server/api/types/Span.py index 33c38ed0d2..13afad5431 100644 --- a/src/phoenix/server/api/types/Span.py +++ b/src/phoenix/server/api/types/Span.py @@ -266,7 +266,7 @@ async def as_example_revision(self, info: Info[Context, None]) -> SpanAsExampleR "score": annotation.score, "explanation": annotation.explanation, "metadata": annotation.metadata, - "annotator_kind": annotation.annotator_kind, + "annotator_kind": annotation.annotator_kind.value, } # Merge annotations into the metadata metadata = {