abdullahmeda commited on
Commit
2a6620e
1 Parent(s): 806f9a8

added callback for streaming response

Browse files
Files changed (1) hide show
  1. callbacks.py +17 -0
callbacks.py ADDED
@@ -0,0 +1,17 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from queue import SimpleQueue
2
+ from langchain.callbacks.base import BaseCallbackHandler
3
+
4
+ job_done = object() # signals the processing is done
5
+
6
+ # https://gist.github.com/mortymike/70711b028311681e5f3c6511031d5d43
7
+ class StreamingGradioCallbackHandler(BaseCallbackHandler):
8
+ """Callback handler for streaming. Only works with LLMs that support streaming."""
9
+
10
+ def __init__(self, q: SimpleQueue):
11
+ self.q = q
12
+
13
+ def on_llm_new_token(self, token: str, **kwargs) -> None:
14
+ self.q.put(token)
15
+
16
+ def on_llm_end(self, *args, **kwargs) -> None:
17
+ self.q.put(job_done)