Spaces:
Running
Running
File size: 561 Bytes
8578df0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 |
# adapted from https://github.com/hwchase17/langchain/issues/2428#issuecomment-1512280045
from queue import Queue
from typing import Any
from langchain.callbacks.base import BaseCallbackHandler
class QueueCallback(BaseCallbackHandler):
"""Callback handler for streaming LLM responses to a queue."""
def __init__(self, queue: Queue):
self.queue = queue
def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
self.queue.put(token)
def on_llm_end(self, *args, **kwargs: Any) -> None:
return self.queue.empty()
|