Update app.py

app.py

@@ -184,7 +184,7 @@ def delete_documents(selected_docs):
     
     return f"Deleted documents: {', '.join(deleted_docs)}", display_documents()
 
-async def generate_chunked_response(prompt, model, max_tokens=10000, num_calls=3, temperature=0.2):
+def generate_chunked_response(prompt, model, max_tokens=10000, num_calls=3, temperature=0.2):
     print(f"Starting generate_chunked_response with model: {model}, num_calls: {num_calls}")
     full_response = ""
     messages = [{"role": "user", "content": prompt}]
@@ -199,27 +199,26 @@ async def generate_chunked_response(prompt, model, max_tokens=10000, num_calls=3
         for i in range(num_calls):
             print(f"Starting Cloudflare API call {i+1}")
             try:
-                async with httpx.AsyncClient() as client:
-                    async with client.stream(
-                        "POST",
-                        f"https://api.cloudflare.com/client/v4/accounts/{ACCOUNT_ID}/ai/run/@cf/meta/llama-3.1-8b-instruct",
-                        json={
-                            "messages": messages,
-                            "stream": True,
-                            "max_tokens": max_tokens,
-                            "temperature": temperature
-                        },
-                        headers=headers
-                    ) as response:
-                        async for line in response.aiter_lines():
-                            if line.startswith("data: "):
-                                try:
-                                    json_data = json.loads(line[6:])
-                                    chunk = json_data.get('response', '')
-                                    full_response += chunk
-                                    yield full_response
-                                except json.JSONDecodeError:
-                                    continue
+                response = requests.post(
+                    f"https://api.cloudflare.com/client/v4/accounts/{ACCOUNT_ID}/ai/run/@cf/meta/llama-3.1-8b-instruct",
+                    json={
+                        "messages": messages,
+                        "stream": True,
+                        "max_tokens": max_tokens,
+                        "temperature": temperature
+                    },
+                    headers=headers,
+                    stream=True
+                )
+                for line in response.iter_lines():
+                    if line:
+                        try:
+                            json_data = json.loads(line.decode('utf-8').split('data: ')[1])
+                            chunk = json_data.get('response', '')
+                            full_response += chunk
+                            yield full_response
+                        except (json.JSONDecodeError, IndexError):
+                            continue
                 print(f"Cloudflare API call {i+1} completed")
             except Exception as e:
                 print(f"Error in generating response from Cloudflare: {str(e)}")
@@ -232,10 +231,11 @@ async def generate_chunked_response(prompt, model, max_tokens=10000, num_calls=3
         for i in range(num_calls):
             print(f"Starting Hugging Face API call {i+1}")
             try:
-                async for message in client.chat_completion_stream(
+                for message in client.chat_completion(
                     messages=messages,
                     max_tokens=max_tokens,
                     temperature=temperature,
+                    stream=True,
                 ):
                     if message.choices and message.choices[0].delta and message.choices[0].delta.content:
                         chunk = message.choices[0].delta.content
@@ -308,60 +308,43 @@ def retry_last_response(history, use_web_search, model, temperature, num_calls):
     
     return chatbot_interface(last_user_msg, history, use_web_search, model, temperature, num_calls)
 
-def respond(message, history, model, temperature, num_calls, use_web_search, selected_docs):
+def respond(message, history, use_web_search, model, temperature, num_calls, selected_docs):
     logging.info(f"User Query: {message}")
     logging.info(f"Model Used: {model}")
     logging.info(f"Search Type: {'Web Search' if use_web_search else 'PDF Search'}")
-
     logging.info(f"Selected Documents: {selected_docs}")
 
     try:
         if use_web_search:
-            for main_content, sources in get_response_with_search(message, model, num_calls=num_calls, temperature=temperature):
-                response = f"{main_content}\n\n{sources}"
-                first_line = response.split('\n')[0] if response else ''
-#                logging.info(f"Generated Response (first line): {first_line}")
-                yield response
+            search_results = duckduckgo_search(message)
+            context = "\n".join(f"{result['title']}\n{result['body']}\nSource: {result['href']}\n" 
+                                for result in search_results if 'body' in result)
+            prompt = f"""Using the following context:
+            {context}
+            Write a detailed and complete research document that fulfills the following user request: '{message}'
+            After writing the document, please provide a list of sources used in your response."""
         else:
             embed = get_embeddings()
             if os.path.exists("faiss_database"):
                 database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
                 retriever = database.as_retriever()
-                
-                # Filter relevant documents based on user selection
-                all_relevant_docs = retriever.get_relevant_documents(message)
-                relevant_docs = [doc for doc in all_relevant_docs if doc.metadata["source"] in selected_docs]
-                
-                if not relevant_docs:
-                    yield "No relevant information found in the selected documents. Please try selecting different documents or rephrasing your query."
-                    return
-
-                context_str = "\n".join([doc.page_content for doc in relevant_docs])
-            else:
-                context_str = "No documents available."
-                yield "No documents available. Please upload PDF documents to answer questions."
-                return
-            
-            if model == "@cf/meta/llama-3.1-8b-instruct":
-                # Use Cloudflare API
-                for partial_response in get_response_from_cloudflare(prompt="", context=context_str, query=message, num_calls=num_calls, temperature=temperature, search_type="pdf"):
-                    first_line = partial_response.split('\n')[0] if partial_response else ''
-#                   logging.info(f"Generated Response (first line): {first_line}")
-                    yield partial_response
+                relevant_docs = retriever.get_relevant_documents(message)
+                filtered_docs = [doc for doc in relevant_docs if doc.metadata["source"] in selected_docs]
+                if not filtered_docs:
+                    return "No relevant information found in the selected documents. Please try selecting different documents or rephrasing your query."
+                context_str = "\n".join([doc.page_content for doc in filtered_docs])
+                prompt = f"""Using the following context from the PDF documents:
+                {context_str}
+                Write a detailed and complete response that answers the following user question: '{message}'"""
             else:
-                # Use Hugging Face API
-                for partial_response in get_response_from_pdf(message, model, selected_docs, num_calls=num_calls, temperature=temperature):
-                    first_line = partial_response.split('\n')[0] if partial_response else ''
-#                    logging.info(f"Generated Response (first line): {first_line}")
-                    yield partial_response
+                return "No documents available. Please upload PDF documents to answer questions."
+
+        for response in generate_chunked_response(prompt, model, num_calls=num_calls, temperature=temperature):
+            yield response
+
     except Exception as e:
         logging.error(f"Error with {model}: {str(e)}")
-        if "microsoft/Phi-3-mini-4k-instruct" in model:
-            logging.info("Falling back to Mistral model due to Phi-3 error")
-            fallback_model = "mistralai/Mistral-7B-Instruct-v0.3"
-            yield from respond(message, history, fallback_model, temperature, num_calls, use_web_search, selected_docs)
-        else:
-            yield f"An error occurred with the {model} model: {str(e)}. Please try again or select a different model."
+        yield f"An error occurred with the {model} model: {str(e)}. Please try again or select a different model."
 
 logging.basicConfig(level=logging.DEBUG)
 
@@ -575,8 +558,7 @@ use_web_search = gr.Checkbox(label="Use Web Search", value=True)
 custom_placeholder = "Ask a question (Note: You can toggle between Web Search and PDF Chat in Additional Inputs below)"
 
 demo = gr.ChatInterface(
-    fn=respond,
-    async_mode=True,
+    respond,
     additional_inputs=[
         gr.Dropdown(choices=MODELS, label="Select Model", value=MODELS[3]),
         gr.Slider(minimum=0.1, maximum=1.0, value=0.2, step=0.1, label="Temperature"),