codekingpro/portable-devtools

1import asyncio2import time3from typing import Any, AsyncIterator, Iterator, List, Mapping, Optional4 5from langchain_core.callbacks import (6    AsyncCallbackManagerForLLMRun,7    CallbackManagerForLLMRun,8)9from langchain_core.language_models import LanguageModelInput10from langchain_core.language_models.llms import LLM11from langchain_core.runnables import RunnableConfig12 13 14class FakeListLLM(LLM):15    """Fake LLM for testing purposes."""16 17    responses: List[str]18    sleep: Optional[float] = None19    i: int = 020 21    @property22    def _llm_type(self) -> str:23        """Return type of llm."""24        return "fake-list"25 26    def _call(27        self,28        prompt: str,29        stop: Optional[List[str]] = None,30        run_manager: Optional[CallbackManagerForLLMRun] = None,31        **kwargs: Any,32    ) -> str:33        """Return next response"""34        response = self.responses[self.i]35        if self.i < len(self.responses) - 1:36            self.i += 137        else:38            self.i = 039        return response40 41    async def _acall(42        self,43        prompt: str,44        stop: Optional[List[str]] = None,45        run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,46        **kwargs: Any,47    ) -> str:48        """Return next response"""49        response = self.responses[self.i]50        if self.i < len(self.responses) - 1:51            self.i += 152        else:53            self.i = 054        return response55 56    @property57    def _identifying_params(self) -> Mapping[str, Any]:58        return {"responses": self.responses}59 60 61class FakeStreamingListLLM(FakeListLLM):62    """Fake streaming list LLM for testing purposes."""63 64    def stream(65        self,66        input: LanguageModelInput,67        config: Optional[RunnableConfig] = None,68        *,69        stop: Optional[List[str]] = None,70        **kwargs: Any,71    ) -> Iterator[str]:72        result = self.invoke(input, config)73        for c in result:74            if self.sleep is not None:75                time.sleep(self.sleep)76            yield c77 78    async def astream(79        self,80        input: LanguageModelInput,81        config: Optional[RunnableConfig] = None,82        *,83        stop: Optional[List[str]] = None,84        **kwargs: Any,85    ) -> AsyncIterator[str]:86        result = await self.ainvoke(input, config)87        for c in result:88            if self.sleep is not None:89                await asyncio.sleep(self.sleep)90            yield c91