Update app.py

app.py

@@ -116,16 +116,13 @@ def update_vectors(files, parser):
         label="Select documents to query"
     )
 
-def generate_chunked_response(prompt, model, max_tokens=1000, num_calls=1, temperature=0.2, should_stop=False, continuation=False):
+def generate_chunked_response(prompt, model, max_tokens=1000, num_calls=3, temperature=0.2, should_stop=False):
     print(f"Starting generate_chunked_response with {num_calls} calls")
     full_response = ""
     messages = [{"role": "user", "content": prompt}]
     
-    if continuation:
-        messages.insert(0, {"role": "system", "content": "This is a continuation of a previous response. Please continue from where you left off, maintaining coherence and avoiding repetition."})
-    
     if model == "@cf/meta/llama-3.1-8b-instruct":
-        # Cloudflare API logic
+        # Cloudflare API
         for i in range(num_calls):
             print(f"Starting Cloudflare API call {i+1}")
             if should_stop:
@@ -136,12 +133,15 @@ def generate_chunked_response(prompt, model, max_tokens=1000, num_calls=1, tempe
                     f"https://api.cloudflare.com/client/v4/accounts/{ACCOUNT_ID}/ai/run/@cf/meta/llama-3.1-8b-instruct",
                     headers={"Authorization": f"Bearer {API_TOKEN}"},
                     json={
-                        "stream": True,
-                        "messages": messages,
+                        "stream": true,
+                        "messages": [
+                            {"role": "system", "content": "You are a friendly assistant"},
+                            {"role": "user", "content": prompt}
+                        ],
                         "max_tokens": max_tokens,
                         "temperature": temperature
                     },
-                    stream=True
+                    stream=true
                 )
                 
                 for line in response.iter_lines():
@@ -153,16 +153,11 @@ def generate_chunked_response(prompt, model, max_tokens=1000, num_calls=1, tempe
                             json_data = json.loads(line.decode('utf-8').split('data: ')[1])
                             chunk = json_data['response']
                             full_response += chunk
-                            yield full_response
                         except json.JSONDecodeError:
                             continue
                 print(f"Cloudflare API call {i+1} completed")
             except Exception as e:
-                if "cancelled" in str(e).lower():
-                    print("Generation cancelled")
-                    return
-                else:
-                    print(f"Error in generating response from Cloudflare: {str(e)}")
+                print(f"Error in generating response from Cloudflare: {str(e)}")
     else:
         # Original Hugging Face API logic
         client = InferenceClient(model, token=huggingface_token)
@@ -185,14 +180,9 @@ def generate_chunked_response(prompt, model, max_tokens=1000, num_calls=1, tempe
                     if message.choices and message.choices[0].delta and message.choices[0].delta.content:
                         chunk = message.choices[0].delta.content
                         full_response += chunk
-                        yield full_response
                 print(f"Hugging Face API call {i+1} completed")
             except Exception as e:
-                if "cancelled" in str(e).lower():
-                    print("Generation cancelled")
-                    return
-                else:
-                    print(f"Error in generating response from Hugging Face: {str(e)}")
+                print(f"Error in generating response from Hugging Face: {str(e)}")
     
     # Clean up the response
     clean_response = re.sub(r'<s>\[INST\].*?\[/INST\]\s*', '', full_response, flags=re.DOTALL)
@@ -214,7 +204,7 @@ def generate_chunked_response(prompt, model, max_tokens=1000, num_calls=1, tempe
     final_response = '\n\n'.join(unique_paragraphs)
     
     print(f"Final clean response: {final_response[:100]}...")
-    yield final_response
+    return final_response
 
 def duckduckgo_search(query):
     with DDGS() as ddgs:
@@ -226,26 +216,6 @@ class CitingSources(BaseModel):
         ...,
         description="List of sources to cite. Should be an URL of the source."
     )
-def chatbot_interface(message, history, model, temperature, num_calls, use_web_search, selected_docs):
-    if not message.strip():
-        return history
-
-    for response in respond(message, history, model, temperature, num_calls, use_web_search, selected_docs):
-        yield response
-
-    # Make the Continue Generation button visible after the response is complete
-    demo.update(visible=True, elem_id="continue_btn")
-
-    try:
-        for response in respond(message, history, model, temperature, num_calls, use_web_search):
-            history[-1] = (message, response)
-            yield history
-    except gr.CancelledError:
-        yield history
-    except Exception as e:
-        logging.error(f"Unexpected error in chatbot_interface: {str(e)}")
-        history[-1] = (message, f"An unexpected error occurred: {str(e)}")
-        yield history
 
 def retry_last_response(history, use_web_search, model, temperature, num_calls):
     if not history:
@@ -257,30 +227,59 @@ 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):
-    if not message.strip():
-        return history
+    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'}")
 
-    history = history + [(message, "")]
+    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}"
-                history[-1] = (message, response)
-                yield history
+                first_line = response.split('\n')[0] if response else ''
+                logging.info(f"Generated Response (first line): {first_line}")
+                yield response
         else:
-            for partial_response in get_response_from_pdf(message, model, selected_docs, num_calls=num_calls, temperature=temperature):
-                history[-1] = (message, partial_response)
-                yield history
+            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
+            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
     except Exception as e:
-        if "cancelled" in str(e).lower():
-            yield history
+        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:
-            logging.error(f"Unexpected error in respond: {str(e)}")
-            history[-1] = (message, f"An unexpected error occurred: {str(e)}")
-            yield history
-
-    return history  # Ensure we always return the history at the end
+            yield f"An error occurred with the {model} model: {str(e)}. Please try again or select a different model."
 
 logging.basicConfig(level=logging.DEBUG)
 
@@ -338,37 +337,6 @@ After writing the document, please provide a list of sources used in your respon
     if not full_response:
         yield "I apologize, but I couldn't generate a response at this time. Please try again later."
 
-def get_response_with_search(query, model, num_calls=3, temperature=0.2):
-    search_results = duckduckgo_search(query)
-    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: '{query}'
-After writing the document, please provide a list of sources used in your response."""
-
-    if model == "@cf/meta/llama-3.1-8b-instruct":
-        # Use Cloudflare API
-        for response in get_response_from_cloudflare(prompt="", context=context, query=query, num_calls=num_calls, temperature=temperature, search_type="web"):
-            yield response, ""  # Yield streaming response without sources
-    else:
-        # 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=1000,
-                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, ""  # Yield partial main content without sources
-
 def get_response_from_pdf(query, model, selected_docs, num_calls=3, temperature=0.2):
     logging.info(f"Entering get_response_from_pdf with query: {query}, model: {model}, selected_docs: {selected_docs}")
     
@@ -386,7 +354,6 @@ def get_response_from_pdf(query, model, selected_docs, num_calls=3, temperature=
     relevant_docs = retriever.get_relevant_documents(query)
     logging.info(f"Number of relevant documents retrieved: {len(relevant_docs)}")
     
-    # Filter relevant_docs based on selected documents
     filtered_docs = [doc for doc in relevant_docs if doc.metadata["source"] in selected_docs]
     logging.info(f"Number of filtered documents: {len(filtered_docs)}")
     
@@ -395,28 +362,24 @@ def get_response_from_pdf(query, model, selected_docs, num_calls=3, temperature=
         yield "No relevant information found in the selected documents. Please try selecting different documents or rephrasing your query."
         return
 
-    for doc in filtered_docs:
-        logging.info(f"Document source: {doc.metadata['source']}")
-        logging.info(f"Document content preview: {doc.page_content[:100]}...")  # Log first 100 characters of each document
-
     context_str = "\n".join([doc.page_content for doc in filtered_docs])
     logging.info(f"Total context length: {len(context_str)}")
 
+    full_response = ""
+
     if model == "@cf/meta/llama-3.1-8b-instruct":
         logging.info("Using Cloudflare API")
-        # Use Cloudflare API with the retrieved context
         for response in get_response_from_cloudflare(prompt="", context=context_str, query=query, num_calls=num_calls, temperature=temperature, search_type="pdf"):
-            yield response
+            full_response += response
+            yield full_response
     else:
         logging.info("Using Hugging Face API")
-        # Use Hugging Face API
         prompt = f"""Using the following context from the PDF documents:
 {context_str}
 Write a detailed and complete response that answers the following user question: '{query}'"""
         
         client = InferenceClient(model, token=huggingface_token)
         
-        response = ""
         for i in range(num_calls):
             logging.info(f"API call {i+1}/{num_calls}")
             for message in client.chat_completion(
@@ -427,44 +390,121 @@ Write a detailed and complete response that answers the following user question:
             ):
                 if message.choices and message.choices[0].delta and message.choices[0].delta.content:
                     chunk = message.choices[0].delta.content
-                    response += chunk
-                    yield response  # Yield partial response
+                    full_response += chunk
+                    yield full_response
+    
+    logging.info("Finished generating initial response")
+
+def get_response_with_search(query, model, num_calls=3, temperature=0.2):
+    search_results = duckduckgo_search(query)
+    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: '{query}'
+After writing the document, please provide a list of sources used in your response."""
+
+    full_response = ""
+
+    if model == "@cf/meta/llama-3.1-8b-instruct":
+        # Use Cloudflare API
+        for response in get_response_from_cloudflare(prompt="", context=context, query=query, num_calls=num_calls, temperature=temperature, search_type="web"):
+            full_response += response
+            yield full_response, ""  # Yield streaming response without sources
+    else:
+        # Use Hugging Face API
+        client = InferenceClient(model, token=huggingface_token)
         
-        logging.info("Finished generating response")
+        for i in range(num_calls):
+            for message in client.chat_completion(
+                messages=[{"role": "user", "content": prompt}],
+                max_tokens=1000,
+                temperature=temperature,
+                stream=True,
+            ):
+                if message.choices and message.choices[0].delta and message.choices[0].delta.content:
+                    chunk = message.choices[0].delta.content
+                    full_response += chunk
+                    yield full_response, ""  # Yield partial main content without sources
+
+    logging.info("Finished generating initial response")
 
 def vote(data: gr.LikeData):
     if data.liked:
         print(f"You upvoted this response: {data.value}")
     else:
         print(f"You downvoted this response: {data.value}")
+def chatbot_interface(message, history, use_web_search, model, temperature, num_calls, selected_docs):
+    if not message.strip():
+        return "", history
+
+    history = history + [(message, "")]
 
-def continue_generation(history, model, temperature, num_calls, use_web_search, 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}"
+                history[-1] = (message, response)
+                yield history
+        else:
+            for response in get_response_from_pdf(message, model, selected_docs, num_calls=num_calls, temperature=temperature):
+                history[-1] = (message, response)
+                yield history
+    except gr.CancelledError:
+        yield history
+    except Exception as e:
+        logging.error(f"Unexpected error in chatbot_interface: {str(e)}")
+        history[-1] = (message, f"An unexpected error occurred: {str(e)}")
+        yield history
+
+def continue_generation(history, use_web_search, model, temperature, selected_docs):
     if not history:
         return history
     
     last_user_msg = history[-1][0]
-    last_ai_response = history[-1][1]
-    
-    continuation_prompt = f"""
-    Previous response: {last_ai_response}
-    
-    Original query: {last_user_msg}
-    
-    Please continue the response from where you left off, maintaining coherence and avoiding repetition.
-    """
+    previous_response = history[-1][1]
     
     try:
-        for response in respond(continuation_prompt, history[:-1], model, temperature, num_calls, use_web_search, selected_docs):
-            new_response = f"{last_ai_response}\n\n{response[-1][1]}"
-            history[-1] = (last_user_msg, new_response)
+        if use_web_search:
+            search_results = duckduckgo_search(last_user_msg)
+            context = "\n".join(f"{result['title']}\n{result['body']}\nSource: {result['href']}\n" 
+                                for result in search_results if 'body' in result)
+        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()
+                relevant_docs = retriever.get_relevant_documents(last_user_msg)
+                filtered_docs = [doc for doc in relevant_docs if doc.metadata["source"] in selected_docs]
+                context = "\n".join([doc.page_content for doc in filtered_docs])
+            else:
+                return history
+        
+        prompt = f"""Using the following context and partial response, please continue and complete the response:
+
+Context:
+{context}
+
+Query: {last_user_msg}
+
+Partial Response:
+{previous_response}
+
+Please continue the response from where it was cut off:"""
+
+        continued_response = previous_response
+        for chunk in get_response_from_cloudflare(prompt=prompt, context="", query="", num_calls=1, temperature=temperature, search_type="continuation"):
+            continued_response += chunk
+            history[-1] = (last_user_msg, continued_response)
             yield history
+    except gr.CancelledError:
+        yield history
     except Exception as e:
-        logging.error(f"Error in continue_generation: {str(e)}")
-        history[-1] = (last_user_msg, f"{last_ai_response}\n\nError continuing generation: {str(e)}")
+        logging.error(f"Unexpected error in continue_generation: {str(e)}")
+        history[-1] = (last_user_msg, f"{previous_response}\n\nAn error occurred while continuing generation: {str(e)}")
         yield history
 
-    return history
-
 css = """
 /* Add your custom CSS here */
 """
@@ -478,9 +518,7 @@ def display_documents():
         label="Select documents to query"
     )
 
-# Define the checkbox outside the demo block
 document_selector = gr.CheckboxGroup(label="Select documents to query")
-
 use_web_search = gr.Checkbox(label="Use Web Search", value=False)
 
 demo = gr.ChatInterface(
@@ -492,6 +530,10 @@ demo = gr.ChatInterface(
         use_web_search,
         document_selector
     ],
+    additional_buttons=[
+        gr.Button("Continue Generation"),
+        gr.Button("Upload Document")
+    ],
     title="AI-powered Web Search and PDF Chat Assistant",
     description="Chat with your PDFs or use web search to answer questions.",
     theme=gr.themes.Soft(
@@ -529,19 +571,22 @@ with demo:
     with gr.Row():
         file_input = gr.Files(label="Upload your PDF documents", file_types=[".pdf"])
         parser_dropdown = gr.Dropdown(choices=["pypdf", "llamaparse"], label="Select PDF Parser", value="llamaparse")
-        update_button = gr.Button("Upload Document")
     
     update_output = gr.Textbox(label="Update Status")
     
     # Update both the output text and the document selector
-    update_button.click(update_vectors, 
-                        inputs=[file_input, parser_dropdown], 
-                        outputs=[update_output, document_selector])
+    demo.additional_buttons[1].click(
+        update_vectors, 
+        inputs=[file_input, parser_dropdown], 
+        outputs=[update_output, document_selector]
+    )
 
-    continue_btn = gr.Button("Continue Generation", visible=False, elem_id="continue_btn")
-    continue_btn.click(continue_generation, 
-                       inputs=[demo.chatbot] + demo.additional_inputs,
-                       outputs=[demo.chatbot])
+    # Set up the continue generation button
+    demo.additional_buttons[0].click(
+        continue_generation,
+        inputs=[demo.chatbot, use_web_search, demo.additional_inputs[0], demo.additional_inputs[1], document_selector],
+        outputs=demo.chatbot
+    )
 
     gr.Markdown(
     """
@@ -552,8 +597,8 @@ with demo:
     4. Ask questions in the chat interface. 
     5. Toggle "Use Web Search" to switch between PDF chat and web search.
     6. Adjust Temperature and Number of API Calls to fine-tune the response generation.
-    7. Use the provided examples or ask your own questions.
-    8. If a response is incomplete, click "Continue Generation" for more information.
+    7. Use the "Continue Generation" button if you want to extend the last response.
+    8. Use the provided examples or ask your own questions.
     """
     )