forked from Archives/langchain
Add AsyncIteratorCallbackHandler (#2329)
parent
6e4e7d2637
commit
6f39e88a2c
@ -0,0 +1,66 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
from typing import Any, AsyncIterator, Dict, List, Literal, Union, cast
|
||||||
|
|
||||||
|
from langchain.callbacks.base import AsyncCallbackHandler
|
||||||
|
from langchain.schema import LLMResult
|
||||||
|
|
||||||
|
# TODO If used by two LLM runs in parallel this won't work as expected
|
||||||
|
|
||||||
|
|
||||||
|
class AsyncIteratorCallbackHandler(AsyncCallbackHandler):
|
||||||
|
"""Callback handler that returns an async iterator."""
|
||||||
|
|
||||||
|
queue: asyncio.Queue[str]
|
||||||
|
|
||||||
|
done: asyncio.Event
|
||||||
|
|
||||||
|
@property
|
||||||
|
def always_verbose(self) -> bool:
|
||||||
|
return True
|
||||||
|
|
||||||
|
def __init__(self) -> None:
|
||||||
|
self.queue = asyncio.Queue()
|
||||||
|
self.done = asyncio.Event()
|
||||||
|
|
||||||
|
async def on_llm_start(
|
||||||
|
self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
|
||||||
|
) -> None:
|
||||||
|
# If two calls are made in a row, this resets the state
|
||||||
|
self.done.clear()
|
||||||
|
|
||||||
|
async def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
|
||||||
|
self.queue.put_nowait(token)
|
||||||
|
|
||||||
|
async def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
|
||||||
|
self.done.set()
|
||||||
|
|
||||||
|
async def on_llm_error(
|
||||||
|
self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
|
||||||
|
) -> None:
|
||||||
|
self.done.set()
|
||||||
|
|
||||||
|
# TODO implement the other methods
|
||||||
|
|
||||||
|
async def aiter(self) -> AsyncIterator[str]:
|
||||||
|
while not self.queue.empty() or not self.done.is_set():
|
||||||
|
# Wait for the next token in the queue,
|
||||||
|
# but stop waiting if the done event is set
|
||||||
|
done, _ = await asyncio.wait(
|
||||||
|
[
|
||||||
|
asyncio.ensure_future(self.queue.get()),
|
||||||
|
asyncio.ensure_future(self.done.wait()),
|
||||||
|
],
|
||||||
|
return_when=asyncio.FIRST_COMPLETED,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Extract the value of the first completed task
|
||||||
|
token_or_done = cast(Union[str, Literal[True]], done.pop().result())
|
||||||
|
|
||||||
|
# If the extracted value is the boolean True, the done event was set
|
||||||
|
if token_or_done is True:
|
||||||
|
break
|
||||||
|
|
||||||
|
# Otherwise, the extracted value is a token, which we yield
|
||||||
|
yield token_or_done
|
Loading…
Reference in New Issue