patch_models.py

"""Patch backend: add new model IDs to AVAILABLE_MODELS + fix OpenRouter routing.

OpenRouter models:
- openai/openrouter/owl-alpha — FREE by default, no :free suffix
- openai/google/gemini-2.0-flash-001 — PAID, no :free
- openai/*/...:free — free tier models, need :free suffix
"""
import ast
import os

AGENT_FILE = "/app/backend/routes/agent.py"
LLM_PARAMS_FILE = "/app/agent/core/llm_params.py"

# ── Step 1: Patch AVAILABLE_MODELS in agent.py ──
with open(AGENT_FILE) as f:
    content = f.read()

target = "AVAILABLE_MODELS = _available_models()\n"
patch = '''
# === PATCH: Add HF Inference + OpenRouter models ===
AVAILABLE_MODELS.extend([
    {"id": "deepseek-ai/DeepSeek-V4-Pro", "label": "DeepSeek V4 Pro", "provider": "huggingface", "tier": "free", "recommended": True},
    {"id": "deepseek-ai/DeepSeek-V4-Flash", "label": "DeepSeek V4 Flash", "provider": "huggingface", "tier": "free"},
    {"id": "openai/deepseek/deepseek-v4-flash:free", "label": "DeepSeek V4 Flash (OR)", "provider": "openrouter", "tier": "free"},
    {"id": "openai/nvidia/nemotron-3-super-120b-a12b:free", "label": "Nemotron 3 Super 120B", "provider": "openrouter", "tier": "free"},
    {"id": "google/gemma-3-1b-it", "label": "Gemma 3 1B", "provider": "huggingface", "tier": "free", "recommended": True},
    {"id": "Qwen/Qwen3-Coder-Next", "label": "Qwen3 Coder Next", "provider": "huggingface", "tier": "free", "recommended": True},
    {"id": "openai/openrouter/owl-alpha", "label": "Owl Alpha", "provider": "openrouter", "tier": "free", "recommended": True},
    {"id": "openai/google/gemini-2.0-flash-001", "label": "Gemini 2.0 Flash", "provider": "openrouter", "tier": "free"},
    {"id": "openai/google/gemma-4-31b-it:free", "label": "Gemma 4 31B", "provider": "openrouter", "tier": "free", "recommended": True},
    {"id": "openai/nvidia/nemotron-3-ultra-550b-a55b:free", "label": "Nemotron 3 Ultra 550B", "provider": "openrouter", "tier": "free", "recommended": True},
    {"id": "openai/poolside/laguna-m.1:free", "label": "Laguna M.1", "provider": "openrouter", "tier": "free", "recommended": True},
    {"id": "openai/poolside/laguna-xs.2:free", "label": "Laguna XS.2", "provider": "openrouter", "tier": "free"},
])
DEFAULT_FREE_MODEL_ID = "deepseek-ai/DeepSeek-V4-Pro"
# === END PATCH ===
'''

if target in content:
    content = content.replace(target, target + patch)
    print("✅ Added models to AVAILABLE_MODELS")
else:
    print("⚠ Could not find target line")

try:
    ast.parse(content)
    print("✅ agent.py syntax OK")
except SyntaxError as e:
    print(f"❌ Syntax error in agent.py: {e}")
    raise

with open(AGENT_FILE, "w") as f:
    f.write(content)

# ── Step 2: Patch _resolve_llm_params to route OpenRouter models ──
with open(LLM_PARAMS_FILE) as f:
    llm_content = f.read()

old_block = """    hf_model = normalized_model
    api_key = _resolve_hf_router_token(session_hf_token)
    params = {
        "model": f"openai/{hf_model}",
        "api_base": HF_ROUTER_BASE_URL,
        "api_key": api_key,
    }"""

new_block = """    hf_model = normalized_model
    api_key = _resolve_hf_router_token(session_hf_token)

    # === PATCH: Route OpenRouter models to openrouter.ai ===
    # OpenRouter expects the model name WITHOUT the openai/ prefix.
    # E.g. "openai/openrouter/owl-alpha" -> send "openrouter/owl-alpha"
    # No :free suffix for Owl Alpha (free by default).
    _or_prefixes = ("openai/google/", "openai/deepseek/", "openai/nvidia/", "openai/poolside/", "openai/openrouter/")
    if normalized_model.startswith(_or_prefixes):
        # Strip the "openai/" prefix so OpenRouter gets the correct model ID
        or_model = normalized_model[len("openai/"):]
        return {
            "model": or_model,
            "api_base": "https://openrouter.ai/api/v1",
            "api_key": os.environ.get("OPENROUTER_API_KEY") or api_key or "",
        }
    # === END PATCH ===

    params = {
        "model": f"openai/{hf_model}",
        "api_base": HF_ROUTER_BASE_URL,
        "api_key": api_key,
    }"""

if old_block in llm_content:
    llm_content = llm_content.replace(old_block, new_block)
    print("✅ Patched _resolve_llm_params: OpenRouter routing with stripped prefix")
else:
    print("⚠ Could not find target block in llm_params.py")
    # Debug: find the actual block
    idx = llm_content.find("hf_model = normalized_model")
    if idx >= 0:
        print(f"Found 'hf_model = normalized_model' at pos {idx}")
        print(repr(llm_content[idx:idx+300]))

try:
    ast.parse(llm_content)
    print("✅ llm_params.py syntax OK")
except SyntaxError as e:
    print(f"❌ Syntax error in llm_params.py: {e}")
    raise

with open(LLM_PARAMS_FILE, "w") as f:
    f.write(llm_content)

print("✅ Backend patched")
print("  - Owl Alpha: sends 'openrouter/owl-alpha' to openrouter.ai (NO :free)")
print("  - Gemini 2.0 Flash: sends 'google/gemini-2.0-flash-001' to openrouter.ai (NO :free)")
print("  - Other OR free models: include :free suffix")