WIP for annotations

burr/tracking/server/backend.py

@@ -15,7 +15,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 +66,67 @@ 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,
+        partition_key: Optional[str],
+        app_id: str,
+        step_sequence_id: int,
+        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 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
+        :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,12 +256,74 @@ 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.
     """
 
+    async def create_annotation(
+        self,
+        annotation: AnnotationCreate,
+        project_id: str,
+        partition_key: Optional[str],
+        app_id: str,
+        step_sequence_id: int,
+    ) -> AnnotationOut:
+        ...
+
+    async def update_annotation(
+        self,
+        annotation: AnnotationUpdate,
+        project_id: str,
+        partition_key: Optional[str],
+        app_id: str,
+        step_sequence_id: int,
+        annotation_id: int,
+    ) -> AnnotationOut:
+        ...
+
+    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]:
+        def generate_fake_annotations(
+            project_id: str, app_id: Optional[str] = None, partition_key: Optional[str] = None
+        ) -> Sequence[AnnotationOut]:
+            import random
+
+            import lorem
+
+            random.seed(42)  # Setting a fixed seed for reproducibility
+            common_tags = ["urgent", "review", "important"]
+            annotations = []
+            for i in range(1, 6):
+                annotation = AnnotationOut(
+                    id=i,
+                    project_id=project_id,
+                    app_id=app_id,
+                    partition_key=partition_key,
+                    step_sequence_id=i,
+                    attributes=[],
+                    span_id=None,
+                    tags=[random.choice(common_tags) for _ in range(random.randint(1, 3))],
+                    note=lorem.paragraph(),
+                    thumbs_up_thumbs_down=random.choice([True, False, None]),
+                )
+                annotations.append(annotation)
+
+            return annotations
+
+        # Example usage
+        annotations = generate_fake_annotations(project_id="proj_123", app_id="app_1")
+
+        # Example usage
+        annotations = generate_fake_annotations(project_id="proj_123", app_id="app_1")
+        return annotations
+
     def __init__(self, path: str = DEFAULT_PATH):
         self.path = path
 

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)
 
@@ -24,6 +29,9 @@
 
     from burr.tracking.server import schema
     from burr.tracking.server.schema import (
+        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,50 @@ async def get_application_logs(
     )
 
 
+@app.post("/api/v0/{project_id}/{app_id}/{partition_key}/annotations", response_model=AnnotationOut)
+async def create_annotation(
+    request: Request, project_id: str, app_id: str, partition_key: str, annotation: AnnotationCreate
+):
+    if partition_key == SENTINEL_PARTITION_KEY:
+        partition_key = None
+    return await backend.create_annotation(request, project_id, app_id, partition_key, annotation)
+
+
+@app.put(
+    "/api/v0/{project_id}/{app_id}/{partition_key}/annotations/{annotation_id}",
+    response_model=AnnotationOut,
+)
+async def update_annotation(
+    request: Request,
+    project_id: str,
+    app_id: str,
+    partition_key: str,
+    annotation_id: int,
+    annotation: AnnotationUpdate,
+):
+    if partition_key == SENTINEL_PARTITION_KEY:
+        partition_key = None
+    return await backend.update_annotation(
+        request, project_id, app_id, partition_key, annotation_id, annotation
+    )
+
+
+@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,
+    annotation_id: Optional[int] = None,
+):
+    # Handle the sentinel value for partition_key
+    if partition_key == SENTINEL_PARTITION_KEY:
+        partition_key = None
+
+    # Logic to retrieve the annotations
+    return await backend.get_annotations(project_id, app_id, partition_key, annotation_id)
+
+
 @app.get("/api/v0/ready")
 async def ready() -> bool:
     return True

burr/tracking/server/schema.py

@@ -182,3 +182,36 @@ class BackendSpec(pydantic.BaseModel):
     indexing: bool
     snapshotting: bool
     supports_demos: bool
+    supports_annotations: bool
+
+
+class AnnotationCreate(pydantic.BaseModel):
+    """Generic link for indexing job -- can be exposed in 'admin mode' in the UI"""
+
+    attributes: List[str]  # list of associated attributes that are part of this
+    span_id: Optional[str]
+    tags: List[str]
+    note: str
+    thumbs_up_thumbs_down: Optional[bool]
+
+
+class AnnotationUpdate(AnnotationCreate):
+    """Generic link for indexing job -- can be exposed in 'admin mode' in the UI"""
+
+    # Identification for association
+    attributes: Optional[List[str]] = None  # list of associated attributes that are part of this
+    span_id: Optional[str] = None
+    tags: Optional[List[str]] = []
+    note: Optional[str] = None
+    thumbs_up_thumbs_down: Optional[bool] = None
+
+
+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

telemetry/ui/package.json

@@ -27,6 +27,7 @@
     "react-query": "^3.39.3",
     "react-router-dom": "^6.22.1",
     "react-scripts": "5.0.1",
+    "react-select": "^5.8.1",
     "react-syntax-highlighter": "^15.5.0",
     "reactflow": "^11.10.4",
     "remark-gfm": "^4.0.0",

telemetry/ui/src/App.tsx

@@ -10,6 +10,7 @@ import { Counter } from './examples/Counter';
 import { EmailAssistantWithTelemetry } from './examples/EmailAssistant';
 import { StreamingChatbotWithTelemetry } from './examples/StreamingChatbot';
 import { AdminView } from './components/routes/AdminView';
+import { AnnotationsViewContainer } from './components/routes/app/AnnotationsView';
 
 /**
  * Basic application. We have an AppContainer -- this has a breadcrumb and a sidebar.
@@ -37,6 +38,7 @@ const App = () => {
             <Route path="/project/:projectId" element={<AppList />} />
             <Route path="/project/:projectId/:partitionKey" element={<AppList />} />
             <Route path="/project/:projectId/:partitionKey/:appId" element={<AppViewContainer />} />
+            <Route path="/annotations/:projectId/" element={<AnnotationsViewContainer />} />
             <Route path="/demos/counter" element={<Counter />} />
             <Route path="/demos/chatbot" element={<ChatbotWithTelemetry />} />
             <Route path="/demos/streaming-chatbot" element={<StreamingChatbotWithTelemetry />} />

telemetry/ui/src/api/index.ts

@@ -8,6 +8,9 @@ export { OpenAPI } from './core/OpenAPI';
 export type { OpenAPIConfig } from './core/OpenAPI';
 
 export type { ActionModel } from './models/ActionModel';
+export type { AnnotationCreate } from './models/AnnotationCreate';
+export type { AnnotationOut } from './models/AnnotationOut';
+export type { AnnotationUpdate } from './models/AnnotationUpdate';
 export type { ApplicationLogs } from './models/ApplicationLogs';
 export type { ApplicationModel } from './models/ApplicationModel';
 export type { ApplicationPage } from './models/ApplicationPage';

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 */
+/**
+ * Generic link for indexing job -- can be exposed in 'admin mode' in the UI
+ */
+export type AnnotationCreate = {
+  attributes: Array<string>;
+  span_id: string | null;
+  tags: Array<string>;
+  note: string;
+  thumbs_up_thumbs_down: boolean | null;
+};

telemetry/ui/src/api/models/AnnotationOut.ts

@@ -0,0 +1,20 @@
+/* generated using openapi-typescript-codegen -- do no edit */
+/* istanbul ignore file */
+/* tslint:disable */
+/* eslint-disable */
+/**
+ * Generic link for indexing job -- can be exposed in 'admin mode' in the UI
+ */
+export type AnnotationOut = {
+  attributes: Array<string>;
+  span_id: string | null;
+  step_name: string;
+  tags: Array<string>;
+  note: string;
+  thumbs_up_thumbs_down: boolean | null;
+  id: number;
+  project_id: string;
+  app_id: string;
+  partition_key: string | null;
+  step_sequence_id: number;
+};

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 */
+/**
+ * Generic link for indexing job -- can be exposed in 'admin mode' in the UI
+ */
+export type AnnotationUpdate = {
+  attributes?: Array<string> | null;
+  span_id?: string | null;
+  tags?: Array<string> | null;
+  note?: string | null;
+  thumbs_up_thumbs_down?: boolean | null;
+};