Update app.py

app.py

@@ -48,6 +48,7 @@ class EnhancedContextDrivenChatbot:
         self.history = []
         self.history_size = history_size
         self.entity_tracker = {}
+        self.conversation_context = ""
 
     def add_to_history(self, text):
         self.history.append(text)
@@ -61,8 +62,12 @@ class EnhancedContextDrivenChatbot:
                 self.entity_tracker[ent.label_] = set()
             self.entity_tracker[ent.label_].add(ent.text)
 
+        # Update conversation context
+        self.conversation_context += f" {text}"
+        self.conversation_context = ' '.join(self.conversation_context.split()[-100:])  # Keep last 100 words
+
     def get_context(self):
-        return " ".join(self.history)
+        return self.conversation_context
 
     def is_follow_up_question(self, question):
         doc = nlp(question.lower())
@@ -87,12 +92,27 @@ class EnhancedContextDrivenChatbot:
         # Calculate similarity
         similarity = cosine_similarity([context_embedding], [question_embedding])[0][0]
         
-        # If similarity is low, it might be a new topic
-        if similarity < 0.3:  # This threshold can be adjusted
-            return question
+        # If similarity is high, it's likely a follow-up question
+        if similarity > 0.5:  # This threshold can be adjusted
+            return f"{combined_context} {question}"
+        
+        # Otherwise, it might be a new topic
+        return question
+
+    def rephrase_query(self, question):
+        prompt = f"""
+        Given the conversation context and the current question, rephrase the question to include relevant context:
         
-        # Otherwise, prepend the context
-        return f"{combined_context} {question}"
+        Conversation context: {self.get_context()}
+        Current question: {question}
+        
+        Rephrased question:
+        """
+        
+        # Use your language model to generate the rephrased question
+        rephrased_question = generate_chunked_response(model, prompt)
+        
+        return rephrased_question.strip()
 
     def process_question(self, question):
         contextualized_question = self.get_most_relevant_context(question)
@@ -100,10 +120,9 @@ class EnhancedContextDrivenChatbot:
         # Extract topics from the question
         topics = self.extract_topics(question)
         
-        # Check if it's a follow-up question
+        # Check if it's a follow-up question and rephrase if necessary
         if self.is_follow_up_question(question):
-            # If it's a follow-up, make sure to include previous context
-            contextualized_question = f"{self.get_context()} {question}"
+            contextualized_question = self.rephrase_query(contextualized_question)
         
         # Add the new question to history
         self.add_to_history(question)
@@ -336,7 +355,7 @@ def ask_question(question, temperature, top_p, repetition_penalty, web_search, c
                 formatted_prompt = prompt_val.format(
                     context=context_str, 
                     conv_context=chatbot.get_context(), 
-                    question=question,
+                    question=contextualized_question,
                     topics=", ".join(topics),
                     entities=json.dumps(serializable_entity_tracker)
                 )
@@ -356,6 +375,9 @@ def ask_question(question, temperature, top_p, repetition_penalty, web_search, c
         sources_section = "\n\nSources:\n" + "\n".join(f"- {source}" for source in sources)
         answer += sources_section
 
+        # Update chatbot context with the answer
+        chatbot.add_to_history(answer)
+
         return answer
 
     else:  # PDF document chat