update final answer extraction

app.py

@@ -14,6 +14,7 @@ from youtube_transcript_api._errors import TranscriptsDisabled, NoTranscriptFoun
 from urllib.parse import urlparse, parse_qs
 import json
 import whisper
+import re
 
 
 
@@ -153,12 +154,43 @@ class BasicAgent:
         )
         print("BasicAgent initialized.")
 
-    def __call__(self, task_id:str, question: str, file_name: str) -> str:
+    def __call__(self, task_id: str, question: str, file_name: str) -> str:
         if file_name:
             question = self.enrich_question_with_associated_file_details(task_id, question, file_name)
 
         final_result = self.code_agent.run(question)
-        return str(final_result)
+
+        # Extract text after "FINAL ANSWER:" (case-insensitive, and trims whitespace)
+        match = re.search(r'final answer:\s*(.*)', str(final_result), re.IGNORECASE | re.DOTALL)
+        if match:
+            return match.group(1).strip()
+
+        # Fallback in case the pattern is not found
+        return str(final_result).strip()
+
+    def enrich_question_with_associated_file_details(self, task_id:str, question: str, file_name: str) -> str:
+        api_url = DEFAULT_API_URL
+        get_associated_files_url = f"{api_url}/files/{task_id}"
+        response = requests.get(get_associated_files_url, timeout=15)
+        response.raise_for_status()
+
+        if file_name.endswith(".mp3"):
+            with tempfile.NamedTemporaryFile(delete=False, suffix=".mp3") as tmp_file:
+                tmp_file.write(response.content)
+                file_path = tmp_file.name
+                return question + "\n\nMentioned .mp3 file local path is: " + file_path
+        elif file_name.endswith(".py"):
+            file_content = response.text
+            return question + "\n\nBelow is mentioned Python file:\n\n```python\n" + file_content + "\n```\n"
+        elif file_name.endswith(".xlsx"):
+            xlsx_io = BytesIO(response.content)
+            df = pd.read_excel(xlsx_io)
+            file_content = df.to_csv(index=False)
+            return question + "\n\nBelow is mentioned excel file in CSV format:\n\n```csv\n" + file_content + "\n```\n"
+        elif file_name.endswith(".png"):
+            base64_str = base64.b64encode(response.content).decode('utf-8')
+            return question + "\n\nBelow is the .png image in base64 format:\n\n```base64\n" + base64_str + "\n```\n"
+
 
     def enrich_question_with_associated_file_details(self, task_id:str, question: str, file_name: str) -> str:
         api_url = DEFAULT_API_URL