fix(langchain): ensure llm spans are created for sync cases (#3201)

Co-authored-by: Abhishek Rao <abhishek.rao@observeinc.com>

Author: ronensc

packages/opentelemetry-instrumentation-langchain/opentelemetry/instrumentation/langchain/__init__.py

@@ -3,6 +3,9 @@
 import logging
 from typing import Collection
 
+from opentelemetry import context as context_api
+
+
 from opentelemetry._events import get_event_logger
 from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
 from opentelemetry.instrumentation.langchain.callback_handler import (
@@ -13,7 +16,7 @@
 from opentelemetry.instrumentation.langchain.version import __version__
 from opentelemetry.instrumentation.utils import unwrap
 from opentelemetry.metrics import get_meter
-from opentelemetry.semconv_ai import Meters
+from opentelemetry.semconv_ai import Meters, SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY
 from opentelemetry.trace import get_tracer
 from opentelemetry.trace.propagation import set_span_in_context
 from opentelemetry.trace.propagation.tracecontext import (
@@ -183,8 +186,8 @@ def _uninstrument(self, **kwargs):
 
 
 class _BaseCallbackManagerInitWrapper:
- def __init__(self, callback_manager: "TraceloopCallbackHandler"):
- self._callback_manager = callback_manager
+ def __init__(self, callback_handler: "TraceloopCallbackHandler"):
+ self._callback_handler = callback_handler
 
  def __call__(
  self,
@@ -195,10 +198,14 @@ def __call__(
  ) -> None:
  wrapped(*args, **kwargs)
  for handler in instance.inheritable_handlers:
- if isinstance(handler, type(self._callback_manager)):
+ if isinstance(handler, type(self._callback_handler)):
  break
  else:
- instance.add_handler(self._callback_manager, True)
+ # Add a property to the handler which indicates the CallbackManager instance.
+ # Since the CallbackHandler only propagates context for sync callbacks,
+ # we need a way to determine the type of CallbackManager being wrapped.
+ self._callback_handler._callback_manager = instance
+ instance.add_handler(self._callback_handler, True)
 
 
 # This class wraps a function call to inject tracing information (trace headers) into
@@ -233,4 +240,10 @@ def __call__(
  # Update kwargs to include the modified headers
  kwargs["extra_headers"] = extra_headers
 
+ # In legacy chains like LLMChain, suppressing model instrumentations
+ # within create_llm_span doesn't work, so this should helps as a fallback
+ context_api.attach(
+ context_api.set_value(SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY, True)
+ )
+
  return wrapped(*args, **kwargs)

packages/opentelemetry-instrumentation-langchain/opentelemetry/instrumentation/langchain/callback_handler.py

@@ -5,6 +5,8 @@
 
 from langchain_core.callbacks import (
  BaseCallbackHandler,
+ CallbackManager,
+ AsyncCallbackManager,
 )
 from langchain_core.messages import (
  AIMessage,
@@ -85,19 +87,6 @@ def _extract_class_name_from_serialized(serialized: Optional[dict[str, Any]]) ->
  return ""
 
 
-def _message_type_to_role(message_type: str) -> str:
- if message_type == "human":
- return "user"
- elif message_type == "system":
- return "system"
- elif message_type == "ai":
- return "assistant"
- elif message_type == "tool":
- return "tool"
- else:
- return "unknown"
-
-
 def _sanitize_metadata_value(value: Any) -> Any:
  """Convert metadata values to OpenTelemetry-compatible types."""
  if value is None:
@@ -163,6 +152,7 @@ def __init__(
  self.token_histogram = token_histogram
  self.spans: dict[UUID, SpanHolder] = {}
  self.run_inline = True
+ self._callback_manager: CallbackManager | AsyncCallbackManager = None
 
  @staticmethod
  def _get_name_from_callback(
@@ -192,6 +182,9 @@ def _end_span(self, span: Span, run_id: UUID) -> None:
  if child_span.end_time is None: # avoid warning on ended spans
  child_span.end()
  span.end()
+ token = self.spans[run_id].token
+ if token:
+ context_api.detach(token)
 
  def _create_span(
  self,
@@ -230,13 +223,17 @@ def _create_span(
  else:
  span = self.tracer.start_span(span_name, kind=kind)
 
+ token = None
+ # TODO: make this unconditional once attach/detach works properly with async callbacks.
+ # Currently, it doesn't work due to this - https://github.com/langchain-ai/langchain/issues/31398
+ # As a sidenote, OTel Python users also report similar issues -
+ # https://github.com/open-telemetry/opentelemetry-python/issues/2606
+ if self._callback_manager and not self._callback_manager.is_async:
+ token = context_api.attach(set_span_in_context(span))
+
  _set_span_attribute(span, SpanAttributes.TRACELOOP_WORKFLOW_NAME, workflow_name)
  _set_span_attribute(span, SpanAttributes.TRACELOOP_ENTITY_PATH, entity_path)
 
- token = context_api.attach(
- context_api.set_value(SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY, True)
- )
-
  self.spans[run_id] = SpanHolder(
  span, token, None, [], workflow_name, entity_name, entity_path
  )
@@ -300,6 +297,16 @@ def _create_llm_span(
  _set_span_attribute(span, SpanAttributes.LLM_SYSTEM, vendor)
  _set_span_attribute(span, SpanAttributes.LLM_REQUEST_TYPE, request_type.value)
 
+ # we already have an LLM span by this point,
+ # so skip any downstream instrumentation from here
+ token = context_api.attach(
+ context_api.set_value(SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY, True)
+ )
+
+ self.spans[run_id] = SpanHolder(
+ span, token, None, [], workflow_name, None, entity_path
+ )
+
  return span
 
  @dont_throw
@@ -464,7 +471,7 @@ def on_llm_end(
  "model_name"
  ) or response.llm_output.get("model_id")
  if model_name is not None:
- _set_span_attribute(span, SpanAttributes.LLM_RESPONSE_MODEL, model_name)
+ _set_span_attribute(span, SpanAttributes.LLM_RESPONSE_MODEL, model_name or "unknown")
 
  if self.spans[run_id].request_model is None:
  _set_span_attribute(