Update app.py

app.py

@@ -88,7 +88,7 @@ def update_vectors(files, parser):
     
     return f"Vector store updated successfully. Processed {total_chunks} chunks from {len(files)} files using {parser}."
 
-def generate_chunked_response(prompt, model, max_tokens=8000, num_calls=3, temperature=0.2, should_stop=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}]
@@ -188,14 +188,14 @@ class CitingSources(BaseModel):
         ...,
         description="List of sources to cite. Should be an URL of the source."
     )
-
-def chatbot_interface(message, history, use_web_search, model, temperature, num_calls, is_interrupted=False, partial_response=""):
+def chatbot_interface(message, history, use_web_search, model, temperature, num_calls):
     if not message.strip():
         return "", history
 
     history = history + [(message, "")]
+
     try:
-        for response in respond(message, history, model, temperature, num_calls, use_web_search, is_interrupted, partial_response):
+        for response in respond(message, history, model, temperature, num_calls, use_web_search):
             history[-1] = (message, response)
             yield history
     except gr.CancelledError:
@@ -214,7 +214,7 @@ 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, is_interrupted=False, partial_response=""):
+def respond(message, history, model, temperature, num_calls, use_web_search):
     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'}")
@@ -223,11 +223,9 @@ def respond(message, history, model, temperature, num_calls, use_web_search, is_
         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}"
-                if is_interrupted:
-                    partial_response += response
-                    yield partial_response
-                else:
-                    yield response
+                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"):
@@ -237,21 +235,19 @@ def respond(message, history, model, temperature, num_calls, use_web_search, is_
                 context_str = "\n".join([doc.page_content for doc in relevant_docs])
             else:
                 context_str = "No documents available."
-
+            
             if model == "@cf/meta/llama-3.1-8b-instruct":
-                for response in get_response_from_cloudflare(prompt="", context=context_str, query=message, num_calls=num_calls, temperature=temperature, search_type="pdf"):
-                    if is_interrupted:
-                        partial_response += response
-                        yield partial_response
-                    else:
-                        yield response
+                # 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:
-                for response in get_response_from_pdf(message, model, num_calls=num_calls, temperature=temperature):
-                    if is_interrupted:
-                        partial_response += response
-                        yield partial_response
-                    else:
-                        yield response
+                # Use Hugging Face API
+                for partial_response in get_response_from_pdf(message, model, 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:
@@ -260,8 +256,6 @@ def respond(message, history, model, temperature, num_calls, use_web_search, is_
             yield from respond(message, history, fallback_model, temperature, num_calls, use_web_search)
         else:
             yield f"An error occurred with the {model} model: {str(e)}. Please try again or select a different model."
-    finally:
-        pass
 
 logging.basicConfig(level=logging.DEBUG)
 
@@ -401,14 +395,12 @@ css = """
 use_web_search = gr.Checkbox(label="Use Web Search", value=False)
 
 demo = gr.ChatInterface(
-    chatbot_interface,
+    respond,
     additional_inputs=[
         gr.Dropdown(choices=MODELS, label="Select Model", value=MODELS[0]),
         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,
-        gr.Checkbox(value=False, label="Is Interrupted"),  # Flag to indicate interruption
-        gr.Textbox(label="Partial Response"),  # Store partial response
+        use_web_search  # Add this line to include the checkbox
     ],
     title="AI-powered Web Search and PDF Chat Assistant",
     description="Chat with your PDFs or use web search to answer questions.",
@@ -430,41 +422,12 @@ demo = gr.ChatInterface(
         color_accent_soft_dark="transparent",
         code_background_fill_dark="#140b0b"
     ),
+    
     css=css,
     examples=[
         ["Tell me about the contents of the uploaded PDFs."],
         ["What are the main topics discussed in the documents?"],
-        ["Can you summarize the key points from the PDFs?"],
-    ],
-    cache_examples=False,
-    analytics_enabled=False,
-)
-
-continue_button = gr.Button("Continue Generation")
-continue_output = gr.Textbox(label="Response")
-
-message_input = gr.Textbox(label="Message")
-
-history = gr.Textbox(label="History", lines=10, visible=False)
-
-continue_button.click(
-    chatbot_interface,
-    inputs=[
-        message_input,
-        history,
-        use_web_search,
-        model_dropdown,
-        temperature_slider,
-        num_calls_slider,
-        gr.Boolean(value=True, label="Is Interrupted"),  # Set is_interrupted to True
-        gr.Textbox(label="Partial Response"),
-    ],
-    outputs=continue_output,  # Add outputs parameter
-    css=css,
-    examples=[
-        ["Tell me about the contents of the uploaded PDFs."],
-        ["What are the main topics discussed in the documents?"],
-        ["Can you summarize the key points from the PDFs?"],
+        ["Can you summarize the key points from the PDFs?"]
     ],
     cache_examples=False,
     analytics_enabled=False,