fix(openai): enrich assistant data if not available

traceloop · Mar 27, 2024 · 2283282 · 2283282
1 parent 7151b27
commit 2283282
Show file tree

Hide file tree

Showing 11 changed files with 2,009 additions and 1,955 deletions.
diff --git a/...ges/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/__init__.py b/...ges/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/__init__.py
@@ -2,6 +2,7 @@
 
 from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
 
+from opentelemetry.instrumentation.openai.shared.config import Config
 from opentelemetry.instrumentation.openai.utils import is_openai_v1
 from opentelemetry.instrumentation.openai.v0 import OpenAIV0Instrumentor
 from opentelemetry.instrumentation.openai.v1 import OpenAIV1Instrumentor
@@ -12,6 +13,10 @@
 class OpenAIInstrumentor(BaseInstrumentor):
     """An instrumentor for OpenAI's client library."""
 
+    def __init__(self, enrich_assistant: bool = False):
+        super().__init__()
+        Config.enrich_assistant = enrich_assistant
+
     def instrumentation_dependencies(self) -> Collection[str]:
         return _instruments
 

diff --git a/...pentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/shared/config.py b/...pentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/shared/config.py
@@ -0,0 +1,2 @@
+class Config:
+    enrich_assistant = False
diff --git a/...etry-instrumentation-openai/opentelemetry/instrumentation/openai/v1/assistant_wrappers.py b/...etry-instrumentation-openai/opentelemetry/instrumentation/openai/v1/assistant_wrappers.py
@@ -11,6 +11,7 @@
 from opentelemetry.semconv.ai import SpanAttributes, LLMRequestTypeValues
 
 from opentelemetry.instrumentation.openai.utils import _with_tracer_wrapper
+from opentelemetry.instrumentation.openai.shared.config import Config
 
 logger = logging.getLogger(__name__)
 
@@ -91,22 +92,30 @@ def messages_list_wrapper(tracer, wrapped, instance, args, kwargs):
     )
 
     i = 0
-    if assistants.get(run["assistant_id"]) is not None:
+    if assistants.get(run["assistant_id"]) is not None or Config.enrich_assistant:
+        if Config.enrich_assistant:
+            assistant = model_as_dict(
+                instance._client.beta.assistants.retrieve(run["assistant_id"])
+            )
+            assistants[run["assistant_id"]] = assistant
+        else:
+            assistant = assistants[run["assistant_id"]]
+
         _set_span_attribute(
             span,
             SpanAttributes.LLM_REQUEST_MODEL,
-            assistants[run["assistant_id"]]["model"],
+            assistant["model"],
         )
         _set_span_attribute(
             span,
             SpanAttributes.LLM_RESPONSE_MODEL,
-            assistants[run["assistant_id"]]["model"],
+            assistant["model"],
         )
         _set_span_attribute(span, f"{SpanAttributes.LLM_PROMPTS}.{i}.role", "system")
         _set_span_attribute(
             span,
             f"{SpanAttributes.LLM_PROMPTS}.{i}.content",
-            assistants[run["assistant_id"]]["instructions"],
+            assistant["instructions"],
         )
         i += 1
     _set_span_attribute(span, f"{SpanAttributes.LLM_PROMPTS}.{i}.role", "system")
@@ -143,7 +152,15 @@ def runs_create_and_stream_wrapper(tracer, wrapped, instance, args, kwargs):
     )
 
     i = 0
-    if assistants.get(assistant_id) is not None:
+    if assistants.get(assistant_id) is not None or Config.enrich_assistant:
+        if Config.enrich_assistant:
+            assistant = model_as_dict(
+                instance._client.beta.assistants.retrieve(assistant_id)
+            )
+            assistants[assistant_id] = assistant
+        else:
+            assistant = assistants[assistant_id]
+
         _set_span_attribute(
             span, SpanAttributes.LLM_REQUEST_MODEL, assistants[assistant_id]["model"]
         )