Skip to content

Commit dee6de1

Browse files
shellmayrantonpirker
shellmayr
and
antonpirker
authored
feat(agents): improve instrumentation of input messages (#4750)
- Improve the instrumentation of input messages in the AI agents instrumentations. Before: <img width="809" height="585" alt="Screenshot 2025-09-03 at 16 30 05" src="https://github.com/user-attachments/assets/39b9a9b5-6599-4fae-b696-ea2cf5930ae9" /> After: <img width="813" height="587" alt="Screenshot 2025-09-03 at 16 30 08" src="https://github.com/user-attachments/assets/9898213c-7525-4b50-87c5-a7b51c4f5cd7" /> Closes TET-1058 --------- Co-authored-by: Anton Pirker <[email protected]>
1 parent ad3c435 commit dee6de1

File tree

7 files changed

+234
-62
lines changed

7 files changed

+234
-62
lines changed

‎sentry_sdk/integrations/langchain.py

Lines changed: 105 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,6 @@
5151
"presence_penalty": SPANDATA.GEN_AI_REQUEST_PRESENCE_PENALTY,
5252
"temperature": SPANDATA.GEN_AI_REQUEST_TEMPERATURE,
5353
"tool_calls": SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS,
54-
"tools": SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS,
5554
"top_k": SPANDATA.GEN_AI_REQUEST_TOP_K,
5655
"top_p": SPANDATA.GEN_AI_REQUEST_TOP_P,
5756
}
@@ -203,8 +202,12 @@ def on_llm_start(
203202
if key in all_params and all_params[key] is not None:
204203
set_data_normalized(span, attribute, all_params[key], unpack=False)
205204

205+
_set_tools_on_span(span, all_params.get("tools"))
206+
206207
if should_send_default_pii() and self.include_prompts:
207-
set_data_normalized(span, SPANDATA.GEN_AI_REQUEST_MESSAGES, prompts)
208+
set_data_normalized(
209+
span, SPANDATA.GEN_AI_REQUEST_MESSAGES, prompts, unpack=False
210+
)
208211

209212
def on_chat_model_start(self, serialized, messages, *, run_id, **kwargs):
210213
# type: (SentryLangchainCallback, Dict[str, Any], List[List[BaseMessage]], UUID, Any) -> Any
@@ -246,14 +249,20 @@ def on_chat_model_start(self, serialized, messages, *, run_id, **kwargs):
246249
if key in all_params and all_params[key] is not None:
247250
set_data_normalized(span, attribute, all_params[key], unpack=False)
248251

252+
_set_tools_on_span(span, all_params.get("tools"))
253+
249254
if should_send_default_pii() and self.include_prompts:
255+
normalized_messages = []
256+
for list_ in messages:
257+
for message in list_:
258+
normalized_messages.append(
259+
self._normalize_langchain_message(message)
260+
)
250261
set_data_normalized(
251262
span,
252263
SPANDATA.GEN_AI_REQUEST_MESSAGES,
253-
[
254-
[self._normalize_langchain_message(x) for x in list_]
255-
for list_ in messages
256-
],
264+
normalized_messages,
265+
unpack=False,
257266
)
258267

259268
def on_chat_model_end(self, response, *, run_id, **kwargs):
@@ -351,9 +360,7 @@ def on_agent_finish(self, finish, *, run_id, **kwargs):
351360

352361
if should_send_default_pii() and self.include_prompts:
353362
set_data_normalized(
354-
span,
355-
SPANDATA.GEN_AI_RESPONSE_TEXT,
356-
finish.return_values.items(),
363+
span, SPANDATA.GEN_AI_RESPONSE_TEXT, finish.return_values.items()
357364
)
358365

359366
self._exit_span(span_data, run_id)
@@ -473,13 +480,11 @@ def _get_token_usage(obj):
473480
if usage is not None:
474481
return usage
475482

476-
# check for usage in the object itself
477483
for name in possible_names:
478484
usage = _get_value(obj, name)
479485
if usage is not None:
480486
return usage
481487

482-
# no usage found anywhere
483488
return None
484489

485490

@@ -531,6 +536,87 @@ def _get_request_data(obj, args, kwargs):
531536
return (agent_name, tools)
532537

533538

539+
def _simplify_langchain_tools(tools):
540+
# type: (Any) -> Optional[List[Any]]
541+
"""Parse and simplify tools into a cleaner format."""
542+
if not tools:
543+
return None
544+
545+
if not isinstance(tools, (list, tuple)):
546+
return None
547+
548+
simplified_tools = []
549+
for tool in tools:
550+
try:
551+
if isinstance(tool, dict):
552+
553+
if "function" in tool and isinstance(tool["function"], dict):
554+
func = tool["function"]
555+
simplified_tool = {
556+
"name": func.get("name"),
557+
"description": func.get("description"),
558+
}
559+
if simplified_tool["name"]:
560+
simplified_tools.append(simplified_tool)
561+
elif "name" in tool:
562+
simplified_tool = {
563+
"name": tool.get("name"),
564+
"description": tool.get("description"),
565+
}
566+
simplified_tools.append(simplified_tool)
567+
else:
568+
name = (
569+
tool.get("name")
570+
or tool.get("tool_name")
571+
or tool.get("function_name")
572+
)
573+
if name:
574+
simplified_tools.append(
575+
{
576+
"name": name,
577+
"description": tool.get("description")
578+
or tool.get("desc"),
579+
}
580+
)
581+
elif hasattr(tool, "name"):
582+
simplified_tool = {
583+
"name": getattr(tool, "name", None),
584+
"description": getattr(tool, "description", None)
585+
or getattr(tool, "desc", None),
586+
}
587+
if simplified_tool["name"]:
588+
simplified_tools.append(simplified_tool)
589+
elif hasattr(tool, "__name__"):
590+
simplified_tools.append(
591+
{
592+
"name": tool.__name__,
593+
"description": getattr(tool, "__doc__", None),
594+
}
595+
)
596+
else:
597+
tool_str = str(tool)
598+
if tool_str and tool_str != "":
599+
simplified_tools.append({"name": tool_str, "description": None})
600+
except Exception:
601+
continue
602+
603+
return simplified_tools if simplified_tools else None
604+
605+
606+
def _set_tools_on_span(span, tools):
607+
# type: (Span, Any) -> None
608+
"""Set available tools data on a span if tools are provided."""
609+
if tools is not None:
610+
simplified_tools = _simplify_langchain_tools(tools)
611+
if simplified_tools:
612+
set_data_normalized(
613+
span,
614+
SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS,
615+
simplified_tools,
616+
unpack=False,
617+
)
618+
619+
534620
def _wrap_configure(f):
535621
# type: (Callable[..., Any]) -> Callable[..., Any]
536622

@@ -601,7 +687,7 @@ def new_configure(
601687
]
602688
elif isinstance(local_callbacks, BaseCallbackHandler):
603689
local_callbacks = [local_callbacks, sentry_handler]
604-
else: # local_callbacks is a list
690+
else:
605691
local_callbacks = [*local_callbacks, sentry_handler]
606692

607693
return f(
@@ -638,10 +724,7 @@ def new_invoke(self, *args, **kwargs):
638724
span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "invoke_agent")
639725
span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, False)
640726

641-
if tools:
642-
set_data_normalized(
643-
span, SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, tools, unpack=False
644-
)
727+
_set_tools_on_span(span, tools)
645728

646729
# Run the agent
647730
result = f(self, *args, **kwargs)
@@ -653,11 +736,7 @@ def new_invoke(self, *args, **kwargs):
653736
and integration.include_prompts
654737
):
655738
set_data_normalized(
656-
span,
657-
SPANDATA.GEN_AI_REQUEST_MESSAGES,
658-
[
659-
input,
660-
],
739+
span, SPANDATA.GEN_AI_REQUEST_MESSAGES, [input], unpack=False
661740
)
662741

663742
output = result.get("output")
@@ -666,7 +745,7 @@ def new_invoke(self, *args, **kwargs):
666745
and should_send_default_pii()
667746
and integration.include_prompts
668747
):
669-
span.set_data(SPANDATA.GEN_AI_RESPONSE_TEXT, output)
748+
set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, output)
670749

671750
return result
672751

@@ -698,10 +777,7 @@ def new_stream(self, *args, **kwargs):
698777
span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "invoke_agent")
699778
span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, True)
700779

701-
if tools:
702-
set_data_normalized(
703-
span, SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, tools, unpack=False
704-
)
780+
_set_tools_on_span(span, tools)
705781

706782
input = args[0].get("input") if len(args) >= 1 else None
707783
if (
@@ -710,11 +786,7 @@ def new_stream(self, *args, **kwargs):
710786
and integration.include_prompts
711787
):
712788
set_data_normalized(
713-
span,
714-
SPANDATA.GEN_AI_REQUEST_MESSAGES,
715-
[
716-
input,
717-
],
789+
span, SPANDATA.GEN_AI_REQUEST_MESSAGES, [input], unpack=False
718790
)
719791

720792
# Run the agent
@@ -737,7 +809,7 @@ def new_iterator():
737809
and should_send_default_pii()
738810
and integration.include_prompts
739811
):
740-
span.set_data(SPANDATA.GEN_AI_RESPONSE_TEXT, output)
812+
set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, output)
741813

742814
span.__exit__(None, None, None)
743815

@@ -756,7 +828,7 @@ async def new_iterator_async():
756828
and should_send_default_pii()
757829
and integration.include_prompts
758830
):
759-
span.set_data(SPANDATA.GEN_AI_RESPONSE_TEXT, output)
831+
set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, output)
760832

761833
span.__exit__(None, None, None)
762834

‎sentry_sdk/integrations/langgraph.py

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -183,7 +183,8 @@ def new_invoke(self, *args, **kwargs):
183183
set_data_normalized(
184184
span,
185185
SPANDATA.GEN_AI_REQUEST_MESSAGES,
186-
safe_serialize(input_messages),
186+
input_messages,
187+
unpack=False,
187188
)
188189

189190
result = f(self, *args, **kwargs)
@@ -232,7 +233,8 @@ async def new_ainvoke(self, *args, **kwargs):
232233
set_data_normalized(
233234
span,
234235
SPANDATA.GEN_AI_REQUEST_MESSAGES,
235-
safe_serialize(input_messages),
236+
input_messages,
237+
unpack=False,
236238
)
237239

238240
result = await f(self, *args, **kwargs)
@@ -305,11 +307,9 @@ def _set_response_attributes(span, input_messages, result, integration):
305307
if llm_response_text:
306308
set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, llm_response_text)
307309
elif new_messages:
308-
set_data_normalized(
309-
span, SPANDATA.GEN_AI_RESPONSE_TEXT, safe_serialize(new_messages)
310-
)
310+
set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, new_messages)
311311
else:
312-
set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, safe_serialize(result))
312+
set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, result)
313313

314314
tool_calls = _extract_tool_calls(new_messages)
315315
if tool_calls:

‎sentry_sdk/integrations/openai.py

Lines changed: 33 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -179,7 +179,9 @@ def _set_input_data(span, kwargs, operation, integration):
179179
and should_send_default_pii()
180180
and integration.include_prompts
181181
):
182-
set_data_normalized(span, SPANDATA.GEN_AI_REQUEST_MESSAGES, messages)
182+
set_data_normalized(
183+
span, SPANDATA.GEN_AI_REQUEST_MESSAGES, messages, unpack=False
184+
)
183185

184186
# Input attributes: Common
185187
set_data_normalized(span, SPANDATA.GEN_AI_SYSTEM, "openai")
@@ -227,25 +229,46 @@ def _set_output_data(span, response, kwargs, integration, finish_span=True):
227229
if should_send_default_pii() and integration.include_prompts:
228230
response_text = [choice.message.dict() for choice in response.choices]
229231
if len(response_text) > 0:
230-
set_data_normalized(
231-
span,
232-
SPANDATA.GEN_AI_RESPONSE_TEXT,
233-
safe_serialize(response_text),
234-
)
232+
set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, response_text)
233+
235234
_calculate_token_usage(messages, response, span, None, integration.count_tokens)
235+
236236
if finish_span:
237237
span.__exit__(None, None, None)
238238

239239
elif hasattr(response, "output"):
240240
if should_send_default_pii() and integration.include_prompts:
241-
response_text = [item.to_dict() for item in response.output]
242-
if len(response_text) > 0:
241+
output_messages = {
242+
"response": [],
243+
"tool": [],
244+
} # type: (dict[str, list[Any]])
245+
246+
for output in response.output:
247+
if output.type == "function_call":
248+
output_messages["tool"].append(output.dict())
249+
elif output.type == "message":
250+
for output_message in output.content:
251+
try:
252+
output_messages["response"].append(output_message.text)
253+
except AttributeError:
254+
# Unknown output message type, just return the json
255+
output_messages["response"].append(output_message.dict())
256+
257+
if len(output_messages["tool"]) > 0:
243258
set_data_normalized(
244259
span,
245-
SPANDATA.GEN_AI_RESPONSE_TEXT,
246-
safe_serialize(response_text),
260+
SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS,
261+
output_messages["tool"],
262+
unpack=False,
263+
)
264+
265+
if len(output_messages["response"]) > 0:
266+
set_data_normalized(
267+
span, SPANDATA.GEN_AI_RESPONSE_TEXT, output_messages["response"]
247268
)
269+
248270
_calculate_token_usage(messages, response, span, None, integration.count_tokens)
271+
249272
if finish_span:
250273
span.__exit__(None, None, None)
251274

‎sentry_sdk/integrations/openai_agents/utils.py

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import sentry_sdk
2+
from sentry_sdk.ai.utils import set_data_normalized
23
from sentry_sdk.consts import SPANDATA
34
from sentry_sdk.integrations import DidNotEnable
45
from sentry_sdk.scope import should_send_default_pii
@@ -127,7 +128,9 @@ def _set_input_data(span, get_response_kwargs):
127128
if len(messages) > 0:
128129
request_messages.append({"role": role, "content": messages})
129130

130-
span.set_data(SPANDATA.GEN_AI_REQUEST_MESSAGES, safe_serialize(request_messages))
131+
set_data_normalized(
132+
span, SPANDATA.GEN_AI_REQUEST_MESSAGES, request_messages, unpack=False
133+
)
131134

132135

133136
def _set_output_data(span, result):
@@ -157,6 +160,6 @@ def _set_output_data(span, result):
157160
)
158161

159162
if len(output_messages["response"]) > 0:
160-
span.set_data(
161-
SPANDATA.GEN_AI_RESPONSE_TEXT, safe_serialize(output_messages["response"])
163+
set_data_normalized(
164+
span, SPANDATA.GEN_AI_RESPONSE_TEXT, output_messages["response"]
162165
)

0 commit comments

Comments
 (0)