From b657f87106d4722be856ccb1db9cb2c7ad090a11 Mon Sep 17 00:00:00 2001 From: elijahbenizzy Date: Tue, 8 Oct 2024 14:33:07 -0700 Subject: [PATCH] Adds data models + BE integrations for annotations 1. Exposes as mixin in the BE 2. Adds GET/POST/PUT endpoints 3. Creates data models Does not work with s3 yet. --- burr/tracking/server/backend.py | 183 ++- burr/tracking/server/run.py | 76 +- burr/tracking/server/schema.py | 51 +- .../ui/src/api/models/AnnotationCreate.ts | 14 + .../src/api/models/AnnotationDataPointer.ts | 15 + .../src/api/models/AnnotationObservation.ts | 10 + telemetry/ui/src/api/models/AnnotationOut.ts | 21 + .../ui/src/api/models/AnnotationUpdate.ts | 14 + telemetry/ui/src/api/models/BackendSpec.ts | 1 + telemetry/ui/src/components/common/drawer.tsx | 58 + .../components/routes/app/AnnotationsView.tsx | 1201 +++++++++++++++++ 11 files changed, 1638 insertions(+), 6 deletions(-) create mode 100644 telemetry/ui/src/api/models/AnnotationCreate.ts create mode 100644 telemetry/ui/src/api/models/AnnotationDataPointer.ts create mode 100644 telemetry/ui/src/api/models/AnnotationObservation.ts create mode 100644 telemetry/ui/src/api/models/AnnotationOut.ts create mode 100644 telemetry/ui/src/api/models/AnnotationUpdate.ts create mode 100644 telemetry/ui/src/components/common/drawer.tsx create mode 100644 telemetry/ui/src/components/routes/app/AnnotationsView.tsx diff --git a/burr/tracking/server/backend.py b/burr/tracking/server/backend.py index 51acf829..2faa06cb 100644 --- a/burr/tracking/server/backend.py +++ b/burr/tracking/server/backend.py @@ -4,6 +4,7 @@ import json import os.path import sys +from datetime import datetime from typing import Any, Optional, Sequence, Tuple, Type, TypeVar import aiofiles @@ -15,7 +16,14 @@ from burr.tracking.common import models from burr.tracking.common.models import ChildApplicationModel from burr.tracking.server import schema -from burr.tracking.server.schema import ApplicationLogs, ApplicationSummary, Step +from burr.tracking.server.schema import ( + AnnotationCreate, + AnnotationOut, + AnnotationUpdate, + ApplicationLogs, + ApplicationSummary, + Step, +) T = TypeVar("T") @@ -59,6 +67,61 @@ async def indexing_jobs( pass +class AnnotationsBackendMixin(abc.ABC): + @abc.abstractmethod + async def create_annotation( + self, + annotation: AnnotationCreate, + project_id: str, + partition_key: Optional[str], + app_id: str, + step_sequence_id: int, + ) -> AnnotationOut: + """Createse an annotation -- annotation has annotation data, the other pointers are given in the parameters. + + :param annotation: Annotation object to create + :param partition_key: Partition key to associate with + :param project_id: Project ID to associate with + :param app_id: App ID to associate with + :param step_sequence_id: Step sequence ID to associate with + :return: + """ + + @abc.abstractmethod + async def update_annotation( + self, + annotation: AnnotationUpdate, + project_id: str, + annotation_id: int, + ) -> AnnotationOut: + """Updates an annotation -- annotation has annotation data, the other pointers are given in the parameters. + + :param annotation: Annotation object to update + :param project_id: Project ID to associate with + :param annotation_id: Annotation ID to update. We include this as we may have multiple... + :return: Updated annotation + """ + + @abc.abstractmethod + async def get_annotations( + self, + project_id: str, + partition_key: Optional[str] = None, + app_id: Optional[str] = None, + step_sequence_id: Optional[int] = None, + ) -> Sequence[AnnotationOut]: + """Returns annotations for a given project, partition_key, app_id, and step sequence ID. + If these are None it does not filter by them. + + :param project_id: Project ID to query for + :param partition_key: Partition key to query for + :param app_id: App ID to query for + :param step_sequence_id: Step sequence ID to query for + :return: Annotations + """ + pass + + class SnapshottingBackendMixin(abc.ABC): """Mixin for backend that conducts snapshotting -- e.g. saves the data to a file or database.""" @@ -188,7 +251,7 @@ def get_uri(project_id: str) -> str: DEFAULT_PATH = os.path.expanduser("~/.burr") -class LocalBackend(BackendBase): +class LocalBackend(BackendBase, AnnotationsBackendMixin): """Quick implementation of a local backend for testing purposes. This is not a production backend. To override the path, set a `burr_path` environment variable to the path you want to use. @@ -197,6 +260,122 @@ class LocalBackend(BackendBase): def __init__(self, path: str = DEFAULT_PATH): self.path = path + def _get_annotation_path(self, project_id: str) -> str: + return os.path.join(self.path, project_id, "annotations.jsonl") + + async def _load_project_annotations(self, project_id: str): + annotations_path = self._get_annotation_path(project_id) + annotations = [] + if os.path.exists(annotations_path): + async with aiofiles.open(annotations_path) as f: + for line in await f.readlines(): + annotations.append(AnnotationOut.parse_raw(line)) + return annotations + + async def create_annotation( + self, + annotation: AnnotationCreate, + project_id: str, + partition_key: Optional[str], + app_id: str, + step_sequence_id: int, + ) -> AnnotationOut: + """Creates an annotation by loading all annotations, finding the max ID, and then appending the new annotation. + This is not efficient but it's OK -- this is the local version and the number of annotations will be unlikely to be + huge. + + :param annotation: Annotation to create + :param project_id: ID of the associated project + :param partition_key: Partition key to associate with + :param app_id: App ID to associate with + :param step_sequence_id: Step sequence ID to associate with + :return: The created annotation, complete with an ID + timestamps + """ + all_annotations = await self._load_project_annotations(project_id) + annotation_id = ( + max([a.id for a in all_annotations], default=-1) + 1 + ) # get the ID, increment + annotation_out = AnnotationOut( + id=annotation_id, + project_id=project_id, + app_id=app_id, + partition_key=partition_key, + step_sequence_id=step_sequence_id, + created=datetime.now(), + updated=datetime.now(), + **annotation.dict(), + ) + annotations_path = self._get_annotation_path(project_id) + async with aiofiles.open(annotations_path, "a") as f: + await f.write(annotation_out.json() + "\n") + return annotation_out + + async def update_annotation( + self, + annotation: AnnotationUpdate, + project_id: str, + annotation_id: int, + ) -> AnnotationOut: + """Updates an annotation by loading all annotations, finding the annotation, updating it, and then writing it back. + Again, inefficient, but this is the local backend and we don't expect huge numbers of annotations. + + :param annotation: Annotation to update -- this is just the update fields to the full annotation + :param project_id: ID of the associated project + :param annotation_id: ID of the associated annotation, created by the backend + :return: The updated annotation, complete with an ID + timestamps + """ + all_annotations = await self._load_project_annotations(project_id) + annotation_out = None + for idx, a in enumerate(all_annotations): + if a.id == annotation_id: + annotation_out = a + all_annotations[idx] = annotation_out.copy( + update={**annotation.dict(), "updated": datetime.now()} + ) + break + if annotation_out is None: + raise fastapi.HTTPException( + status_code=404, + detail=f"Annotation: {annotation_id} from project: {project_id} not found", + ) + annotations_path = self._get_annotation_path(project_id) + async with aiofiles.open(annotations_path, "w") as f: + for a in all_annotations: + await f.write(a.json() + "\n") + return annotation_out + + async def get_annotations( + self, + project_id: str, + partition_key: Optional[str] = None, + app_id: Optional[str] = None, + step_sequence_id: Optional[int] = None, + ) -> Sequence[AnnotationOut]: + """Gets the annotation by loading all annotations and filtering by the parameters. Will return all annotations + that match. Only project is required. + + + :param project_id: + :param partition_key: + :param app_id: + :param step_sequence_id: + :return: + """ + annotation_path = self._get_annotation_path(project_id) + if not os.path.exists(annotation_path): + return [] + annotations = [] + async with aiofiles.open(annotation_path) as f: + for line in await f.readlines(): + parsed = AnnotationOut.parse_raw(line) + if ( + (partition_key is None or parsed.partition_key == partition_key) + and (app_id is None or parsed.app_id == app_id) + and (step_sequence_id is None or parsed.step_sequence_id == step_sequence_id) + ): + annotations.append(parsed) + return annotations + async def list_projects(self, request: fastapi.Request) -> Sequence[schema.Project]: out = [] if not os.path.exists(self.path): diff --git a/burr/tracking/server/run.py b/burr/tracking/server/run.py index f99c4602..6b5d446e 100644 --- a/burr/tracking/server/run.py +++ b/burr/tracking/server/run.py @@ -3,13 +3,18 @@ import os from contextlib import asynccontextmanager from importlib.resources import files -from typing import Sequence +from typing import Optional, Sequence from starlette import status # TODO -- remove this, just for testing from burr.log_setup import setup_logging -from burr.tracking.server.backend import BackendBase, IndexingBackendMixin, SnapshottingBackendMixin +from burr.tracking.server.backend import ( + AnnotationsBackendMixin, + BackendBase, + IndexingBackendMixin, + SnapshottingBackendMixin, +) setup_logging(logging.INFO) @@ -23,7 +28,10 @@ from starlette.templating import Jinja2Templates from burr.tracking.server import schema - from burr.tracking.server.schema import ( + from burr.tracking.server.schema import ( # AnnotationUpdate, + AnnotationCreate, + AnnotationOut, + AnnotationUpdate, ApplicationLogs, ApplicationPage, BackendSpec, @@ -130,11 +138,13 @@ def is_ready(): def get_app_spec(): is_indexing_backend = isinstance(backend, IndexingBackendMixin) is_snapshotting_backend = isinstance(backend, SnapshottingBackendMixin) + is_annotations_backend = isinstance(backend, AnnotationsBackendMixin) supports_demos = backend.supports_demos() return BackendSpec( indexing=is_indexing_backend, snapshotting=is_snapshotting_backend, supports_demos=supports_demos, + supports_annotations=is_annotations_backend, ) @@ -217,6 +227,66 @@ async def get_application_logs( ) +@app.post( + "/api/v0/{project_id}/{app_id}/{partition_key}/{sequence_id}/annotations", + response_model=AnnotationOut, +) +async def create_annotation( + request: Request, + project_id: str, + app_id: str, + partition_key: str, + sequence_id: int, + annotation: AnnotationCreate, +): + if partition_key == SENTINEL_PARTITION_KEY: + partition_key = None + spec = get_app_spec() + if not spec.supports_annotations: + return [] # empty default -- the case that we don't support annotations + return await backend.create_annotation( + annotation, project_id, partition_key, app_id, sequence_id + ) + + +# +# # TODO -- take out these parameters cause we have the annotation ID +@app.put( + "/api/v0/{project_id}/{annotation_id}/update_annotations", + response_model=AnnotationOut, +) +async def update_annotation( + request: Request, + project_id: str, + annotation_id: int, + annotation: AnnotationUpdate, +): + return await backend.update_annotation( + annotation_id=annotation_id, annotation=annotation, project_id=project_id + ) + + +@app.get("/api/v0/{project_id}/annotations", response_model=Sequence[AnnotationOut]) +async def get_annotations( + request: Request, + project_id: str, + app_id: Optional[str] = None, + partition_key: Optional[str] = None, + step_sequence_id: Optional[int] = None, +): + # Handle the sentinel value for partition_key + if partition_key == SENTINEL_PARTITION_KEY: + partition_key = None + backend_spec = get_app_spec() + + if not backend_spec.supports_annotations: + # makes it easier to wire through to the FE + return [] + + # Logic to retrieve the annotations + return await backend.get_annotations(project_id, partition_key, app_id, step_sequence_id) + + @app.get("/api/v0/ready") async def ready() -> bool: return True diff --git a/burr/tracking/server/schema.py b/burr/tracking/server/schema.py index 7bc1d30f..d26eace8 100644 --- a/burr/tracking/server/schema.py +++ b/burr/tracking/server/schema.py @@ -1,6 +1,6 @@ import collections import datetime -from typing import Any, Dict, List, Optional, Union +from typing import Any, Dict, List, Literal, Optional, Union import pydantic from pydantic import fields @@ -182,3 +182,52 @@ class BackendSpec(pydantic.BaseModel): indexing: bool snapshotting: bool supports_demos: bool + supports_annotations: bool + + +class AnnotationDataPointer(pydantic.BaseModel): + type: Literal["state_field", "attribute"] + field_name: str # key of attribute/state field + span_id: Optional[ + str + ] # span_id if it's associated with a span, otherwise it's associated with an action + + +AllowedDataField = Literal["note", "ground_truth"] + + +class AnnotationObservation(pydantic.BaseModel): + data_fields: dict[str, Any] + thumbs_up_thumbs_down: Optional[bool] + data_pointers: List[AnnotationDataPointer] + + +class AnnotationCreate(pydantic.BaseModel): + """Generic link for indexing job -- can be exposed in 'admin mode' in the UI""" + + span_id: Optional[str] + step_name: str # Should be able to look it up but including for now + tags: List[str] + observations: List[AnnotationObservation] + + +class AnnotationUpdate(AnnotationCreate): + """Generic link for indexing job -- can be exposed in 'admin mode' in the UI""" + + # Identification for association + span_id: Optional[str] = None + tags: Optional[List[str]] = [] + observations: List[AnnotationObservation] + + +class AnnotationOut(AnnotationCreate): + """Generic link for indexing job -- can be exposed in 'admin mode' in the UI""" + + id: int + # Identification for association + project_id: str # associated project ID + app_id: str + partition_key: Optional[str] + step_sequence_id: int + created: datetime.datetime + updated: datetime.datetime diff --git a/telemetry/ui/src/api/models/AnnotationCreate.ts b/telemetry/ui/src/api/models/AnnotationCreate.ts new file mode 100644 index 00000000..5c989a00 --- /dev/null +++ b/telemetry/ui/src/api/models/AnnotationCreate.ts @@ -0,0 +1,14 @@ +/* generated using openapi-typescript-codegen -- do no edit */ +/* istanbul ignore file */ +/* tslint:disable */ +/* eslint-disable */ +import type { AnnotationObservation } from './AnnotationObservation'; +/** + * Generic link for indexing job -- can be exposed in 'admin mode' in the UI + */ +export type AnnotationCreate = { + span_id: string | null; + step_name: string; + tags: Array; + observations: Array; +}; diff --git a/telemetry/ui/src/api/models/AnnotationDataPointer.ts b/telemetry/ui/src/api/models/AnnotationDataPointer.ts new file mode 100644 index 00000000..10ecfa5e --- /dev/null +++ b/telemetry/ui/src/api/models/AnnotationDataPointer.ts @@ -0,0 +1,15 @@ +/* generated using openapi-typescript-codegen -- do no edit */ +/* istanbul ignore file */ +/* tslint:disable */ +/* eslint-disable */ +export type AnnotationDataPointer = { + type: AnnotationDataPointer.type; + field_name: string; + span_id: string | null; +}; +export namespace AnnotationDataPointer { + export enum type { + STATE_FIELD = 'state_field', + ATTRIBUTE = 'attribute' + } +} diff --git a/telemetry/ui/src/api/models/AnnotationObservation.ts b/telemetry/ui/src/api/models/AnnotationObservation.ts new file mode 100644 index 00000000..5412e465 --- /dev/null +++ b/telemetry/ui/src/api/models/AnnotationObservation.ts @@ -0,0 +1,10 @@ +/* generated using openapi-typescript-codegen -- do no edit */ +/* istanbul ignore file */ +/* tslint:disable */ +/* eslint-disable */ +import type { AnnotationDataPointer } from './AnnotationDataPointer'; +export type AnnotationObservation = { + data_fields: Record; + thumbs_up_thumbs_down: boolean | null; + data_pointers: Array; +}; diff --git a/telemetry/ui/src/api/models/AnnotationOut.ts b/telemetry/ui/src/api/models/AnnotationOut.ts new file mode 100644 index 00000000..bab7cdf2 --- /dev/null +++ b/telemetry/ui/src/api/models/AnnotationOut.ts @@ -0,0 +1,21 @@ +/* generated using openapi-typescript-codegen -- do no edit */ +/* istanbul ignore file */ +/* tslint:disable */ +/* eslint-disable */ +import type { AnnotationObservation } from './AnnotationObservation'; +/** + * Generic link for indexing job -- can be exposed in 'admin mode' in the UI + */ +export type AnnotationOut = { + span_id: string | null; + step_name: string; + tags: Array; + observations: Array; + id: number; + project_id: string; + app_id: string; + partition_key: string | null; + step_sequence_id: number; + created: string; + updated: string; +}; diff --git a/telemetry/ui/src/api/models/AnnotationUpdate.ts b/telemetry/ui/src/api/models/AnnotationUpdate.ts new file mode 100644 index 00000000..da6981fa --- /dev/null +++ b/telemetry/ui/src/api/models/AnnotationUpdate.ts @@ -0,0 +1,14 @@ +/* generated using openapi-typescript-codegen -- do no edit */ +/* istanbul ignore file */ +/* tslint:disable */ +/* eslint-disable */ +import type { AnnotationObservation } from './AnnotationObservation'; +/** + * Generic link for indexing job -- can be exposed in 'admin mode' in the UI + */ +export type AnnotationUpdate = { + span_id?: string | null; + step_name: string; + tags?: Array | null; + observations: Array; +}; diff --git a/telemetry/ui/src/api/models/BackendSpec.ts b/telemetry/ui/src/api/models/BackendSpec.ts index f55f6384..d3124684 100644 --- a/telemetry/ui/src/api/models/BackendSpec.ts +++ b/telemetry/ui/src/api/models/BackendSpec.ts @@ -9,4 +9,5 @@ export type BackendSpec = { indexing: boolean; snapshotting: boolean; supports_demos: boolean; + supports_annotations: boolean; }; diff --git a/telemetry/ui/src/components/common/drawer.tsx b/telemetry/ui/src/components/common/drawer.tsx new file mode 100644 index 00000000..32314bad --- /dev/null +++ b/telemetry/ui/src/components/common/drawer.tsx @@ -0,0 +1,58 @@ +import { + Dialog, + DialogBackdrop, + DialogPanel, + DialogTitle, + TransitionChild +} from '@headlessui/react'; +import { XMarkIcon } from '@heroicons/react/24/outline'; + +export const Drawer = (props: { + open: boolean; + close: () => void; + children: React.ReactNode; + title: string; +}) => { + const { open, close } = props; + return ( + + + +
+
+
+ + +
+ +
+ +
+
+ + {props.title} + +
+
{props.children}
+
+ +
+
+
+ + ); +}; diff --git a/telemetry/ui/src/components/routes/app/AnnotationsView.tsx b/telemetry/ui/src/components/routes/app/AnnotationsView.tsx new file mode 100644 index 00000000..5a7848ff --- /dev/null +++ b/telemetry/ui/src/components/routes/app/AnnotationsView.tsx @@ -0,0 +1,1201 @@ +import { useContext, useEffect, useMemo, useState } from 'react'; +import { + AnnotationCreate, + AnnotationDataPointer, + AnnotationObservation, + AnnotationOut, + AnnotationUpdate, + DefaultService, + Step +} from '../../../api'; +import { AppView, AppContext } from './AppView'; + +import Select, { components } from 'react-select'; +import CreatableSelect from 'react-select/creatable'; +import { FaClipboardList, FaExternalLinkAlt, FaThumbsDown, FaThumbsUp } from 'react-icons/fa'; +import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from '../../common/table'; +import { Chip } from '../../common/chip'; +import { Link, useParams } from 'react-router-dom'; +import { useMutation, useQuery } from 'react-query'; +import { Loading } from '../../common/loading'; +import { + ChevronDownIcon, + PencilIcon, + PencilSquareIcon, + PlusIcon, + XMarkIcon +} from '@heroicons/react/24/outline'; +import { classNames } from '../../../utils/tailwind'; +import { DateTimeDisplay } from '../../common/dates'; +import { Drawer } from '../../common/drawer'; + +export const InlineAppView = (props: { + projectId: string; + partitionKey: string | null; + appId: string; + sequenceID: number; +}) => { + return ( +
+ +
+ ); +}; + +const getAnnotationsTarget = (steps: Step[]) => { + return steps.map((step) => ({ + sequenceId: step.step_start_log.sequence_id, + spanId: null, // TODO -- allow annotations at the attribute/trace level! + actionName: step.step_start_log.action + })); +}; + +const getPossibleDataTargets = (step: Step) => { + return [ + ...step.attributes.map((attribute) => ({ + type: AnnotationDataPointer.type.ATTRIBUTE, + field_name: attribute.key, + span_id: attribute.span_id, + value: JSON.stringify(attribute.value) + })), + ...Object.keys(step.step_end_log?.state || {}) + .map((key) => ({ + type: AnnotationDataPointer.type.STATE_FIELD, + field_name: key, + span_id: null, + value: JSON.stringify(step.step_end_log?.state[key]) + })) + .filter((data) => !data.field_name.startsWith('__')) + ]; +}; + +export const AnnotationsView = (props: { + currentStep: Step | undefined; + appId: string; + partitionKey: string | null; + projectId: string; + allAnnotations: AnnotationOut[]; + allSteps: Step[]; // List of steps so we can select the possible targets to annotate +}) => { + const { + currentEditingAnnotationContext, + setCurrentEditingAnnotationContext, + setCurrentHoverIndex, + setCurrentSelectedIndex, + currentSelectedIndex, + createAnnotation, + updateAnnotation, + refreshAnnotationData + } = useContext(AppContext); + const tags = Array.from(new Set(props.allAnnotations.flatMap((annotation) => annotation.tags))); + const annotationTargets = getAnnotationsTarget(props.allSteps); + const selectedAnnotationTarget = + currentEditingAnnotationContext === undefined + ? undefined + : { + sequenceId: currentEditingAnnotationContext.sequenceId, + spanId: null, + actionName: + props.allSteps.find( + (step) => + step.step_start_log.sequence_id === currentEditingAnnotationContext.sequenceId + )?.step_start_log.action || '' + }; + const existingAnnotation = props.allAnnotations.find( + (annotation) => + annotation.step_sequence_id === currentEditingAnnotationContext?.sequenceId && + annotation.app_id === props.appId && + annotation.span_id === currentEditingAnnotationContext.spanId + ); + + const step = existingAnnotation + ? props.allSteps.find( + (step) => step.step_start_log.sequence_id === existingAnnotation?.step_sequence_id + ) + : props.allSteps.find( + (step) => step.step_start_log.sequence_id === currentEditingAnnotationContext?.sequenceId + ); + + const allPossibleDataTargets: AnnnotationDataPointerWithValue[] = step + ? getPossibleDataTargets(step) + : []; + + return ( +
+ {currentEditingAnnotationContext && ( + { + setCurrentEditingAnnotationContext(undefined); + }} + createAnnotation={(annotation) => { + createAnnotation( + props.projectId, + props.partitionKey, + props.appId, + currentEditingAnnotationContext.sequenceId, + currentEditingAnnotationContext.spanId || undefined, + annotation + ).then(() => { + refreshAnnotationData(); + setCurrentEditingAnnotationContext(undefined); + }); + }} + allPossibleDataTargets={allPossibleDataTargets} + updateAnnotation={(annotationID, annotationUpdate) => { + updateAnnotation(annotationID, annotationUpdate).then(() => { + refreshAnnotationData(); + setCurrentEditingAnnotationContext(undefined); + }); + }} + /> + )} + { + // TODO -- ensure that the indices are aligned/set correctly + setCurrentSelectedIndex(annotation.step_sequence_id); + }} + onHover={(annotation) => { + setCurrentHoverIndex(annotation.step_sequence_id); + }} + displayProjectLevelAnnotationsLink={true} // we want to link back to the project level view + projectId={props.projectId} + highlightedSequence={currentSelectedIndex} + /> +
+ ); +}; + +type DownloadButtonProps = { + data: AnnotationOut[]; + fileName: string; +}; +const annotationsToCSV = (annotations: AnnotationOut[]): string => { + // Define the CSV headers + const headers = [ + 'id', + 'project_id', + 'app_id', + 'span_id', + 'step_sequence_id', + 'step_name', + 'tags', + 'observation_number', + 'note', + 'ground_truth', + 'thumbs_up_thumbs_down', + 'data_field_type', + 'data_field_name', + 'partition_key', + 'created', + 'updated' + ]; + + // Helper function to escape fields for CSV format + const escapeCSV = (value: string) => { + if (value === undefined) { + return ''; + } + if (value.includes(',') || value.includes('"') || value.includes('\n')) { + return `"${value.replace(/"/g, '""')}"`; + } + return value; + }; + + // Convert the annotations to CSV format + const rows = annotations.flatMap((annotation) => { + return annotation.observations.map((observation, i) => { + return { + id: annotation.id.toString(), + project_id: annotation.project_id, + app_id: annotation.app_id, + span_id: annotation.span_id || '', + step_sequence_id: annotation.step_sequence_id.toString(), + step_name: annotation.step_name, + tags: annotation.tags.join(' '), + observation_number: i.toString(), + note: observation.data_fields['note'] || '', + ground_truth: observation.data_fields['ground_truth'] || '', + thumbs_up_thumbs_down: + observation.thumbs_up_thumbs_down !== null + ? observation.thumbs_up_thumbs_down.toString() + : '', + data_field_type: observation.data_pointers.map((dp) => dp.type).join(' '), + data_field_name: observation.data_pointers.map((dp) => dp.field_name).join(' '), + partition_key: annotation.partition_key || '', + created: new Date(annotation.created).toISOString(), + updated: new Date(annotation.updated).toISOString() + }; + }); + }); + + // Construct the CSV string -- TODO -- use a library for this + const csvContent = [ + headers.join(','), // Add headers + ...rows.map((row) => + headers.map((header) => escapeCSV(row[header as keyof typeof row])).join(',') + ) // Add data rows + ].join('\n'); + + return csvContent; +}; + +const DownloadAnnotationsButton: React.FC = ({ data, fileName }) => { + const handleDownload = () => { + const csvData = annotationsToCSV(data); + const blob = new Blob([csvData], { type: 'text/csv;charset=utf-8;' }); + const url = URL.createObjectURL(blob); + const link = document.createElement('a'); + link.href = url; + link.setAttribute('download', fileName); + document.body.appendChild(link); + link.click(); + document.body.removeChild(link); + URL.revokeObjectURL(url); + }; + + return ( + + ); +}; + +type TagOptionType = { + value: string; + label: string; +}; + +type AnnotationTarget = { + sequenceId: number; + actionName: string; + spanId: string | null; +}; + +type TargetOptionType = { + value: AnnotationTarget; + label: string; +}; + +const getLabelForTarget = (target: AnnotationTarget) => + `${target.actionName}:${target.sequenceId}` + (target.spanId ? `-${target.spanId}` : ''); + +const ObservationsView = (props: { observations: AnnotationObservation[] }) => { + // Observation -- this will be a list of notes with data attached + // Data -- a chip with the type of the data (state, attribute) + key, with link + // Thumbs up/down -- a thumbs up or thumbs down icon + // Notes -- free-form text + // This is going to expand to show the data and the notes, otherwise it'll just be truncated + const [isExpanded, setIsExpanded] = useState(false); + const observationsToShow = isExpanded ? props.observations : props.observations.slice(0, 1); + return ( +
{ + setIsExpanded((expanded) => !expanded); + e.preventDefault(); + }} + > + {observationsToShow.map((observation, i) => { + const Icon = observation.thumbs_up_thumbs_down ? FaThumbsUp : FaThumbsDown; + const iconColor = observation.thumbs_up_thumbs_down ? 'text-green-500' : 'text-dwred'; + return ( +
+
+ {observation.thumbs_up_thumbs_down !== undefined && ( +
+ +
+ )} +
+ {observation.data_pointers.map((dataPointer, i) => ( +
+ + 
.{dataPointer.field_name}
+
+ ))} +
+
+
+ {observation.data_fields['note'] && ( +
+ {' '} + Note + {observation.data_fields['note']} +
+ )} + {observation.data_fields['ground_truth'] && ( +
+ {' '} + Ground Truth: + {observation.data_fields['ground_truth']} +
+ )} + {!isExpanded && props.observations.length > 1 && ( + + +{`${props.observations.length - 1} more`} + + )} +
+
+ ); + })} +
+ ); +}; + +type Filters = { + tags?: string[]; + actionNames?: string[]; +}; + +type SearchBarProps = { + filters: Filters; + setFilters: (filters: Filters) => void; + data: AnnotationOut[]; +}; + +const SearchBar: React.FC = ({ filters, setFilters, data }) => { + // Options for react-select derived from the data + const options = useMemo(() => { + // Use Sets to ensure uniqueness + const tagSet = new Set(); + const actionNameSet = new Set(); + + // Populate the sets with unique tags and action names + data.forEach((annotation) => { + annotation.tags.forEach((tag) => tagSet.add(tag)); + actionNameSet.add(annotation.step_name); + }); + + // Convert sets to the format required for react-select + const tagOptions = Array.from(tagSet).map((tag) => ({ value: tag, label: tag, type: 'tag' })); + const actionNameOptions = Array.from(actionNameSet).map((name) => ({ + value: name, + label: name, + type: 'actionName' + })); + + return [...tagOptions, ...actionNameOptions]; + }, [data]); + + // Handle selection from react-select + // TODO -- remove anys here + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const handleChange = (selectedOptions: any) => { + const newFilters: Filters = { + tags: selectedOptions + + // eslint-disable-next-line @typescript-eslint/no-explicit-any + .filter((option: any) => option.type === 'tag') + + // eslint-disable-next-line @typescript-eslint/no-explicit-any + .map((option: any) => option.value), + actionNames: selectedOptions + + // eslint-disable-next-line @typescript-eslint/no-explicit-any + .filter((option: any) => option.type === 'actionName') + + // eslint-disable-next-line @typescript-eslint/no-explicit-any + .map((option: any) => option.value) + }; + setFilters(newFilters); + }; + const selectedOptions = options.filter((option) => { + return ( + (filters.tags && filters.tags.includes(option.value)) || + (filters.actionNames && filters.actionNames.includes(option.value)) + ); + }); + + const OptionType = (props: { type: 'actionName' | 'tag'; value: string }) => { + const { type, value } = props; + return ( +
+ + {value} +
+ ); + }; + return ( + + option.value.field_name === observation.data_pointers[0]?.field_name && + option.value.type === observation.data_pointers[0]?.type + ) || null + } + onChange={(selectedOption) => { + if (selectedOption === null) return; + observation.data_pointers = [selectedOption.value]; + props.setObservation(observation); + }} + placeholder="Select data (fields in state/attributes) associated with your observation" + className="w-full" + // TODO: Create and use a custom component for the menu items + components={{ + Option: (props) => { + // TODO: Customize the menu item component here + return ( + + {' '} +
+ + {props.data.value.field_name} + 
{props.data.value.value}
+
+ + ); + }, + SingleValue: (props) => { + // TODO: Customize the selected item component here + return ( + +
+ + {props.data.value.field_name} + {/* 
{props.data.value.value}
*/} +
+ + ); + } + }} + /> + + + + {props.allowDelete && ( + + )} + + {observation.data_pointers.length > 0 && ( + 
+          {
+            possibleDataTargetValues.find(
+              (option) =>
+                option.value.field_name === observation.data_pointers[0]?.field_name &&
+                option.value.type === observation.data_pointers[0]?.type
+            )?.value.value
+          }
+
+ )} +
+
Notes
+ +
+ )} + + ); +}; +type AnnnotationDataPointerWithValue = AnnotationDataPointer & { value: string }; + +export const AnnotateButton = (props: { + sequenceID: number; + spanID?: string; + attribute?: string; // TODO -- consider whether we want to remove, we generally annotate at the step level + // But we might want to prepopulate the attribute if we are annotating a specific attribute (in the observations field) + existingAnnotation: AnnotationOut | undefined; + setCurrentEditingAnnotationContext: (context: { + sequenceId: number; + attributeName: string | undefined; + spanId: string | null; + existingAnnotation: AnnotationOut | undefined; + }) => void; + // setTab: (tab: string) => void; // used if we want to change tab for view +}) => { + const Icon = props.existingAnnotation ? PencilSquareIcon : PlusIcon; + return ( + { + props.setCurrentEditingAnnotationContext({ + sequenceId: props.sequenceID, + attributeName: props.attribute, + spanId: props.spanID || null, + existingAnnotation: props.existingAnnotation + }); + e.stopPropagation(); + e.preventDefault(); + }} + /> + ); +}; + +const DEFAULT_TAG_OPTIONS = [ + 'to-review', + 'hallucination', + 'incomplete', + 'incorrect', + 'correct', + 'ambiguous', + 'user-error', + 'intentional-user-error' +]; +const AnnotationEditCreateForm = (props: { + tagOptions: string[]; + allAnnotationTargets: AnnotationTarget[]; + selectedAnnotationTarget: AnnotationTarget | undefined; + existingAnnotation: AnnotationOut | undefined; // Only there if we are editing an existing annotation + resetAnnotationContext: () => void; + createAnnotation: (annotation: AnnotationCreate) => void; + updateAnnotation: (annotationID: number, annotation: AnnotationUpdate) => void; + allPossibleDataTargets: AnnnotationDataPointerWithValue[]; +}) => { + const [targetValue, setTargetValue] = useState(null); + + const [tags, setTags] = useState([]); + + // const [attribute, setAttribute] = useState(''); + const [observations, setObservations] = useState([ + getEmptyObservation() + ]); + + // Define options for the select components + + const tagOptions: TagOptionType[] = [...DEFAULT_TAG_OPTIONS, ...props.tagOptions].map((tag) => ({ + value: tag, + label: tag + })); + + const allTargets = props.allAnnotationTargets.map((target) => ({ + value: target, + label: getLabelForTarget(target) + })); + + useEffect(() => { + // Reset to the selected annotation if it exists + if (props.existingAnnotation) { + setTargetValue({ + value: { + sequenceId: props.existingAnnotation.step_sequence_id, + actionName: props.existingAnnotation.step_name, + spanId: props.existingAnnotation.span_id + }, + label: getLabelForTarget({ + sequenceId: props.existingAnnotation.step_sequence_id, + actionName: props.existingAnnotation.step_name, + spanId: props.existingAnnotation.span_id + }) + }); + setObservations(props.existingAnnotation.observations); + setTags( + props.existingAnnotation.tags.map((tag) => ({ + value: tag, + label: tag + })) + ); + // Otherwise, create a new one + // } else if (props.selectedAnnotationTarget) { + // setTargetValue({ + // value: props.selectedAnnotationTarget, + // label: getLabelForTarget(props.selectedAnnotationTarget) + // }); + } + }, [props.existingAnnotation]); + + useEffect(() => { + if (props.selectedAnnotationTarget && !props.existingAnnotation) { + setTargetValue({ + value: props.selectedAnnotationTarget, + label: getLabelForTarget(props.selectedAnnotationTarget) + }); + setTags([]); + setObservations([getEmptyObservation()]); + } + }, [props.selectedAnnotationTarget?.sequenceId]); + const TagChipOption = (props: { label: string; chipType: string }) => { + return ( +
+ +
+ ); + }; + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const handleTagChange = (selectedOptions: any) => { + setTags(selectedOptions ? selectedOptions.map((option: TagOptionType) => option) : []); + }; + const mode = props.existingAnnotation !== undefined ? 'edit' : 'create'; + + return ( +
+