Update app.py

app.py

@@ -7,7 +7,7 @@ import pandas as pd
 from langchain_core.messages import HumanMessage
 #from agent import build_graph
 from agent_simple import build_graph
-
+import re
 
 
 # (Keep Constants as is)
@@ -24,14 +24,51 @@ class BasicAgent:
         print("BasicAgent initialized.")
         self.graph = build_graph()
 
+    def extract_final_answer(self, response: str) -> str:
+        """Extract the final answer from the agent's response."""
+        # Look for patterns like "Answer: X" or "**Answer:** X"
+        answer_patterns = [
+            r'\*\*Answer:\*\*\s*(.+?)(?:\n|$)',
+            r'Answer:\s*(.+?)(?:\n|$)',
+            r'\*\*(.+?)\*\*\s*(?:studio albums?|albums?)',
+            r'(?:the answer is|therefore|thus|so)\s*[:\-]?\s*\*\*?(\d+)\*\*?',
+            r'(?:count is|total (?:would be|is))\s*\*\*?(\d+)\*\*?',
+        ]
+        
+        for pattern in answer_patterns:
+            match = re.search(pattern, response, re.IGNORECASE | re.MULTILINE)
+            if match:
+                answer = match.group(1).strip()
+                # Clean up common formatting
+                answer = re.sub(r'\*\*', '', answer)
+                answer = re.sub(r'^"|"$', '', answer)  # Remove quotes
+                return answer
+        
+        # If no pattern matches, try to extract the last meaningful sentence
+        sentences = response.split('.')
+        for sentence in reversed(sentences):
+            sentence = sentence.strip()
+            if sentence and len(sentence) < 100:  # Reasonable answer length
+                # Look for numbers or short phrases
+                if re.search(r'\d+|[a-zA-Z]{1,20}', sentence):
+                    return sentence
+        
+        # Fallback: return the last 50 characters, cleaned up
+        fallback = response.strip()[-50:].strip()
+        return fallback
+
     def __call__(self, question: str) -> str:
         print(f"Agent received question (first 50 chars): {question[:50]}...")
         # Wrap the question in a HumanMessage from langchain_core
         messages = [HumanMessage(content=question)]
         messages = self.graph.invoke({"messages": messages})
-        answer = messages['messages'][-1].content
-        return answer[14:]
-
+        full_response = messages['messages'][-1].content
+        
+        # Extract just the final answer
+        final_answer = self.extract_final_answer(full_response)
+        print(f"Extracted answer: {final_answer}")
+        
+        return final_answer
 
 def run_and_submit_all(nb_questions: int, profile: gr.OAuthProfile | None):
     """