Update app.py

app.py

@@ -13,6 +13,7 @@ from PIL import Image
 from datetime import datetime
 import tempfile
 import helium
+import gradio as gr
 from smolagents import CodeAgent, LiteLLMModel
 from smolagents.agents import ActionStep
 from tools.search_item_ctrl_f import SearchItemCtrlFTool
@@ -115,40 +116,57 @@ for idx, tool in enumerate(tools):
     except Exception as e:
         logger.error(f"Failed to register tool {idx}: {str(e)}")
 
-# Initialize model
-model = LiteLLMModel("gemini/gemini-2.0-flash", api_key=default_gemini_api_key)
+# Initialize model with API key handling
+def initialize_model(gemini_api_key=None):
+    # Log the API key being used (mask last 4 chars for security)
+    if gemini_api_key and gemini_api_key.strip():
+        logger.debug(f"Using user-provided API key: {gemini_api_key[:-4] + '****'}")
+    else:
+        logger.debug(f"Using default API key: {default_gemini_api_key[:-4] + '****' if default_gemini_api_key else 'None'}")
+
+    try:
+        api_key = gemini_api_key.strip() if gemini_api_key and gemini_api_key.strip() else default_gemini_api_key
+        if not api_key:
+            raise ValueError("No valid API key provided and GOOGLE_API_KEY not set in environment")
+        return LiteLLMModel("gemini/gemini-2.0-flash", api_key=api_key)
+    except Exception as e:
+        logger.error(f"Failed to initialize LiteLLMModel: {str(e)}")
+        raise gr.Error(f"API Key Error: {str(e)}", duration=5)
 
 # Initialize agent
-agent = CodeAgent(
-    model=model,
-    tools=tools,
-    max_steps=20,
-    verbosity_level=2,
-    prompt_templates=prompt_templates,
-    step_callbacks=[save_screenshot],
-    additional_authorized_imports=[
-        "helium",
-        "unicodedata",
-        "stat",
-        "datetime",
-        "random",
-        "pandas",
-        "itertools",
-        "math",
-        "statistics",
-        "queue",
-        "time",
-        "collections",
-        "re",
-        "cv2",
-        "numpy"
-    ]
-)
-agent.python_executor("from helium import *")
+def initialize_agent(gemini_api_key=None):
+    model = initialize_model(gemini_api_key)
+    agent = CodeAgent(
+        model=model,
+        tools=tools,
+        max_steps=20,
+        verbosity_level=2,
+        prompt_templates=prompt_templates,
+        step_callbacks=[save_screenshot],
+        additional_authorized_imports=[
+            "helium",
+            "unicodedata",
+            "stat",
+            "datetime",
+            "random",
+            "pandas",
+            "itertools",
+            "math",
+            "statistics",
+            "queue",
+            "time",
+            "collections",
+            "re",
+            "cv2",
+            "numpy"
+        ]
+    )
+    agent.python_executor("from helium import *")
+    return agent
 
-# Launch Gradio UI
+# Launch Gradio UI with API key support
 try:
-    GradioUI(agent).launch()
+    GradioUI(initialize_agent).launch()
 except KeyboardInterrupt:
     driver.quit()
-    logger.info("Chrome driver closed on exit.") 
+    logger.info("Chrome driver closed on exit.")