langchain-core 0.3.79__py3-none-any.whl → 1.0.0__py3-none-any.whl

langchain_core/messages/ai.py

@@ -3,57 +3,52 @@
 import json
 import logging
 import operator
-from typing import Any, Literal, Optional, Union, cast
+from collections.abc import Sequence
+from typing import Any, Literal, cast, overload
 
 from pydantic import model_validator
 from typing_extensions import NotRequired, Self, TypedDict, override
 
+from langchain_core.messages import content as types
 from langchain_core.messages.base import (
     BaseMessage,
     BaseMessageChunk,
+    _extract_reasoning_from_additional_kwargs,
     merge_content,
 )
+from langchain_core.messages.content import InvalidToolCall
 from langchain_core.messages.tool import (
-    InvalidToolCall,
     ToolCall,
     ToolCallChunk,
     default_tool_chunk_parser,
     default_tool_parser,
 )
-from langchain_core.messages.tool import (
-    invalid_tool_call as create_invalid_tool_call,
-)
-from langchain_core.messages.tool import (
-    tool_call as create_tool_call,
-)
-from langchain_core.messages.tool import (
-    tool_call_chunk as create_tool_call_chunk,
-)
+from langchain_core.messages.tool import invalid_tool_call as create_invalid_tool_call
+from langchain_core.messages.tool import tool_call as create_tool_call
+from langchain_core.messages.tool import tool_call_chunk as create_tool_call_chunk
 from langchain_core.utils._merge import merge_dicts, merge_lists
 from langchain_core.utils.json import parse_partial_json
 from langchain_core.utils.usage import _dict_int_op
+from langchain_core.utils.utils import LC_AUTO_PREFIX, LC_ID_PREFIX
 
 logger = logging.getLogger(__name__)
 
 
-_LC_ID_PREFIX = "run-"
-
-
 class InputTokenDetails(TypedDict, total=False):
     """Breakdown of input token counts.
 
     Does *not* need to sum to full input token count. Does *not* need to have all keys.
 
     Example:
-        .. code-block:: python
-
-            {
-                "audio": 10,
-                "cache_creation": 200,
-                "cache_read": 100,
-            }
+        ```python
+        {
+            "audio": 10,
+            "cache_creation": 200,
+            "cache_read": 100,
+        }
+        ```
 
-    .. versionadded:: 0.3.9
+    !!! version-added "Added in version 0.3.9"
 
     May also hold extra provider-specific keys.
 
@@ -81,14 +76,14 @@ class OutputTokenDetails(TypedDict, total=False):
     Does *not* need to sum to full output token count. Does *not* need to have all keys.
 
     Example:
-        .. code-block:: python
-
-            {
-                "audio": 10,
-                "reasoning": 200,
-            }
+        ```python
+        {
+            "audio": 10,
+            "reasoning": 200,
+        }
+        ```
 
-    .. versionadded:: 0.3.9
+    !!! version-added "Added in version 0.3.9"
 
     """
 
@@ -109,26 +104,25 @@ class UsageMetadata(TypedDict):
     This is a standard representation of token usage that is consistent across models.
 
     Example:
-        .. code-block:: python
-
-            {
-                "input_tokens": 350,
-                "output_tokens": 240,
-                "total_tokens": 590,
-                "input_token_details": {
-                    "audio": 10,
-                    "cache_creation": 200,
-                    "cache_read": 100,
-                },
-                "output_token_details": {
-                    "audio": 10,
-                    "reasoning": 200,
-                },
-            }
-
-    .. versionchanged:: 0.3.9
+        ```python
+        {
+            "input_tokens": 350,
+            "output_tokens": 240,
+            "total_tokens": 590,
+            "input_token_details": {
+                "audio": 10,
+                "cache_creation": 200,
+                "cache_read": 100,
+            },
+            "output_token_details": {
+                "audio": 10,
+                "reasoning": 200,
+            },
+        }
+        ```
 
-        Added ``input_token_details`` and ``output_token_details``.
+    !!! warning "Behavior changed in 0.3.9"
+        Added `input_token_details` and `output_token_details`.
 
     """
 
@@ -154,56 +148,143 @@ class UsageMetadata(TypedDict):
 class AIMessage(BaseMessage):
     """Message from an AI.
 
-    AIMessage is returned from a chat model as a response to a prompt.
+    An `AIMessage` is returned from a chat model as a response to a prompt.
 
     This message represents the output of the model and consists of both
-    the raw output as returned by the model together standardized fields
+    the raw output as returned by the model and standardized fields
     (e.g., tool calls, usage metadata) added by the LangChain framework.
 
     """
 
-    example: bool = False
-    """Use to denote that a message is part of an example conversation.
-
-    At the moment, this is ignored by most models. Usage is discouraged.
-
-    """
-
     tool_calls: list[ToolCall] = []
-    """If provided, tool calls associated with the message."""
+    """If present, tool calls associated with the message."""
     invalid_tool_calls: list[InvalidToolCall] = []
-    """If provided, tool calls with parsing errors associated with the message."""
-    usage_metadata: Optional[UsageMetadata] = None
-    """If provided, usage metadata for a message, such as token counts.
+    """If present, tool calls with parsing errors associated with the message."""
+    usage_metadata: UsageMetadata | None = None
+    """If present, usage metadata for a message, such as token counts.
 
     This is a standard representation of token usage that is consistent across models.
-
     """
 
     type: Literal["ai"] = "ai"
-    """The type of the message (used for deserialization). Defaults to ``'ai'``."""
+    """The type of the message (used for deserialization)."""
+
+    @overload
+    def __init__(
+        self,
+        content: str | list[str | dict],
+        **kwargs: Any,
+    ) -> None: ...
 
+    @overload
     def __init__(
         self,
-        content: Union[str, list[Union[str, dict]]],
+        content: str | list[str | dict] | None = None,
+        content_blocks: list[types.ContentBlock] | None = None,
+        **kwargs: Any,
+    ) -> None: ...
+
+    def __init__(
+        self,
+        content: str | list[str | dict] | None = None,
+        content_blocks: list[types.ContentBlock] | None = None,
         **kwargs: Any,
     ) -> None:
-        """Initialize ``AIMessage``.
+        """Initialize an `AIMessage`.
+
+        Specify `content` as positional arg or `content_blocks` for typing.
 
         Args:
             content: The content of the message.
-            kwargs: Additional arguments to pass to the parent class.
+            content_blocks: Typed standard content.
+            **kwargs: Additional arguments to pass to the parent class.
         """
-        super().__init__(content=content, **kwargs)
+        if content_blocks is not None:
+            # If there are tool calls in content_blocks, but not in tool_calls, add them
+            content_tool_calls = [
+                block for block in content_blocks if block.get("type") == "tool_call"
+            ]
+            if content_tool_calls and "tool_calls" not in kwargs:
+                kwargs["tool_calls"] = content_tool_calls
+
+            super().__init__(
+                content=cast("str | list[str | dict]", content_blocks),
+                **kwargs,
+            )
+        else:
+            super().__init__(content=content, **kwargs)
 
     @property
     def lc_attributes(self) -> dict:
-        """Attrs to be serialized even if they are derived from other init args."""
+        """Attributes to be serialized.
+
+        Includes all attributes, even if they are derived from other initialization
+        arguments.
+        """
         return {
             "tool_calls": self.tool_calls,
             "invalid_tool_calls": self.invalid_tool_calls,
         }
 
+    @property
+    def content_blocks(self) -> list[types.ContentBlock]:
+        """Return standard, typed `ContentBlock` dicts from the message.
+
+        If the message has a known model provider, use the provider-specific translator
+        first before falling back to best-effort parsing. For details, see the property
+        on `BaseMessage`.
+        """
+        if self.response_metadata.get("output_version") == "v1":
+            return cast("list[types.ContentBlock]", self.content)
+
+        model_provider = self.response_metadata.get("model_provider")
+        if model_provider:
+            from langchain_core.messages.block_translators import (  # noqa: PLC0415
+                get_translator,
+            )
+
+            translator = get_translator(model_provider)
+            if translator:
+                try:
+                    return translator["translate_content"](self)
+                except NotImplementedError:
+                    pass
+
+        # Otherwise, use best-effort parsing
+        blocks = super().content_blocks
+
+        if self.tool_calls:
+            # Add from tool_calls if missing from content
+            content_tool_call_ids = {
+                block.get("id")
+                for block in self.content
+                if isinstance(block, dict) and block.get("type") == "tool_call"
+            }
+            for tool_call in self.tool_calls:
+                if (id_ := tool_call.get("id")) and id_ not in content_tool_call_ids:
+                    tool_call_block: types.ToolCall = {
+                        "type": "tool_call",
+                        "id": id_,
+                        "name": tool_call["name"],
+                        "args": tool_call["args"],
+                    }
+                    if "index" in tool_call:
+                        tool_call_block["index"] = tool_call["index"]  # type: ignore[typeddict-item]
+                    if "extras" in tool_call:
+                        tool_call_block["extras"] = tool_call["extras"]  # type: ignore[typeddict-item]
+                    blocks.append(tool_call_block)
+
+        # Best-effort reasoning extraction from additional_kwargs
+        # Only add reasoning if not already present
+        # Insert before all other blocks to keep reasoning at the start
+        has_reasoning = any(block.get("type") == "reasoning" for block in blocks)
+        if not has_reasoning and (
+            reasoning_block := _extract_reasoning_from_additional_kwargs(self)
+        ):
+            blocks.insert(0, reasoning_block)
+
+        return blocks
+
     # TODO: remove this logic if possible, reducing breaking nature of changes
     @model_validator(mode="before")
     @classmethod
@@ -232,7 +313,9 @@ class AIMessage(BaseMessage):
         # Ensure "type" is properly set on all tool call-like dicts.
         if tool_calls := values.get("tool_calls"):
             values["tool_calls"] = [
-                create_tool_call(**{k: v for k, v in tc.items() if k != "type"})
+                create_tool_call(
+                    **{k: v for k, v in tc.items() if k not in ("type", "extras")}
+                )
                 for tc in tool_calls
             ]
         if invalid_tool_calls := values.get("invalid_tool_calls"):
@@ -251,11 +334,10 @@ class AIMessage(BaseMessage):
 
     @override
     def pretty_repr(self, html: bool = False) -> str:
-        """Return a pretty representation of the message.
+        """Return a pretty representation of the message for display.
 
         Args:
             html: Whether to return an HTML-formatted string.
-                 Defaults to False.
 
         Returns:
             A pretty representation of the message.
@@ -264,7 +346,7 @@ class AIMessage(BaseMessage):
         base = super().pretty_repr(html=html)
         lines = []
 
-        def _format_tool_args(tc: Union[ToolCall, InvalidToolCall]) -> list[str]:
+        def _format_tool_args(tc: ToolCall | InvalidToolCall) -> list[str]:
             lines = [
                 f"  {tc.get('name', 'Tool')} ({tc.get('id')})",
                 f" Call ID: {tc.get('id')}",
@@ -292,29 +374,86 @@ class AIMessage(BaseMessage):
 
 
 class AIMessageChunk(AIMessage, BaseMessageChunk):
-    """Message chunk from an AI."""
+    """Message chunk from an AI (yielded when streaming)."""
 
     # Ignoring mypy re-assignment here since we're overriding the value
     # to make sure that the chunk variant can be discriminated from the
     # non-chunk variant.
     type: Literal["AIMessageChunk"] = "AIMessageChunk"  # type: ignore[assignment]
-    """The type of the message (used for deserialization).
-
-    Defaults to ``AIMessageChunk``.
-
-    """
+    """The type of the message (used for deserialization)."""
 
     tool_call_chunks: list[ToolCallChunk] = []
     """If provided, tool call chunks associated with the message."""
 
+    chunk_position: Literal["last"] | None = None
+    """Optional span represented by an aggregated `AIMessageChunk`.
+
+    If a chunk with `chunk_position="last"` is aggregated into a stream,
+    `tool_call_chunks` in message content will be parsed into `tool_calls`.
+    """
+
     @property
     def lc_attributes(self) -> dict:
-        """Attrs to be serialized even if they are derived from other init args."""
+        """Attributes to be serialized, even if they are derived from other initialization args."""  # noqa: E501
         return {
             "tool_calls": self.tool_calls,
             "invalid_tool_calls": self.invalid_tool_calls,
         }
 
+    @property
+    def content_blocks(self) -> list[types.ContentBlock]:
+        """Return standard, typed `ContentBlock` dicts from the message."""
+        if self.response_metadata.get("output_version") == "v1":
+            return cast("list[types.ContentBlock]", self.content)
+
+        model_provider = self.response_metadata.get("model_provider")
+        if model_provider:
+            from langchain_core.messages.block_translators import (  # noqa: PLC0415
+                get_translator,
+            )
+
+            translator = get_translator(model_provider)
+            if translator:
+                try:
+                    return translator["translate_content_chunk"](self)
+                except NotImplementedError:
+                    pass
+
+        # Otherwise, use best-effort parsing
+        blocks = super().content_blocks
+
+        if (
+            self.tool_call_chunks
+            and not self.content
+            and self.chunk_position != "last"  # keep tool_calls if aggregated
+        ):
+            blocks = [
+                block
+                for block in blocks
+                if block["type"] not in ("tool_call", "invalid_tool_call")
+            ]
+            for tool_call_chunk in self.tool_call_chunks:
+                tc: types.ToolCallChunk = {
+                    "type": "tool_call_chunk",
+                    "id": tool_call_chunk.get("id"),
+                    "name": tool_call_chunk.get("name"),
+                    "args": tool_call_chunk.get("args"),
+                }
+                if (idx := tool_call_chunk.get("index")) is not None:
+                    tc["index"] = idx
+                blocks.append(tc)
+
+        # Best-effort reasoning extraction from additional_kwargs
+        # Only add reasoning if not already present
+        # Insert before all other blocks to keep reasoning at the start
+        has_reasoning = any(block.get("type") == "reasoning" for block in blocks)
+        if not has_reasoning and (
+            reasoning_block := _extract_reasoning_from_additional_kwargs(self)
+        ):
+            blocks.insert(0, reasoning_block)
+
+        return blocks
+
     @model_validator(mode="after")
     def init_tool_calls(self) -> Self:
         """Initialize tool calls from tool call chunks.
@@ -379,10 +518,73 @@ class AIMessageChunk(AIMessage, BaseMessageChunk):
                 add_chunk_to_invalid_tool_calls(chunk)
         self.tool_calls = tool_calls
         self.invalid_tool_calls = invalid_tool_calls
+
+        if (
+            self.chunk_position == "last"
+            and self.tool_call_chunks
+            and self.response_metadata.get("output_version") == "v1"
+            and isinstance(self.content, list)
+        ):
+            id_to_tc: dict[str, types.ToolCall] = {
+                cast("str", tc.get("id")): {
+                    "type": "tool_call",
+                    "name": tc["name"],
+                    "args": tc["args"],
+                    "id": tc.get("id"),
+                }
+                for tc in self.tool_calls
+                if "id" in tc
+            }
+            for idx, block in enumerate(self.content):
+                if (
+                    isinstance(block, dict)
+                    and block.get("type") == "tool_call_chunk"
+                    and (call_id := block.get("id"))
+                    and call_id in id_to_tc
+                ):
+                    self.content[idx] = cast("dict[str, Any]", id_to_tc[call_id])
+                    if "extras" in block:
+                        # mypy does not account for instance check for dict above
+                        self.content[idx]["extras"] = block["extras"]  # type: ignore[index]
+
+        return self
+
+    @model_validator(mode="after")
+    def init_server_tool_calls(self) -> Self:
+        """Parse `server_tool_call_chunks`."""
+        if (
+            self.chunk_position == "last"
+            and self.response_metadata.get("output_version") == "v1"
+            and isinstance(self.content, list)
+        ):
+            for idx, block in enumerate(self.content):
+                if (
+                    isinstance(block, dict)
+                    and block.get("type")
+                    in ("server_tool_call", "server_tool_call_chunk")
+                    and (args_str := block.get("args"))
+                    and isinstance(args_str, str)
+                ):
+                    try:
+                        args = json.loads(args_str)
+                        if isinstance(args, dict):
+                            self.content[idx]["type"] = "server_tool_call"  # type: ignore[index]
+                            self.content[idx]["args"] = args  # type: ignore[index]
+                    except json.JSONDecodeError:
+                        pass
         return self
 
+    @overload  # type: ignore[override]  # summing BaseMessages gives ChatPromptTemplate
+    def __add__(self, other: "AIMessageChunk") -> "AIMessageChunk": ...
+
+    @overload
+    def __add__(self, other: Sequence["AIMessageChunk"]) -> "AIMessageChunk": ...
+
+    @overload
+    def __add__(self, other: Any) -> BaseMessageChunk: ...
+
     @override
-    def __add__(self, other: Any) -> BaseMessageChunk:  # type: ignore[override]
+    def __add__(self, other: Any) -> BaseMessageChunk:
         if isinstance(other, AIMessageChunk):
             return add_ai_message_chunks(self, other)
         if isinstance(other, (list, tuple)) and all(
@@ -395,23 +597,16 @@ class AIMessageChunk(AIMessage, BaseMessageChunk):
 def add_ai_message_chunks(
     left: AIMessageChunk, *others: AIMessageChunk
 ) -> AIMessageChunk:
-    """Add multiple ``AIMessageChunk``s together.
+    """Add multiple `AIMessageChunk`s together.
 
     Args:
-        left: The first ``AIMessageChunk``.
-        *others: Other ``AIMessageChunk``s to add.
-
-    Raises:
-        ValueError: If the example values of the chunks are not the same.
+        left: The first `AIMessageChunk`.
+        *others: Other `AIMessageChunk`s to add.
 
     Returns:
-        The resulting ``AIMessageChunk``.
+        The resulting `AIMessageChunk`.
 
     """
-    if any(left.example != o.example for o in others):
-        msg = "Cannot concatenate AIMessageChunks with different example values."
-        raise ValueError(msg)
-
     content = merge_content(left.content, *(o.content for o in others))
     additional_kwargs = merge_dicts(
         left.additional_kwargs, *(o.additional_kwargs for o in others)
@@ -438,7 +633,7 @@ def add_ai_message_chunks(
 
     # Token usage
     if left.usage_metadata or any(o.usage_metadata is not None for o in others):
-        usage_metadata: Optional[UsageMetadata] = left.usage_metadata
+        usage_metadata: UsageMetadata | None = left.usage_metadata
         for other in others:
             usage_metadata = add_usage(usage_metadata, other.usage_metadata)
     else:
@@ -446,72 +641,83 @@ def add_ai_message_chunks(
 
     chunk_id = None
     candidates = [left.id] + [o.id for o in others]
-    # first pass: pick the first non-run-* id
+    # first pass: pick the first provider-assigned id (non-run-* and non-lc_*)
     for id_ in candidates:
-        if id_ and not id_.startswith(_LC_ID_PREFIX):
+        if (
+            id_
+            and not id_.startswith(LC_ID_PREFIX)
+            and not id_.startswith(LC_AUTO_PREFIX)
+        ):
             chunk_id = id_
             break
     else:
-        # second pass: no provider-assigned id found, just take the first non-null
+        # second pass: prefer lc_run-* ids over lc_* ids
         for id_ in candidates:
-            if id_:
+            if id_ and id_.startswith(LC_ID_PREFIX):
                 chunk_id = id_
                 break
+        else:
+            # third pass: take any remaining id (auto-generated lc_* ids)
+            for id_ in candidates:
+                if id_:
+                    chunk_id = id_
+                    break
+
+    chunk_position: Literal["last"] | None = (
+        "last" if any(x.chunk_position == "last" for x in [left, *others]) else None
+    )
 
     return left.__class__(
-        example=left.example,
         content=content,
         additional_kwargs=additional_kwargs,
         tool_call_chunks=tool_call_chunks,
         response_metadata=response_metadata,
         usage_metadata=usage_metadata,
         id=chunk_id,
+        chunk_position=chunk_position,
     )
 
 
-def add_usage(
-    left: Optional[UsageMetadata], right: Optional[UsageMetadata]
-) -> UsageMetadata:
+def add_usage(left: UsageMetadata | None, right: UsageMetadata | None) -> UsageMetadata:
     """Recursively add two UsageMetadata objects.
 
     Example:
-        .. code-block:: python
-
-            from langchain_core.messages.ai import add_usage
-
-            left = UsageMetadata(
-                input_tokens=5,
-                output_tokens=0,
-                total_tokens=5,
-                input_token_details=InputTokenDetails(cache_read=3),
-            )
-            right = UsageMetadata(
-                input_tokens=0,
-                output_tokens=10,
-                total_tokens=10,
-                output_token_details=OutputTokenDetails(reasoning=4),
-            )
+        ```python
+        from langchain_core.messages.ai import add_usage
+
+        left = UsageMetadata(
+            input_tokens=5,
+            output_tokens=0,
+            total_tokens=5,
+            input_token_details=InputTokenDetails(cache_read=3),
+        )
+        right = UsageMetadata(
+            input_tokens=0,
+            output_tokens=10,
+            total_tokens=10,
+            output_token_details=OutputTokenDetails(reasoning=4),
+        )
 
-            add_usage(left, right)
+        add_usage(left, right)
+        ```
 
         results in
 
-        .. code-block:: python
-
-            UsageMetadata(
-                input_tokens=5,
-                output_tokens=10,
-                total_tokens=15,
-                input_token_details=InputTokenDetails(cache_read=3),
-                output_token_details=OutputTokenDetails(reasoning=4),
-            )
-
+        ```python
+        UsageMetadata(
+            input_tokens=5,
+            output_tokens=10,
+            total_tokens=15,
+            input_token_details=InputTokenDetails(cache_read=3),
+            output_token_details=OutputTokenDetails(reasoning=4),
+        )
+        ```
     Args:
-        left: The first ``UsageMetadata`` object.
-        right: The second ``UsageMetadata`` object.
+        left: The first `UsageMetadata` object.
+        right: The second `UsageMetadata` object.
 
     Returns:
-        The sum of the two ``UsageMetadata`` objects.
+        The sum of the two `UsageMetadata` objects.
 
     """
     if not (left or right):
@@ -532,50 +738,49 @@ def add_usage(
 
 
 def subtract_usage(
-    left: Optional[UsageMetadata], right: Optional[UsageMetadata]
+    left: UsageMetadata | None, right: UsageMetadata | None
 ) -> UsageMetadata:
-    """Recursively subtract two ``UsageMetadata`` objects.
+    """Recursively subtract two `UsageMetadata` objects.
 
-    Token counts cannot be negative so the actual operation is ``max(left - right, 0)``.
+    Token counts cannot be negative so the actual operation is `max(left - right, 0)`.
 
     Example:
-        .. code-block:: python
-
-            from langchain_core.messages.ai import subtract_usage
-
-            left = UsageMetadata(
-                input_tokens=5,
-                output_tokens=10,
-                total_tokens=15,
-                input_token_details=InputTokenDetails(cache_read=4),
-            )
-            right = UsageMetadata(
-                input_tokens=3,
-                output_tokens=8,
-                total_tokens=11,
-                output_token_details=OutputTokenDetails(reasoning=4),
-            )
+        ```python
+        from langchain_core.messages.ai import subtract_usage
+
+        left = UsageMetadata(
+            input_tokens=5,
+            output_tokens=10,
+            total_tokens=15,
+            input_token_details=InputTokenDetails(cache_read=4),
+        )
+        right = UsageMetadata(
+            input_tokens=3,
+            output_tokens=8,
+            total_tokens=11,
+            output_token_details=OutputTokenDetails(reasoning=4),
+        )
 
-            subtract_usage(left, right)
+        subtract_usage(left, right)
+        ```
 
         results in
 
-        .. code-block:: python
-
-            UsageMetadata(
-                input_tokens=2,
-                output_tokens=2,
-                total_tokens=4,
-                input_token_details=InputTokenDetails(cache_read=4),
-                output_token_details=OutputTokenDetails(reasoning=0),
-            )
-
+        ```python
+        UsageMetadata(
+            input_tokens=2,
+            output_tokens=2,
+            total_tokens=4,
+            input_token_details=InputTokenDetails(cache_read=4),
+            output_token_details=OutputTokenDetails(reasoning=0),
+        )
+        ```
     Args:
-        left: The first ``UsageMetadata`` object.
-        right: The second ``UsageMetadata`` object.
+        left: The first `UsageMetadata` object.
+        right: The second `UsageMetadata` object.
 
     Returns:
-        The resulting ``UsageMetadata`` after subtraction.
+        The resulting `UsageMetadata` after subtraction.
 
     """
     if not (left or right):