Skip to content

Commit d3ba332

Browse files
xrmxxrmx
authored
elastic-opentelemetry-instrumentation-openai: make OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT work like upstream (elastic#42)
So that it still sends events but clears the message content key.
1 parent 5a4c151 commit d3ba332

File tree

4 files changed

+165
-72
lines changed

4 files changed

+165
-72
lines changed

‎instrumentation/elastic-opentelemetry-instrumentation-openai/src/opentelemetry/instrumentation/openai/__init__.py‎

Lines changed: 26 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -145,11 +145,13 @@ def _chat_completion_wrapper(self, wrapped, instance, args, kwargs):
145145
# this is important to avoid having the span closed before ending the stream
146146
end_on_exit=False,
147147
) as span:
148-
# TODO: more fine grained depending on the message.role?
149-
if self.capture_message_content:
150-
messages = kwargs.get("messages", [])
151-
152-
_send_log_events_from_messages(self.event_logger, messages=messages, attributes=event_attributes)
148+
messages = kwargs.get("messages", [])
149+
_send_log_events_from_messages(
150+
self.event_logger,
151+
messages=messages,
152+
attributes=event_attributes,
153+
capture_message_content=self.capture_message_content,
154+
)
153155

154156
start_time = default_timer()
155157
try:
@@ -183,8 +185,12 @@ def _chat_completion_wrapper(self, wrapped, instance, args, kwargs):
183185
_record_token_usage_metrics(self.token_usage_metric, span, result.usage)
184186
_record_operation_duration_metric(self.operation_duration_metric, span, start_time)
185187

186-
if self.capture_message_content:
187-
_send_log_events_from_choices(self.event_logger, choices=result.choices, attributes=event_attributes)
188+
_send_log_events_from_choices(
189+
self.event_logger,
190+
choices=result.choices,
191+
attributes=event_attributes,
192+
capture_message_content=self.capture_message_content,
193+
)
188194

189195
span.end()
190196

@@ -204,9 +210,13 @@ async def _async_chat_completion_wrapper(self, wrapped, instance, args, kwargs):
204210
# this is important to avoid having the span closed before ending the stream
205211
end_on_exit=False,
206212
) as span:
207-
if self.capture_message_content:
208-
messages = kwargs.get("messages", [])
209-
_send_log_events_from_messages(self.event_logger, messages=messages, attributes=event_attributes)
213+
messages = kwargs.get("messages", [])
214+
_send_log_events_from_messages(
215+
self.event_logger,
216+
messages=messages,
217+
attributes=event_attributes,
218+
capture_message_content=self.capture_message_content,
219+
)
210220

211221
start_time = default_timer()
212222
try:
@@ -240,8 +250,12 @@ async def _async_chat_completion_wrapper(self, wrapped, instance, args, kwargs):
240250
_record_token_usage_metrics(self.token_usage_metric, span, result.usage)
241251
_record_operation_duration_metric(self.operation_duration_metric, span, start_time)
242252

243-
if self.capture_message_content:
244-
_send_log_events_from_choices(self.event_logger, choices=result.choices, attributes=event_attributes)
253+
_send_log_events_from_choices(
254+
self.event_logger,
255+
choices=result.choices,
256+
attributes=event_attributes,
257+
capture_message_content=self.capture_message_content,
258+
)
245259

246260
span.end()
247261

‎instrumentation/elastic-opentelemetry-instrumentation-openai/src/opentelemetry/instrumentation/openai/helpers.py‎

Lines changed: 20 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -245,18 +245,22 @@ def _serialize_tool_calls_for_event(tool_calls):
245245
]
246246

247247

248-
def _send_log_events_from_messages(event_logger: EventLogger, messages, attributes: Attributes):
248+
def _send_log_events_from_messages(
249+
event_logger: EventLogger, messages, attributes: Attributes, capture_message_content: bool
250+
):
249251
for message in messages:
252+
body = {}
253+
if capture_message_content:
254+
content = message.get("content")
255+
if content:
256+
body["content"] = content
250257
if message["role"] == "system":
251-
event = Event(name=EVENT_GEN_AI_SYSTEM_MESSAGE, body={"content": message["content"]}, attributes=attributes)
258+
event = Event(name=EVENT_GEN_AI_SYSTEM_MESSAGE, body=body, attributes=attributes)
252259
event_logger.emit(event)
253260
elif message["role"] == "user":
254-
event = Event(name=EVENT_GEN_AI_USER_MESSAGE, body={"content": message["content"]}, attributes=attributes)
261+
event = Event(name=EVENT_GEN_AI_USER_MESSAGE, body=body, attributes=attributes)
255262
event_logger.emit(event)
256263
elif message["role"] == "assistant":
257-
body = {}
258-
if content := message.get("content"):
259-
body["content"] = content
260264
tool_calls = _serialize_tool_calls_for_event(message.get("tool_calls", []))
261265
if tool_calls:
262266
body["tool_calls"] = tool_calls
@@ -267,28 +271,33 @@ def _send_log_events_from_messages(event_logger: EventLogger, messages, attribut
267271
)
268272
event_logger.emit(event)
269273
elif message["role"] == "tool":
274+
body["id"] = message["tool_call_id"]
270275
event = Event(
271276
name=EVENT_GEN_AI_TOOL_MESSAGE,
272-
body={"content": message["content"], "id": message["tool_call_id"]},
277+
body=body,
273278
attributes=attributes,
274279
)
275280
event_logger.emit(event)
276281

277282

278-
def _send_log_events_from_choices(event_logger: EventLogger, choices, attributes: Attributes):
283+
def _send_log_events_from_choices(
284+
event_logger: EventLogger, choices, attributes: Attributes, capture_message_content: bool
285+
):
279286
for choice in choices:
280287
tool_calls = _serialize_tool_calls_for_event(choice.message.tool_calls or [])
281288
body = {"finish_reason": choice.finish_reason, "index": choice.index, "message": {}}
282289
if tool_calls:
283290
body["message"]["tool_calls"] = tool_calls
284-
if choice.message.content:
291+
if capture_message_content and choice.message.content:
285292
body["message"]["content"] = choice.message.content
286293

287294
event = Event(name=EVENT_GEN_AI_CHOICE, body=body, attributes=attributes)
288295
event_logger.emit(event)
289296

290297

291-
def _send_log_events_from_stream_choices(event_logger: EventLogger, choices, span: Span, attributes: Attributes):
298+
def _send_log_events_from_stream_choices(
299+
event_logger: EventLogger, choices, span: Span, attributes: Attributes, capture_message_content: bool
300+
):
292301
body = {}
293302
message = {}
294303
message_content = ""
@@ -312,7 +321,7 @@ def _send_log_events_from_stream_choices(event_logger: EventLogger, choices, spa
312321
body["finish_reason"] = choice.finish_reason
313322
body["index"] = choice.index
314323

315-
if message_content:
324+
if capture_message_content and message_content:
316325
message["content"] = message_content
317326
if tool_calls:
318327
message["tool_calls"] = [call for _, call in sorted(tool_calls.items())]

‎instrumentation/elastic-opentelemetry-instrumentation-openai/src/opentelemetry/instrumentation/openai/wrappers.py‎

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -79,10 +79,13 @@ def end(self, exc=None):
7979
if self.usage:
8080
_record_token_usage_metrics(self.token_usage_metric, self.span, self.usage)
8181

82-
if self.capture_message_content:
83-
_send_log_events_from_stream_choices(
84-
self.event_logger, choices=self.choices, span=self.span, attributes=self.event_attributes
85-
)
82+
_send_log_events_from_stream_choices(
83+
self.event_logger,
84+
choices=self.choices,
85+
span=self.span,
86+
attributes=self.event_attributes,
87+
capture_message_content=self.capture_message_content,
88+
)
8689

8790
self.span.end()
8891

0 commit comments

Comments
 (0)