Update agent.py

agent.py

@@ -76,32 +76,63 @@ def build_graph():
     # Bind tools to LLM
     llm_with_tools = llm.bind_tools(tools)
 
-    def clean_answer(text):
-        """Extract clean answer from LLM response"""
-        if not text:
-            return ""
+    def clean_answer(answer: any) -> str:
+        """
+        Clean up the answer to remove common prefixes and formatting
+        that models often add but that can cause exact match failures.
         
-        # Remove common prefixes and suffixes
-        text = text.strip()
+        Args:
+            answer: The raw answer from the model
+            
+        Returns:
+            The cleaned answer as a string
+        """
+        # Convert non-string types to strings
+        if not isinstance(answer, str):
+            # Handle numeric types (float, int)
+            if isinstance(answer, float):
+                # Format floating point numbers properly
+                # Check if it's an integer value in float form (e.g., 12.0)
+                if answer.is_integer():
+                    formatted_answer = str(int(answer))
+                else:
+                    # For currency values that might need formatting
+                    if abs(answer) >= 1000:
+                        formatted_answer = f"${answer:,.2f}"
+                    else:
+                        formatted_answer = str(answer)
+                return formatted_answer
+            elif isinstance(answer, int):
+                return str(answer)
+            else:
+                # For any other type
+                return str(answer)
+        
+        # Now we know answer is a string, so we can safely use string methods
+        # Normalize whitespace
+        answer = answer.strip()
         
-        # Remove common response patterns
-        patterns_to_remove = [
-            r'^(The answer is:?\s*)',
-            r'^(Answer:?\s*)',
-            r'^(Final answer:?\s*)',
-            r'^(Result:?\s*)',
-            r'(\s*is the answer\.?)$',
-            r'(\s*\.)$'
+        # Remove common prefixes and formatting that models add
+        prefixes_to_remove = [
+            "The answer is ", 
+            "Answer: ",
+            "Final answer: ",
+            "The result is ",
+            "To answer this question: ",
+            "Based on the information provided, ",
+            "According to the information: ",
         ]
         
-        for pattern in patterns_to_remove:
-            text = re.sub(pattern, '', text, flags=re.IGNORECASE)
+        for prefix in prefixes_to_remove:
+            if answer.startswith(prefix):
+                answer = answer[len(prefix):].strip()
         
-        # Take only the first line if multiple lines
-        first_line = text.split('\n')[0].strip()
+        # Remove quotes if they wrap the entire answer
+        if (answer.startswith('"') and answer.endswith('"')) or (answer.startswith("'") and answer.endswith("'")):
+            answer = answer[1:-1].strip()
         
-        return first_line
-
+        return answer
+    
     def assistant(state: MessagesState):
         messages_with_system_prompt = [sys_msg] + state["messages"]
         llm_response = llm_with_tools.invoke(messages_with_system_prompt)
@@ -127,43 +158,6 @@ def build_graph():
     # Compile graph
     return builder.compile()
 
-
-def is_valid_agent_output(output):
-    """
-    Checks if the output matches the required format:
-    [{"task_id": ..., "submitted_answer": ...}]
-    """
-    try:
-        parsed = json.loads(output.strip())
-        if not isinstance(parsed, list):
-            return False
-        
-        for item in parsed:
-            if not isinstance(item, dict):
-                return False
-            if "task_id" not in item or "submitted_answer" not in item:
-                return False
-        return True
-    except:
-        return False
-
-
-def extract_flat_answer(output):
-    """Extract properly formatted answer from output"""
-    try:
-        # Try to parse as JSON first
-        parsed = json.loads(output.strip())
-        if isinstance(parsed, list) and len(parsed) > 0:
-            first_item = parsed[0]
-            if isinstance(first_item, dict) and "task_id" in first_item and "submitted_answer" in first_item:
-                return output  # Already properly formatted
-    except:
-        pass
-    
-    # If not properly formatted, return as-is (fallback)
-    return output
-
-
 # test
 if __name__ == "__main__":
     question = "What is 2 + 2?"