scripts/patch_aide.py

"""Patch AIDE for Gemini compatibility + auto token refresh. Run inside the AIDE venv."""
import aide.backend.utils as utils_mod
import aide.backend.backend_openai as backend_mod

# Patch 1: Empty user message fix
p1 = utils_mod.__file__
c1 = open(p1).read()
if "Complete the task" not in c1:
    c1 = c1.replace(
        "    return messages",
        '    if messages and not any(m["role"] == "user" for m in messages):\n        messages.append({"role": "user", "content": "Complete the task described above."})\n    return messages',
        1
    )
    open(p1, "w").write(c1)
    print("Patched utils.py")
else:
    print("utils.py already patched")

# Patch 2: max_tokens fix + auto token refresh
p2 = backend_mod.__file__
c2 = open(p2).read()

# Fix max_tokens rename
old = '    if "max_tokens" in filtered_kwargs:\n        filtered_kwargs["max_output_tokens"] = filtered_kwargs.pop("max_tokens")'
if old in c2:
    c2 = c2.replace(old, "    # max_tokens rename deferred")
    c2 = c2.replace(
        '    use_chat_api = os.getenv("OPENAI_BASE_URL") is not None and not is_openai_model',
        '    use_chat_api = os.getenv("OPENAI_BASE_URL") is not None and not is_openai_model\n\n    if "max_tokens" in filtered_kwargs and not use_chat_api:\n        filtered_kwargs["max_output_tokens"] = filtered_kwargs.pop("max_tokens")'
    )
    print("Patched max_tokens")
elif "max_tokens rename deferred" in c2:
    print("max_tokens already patched")

# Patch 3: Auto-refresh GCP token before each API call
# Replace the @once decorator on _setup_custom_client so it recreates the client with fresh token
if "_token_refresh" not in c2:
    # Add token refresh imports at top
    c2 = c2.replace(
        "import openai\n",
        "import openai\nimport subprocess\n"
    )

    # Replace _setup_custom_client to refresh token each time
    old_setup = '''@once
def _setup_custom_client():
    global _custom_client
    # Only create custom client if base URL is set
    base_url = os.getenv("OPENAI_BASE_URL")
    api_key = os.getenv("OPENAI_API_KEY")
    if base_url:
        _custom_client = openai.OpenAI(
            api_key=api_key, base_url=base_url, max_retries=0
        )'''

    new_setup = '''_token_refresh_time = 0

def _setup_custom_client():
    global _custom_client, _token_refresh_time
    import time
    base_url = os.getenv("OPENAI_BASE_URL")
    api_key = os.getenv("OPENAI_API_KEY", "")
    now = time.time()
    # Only refresh if using OAuth tokens (not API keys)
    is_api_key = api_key.startswith("AIza")
    if _custom_client is None or (not is_api_key and (now - _token_refresh_time) > 1800):
        if not is_api_key:
            try:
                import urllib.request
                req = urllib.request.Request(
                    "http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/token",
                    headers={"Metadata-Flavor": "Google"}
                )
                resp = urllib.request.urlopen(req, timeout=5)
                import json as _json
                token_data = _json.loads(resp.read())
                api_key = token_data["access_token"]
                os.environ["OPENAI_API_KEY"] = api_key
                logger.info("Refreshed GCP access token")
            except Exception:
                pass
        if base_url:
            _custom_client = openai.OpenAI(
                api_key=api_key, base_url=base_url, max_retries=0
            )
            _token_refresh_time = now'''

    if old_setup in c2:
        c2 = c2.replace(old_setup, new_setup)
        print("Patched token auto-refresh")
    else:
        print("Could not find _setup_custom_client to patch (may already be patched)")

open(p2, "w").write(c2)
print("Patched backend_openai.py")

# Patch 3b: Fix backoff to refresh token on auth errors
import aide.backend.utils as utils_mod2
p3b = utils_mod2.__file__
c3b = open(p3b).read()
if "AuthenticationError" not in c3b:
    # Add auth error to retry exceptions and force client refresh
    c3b = c3b.replace(
        "    except retry_exceptions as e:",
        """    except retry_exceptions as e:
        # Force token refresh on next API call
        try:
            import aide.backend.backend_openai as _bmod
            _bmod._token_refresh_time = 0  # Force refresh on next call
            _bmod._setup_custom_client()   # Recreate client with fresh token
        except Exception:
            pass"""
    )
    open(p3b, "w").write(c3b)
    print("Patched utils.py (token refresh on retry)")
else:
    print("utils.py retry patch already applied")

# Patch 4: Remove Kaggle references from agent prompts
import aide.agent as agent_mod
p4 = agent_mod.__file__
c4 = open(p4).read()
if "Kaggle grandmaster" in c4:
    c4 = c4.replace(
        "You are a Kaggle grandmaster attending a competition. ",
        "You are an expert ML engineer. "
    )
    open(p4, "w").write(c4)
    print("Patched agent.py (removed Kaggle references)")
else:
    print("agent.py already patched")