feat(langchain): add support for tracing RunnableLambda (#15477)

## Description

Adds APM and LLM Observability support for tracing `RunnableLambda`
operations. Additionally, makes sure span linking happens automatically
as well.

## Testing

Added unit tests, verified locally as well.

## Risks

None.

Author: sabrenner

ddtrace/contrib/internal/langchain/patch.py

@@ -612,6 +612,78 @@ def patched_vectorstore_init_subclass(func, instance, args, kwargs):
         log.warning("Unable to patch LangChain VectorStore class %s", str(cls))
 
 
+def traced_runnable_lambda_operation(is_batch: bool = False):
+    @with_traced_module
+    def _traced_runnable_lambda_impl(langchain_core, pin, func, instance, args, kwargs):
+        integration: LangChainIntegration = langchain_core._datadog_integration
+
+        instance_name = getattr(instance, "name", None)
+        default_name = f"{instance.__class__.__name__}.{func.__name__}"
+        if is_batch:
+            span_name = f"{instance_name}_batch" if instance_name else default_name
+        else:
+            span_name = instance_name or default_name
+
+        span = integration.trace(
+            pin,
+            span_name,
+            submit_to_llmobs=True,
+            instance=instance,
+        )
+
+        integration.record_instance(instance, span)
+
+        result = None
+
+        try:
+            result = func(*args, **kwargs)
+            return result
+        except Exception:
+            span.set_exc_info(*sys.exc_info())
+            raise
+        finally:
+            integration.llmobs_set_tags(span, args=args, kwargs=kwargs, response=result, operation="runnable_lambda")
+            span.finish()
+
+    return _traced_runnable_lambda_impl
+
+
+def traced_runnable_lambda_operation_async(is_batch: bool = False):
+    @with_traced_module
+    async def _traced_runnable_lambda_impl(langchain_core, pin, func, instance, args, kwargs):
+        integration: LangChainIntegration = langchain_core._datadog_integration
+
+        instance_name = getattr(instance, "name", None)
+        default_name = f"{instance.__class__.__name__}.{func.__name__}"
+        if is_batch:
+            span_name = f"{instance_name}_batch" if instance_name else default_name
+        else:
+            span_name = instance_name or default_name
+
+        span = integration.trace(
+            pin,
+            span_name,
+            submit_to_llmobs=True,
+            instance=instance,
+        )
+
+        integration.record_instance(instance, span)
+
+        result = None
+
+        try:
+            result = await func(*args, **kwargs)
+            return result
+        except Exception:
+            span.set_exc_info(*sys.exc_info())
+            raise
+        finally:
+            integration.llmobs_set_tags(span, args=args, kwargs=kwargs, response=result, operation="runnable_lambda")
+            span.finish()
+
+    return _traced_runnable_lambda_impl
+
+
 def patch():
     if getattr(langchain_core, "_datadog_patch", False):
         return
@@ -626,6 +698,7 @@ def patch():
     from langchain_core.language_models.chat_models import BaseChatModel
     from langchain_core.language_models.llms import BaseLLM
     from langchain_core.prompts.base import BasePromptTemplate
+    from langchain_core.runnables.base import RunnableLambda
     from langchain_core.runnables.base import RunnableSequence
     from langchain_core.tools import BaseTool
     from langchain_core.vectorstores import VectorStore
@@ -651,6 +724,11 @@ def patch():
     wrap(RunnableSequence, "stream", traced_chain_stream(langchain_core))
     wrap(RunnableSequence, "astream", traced_chain_stream(langchain_core))
 
+    wrap(RunnableLambda, "invoke", traced_runnable_lambda_operation(is_batch=False)(langchain_core))
+    wrap(RunnableLambda, "ainvoke", traced_runnable_lambda_operation_async(is_batch=False)(langchain_core))
+    wrap(RunnableLambda, "batch", traced_runnable_lambda_operation(is_batch=True)(langchain_core))
+    wrap(RunnableLambda, "abatch", traced_runnable_lambda_operation_async(is_batch=True)(langchain_core))
+
     wrap(BasePromptTemplate, "invoke", patched_base_prompt_template_invoke(langchain_core))
     wrap(BasePromptTemplate, "ainvoke", patched_base_prompt_template_ainvoke(langchain_core))
 
@@ -683,6 +761,10 @@ def unpatch():
     unwrap(langchain_core.runnables.base.RunnableSequence, "abatch")
     unwrap(langchain_core.runnables.base.RunnableSequence, "stream")
     unwrap(langchain_core.runnables.base.RunnableSequence, "astream")
+    unwrap(langchain_core.runnables.base.RunnableLambda, "invoke")
+    unwrap(langchain_core.runnables.base.RunnableLambda, "ainvoke")
+    unwrap(langchain_core.runnables.base.RunnableLambda, "batch")
+    unwrap(langchain_core.runnables.base.RunnableLambda, "abatch")
     unwrap(langchain_core.language_models.chat_models.BaseChatModel, "stream")
     unwrap(langchain_core.language_models.chat_models.BaseChatModel, "astream")
     unwrap(langchain_core.language_models.llms.BaseLLM, "stream")

ddtrace/llmobs/_integrations/langchain.py

@@ -71,7 +71,7 @@
     "system": "system",
 }
 
-SUPPORTED_OPERATIONS = ["llm", "chat", "chain", "embedding", "retrieval", "tool"]
+SUPPORTED_OPERATIONS = ["llm", "chat", "chain", "embedding", "retrieval", "tool", "runnable_lambda"]
 LANGCHAIN_BASE_URL_FIELDS = [
     "api_base",
     "api_host",
@@ -165,7 +165,7 @@ def _llmobs_set_tags(
         args: List[Any],
         kwargs: Dict[str, Any],
         response: Optional[Any] = None,
-        operation: str = "",  # oneof "llm","chat","chain","embedding","retrieval","tool"
+        operation: str = "",  # oneof SUPPORTED_OPERATIONS
     ) -> None:
         """Sets meta tags and metrics for span events to be sent to LLMObs."""
         if not self.llmobs_enabled:
@@ -211,6 +211,8 @@ def _llmobs_set_tags(
             self._llmobs_set_meta_tags_from_similarity_search(span, args, kwargs, response, is_workflow=is_workflow)
         elif operation == "tool":
             self._llmobs_set_meta_tags_from_tool(span, tool_inputs=kwargs, tool_output=response)
+        elif operation == "runnable_lambda":
+            self._llmobs_set_meta_tags_from_runnable_lambda(span, args, kwargs, response)
 
     def _set_links(self, span: Span) -> None:
         """
@@ -763,6 +765,19 @@ def _llmobs_set_meta_tags_from_tool(self, span: Span, tool_inputs: Dict[str, Any
             }
         )
 
+    def _llmobs_set_meta_tags_from_runnable_lambda(
+        self, span: Span, args: List[Any], kwargs: Dict[str, Any], response: Any
+    ) -> None:
+        inputs = get_argument_value(args, kwargs, 0, "inputs")
+
+        span._set_ctx_items(
+            {
+                SPAN_KIND: "task",
+                INPUT_VALUE: safe_json(inputs),
+                OUTPUT_VALUE: safe_json(response),
+            }
+        )
+
     def _set_base_span_tags(
         self,
         span: Span,

releasenotes/notes/langchain-runnable-lambda-cec8e99164d081d1.yaml

@@ -0,0 +1,4 @@
+---
+features:
+  - |
+    langchain: Adds support for tracing ``RunnableLambda`` instances.

tests/contrib/langchain/test_langchain.py

@@ -553,3 +553,43 @@ def test_streamed_chat_model_with_no_output(langchain_openai, openai_url):
     except Exception as e:
         if not isinstance(e, APITimeoutError):
             assert False, f"Expected APITimeoutError, got {e}"
+
+
+@pytest.mark.snapshot(ignores=IGNORE_FIELDS)
+def test_runnable_lambda_invoke(langchain_core):
+    def add(inputs: dict) -> int:
+        return inputs["a"] + inputs["b"]
+
+    runnable_lambda = langchain_core.runnables.RunnableLambda(add)
+    result = runnable_lambda.invoke(dict(a=1, b=2))
+    assert result == 3
+
+
+@pytest.mark.snapshot(ignores=IGNORE_FIELDS)
+async def test_runnable_lambda_ainvoke(langchain_core):
+    async def add(inputs: dict) -> int:
+        return inputs["a"] + inputs["b"]
+
+    runnable_lambda = langchain_core.runnables.RunnableLambda(add)
+    result = await runnable_lambda.ainvoke(dict(a=1, b=2))
+    assert result == 3
+
+
+@pytest.mark.snapshot(ignores=IGNORE_FIELDS)
+def test_runnable_lambda_batch(langchain_core):
+    def add(inputs: dict) -> int:
+        return inputs["a"] + inputs["b"]
+
+    runnable_lambda = langchain_core.runnables.RunnableLambda(add)
+    result = runnable_lambda.batch([dict(a=1, b=2), dict(a=3, b=4), dict(a=5, b=6)])
+    assert result == [3, 7, 11]
+
+
+@pytest.mark.snapshot(ignores=IGNORE_FIELDS)
+async def test_runnable_lambda_abatch(langchain_core):
+    async def add(inputs: dict) -> int:
+        return inputs["a"] + inputs["b"]
+
+    runnable_lambda = langchain_core.runnables.RunnableLambda(add)
+    result = await runnable_lambda.abatch([dict(a=1, b=2), dict(a=3, b=4), dict(a=5, b=6)])
+    assert result == [3, 7, 11]