post-processing step

app.py

@@ -20,6 +20,7 @@ from tools import (
     file_type_detection_tool,
     read_file_tool
 )
+import re
 
 # (Keep Constants as is)
 # --- Constants ---
@@ -104,6 +105,26 @@ def process_question_with_files(question_data: dict) -> str:
         return f"{question_text}\n\n[Note: Error processing attached file {file_name}: {str(e)}]"
 
 
+def extract_final_answer(text: str) -> str:
+    # Remove common prefixes
+    text = re.sub(r'(?i)(answer:|final answer:|the answer is:)', '', text)
+    # Remove repeated question lines
+    lines = [line for line in text.strip().split(
+        '\n') if not line.strip().endswith('?')]
+    # If the answer is a number at the end, return it
+    match = re.search(r'\b\d+\b$', text.strip())
+    if match:
+        return match.group(0)
+    # If the answer is a comma-separated list, return it
+    if ',' in text and len(text.split(',')) <= 10:
+        return ','.join([x.strip() for x in text.split(',') if x.strip()])
+    # Otherwise, return the last non-empty line
+    for line in reversed(lines):
+        if line.strip():
+            return line.strip()
+    return text.strip()
+
+
 class AgentWrapper:
     def __init__(self):
         print("AgentWrapper initialized.")
@@ -122,8 +143,9 @@ class AgentWrapper:
                 {"messages": [HumanMessage(content=question_text)]})
             last_message = result["messages"][-1]
             answer = last_message.content
-            print(f"Agent returning answer: {answer}")
-            return answer
+            final_answer = extract_final_answer(answer)
+            print(f"Agent returning answer: {final_answer}")
+            return final_answer
         except Exception as e:
             print(f"Error in agent processing: {e}")
             return f"Error processing question: {e}"