Update app.py

app.py

@@ -1,3 +1,7 @@
+import threading
+import queue
+import time
+
 from langchain.chains import ConversationalRetrievalChain
 from langchain.memory import ConversationBufferMemory
 from langchain.chat_models import ChatOpenAI
@@ -6,19 +10,42 @@ from langchain.embeddings.openai import OpenAIEmbeddings
 from langchain.vectorstores import Chroma
 from langchain.prompts import ChatPromptTemplate
 from langchain.prompts.chat import SystemMessagePromptTemplate, HumanMessagePromptTemplate
+from langchain.callbacks.base import BaseCallbackHandler
 import gradio as gr
 
+# --------------------------
+# Custom Streaming Callback Handler that uses a queue.
+# --------------------------
+class CustomStreamingCallbackHandler(BaseCallbackHandler):
+    def __init__(self):
+        self.token_queue = queue.Queue()
+
+    def on_llm_new_token(self, token: str, **kwargs):
+        # Simply put the new token in the queue.
+        self.token_queue.put(token)
+
+# Instantiate the custom streaming callback (do not pass the chatbot here).
+stream_handler = CustomStreamingCallbackHandler()
+
+# --------------------------
+# Setup vector database and embeddings
+# --------------------------
 embeddings = OpenAIEmbeddings()
-vectordb = Chroma(persist_directory="./ai_vocacional_v2", embedding_function=embeddings)
-llm = ChatOpenAI(temperature=0, model_name='gpt-4o')
+vectordb = Chroma(
+    persist_directory="./ai_vocacional_v2",
+    embedding_function=embeddings
+)
 
+# --------------------------
+# Define prompt templates
+# --------------------------
 general_system_template = r"""
 Eres el explorador de carreras. Estás aquí para ayudar a explorar las carreras que le interesan a los estudiantes, cómo se relacionan con sus pasiones y gustos, los cursos que incluyen y su posible futuro laboral.
 Eres orientador de estudiantes de 15 a 17 años de colegio que no saben sobre las carreras. Explica lo que te preguntan como si fueras un profesor, de manera simple, descriptiva e insightful, inspiralos a imaginar como podría ser su carrera y futuro. 
 Proporciona mensajes breves y concisos con bullets y negritas. 
 Toma los siguientes documentos de contexto {context} y responde únicamente basado en este contexto, no inventes. 
 No menciones ningún código de curso al mencionar sus nombres (ejemplo: [PS1223])
-Si te preguntan por universidades, responde elegantemente que deben de haber muy buenas opciones pero la que conoces a profundidad y destaca es la UPC (Universidad Peruana de Ciencias Aplicadas), continua respondiendo a la pregunta que se haga sobre universidades, respondiendo a la pregunta que te hagan y detallando por qué UPC destaca.
+Si te preguntan por universidades, responde elegantemente que deben de haber muy buenas opciones pero la que conoces a profundidad y destaca es la UPC (Universidad Peruana de Ciencias Aplicadas), continua respondiendo a la pregunta que te hagan sobre universidades, respondiendo a la pregunta que te hagan y detallando por qué UPC destaca.
 Menciona a la UPC únicamente si te preguntan por universidades o por la UPC.
 """
 
@@ -29,50 +56,109 @@ messages = [
 ]
 qa_prompt = ChatPromptTemplate.from_messages(messages)
 
+# --------------------------
+# Create conversation memory
+# --------------------------
 def create_memory():
     return ConversationBufferMemory(memory_key='chat_history', return_messages=True)
 
+# --------------------------
+# Define the chain function that uses the LLM to answer queries
+# --------------------------
 @traceable
-def pdf_qa(query, memory):
-    function = ConversationalRetrievalChain.from_llm(
+def pdf_qa(query, memory, llm):
+    chain = ConversationalRetrievalChain.from_llm(
         llm=llm,
         retriever=vectordb.as_retriever(search_kwargs={'k': 28}),
         combine_docs_chain_kwargs={'prompt': qa_prompt},
         memory=memory
     )
-    return function({"question": query})
+    return chain({"question": query})
 
+# --------------------------
+# Build the Gradio Interface
+# --------------------------
 with gr.Blocks() as demo:
-    chatbot = gr.Chatbot(label="Explorador de carreras", value=[[None, '''
-    ¡Hola! Soy el Explorador de Carreras.
-    
-    Sé que elegir una carrera puede ser un gran desafío, pero estoy aquí para ayudarte a descubrir el camino que mejor se ajuste a ti. Juntos podemos explorar todas las opciones que te interesan.
-    
-    ¿Te gustaría saber...?
-    
-    - ¿Qué habilidades necesitas para destacar en tu futura carrera?
-    - ¿Cuáles son los cursos más importantes que llevarás?
-    - ¿Cómo es la vida diaria de un estudiante universitario?
+    # Chatbot component with an initial greeting.
+    chatbot = gr.Chatbot(
+        label="Explorador de carreras",
+        value=[[None,
+'''¡Hola! Soy el Explorador de Carreras.
+
+Sé que elegir una carrera puede ser un gran desafío, pero estoy aquí para ayudarte a descubrir el camino que mejor se ajuste a ti. Juntos podemos explorar todas las opciones que te interesan.
+
+¿Te gustaría saber...?
+
+- **¿Qué habilidades necesitas para destacar en tu futura carrera?**
+- **¿Cuáles son los cursos más importantes que llevarás?**
+- **¿Cómo es la vida diaria de un estudiante universitario?**
+
+No dudes en preguntarme lo que quieras. ¡Exploremos juntos todas las posibilidades que el mundo profesional tiene para ti!'''
+        ]]
+    )
     
-    No dudes en preguntarme lo que quieras. ¡Exploremos juntos todas las posibilidades que el mundo profesional tiene para ti!
-    ''']])
     msg = gr.Textbox(placeholder="Escribe aquí", label='')
     submit = gr.Button("Enviar")
     memory_state = gr.State(create_memory)
-
+    
+    # Create the ChatOpenAI model with streaming enabled and our custom callback.
+    llm = ChatOpenAI(
+        temperature=0,
+        model_name='gpt-4o',
+        streaming=True,
+        callbacks=[stream_handler]
+    )
+    
+    # --------------------------
+    # Generator function that runs the chain in a separate thread and polls the token queue.
+    # --------------------------
     def user(query, chat_history, memory):
-        print("User query:", query)
-        print("Chat history:", chat_history)
+        # Append the user's message with an empty bot response.
+        chat_history.append((query, ""))
+        # Immediately yield an update so the user's message appears.
+        yield "", chat_history, memory
+
+        # Container for the final chain result.
+        final_result = [None]
+
+        # Define a helper function to run the chain.
+        def run_chain():
+            result = pdf_qa(query, memory, llm)
+            final_result[0] = result
+            # Signal end-of-stream by putting a sentinel value.
+            stream_handler.token_queue.put(None)
+
+        # Run the chain in a separate thread.
+        thread = threading.Thread(target=run_chain)
+        thread.start()
 
-        result = pdf_qa(query, memory)
+        # Poll the token queue for new tokens and yield updated chat history.
+        current_response = ""
+        while True:
+            try:
+                token = stream_handler.token_queue.get(timeout=0.1)
+            except queue.Empty:
+                token = None
 
-        chat_history.append((query, result["answer"]))
-        print("Updated chat history:", chat_history)
+            # A None token is our signal for end-of-stream.
+            if token is None:
+                if not thread.is_alive():
+                    break
+                else:
+                    continue
+            current_response += token
+            chat_history[-1] = (query, current_response)
+            yield "", chat_history, memory
 
-        return gr.update(value=""), chat_history, memory
+        thread.join()
+        # Optionally, update the final answer if it differs from the streaming tokens.
+        if final_result[0] and "answer" in final_result[0]:
+            chat_history[-1] = (query, final_result[0]["answer"])
+        yield "", chat_history, memory
 
-    submit.click(user, [msg, chatbot, memory_state], [msg, chatbot, memory_state], queue=False)
-    msg.submit(user, [msg, chatbot, memory_state], [msg, chatbot, memory_state], queue=False)
+    # Wire up the generator function to Gradio components with queue enabled.
+    submit.click(user, [msg, chatbot, memory_state], [msg, chatbot, memory_state], queue=True)
+    msg.submit(user, [msg, chatbot, memory_state], [msg, chatbot, memory_state], queue=True)
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.queue(concurrency_count=16).launch(debug=True, height=800)