langchain-core 0.4.0.dev0__py3-none-any.whl → 1.0.0__py3-none-any.whl

langchain_core/tracers/event_stream.py

@@ -8,22 +8,22 @@ import logging
 from typing import (
     TYPE_CHECKING,
     Any,
-    Optional,
+    TypedDict,
     TypeVar,
-    Union,
     cast,
 )
 from uuid import UUID, uuid4
 
-from typing_extensions import NotRequired, TypedDict, override
+from typing_extensions import NotRequired, override
 
-from langchain_core.callbacks.base import AsyncCallbackHandler
+from langchain_core.callbacks.base import AsyncCallbackHandler, BaseCallbackManager
 from langchain_core.messages import AIMessageChunk, BaseMessage, BaseMessageChunk
 from langchain_core.outputs import (
     ChatGenerationChunk,
     GenerationChunk,
     LLMResult,
 )
+from langchain_core.runnables import ensure_config
 from langchain_core.runnables.schema import (
     CustomStreamEvent,
     EventData,
@@ -36,9 +36,13 @@ from langchain_core.runnables.utils import (
     _RootEventFilter,
 )
 from langchain_core.tracers._streaming import _StreamingCallbackHandler
+from langchain_core.tracers.log_stream import (
+    LogStreamCallbackHandler,
+    RunLog,
+    _astream_log_implementation,
+)
 from langchain_core.tracers.memory_stream import _MemoryStream
 from langchain_core.utils.aiter import aclosing, py_anext
-from langchain_core.v1.messages import MessageV1
 
 if TYPE_CHECKING:
     from collections.abc import AsyncIterator, Iterator, Sequence
@@ -46,8 +50,6 @@ if TYPE_CHECKING:
     from langchain_core.documents import Document
     from langchain_core.runnables import Runnable, RunnableConfig
     from langchain_core.tracers.log_stream import LogEntry
-    from langchain_core.v1.messages import AIMessage as AIMessageV1
-    from langchain_core.v1.messages import AIMessageChunk as AIMessageChunkV1
 
 logger = logging.getLogger(__name__)
 
@@ -56,25 +58,23 @@ class RunInfo(TypedDict):
     """Information about a run.
 
     This is used to keep track of the metadata associated with a run.
-
-    Parameters:
-        name: The name of the run.
-        tags: The tags associated with the run.
-        metadata: The metadata associated with the run.
-        run_type: The type of the run.
-        inputs: The inputs to the run.
-        parent_run_id: The ID of the parent run.
     """
 
     name: str
+    """The name of the run."""
     tags: list[str]
+    """The tags associated with the run."""
     metadata: dict[str, Any]
+    """The metadata associated with the run."""
     run_type: str
+    """The type of the run."""
     inputs: NotRequired[Any]
-    parent_run_id: Optional[UUID]
+    """The inputs to the run."""
+    parent_run_id: UUID | None
+    """The ID of the parent run."""
 
 
-def _assign_name(name: Optional[str], serialized: Optional[dict[str, Any]]) -> str:
+def _assign_name(name: str | None, serialized: dict[str, Any] | None) -> str:
     """Assign a name to a run."""
     if name is not None:
         return name
@@ -95,12 +95,12 @@ class _AstreamEventsCallbackHandler(AsyncCallbackHandler, _StreamingCallbackHand
     def __init__(
         self,
         *args: Any,
-        include_names: Optional[Sequence[str]] = None,
-        include_types: Optional[Sequence[str]] = None,
-        include_tags: Optional[Sequence[str]] = None,
-        exclude_names: Optional[Sequence[str]] = None,
-        exclude_types: Optional[Sequence[str]] = None,
-        exclude_tags: Optional[Sequence[str]] = None,
+        include_names: Sequence[str] | None = None,
+        include_types: Sequence[str] | None = None,
+        include_tags: Sequence[str] | None = None,
+        exclude_names: Sequence[str] | None = None,
+        exclude_types: Sequence[str] | None = None,
+        exclude_tags: Sequence[str] | None = None,
         **kwargs: Any,
     ) -> None:
         """Initialize the tracer."""
@@ -114,7 +114,7 @@ class _AstreamEventsCallbackHandler(AsyncCallbackHandler, _StreamingCallbackHand
         # of a child run, which results in clean up of run_map.
         # So we keep track of the mapping between children and parent run IDs
         # in a separate container. This container is GCed when the tracer is GCed.
-        self.parent_map: dict[UUID, Optional[UUID]] = {}
+        self.parent_map: dict[UUID, UUID | None] = {}
 
         self.is_tapped: dict[UUID, Any] = {}
 
@@ -128,7 +128,10 @@ class _AstreamEventsCallbackHandler(AsyncCallbackHandler, _StreamingCallbackHand
             exclude_tags=exclude_tags,
         )
 
-        loop = asyncio.get_event_loop()
+        try:
+            loop = asyncio.get_event_loop()
+        except RuntimeError:
+            loop = asyncio.new_event_loop()
         memory_stream = _MemoryStream[StreamEvent](loop)
         self.send_stream = memory_stream.get_send_stream()
         self.receive_stream = memory_stream.get_receive_stream()
@@ -158,7 +161,11 @@ class _AstreamEventsCallbackHandler(AsyncCallbackHandler, _StreamingCallbackHand
             self.send_stream.send_nowait(event)
 
     def __aiter__(self) -> AsyncIterator[Any]:
-        """Iterate over the receive stream."""
+        """Iterate over the receive stream.
+
+        Returns:
+            An async iterator over the receive stream.
+        """
         return self.receive_stream.__aiter__()
 
     async def tap_output_aiter(
@@ -175,7 +182,7 @@ class _AstreamEventsCallbackHandler(AsyncCallbackHandler, _StreamingCallbackHand
             output: The output of the Runnable.
 
         Yields:
-            T: The output of the Runnable.
+            The output of the Runnable.
         """
         sentinel = object()
         # atomic check and set
@@ -218,14 +225,14 @@ class _AstreamEventsCallbackHandler(AsyncCallbackHandler, _StreamingCallbackHand
                 yield chunk
 
     def tap_output_iter(self, run_id: UUID, output: Iterator[T]) -> Iterator[T]:
-        """Tap the output aiter.
+        """Tap the output iter.
 
         Args:
             run_id: The ID of the run.
             output: The output of the Runnable.
 
         Yields:
-            T: The output of the Runnable.
+            The output of the Runnable.
         """
         sentinel = object()
         # atomic check and set
@@ -271,9 +278,9 @@ class _AstreamEventsCallbackHandler(AsyncCallbackHandler, _StreamingCallbackHand
         self,
         run_id: UUID,
         *,
-        tags: Optional[list[str]],
-        metadata: Optional[dict[str, Any]],
-        parent_run_id: Optional[UUID],
+        tags: list[str] | None,
+        metadata: dict[str, Any] | None,
+        parent_run_id: UUID | None,
         name_: str,
         run_type: str,
         **kwargs: Any,
@@ -300,18 +307,16 @@ class _AstreamEventsCallbackHandler(AsyncCallbackHandler, _StreamingCallbackHand
     async def on_chat_model_start(
         self,
         serialized: dict[str, Any],
-        messages: Union[list[list[BaseMessage]], list[MessageV1]],
+        messages: list[list[BaseMessage]],
         *,
         run_id: UUID,
-        tags: Optional[list[str]] = None,
-        parent_run_id: Optional[UUID] = None,
-        metadata: Optional[dict[str, Any]] = None,
-        name: Optional[str] = None,
+        tags: list[str] | None = None,
+        parent_run_id: UUID | None = None,
+        metadata: dict[str, Any] | None = None,
+        name: str | None = None,
         **kwargs: Any,
     ) -> None:
-        """Start a trace for an LLM run."""
-        # below cast is because type is converted in handle_event
-        messages = cast("list[list[BaseMessage]]", messages)
+        """Start a trace for a chat model run."""
         name_ = _assign_name(name, serialized)
         run_type = "chat_model"
 
@@ -347,13 +352,13 @@ class _AstreamEventsCallbackHandler(AsyncCallbackHandler, _StreamingCallbackHand
         prompts: list[str],
         *,
         run_id: UUID,
-        tags: Optional[list[str]] = None,
-        parent_run_id: Optional[UUID] = None,
-        metadata: Optional[dict[str, Any]] = None,
-        name: Optional[str] = None,
+        tags: list[str] | None = None,
+        parent_run_id: UUID | None = None,
+        metadata: dict[str, Any] | None = None,
+        name: str | None = None,
         **kwargs: Any,
     ) -> None:
-        """Start a trace for an LLM run."""
+        """Start a trace for a (non-chat model) LLM run."""
         name_ = _assign_name(name, serialized)
         run_type = "llm"
 
@@ -391,8 +396,8 @@ class _AstreamEventsCallbackHandler(AsyncCallbackHandler, _StreamingCallbackHand
         data: Any,
         *,
         run_id: UUID,
-        tags: Optional[list[str]] = None,
-        metadata: Optional[dict[str, Any]] = None,
+        tags: list[str] | None = None,
+        metadata: dict[str, Any] | None = None,
         **kwargs: Any,
     ) -> None:
         """Generate a custom astream event."""
@@ -412,19 +417,17 @@ class _AstreamEventsCallbackHandler(AsyncCallbackHandler, _StreamingCallbackHand
         self,
         token: str,
         *,
-        chunk: Optional[
-            Union[GenerationChunk, ChatGenerationChunk, AIMessageChunkV1]
-        ] = None,
+        chunk: GenerationChunk | ChatGenerationChunk | None = None,
         run_id: UUID,
-        parent_run_id: Optional[UUID] = None,
+        parent_run_id: UUID | None = None,
         **kwargs: Any,
     ) -> None:
-        """Run on new LLM token. Only available when streaming is enabled."""
+        """Run on new output token. Only available when streaming is enabled.
+
+        For both chat models and non-chat models (legacy LLMs).
+        """
         run_info = self.run_map.get(run_id)
-        chunk = cast(
-            "Optional[Union[GenerationChunk, ChatGenerationChunk]]", chunk
-        )  # converted in handle_event
-        chunk_: Union[GenerationChunk, BaseMessageChunk]
+        chunk_: GenerationChunk | BaseMessageChunk
 
         if run_info is None:
             msg = f"Run ID {run_id} not found in run map."
@@ -466,15 +469,20 @@ class _AstreamEventsCallbackHandler(AsyncCallbackHandler, _StreamingCallbackHand
 
     @override
     async def on_llm_end(
-        self, response: Union[LLMResult, AIMessageV1], *, run_id: UUID, **kwargs: Any
+        self, response: LLMResult, *, run_id: UUID, **kwargs: Any
     ) -> None:
-        """End a trace for an LLM run."""
-        response = cast("LLMResult", response)  # converted in handle_event
+        """End a trace for a model run.
+
+        For both chat models and non-chat models (legacy LLMs).
+
+        Raises:
+            ValueError: If the run type is not `'llm'` or `'chat_model'`.
+        """
         run_info = self.run_map.pop(run_id)
-        inputs_ = run_info["inputs"]
+        inputs_ = run_info.get("inputs")
 
-        generations: Union[list[list[GenerationChunk]], list[list[ChatGenerationChunk]]]
-        output: Union[dict, BaseMessage] = {}
+        generations: list[list[GenerationChunk]] | list[list[ChatGenerationChunk]]
+        output: dict | BaseMessage = {}
 
         if run_info["run_type"] == "chat_model":
             generations = cast("list[list[ChatGenerationChunk]]", response.generations)
@@ -526,11 +534,11 @@ class _AstreamEventsCallbackHandler(AsyncCallbackHandler, _StreamingCallbackHand
         inputs: dict[str, Any],
         *,
         run_id: UUID,
-        tags: Optional[list[str]] = None,
-        parent_run_id: Optional[UUID] = None,
-        metadata: Optional[dict[str, Any]] = None,
-        run_type: Optional[str] = None,
-        name: Optional[str] = None,
+        tags: list[str] | None = None,
+        parent_run_id: UUID | None = None,
+        metadata: dict[str, Any] | None = None,
+        run_type: str | None = None,
+        name: str | None = None,
         **kwargs: Any,
     ) -> None:
         """Start a trace for a chain run."""
@@ -574,7 +582,7 @@ class _AstreamEventsCallbackHandler(AsyncCallbackHandler, _StreamingCallbackHand
         outputs: dict[str, Any],
         *,
         run_id: UUID,
-        inputs: Optional[dict[str, Any]] = None,
+        inputs: dict[str, Any] | None = None,
         **kwargs: Any,
     ) -> None:
         """End a trace for a chain run."""
@@ -603,6 +611,28 @@ class _AstreamEventsCallbackHandler(AsyncCallbackHandler, _StreamingCallbackHand
             run_type,
         )
 
+    def _get_tool_run_info_with_inputs(self, run_id: UUID) -> tuple[RunInfo, Any]:
+        """Get run info for a tool and extract inputs, with validation.
+
+        Args:
+            run_id: The run ID of the tool.
+
+        Returns:
+            A tuple of (run_info, inputs).
+
+        Raises:
+            AssertionError: If the run ID is a tool call and does not have inputs.
+        """
+        run_info = self.run_map.pop(run_id)
+        if "inputs" not in run_info:
+            msg = (
+                f"Run ID {run_id} is a tool call and is expected to have "
+                f"inputs associated with it."
+            )
+            raise AssertionError(msg)
+        inputs = run_info["inputs"]
+        return run_info, inputs
+
     @override
     async def on_tool_start(
         self,
@@ -610,11 +640,11 @@ class _AstreamEventsCallbackHandler(AsyncCallbackHandler, _StreamingCallbackHand
         input_str: str,
         *,
         run_id: UUID,
-        tags: Optional[list[str]] = None,
-        parent_run_id: Optional[UUID] = None,
-        metadata: Optional[dict[str, Any]] = None,
-        name: Optional[str] = None,
-        inputs: Optional[dict[str, Any]] = None,
+        tags: list[str] | None = None,
+        parent_run_id: UUID | None = None,
+        metadata: dict[str, Any] | None = None,
+        name: str | None = None,
+        inputs: dict[str, Any] | None = None,
         **kwargs: Any,
     ) -> None:
         """Start a trace for a tool run."""
@@ -645,17 +675,43 @@ class _AstreamEventsCallbackHandler(AsyncCallbackHandler, _StreamingCallbackHand
             "tool",
         )
 
+    @override
+    async def on_tool_error(
+        self,
+        error: BaseException,
+        *,
+        run_id: UUID,
+        parent_run_id: UUID | None = None,
+        tags: list[str] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Run when tool errors."""
+        run_info, inputs = self._get_tool_run_info_with_inputs(run_id)
+
+        self._send(
+            {
+                "event": "on_tool_error",
+                "data": {
+                    "error": error,
+                    "input": inputs,
+                },
+                "run_id": str(run_id),
+                "name": run_info["name"],
+                "tags": run_info["tags"],
+                "metadata": run_info["metadata"],
+                "parent_ids": self._get_parent_ids(run_id),
+            },
+            "tool",
+        )
+
     @override
     async def on_tool_end(self, output: Any, *, run_id: UUID, **kwargs: Any) -> None:
-        """End a trace for a tool run."""
-        run_info = self.run_map.pop(run_id)
-        if "inputs" not in run_info:
-            msg = (
-                f"Run ID {run_id} is a tool call and is expected to have "
-                f"inputs associated with it."
-            )
-            raise AssertionError(msg)
-        inputs = run_info["inputs"]
+        """End a trace for a tool run.
+
+        Raises:
+            AssertionError: If the run ID is a tool call and does not have inputs
+        """
+        run_info, inputs = self._get_tool_run_info_with_inputs(run_id)
 
         self._send(
             {
@@ -680,10 +736,10 @@ class _AstreamEventsCallbackHandler(AsyncCallbackHandler, _StreamingCallbackHand
         query: str,
         *,
         run_id: UUID,
-        parent_run_id: Optional[UUID] = None,
-        tags: Optional[list[str]] = None,
-        metadata: Optional[dict[str, Any]] = None,
-        name: Optional[str] = None,
+        parent_run_id: UUID | None = None,
+        tags: list[str] | None = None,
+        metadata: dict[str, Any] | None = None,
+        name: str | None = None,
         **kwargs: Any,
     ) -> None:
         """Run when Retriever starts running."""
@@ -729,7 +785,7 @@ class _AstreamEventsCallbackHandler(AsyncCallbackHandler, _StreamingCallbackHand
                 "event": "on_retriever_end",
                 "data": {
                     "output": documents,
-                    "input": run_info["inputs"],
+                    "input": run_info.get("inputs"),
                 },
                 "run_id": str(run_id),
                 "name": run_info["name"],
@@ -741,35 +797,27 @@ class _AstreamEventsCallbackHandler(AsyncCallbackHandler, _StreamingCallbackHand
         )
 
     def __deepcopy__(self, memo: dict) -> _AstreamEventsCallbackHandler:
-        """Deepcopy the tracer."""
+        """Return self."""
         return self
 
     def __copy__(self) -> _AstreamEventsCallbackHandler:
-        """Copy the tracer."""
+        """Return self."""
         return self
 
 
 async def _astream_events_implementation_v1(
     runnable: Runnable[Input, Output],
     value: Any,
-    config: Optional[RunnableConfig] = None,
+    config: RunnableConfig | None = None,
     *,
-    include_names: Optional[Sequence[str]] = None,
-    include_types: Optional[Sequence[str]] = None,
-    include_tags: Optional[Sequence[str]] = None,
-    exclude_names: Optional[Sequence[str]] = None,
-    exclude_types: Optional[Sequence[str]] = None,
-    exclude_tags: Optional[Sequence[str]] = None,
+    include_names: Sequence[str] | None = None,
+    include_types: Sequence[str] | None = None,
+    include_tags: Sequence[str] | None = None,
+    exclude_names: Sequence[str] | None = None,
+    exclude_types: Sequence[str] | None = None,
+    exclude_tags: Sequence[str] | None = None,
     **kwargs: Any,
 ) -> AsyncIterator[StandardStreamEvent]:
-    from langchain_core.runnables import ensure_config
-    from langchain_core.runnables.utils import _RootEventFilter
-    from langchain_core.tracers.log_stream import (
-        LogStreamCallbackHandler,
-        RunLog,
-        _astream_log_implementation,
-    )
-
     stream = LogStreamCallbackHandler(
         auto_close=False,
         include_names=include_names,
@@ -849,12 +897,12 @@ async def _astream_events_implementation_v1(
                 # Usually they will NOT be available for components that operate
                 # on streams, since those components stream the input and
                 # don't know its final value until the end of the stream.
-                inputs = log_entry["inputs"]
+                inputs = log_entry.get("inputs")
                 if inputs is not None:
                     data["input"] = inputs
 
             if event_type == "end":
-                inputs = log_entry["inputs"]
+                inputs = log_entry.get("inputs")
                 if inputs is not None:
                     data["input"] = inputs
 
@@ -936,20 +984,17 @@ async def _astream_events_implementation_v1(
 async def _astream_events_implementation_v2(
     runnable: Runnable[Input, Output],
     value: Any,
-    config: Optional[RunnableConfig] = None,
+    config: RunnableConfig | None = None,
     *,
-    include_names: Optional[Sequence[str]] = None,
-    include_types: Optional[Sequence[str]] = None,
-    include_tags: Optional[Sequence[str]] = None,
-    exclude_names: Optional[Sequence[str]] = None,
-    exclude_types: Optional[Sequence[str]] = None,
-    exclude_tags: Optional[Sequence[str]] = None,
+    include_names: Sequence[str] | None = None,
+    include_types: Sequence[str] | None = None,
+    include_tags: Sequence[str] | None = None,
+    exclude_names: Sequence[str] | None = None,
+    exclude_types: Sequence[str] | None = None,
+    exclude_tags: Sequence[str] | None = None,
     **kwargs: Any,
 ) -> AsyncIterator[StandardStreamEvent]:
     """Implementation of the astream events API for V2 runnables."""
-    from langchain_core.callbacks.base import BaseCallbackManager
-    from langchain_core.runnables import ensure_config
-
     event_streamer = _AstreamEventsCallbackHandler(
         include_names=include_names,
         include_types=include_types,
@@ -1010,7 +1055,7 @@ async def _astream_events_implementation_v2(
                 continue
 
             # If it's the end event corresponding to the root runnable
-            # we dont include the input in the event since it's guaranteed
+            # we don't include the input in the event since it's guaranteed
             # to be included in the first event.
             if (
                 event["run_id"] == first_event_run_id