File size: 5,022 Bytes
9f7ad84 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 | """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")
|