Upload folder using huggingface_hub

app.py

@@ -16,7 +16,7 @@ graph = generate_graph()
 
 def init():
     load_dotenv(override=True)
-    
+
     if(os.path.isdir("./chromadb-store")):
         logger.warning("database exists locally. not downloading.")
         return
@@ -39,6 +39,17 @@ def init_session():
     return str(uuid.uuid4())
 
 
+async def chat_wrapper(message, history, thread_id, debug):
+    if debug:
+        # Yield each message as it's generated
+        async for chunk in chat_streaming(message, history, thread_id):
+            yield chunk
+    else:
+        # Yield full response once
+        response = chat(message, history, thread_id)
+        yield response
+
+
 def chat(message, history, thread_id):
     print("received chat message for thread:", thread_id)
     config = {"configurable": {"thread_id": thread_id}}
@@ -47,6 +58,15 @@ def chat(message, history, thread_id):
     )
     return response["messages"][-1].content
 
+async def chat_streaming(message, history, thread_id):
+    state = {"messages": (history or []) + [{"role": "user", "content": message}]}
+    config = {"configurable": {"thread_id": thread_id}}
+    async for step in graph.astream(state,config=config):
+        # LangGraph yields steps like {"respond": ChatState, "post_process": ChatState}
+        for node_output in step.values():
+            yield node_output["messages"][-1].content
+
+    print("received chat message for thread:", thread_id)
 
 thread_id = gr.State(init_session)
 supported_scriptures = "\n - ".join(
@@ -87,9 +107,17 @@ chatbot = gr.Chatbot(
     height=700,
     render_markdown=True,
 )
+
+debug_checkbox = gr.Checkbox(label="Debug (Streaming)", value=False)
+
 chatInterface = gr.ChatInterface(
-    title="Sanatan-AI", fn=chat, additional_inputs=[thread_id], chatbot=chatbot
+    title="Sanatan-AI",
+    fn=chat_wrapper,
+    additional_inputs=[thread_id, debug_checkbox],
+    chatbot=chatbot,
 )
+
+
 # initializze database
 init()
 chatInterface.launch()