forked from Archives/langchain
You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
67 lines
2.1 KiB
Python
67 lines
2.1 KiB
Python
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
|