Update app.py

app.py

@@ -1,6 +1,7 @@
 import os
 import logging
 import asyncio
+from typing import AsyncGenerator, Tuple
 import gradio as gr
 from huggingface_hub import InferenceClient
 from langchain.embeddings import HuggingFaceEmbeddings
@@ -55,7 +56,16 @@ def create_web_search_vectors(search_results):
     logging.info(f"Created vectors for {len(documents)} search results.")
     return FAISS.from_documents(documents, embed)
 
-async def get_response_with_search(query, system_prompt, model, use_embeddings, history=None, num_calls=3, temperature=0.2):
+def create_context(search_results, use_embeddings, query):
+    if use_embeddings:
+        web_search_database = create_web_search_vectors(search_results)
+        retriever = web_search_database.as_retriever(search_kwargs={"k": 5})
+        relevant_docs = retriever.get_relevant_documents(query)
+        return "\n".join([doc.page_content for doc in relevant_docs])
+    else:
+        return "\n".join([f"{result['title']}\n{result['body']}" for result in search_results])
+
+async def get_response_with_search(query: str, system_prompt: str, model: str, use_embeddings: bool, history=None, num_calls: int = 3, temperature: float = 0.2) -> AsyncGenerator[Tuple[str, str], None]:
     search_results = duckduckgo_search(query)
     
     if not search_results:
@@ -66,14 +76,7 @@ async def get_response_with_search(query, system_prompt, model, use_embeddings,
     sources = [result['href'] for result in search_results if 'href' in result]
     source_list_str = "\n".join(sources)
 
-    if use_embeddings:
-        web_search_database = create_web_search_vectors(search_results)
-        retriever = web_search_database.as_retriever(search_kwargs={"k": 5})
-        relevant_docs = retriever.get_relevant_documents(query)
-        context = "\n".join([doc.page_content for doc in relevant_docs])
-    else:
-        context = "\n".join([f"{result['title']}\n{result['body']}" for result in search_results])
-
+    context = create_context(search_results, use_embeddings, query)
     logging.info(f"Context created for query: {query}")
 
     user_message = f"""Using the following context from web search results:
@@ -81,9 +84,6 @@ async def get_response_with_search(query, system_prompt, model, use_embeddings,
 
 Write a detailed and complete research document that fulfills the following user request: '{query}'."""
 
-    client = InferenceClient(model, token=huggingface_token)
-    full_response = ""
-    
     messages = [
         {"role": "system", "content": system_prompt},
         {"role": "user", "content": user_message}
@@ -92,50 +92,38 @@ Write a detailed and complete research document that fulfills the following user
     if history:
         messages = history + messages
 
-    try:
-        for call in range(num_calls):
-            try:
-                response_stream = client.chat_completion(
-                    messages=messages,
-                    max_tokens=6000,
-                    temperature=temperature,
-                    stream=True,
-                    top_p=0.8,
-                )
-                
-                if response_stream is None:
-                    logging.error(f"API call {call + 1} returned None")
-                    yield "The API returned an empty response. Please try again.", ""
-                    continue
-
-                for response in response_stream:
-                    if isinstance(response, dict) and "choices" in response:
-                        for choice in response["choices"]:
-                            if "delta" in choice and "content" in choice["delta"]:
-                                chunk = choice["delta"]["content"]
-                                full_response += chunk
-                                yield full_response, ""
-                    else:
-                        logging.error(f"Unexpected response format in API call {call + 1}: {response}")
-                
-                if full_response:
-                    break  # If we got a valid response, exit the loop
-                
-            except Exception as e:
-                logging.error(f"Error in API call {call + 1}: {str(e)}")
-                if "422 Client Error" in str(e):
-                    logging.warning("Received 422 Client Error. Adjusting request parameters.")
-                    # You might want to adjust parameters here, e.g., reduce max_tokens
-                yield f"An error occurred during API call {call + 1}. Retrying...", ""
-            
-            await asyncio.sleep(1)  # 1 second delay between calls
+    client = InferenceClient(model, token=huggingface_token)
+    full_response = ""
 
-    except asyncio.CancelledError:
-        logging.warning("The operation was cancelled.")
-        yield "The operation was cancelled. Please try again.", ""
-    except Exception as e:
-        logging.error(f"Unexpected error in get_response_with_search: {str(e)}")
-        yield f"An unexpected error occurred: {str(e)}", ""
+    for call in range(num_calls):
+        try:
+            response = await asyncio.to_thread(
+                client.chat_completion,
+                messages=messages,
+                max_tokens=6000,
+                temperature=temperature,
+                top_p=0.8,
+            )
+
+            if response is None or not isinstance(response, dict) or 'choices' not in response:
+                logging.error(f"API call {call + 1} returned an invalid response: {response}")
+                if call == num_calls - 1:
+                    yield "The API returned an invalid response. Please try again later.", ""
+                continue
+
+            new_content = response['choices'][0]['message']['content']
+            full_response += new_content
+            yield full_response, ""
+
+            if full_response:
+                break  # If we got a valid response, exit the loop
+
+        except Exception as e:
+            logging.error(f"Error in API call {call + 1}: {str(e)}")
+            if call == num_calls - 1:
+                yield f"An error occurred during API calls: {str(e)}. Please try again later.", ""
+        
+        await asyncio.sleep(1)  # 1 second delay between calls
 
     if not full_response:
         logging.warning("No response generated from the model")
@@ -143,18 +131,12 @@ Write a detailed and complete research document that fulfills the following user
     else:
         yield f"{full_response}\n\nSources:\n{source_list_str}", ""
 
-async def respond(message, system_prompt, history, model, temperature, num_calls, use_embeddings):
-    logging.info(f"User Query: {message}")
-    logging.info(f"Model Used: {model}")
-    logging.info(f"Temperature: {temperature}")
-    logging.info(f"Number of API Calls: {num_calls}")
-    logging.info(f"Use Embeddings: {use_embeddings}")
-    logging.info(f"System Prompt: {system_prompt}")
-    logging.info(f"History: {history}")  # Log the history for debugging
-
-    # Convert gradio history to the format expected by get_response_with_search
+def process_history(history):
     chat_history = []
-    if history:
+    if isinstance(history, str):
+        # If history is a string (like the system prompt), add it as a system message
+        chat_history.append({"role": "system", "content": history})
+    elif isinstance(history, list):
         for entry in history:
             if isinstance(entry, (list, tuple)) and len(entry) == 2:
                 human, assistant = entry
@@ -164,10 +146,20 @@ async def respond(message, system_prompt, history, model, temperature, num_calls
             elif isinstance(entry, str):
                 # If it's a string, assume it's a user message
                 chat_history.append({"role": "user", "content": entry})
-            # Ignore any other formats
+    return chat_history
+
+async def respond(message, system_prompt, history, model, temperature, num_calls, use_embeddings):
+    logging.info(f"User Query: {message}")
+    logging.info(f"Model Used: {model}")
+    logging.info(f"Temperature: {temperature}")
+    logging.info(f"Number of API Calls: {num_calls}")
+    logging.info(f"Use Embeddings: {use_embeddings}")
+    logging.info(f"System Prompt: {system_prompt}")
+    logging.info(f"History: {history}")
+
+    chat_history = process_history(history)
 
     try:
-        full_response = ""
         async for main_content, sources in get_response_with_search(
             message, 
             system_prompt, 
@@ -177,16 +169,8 @@ async def respond(message, system_prompt, history, model, temperature, num_calls
             num_calls=num_calls, 
             temperature=temperature
         ):
-            if "error" in main_content.lower() or "no response" in main_content.lower():
-                # If it's an error message, yield it as is
-                yield main_content
-            else:
-                # Otherwise, yield only the new content
-                new_content = main_content[len(full_response):]
-                full_response = main_content
-                yield new_content
-
-        # Yield the sources as a separate message
+            yield main_content
+
         if sources:
             yield f"\n\nSources:\n{sources}"
 
@@ -213,16 +197,8 @@ css = """
 def create_gradio_interface():
     custom_placeholder = "Enter your question here for web search."
 
-    async def wrapped_respond(*args):
-        try:
-            async for response in respond(*args):
-                yield response
-        except Exception as e:
-            logging.error(f"Error in wrapped_respond: {str(e)}")
-            yield f"An error occurred: {str(e)}"
-
     demo = gr.ChatInterface(
-        fn=wrapped_respond,  # Use the wrapped version
+        fn=respond,
         additional_inputs_accordion=gr.Accordion(label="⚙️ Parameters", open=True, render=False),
         additional_inputs=[
             gr.Textbox(value=DEFAULT_SYSTEM_PROMPT, lines=6, label="System Prompt", placeholder="Enter your system prompt here"),
@@ -268,4 +244,4 @@ def create_gradio_interface():
 
 if __name__ == "__main__":
     demo = create_gradio_interface()
-    demo.launch(share=True)
+    demo.launch(share=True)