replace LiteLLMModel with LiteLLMRouterModel to add fallback feature

app.py

@@ -8,7 +8,7 @@ import pandas as pd
 from smolagents import (
     Tool,
     CodeAgent, 
-    LiteLLMModel, 
+    LiteLLMRouterModel, 
     PythonInterpreterTool,
     VisitWebpageTool,
     GoogleSearchTool, 
@@ -36,7 +36,7 @@ CUSTOM_PROMPT = """
         and separete them with a comma and a whitespace like "a, b, c" not "a,b,c".
 
     Attached files can be downloaded via the Internet.
-    When downloading files, try to download them in python code before visiting the website.
+    When downloading files, try to download them in your python code before visiting the website.
 """
 
 rate_limiter = RateLimiter()
@@ -45,9 +45,43 @@ rate_limiter = RateLimiter()
 # ----- THIS IS WERE YOU CAN BUILD WHAT YOU WANT ------
 class BasicAgent:
     def __init__(self):
-        model = LiteLLMModel(
-            model_id=os.getenv("MODEL_ID"), 
-            api_key=os.getenv("API_KEY")
+        model = LiteLLMRouterModel(
+            model_id="gemini-flash-family",
+            model_list=[
+                {
+                    "model_name": "gemini-flash-family",
+                    "litellm_params": {
+                        "model": "gemini/gemini-2.5-flash",
+                        "api_key": os.getenv("GEMINI_API_KEY"),
+                        "rpm": 10,
+                        "tpm": 250_000,
+                    },
+                    "fallbacks": ["gemini/gemini-2.0-flash"],
+                },
+                {
+                    "model_name": "gemini-flash-family",
+                    "litellm_params": {
+                        "model": "gemini/gemini-2.0-flash",
+                        "api_key": os.getenv("GEMINI_API_KEY"),
+                        "rpm": 15,
+                        "tpm": 1_000_000,
+                    },
+                    "fallbacks": ["gemini/gemini-2.0-flash-lite"]
+                },
+                {
+                    "model_name": "gemini-flash-family",
+                    "litellm_params": {
+                        "model": "gemini/gemini-2.0-flash-lite",
+                        "api_key": os.getenv("GEMINI_API_KEY"),
+                        "rpm": 30,
+                        "tpm": 1_000_000,
+                    },
+                },
+            ],
+            client_kwargs={
+                "routing_strategy": "usage-based-routing-v2"
+            },
+            num_retries=5,
         )
         self.agent = CodeAgent(
             tools=[