utils.py

import re


def extract_final_answer(output: str) -> str:
    """
    Extracts the text after 'FINAL ANSWER:' in the model's output.
    Strips whitespace and ensures clean formatting.
    If the answer is a comma-separated list, ensures a space after each comma.
    """
    output = str(output)
    marker = "FINAL ANSWER:"
    lower_output = output.lower()

    if marker.lower() in lower_output:
        # Find actual case version in original output (for safety)
        idx = lower_output.rfind(marker.lower())
        raw_answer = output[idx + len(marker) :].strip()

        # Normalize comma-separated lists: ensure single space after commas
        cleaned_answer = re.sub(r",\s*", ", ", raw_answer)
        return cleaned_answer

    return output


def replace_tool_mentions(prompt: str) -> str:
    # Replace tool mentions in backticks: `search` -> `web_search`, `wiki` -> `wikipedia_search`
    prompt = re.sub(r"(?<!\w)`search`(?!\w)", "`web_search`", prompt)
    prompt = re.sub(r"(?<!\w)`wiki`(?!\w)", "`wikipedia_search`", prompt)

    # Replace function calls: search(...) -> web_search(...), wiki(...) -> wikipedia_search(...)
    # This ensures we only catch function calls (not words like arxiv_search)
    prompt = re.sub(r"(?<!\w)(?<!_)search\(", "web_search(", prompt)
    prompt = re.sub(r"(?<!\w)(?<!_)wiki\(", "wikipedia_search(", prompt)

    return prompt