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

langchain_core/language_models/fake.py

@@ -3,7 +3,7 @@
 import asyncio
 import time
 from collections.abc import AsyncIterator, Iterator, Mapping
-from typing import Any, Optional
+from typing import Any
 
 from typing_extensions import override
 
@@ -23,7 +23,7 @@ class FakeListLLM(LLM):
     """List of responses to return in order."""
     # This parameter should be removed from FakeListLLM since
     # it's only used by sub-classes.
-    sleep: Optional[float] = None
+    sleep: float | None = None
     """Sleep time in seconds between responses.
 
     Ignored by FakeListLLM, but used by sub-classes.
@@ -44,8 +44,8 @@ class FakeListLLM(LLM):
     def _call(
         self,
         prompt: str,
-        stop: Optional[list[str]] = None,
-        run_manager: Optional[CallbackManagerForLLMRun] = None,
+        stop: list[str] | None = None,
+        run_manager: CallbackManagerForLLMRun | None = None,
         **kwargs: Any,
     ) -> str:
         """Return next response."""
@@ -60,8 +60,8 @@ class FakeListLLM(LLM):
     async def _acall(
         self,
         prompt: str,
-        stop: Optional[list[str]] = None,
-        run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
+        stop: list[str] | None = None,
+        run_manager: AsyncCallbackManagerForLLMRun | None = None,
         **kwargs: Any,
     ) -> str:
         """Return next response."""
@@ -91,16 +91,16 @@ class FakeStreamingListLLM(FakeListLLM):
     chunks in a streaming implementation.
     """
 
-    error_on_chunk_number: Optional[int] = None
+    error_on_chunk_number: int | None = None
     """If set, will raise an exception on the specified chunk number."""
 
     @override
     def stream(
         self,
         input: LanguageModelInput,
-        config: Optional[RunnableConfig] = None,
+        config: RunnableConfig | None = None,
         *,
-        stop: Optional[list[str]] = None,
+        stop: list[str] | None = None,
         **kwargs: Any,
     ) -> Iterator[str]:
         result = self.invoke(input, config)
@@ -119,9 +119,9 @@ class FakeStreamingListLLM(FakeListLLM):
     async def astream(
         self,
         input: LanguageModelInput,
-        config: Optional[RunnableConfig] = None,
+        config: RunnableConfig | None = None,
         *,
-        stop: Optional[list[str]] = None,
+        stop: list[str] | None = None,
         **kwargs: Any,
     ) -> AsyncIterator[str]:
         result = await self.ainvoke(input, config)

langchain_core/language_models/fake_chat_models.py

@@ -1,10 +1,10 @@
-"""Fake ChatModel for testing purposes."""
+"""Fake chat model for testing purposes."""
 
 import asyncio
 import re
 import time
-from collections.abc import AsyncIterator, Iterable, Iterator
-from typing import Any, Optional, Union, cast
+from collections.abc import AsyncIterator, Iterator
+from typing import Any, Literal, cast
 
 from typing_extensions import override
 
@@ -16,18 +16,14 @@ from langchain_core.language_models.chat_models import BaseChatModel, SimpleChat
 from langchain_core.messages import AIMessage, AIMessageChunk, BaseMessage
 from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult
 from langchain_core.runnables import RunnableConfig
-from langchain_core.v1.chat_models import BaseChatModel as BaseChatModelV1
-from langchain_core.v1.messages import AIMessage as AIMessageV1
-from langchain_core.v1.messages import AIMessageChunk as AIMessageChunkV1
-from langchain_core.v1.messages import MessageV1
 
 
 class FakeMessagesListChatModel(BaseChatModel):
-    """Fake ChatModel for testing purposes."""
+    """Fake chat model for testing purposes."""
 
     responses: list[BaseMessage]
     """List of responses to **cycle** through in order."""
-    sleep: Optional[float] = None
+    sleep: float | None = None
     """Sleep time in seconds between responses."""
     i: int = 0
     """Internally incremented after every model invocation."""
@@ -36,8 +32,8 @@ class FakeMessagesListChatModel(BaseChatModel):
     def _generate(
         self,
         messages: list[BaseMessage],
-        stop: Optional[list[str]] = None,
-        run_manager: Optional[CallbackManagerForLLMRun] = None,
+        stop: list[str] | None = None,
+        run_manager: CallbackManagerForLLMRun | None = None,
         **kwargs: Any,
     ) -> ChatResult:
         if self.sleep is not None:
@@ -61,14 +57,14 @@ class FakeListChatModelError(Exception):
 
 
 class FakeListChatModel(SimpleChatModel):
-    """Fake ChatModel for testing purposes."""
+    """Fake chat model for testing purposes."""
 
     responses: list[str]
     """List of responses to **cycle** through in order."""
-    sleep: Optional[float] = None
+    sleep: float | None = None
     i: int = 0
     """Internally incremented after every model invocation."""
-    error_on_chunk_number: Optional[int] = None
+    error_on_chunk_number: int | None = None
     """If set, raise an error on the specified chunk number during streaming."""
 
     @property
@@ -79,12 +75,13 @@ class FakeListChatModel(SimpleChatModel):
     @override
     def _call(
         self,
-        messages: list[BaseMessage],
-        stop: Optional[list[str]] = None,
-        run_manager: Optional[CallbackManagerForLLMRun] = None,
+        *args: Any,
         **kwargs: Any,
     ) -> str:
-        """First try to lookup in queries, else return 'foo' or 'bar'."""
+        """Return the next response in the list.
+
+        Cycle back to the start if at the end.
+        """
         if self.sleep is not None:
             time.sleep(self.sleep)
         response = self.responses[self.i]
@@ -98,8 +95,8 @@ class FakeListChatModel(SimpleChatModel):
     def _stream(
         self,
         messages: list[BaseMessage],
-        stop: Union[list[str], None] = None,
-        run_manager: Union[CallbackManagerForLLMRun, None] = None,
+        stop: list[str] | None = None,
+        run_manager: CallbackManagerForLLMRun | None = None,
         **kwargs: Any,
     ) -> Iterator[ChatGenerationChunk]:
         response = self.responses[self.i]
@@ -116,14 +113,19 @@ class FakeListChatModel(SimpleChatModel):
             ):
                 raise FakeListChatModelError
 
-            yield ChatGenerationChunk(message=AIMessageChunk(content=c))
+            chunk_position: Literal["last"] | None = (
+                "last" if i_c == len(response) - 1 else None
+            )
+            yield ChatGenerationChunk(
+                message=AIMessageChunk(content=c, chunk_position=chunk_position)
+            )
 
     @override
     async def _astream(
         self,
         messages: list[BaseMessage],
-        stop: Union[list[str], None] = None,
-        run_manager: Union[AsyncCallbackManagerForLLMRun, None] = None,
+        stop: list[str] | None = None,
+        run_manager: AsyncCallbackManagerForLLMRun | None = None,
         **kwargs: Any,
     ) -> AsyncIterator[ChatGenerationChunk]:
         response = self.responses[self.i]
@@ -139,7 +141,12 @@ class FakeListChatModel(SimpleChatModel):
                 and i_c == self.error_on_chunk_number
             ):
                 raise FakeListChatModelError
-            yield ChatGenerationChunk(message=AIMessageChunk(content=c))
+            chunk_position: Literal["last"] | None = (
+                "last" if i_c == len(response) - 1 else None
+            )
+            yield ChatGenerationChunk(
+                message=AIMessageChunk(content=c, chunk_position=chunk_position)
+            )
 
     @property
     @override
@@ -151,27 +158,33 @@ class FakeListChatModel(SimpleChatModel):
     def batch(
         self,
         inputs: list[Any],
-        config: Optional[Union[RunnableConfig, list[RunnableConfig]]] = None,
+        config: RunnableConfig | list[RunnableConfig] | None = None,
         *,
         return_exceptions: bool = False,
         **kwargs: Any,
-    ) -> list[BaseMessage]:
+    ) -> list[AIMessage]:
         if isinstance(config, list):
-            return [self.invoke(m, c, **kwargs) for m, c in zip(inputs, config)]
+            return [
+                self.invoke(m, c, **kwargs)
+                for m, c in zip(inputs, config, strict=False)
+            ]
         return [self.invoke(m, config, **kwargs) for m in inputs]
 
     @override
     async def abatch(
         self,
         inputs: list[Any],
-        config: Optional[Union[RunnableConfig, list[RunnableConfig]]] = None,
+        config: RunnableConfig | list[RunnableConfig] | None = None,
         *,
         return_exceptions: bool = False,
         **kwargs: Any,
-    ) -> list[BaseMessage]:
+    ) -> list[AIMessage]:
         if isinstance(config, list):
             # do Not use an async iterator here because need explicit ordering
-            return [await self.ainvoke(m, c, **kwargs) for m, c in zip(inputs, config)]
+            return [
+                await self.ainvoke(m, c, **kwargs)
+                for m, c in zip(inputs, config, strict=False)
+            ]
         # do Not use an async iterator here because need explicit ordering
         return [await self.ainvoke(m, config, **kwargs) for m in inputs]
 
@@ -183,8 +196,8 @@ class FakeChatModel(SimpleChatModel):
     def _call(
         self,
         messages: list[BaseMessage],
-        stop: Optional[list[str]] = None,
-        run_manager: Optional[CallbackManagerForLLMRun] = None,
+        stop: list[str] | None = None,
+        run_manager: CallbackManagerForLLMRun | None = None,
         **kwargs: Any,
     ) -> str:
         return "fake response"
@@ -193,8 +206,8 @@ class FakeChatModel(SimpleChatModel):
     async def _agenerate(
         self,
         messages: list[BaseMessage],
-        stop: Optional[list[str]] = None,
-        run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
+        stop: list[str] | None = None,
+        run_manager: AsyncCallbackManagerForLLMRun | None = None,
         **kwargs: Any,
     ) -> ChatResult:
         output_str = "fake response"
@@ -215,34 +228,36 @@ class GenericFakeChatModel(BaseChatModel):
     """Generic fake chat model that can be used to test the chat model interface.
 
     * Chat model should be usable in both sync and async tests
-    * Invokes on_llm_new_token to allow for testing of callback related code for new
-      tokens.
+    * Invokes `on_llm_new_token` to allow for testing of callback related code for new
+        tokens.
     * Includes logic to break messages into message chunk to facilitate testing of
-      streaming.
+        streaming.
+
     """
 
-    messages: Iterator[Union[AIMessage, str]]
+    messages: Iterator[AIMessage | str]
     """Get an iterator over messages.
 
     This can be expanded to accept other types like Callables / dicts / strings
     to make the interface more generic if needed.
 
-    Note: if you want to pass a list, you can use `iter` to convert it to an iterator.
+    !!! note
+        if you want to pass a list, you can use `iter` to convert it to an iterator.
+
+    !!! warning
+        Streaming is not implemented yet. We should try to implement it in the future by
+        delegating to invoke and then breaking the resulting output into message chunks.
 
-    Please note that streaming is not implemented yet. We should try to implement it
-    in the future by delegating to invoke and then breaking the resulting output
-    into message chunks.
     """
 
     @override
     def _generate(
         self,
         messages: list[BaseMessage],
-        stop: Optional[list[str]] = None,
-        run_manager: Optional[CallbackManagerForLLMRun] = None,
+        stop: list[str] | None = None,
+        run_manager: CallbackManagerForLLMRun | None = None,
         **kwargs: Any,
     ) -> ChatResult:
-        """Top Level call."""
         message = next(self.messages)
         message_ = AIMessage(content=message) if isinstance(message, str) else message
         generation = ChatGeneration(message=message_)
@@ -251,11 +266,10 @@ class GenericFakeChatModel(BaseChatModel):
     def _stream(
         self,
         messages: list[BaseMessage],
-        stop: Optional[list[str]] = None,
-        run_manager: Optional[CallbackManagerForLLMRun] = None,
+        stop: list[str] | None = None,
+        run_manager: CallbackManagerForLLMRun | None = None,
         **kwargs: Any,
     ) -> Iterator[ChatGenerationChunk]:
-        """Stream the output of the model."""
         chat_result = self._generate(
             messages, stop=stop, run_manager=run_manager, **kwargs
         )
@@ -286,10 +300,16 @@ class GenericFakeChatModel(BaseChatModel):
 
             content_chunks = cast("list[str]", re.split(r"(\s)", content))
 
-            for token in content_chunks:
+            for idx, token in enumerate(content_chunks):
                 chunk = ChatGenerationChunk(
                     message=AIMessageChunk(content=token, id=message.id)
                 )
+                if (
+                    idx == len(content_chunks) - 1
+                    and isinstance(chunk.message, AIMessageChunk)
+                    and not message.additional_kwargs
+                ):
+                    chunk.message.chunk_position = "last"
                 if run_manager:
                     run_manager.on_llm_new_token(token, chunk=chunk)
                 yield chunk
@@ -355,85 +375,19 @@ class ParrotFakeChatModel(BaseChatModel):
     """Generic fake chat model that can be used to test the chat model interface.
 
     * Chat model should be usable in both sync and async tests
+
     """
 
     @override
     def _generate(
         self,
         messages: list[BaseMessage],
-        stop: Optional[list[str]] = None,
-        run_manager: Optional[CallbackManagerForLLMRun] = None,
+        stop: list[str] | None = None,
+        run_manager: CallbackManagerForLLMRun | None = None,
         **kwargs: Any,
     ) -> ChatResult:
-        """Top Level call."""
         return ChatResult(generations=[ChatGeneration(message=messages[-1])])
 
     @property
     def _llm_type(self) -> str:
         return "parrot-fake-chat-model"
-
-
-class GenericFakeChatModelV1(BaseChatModelV1):
-    """Generic fake chat model that can be used to test the chat model interface."""
-
-    messages: Optional[Iterator[Union[AIMessageV1, str]]] = None
-    message_chunks: Optional[Iterable[Union[AIMessageChunkV1, str]]] = None
-
-    @override
-    def _invoke(
-        self,
-        messages: list[MessageV1],
-        run_manager: Optional[CallbackManagerForLLMRun] = None,
-        **kwargs: Any,
-    ) -> AIMessageV1:
-        """Top Level call."""
-        if self.messages is None:
-            error_msg = "Messages iterator is not set."
-            raise ValueError(error_msg)
-        message = next(self.messages)
-        return AIMessageV1(content=message) if isinstance(message, str) else message
-
-    @override
-    def _stream(
-        self,
-        messages: list[MessageV1],
-        run_manager: Optional[CallbackManagerForLLMRun] = None,
-        **kwargs: Any,
-    ) -> Iterator[AIMessageChunkV1]:
-        """Top Level call."""
-        if self.message_chunks is None:
-            error_msg = "Message chunks iterator is not set."
-            raise ValueError(error_msg)
-        for chunk in self.message_chunks:
-            if isinstance(chunk, str):
-                yield AIMessageChunkV1(chunk)
-            else:
-                yield chunk
-
-    @property
-    def _llm_type(self) -> str:
-        return "generic-fake-chat-model"
-
-
-class ParrotFakeChatModelV1(BaseChatModelV1):
-    """Generic fake chat model that can be used to test the chat model interface.
-
-    * Chat model should be usable in both sync and async tests
-    """
-
-    @override
-    def _invoke(
-        self,
-        messages: list[MessageV1],
-        stop: Optional[list[str]] = None,
-        run_manager: Optional[CallbackManagerForLLMRun] = None,
-        **kwargs: Any,
-    ) -> AIMessageV1:
-        """Top Level call."""
-        if isinstance(messages[-1], AIMessageV1):
-            return messages[-1]
-        return AIMessageV1(content=messages[-1].content)
-
-    @property
-    def _llm_type(self) -> str:
-        return "parrot-fake-chat-model"