Update app.py

app.py

@@ -35,6 +35,7 @@ class Agent1:
         self.question_words = set(["what", "when", "where", "who", "whom", "which", "whose", "why", "how"])
         self.conjunctions = set(["and", "or"])
         self.pronouns = set(["it", "its", "they", "their", "them", "he", "his", "him", "she", "her", "hers"])
+        self.context = {}  # Store multiple context elements
 
     def is_question(self, text: str) -> bool:
         words = word_tokenize(text.lower())
@@ -99,12 +100,42 @@ class Agent1:
 
         return questions
 
+    def update_context(self, query: str):
+        tokens = nltk.pos_tag(word_tokenize(query))
+        for word, tag in tokens:
+            if tag.startswith('NN'):  # Noun
+                self.context['subject'] = word
+            elif tag.startswith('JJ'):  # Adjective
+                self.context['attribute'] = word
+            elif tag == 'NNP':  # Proper noun (e.g., country names)
+                self.context['location'] = word
+
+    def apply_context(self, query: str) -> str:
+        words = word_tokenize(query.lower())
+        
+        # Check if the query is short or contains pronouns
+        if len(words) <= 5 or any(word in self.pronouns for word in words):
+            # Apply relevant context
+            if 'location' in self.context and 'location' not in query:
+                query += f" in {self.context['location']}"
+            if 'subject' in self.context and self.context['subject'] not in query:
+                query = f"{self.context['subject']} {query}"
+        
+        return query
+
     def process(self, user_input: str) -> tuple[List[str], Dict[str, List[Dict[str, str]]]]:
-        queries = self.rephrase_and_split(user_input)
+        # First, update context with the new input
+        self.update_context(user_input)
+        
+        # Then apply context and split queries
+        contextualized_input = self.apply_context(user_input)
+        queries = self.rephrase_and_split(contextualized_input)
         print("Identified queries:", queries)
+        
         results = {}
         for query in queries:
             results[query] = google_search(query)
+        
         return queries, results
         
 def load_document(file: NamedTemporaryFile) -> List[Document]: