add callback

app.py

@@ -1,4 +1,6 @@
 # from typing import Any, Coroutine
+from uuid import UUID
+from langchain.schema.agent import AgentAction, AgentFinish
 import openai
 import os
 # from langchain.vectorstores import Chroma
@@ -40,7 +42,7 @@ import gradio as gr
 import time
 
 import glob
-from typing import Any, List, Mapping, Optional
+from typing import Any, Dict, List, Mapping, Optional
 from multiprocessing import Pool
 from tqdm import tqdm
 from pygame import mixer
@@ -94,6 +96,59 @@ from interpreter.code_interpreter import CodeInterpreter
 
 import regex
 
+from langchain.callbacks.base import BaseCallbackHandler
+from collections.abc import Generator
+from queue import Queue, Empty
+from threading import Thread
+
+
+class QueueCallback(BaseCallbackHandler):
+    """Callback handler for streaming LLM responses to a queue."""
+
+    def __init__(self, q):
+        self.q = q
+
+    def on_llm_new_token(self, token: str, **kwargs: any) -> None:
+        self.q.put(token)
+    
+    def on_agent_finish(self, finish: AgentFinish, *, run_id: UUID, parent_run_id: UUID | None = None, **kwargs: Any) -> Any:
+        self.q.put(super().on_agent_finish(finish, run_id=run_id, parent_run_id=parent_run_id, **kwargs))
+
+    def on_agent_action(self, action: AgentAction, *, run_id: UUID, parent_run_id: UUID | None = None, **kwargs: Any) -> Any:
+        self.q.put(super().on_agent_action(action, run_id=run_id, parent_run_id=parent_run_id, **kwargs))
+
+    def on_llm_end(self, *args, **kwargs: any) -> None:
+        return self.q.empty()
+
+def stream(input_text) -> Generator:
+    # Create a Queue
+    q = Queue()
+    job_done = object()
+
+    # Create a funciton to call - this will run in a thread
+
+    def task():
+        resp = agent.run(input_text, callbacks=QueueCallback(q))
+        q.put(job_done)
+
+    # Create a thread and start the function
+    t = Thread(target=task)
+    t.start()
+
+    content = ""
+
+    # Get each new token from the queue and yield for our generator
+    while True:
+        try:
+            next_token = q.get(True, timeout=1)
+            if next_token is job_done:
+                break
+            content += next_token
+            yield next_token, content
+        except Empty:
+            continue
+
+
 global CurrentAgent
 CurrentAgent = os.environ["agent_type"]
 
@@ -1549,13 +1604,94 @@ def chathmi4(message, history2):
     # yield ["", history2, "SUBMIT", "STOP"]
     try:
         if agent is not None:
+
             response = agent.run(message)
+            # test callback
+            
+            # temp = []
+            # for next_token, content in stream(message):
+            #     temp = temp + content
+            #     history_int = history2 + [(None, temp)]
+            #     yield(None, history_int, None, None)
         elif agent is None:
             response = asyncio.run(start_playwright(message))
 
         time.sleep(0.1)
         history2 = history2 + [(None, response)]
         yield ["", history2, gr.update(visible = True), gr.update(visible = False)]
+
+
+        # yield ["", history2, None, None]
+        print ("response of chatbot:", response)
+        # real_content = response[-1:]
+        # print("real_content", real_content)
+        # try:
+        #     # temp = response.split("(sandbox:/")[1] # (sandbox:/sample-20230805-0807.wav)
+        #     file_name = CheckFileinResp(response)
+        #     print("file_name:", file_name)
+        #     if file_name != "N/A":
+        #         history2 = history2 + [(None, (file_name,))]
+        #         Filename_Chatbot = file_name    
+        #         yield ["", history2, "SUBMIT", "STOP"]        
+        # except Exception as e:
+        #     print("No need to add file in chatbot:", e)        
+
+        try:
+            # temp = response.split("(sandbox:/")[1] # (sandbox:/sample-20230805-0807.wav)
+            file_names = CheckFileinResp(response)
+            print("file_name:", file_names)
+            if file_names != []:
+                for file_name in file_names:
+                    if file_name != "":
+                        history2 = history2 + [(None, (file_name, file_name))]
+                        Filename_Chatbot = file_name
+                        yield ["", history2, "SUBMIT", "STOP"]
+            else:
+                print("No File Found in Response")
+        except Exception as e:
+            print("No need to add file in chatbot:", e)    
+
+
+
+    except Exception as e:
+        print("chathmi4 error:", e)
+
+    # history = history + [(message, None)]
+        
+    print("History2: ", history2)
+    print("-" * 20)
+    print("-" * 20)
+
+def chathmi5(message, history2):
+    global last_request
+    global Filename_Chatbot
+    global agent
+    # global ChatbotHistory
+    print("Input Message:", message)
+    
+    last_request = message
+    history2 = history2 + [(message, None)]
+    yield ["", history2, gr.update(visible = False), gr.update(visible = True)]
+    # yield ["", history2, "SUBMIT", "STOP"]
+    try:
+        if agent is not None:
+
+            # response = agent.run(message)
+            # test callback
+            
+            temp = []
+            for next_token, content in stream(message):
+                temp = temp + content
+                history_int = history2 + [(None, temp)]
+                yield(None, history_int, None, None)
+        elif agent is None:
+            response = asyncio.run(start_playwright(message))
+
+        time.sleep(0.1)
+        history2 = history_int
+        yield ["", history2, gr.update(visible = True), gr.update(visible = False)]
+
+        
         # yield ["", history2, None, None]
         print ("response of chatbot:", response)
         # real_content = response[-1:]
@@ -1597,6 +1733,7 @@ def chathmi4(message, history2):
     print("-" * 20)
     print("-" * 20)
 
+
 def chatremote(message, history2):
     global last_request
     global Filename_Chatbot
@@ -1987,7 +2124,8 @@ with gr.Blocks() as demo:
     #     then(HMI_Wait, None, [submit_button, stop_button])
     # inf4 = inputtext.submit(chathmi4, [inputtext, chatbot], [inputtext, chatbot, submit_button, stop_button])
     ''' open ai | new'''
-    inf4 = inputtext.submit(chathmi4, [inputtext, chatbot], [inputtext, chatbot, submit_button, stop_button]).\
+    # chathmi4 = normal, chathmi5 = callback
+    inf4 = inputtext.submit(chathmi5, [inputtext, chatbot], [inputtext, chatbot, submit_button, stop_button]).\
         success(playsound1, None, voice_output, queue=True)#.\
         # success(ClearAudio, None, voice_output)