Update functions.py

functions.py

@@ -91,50 +91,37 @@ def create_retriever_from_chroma(vectorstore_path="docs/chroma/", search_type='m
     return  retriever
 
 
-def handle_userinput(user_question, custom_graph):
+async def handle_userinput(user_question, custom_graph):
     # Add the user's question to the chat history and display it in the UI
     st.session_state.messages.append({"role": "user", "content": user_question})
     st.chat_message("user").write(user_question)
 
-    # Generate a unique thread ID for the graph's state
+    # Config setup (if required for the graph)
     config = {"configurable": {"thread_id": str(uuid.uuid4())}}
 
     try:
-        # Invoke the custom graph with the input question
-        state_dict =  custom_graph.invoke(
-            {"question": user_question, "steps": []}, config
-        )
+        # Await the asynchronous invocation of the custom graph
+        state_dict = await custom_graph.ainvoke({"question": user_question, "steps": []}, config)
 
-        docs = state_dict["documents"]
+        # Extract documents from the state dictionary
+        docs = state_dict.get("documents", [])
         with st.sidebar:
             st.subheader("Dokumentai, kuriuos Birutė gavo kaip kontekstą")
-            with st.spinner("Processing"):
+            with st.spinner("Kraunama..."):
                 for doc in docs:
-                    # Extract document content
-                    content = doc
-                    
-                    # Extract document metadata if available
-                    #metadata =doc.metadata.get('original_doc_name', 'unknown')
-                    # Display content and metadata
-                    st.write(f"Documentas: {content}")
-                    
-                   
-        
-
-        # Check if a response (generation) was produced by the graph
-        if 'generation' in state_dict and state_dict['generation']:
-            response = state_dict["generation"]
-
-            # Add the assistant's response to the chat history and display it
+                    # Display each document
+                    st.write(f"Dokumentas: {doc}")
+
+        # Check for and display the assistant's response
+        response = state_dict.get("generation")
+        if response:
             st.session_state.messages.append({"role": "assistant", "content": response})
             st.chat_message("assistant").write(response)
-        else:
-            st.chat_message("assistant").write("Your question violates toxicity rules or contains sensitive information.")
-
+        
+            
     except Exception as e:
         # Display an error message in case of failure
         st.chat_message("assistant").write("Klaida: Arba per didelis kontekstas suteiktas modeliui, arba užklausų serveryje yra per daug")
-