wrap agent executor

Author: shellmayr

sentry_sdk/integrations/langchain.py

@@ -27,6 +27,8 @@
         Callbacks,
     )
     from langchain_core.agents import AgentAction, AgentFinish
+    from langchain.agents import AgentExecutor
+
 except ImportError:
     raise DidNotEnable("langchain not installed")
 
@@ -71,6 +73,10 @@ def setup_once():
         # type: () -> None
         manager._configure = _wrap_configure(manager._configure)
 
+        if AgentExecutor is not None:
+            AgentExecutor.invoke = _wrap_agent_executor_invoke(AgentExecutor.invoke)
+            AgentExecutor.stream = _wrap_agent_executor_stream(AgentExecutor.stream)
+
 
 class WatchedSpan:
     span = None  # type: Span
@@ -265,7 +271,7 @@ def on_llm_start(
 
             watched_span = self._create_span(
                 run_id,
-                kwargs.get("parent_run_id"),
+                parent_id=parent_run_id,
                 op=OP.GEN_AI_PIPELINE,
                 name=kwargs.get("name") or "Langchain LLM call",
                 origin=LangchainIntegration.origin,
@@ -293,7 +299,7 @@ def on_chat_model_start(self, serialized, messages, *, run_id, **kwargs):
             all_params.update(serialized.get("kwargs", {}))
             watched_span = self._create_span(
                 run_id,
-                kwargs.get("parent_run_id"),
+                parent_id=kwargs.get("parent_run_id"),
                 op=OP.GEN_AI_CHAT,
                 name=kwargs.get("name") or "Langchain Chat Model",
                 origin=LangchainIntegration.origin,
@@ -483,22 +489,7 @@ def on_chain_error(self, error, *, run_id, **kwargs):
 
     def on_agent_action(self, action, *, run_id, **kwargs):
         # type: (SentryLangchainCallback, AgentAction, UUID, Any) -> Any
-        with capture_internal_exceptions():
-            if not run_id:
-                return
-            watched_span = self._create_span(
-                run_id,
-                kwargs.get("parent_run_id"),
-                op=OP.GEN_AI_INVOKE_AGENT,
-                name=action.tool or "AI tool usage",
-                origin=LangchainIntegration.origin,
-            )
-            if action.tool_input and should_send_default_pii() and self.include_prompts:
-                set_data_normalized(
-                    watched_span.span,
-                    SPANDATA.GEN_AI_REQUEST_MESSAGES,
-                    action.tool_input,
-                )
+        pass
 
     def on_agent_finish(self, finish, *, run_id, **kwargs):
         # type: (SentryLangchainCallback, AgentFinish, UUID, Any) -> Any
@@ -528,7 +519,7 @@ def on_tool_start(self, serialized, input_str, *, run_id, **kwargs):
 
             watched_span = self._create_span(
                 run_id,
-                kwargs.get("parent_run_id"),
+                parent_id=kwargs.get("parent_run_id"),
                 op=OP.GEN_AI_EXECUTE_TOOL,
                 name=f"execute_tool {tool_name}",
                 origin=LangchainIntegration.origin,
@@ -657,3 +648,63 @@ def new_configure(
         )
 
     return new_configure
+
+
+def _wrap_agent_executor_invoke(f):
+    # type: (Callable[..., Any]) -> Callable[..., Any]
+
+    @wraps(f)
+    def new_invoke(self, *args, **kwargs):
+        # type: (Any, Any, Any) -> Any
+
+        integration = sentry_sdk.get_client().get_integration(LangchainIntegration)
+        if integration is None:
+            return f(self, *args, **kwargs)
+
+        # Create a span that will act as the parent for all callback-generated spans
+        with sentry_sdk.start_span(
+            op=OP.GEN_AI_INVOKE_AGENT,
+            name="AgentExecutor.invoke",
+            origin=LangchainIntegration.origin,
+        ) as span:
+            span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "invoke_agent")
+            if hasattr(self, "agent") and hasattr(self.agent, "llm"):
+                model_name = getattr(self.agent.llm, "model_name", None) or getattr(
+                    self.agent.llm, "model", None
+                )
+                if model_name:
+                    span.set_data(SPANDATA.GEN_AI_REQUEST_MODEL, model_name)
+
+            return f(self, *args, **kwargs)
+
+    return new_invoke
+
+
+def _wrap_agent_executor_stream(f):
+    # type: (Callable[..., Any]) -> Callable[..., Any]
+
+    @wraps(f)
+    def new_stream(self, *args, **kwargs):
+        # type: (Any, Any, Any) -> Any
+
+        integration = sentry_sdk.get_client().get_integration(LangchainIntegration)
+        if integration is None:
+            return f(self, *args, **kwargs)
+
+        # Create a span that will act as the parent for all callback-generated spans
+        with sentry_sdk.start_span(
+            op=OP.GEN_AI_INVOKE_AGENT,
+            name="AgentExecutor.stream",
+            origin=LangchainIntegration.origin,
+        ) as span:
+            span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "invoke_agent")
+            if hasattr(self, "agent") and hasattr(self.agent, "llm"):
+                model_name = getattr(self.agent.llm, "model_name", None) or getattr(
+                    self.agent.llm, "model", None
+                )
+                if model_name:
+                    span.set_data(SPANDATA.GEN_AI_REQUEST_MODEL, model_name)
+
+            return f(self, *args, **kwargs)
+
+    return new_stream