From f9883e30c911b126cf523cc8b73c88bf554a421f Mon Sep 17 00:00:00 2001
From: elijahbenizzy <elijah_ben_izzy@alumni.brown.edu>
Date: Sat, 10 Aug 2024 23:14:05 -0700
Subject: [PATCH] Adds OpenTelemetry integration.

Two approaches:
1. Log Burr -> Otel -- simple adapter
2. Log any Otel traces inside a burr step to the Burr UI

This is overall fairly sound, although there is a bit of global state
that might be dealt with.

Design decisions:

1. Use a custom span processor
2. Store state of a map of span_id -> step
3. Propogate that when processing spans

It's a bit odd as the span processor should just be writing/queuing, but
the way I'm thinking of this is that the writing is converting OTel
spans to Burr spans, so we have to do that somewhere. Why not when
processing? Then we can just cascade down.

Any spans outside of a Burr step are ignored, as we are not a full OTel
provider.
---
 burr/cli/demo_data.py    | 11 ++++++-----
 burr/core/application.py | 10 ++++++++--
 burr/tracking/base.py    |  2 ++
 pyproject.toml           |  4 ++++
 4 files changed, 20 insertions(+), 7 deletions(-)

diff --git a/burr/cli/demo_data.py b/burr/cli/demo_data.py
index 1d15f278..72bd9a27 100644
--- a/burr/cli/demo_data.py
+++ b/burr/cli/demo_data.py
@@ -66,16 +66,20 @@ def _modify(app_id: str) -> str:
 
     def _run_conversation(app_id, prompts):
         tracker = (
-            LocalTrackingClient(project=project_id, storage_dir=data_dir)
+            LocalTrackingClient(project=project_id + "_otel", storage_dir=data_dir)
             if not s3_bucket
             else S3TrackingClient(project=project_id, bucket=s3_bucket)
         )
+        from opentelemetry.instrumentation.openai import OpenAIInstrumentor
+
+        # TODO -- get this auto-registered
+        OpenAIInstrumentor().instrument()
         graph = (chatbot_application_with_traces if use_traces else chatbot_application).graph
         app = (
             ApplicationBuilder()
             .with_graph(graph)
             .with_identifiers(app_id=app_id)
-            .with_tracker(tracker)
+            .with_tracker(tracker, use_otel_tracing=True)
             .with_entrypoint("prompt")
             .build()
         )
@@ -201,6 +205,3 @@ def generate_all(
     generate_counter_data(data_dir=data_dir, s3_bucket=s3_bucket, unique_app_names=unique_app_names)
     logger.info("Generating RAG data")
     generate_rag_data(data_dir=data_dir, s3_bucket=s3_bucket, unique_app_names=unique_app_names)
-
-
-#
diff --git a/burr/core/application.py b/burr/core/application.py
index 5f0140fb..219f246a 100644
--- a/burr/core/application.py
+++ b/burr/core/application.py
@@ -1812,6 +1812,7 @@ def with_tracker(
         tracker: Union[Literal["local"], "TrackingClient"] = "local",
         project: str = "default",
         params: Dict[str, Any] = None,
+        use_otel_tracing: bool = False,
     ):
         """Adds a "tracker" to the application. The tracker specifies
         a project name (used for disambiguating groups of tracers), and plugs into the
@@ -1826,6 +1827,8 @@ def with_tracker(
         :param tracker: Tracker to use. ``local`` creates one, else pass one in.
         :param project: Project name -- used if the tracker is string-specified (local).
         :param params: Parameters to pass to the tracker if it's string-specified (local).
+        :param use_otel_tracing: Whether to log opentelemetry traces to the Burr UI. This is experimental but we will be adding
+            full support shortly. This requires burr[opentelemetry] installed. Note you can also log burr to OpenTelemetry using the OpenTelemetry adapter
         :return: The application builder for future chaining.
         """
         # if it's a lifecycle adapter, just add it
@@ -1839,15 +1842,18 @@ def with_tracker(
                 kwargs = {"project": project}
                 kwargs.update(params)
                 instantiated_tracker = LocalTrackingClient(**kwargs)
-                self.lifecycle_adapters.append(instantiated_tracker)
             else:
                 raise ValueError(f"Tracker {tracker}:{project} not supported")
         else:
-            self.lifecycle_adapters.append(instantiated_tracker)
             if params is not None:
                 raise ValueError(
                     "Params are not supported for object-specified trackers, these are already initialized!"
                 )
+        if use_otel_tracing:
+            from burr.integrations.opentelemetry import OpenTelemetryTracker
+
+            instantiated_tracker = OpenTelemetryTracker(burr_tracker=instantiated_tracker)
+        self.lifecycle_adapters.append(instantiated_tracker)
         self.tracker = instantiated_tracker
         return self
 
diff --git a/burr/tracking/base.py b/burr/tracking/base.py
index cc87d792..635ddc07 100644
--- a/burr/tracking/base.py
+++ b/burr/tracking/base.py
@@ -7,6 +7,7 @@
     PreRunStepHook,
     PreStartSpanHook,
 )
+from burr.lifecycle.base import DoLogAttributeHook
 
 
 class SyncTrackingClient(
@@ -15,6 +16,7 @@ class SyncTrackingClient(
     PostRunStepHook,
     PreStartSpanHook,
     PostEndSpanHook,
+    DoLogAttributeHook,
     abc.ABC,
 ):
     """Base class for synchronous tracking clients. All tracking clients must implement from this
diff --git a/pyproject.toml b/pyproject.toml
index b972b228..f7328186 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -144,6 +144,10 @@ developer = [
   "pre-commit",
 ]
 
+opentelemetry = [
+  "opentelemetry-api",
+  "opentelemetry-sdk",
+]
 [tool.setuptools]
 include-package-data = true