Update app.py

app.py

@@ -102,24 +102,41 @@ class Agent1:
 
     def update_context(self, query: str):
         tokens = nltk.pos_tag(word_tokenize(query))
+        important_phrases = []
+        current_phrase = []
+
         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
+            if tag.startswith('NN') or tag.startswith('JJ') or tag == 'NNP':
+                current_phrase.append(word)
+            else:
+                if current_phrase:
+                    important_phrases.append(' '.join(current_phrase))
+                    current_phrase = []
+
+        if current_phrase:
+            important_phrases.append(' '.join(current_phrase))
+
+        if important_phrases:
+            self.context['main_topic'] = important_phrases[0]  # Use the first important phrase as main topic
+            self.context['related_topics'] = important_phrases[1:]  # Store other phrases as related topics
 
     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}"
+        # Check if the query is short, contains pronouns, or doesn't contain the main topic
+        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())):
+            
+            # Apply main topic context
+            if 'main_topic' in self.context:
+                query = f"{self.context['main_topic']} {query}"
+            
+            # Apply related topics if query is very short
+            if len(words) <= 3 and 'related_topics' in self.context:
+                for topic in self.context['related_topics']:
+                    if topic.lower() not in query.lower():
+                        query += f" {topic}"
         
         return query