fix(langchain): add experimental context propagation for langchain

python/instrumentation/openinference-instrumentation-langchain/examples/mcp/README.md

@@ -0,0 +1 @@
+# LangGraph + MCP composability Example

python/instrumentation/openinference-instrumentation-langchain/examples/mcp/client.py

@@ -0,0 +1,43 @@
+# Create server parameters for stdio connection
+import asyncio
+
+from langchain_mcp_adapters.tools import load_mcp_tools
+from langchain_openai import ChatOpenAI
+from langgraph.prebuilt import create_react_agent
+from mcp import ClientSession, StdioServerParameters
+from mcp.client.stdio import stdio_client
+from phoenix.otel import register
+
+from openinference.instrumentation.langchain import LangChainInstrumentor
+from openinference.instrumentation.mcp import MCPInstrumentor
+
+tracer_provider = register(project_name="langchain-mcp-adapters")
+MCPInstrumentor().instrument(tracer_provider=tracer_provider)
+LangChainInstrumentor().instrument(tracer_provider=tracer_provider, attach_otel_context=True)
+
+
+async def main():
+    model = ChatOpenAI(model="gpt-4")
+
+    server_params = StdioServerParameters(
+        command="python",
+        # Make sure to update to the full absolute path to your math_server.py file
+        args=["./math_server.py"],
+    )
+
+    async with stdio_client(server_params) as (read, write):
+        async with ClientSession(read, write) as session:
+            # Initialize the connection
+            await session.initialize()
+
+            # Get tools
+            tools = await load_mcp_tools(session)
+
+            # Create and run the agent
+            agent = create_react_agent(model, tools)
+            agent_response = await agent.ainvoke({"messages": "what's (3 + 5) x 12?"})
+            print(agent_response)
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

python/instrumentation/openinference-instrumentation-langchain/examples/mcp/math_server.py

@@ -0,0 +1,30 @@
+# math_server.py
+from mcp.server.fastmcp import FastMCP
+from phoenix.otel import register
+
+from openinference.instrumentation.mcp import MCPInstrumentor
+
+tracer_provider = register(project_name="langchain-mcp-adapters")
+MCPInstrumentor().instrument(tracer_provider=tracer_provider)
+
+tracer = tracer_provider.get_tracer("math-mcp-server")
+
+mcp = FastMCP("Math")
+
+
+@mcp.tool()
+@tracer.tool()
+def add(a: int, b: int) -> int:
+    """Add two numbers"""
+    return a + b
+
+
+@mcp.tool()
+@tracer.tool()
+def multiply(a: int, b: int) -> int:
+    """Multiply two numbers"""
+    return a * b
+
+
+if __name__ == "__main__":
+    mcp.run(transport="stdio")

python/instrumentation/openinference-instrumentation-langchain/examples/mcp/requirements.txt

@@ -0,0 +1,6 @@
+langchain-mcp-adapters >=0.0.9
+langchain-mcp-adapters==0.0.9
+langchain-openai==0.3.9
+langgraph==0.3.15
+arize-phoenix-otel==0.9.2
+openinference-instrumentation-mcp==1.1.0

python/instrumentation/openinference-instrumentation-langchain/src/openinference/instrumentation/langchain/__init__.py

@@ -49,6 +49,7 @@ def _instrument(self, **kwargs: Any) -> None:
         self._tracer: Optional[OpenInferenceTracer] = OpenInferenceTracer(
             tracer,
             bool(kwargs.get("separate_trace_from_runtime_context")),
+            bool(kwargs.get("attach_otel_context")),
         )
         self._original_callback_manager_init = langchain_core.callbacks.BaseCallbackManager.__init__
         wrap_function_wrapper(

python/instrumentation/openinference-instrumentation-langchain/src/openinference/instrumentation/langchain/_tracer.py

@@ -4,6 +4,7 @@
 import re
 import time
 import traceback
+from contextvars import Context, Token
 from copy import deepcopy
 from datetime import datetime, timezone
 from enum import Enum
@@ -118,22 +119,28 @@ class OpenInferenceTracer(BaseTracer):
     __slots__ = (
         "_tracer",
         "_separate_trace_from_runtime_context",
+        "_attach_otel_context",
         "_spans_by_run",
     )
 
     def __init__(
         self,
         tracer: trace_api.Tracer,
-        separate_trace_from_runtime_context: bool,
+        separate_trace_from_runtime_context: bool = False,
+        attach_otel_context: bool = False,
         *args: Any,
         **kwargs: Any,
     ) -> None:
         """Initialize the OpenInferenceTracer.
 
         Args:
             tracer (trace_api.Tracer): The OpenTelemetry tracer for creating spans.
-            separate_trace_from_runtime_context (bool): When True, always start a new trace for each
-                span without a parent, isolating it from any existing trace in the runtime context.
+            separate_trace_from_runtime_context (bool, optional): When True, always start a new
+                trace for each span without a parent, isolating it from any existing trace in the
+                runtime context. Defaults to True.
+            attach_otel_context (bool, optional): When True, attach the OpenTelemetry context to the
+                spans. This defaults to False due to LangChain's callback manager design and should
+                be used with caution until a wholistic solution is implemented within LangChain.
             *args (Any): Positional arguments for BaseTracer.
             **kwargs (Any): Keyword arguments for BaseTracer.
         """
@@ -144,7 +151,9 @@ def __init__(
         self.run_map = _DictWithLock[str, Run](self.run_map)
         self._tracer = tracer
         self._separate_trace_from_runtime_context = separate_trace_from_runtime_context
+        self._attach_otel_context = attach_otel_context
         self._spans_by_run: Dict[UUID, Span] = _DictWithLock[UUID, Span]()
+        self._context_tokens_by_run: Dict[UUID, Token] = _DictWithLock[UUID, Token]()
         self._lock = RLock()  # handlers may be run in a thread by langchain
 
     def get_span(self, run_id: UUID) -> Optional[Span]:
@@ -171,15 +180,24 @@ def _start_trace(self, run: Run) -> None:
             start_time=start_time_utc_nano,
         )
 
-        # The following line of code is commented out to serve as a reminder that in a system
-        # of callbacks, attaching the context can be hazardous because there is no guarantee
+        # The following line of code is disabled by default due to LangChain's callback manager design
+        # and should be used with caution until a wholistic solution is implemented.
+        #
+        # Note: attaching the context can be hazardous because there is no guarantee
         # that the context will be detached. An error could happen between callbacks leaving
         # the context attached forever, and all future spans will use it as parent. What's
         # worse is that the error could have also prevented the span from being exported,
         # leaving all future spans as orphans. That is a very bad scenario.
-        # token = context_api.attach(context)
+        token = None
+        if self._attach_otel_context:
+            print("attaching otel context")
+            context = trace_api.set_span_in_context(span)
+            token = context_api.attach(context)
         with self._lock:
             self._spans_by_run[run.id] = span
+            if token:
+                print("attaching token")
+                self._context_tokens_by_run[run.id] = token
 
     @audit_timing  # type: ignore
     def _end_trace(self, run: Run) -> None:
@@ -197,6 +215,11 @@ def _end_trace(self, run: Run) -> None:
             end_time_utc_nano = _as_utc_nano(run.end_time) if run.end_time else None
             span.end(end_time=end_time_utc_nano)
 
+        # Detach the context token if it exists
+        token = self._context_tokens_by_run.pop(run.id, None)
+        if token:
+            context_api.detach(token)
+
     def _persist_run(self, run: Run) -> None:
         pass