Update app.py

app.py

@@ -97,48 +97,52 @@ class Agent1:
         return questions
 
     def update_context(self, query: str):
-        tokens = nltk.pos_tag(word_tokenize(query))
-        noun_phrases = []
-        current_phrase = []
-
-        for word, tag in tokens:
-            if tag.startswith('NN') or tag.startswith('JJ'):
-                current_phrase.append(word)
-            else:
-                if current_phrase:
-                    noun_phrases.append(' '.join(current_phrase))
-                    current_phrase = []
-
-        if current_phrase:
-            noun_phrases.append(' '.join(current_phrase))
-
-        if noun_phrases:
-            self.context['main_topic'] = noun_phrases[0]
-            self.context['related_topics'] = noun_phrases[1:]
+       tokens = nltk.pos_tag(word_tokenize(query))
+       noun_phrases = []
+       current_phrase = []
+
+       for word, tag in tokens:
+           if tag.startswith('NN') or tag.startswith('JJ'):
+               current_phrase.append(word)
+           else:
+               if current_phrase:
+                   noun_phrases.append(' '.join(current_phrase))
+                   current_phrase = []
+
+       if current_phrase:
+           noun_phrases.append(' '.join(current_phrase))
+
+       if noun_phrases:
+           self.context['main_topic'] = noun_phrases[0]
+           self.context['related_topics'] = noun_phrases[1:]
+           self.context['last_query'] = query
 
     def apply_context(self, query: str) -> str:
-        words = word_tokenize(query.lower())
-        
-        if (len(words) <= 5 or 
-            any(word in self.pronouns for word in words) or
-            (self.context.get('main_topic') and self.context['main_topic'].lower() not in query.lower())):
-            
-            new_query_parts = []
-            main_topic_added = False
+       words = word_tokenize(query.lower())
+       
+       if (len(words) <= 5 or 
+           any(word in self.pronouns for word in words) or
+           (self.context.get('main_topic') and self.context['main_topic'].lower() not in query.lower())):
+           
+           new_query_parts = []
+           main_topic_added = False
 
-            for word in words:
-                if word in self.pronouns and self.context.get('main_topic'):
-                    new_query_parts.append(self.context['main_topic'])
-                    main_topic_added = True
-                else:
-                    new_query_parts.append(word)
+           for word in words:
+               if word in self.pronouns and self.context.get('main_topic'):
+                   new_query_parts.append(self.context['main_topic'])
+                   main_topic_added = True
+               else:
+                   new_query_parts.append(word)
+
+           if not main_topic_added and self.context.get('main_topic'):
+               new_query_parts.append(f"in the context of {self.context['main_topic']}")
 
-            if not main_topic_added and self.context.get('main_topic'):
-                new_query_parts.append(f"of {self.context['main_topic']}")
+           query = ' '.join(new_query_parts)
 
-            query = ' '.join(new_query_parts)
+       if self.context.get('last_query'):
+           query = f"{self.context['last_query']} and now {query}"
 
-        return query
+       return query
 
     def process(self, user_input: str) -> tuple[List[str], Dict[str, List[Dict[str, str]]]]:
         self.update_context(user_input)
@@ -306,13 +310,15 @@ def google_search(term, num_results=5, lang="en", timeout=5, safe="active", ssl_
 
     return all_results
 
-def ask_question(question, temperature, top_p, repetition_penalty, web_search):
+def ask_question(question, temperature, top_p, repetition_penalty, web_search, agent1=None):
     if not question:
         return "Please enter a question."
 
+    if agent1 is None:
+        agent1 = Agent1()
+
     model = get_model(temperature, top_p, repetition_penalty)
     embed = get_embeddings()
-    agent1 = Agent1()  # Create Agent1 without passing a model
 
     if os.path.exists("faiss_database"):
         database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
@@ -322,8 +328,11 @@ def ask_question(question, temperature, top_p, repetition_penalty, web_search):
     max_attempts = 3
     context_reduction_factor = 0.7
 
+    agent1.update_context(question)
+    contextualized_question = agent1.apply_context(question)
+
     if web_search:
-        queries, search_results = agent1.process(question)
+        queries, search_results = agent1.process(contextualized_question)
         all_answers = []
 
         for query in queries:
@@ -395,7 +404,7 @@ def ask_question(question, temperature, top_p, repetition_penalty, web_search):
                     return "No documents available. Please upload documents or enable web search to answer questions."
 
                 retriever = database.as_retriever()
-                relevant_docs = retriever.get_relevant_documents(question)
+                relevant_docs = retriever.get_relevant_documents(contextualized_question)
                 context_str = "\n".join([doc.page_content for doc in relevant_docs])
 
                 if attempt > 0:
@@ -413,7 +422,7 @@ def ask_question(question, temperature, top_p, repetition_penalty, web_search):
                 """
 
                 prompt_val = ChatPromptTemplate.from_template(prompt_template)
-                formatted_prompt = prompt_val.format(context=context_str, question=question)
+                formatted_prompt = prompt_val.format(context=context_str, question=contextualized_question)
 
                 full_response = generate_chunked_response(model, formatted_prompt)
 
@@ -466,8 +475,10 @@ with gr.Blocks() as demo:
             repetition_penalty_slider = gr.Slider(label="Repetition Penalty", minimum=1.0, maximum=2.0, value=1.0, step=0.1)
             web_search_checkbox = gr.Checkbox(label="Enable Web Search", value=False)
 
+    agent1 = Agent1()
+
     def chat(question, history, temperature, top_p, repetition_penalty, web_search):
-        answer = ask_question(question, temperature, top_p, repetition_penalty, web_search)
+        answer = ask_question(question, temperature, top_p, repetition_penalty, web_search, agent1)
         history.append((question, answer))
         return "", history