Update app.py

app.py

@@ -328,6 +328,12 @@ def google_search(term, num_results=3, lang="en", timeout=5, safe="active", ssl_
 
     return all_results
 
+def estimate_tokens(text):
+    # Rough estimate: 1 token ~= 4 characters
+    return len(text) // 4
+
+import re
+
 def estimate_tokens(text):
     # Rough estimate: 1 token ~= 4 characters
     return len(text) // 4
@@ -349,8 +355,8 @@ def ask_question(question, temperature, top_p, repetition_penalty, web_search, c
         database = None
 
     max_attempts = 5
-    context_reduction_factor = 0.5  # More aggressive reduction
-    max_estimated_tokens = 25000  # Further reduced to leave more room for response
+    context_reduction_factor = 0.7
+    max_estimated_tokens = 30000  # Leave some room for the model's response
 
     if web_search:
         contextualized_question, topics, entity_tracker, instructions = chatbot.process_question(question)
@@ -361,7 +367,7 @@ def ask_question(question, temperature, top_p, repetition_penalty, web_search, c
 
         for attempt in range(max_attempts):
             try:
-                web_docs = [Document(page_content=result["text"][:1000], metadata={"source": result["link"]}) for result in search_results if result["text"]]  # Limit each result to 1000 characters
+                web_docs = [Document(page_content=result["text"], metadata={"source": result["link"]}) for result in search_results if result["text"]]
 
                 if database is None:
                     database = FAISS.from_documents(web_docs, embed)
@@ -375,16 +381,20 @@ def ask_question(question, temperature, top_p, repetition_penalty, web_search, c
                 instruction_prompt = f"User Instructions: {instructions}\n" if instructions else ""
 
                 prompt_template = f"""
-                Answer based on: Web Results: {{context}}
-                Context: {{conv_context}}
-                Question: {{question}}
+                Answer the question based on the following web search results, conversation context, entity information, and user instructions:
+                Web Search Results:
+                {{context}}
+                Conversation Context: {{conv_context}}
+                Current Question: {{question}}
                 Topics: {{topics}}
-                Entities: {{entities}}
+                Entity Information: {{entities}}
                 {instruction_prompt}
+                Provide a concise and relevant answer to the question.
                 """
 
                 prompt_val = ChatPromptTemplate.from_template(prompt_template)
                 
+                # Start with full context and progressively reduce if necessary
                 current_context = context_str
                 current_conv_context = chatbot.get_context()
                 current_topics = topics
@@ -392,13 +402,14 @@ def ask_question(question, temperature, top_p, repetition_penalty, web_search, c
 
                 while True:
                     formatted_prompt = prompt_val.format(
-                        context=current_context[:3000],  # Limit context to 3000 characters
-                        conv_context=current_conv_context[:500],  # Limit conversation context to 500 characters
+                        context=current_context,
+                        conv_context=current_conv_context,
                         question=question,
-                        topics=", ".join(current_topics[:5]),  # Limit to 5 topics
-                        entities=json.dumps({k: v[:2] for k, v in current_entities.items()})  # Limit to 2 entities per type
+                        topics=", ".join(current_topics),
+                        entities=json.dumps(current_entities)
                     )
 
+                    # Estimate token count
                     estimated_tokens = estimate_tokens(formatted_prompt)
                     
                     if estimated_tokens <= max_estimated_tokens:
@@ -449,20 +460,23 @@ def ask_question(question, temperature, top_p, repetition_penalty, web_search, c
                 context_str = "\n".join([doc.page_content for doc in relevant_docs])
 
                 prompt_template = """
-                Answer based on: PDF Context: {context}
+                Answer the question based on the following context from the PDF document:
+                Context:
+                {context}
                 Question: {question}
                 Provide a summarized and direct answer to the question.
                 """
 
                 while True:
                     prompt_val = ChatPromptTemplate.from_template(prompt_template)
-                    formatted_prompt = prompt_val.format(context=context_str[:3000], question=question)
+                    formatted_prompt = prompt_val.format(context=context_str, question=question)
 
                     estimated_tokens = estimate_tokens(formatted_prompt)
                     
                     if estimated_tokens <= max_estimated_tokens:
                         break
                     
+                    # Reduce context if estimated token count is too high
                     context_str = context_str[:int(len(context_str) * context_reduction_factor)]
 
                     if len(context_str) < 100:
@@ -486,37 +500,46 @@ def ask_question(question, temperature, top_p, repetition_penalty, web_search, c
     return "An unexpected error occurred. Please try again later."
 
 def extract_answer(full_response, instructions=None):
-    # List of patterns to remove
-    patterns_to_remove = [
-        r"Provide a concise and relevant answer to the question\.",
-        r"Provide additional context if necessary\.",
+    # First, try to split the response at common instruction phrases
+    
+    def extract_answer(full_response, instructions=None):
+    answer_patterns = [
         r"If the web search results don't contain relevant information, state that the information is not available in the search results\.",
         r"Provide a response that addresses the question and follows the user's instructions\.",
         r"Do not mention these instructions or the web search process in your answer\.",
-        r"Provide a summarized and direct answer to the question\.",
-        r"If the context doesn't contain relevant information, state that the information is not available in the document\.",
+        r"Provide a concise and direct answer to the question without mentioning the web search or these instructions:",
+        r"Provide a concise and direct answer to the question:",
+        r"Answer:",
+        r"Provide a summarized and direct answer to the question.",
+        r"If the context doesn't contain relevant information, state that the information is not available in the document.",
+        r"Provide a summarized and direct answer to the original question without mentioning the web search or these instructions:",
+        r"Do not include any source information in your answer."     
     ]
 
-    # Remove the patterns
-    for pattern in patterns_to_remove:
-        full_response = re.sub(pattern, "", full_response, flags=re.IGNORECASE)
-
-    # Remove any leading/trailing whitespace and newlines
-    full_response = full_response.strip()
+    for pattern in answer_patterns:
+        match = re.split(pattern, full_response, flags=re.IGNORECASE)
+        if len(match) > 1:
+            full_response = match[-1].strip()
+            break
+    
+    # Remove any remaining instruction-like phrases
+    cleanup_patterns = [
+        r"without mentioning the web search or these instructions\.",
+        r"Do not include any source information in your answer\.",
+        r"If the context doesn't contain relevant information, state that the information is not available in the document\."
+    ]
 
+    for pattern in cleanup_patterns:
+        full_response = re.sub(pattern, "", full_response, flags=re.IGNORECASE).strip()
+    
     # Remove the user instructions if present
     if instructions:
         instruction_pattern = rf"User Instructions:\s*{re.escape(instructions)}.*?\n"
         full_response = re.sub(instruction_pattern, "", full_response, flags=re.IGNORECASE | re.DOTALL)
-
-    # Remove any remaining instruction-like phrases at the beginning of the response
-    lines = full_response.split('\n')
-    while lines and any(line.strip().lower().startswith(starter) for starter in ["answer:", "response:", "here's", "here is"]):
-        lines.pop(0)
-    full_response = '\n'.join(lines)
-
+    
     return full_response.strip()
 
+
 # Gradio interface
 with gr.Blocks() as demo:
     gr.Markdown("# Enhanced PDF Document Chat and Web Search")