utils/simple_answer_formatter.py

"""
Simple Answer Formatter - Following 100% Successful GAIA Space Patterns

This implementation abandons complex hardcoded pattern matching in favor of:
1. Trust in core agent output with proper prompting
2. Simple extraction of FINAL ANSWER format
3. Minimal post-processing focused on GAIA exact match requirements

Based on analysis of successful spaces: fisherman611/gaia-agent, baixianger/RobotPai, ZeroTimo/RobotPai
"""

import re
import logging
from typing import Optional

logger = logging.getLogger(__name__)

class SimpleGAIAAnswerFormatter:
    """
    Simple answer formatter following successful GAIA space patterns.
    
    Key principles:
    1. Trust the agent's output when properly prompted
    2. Extract FINAL ANSWER format cleanly
    3. Apply minimal GAIA-specific formatting rules
    4. No complex hardcoded pattern matching
    """
    
    def __init__(self):
        """Initialize the simple formatter."""
        logger.info("✅ Simple GAIA Answer Formatter initialized")
    
    def format_answer(self, raw_answer: str) -> str:
        """
        Format answer following successful space patterns.
        
        Args:
            raw_answer: The raw answer from the agent
            
        Returns:
            Formatted answer for GAIA evaluation
        """
        try:
            # Step 1: Extract FINAL ANSWER using simple slicing (matches 100% successful spaces)
            if "FINAL ANSWER:" in raw_answer:
                # Use simple slicing like top 100% spaces: answer[14:]
                final_answer_index = raw_answer.rfind("FINAL ANSWER:")
                if final_answer_index != -1:
                    answer = raw_answer[final_answer_index + 14:].strip()
                    # Take only the first line if multi-line
                    answer = answer.split('\n')[0].strip()
                    logger.info(f"✅ Extracted FINAL ANSWER: {answer}")
                else:
                    answer = raw_answer.strip()
                    logger.warning("⚠️ FINAL ANSWER found but extraction failed, using raw answer")
            else:
                # Fallback: use the raw answer
                answer = raw_answer.strip()
                logger.warning("⚠️ No FINAL ANSWER format found, using raw answer")
            
            # Step 2: Apply enhanced GAIA formatting rules
            formatted_answer = self._apply_enhanced_gaia_rules(answer)
            
            logger.info(f"✅ Final formatted answer: {formatted_answer}")
            return formatted_answer
            
        except Exception as e:
            logger.error(f"❌ Error formatting answer: {e}")
            return raw_answer.strip()
    
    def _apply_enhanced_gaia_rules(self, answer: str) -> str:
        """
        Apply enhanced GAIA formatting rules based on 100% successful spaces analysis.
        
        GAIA exact match requirements:
        - Numbers: No commas, no units (unless specified), consistent decimal formatting
        - Strings: No articles, no abbreviations, digits in plain text
        - Lists: Comma-separated, alphabetically sorted when appropriate
        """
        answer = answer.strip()
        
        # Rule 1: Remove common formatting artifacts
        answer = self._clean_basic_artifacts(answer)
        
        # Rule 2: Enhanced number handling (critical fix)
        answer = self._format_numbers_enhanced(answer)
        
        # Rule 3: Enhanced list processing
        answer = self._format_lists_enhanced(answer)
        
        # Rule 4: Handle common string issues
        answer = self._format_strings(answer)
        
        return answer
    
    def _clean_basic_artifacts(self, answer: str) -> str:
        """Remove basic formatting artifacts."""
        # Remove quotes around single answers
        if answer.startswith('"') and answer.endswith('"'):
            answer = answer[1:-1]
        if answer.startswith("'") and answer.endswith("'"):
            answer = answer[1:-1]
        
        # Remove trailing periods for single word/number answers
        # But preserve periods in full sentences
        words = answer.split()
        if len(words) <= 2 and not ',' in answer and answer.endswith('.'):
            answer = answer[:-1]
        
        return answer.strip()
    
    def _format_numbers_enhanced(self, answer: str) -> str:
        """Enhanced number formatting based on successful GAIA spaces analysis."""
        if ',' not in answer and '.' not in answer:
            return answer
        
        result = answer
        
        # Remove commas from large numbers (e.g., "1,234" -> "1234")
        # But preserve commas that separate list items
        while re.search(r'(\d),(\d)', result):
            result = re.sub(r'(\d),(\d)', r'\1\2', result)
        
        # Handle decimal formatting consistency
        # Ensure consistent decimal representation (no trailing zeros unless significant)
        def clean_decimal(match):
            number = match.group(0)
            try:
                # Convert to float and back to remove unnecessary trailing zeros
                float_val = float(number)
                # If it's a whole number, return as integer
                if float_val.is_integer():
                    return str(int(float_val))
                else:
                    # Remove trailing zeros after decimal point
                    return str(float_val).rstrip('0').rstrip('.')
            except ValueError:
                return number
        
        # Apply decimal cleaning to standalone numbers
        result = re.sub(r'\b\d+\.\d+\b', clean_decimal, result)
        
        return result
    
    def _format_strings(self, answer: str) -> str:
        """Format strings according to GAIA rules."""
        # This is intentionally minimal - successful spaces trust the agent
        # to provide properly formatted answers when prompted correctly
        
        # Only remove articles if this looks like a single entity name
        # Don't remove articles from full sentences
        words = answer.split()
        if len(words) <= 3:  # Only for short answers
            if answer.lower().startswith('the '):
                answer = answer[4:]
            elif answer.lower().startswith('a '):
                answer = answer[2:]
            elif answer.lower().startswith('an '):
                answer = answer[3:]
        
        return answer.strip()
    
    def _format_lists_enhanced(self, answer: str) -> str:
        """Enhanced list processing based on successful GAIA spaces analysis."""
        # Check if this looks like a comma-separated list
        if ',' in answer and len(answer.split(',')) > 1:
            items = [item.strip() for item in answer.split(',')]
            
            # Remove "and" from the last item if present
            if len(items) > 1 and items[-1].lower().startswith('and '):
                items[-1] = items[-1][4:].strip()
            
            # Check if all items are simple strings (not complex phrases)
            # Only sort if they appear to be simple names/entities
            if all(len(item.split()) <= 3 for item in items) and len(items) <= 10:
                # Sort alphabetically for consistency (common GAIA requirement)
                items.sort()
            
            return ', '.join(items)
        
        return answer

def create_simple_formatter() -> SimpleGAIAAnswerFormatter:
    """Create a simple GAIA answer formatter instance."""
    return SimpleGAIAAnswerFormatter()

# Enhanced system prompt for GAIA (following successful space patterns)
GAIA_SYSTEM_PROMPT = """You are a helpful assistant tasked with answering questions using a set of tools.

Available tools:
- Mathematical operations (add, subtract, multiply, divide, power, square_root, factorial, etc.)
- Code execution (execute_python, execute_sql, execute_bash)
- Web search (web_search, wikipedia_search, arxiv_search)
- Text processing (extract_numbers, count_words, count_characters)

For computational questions, use the code execution tools.
For current information, use the search tools.
For basic math, use the mathematical operation tools.

CRITICAL: Always end your response with the following template:
FINAL ANSWER: [YOUR FINAL ANSWER]

FORMATTING RULES (CRITICAL FOR EVALUATION):

For NUMBERS:
- Provide just the number without commas: "42" not "42,000"
- No units unless specifically requested: "3.14159" not "3.14159 meters"
- No trailing zeros: "3.5" not "3.50"
- Examples: "42", "3.14159", "1000000"

For STRINGS:
- No articles (a, an, the): "Paris" not "The Paris"
- No abbreviations: "New York" not "NY"
- Write digits in plain text: "twenty one" not "21"
- Examples: "Paris", "Albert Einstein", "twenty one"

For LISTS:
- Comma-separated values: "apple, banana, orange"
- Apply number/string rules to each element
- Alphabetical order when appropriate
- No "and" before last item: "red, blue, green" not "red, blue, and green"
- Examples: "1, 2, 3", "apple, banana, orange", "Einstein, Newton, Tesla"

EXAMPLES OF CORRECT FORMATTING:
- Question: "What is 1,234 + 5,678?" → FINAL ANSWER: 6912
- Question: "Name the capital of France" → FINAL ANSWER: Paris
- Question: "List three colors" → FINAL ANSWER: blue, green, red
- Question: "How many sides does a triangle have?" → FINAL ANSWER: 3

Your answer should ONLY start with "FINAL ANSWER: " followed by the properly formatted answer."""