Add asyncio support for LLM (OpenAI), Chain (LLMChain, LLMMathChain), and Agent #841
Conversation
agola11
changed the title
langchain/agents/agent.py
Outdated
| observation = ( | ||
| await tool.coroutine(output.tool_input) | ||
| if tool.coroutine | ||
| else tool.func(output.tool_input) |
There was a problem hiding this comment.
This should probably be called with one of the asyncio methods for executing sync functions in another thread. If left like this it blocks the event loop for the duration of the tool call
There was a problem hiding this comment.
asyncio.get_event_loop().run_in_executor(None, tool.func, output.tool_input)| color = color_mapping[output.tool] | ||
| return_direct = tool.return_direct | ||
| except (KeyboardInterrupt, Exception) as e: | ||
| self.callback_manager.on_tool_error(e, verbose=self.verbose) |
There was a problem hiding this comment.
Should the callback manager have async methods? If I'm using an async agent I might also want to use an async callback handler
There was a problem hiding this comment.
Rn, callbacks are all synchronous but on our roadmap to change soon
|
|
||
| def apply(self, input_list: List[Dict[str, Any]]) -> List[Dict[str, str]]: | ||
| """Utilize the LLM generate method for speed gains.""" | ||
| response = self.generate(input_list) | ||
| return self.create_outputs(response) | ||
|
|
||
| async def aapply(self, input_list: List[Dict[str, Any]]) -> List[Dict[str, str]]: |
There was a problem hiding this comment.
An unfortunate name ha, but I guess it is what it is
langchain/serpapi.py
Outdated
| url = "https://serpapi.com/search" | ||
| return url, params | ||
|
|
||
| async with aiohttp.ClientSession() as session: |
There was a problem hiding this comment.
This is not recommended, sessions should be reused if possible
There was a problem hiding this comment.
Also, the docs for this feature should probably mention this part of the openai python docs https://github.com/openai/openai-python#async-api
|
Making the async methods live in the same class makes reusing helper methods across both implementations easier |
|
How is this compatible with the global import time
import asyncio
from langchain.callbacks import get_openai_callback
from langchain.llms import OpenAI
async def async_generate(llm):
with get_openai_callback() as cb:
resp = await llm.agenerate(["Hello, how are you?"])
print(resp.generations[0][0].text)
# not consistent as async execution will share the global handler
print(cb.total_tokens)
async def generate_concurrently():
llm = OpenAI(temperature=0.9)
tasks = [async_generate(llm) for _ in range(10)]
await asyncio.gather(*tasks)I am not sure if this implementation is complete without handling callbacks properly |
… and Agent (langchain-ai#841) Supporting asyncio in langchain primitives allows for users to run them concurrently and creates more seamless integration with asyncio-supported frameworks (FastAPI, etc.) Summary of changes: **LLM** * Add `agenerate` and `_agenerate` * Implement in OpenAI by leveraging `client.Completions.acreate` **Chain** * Add `arun`, `acall`, `_acall` * Implement them in `LLMChain` and `LLMMathChain` for now **Agent** * Refactor and leverage async chain and llm methods * Add ability for `Tools` to contain async coroutine * Implement async SerpaPI `arun` Create demo notebook. Open questions: * Should all the async stuff go in separate classes? I've seen both patterns (keeping the same class and having async and sync methods vs. having class separation)
Supporting asyncio in langchain primitives allows for users to run them concurrently and creates more seamless integration with asyncio-supported frameworks (FastAPI, etc.)
Summary of changes:
LLM
agenerateand_agenerateclient.Completions.acreateChain
arun,acall,_acallLLMChainandLLMMathChainfor nowAgent
Toolsto contain async coroutinearunCreate demo notebook.
Open questions: