normalizer.py

import re

def remove_think_tags(text: str) -> str:
    """
    Removes <think> XML tags and their contents from a string.
    Handles all variations of think tags.
    
    Args:
        text (str): The input string containing potential <think> tags
        
    Returns:
        str: The input string with all <think> tags and their contents removed
    """
    # First normalize all newlines to \n
    text = text.replace('\r\n', '\n')
    
    # More flexible pattern that handles various think tag formats
    pattern = r'<think>[\s\S]*?</think>'
    cleaned = re.sub(pattern, '', text)
    
    # Clean up any resulting whitespace/newlines
    cleaned = re.sub(r'\s*\n\s*', '\n', cleaned)
    cleaned = re.sub(r'\n{2,}', '\n', cleaned)
    cleaned = re.sub(r'^\s+|\s+$', '', cleaned, flags=re.MULTILINE)
    
    return cleaned.strip()

def is_valid_think_tags(text: str) -> bool:
    """
    Validates that think tags are properly formatted.
    
    Args:
        text (str): The input string to validate
        
    Returns:
        bool: True if tags are valid, False otherwise
    """
    # Count opening and closing tags
    open_tags = text.count('<think>')
    close_tags = text.count('</think>')
    
    # Check for proper nesting
    if open_tags != close_tags:
        return False
        
    # Check for proper ordering
    pos = 0
    stack = []
    while True:
        open_pos = text.find('<think>', pos)
        close_pos = text.find('</think>', pos)
        
        if open_pos == -1 and close_pos == -1:
            break
            
        if open_pos != -1 and (close_pos == -1 or open_pos < close_pos):
            stack.append(open_pos)
            pos = open_pos + 7
        elif close_pos != -1:
            if not stack:
                return False
            stack.pop()
            pos = close_pos + 8
            
    return len(stack) == 0

# Test cases
def test_remove_think_tags():
    test_cases = [
        # Basic case
        (
            'SELECT * FROM users\n<think>\nthinking...\n</think>\nWHERE id = 1;',
            'SELECT * FROM users\nWHERE id = 1;'
        ),
        # Inline case
        (
            'SELECT * FROM users <think>thinking...</think> WHERE id = 1;',
            'SELECT * FROM users WHERE id = 1;'
        ),
        # Multiple think tags
        (
            'SELECT * <think>think1</think> FROM users <think>think2</think> WHERE id = 1;',
            'SELECT * FROM users WHERE id = 1;'
        ),
        # Nested think tags
        (
            'SELECT * FROM <think>Let me <think>deeply</think> think</think> users;',
            'SELECT * FROM users;'
        ),
        # No think tags
        (
            'SELECT * FROM users WHERE id = 1;',
            'SELECT * FROM users WHERE id = 1;'
        ),
    ]
    
    for input_text, expected in test_cases:
        result = remove_think_tags(input_text)
        assert result == expected, f'Failed: {input_text} -> {result} != {expected}'

if __name__ == "__main__":
    test_remove_think_tags()
    test_string = '''SELECT * FROM users
<think>
Let me think about the conditions...
</think>
WHERE active = true;'''
    
    result = remove_think_tags(test_string)
    print("Original:")
    print(test_string)
    print("\nProcessed:")
    print(result)