Override backend/main.py with OpenRouter patch at startup

backend/main.py

@@ -0,0 +1,126 @@
+"""FastAPI application for HF Agent web interface — patched for OpenRouter."""
+
+import logging
+import os
+import sys
+from contextlib import asynccontextmanager
+from pathlib import Path
+
+from dotenv import load_dotenv
+
+load_dotenv(Path(__file__).parent.parent / ".env")
+
+# ===== OPENROUTER PATCH =====
+# Patch _resolve_llm_params BEFORE importing routes
+sys.path.insert(0, str(Path(__file__).parent.parent))
+
+def _apply_openrouter_patch():
+    """Patch litellm routing to support OpenRouter."""
+    from agent.core import llm_params as lp
+    _original = lp._resolve_llm_params
+
+    def _patched(model_name, session_hf_token=None, reasoning_effort=None, strict=False):
+        openrouter_key = os.environ.get("OPENROUTER_API_KEY") or os.environ.get("OPENAI_API_KEY")
+        openrouter_base = os.environ.get("OPENAI_API_BASE", "")
+
+        # If OpenRouter is configured and model is not anthropic/bedrock direct
+        if openrouter_key and "openrouter" in openrouter_base:
+            if model_name.startswith("openai/"):
+                raw_model = model_name[len("openai/"):]
+                return {"model": f"openai/{raw_model}", "api_base": openrouter_base, "api_key": openrouter_key}
+            elif not model_name.startswith(("anthropic/", "bedrock/")):
+                return {"model": f"openai/{model_name}", "api_base": openrouter_base, "api_key": openrouter_key}
+
+        return _original(model_name, session_hf_token, reasoning_effort, strict)
+
+    lp._resolve_llm_params = _patched
+    print("[patch] ✅ OpenRouter LLM routing patched")
+
+try:
+    _apply_openrouter_patch()
+except Exception as e:
+    print(f"[patch] ⚠️ OpenRouter patch failed: {e}")
+# ===== END PATCH =====
+
+from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.staticfiles import StaticFiles
+
+from routes.agent import router as agent_router
+from routes.auth import router as auth_router
+from session_manager import session_manager
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
+logger = logging.getLogger(__name__)
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    logger.info("Starting HF Agent backend...")
+    await session_manager.start()
+    try:
+        import kpis_scheduler
+        kpis_scheduler.start()
+    except Exception as e:
+        logger.warning("KPI scheduler failed to start: %s", e)
+    yield
+    logger.info("Shutting down HF Agent backend...")
+    try:
+        import kpis_scheduler
+        await kpis_scheduler.shutdown()
+    except Exception as e:
+        logger.warning("KPI scheduler shutdown failed: %s", e)
+    try:
+        for sid, agent_session in list(session_manager.sessions.items()):
+            sess = agent_session.session
+            if sess.config.save_sessions:
+                try:
+                    sess.save_and_upload_detached(sess.config.session_dataset_repo)
+                except Exception:
+                    pass
+    except Exception:
+        pass
+    await session_manager.close()
+
+
+_DOCS_DISABLED = os.environ.get("SPACE_ID") is not None
+
+app = FastAPI(
+    title="HF Agent", version="1.0.0", lifespan=lifespan,
+    docs_url=None if _DOCS_DISABLED else "/docs",
+    redoc_url=None if _DOCS_DISABLED else "/redoc",
+    openapi_url=None if _DOCS_DISABLED else "/openapi.json",
+)
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["http://localhost:5173", "http://localhost:3000", "http://127.0.0.1:5173"],
+    allow_credentials=True, allow_methods=["*"], allow_headers=["*"],
+)
+
+app.include_router(agent_router)
+app.include_router(auth_router)
+
+# ===== PATCH AVAILABLE MODELS AFTER IMPORT =====
+try:
+    from routes import agent as _agent_routes
+    _agent_routes.AVAILABLE_MODELS = [
+        {"id": "deepseek/deepseek-chat-v3-0324", "label": "DeepSeek V3", "provider": "openrouter", "tier": "free", "recommended": True},
+        {"id": "google/gemini-2.0-flash-001", "label": "Gemini 2.0 Flash", "provider": "openrouter", "tier": "free"},
+        {"id": "qwen/qwen3-30b-a3b", "label": "Qwen3 30B", "provider": "openrouter", "tier": "free"},
+        {"id": "anthropic/claude-sonnet-4", "label": "Claude Sonnet 4", "provider": "openrouter", "tier": "pro"},
+    ]
+    _agent_routes.PREMIUM_MODEL_IDS = {"anthropic/claude-sonnet-4", "anthropic/claude-opus-4"}
+    _agent_routes.DEFAULT_FREE_MODEL_ID = "deepseek/deepseek-chat-v3-0324"
+    print("[patch] ✅ Available models patched (DeepSeek V3, Gemini Flash, Qwen3)")
+except Exception as e:
+    print(f"[patch] ⚠️ Models patch failed: {e}")
+# ===== END MODELS PATCH =====
+
+static_path = Path(__file__).parent.parent / "static"
+if static_path.exists():
+    app.mount("/", StaticFiles(directory=str(static_path), html=True), name="static")
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=int(os.environ.get("PORT", 7860)))