Update app.py

app.py

@@ -16,7 +16,6 @@ from huggingface_hub import InferenceClient
 import inspect
 import logging
 
-
 # Set up basic configuration for logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
@@ -88,7 +87,6 @@ def update_vectors(files, parser):
             logging.info(f"Loaded {len(data)} chunks from {file.name}")
             all_data.extend(data)
             total_chunks += len(data)
-            # Append new documents instead of replacing
             if not any(doc["name"] == file.name for doc in uploaded_documents):
                 uploaded_documents.append({"name": file.name, "selected": True})
                 logging.info(f"Added new document to uploaded_documents: {file.name}")
@@ -116,96 +114,6 @@ def update_vectors(files, parser):
         label="Select documents to query"
     )
 
-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 model == "@cf/meta/llama-3.1-8b-instruct":
-        # Cloudflare API
-        for i in range(num_calls):
-            print(f"Starting Cloudflare API call {i+1}")
-            if should_stop:
-                print("Stop clicked, breaking loop")
-                break
-            try:
-                response = requests.post(
-                    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": [
-                            {"role": "system", "content": "You are a friendly assistant"},
-                            {"role": "user", "content": prompt}
-                        ],
-                        "max_tokens": max_tokens,
-                        "temperature": temperature
-                    },
-                    stream=true
-                )
-                
-                for line in response.iter_lines():
-                    if should_stop:
-                        print("Stop clicked during streaming, breaking")
-                        break
-                    if line:
-                        try:
-                            json_data = json.loads(line.decode('utf-8').split('data: ')[1])
-                            chunk = json_data['response']
-                            full_response += chunk
-                        except json.JSONDecodeError:
-                            continue
-                print(f"Cloudflare API call {i+1} completed")
-            except Exception as e:
-                print(f"Error in generating response from Cloudflare: {str(e)}")
-    else:
-        # Original Hugging Face API logic
-        client = InferenceClient(model, token=huggingface_token)
-        
-        for i in range(num_calls):
-            print(f"Starting Hugging Face API call {i+1}")
-            if should_stop:
-                print("Stop clicked, breaking loop")
-                break
-            try:
-                for message in client.chat_completion(
-                    messages=messages,
-                    max_tokens=max_tokens,
-                    temperature=temperature,
-                    stream=True,
-                ):
-                    if should_stop:
-                        print("Stop clicked during streaming, breaking")
-                        break
-                    if message.choices and message.choices[0].delta and message.choices[0].delta.content:
-                        chunk = message.choices[0].delta.content
-                        full_response += chunk
-                print(f"Hugging Face API call {i+1} completed")
-            except Exception as 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)
-    clean_response = clean_response.replace("Using the following context:", "").strip()
-    clean_response = clean_response.replace("Using the following context from the PDF documents:", "").strip()
-    
-    # Remove duplicate paragraphs and sentences
-    paragraphs = clean_response.split('\n\n')
-    unique_paragraphs = []
-    for paragraph in paragraphs:
-        if paragraph not in unique_paragraphs:
-            sentences = paragraph.split('. ')
-            unique_sentences = []
-            for sentence in sentences:
-                if sentence not in unique_sentences:
-                    unique_sentences.append(sentence)
-            unique_paragraphs.append('. '.join(unique_sentences))
-    
-    final_response = '\n\n'.join(unique_paragraphs)
-    
-    print(f"Final clean response: {final_response[:100]}...")
-    return final_response
-
 def duckduckgo_search(query):
     with DDGS() as ddgs:
         results = ddgs.text(query, max_results=5)
@@ -217,72 +125,6 @@ class CitingSources(BaseModel):
         description="List of sources to cite. Should be an URL of the source."
     )
 
-def retry_last_response(history, use_web_search, model, temperature, num_calls):
-    if not history:
-        return history
-    
-    last_user_msg = history[-1][0]
-    history = history[:-1]  # Remove the last response
-    
-    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):
-    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
-        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
-            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:
-        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."
-
-logging.basicConfig(level=logging.DEBUG)
-
 def get_response_from_cloudflare(prompt, context, query, num_calls=3, temperature=0.2, search_type="pdf"):
     headers = {
         "Authorization": f"Bearer {API_TOKEN}",
@@ -337,6 +179,37 @@ 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}")
     
@@ -354,6 +227,7 @@ 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)}")
     
@@ -362,24 +236,28 @@ 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"):
-            full_response += response
-            yield full_response
+            yield 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(
@@ -390,51 +268,36 @@ 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
-                    full_response += chunk
-                    yield full_response
-    
-    logging.info("Finished generating initial response")
+                    response += chunk
+                    yield response  # Yield partial response
+        
+        logging.info("Finished generating 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:
+def continue_response(last_response, context, query, model, temperature):
+    prompt = f"""Using the following context and partial response:
+
+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 = ""
+Partial Response:
+{last_response}
+
+Continue the response to fully answer the query: '{query}'
+Make sure the continuation flows smoothly from the previous part."""
 
     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
+        return get_response_from_cloudflare(prompt="", context=context, query=prompt, num_calls=1, temperature=temperature, search_type="pdf")
     else:
-        # Use Hugging Face API
         client = InferenceClient(model, token=huggingface_token)
-        
-        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")
+        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:
+                yield message.choices[0].delta.content
 
-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
@@ -442,15 +305,16 @@ def chatbot_interface(message, history, use_web_search, model, temperature, num_
     history = history + [(message, "")]
 
     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
+        last_response = ""
+        for response in respond(message, history, model, temperature, num_calls, use_web_search, selected_docs):
+            last_response = response
+            history[-1] = (message, response)
+            yield history
+
+        # Check if the response seems truncated
+        if not last_response.strip().endswith((".", "!", "?")):
+            history.append((None, "Response may be incomplete. Type 'continue' to generate more."))
+            yield history
     except gr.CancelledError:
         yield history
     except Exception as e:
@@ -460,50 +324,88 @@ def chatbot_interface(message, history, use_web_search, model, temperature, num_
 
 def continue_generation(history, use_web_search, model, temperature, selected_docs):
     if not history:
-        return history
-    
-    last_user_msg = history[-1][0]
-    previous_response = history[-1][1]
-    
-    try:
-        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:
+        return history, gr.Button.update(visible=False)
 
-Context:
-{context}
+    last_message = history[-1][0]
+    last_response = history[-1][1]
+
+    if use_web_search:
+        search_results = duckduckgo_search(last_message)
+        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()
+        database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
+        retriever = database.as_retriever()
+        relevant_docs = retriever.get_relevant_documents(last_message)
+        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])
+
+    continuation = ""
+    for chunk in continue_response(last_response, context, last_message, model, temperature):
+        continuation += chunk
+        history[-1] = (last_message, last_response + continuation)
+        yield history, gr.Button.update(visible=True)
+
+    if not (last_response + continuation).strip().endswith((".", "!", "?")):
+        yield history, gr.Button.update(visible=True, text="Continue Generation")
+    else:
+        yield history, gr.Button.update(visible=False)
 
-Query: {last_user_msg}
+def respond(message, history, model, temperature, num_calls, use_web_search, 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}")
 
-Partial Response:
-{previous_response}
+    # Check if the user wants to continue the previous response
+    if message.strip().lower() == "continue" and history:
+        last_message = history[-2][0]  # Get the last user message
+        last_response = history[-2][1]  # Get the last bot response
+        context = get_context(last_message, use_web_search, selected_docs)
+        for continuation in continue_response(last_response, context, last_message, model, temperature):
+            yield last_response + continuation
+    else:
+        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
+            else:
+                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:
+            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."
 
-Please continue the response from where it was cut off:"""
+def get_context(message, use_web_search, selected_docs):
+    if use_web_search:
+        search_results = duckduckgo_search(message)
+        return "\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()
+        database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
+        retriever = database.as_retriever()
+        relevant_docs = retriever.get_relevant_documents(message)
+        filtered_docs = [doc for doc in relevant_docs if doc.metadata["source"] in selected_docs]
+        return "\n".join([doc.page_content for doc in filtered_docs])
 
-        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"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
+
+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}")
 
 css = """
 /* Add your custom CSS here */
@@ -518,7 +420,9 @@ 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(
@@ -528,14 +432,10 @@ demo = gr.ChatInterface(
         gr.Slider(minimum=0.1, maximum=1.0, value=0.2, step=0.1, label="Temperature"),
         gr.Slider(minimum=1, maximum=5, value=1, step=1, label="Number of API Calls"),
         use_web_search,
-        document_selector
-    ],
-    additional_buttons=[
-        gr.Button("Continue Generation"),
-        gr.Button("Upload Document")
+        document_selector  # Add the document selector to the chat interface
     ],
     title="AI-powered Web Search and PDF Chat Assistant",
-    description="Chat with your PDFs or use web search to answer questions.",
+    description="Chat with your PDFs or use web search to answer questions. Type 'continue' to generate more if a response seems incomplete.",
     theme=gr.themes.Soft(
         primary_hue="orange",
         secondary_hue="amber",
@@ -567,26 +467,17 @@ demo = gr.ChatInterface(
 # Add file upload functionality
 with demo:
     gr.Markdown("## Upload PDF Documents")
-
     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
-    demo.additional_buttons[1].click(
-        update_vectors, 
-        inputs=[file_input, parser_dropdown], 
-        outputs=[update_output, document_selector]
-    )
-
-    # 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
-    )
+    update_button.click(update_vectors, 
+                        inputs=[file_input, parser_dropdown], 
+                        outputs=[update_output, document_selector])
 
     gr.Markdown(
     """
@@ -597,8 +488,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 "Continue Generation" button if you want to extend the last response.
-    8. Use the provided examples or ask your own questions.
+    7. Use the provided examples or ask your own questions.
+    8. If a response seems incomplete, type 'continue' to generate more.
     """
     )