From 0083ff9c7dbcaeae00d9f3ae19eff937ee522b2d Mon Sep 17 00:00:00 2001 From: Ankush Gola Date: Sat, 10 Jun 2023 16:41:43 -0700 Subject: [PATCH] change _agenerate in to use gather --- langchain/llms/base.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/langchain/llms/base.py b/langchain/llms/base.py index 884c5c9c73b..ae8bb97fd7a 100644 --- a/langchain/llms/base.py +++ b/langchain/llms/base.py @@ -494,13 +494,17 @@ class LLM(BaseLLM): run_manager: Optional[AsyncCallbackManagerForLLMRun] = None, ) -> LLMResult: """Run the LLM on the given prompt and input.""" - generations = [] new_arg_supported = inspect.signature(self._acall).parameters.get("run_manager") - for prompt in prompts: - text = ( + + async def get_completion(prompt): + return ( await self._acall(prompt, stop=stop, run_manager=run_manager) if new_arg_supported else await self._acall(prompt, stop=stop) ) - generations.append([Generation(text=text)]) - return LLMResult(generations=generations) + + generations = await asyncio.gather( + *(get_completion(prompt) for prompt in prompts) + ) + + return LLMResult(generations=[[Generation(text=text)] for text in generations])