refactor tool tags list

DataDog · Sep 2, 2024 · 7781f30 · 7781f30
1 parent af98c82
commit 7781f30
Showing 1 changed file with 18 additions and 24 deletions.
diff --git a/ddtrace/contrib/internal/langchain/patch.py b/ddtrace/contrib/internal/langchain/patch.py
@@ -990,22 +990,19 @@ def traced_base_tool_invoke(langchain, pin, func, instance, args, kwargs):
 
     tool_output = None
     try:
-        tool_attributes = [
-            "name",
-            "description",
-        ]
+        tool_attributes = ["name", "description", "metadata", "tags"]
+
         for attribute in tool_attributes:
             value = getattr(instance, attribute, None)
-            if value:
+            if isinstance(value, dict):
+                for key, meta_value in value.items():
+                    span.set_tag_str("langchain.request.tool.%s.%s" % (attribute, key), str(value))
+            elif isinstance(value, list):
+                for idx, tag in enumerate(value):
+                    span.set_tag_str("langchain.request.tool.%s.%d" % (attribute, idx), str(value))
+            elif value is not None:
                 span.set_tag_str("langchain.request.tool.%s" % attribute, str(value))
 
-        if getattr(instance, "metadata", None):
-            for key, value in instance.metadata.items():
-                span.set_tag_str("langchain.request.tool.metadata.%s" % key, str(value))
-        if getattr(instance, "tags", None):
-            for idx, tag in enumerate(instance.tags):
-                span.set_tag_str("langchain.request.tool.tags.%d" % idx, str(tag))
-
         if integration.is_pc_sampled_span(span):
             if tool_input:
                 span.set_tag_str("langchain.request.input", integration.trunc(str(tool_input)))
@@ -1046,22 +1043,19 @@ async def traced_base_tool_ainvoke(langchain, pin, func, instance, args, kwargs)
 
     tool_output = None
     try:
-        tool_attributes = [
-            "name",
-            "description",
-        ]
+        tool_attributes = ["name", "description", "metadata", "tags"]
+
         for attribute in tool_attributes:
             value = getattr(instance, attribute, None)
-            if value:
+            if isinstance(value, dict):
+                for key, meta_value in value.items():
+                    span.set_tag_str("langchain.request.tool.%s.%s" % (attribute, key), str(value))
+            elif isinstance(value, list):
+                for idx, tag in enumerate(value):
+                    span.set_tag_str("langchain.request.tool.%s.%d" % (attribute, idx), str(value))
+            elif value is not None:
                 span.set_tag_str("langchain.request.tool.%s" % attribute, str(value))
 
-        if getattr(instance, "metadata", None):
-            for key, value in instance.metadata.items():
-                span.set_tag_str("langchain.request.tool.metadata.%s" % key, str(value))
-        if getattr(instance, "tags", None):
-            for idx, tag in enumerate(instance.tags):
-                span.set_tag_str("langchain.request.tool.tags.%d" % idx, str(tag))
-
         if integration.is_pc_sampled_span(span):
             if tool_input:
                 span.set_tag_str("langchain.request.input", integration.trunc(str(tool_input)))