Update app.py

app.py

@@ -1,5 +1,6 @@
 import os
 import logging
+import asyncio
 import gradio as gr
 from huggingface_hub import InferenceClient
 from langchain.embeddings import HuggingFaceEmbeddings
@@ -13,9 +14,9 @@ huggingface_token = os.environ.get("HUGGINGFACE_TOKEN")
 MODELS = [
     "mistralai/Mistral-7B-Instruct-v0.3",
     "mistralai/Mixtral-8x7B-Instruct-v0.1",
+    "mistralai/Mistral-Nemo-Instruct-2407",
     "meta-llama/Meta-Llama-3.1-8B-Instruct",
-    "meta-llama/Meta-Llama-3.1-70B-Instruct",
-    "mistralai/Mistral-Nemo-Instruct-2407"
+    "meta-llama/Meta-Llama-3.1-70B-Instruct"
 ]
 
 def get_embeddings():
@@ -35,7 +36,7 @@ def create_web_search_vectors(search_results):
             documents.append(Document(page_content=content, metadata={"source": result['href']}))
     return FAISS.from_documents(documents, embed)
 
-def get_response_with_search(query, model, use_embeddings, num_calls=3, temperature=0.2):
+async def get_response_with_search(query, model, use_embeddings, num_calls=3, temperature=0.2):
     search_results = duckduckgo_search(query)
     
     if not search_results:
@@ -58,19 +59,23 @@ After writing the document, please provide a list of sources used in your respon
     # Use Hugging Face API
     client = InferenceClient(model, token=huggingface_token)
     main_content = ""
-    for i in range(num_calls):
-        for message in client.chat_completion(
-            messages=[{"role": "user", "content": prompt}],
-            max_tokens=10000,
-            temperature=temperature,
-            stream=True,
-        ):
-            if message.choices and message.choices[0].delta and message.choices[0].delta.content:
-                chunk = message.choices[0].delta.content
-                main_content += chunk
-                yield main_content, ""
-
-def respond(message, history, model, temperature, num_calls, use_embeddings):
+    try:
+        for i in range(num_calls):
+            async for message in client.chat_completion(
+                messages=[{"role": "user", "content": prompt}],
+                max_tokens=10000,
+                temperature=temperature,
+                stream=True,
+            ):
+                if message.choices and message.choices[0].delta and message.choices[0].delta.content:
+                    chunk = message.choices[0].delta.content
+                    main_content += chunk
+                    yield main_content, ""
+    except Exception as e:
+        logging.error(f"Error in get_response_with_search: {str(e)}")
+        yield f"An error occurred while processing your request: {str(e)}", ""
+
+async def respond(message, history, model, temperature, num_calls, use_embeddings):
     logging.info(f"User Query: {message}")
     logging.info(f"Model Used: {model}")
     logging.info(f"Temperature: {temperature}")
@@ -78,9 +83,11 @@ def respond(message, history, model, temperature, num_calls, use_embeddings):
     logging.info(f"Use Embeddings: {use_embeddings}")
 
     try:
-        for main_content, sources in get_response_with_search(message, model, use_embeddings, num_calls=num_calls, temperature=temperature):
+        async for main_content, sources in get_response_with_search(message, model, use_embeddings, num_calls=num_calls, temperature=temperature):
             response = f"{main_content}\n\n{sources}"
             yield response
+    except asyncio.CancelledError:
+        yield "The operation was cancelled. Please try again."
     except Exception as e:
         logging.error(f"Error in respond function: {str(e)}")
         yield f"An error occurred: {str(e)}"