diff --git a/langchain/llms/fake.py b/langchain/llms/fake.py index 15fbab5e..5700e82f 100644 --- a/langchain/llms/fake.py +++ b/langchain/llms/fake.py @@ -1,7 +1,10 @@ """Fake LLM wrapper for testing purposes.""" from typing import Any, List, Mapping, Optional -from langchain.callbacks.manager import CallbackManagerForLLMRun +from langchain.callbacks.manager import ( + AsyncCallbackManagerForLLMRun, + CallbackManagerForLLMRun, +) from langchain.llms.base import LLM @@ -22,7 +25,18 @@ class FakeListLLM(LLM): stop: Optional[List[str]] = None, run_manager: Optional[CallbackManagerForLLMRun] = None, ) -> str: - """First try to lookup in queries, else return 'foo' or 'bar'.""" + """Return next response""" + response = self.responses[self.i] + self.i += 1 + return response + + async def _acall( + self, + prompt: str, + stop: Optional[List[str]] = None, + run_manager: Optional[AsyncCallbackManagerForLLMRun] = None, + ) -> str: + """Return next response""" response = self.responses[self.i] self.i += 1 return response