Update app.py

app.py

@@ -1,45 +1,56 @@
 import gradio as gr
 import random
-from smolagents import GradioUI, CodeAgent, HfApiModel, LiteLLMModel
-
-# Import our custom tools from their modules
+from smolagents import GradioUI, CodeAgent, LiteLLMModel
 from tools import DuckDuckGoSearchTool, WeatherInfoTool, HubStatsTool
 from retriever import load_guest_dataset
 import os
-
-#model = HfApiModel(
-#    max_tokens=2096,
-#    temperature=0,
-#    model_id='mistralai/Mistral-7B-Instruct-v0.3',
-#    custom_role_conversions=None,
-#)
-
-
-
-model = LiteLLMModel(
-    model_id="gemini/gemini-2.0-flash-lite", 
-    api_key=os.environ['GEMINI_API_TOKEN']
-) 
-
-# Initialize the web search tool
+import logging
+from tenacity import retry, stop_after_attempt, wait_exponential, retry_if_exception_type
+import litellm
+
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
+class RetryLiteLLMModel(LiteLLMModel):
+    @retry(
+        stop=stop_after_attempt(3),  # Максимум 3 попытки
+        wait=wait_exponential(multiplier=1, min=1, max=10),  # Задержка: 1с, 2с, 4с (макс. 10с)
+        retry=retry_if_exception_type(litellm.InternalServerError),  # Повторять только для InternalServerError
+        before=lambda _: logger.info("Retrying due to 503 error..."),
+        after=lambda retry_state: logger.info(f"Retry attempt {retry_state.attempt_number} completed")
+    )
+    def generate(self, *args, **kwargs):
+        """Переопределение метода generate с повторными попытками"""
+        return super().generate(*args, **kwargs)
+
+# Model wrapper
+try:
+    model = RetryLiteLLMModel(
+        model_id="gemini/gemini-2.0-flash-lite",
+        api_key=os.environ['GEMINI_API_TOKEN']
+    )
+    logger.info("Successfully initialized LiteLLM model")
+except KeyError as e:
+    logger.error("GEMINI_API_TOKEN not set in environment variables")
+    raise
+except Exception as e:
+    logger.error(f"Failed to initialize model: {str(e)}")
+    raise
+
+# Tools initialization
 search_tool = DuckDuckGoSearchTool()
-
-# Initialize the weather tool
 weather_info_tool = WeatherInfoTool()
-
-# Initialize the Hub stats tool
 hub_stats_tool = HubStatsTool()
-
-# Load the guest dataset and initialize the guest info tool
 guest_info_tool = load_guest_dataset()
 
-# Create Alfred with all the tools
+# Создание агента
 alfred = CodeAgent(
-    tools=[guest_info_tool, weather_info_tool, hub_stats_tool, search_tool], 
+    tools=[guest_info_tool, weather_info_tool, hub_stats_tool, search_tool],
     model=model,
-    add_base_tools=True,  # Add any additional base tools
-    planning_interval=3   # Enable planning every 3 steps
+    add_base_tools=True,
+    planning_interval=3
 )
 
 if __name__ == "__main__":
+    logger.info("Launching Gradio UI...")
     GradioUI(alfred).launch()