Update app.py

app.py

@@ -35,7 +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
+        self.context = {}
 
     def is_question(self, text: str) -> bool:
         words = word_tokenize(text.lower())
@@ -58,13 +58,11 @@ class Agent1:
         if len(questions) < 2:
             return questions
 
-        # Find the subject in the first question
         subject = self.find_subject(questions[0])
         
         if not subject:
             return questions
 
-        # Replace pronouns in subsequent questions
         for i in range(1, len(questions)):
             words = word_tokenize(questions[i])
             for j, word in enumerate(words):
@@ -91,60 +89,59 @@ class Agent1:
             if self.is_question(' '.join(current_question)):
                 questions.append(' '.join(current_question))
 
-        # If no questions identified, return the original input
         if not questions:
             return [user_input]
 
-        # Handle pronoun replacement
         questions = self.replace_pronoun(questions)
 
         return questions
 
     def update_context(self, query: str):
         tokens = nltk.pos_tag(word_tokenize(query))
-        important_phrases = []
+        noun_phrases = []
         current_phrase = []
 
         for word, tag in tokens:
-            if tag.startswith('NN') or tag.startswith('JJ') or tag == 'NNP':
+            if tag.startswith('NN') or tag.startswith('JJ'):
                 current_phrase.append(word)
             else:
                 if current_phrase:
-                    important_phrases.append(' '.join(current_phrase))
+                    noun_phrases.append(' '.join(current_phrase))
                     current_phrase = []
 
         if current_phrase:
-            important_phrases.append(' '.join(current_phrase))
+            noun_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
+        if noun_phrases:
+            self.context['main_topic'] = noun_phrases[0]
+            self.context['related_topics'] = noun_phrases[1:]
 
     def apply_context(self, query: str) -> str:
         words = word_tokenize(query.lower())
         
-        # 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}"
-        
+            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)
+
+            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)
+
         return query
 
     def process(self, user_input: str) -> tuple[List[str], Dict[str, List[Dict[str, str]]]]:
-        # 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)