feat(llmobs): submit langchain tool.invoke tool spans

ddtrace/contrib/internal/langchain/patch.py

@@ -1,3 +1,4 @@
+import json
 import os
 import sys
 from typing import Any
@@ -974,6 +975,118 @@ def traced_similarity_search(langchain, pin, func, instance, args, kwargs):
     return documents
 
 
+@with_traced_module
+def traced_base_tool_invoke(langchain, pin, func, instance, args, kwargs):
+    integration = langchain._datadog_integration
+    tool_input = get_argument_value(args, kwargs, 0, "input")
+    config = get_argument_value(args, kwargs, 1, "config", optional=True)
+
+    span = integration.trace(
+        pin,
+        "%s.%s.%s.%s" % (func.__module__, func.__class__.__name__, func.__name__, func.__self__.name),
+        interface_type="tool",
+        submit_to_llmobs=True,
+    )
+
+    tool_output = None
+    try:
+        tool_attributes = [
+            "name",
+            "description",
+        ]
+        for attribute in tool_attributes:
+            value = getattr(instance, attribute, None)
+            if value:
+                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)))
+            if config:
+                span.set_tag_str("langchain.request.config", json.dumps(config))
+        tool_output = func(*args, **kwargs)
+        if tool_output is not None:
+            if integration.is_pc_sampled_span(span):
+                span.set_tag_str("langchain.response.output", integration.trunc(str(tool_output)))
+    except Exception:
+        span.set_exc_info(*sys.exc_info())
+        raise
+    finally:
+        if integration.is_pc_sampled_llmobs(span):
+            integration.llmobs_set_tags(
+                "tool",
+                span,
+                tool_input,
+                tool_output,
+                error=bool(span.error),
+            )
+        span.finish()
+    return tool_output
+
+
+@with_traced_module
+async def traced_base_tool_ainvoke(langchain, pin, func, instance, args, kwargs):
+    integration = langchain._datadog_integration
+    tool_input = get_argument_value(args, kwargs, 0, "input")
+    tool_config = get_argument_value(args, kwargs, 1, "config", optional=True)
+
+    span = integration.trace(
+        pin,
+        "%s" % func.__self__.name,
+        interface_type="tool",
+        submit_to_llmobs=True,
+    )
+
+    tool_output = None
+    try:
+        tool_attributes = [
+            "name",
+            "description",
+        ]
+        for attribute in tool_attributes:
+            value = getattr(instance, attribute, None)
+            if value:
+                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)))
+            if tool_config:
+                span.set_tag_str("langchain.request.config", json.dumps(tool_config))
+        tool_output = await func(*args, **kwargs)
+        if tool_output is not None:
+            if integration.is_pc_sampled_span(span):
+                span.set_tag_str("langchain.response.output", integration.trunc(str(tool_output)))
+    except Exception:
+        span.set_exc_info(*sys.exc_info())
+        raise
+    finally:
+        if integration.is_pc_sampled_llmobs(span):
+            integration.llmobs_set_tags(
+                "tool",
+                span,
+                tool_input,
+                tool_output,
+                error=bool(span.error),
+            )
+        span.finish()
+    return tool_output
+
+
 def _patch_embeddings_and_vectorstores():
     """
     Text embedding models override two abstract base methods instead of super calls,
@@ -1101,6 +1214,8 @@ def patch():
         )
         wrap("langchain_core", "runnables.base.RunnableSequence.batch", traced_lcel_runnable_sequence(langchain))
         wrap("langchain_core", "runnables.base.RunnableSequence.abatch", traced_lcel_runnable_sequence_async(langchain))
+        wrap("langchain_core", "tools.BaseTool.invoke", traced_base_tool_invoke(langchain))
+        wrap("langchain_core", "tools.BaseTool.ainvoke", traced_base_tool_ainvoke(langchain))
         if langchain_openai:
             wrap("langchain_openai", "OpenAIEmbeddings.embed_documents", traced_embedding(langchain))
         if langchain_pinecone:

ddtrace/llmobs/_integrations/langchain.py

@@ -45,15 +45,15 @@
     "system": "system",
 }
 
-SUPPORTED_OPERATIONS = ["llm", "chat", "chain", "embedding", "retrieval"]
+SUPPORTED_OPERATIONS = ["llm", "chat", "chain", "embedding", "retrieval", "tool"]
 
 
 class LangChainIntegration(BaseLLMIntegration):
     _integration_name = "langchain"
 
     def llmobs_set_tags(
         self,
-        operation: str,  # oneof "llm","chat","chain","embedding","retrieval"
+        operation: str,  # oneof "llm","chat","chain","embedding","retrieval","tool"
         span: Span,
         inputs: Any,
         response: Any = None,
@@ -92,6 +92,8 @@ def llmobs_set_tags(
             self._llmobs_set_meta_tags_from_embedding(span, inputs, response, error, is_workflow=is_workflow)
         elif operation == "retrieval":
             self._llmobs_set_meta_tags_from_similarity_search(span, inputs, response, error, is_workflow=is_workflow)
+        elif operation == "tool":
+            self._llmobs_set_meta_tags_from_tool(span, inputs, response, error)
         span.set_tag_str(METRICS, json.dumps({}))
 
     def _llmobs_set_metadata(self, span: Span, model_provider: Optional[str] = None) -> None:
@@ -304,6 +306,35 @@ def _llmobs_set_meta_tags_from_similarity_search(
                 except TypeError:
                     log.warning("Failed to serialize similarity output documents to JSON")
 
+    def _llmobs_set_meta_tags_from_tool(
+        self,
+        span: Span,
+        tool_input: Union[str, Dict[str, object], object],
+        tool_output: object,
+        error: bool,
+    ) -> None:
+        span.set_tag_str(SPAN_KIND, "tool")
+        if tool_input is not None:
+            try:
+                formatted_inputs = self.format_io(tool_input)
+                if isinstance(formatted_inputs, str):
+                    span.set_tag_str(INPUT_VALUE, formatted_inputs)
+                else:
+                    span.set_tag_str(INPUT_VALUE, json.dumps(self.format_io(tool_input)))
+            except TypeError:
+                log.warning("Failed to serialize tool input data to JSON")
+        if error:
+            span.set_tag_str(OUTPUT_VALUE, "")
+        elif tool_output is not None:
+            try:
+                formatted_outputs = self.format_io(tool_output)
+                if isinstance(formatted_outputs, str):
+                    span.set_tag_str(OUTPUT_VALUE, formatted_outputs)
+                else:
+                    span.set_tag_str(OUTPUT_VALUE, json.dumps(self.format_io(tool_output)))
+            except TypeError:
+                log.warning("Failed to serialize tool output data to JSON")
+
     def _set_base_span_tags(  # type: ignore[override]
         self,
         span: Span,

releasenotes/notes/feat-llmobs-submit-langchain-tool-spans-d7274ae79d3808e9.yaml

@@ -0,0 +1,4 @@
+---
+features:
+  - |
+    LLM Observability: The LangChain integration now submits tool spans to LLM Observability.

tests/contrib/langchain/test_langchain_community.py

@@ -1274,3 +1274,65 @@ def test_faiss_vectorstore_retrieval(langchain_community, langchain_openai, requ
             retriever = faiss.as_retriever()
         with request_vcr.use_cassette("openai_retrieval_embedding.yaml"):
             retriever.invoke("What was the message of the last test query?")
+
+
+@pytest.mark.snapshot(
+    ignores=["meta.langchain.request.tool.description"],
+    token="tests.contrib.langchain.test_langchain_community.test_base_tool_invoke",
+)
+def test_base_tool_invoke(langchain_core, request_vcr):
+    """
+    Test that invoking a tool with langchain will
+    result in a 1-span trace with a tool span.
+    """
+    if langchain_core is None:
+        pytest.skip("langchain-core not installed which is required for this test.")
+
+    from math import pi
+
+    from langchain_core.tools import StructuredTool
+
+    def circumference_tool(radius: float) -> float:
+        return float(radius) * 2.0 * pi
+
+    calculator = StructuredTool.from_function(
+        func=circumference_tool,
+        name="Circumference calculator",
+        description="Use this tool when you need to calculate a circumference using the radius of a circle",
+        return_direct=True,
+        response_format="content",
+    )
+
+    calculator.invoke("2")
+
+
+@pytest.mark.asyncio
+@pytest.mark.snapshot(
+    ignores=["meta.langchain.request.tool.description"],
+    token="tests.contrib.langchain.test_langchain_community.test_base_tool_ainvoke",
+)
+async def test_base_tool_ainvoke(langchain_core, request_vcr):
+    """
+    Test that invoking a tool with langchain will
+    result in a 1-span trace with a tool span. Async mode
+    """
+
+    if langchain_core is None:
+        pytest.skip("langchain-core not installed which is required for this test.")
+
+    from math import pi
+
+    from langchain_core.tools import StructuredTool
+
+    def circumference_tool(radius: float) -> float:
+        return float(radius) * 2.0 * pi
+
+    calculator = StructuredTool.from_function(
+        func=circumference_tool,
+        name="Circumference calculator",
+        description="Use this tool when you need to calculate a circumference using the radius of a circle",
+        return_direct=True,
+        response_format="content",
+    )
+
+    await calculator.ainvoke("2")

tests/contrib/langchain/test_langchain_llmobs.py

@@ -172,6 +172,18 @@ def _similarity_search(cls, pinecone, pinecone_vector_store, embedding_model, qu
         LLMObs.disable()
         return mock_tracer.pop_traces()[0]
 
+    @classmethod
+    def _invoke_tool(cls, tool, tool_input, mock_tracer, cassette_name):
+        LLMObs.enable(ml_app=cls.ml_app, integrations_enabled=False, _tracer=mock_tracer)
+        if LANGCHAIN_VERSION > (0, 1):
+            if cassette_name is not None:
+                with get_request_vcr(subdirectory_name=cls.cassette_subdirectory_name).use_cassette(cassette_name):
+                    tool.invoke(tool_input)
+            else:
+                tool.invoke(tool_input)
+        LLMObs.disable()
+        return mock_tracer.pop_traces()[0][0]
+
 
 @pytest.mark.skipif(LANGCHAIN_VERSION >= (0, 1), reason="These tests are for langchain < 0.1.0")
 class TestLLMObsLangchain(BaseTestLLMObsLangchain):
@@ -707,6 +719,46 @@ def test_llmobs_similarity_search(self, langchain_openai, langchain_pinecone, mo
         )
         mock_llmobs_span_writer.enqueue.assert_any_call(expected_span)
 
+    def test_llmobs_base_tool_invoke(self, langchain_core, mock_llmobs_span_writer, mock_tracer):
+        if langchain_core is None:
+            pytest.skip("langchain-core not installed which is required for this test.")
+
+        if langchain_core is None:
+            pytest.skip("langchain-core not installed which is required for this test.")
+
+        from math import pi
+
+        from langchain_core.tools import StructuredTool
+
+        def circumference_tool(radius: float) -> float:
+            return float(radius) * 2.0 * pi
+
+        calculator = StructuredTool.from_function(
+            func=circumference_tool,
+            name="Circumference calculator",
+            description="Use this tool when you need to calculate a circumference using the radius of a circle",
+            return_direct=True,
+            response_format="content",
+        )
+
+        span = self._invoke_tool(
+            tool=calculator,
+            tool_input="2",
+            mock_tracer=mock_tracer,
+            cassette_name=None,
+        )
+        assert mock_llmobs_span_writer.enqueue.call_count == 1
+        mock_llmobs_span_writer.enqueue.assert_called_with(
+            _expected_llmobs_non_llm_span_event(
+                span,
+                span_kind="tool",
+                input_value="2",
+                output_value="12.566370614359172",
+                tags={"ml_app": "langchain_test"},
+                integration="langchain",
+            )
+        )
+
 
 @pytest.mark.skipif(LANGCHAIN_VERSION < (0, 1), reason="These tests are for langchain >= 0.1.0")
 class TestTraceStructureWithLLMIntegrations(SubprocessTestCase):

tests/snapshots/tests.contrib.langchain.test_langchain_community.test_base_tool_ainvoke.json

@@ -0,0 +1,53 @@
+[[
+  {
+    "name": "langchain.request",
+    "service": "",
+    "resource": "langchain_core.tools.method.invoke.Circumference calculator",
+    "trace_id": 0,
+    "span_id": 1,
+    "parent_id": 0,
+    "type": "",
+    "error": 0,
+    "meta": {
+      "_dd.p.dm": "-0",
+      "_dd.p.tid": "66cf6d8d00000000",
+      "langchain.request.input": "2",
+      "langchain.request.tool.description": "Circumference calculator(radius: float) -> float - Use this tool when you need to calculate a circumference using the radius of a circle",
+      "langchain.request.tool.name": "Circumference calculator",
+      "langchain.request.type": "tool",
+      "langchain.response.output": "12.566370614359172",
+      "language": "python",
+      "runtime-id": "fdfa007d5e604a6a880c73034778aa7f"
+    },
+    "metrics": {
+      "_dd.measured": 1,
+      "_dd.top_level": 1,
+      "_dd.tracer_kr": 1.0,
+      "_sampling_priority_v1": 1,
+      "process_id": 73484
+    },
+    "duration": 382000,
+    "start": 1724870029728802000
+  },
+     {
+       "name": "langchain.request",
+       "service": "",
+       "resource": "langchain_core.tools.method.invoke.Circumference calculator",
+       "trace_id": 0,
+       "span_id": 2,
+       "parent_id": 1,
+       "type": "",
+       "error": 0,
+       "meta": {
+         "langchain.request.input": "2",
+         "langchain.request.tool.description": "Circumference calculator(radius: float) -> float - Use this tool when you need to calculate a circumference using the radius of a circle",
+         "langchain.request.tool.name": "Circumference calculator",
+         "langchain.request.type": "tool",
+         "langchain.response.output": "12.566370614359172"
+       },
+       "metrics": {
+         "_dd.measured": 1
+       },
+       "duration": 247000,
+       "start": 1724870029728902000
+     }]]