feat(langchain): modernize callback for v1 (tokens, provider, langgraph, content blocks)

- _extract_token_info reads AIMessage.usage_metadata first (the v1-standardized
  source) before llm_output, and captures input/output token details
  (cache_read, cache_creation, reasoning) under step metadata token_details.
- Use metadata["ls_provider"] as the primary provider source (mapped to
  Openlayer provider names), falling back to the _type map and LiteLLM prefixes.
- Name chain steps from metadata["langgraph_node"] and auto-map
  metadata["thread_id"] to session_id (opt-out via map_thread_id_to_session),
  never clobbering an explicitly set session.
- _message_to_dict normalizes v1 content blocks: text joined into content,
  non-text blocks preserved under content_blocks.
- Drop the dead langchain.schema / langchain.callbacks.base import fallback
  (removed in LangChain v1); import from langchain_core only.

Implements the medium/low items of OPEN-11315.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Author: viniciusdsmello

src/openlayer/lib/integrations/langchain_callback.py

@@ -7,15 +7,14 @@
 from uuid import UUID
 
 try:
-    try:
-        from langchain_core import messages as langchain_schema
-        from langchain_core.callbacks.base import (
-            AsyncCallbackHandler,
-            BaseCallbackHandler,
-        )
-    except ImportError:
-        from langchain import schema as langchain_schema
-        from langchain.callbacks.base import AsyncCallbackHandler, BaseCallbackHandler
+    # langchain-core is sufficient for this handler. The legacy
+    # ``langchain.schema`` / ``langchain.callbacks.base`` import paths were
+    # removed in LangChain v1, so we import from ``langchain_core`` directly.
+    from langchain_core import messages as langchain_schema
+    from langchain_core.callbacks.base import (
+        AsyncCallbackHandler,
+        BaseCallbackHandler,
+    )
 
     HAVE_LANGCHAIN = True
 except ImportError:
@@ -33,6 +32,34 @@
     "amazon_bedrock_converse_chat": "Bedrock",
 }
 
+# LangChain v1 injects a standardized ``metadata["ls_provider"]`` (e.g.
+# "openai", "anthropic", "google_genai", "groq"...). This maps those values to
+# Openlayer's canonical provider names. Values without an explicit entry fall
+# back to a title-cased version of the ls_provider string.
+LS_PROVIDER_TO_OPENLAYER_MAP = {
+    "openai": "OpenAI",
+    "azure": "Azure",
+    "azure_openai": "Azure",
+    "anthropic": "Anthropic",
+    "google": "Google",
+    "google_genai": "Google",
+    "google_vertexai": "Google",
+    "vertexai": "Google",
+    "cohere": "Cohere",
+    "mistralai": "Mistral",
+    "mistral": "Mistral",
+    "bedrock": "Bedrock",
+    "amazon_bedrock": "Bedrock",
+    "ollama": "Ollama",
+    "huggingface": "Hugging Face",
+    "replicate": "Replicate",
+    "together": "Together AI",
+    "groq": "Groq",
+    "deepseek": "DeepSeek",
+    "fireworks": "Fireworks AI",
+    "perplexity": "Perplexity",
+}
+
 # LiteLLM model prefixes to provider names.
 # When models are accessed via a LiteLLM proxy (e.g. "gemini/gemini-2.5-flash"),
 # the LangChain _type is "openai-chat" which incorrectly maps to "OpenAI".
@@ -73,7 +100,7 @@ def __init__(self, **kwargs: Any) -> None:
         if not HAVE_LANGCHAIN:
             raise ImportError(
                 "LangChain library is not installed. Please install it with: pip "
-                "install langchain"
+                "install langchain-core"
             )
         super().__init__()
         self.metadata: Dict[str, Any] = kwargs or {}
@@ -85,6 +112,9 @@ def __init__(self, **kwargs: Any) -> None:
         self._inference_id = kwargs.get("inference_id")
         # Extract metadata_transformer from kwargs if provided
         self._metadata_transformer = kwargs.get("metadata_transformer")
+        # Opt-out flag: auto-map LangGraph metadata["thread_id"] to the trace's
+        # session_id when the user has not explicitly provided one.
+        self._map_thread_id_to_session = kwargs.get("map_thread_id_to_session", True)
 
     def _start_step(
         self,
@@ -374,7 +404,14 @@ def _message_to_dict(
         elif message_type == "system":
             role = "system"
 
-        result: Dict[str, Any] = {"role": role, "content": str(message.content)}
+        content_str, content_blocks = self._normalize_content(message.content)
+        result: Dict[str, Any] = {"role": role, "content": content_str}
+
+        # With ``LC_OUTPUT_VERSION=v1`` the content can be a list of typed
+        # blocks. Preserve any non-text blocks (reasoning, tool_call, image...)
+        # structurally rather than stringifying the whole list.
+        if content_blocks:
+            result["content_blocks"] = content_blocks
 
         # Preserve tool calls on assistant messages (e.g. AIMessage.tool_calls).
         # Without this, agent turns that only emit tool calls would drop them
@@ -386,6 +423,41 @@ def _message_to_dict(
 
         return result
 
+    @staticmethod
+    def _normalize_content(content: Any) -> tuple:
+        """Normalize message content into a (text, non_text_blocks) pair.
+
+        Handles the three shapes ``message.content`` can take:
+          * a plain string (returned unchanged, no blocks);
+          * a list of strings (joined into the text);
+          * a list of typed blocks (dicts with a ``type`` key) as produced
+            under ``LC_OUTPUT_VERSION=v1`` -- text is extracted from
+            ``text``-type blocks while non-text blocks are preserved
+            structurally.
+
+        Defensive: never raises on unexpected shapes.
+        """
+        if isinstance(content, str):
+            return content, []
+
+        if isinstance(content, list):
+            text_parts: List[str] = []
+            non_text_blocks: List[Any] = []
+            for block in content:
+                if isinstance(block, str):
+                    text_parts.append(block)
+                elif isinstance(block, dict):
+                    if block.get("type") == "text" and "text" in block:
+                        text_parts.append(str(block.get("text", "")))
+                    else:
+                        non_text_blocks.append(block)
+                else:
+                    non_text_blocks.append(block)
+            return "".join(text_parts), non_text_blocks
+
+        # Fallback for any other shape.
+        return str(content), []
+
     def _messages_to_prompt_format(
         self, messages: List[List["langchain_schema.BaseMessage"]]
     ) -> List[Dict[str, str]]:
@@ -409,9 +481,19 @@ def _extract_model_info(
         invocation_params = invocation_params or {}
         metadata = metadata or {}
 
-        provider = invocation_params.get("_type")
-        if provider in LANGCHAIN_TO_OPENLAYER_PROVIDER_MAP:
-            provider = LANGCHAIN_TO_OPENLAYER_PROVIDER_MAP[provider]
+        # Provider resolution order:
+        #   1. metadata["ls_provider"] (standardized by LangChain v1)
+        #   2. invocation_params["_type"] via the legacy _type map
+        #   3. LiteLLM model prefix (handled below)
+        ls_provider = metadata.get("ls_provider")
+        if ls_provider:
+            provider = LS_PROVIDER_TO_OPENLAYER_MAP.get(
+                ls_provider, str(ls_provider).replace("_", " ").title()
+            )
+        else:
+            provider = invocation_params.get("_type")
+            if provider in LANGCHAIN_TO_OPENLAYER_PROVIDER_MAP:
+                provider = LANGCHAIN_TO_OPENLAYER_PROVIDER_MAP[provider]
 
         model = (
             invocation_params.get("model_name")
@@ -443,15 +525,49 @@ def _extract_model_info(
     def _extract_token_info(
         self, response: "langchain_schema.LLMResult"
     ) -> Dict[str, Any]:
-        """Extract token information generically from LLM response."""
-        llm_output = response.llm_output or {}
+        """Extract token information generically from LLM response.
+
+        In LangChain v1, ``AIMessage.usage_metadata`` is the guaranteed,
+        standardized token source, so it is read FIRST. The provider-specific
+        ``llm_output`` / ``generation_info`` paths (Ollama ``prompt_eval_count``,
+        Google ``usage_metadata`` in generation_info) remain as fallbacks so
+        nothing regresses.
+
+        When ``usage_metadata`` carries ``input_token_details`` /
+        ``output_token_details`` (e.g. ``cache_creation``, ``cache_read``,
+        ``reasoning``, ``audio``), those are surfaced under a ``token_details``
+        key so cost is accurate for prompt-caching / reasoning models.
+        """
+        token_usage: Dict[str, Any] = {}
+        token_details: Dict[str, Any] = {}
 
-        # Try standard token_usage location first
-        token_usage = (
-            llm_output.get("token_usage") or llm_output.get("estimatedTokens") or {}
-        )
+        # 1. usage_metadata on the message (standardized, v1-first).
+        if response.generations:
+            gen = response.generations[0][0]
+            usage = getattr(getattr(gen, "message", None), "usage_metadata", None)
+            if usage:
+                token_usage = {
+                    "prompt_tokens": usage.get("input_tokens", 0),
+                    "completion_tokens": usage.get("output_tokens", 0),
+                    "total_tokens": usage.get("total_tokens", 0),
+                }
+                input_details = usage.get("input_token_details")
+                if input_details:
+                    token_details["input_token_details"] = dict(input_details)
+                output_details = usage.get("output_token_details")
+                if output_details:
+                    token_details["output_token_details"] = dict(output_details)
+
+        # 2. Fall back to llm_output for providers that surface it there.
+        if not token_usage:
+            llm_output = response.llm_output or {}
+            token_usage = (
+                llm_output.get("token_usage")
+                or llm_output.get("estimatedTokens")
+                or {}
+            )
 
-        # Fallback to generation info for providers like Ollama/Google
+        # 3. Fall back to generation_info for providers like Ollama/Google.
         if not token_usage and response.generations:
             gen = response.generations[0][0]
             generation_info = gen.generation_info or {}
@@ -473,21 +589,15 @@ def _extract_token_info(
                     "completion_tokens": usage.get("candidates_token_count", 0),
                     "total_tokens": usage.get("total_token_count", 0),
                 }
-            # AWS Bedrock / newer LangChain style - usage_metadata on the message
-            elif hasattr(gen, "message") and hasattr(gen.message, "usage_metadata"):
-                usage = gen.message.usage_metadata
-                if usage:
-                    token_usage = {
-                        "prompt_tokens": usage.get("input_tokens", 0),
-                        "completion_tokens": usage.get("output_tokens", 0),
-                        "total_tokens": usage.get("total_tokens", 0),
-                    }
 
-        return {
+        result = {
             "prompt_tokens": token_usage.get("prompt_tokens", 0),
             "completion_tokens": token_usage.get("completion_tokens", 0),
             "tokens": token_usage.get("total_tokens", 0),
         }
+        if token_details:
+            result["token_details"] = token_details
+        return result
 
     def _extract_output(self, response: "langchain_schema.LLMResult") -> str:
         """Extract output text from LLM response.
@@ -570,6 +680,7 @@ def _handle_llm_start(
             metadata={"tags": tags} if tags else None,
             **model_info,
         )
+        self._apply_thread_id_session(run_id, metadata)
 
     def _handle_chat_model_start(
         self,
@@ -603,6 +714,7 @@ def _handle_chat_model_start(
             metadata={"tags": tags} if tags else None,
             **model_info,
         )
+        self._apply_thread_id_session(run_id, metadata)
 
     def _handle_llm_end(
         self,
@@ -629,6 +741,13 @@ def _handle_llm_end(
         ):
             token_info = self._extract_token_info(response)
 
+        # ChatCompletionStep has no dedicated field for fine-grained token
+        # details (cache_read/creation, reasoning, audio...), so surface them
+        # under step metadata where cost-relevant breakdowns can be read.
+        token_details = token_info.pop("token_details", None)
+        if token_details:
+            step.metadata = {**step.metadata, "token_details": token_details}
+
         self._end_step(
             run_id=run_id,
             parent_run_id=parent_run_id,
@@ -663,8 +782,15 @@ def _handle_chain_start(
         # Extract chain name from serialized data or use provided name
         # Handle case where serialized can be None
         serialized = serialized or {}
+        metadata = metadata or {}
+
+        # LangGraph injects ``metadata["langgraph_node"]`` for each graph node.
+        # Prefer it so graph nodes are identifiable in the trace rather than
+        # showing the generic chain/runnable name.
+        langgraph_node = metadata.get("langgraph_node")
         chain_name = (
-            name
+            langgraph_node
+            or name
             or (serialized.get("id", [])[-1] if serialized.get("id") else None)
             or "Chain"
         )
@@ -687,10 +813,11 @@ def _handle_chain_start(
             metadata={
                 "tags": tags,
                 "serialized": serialized,
-                **(metadata or {}),
+                **metadata,
                 **kwargs,
             },
         )
+        self._apply_thread_id_session(run_id, metadata)
 
     def _handle_chain_end(
         self,
@@ -1005,6 +1132,44 @@ def _handle_llm_new_token(self, token: str, **kwargs: Any) -> Any:
                         step.log(**token_info)
         return
 
+    def _apply_thread_id_session(
+        self, run_id: UUID, metadata: Optional[Dict[str, Any]]
+    ) -> None:
+        """Map LangGraph ``metadata["thread_id"]`` to the trace's session_id.
+
+        Any LangGraph app with a checkpointer injects ``thread_id`` into
+        callback metadata. When ``map_thread_id_to_session`` is enabled (the
+        default) and the user has not explicitly set a session via
+        ``UserSessionContext``, the thread_id is written to the trace's
+        ``session_id`` metadata so it is promoted to the ``session_id`` column
+        by ``post_process_trace``.
+
+        An explicitly-set session_id is never clobbered: ``post_process_trace``
+        re-applies the ``UserSessionContext`` value after merging trace
+        metadata, and this method additionally skips writing when a context
+        session_id is present.
+        """
+        if not self._map_thread_id_to_session or not metadata:
+            return
+
+        thread_id = metadata.get("thread_id")
+        if not thread_id:
+            return
+
+        # Do not clobber an explicitly-provided session_id.
+        if tracer.UserSessionContext.get_session_id() is not None:
+            return
+
+        target_trace = tracer.get_current_trace() or self._find_owning_trace(run_id)
+        if target_trace is None:
+            return
+
+        existing = (target_trace.metadata or {}).get("session_id")
+        if existing:
+            return
+
+        target_trace.update_metadata(session_id=str(thread_id))
+
 
 class OpenlayerHandler(OpenlayerHandlerMixin, BaseCallbackHandlerClass):  # type: ignore[misc]
     """LangChain callback handler that logs to Openlayer."""
@@ -1020,13 +1185,15 @@ def __init__(
         metadata_transformer: Optional[
             Callable[[Dict[str, Any]], Dict[str, Any]]
         ] = None,
+        map_thread_id_to_session: bool = True,
         **kwargs: Any,
     ) -> None:
         # Add both inference_id and metadata_transformer to kwargs so they get passed to mixin
         if inference_id is not None:
             kwargs["inference_id"] = inference_id
         if metadata_transformer is not None:
             kwargs["metadata_transformer"] = metadata_transformer
+        kwargs["map_thread_id_to_session"] = map_thread_id_to_session
         super().__init__(**kwargs)
         # Store the ignore flags as instance variables
         self._ignore_llm = ignore_llm
@@ -1191,13 +1358,15 @@ def __init__(
         metadata_transformer: Optional[
             Callable[[Dict[str, Any]], Dict[str, Any]]
         ] = None,
+        map_thread_id_to_session: bool = True,
         **kwargs: Any,
     ) -> None:
         # Add both inference_id and metadata_transformer to kwargs so they get passed to mixin
         if inference_id is not None:
             kwargs["inference_id"] = inference_id
         if metadata_transformer is not None:
             kwargs["metadata_transformer"] = metadata_transformer
+        kwargs["map_thread_id_to_session"] = map_thread_id_to_session
         super().__init__(**kwargs)
         # Store the ignore flags as instance variables
         self._ignore_llm = ignore_llm