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

langchain_core/callbacks/manager.py

@@ -11,15 +11,7 @@ from abc import ABC, abstractmethod
 from concurrent.futures import ThreadPoolExecutor
 from contextlib import asynccontextmanager, contextmanager
 from contextvars import copy_context
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Callable,
-    Optional,
-    TypeVar,
-    Union,
-    cast,
-)
+from typing import TYPE_CHECKING, Any, Callable, Optional, TypeVar, Union, cast
 from uuid import UUID
 
 from langsmith.run_helpers import get_tracing_context
@@ -37,16 +29,8 @@ from langchain_core.callbacks.base import (
 )
 from langchain_core.callbacks.stdout import StdOutCallbackHandler
 from langchain_core.messages import BaseMessage, get_buffer_string
-from langchain_core.messages.utils import convert_from_v1_message
-from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, LLMResult
 from langchain_core.tracers.schemas import Run
 from langchain_core.utils.env import env_var_is_set
-from langchain_core.v1.messages import (
-    AIMessage,
-    AIMessageChunk,
-    MessageV1,
-    MessageV1Types,
-)
 
 if TYPE_CHECKING:
     from collections.abc import AsyncGenerator, Coroutine, Generator, Sequence
@@ -55,7 +39,7 @@ if TYPE_CHECKING:
 
     from langchain_core.agents import AgentAction, AgentFinish
     from langchain_core.documents import Document
-    from langchain_core.outputs import GenerationChunk
+    from langchain_core.outputs import ChatGenerationChunk, GenerationChunk, LLMResult
     from langchain_core.runnables.config import RunnableConfig
 
 logger = logging.getLogger(__name__)
@@ -100,7 +84,8 @@ def trace_as_chain_group(
         metadata (dict[str, Any], optional): The metadata to apply to all runs.
             Defaults to None.
 
-    Note: must have LANGCHAIN_TRACING_V2 env var set to true to see the trace in LangSmith.
+    .. note:
+        Must have ``LANGCHAIN_TRACING_V2`` env var set to true to see the trace in LangSmith.
 
     Returns:
         CallbackManagerForChainGroup: The callback manager for the chain group.
@@ -185,7 +170,8 @@ async def atrace_as_chain_group(
     Returns:
         AsyncCallbackManager: The async callback manager for the chain group.
 
-    Note: must have LANGCHAIN_TRACING_V2 env var set to true to see the trace in LangSmith.
+    .. note:
+        Must have ``LANGCHAIN_TRACING_V2`` env var set to true to see the trace in LangSmith.
 
     Example:
         .. code-block:: python
@@ -242,6 +228,7 @@ def shielded(func: Func) -> Func:
 
     Returns:
         Callable: The shielded function
+
     """
 
     @functools.wraps(func)
@@ -251,46 +238,6 @@ def shielded(func: Func) -> Func:
     return cast("Func", wrapped)
 
 
-def _convert_llm_events(
-    event_name: str, args: tuple[Any, ...], kwargs: dict[str, Any]
-) -> tuple[tuple[Any, ...], dict[str, Any]]:
-    args_list = list(args)
-    if (
-        event_name == "on_chat_model_start"
-        and isinstance(args_list[1], list)
-        and args_list[1]
-        and isinstance(args_list[1][0], MessageV1Types)
-    ):
-        batch = [
-            convert_from_v1_message(item)
-            for item in args_list[1]
-            if isinstance(item, MessageV1Types)
-        ]
-        args_list[1] = [batch]
-    elif (
-        event_name == "on_llm_new_token"
-        and "chunk" in kwargs
-        and isinstance(kwargs["chunk"], MessageV1Types)
-    ):
-        chunk = kwargs["chunk"]
-        kwargs["chunk"] = ChatGenerationChunk(text=chunk.text, message=chunk)
-    elif event_name == "on_llm_end" and isinstance(args_list[0], MessageV1Types):
-        args_list[0] = LLMResult(
-            generations=[
-                [
-                    ChatGeneration(
-                        text=args_list[0].text,
-                        message=convert_from_v1_message(args_list[0]),
-                    )
-                ]
-            ]
-        )
-    else:
-        pass
-
-    return tuple(args_list), kwargs
-
-
 def handle_event(
     handlers: list[BaseCallbackHandler],
     event_name: str,
@@ -300,15 +247,17 @@ def handle_event(
 ) -> None:
     """Generic event handler for CallbackManager.
 
-    Note: This function is used by LangServe to handle events.
+    .. note::
+        This function is used by ``LangServe`` to handle events.
 
     Args:
         handlers: The list of handlers that will handle the event.
-        event_name: The name of the event (e.g., "on_llm_start").
+        event_name: The name of the event (e.g., ``'on_llm_start'``).
         ignore_condition_name: Name of the attribute defined on handler
             that if True will cause the handler to be skipped for the given event.
         *args: The arguments to pass to the event handler.
         **kwargs: The keyword arguments to pass to the event handler
+
     """
     coros: list[Coroutine[Any, Any, Any]] = []
 
@@ -319,8 +268,6 @@ def handle_event(
                 if ignore_condition_name is None or not getattr(
                     handler, ignore_condition_name
                 ):
-                    if not handler.accepts_new_messages:
-                        args, kwargs = _convert_llm_events(event_name, args, kwargs)
                     event = getattr(handler, event_name)(*args, **kwargs)
                     if asyncio.iscoroutine(event):
                         coros.append(event)
@@ -415,8 +362,6 @@ async def _ahandle_event_for_handler(
 ) -> None:
     try:
         if ignore_condition_name is None or not getattr(handler, ignore_condition_name):
-            if not handler.accepts_new_messages:
-                args, kwargs = _convert_llm_events(event_name, args, kwargs)
             event = getattr(handler, event_name)
             if asyncio.iscoroutinefunction(event):
                 await event(*args, **kwargs)
@@ -467,17 +412,19 @@ async def ahandle_event(
     *args: Any,
     **kwargs: Any,
 ) -> None:
-    """Async generic event handler for AsyncCallbackManager.
+    """Async generic event handler for ``AsyncCallbackManager``.
 
-    Note: This function is used by LangServe to handle events.
+    .. note::
+        This function is used by ``LangServe`` to handle events.
 
     Args:
         handlers: The list of handlers that will handle the event.
-        event_name: The name of the event (e.g., "on_llm_start").
+        event_name: The name of the event (e.g., ``'on_llm_start'``).
         ignore_condition_name: Name of the attribute defined on handler
             that if True will cause the handler to be skipped for the given event.
         *args: The arguments to pass to the event handler.
         **kwargs: The keyword arguments to pass to the event handler.
+
     """
     for handler in [h for h in handlers if h.run_inline]:
         await _ahandle_event_for_handler(
@@ -529,6 +476,7 @@ class BaseRunManager(RunManagerMixin):
                 Defaults to None.
             inheritable_metadata (Optional[dict[str, Any]]): The inheritable metadata.
                 Defaults to None.
+
         """
         self.run_id = run_id
         self.handlers = handlers
@@ -545,6 +493,7 @@ class BaseRunManager(RunManagerMixin):
 
         Returns:
             BaseRunManager: The noop manager.
+
         """
         return cls(
             run_id=uuid.uuid4(),
@@ -597,6 +546,7 @@ class RunManager(BaseRunManager):
         Args:
             retry_state (RetryCallState): The retry state.
             **kwargs (Any): Additional keyword arguments.
+
         """
         if not self.handlers:
             return
@@ -624,6 +574,7 @@ class ParentRunManager(RunManager):
 
         Returns:
             CallbackManager: The child callback manager.
+
         """
         manager = CallbackManager(handlers=[], parent_run_id=self.run_id)
         manager.set_handlers(self.inheritable_handlers)
@@ -643,6 +594,7 @@ class AsyncRunManager(BaseRunManager, ABC):
 
         Returns:
             RunManager: The sync RunManager.
+
         """
 
     async def on_text(
@@ -658,6 +610,7 @@ class AsyncRunManager(BaseRunManager, ABC):
 
         Returns:
             Any: The result of the callback.
+
         """
         if not self.handlers:
             return
@@ -682,6 +635,7 @@ class AsyncRunManager(BaseRunManager, ABC):
         Args:
             retry_state (RetryCallState): The retry state.
             **kwargs (Any): Additional keyword arguments.
+
         """
         if not self.handlers:
             return
@@ -709,6 +663,7 @@ class AsyncParentRunManager(AsyncRunManager):
 
         Returns:
             AsyncCallbackManager: The child callback manager.
+
         """
         manager = AsyncCallbackManager(handlers=[], parent_run_id=self.run_id)
         manager.set_handlers(self.inheritable_handlers)
@@ -726,9 +681,7 @@ class CallbackManagerForLLMRun(RunManager, LLMManagerMixin):
         self,
         token: str,
         *,
-        chunk: Optional[
-            Union[GenerationChunk, ChatGenerationChunk, AIMessageChunk]
-        ] = None,
+        chunk: Optional[Union[GenerationChunk, ChatGenerationChunk]] = None,
         **kwargs: Any,
     ) -> None:
         """Run when LLM generates a new token.
@@ -738,6 +691,7 @@ class CallbackManagerForLLMRun(RunManager, LLMManagerMixin):
             chunk (Optional[Union[GenerationChunk, ChatGenerationChunk]], optional):
                 The chunk. Defaults to None.
             **kwargs (Any): Additional keyword arguments.
+
         """
         if not self.handlers:
             return
@@ -753,12 +707,13 @@ class CallbackManagerForLLMRun(RunManager, LLMManagerMixin):
             **kwargs,
         )
 
-    def on_llm_end(self, response: Union[LLMResult, AIMessage], **kwargs: Any) -> None:
+    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
         """Run when LLM ends running.
 
         Args:
-            response (LLMResult | AIMessage): The LLM result.
+            response (LLMResult): The LLM result.
             **kwargs (Any): Additional keyword arguments.
+
         """
         if not self.handlers:
             return
@@ -783,8 +738,8 @@ class CallbackManagerForLLMRun(RunManager, LLMManagerMixin):
         Args:
             error (Exception or KeyboardInterrupt): The error.
             kwargs (Any): Additional keyword arguments.
-                - response (LLMResult | AIMessage): The response which was generated
-                  before the error occurred.
+                - response (LLMResult): The response which was generated before
+                    the error occurred.
         """
         if not self.handlers:
             return
@@ -808,6 +763,7 @@ class AsyncCallbackManagerForLLMRun(AsyncRunManager, LLMManagerMixin):
 
         Returns:
             CallbackManagerForLLMRun: The sync RunManager.
+
         """
         return CallbackManagerForLLMRun(
             run_id=self.run_id,
@@ -824,9 +780,7 @@ class AsyncCallbackManagerForLLMRun(AsyncRunManager, LLMManagerMixin):
         self,
         token: str,
         *,
-        chunk: Optional[
-            Union[GenerationChunk, ChatGenerationChunk, AIMessageChunk]
-        ] = None,
+        chunk: Optional[Union[GenerationChunk, ChatGenerationChunk]] = None,
         **kwargs: Any,
     ) -> None:
         """Run when LLM generates a new token.
@@ -836,6 +790,7 @@ class AsyncCallbackManagerForLLMRun(AsyncRunManager, LLMManagerMixin):
             chunk (Optional[Union[GenerationChunk, ChatGenerationChunk]], optional):
                 The chunk. Defaults to None.
             **kwargs (Any): Additional keyword arguments.
+
         """
         if not self.handlers:
             return
@@ -852,14 +807,13 @@ class AsyncCallbackManagerForLLMRun(AsyncRunManager, LLMManagerMixin):
         )
 
     @shielded
-    async def on_llm_end(
-        self, response: Union[LLMResult, AIMessage], **kwargs: Any
-    ) -> None:
+    async def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
         """Run when LLM ends running.
 
         Args:
-            response (LLMResult | AIMessage): The LLM result.
+            response (LLMResult): The LLM result.
             **kwargs (Any): Additional keyword arguments.
+
         """
         if not self.handlers:
             return
@@ -885,8 +839,11 @@ class AsyncCallbackManagerForLLMRun(AsyncRunManager, LLMManagerMixin):
         Args:
             error (Exception or KeyboardInterrupt): The error.
             kwargs (Any): Additional keyword arguments.
-                - response (LLMResult | AIMessage): The response which was generated
-                  before the error occurred.
+                - response (LLMResult): The response which was generated before
+                    the error occurred.
+
+
+
         """
         if not self.handlers:
             return
@@ -911,6 +868,7 @@ class CallbackManagerForChainRun(ParentRunManager, ChainManagerMixin):
         Args:
             outputs (Union[dict[str, Any], Any]): The outputs of the chain.
             **kwargs (Any): Additional keyword arguments.
+
         """
         if not self.handlers:
             return
@@ -935,6 +893,7 @@ class CallbackManagerForChainRun(ParentRunManager, ChainManagerMixin):
         Args:
             error (Exception or KeyboardInterrupt): The error.
             **kwargs (Any): Additional keyword arguments.
+
         """
         if not self.handlers:
             return
@@ -958,6 +917,7 @@ class CallbackManagerForChainRun(ParentRunManager, ChainManagerMixin):
 
         Returns:
             Any: The result of the callback.
+
         """
         if not self.handlers:
             return
@@ -981,6 +941,7 @@ class CallbackManagerForChainRun(ParentRunManager, ChainManagerMixin):
 
         Returns:
             Any: The result of the callback.
+
         """
         if not self.handlers:
             return
@@ -1025,6 +986,7 @@ class AsyncCallbackManagerForChainRun(AsyncParentRunManager, ChainManagerMixin):
         Args:
             outputs (Union[dict[str, Any], Any]): The outputs of the chain.
             **kwargs (Any): Additional keyword arguments.
+
         """
         if not self.handlers:
             return
@@ -1050,6 +1012,7 @@ class AsyncCallbackManagerForChainRun(AsyncParentRunManager, ChainManagerMixin):
         Args:
             error (Exception or KeyboardInterrupt): The error.
             **kwargs (Any): Additional keyword arguments.
+
         """
         if not self.handlers:
             return
@@ -1073,6 +1036,7 @@ class AsyncCallbackManagerForChainRun(AsyncParentRunManager, ChainManagerMixin):
 
         Returns:
             Any: The result of the callback.
+
         """
         if not self.handlers:
             return
@@ -1096,6 +1060,7 @@ class AsyncCallbackManagerForChainRun(AsyncParentRunManager, ChainManagerMixin):
 
         Returns:
             Any: The result of the callback.
+
         """
         if not self.handlers:
             return
@@ -1124,6 +1089,7 @@ class CallbackManagerForToolRun(ParentRunManager, ToolManagerMixin):
         Args:
             output (Any): The output of the tool.
             **kwargs (Any): The keyword arguments to pass to the event handler
+
         """
         if not self.handlers:
             return
@@ -1148,6 +1114,7 @@ class CallbackManagerForToolRun(ParentRunManager, ToolManagerMixin):
         Args:
             error (Exception or KeyboardInterrupt): The error.
             **kwargs (Any): Additional keyword arguments.
+
         """
         if not self.handlers:
             return
@@ -1189,6 +1156,7 @@ class AsyncCallbackManagerForToolRun(AsyncParentRunManager, ToolManagerMixin):
         Args:
             output (Any): The output of the tool.
             **kwargs (Any): Additional keyword arguments.
+
         """
         if not self.handlers:
             return
@@ -1213,6 +1181,7 @@ class AsyncCallbackManagerForToolRun(AsyncParentRunManager, ToolManagerMixin):
         Args:
             error (Exception or KeyboardInterrupt): The error.
             **kwargs (Any): Additional keyword arguments.
+
         """
         if not self.handlers:
             return
@@ -1241,6 +1210,7 @@ class CallbackManagerForRetrieverRun(ParentRunManager, RetrieverManagerMixin):
         Args:
             documents (Sequence[Document]): The retrieved documents.
             **kwargs (Any): Additional keyword arguments.
+
         """
         if not self.handlers:
             return
@@ -1265,6 +1235,7 @@ class CallbackManagerForRetrieverRun(ParentRunManager, RetrieverManagerMixin):
         Args:
             error (BaseException): The error.
             **kwargs (Any): Additional keyword arguments.
+
         """
         if not self.handlers:
             return
@@ -1291,6 +1262,7 @@ class AsyncCallbackManagerForRetrieverRun(
 
         Returns:
             CallbackManagerForRetrieverRun: The sync RunManager.
+
         """
         return CallbackManagerForRetrieverRun(
             run_id=self.run_id,
@@ -1312,6 +1284,7 @@ class AsyncCallbackManagerForRetrieverRun(
         Args:
             documents (Sequence[Document]): The retrieved documents.
             **kwargs (Any): Additional keyword arguments.
+
         """
         if not self.handlers:
             return
@@ -1337,6 +1310,7 @@ class AsyncCallbackManagerForRetrieverRun(
         Args:
             error (BaseException): The error.
             **kwargs (Any): Additional keyword arguments.
+
         """
         if not self.handlers:
             return
@@ -1373,6 +1347,7 @@ class CallbackManager(BaseCallbackManager):
         Returns:
             list[CallbackManagerForLLMRun]: A callback manager for each
                 prompt as an LLM run.
+
         """
         managers = []
         for i, prompt in enumerate(prompts):
@@ -1409,7 +1384,7 @@ class CallbackManager(BaseCallbackManager):
     def on_chat_model_start(
         self,
         serialized: dict[str, Any],
-        messages: Union[list[list[BaseMessage]], list[MessageV1]],
+        messages: list[list[BaseMessage]],
         run_id: Optional[UUID] = None,
         **kwargs: Any,
     ) -> list[CallbackManagerForLLMRun]:
@@ -1417,40 +1392,15 @@ class CallbackManager(BaseCallbackManager):
 
         Args:
             serialized (dict[str, Any]): The serialized LLM.
-            messages (list[list[BaseMessage | MessageV1]]): The list of messages.
+            messages (list[list[BaseMessage]]): The list of messages.
             run_id (UUID, optional): The ID of the run. Defaults to None.
             **kwargs (Any): Additional keyword arguments.
 
         Returns:
             list[CallbackManagerForLLMRun]: A callback manager for each
                 list of messages as an LLM run.
+
         """
-        if messages and isinstance(messages[0], MessageV1Types):
-            run_id_ = run_id if run_id is not None else uuid.uuid4()
-            handle_event(
-                self.handlers,
-                "on_chat_model_start",
-                "ignore_chat_model",
-                serialized,
-                messages,
-                run_id=run_id_,
-                parent_run_id=self.parent_run_id,
-                tags=self.tags,
-                metadata=self.metadata,
-                **kwargs,
-            )
-            return [
-                CallbackManagerForLLMRun(
-                    run_id=run_id_,
-                    handlers=self.handlers,
-                    inheritable_handlers=self.inheritable_handlers,
-                    parent_run_id=self.parent_run_id,
-                    tags=self.tags,
-                    inheritable_tags=self.inheritable_tags,
-                    metadata=self.metadata,
-                    inheritable_metadata=self.inheritable_metadata,
-                )
-            ]
         managers = []
         for message_list in messages:
             if run_id is not None:
@@ -1503,6 +1453,7 @@ class CallbackManager(BaseCallbackManager):
 
         Returns:
             CallbackManagerForChainRun: The callback manager for the chain run.
+
         """
         if run_id is None:
             run_id = uuid.uuid4()
@@ -1557,6 +1508,7 @@ class CallbackManager(BaseCallbackManager):
 
         Returns:
             CallbackManagerForToolRun: The callback manager for the tool run.
+
         """
         if run_id is None:
             run_id = uuid.uuid4()
@@ -1603,6 +1555,7 @@ class CallbackManager(BaseCallbackManager):
             run_id (UUID, optional): The ID of the run. Defaults to None.
             parent_run_id (UUID, optional): The ID of the parent run. Defaults to None.
             **kwargs (Any): Additional keyword arguments.
+
         """
         if run_id is None:
             run_id = uuid.uuid4()
@@ -1650,6 +1603,7 @@ class CallbackManager(BaseCallbackManager):
             run_id: The ID of the run. Defaults to None.
 
         .. versionadded:: 0.2.14
+
         """
         if not self.handlers:
             return
@@ -1704,6 +1658,7 @@ class CallbackManager(BaseCallbackManager):
 
         Returns:
             CallbackManager: The configured callback manager.
+
         """
         return _configure(
             cls,
@@ -1738,6 +1693,7 @@ class CallbackManagerForChainGroup(CallbackManager):
             parent_run_id (Optional[UUID]): The ID of the parent run. Defaults to None.
             parent_run_manager (CallbackManagerForChainRun): The parent run manager.
             **kwargs (Any): Additional keyword arguments.
+
         """
         super().__init__(
             handlers,
@@ -1826,6 +1782,7 @@ class CallbackManagerForChainGroup(CallbackManager):
         Args:
             outputs (Union[dict[str, Any], Any]): The outputs of the chain.
             **kwargs (Any): Additional keyword arguments.
+
         """
         self.ended = True
         return self.parent_run_manager.on_chain_end(outputs, **kwargs)
@@ -1840,6 +1797,7 @@ class CallbackManagerForChainGroup(CallbackManager):
         Args:
             error (Exception or KeyboardInterrupt): The error.
             **kwargs (Any): Additional keyword arguments.
+
         """
         self.ended = True
         return self.parent_run_manager.on_chain_error(error, **kwargs)
@@ -1945,7 +1903,7 @@ class AsyncCallbackManager(BaseCallbackManager):
     async def on_chat_model_start(
         self,
         serialized: dict[str, Any],
-        messages: Union[list[list[BaseMessage]], list[MessageV1]],
+        messages: list[list[BaseMessage]],
         run_id: Optional[UUID] = None,
         **kwargs: Any,
     ) -> list[AsyncCallbackManagerForLLMRun]:
@@ -1953,7 +1911,7 @@ class AsyncCallbackManager(BaseCallbackManager):
 
         Args:
             serialized (dict[str, Any]): The serialized LLM.
-            messages (list[list[BaseMessage | MessageV1]]): The list of messages.
+            messages (list[list[BaseMessage]]): The list of messages.
             run_id (UUID, optional): The ID of the run. Defaults to None.
             **kwargs (Any): Additional keyword arguments.
 
@@ -1962,51 +1920,10 @@ class AsyncCallbackManager(BaseCallbackManager):
                 async callback managers, one for each LLM Run
                 corresponding to each inner  message list.
         """
-        if messages and isinstance(messages[0], MessageV1Types):
-            run_id_ = run_id if run_id is not None else uuid.uuid4()
-            inline_tasks = []
-            non_inline_tasks = []
-            for handler in self.handlers:
-                task = ahandle_event(
-                    [handler],
-                    "on_chat_model_start",
-                    "ignore_chat_model",
-                    serialized,
-                    messages,
-                    run_id=run_id_,
-                    parent_run_id=self.parent_run_id,
-                    tags=self.tags,
-                    metadata=self.metadata,
-                    **kwargs,
-                )
-                if handler.run_inline:
-                    inline_tasks.append(task)
-                else:
-                    non_inline_tasks.append(task)
-            managers = [
-                AsyncCallbackManagerForLLMRun(
-                    run_id=run_id_,
-                    handlers=self.handlers,
-                    inheritable_handlers=self.inheritable_handlers,
-                    parent_run_id=self.parent_run_id,
-                    tags=self.tags,
-                    inheritable_tags=self.inheritable_tags,
-                    metadata=self.metadata,
-                    inheritable_metadata=self.inheritable_metadata,
-                )
-            ]
-            # Run inline tasks sequentially
-            for task in inline_tasks:
-                await task
-
-            # Run non-inline tasks concurrently
-            if non_inline_tasks:
-                await asyncio.gather(*non_inline_tasks)
-
-            return managers
         inline_tasks = []
         non_inline_tasks = []
         managers = []
+
         for message_list in messages:
             if run_id is not None:
                 run_id_ = run_id

langchain_core/callbacks/streaming_stdout.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 import sys
-from typing import TYPE_CHECKING, Any, Union
+from typing import TYPE_CHECKING, Any
 
 from typing_extensions import override
 
@@ -13,7 +13,6 @@ if TYPE_CHECKING:
     from langchain_core.agents import AgentAction, AgentFinish
     from langchain_core.messages import BaseMessage
     from langchain_core.outputs import LLMResult
-    from langchain_core.v1.messages import AIMessage, MessageV1
 
 
 class StreamingStdOutCallbackHandler(BaseCallbackHandler):
@@ -33,7 +32,7 @@ class StreamingStdOutCallbackHandler(BaseCallbackHandler):
     def on_chat_model_start(
         self,
         serialized: dict[str, Any],
-        messages: Union[list[list[BaseMessage]], list[MessageV1]],
+        messages: list[list[BaseMessage]],
         **kwargs: Any,
     ) -> None:
         """Run when LLM starts running.
@@ -55,7 +54,7 @@ class StreamingStdOutCallbackHandler(BaseCallbackHandler):
         sys.stdout.write(token)
         sys.stdout.flush()
 
-    def on_llm_end(self, response: Union[LLMResult, AIMessage], **kwargs: Any) -> None:
+    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
         """Run when LLM ends running.
 
         Args: