fix(integrations): ensure that GEN_AI_AGENT_NAME is properly set for GEN_AI spans under an invoke_agent span

Author: constantinius

sentry_sdk/integrations/google_genai/__init__.py

@@ -92,6 +92,7 @@ def new_generate_content_stream(self, *args, **kwargs):
         chat_span.set_data(SPANDATA.GEN_AI_REQUEST_MODEL, model_name)
         set_span_data_for_request(chat_span, integration, model_name, contents, kwargs)
         chat_span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, True)
+        chat_span.set_data(SPANDATA.GEN_AI_AGENT_NAME, model_name)
 
         try:
             stream = f(self, *args, **kwargs)
@@ -165,6 +166,7 @@ async def new_async_generate_content_stream(self, *args, **kwargs):
         chat_span.set_data(SPANDATA.GEN_AI_REQUEST_MODEL, model_name)
         set_span_data_for_request(chat_span, integration, model_name, contents, kwargs)
         chat_span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, True)
+        chat_span.set_data(SPANDATA.GEN_AI_AGENT_NAME, model_name)
 
         try:
             stream = await f(self, *args, **kwargs)
@@ -233,6 +235,7 @@ def new_generate_content(self, *args, **kwargs):
                 chat_span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "chat")
                 chat_span.set_data(SPANDATA.GEN_AI_SYSTEM, GEN_AI_SYSTEM)
                 chat_span.set_data(SPANDATA.GEN_AI_REQUEST_MODEL, model_name)
+                chat_span.set_data(SPANDATA.GEN_AI_AGENT_NAME, model_name)
                 set_span_data_for_request(
                     chat_span, integration, model_name, contents, kwargs
                 )

sentry_sdk/integrations/langchain.py

@@ -276,6 +276,14 @@ def on_chat_model_start(self, serialized, messages, *, run_id, **kwargs):
             elif "openai" in ai_type:
                 span.set_data(SPANDATA.GEN_AI_SYSTEM, "openai")
 
+            agent_name = (
+                sentry_sdk.get_current_scope()
+                ._contexts.get("langchain_agent", {})
+                .get("agent_name")
+            )
+            if agent_name:
+                span.set_data(SPANDATA.GEN_AI_AGENT_NAME, agent_name)
+
             for key, attribute in DATA_FIELDS.items():
                 if key in all_params and all_params[key] is not None:
                     set_data_normalized(span, attribute, all_params[key], unpack=False)
@@ -428,6 +436,14 @@ def on_tool_start(self, serialized, input_str, *, run_id, **kwargs):
             if tool_description is not None:
                 span.set_data(SPANDATA.GEN_AI_TOOL_DESCRIPTION, tool_description)
 
+            agent_name = (
+                sentry_sdk.get_current_scope()
+                ._contexts.get("langchain_agent", {})
+                .get("agent_name")
+            )
+            if agent_name:
+                span.set_data(SPANDATA.GEN_AI_AGENT_NAME, agent_name)
+
             if should_send_default_pii() and self.include_prompts:
                 set_data_normalized(
                     span,
@@ -756,6 +772,9 @@ def new_invoke(self, *args, **kwargs):
             name=f"invoke_agent {agent_name}" if agent_name else "invoke_agent",
             origin=LangchainIntegration.origin,
         ) as span:
+            sentry_sdk.get_current_scope().set_context(
+                "langchain_agent", {"agent_name": agent_name, "tools": tools}
+            )
             if agent_name:
                 span.set_data(SPANDATA.GEN_AI_AGENT_NAME, agent_name)
 
@@ -794,6 +813,8 @@ def new_invoke(self, *args, **kwargs):
             ):
                 set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, output)
 
+            sentry_sdk.get_current_scope().remove_context("langchain_agent")
+
             return result
 
     return new_invoke
@@ -814,11 +835,15 @@ def new_stream(self, *args, **kwargs):
 
         span = start_span_function(
             op=OP.GEN_AI_INVOKE_AGENT,
-            name=f"invoke_agent {agent_name}".strip(),
+            name=f"invoke_agent {agent_name}" if agent_name else "invoke_agent",
             origin=LangchainIntegration.origin,
         )
         span.__enter__()
 
+        sentry_sdk.get_current_scope().set_context(
+            "langchain_agent", {"agent_name": agent_name, "tools": tools}
+        )
+
         if agent_name:
             span.set_data(SPANDATA.GEN_AI_AGENT_NAME, agent_name)
 
@@ -868,6 +893,8 @@ def new_iterator():
             ):
                 set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, output)
 
+            sentry_sdk.get_current_scope().remove_context("langchain_agent")
+
             span.__exit__(None, None, None)
 
         async def new_iterator_async():
@@ -887,6 +914,8 @@ async def new_iterator_async():
             ):
                 set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, output)
 
+            sentry_sdk.get_current_scope().remove_context("langchain_agent")
+
             span.__exit__(None, None, None)
 
         if str(type(result)) == "<class 'async_generator'>":