From ad564954b6ad32cf2be7a05e19ae629187607d6c Mon Sep 17 00:00:00 2001
From: Sam Brenner <sam.brenner@datadoghq.com>
Date: Thu, 6 Feb 2025 12:58:45 -0500
Subject: [PATCH] fix llmobs traces not submitting for langgraph errors

---
 ddtrace/contrib/internal/langgraph/patch.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/ddtrace/contrib/internal/langgraph/patch.py b/ddtrace/contrib/internal/langgraph/patch.py
index ab82271e83e..9fc198fcfec 100644
--- a/ddtrace/contrib/internal/langgraph/patch.py
+++ b/ddtrace/contrib/internal/langgraph/patch.py
@@ -121,6 +121,7 @@ def traced_pregel_stream(langgraph, pin, func, instance, args, kwargs):
         result = func(*args, **kwargs)
     except Exception:
         span.set_exc_info(*sys.exc_info())
+        integration.llmobs_set_tags(span, args=args, kwargs={**kwargs, "name": name}, response=None, operation="graph")
         span.finish()
         raise
 
@@ -139,6 +140,9 @@ def _stream():
                 break
             except Exception:
                 span.set_exc_info(*sys.exc_info())
+                integration.llmobs_set_tags(
+                    span, args=args, kwargs={**kwargs, "name": name}, response=None, operation="graph"
+                )
                 span.finish()
                 raise
 
@@ -160,6 +164,7 @@ def traced_pregel_astream(langgraph, pin, func, instance, args, kwargs):
         result = func(*args, **kwargs)
     except Exception:
         span.set_exc_info(*sys.exc_info())
+        integration.llmobs_set_tags(span, args=args, kwargs={**kwargs, "name": name}, response=None, operation="graph")
         span.finish()
         raise
 
@@ -178,6 +183,9 @@ async def _astream():
                 break
             except Exception:
                 span.set_exc_info(*sys.exc_info())
+                integration.llmobs_set_tags(
+                    span, args=args, kwargs={**kwargs, "name": name}, response=None, operation="graph"
+                )
                 span.finish()
                 raise