iKING-ROC commited on
Commit
40a5cfe
1 Parent(s): 7e38677

Upload callback.py

Browse files
Files changed (1) hide show
  1. callback.py +18 -0
callback.py ADDED
@@ -0,0 +1,18 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # adapted from https://github.com/hwchase17/langchain/issues/2428#issuecomment-1512280045
2
+ from queue import Queue
3
+ from typing import Any
4
+
5
+ from langchain.callbacks.base import BaseCallbackHandler
6
+
7
+
8
+ class QueueCallback(BaseCallbackHandler):
9
+ """Callback handler for streaming LLM responses to a queue."""
10
+
11
+ def __init__(self, queue: Queue):
12
+ self.queue = queue
13
+
14
+ def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
15
+ self.queue.put(token)
16
+
17
+ def on_llm_end(self, *args, **kwargs: Any) -> None:
18
+ return self.queue.empty()