Toggle Menu
Prev Up Next

# Source code for langchain_core.language_models.fake

    import asyncio
    import time
    from typing import Any, AsyncIterator, Iterator, List, Mapping, Optional

    from langchain_core.callbacks import (
        AsyncCallbackManagerForLLMRun,
        CallbackManagerForLLMRun,
    )
    from langchain_core.language_models import LanguageModelInput
    from langchain_core.language_models.llms import LLM
    from langchain_core.runnables import RunnableConfig

    class FakeListLLM(LLM):
        """Fake LLM for testing purposes."""

        responses: List[str]
        sleep: Optional[float] = None
        i: int = 0

        @property
        def _llm_type(self) -> str:
            """Return type of llm."""
            return "fake-list"

        def _call(
            self,
            prompt: str,
            stop: Optional[List[str]] = None,
            run_manager: Optional[CallbackManagerForLLMRun] = None,
            **kwargs: Any,
        ) -> str:
            """Return next response"""
            response = self.responses[self.i]
            if self.i < len(self.responses) - 1:
                self.i += 1
            else:
                self.i = 0
            return response

        async def _acall(
            self,
            prompt: str,
            stop: Optional[List[str]] = None,
            run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
            **kwargs: Any,
        ) -> str:
            """Return next response"""
            response = self.responses[self.i]
            if self.i < len(self.responses) - 1:
                self.i += 1
            else:
                self.i = 0
            return response

        @property
        def _identifying_params(self) -> Mapping[str, Any]:
            return {"responses": self.responses}

    class FakeStreamingListLLM(FakeListLLM):
        """Fake streaming list LLM for testing purposes."""

        error_on_chunk_number: Optional[int] = None

        def stream(
            self,
            input: LanguageModelInput,
            config: Optional[RunnableConfig] = None,
            *,
            stop: Optional[List[str]] = None,
            **kwargs: Any,
        ) -> Iterator[str]:
            result = self.invoke(input, config)
            for i_c, c in enumerate(result):
                if self.sleep is not None:
                    time.sleep(self.sleep)

                if (
                    self.error_on_chunk_number is not None
                    and i_c == self.error_on_chunk_number
                ):
                    raise Exception("Fake error")
                yield c

        async def astream(
            self,
            input: LanguageModelInput,
            config: Optional[RunnableConfig] = None,
            *,
            stop: Optional[List[str]] = None,
            **kwargs: Any,
        ) -> AsyncIterator[str]:
            result = await self.ainvoke(input, config)
            for i_c, c in enumerate(result):
                if self.sleep is not None:
                    await asyncio.sleep(self.sleep)

                if (
                    self.error_on_chunk_number is not None
                    and i_c == self.error_on_chunk_number
                ):
                    raise Exception("Fake error")
                yield c