#!/usr/bin/env python3
"""
Patch litellm to properly detect Together.ai context window exceeded errors.

This script modifies the exception_mapping_utils.py file in the litellm package
to add the Together.ai specific error message pattern.
"""
import sys
from pathlib import Path


def find_litellm_exception_mapping():
    """Find the litellm exception_mapping_utils.py file."""
    try:
        import litellm
        litellm_path = Path(litellm.__file__).parent
        exception_mapping_path = litellm_path / "litellm_core_utils" / "exception_mapping_utils.py"
        return exception_mapping_path
    except ImportError:
        print("Error: litellm not found", file=sys.stderr)
        return None


def patch_context_window_detection(file_path):
    """Add the Together.ai error pattern to context window detection."""
    content = file_path.read_text()
    
    # Check if already patched
    if '"`inputs` tokens + `max_new_tokens` must be"' in content:
        print(f"Already patched: {file_path}")
        return True
    
    # Find the is_context_window_exceeded_error method
    search_pattern = '"is longer than the model\'s context length",'
    if search_pattern not in content:
        print(f"Error: Could not find pattern to patch in {file_path}", file=sys.stderr)
        return False
    
    # Add the new pattern after the existing one
    new_pattern = '"is longer than the model\'s context length",\n        "`inputs` tokens + `max_new_tokens` must be",'
    content = content.replace(search_pattern, new_pattern)
    
    # Write the patched content back
    file_path.write_text(content)
    print(f"Successfully patched: {file_path}")
    return True


def main():
    exception_mapping_path = find_litellm_exception_mapping()
    if not exception_mapping_path:
        sys.exit(1)
    
    if not exception_mapping_path.exists():
        print(f"Error: File not found: {exception_mapping_path}", file=sys.stderr)
        sys.exit(1)
    
    if not patch_context_window_detection(exception_mapping_path):
        sys.exit(1)
    
    print("Patch completed successfully!")


if __name__ == "__main__":
    main()

