Source code for langchain.chat_models.fake
"""Fake ChatModel for testing purposes."""
import asyncio
import time
from typing import Any, AsyncIterator, Dict, Iterator, List, Optional, Union
from langchain.callbacks.manager import (
    AsyncCallbackManagerForLLMRun,
    CallbackManagerForLLMRun,
)
from langchain.chat_models.base import BaseChatModel, SimpleChatModel
from langchain.schema import ChatResult
from langchain.schema.messages import AIMessageChunk, BaseMessage
from langchain.schema.output import ChatGeneration, ChatGenerationChunk
[docs]class FakeMessagesListChatModel(BaseChatModel):
    responses: List[BaseMessage]
    sleep: Optional[float] = None
    i: int = 0
    def _generate(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> ChatResult:
        response = self.responses[self.i]
        if self.i < len(self.responses) - 1:
            self.i += 1
        else:
            self.i = 0
        generation = ChatGeneration(message=response)
        return ChatResult(generations=[generation])
    @property
    def _llm_type(self) -> str:
        return "fake-messages-list-chat-model"
[docs]class FakeListChatModel(SimpleChatModel):
    """Fake ChatModel for testing purposes."""
    responses: List
    sleep: Optional[float] = None
    i: int = 0
    @property
    def _llm_type(self) -> str:
        return "fake-list-chat-model"
    def _call(
        self,
        messages: List[BaseMessage],