Enhance Google Gemini integration and environment configuration

- Added support for Google Gemini API in `test_gemini_connection.py`, including model validation and content generation.
- Updated `.gitignore` to include `.env.google` for environment variable management.
- Modified `browser_use_cli.py` to dynamically set the model name from environment variables.
- Improved `get_llm_model` function in `utils.py` to handle model names and API keys more securely using `SecretStr`.
- Updated tests in `test_browser_cli.py` to set the Google model via environment variable for better test isolation.

Author: drumnation

.gitignore

@@ -193,3 +193,5 @@ data/
 
 # Brain directory
 .brain/
+
+.env.google

cli/browser_use_cli.py

@@ -121,8 +121,9 @@ async def run_browser_task(
     # Get LLM model
     llm = utils.get_llm_model(
         provider="deepseek" if model == "deepseek-chat" else model,
-        model_name=model,
-        temperature=0.8
+        model_name=os.getenv("GOOGLE_API_MODEL", "gemini-pro-vision") if model == "gemini" else model,
+        temperature=0.8,
+        vision=vision
     )
 
     # Update context with runtime options if needed

src/utils/utils.py

@@ -9,6 +9,7 @@
 import time
 from pathlib import Path
 from typing import Dict, Optional
+from pydantic import SecretStr
 
 from langchain_anthropic import ChatAnthropic
 from langchain_google_genai import ChatGoogleGenerativeAI
@@ -30,15 +31,17 @@ def get_llm_model(provider: str, **kwargs):
             base_url = kwargs.get("base_url")
 
         if not kwargs.get("api_key", ""):
-            api_key = os.getenv("ANTHROPIC_API_KEY", "")
+            api_key = SecretStr(os.getenv("ANTHROPIC_API_KEY") or "")
         else:
-            api_key = kwargs.get("api_key")
+            api_key = SecretStr(kwargs.get("api_key") or "")
 
         return ChatAnthropic(
             model_name=kwargs.get("model_name", "claude-3-5-sonnet-20240620"),
             temperature=kwargs.get("temperature", 0.0),
             base_url=base_url,
             api_key=api_key,
+            timeout=kwargs.get("timeout", 60),
+            stop=kwargs.get("stop", None)
         )
     elif provider == "openai":
         if not kwargs.get("base_url", ""):
@@ -47,15 +50,16 @@ def get_llm_model(provider: str, **kwargs):
             base_url = kwargs.get("base_url")
 
         if not kwargs.get("api_key", ""):
-            api_key = os.getenv("OPENAI_API_KEY", "")
+            api_key = SecretStr(os.getenv("OPENAI_API_KEY") or "")
         else:
-            api_key = kwargs.get("api_key")
+            api_key = SecretStr(kwargs.get("api_key") or "")
 
         return ChatOpenAI(
-            model=kwargs.get("model_name", "gpt-4o"),
+            model=kwargs.get("model_name", "gpt-4"),
             temperature=kwargs.get("temperature", 0.0),
             base_url=base_url,
             api_key=api_key,
+            timeout=kwargs.get("timeout", 60),
         )
     elif provider == "deepseek":
         if not kwargs.get("base_url", ""):
@@ -64,25 +68,37 @@ def get_llm_model(provider: str, **kwargs):
             base_url = kwargs.get("base_url")
 
         if not kwargs.get("api_key", ""):
-            api_key = os.getenv("DEEPSEEK_API_KEY", "")
+            api_key = SecretStr(os.getenv("DEEPSEEK_API_KEY") or "")
         else:
-            api_key = kwargs.get("api_key")
+            api_key = SecretStr(kwargs.get("api_key") or "")
 
         return ChatOpenAI(
             model=kwargs.get("model_name", "deepseek-chat"),
             temperature=kwargs.get("temperature", 0.0),
             base_url=base_url,
             api_key=api_key,
+            timeout=kwargs.get("timeout", 60),
         )
     elif provider == "gemini":
         if not kwargs.get("api_key", ""):
-            api_key = os.getenv("GOOGLE_API_KEY", "")
+            api_key = SecretStr(os.getenv("GOOGLE_API_KEY") or "")
         else:
-            api_key = kwargs.get("api_key")
+            api_key = SecretStr(kwargs.get("api_key") or "")
+        
+        # Get model name from environment or kwargs
+        model_name = kwargs.get("model_name")
+        if not model_name:
+            if kwargs.get("vision"):
+                model_name = os.getenv("GOOGLE_API_MODEL", "gemini-1.5-flash")
+            else:
+                model_name = os.getenv("GOOGLE_API_TYPE", "gemini-1.5-flash")
+
         return ChatGoogleGenerativeAI(
-            model=kwargs.get("model_name", "gemini-2.0-flash-exp"),
+            model=model_name,
             temperature=kwargs.get("temperature", 0.0),
-            google_api_key=api_key,
+            api_key=api_key,
+            timeout=kwargs.get("timeout", 60),
+            convert_system_message_to_human=True
         )
     elif provider == "ollama":
         return ChatOllama(
@@ -97,9 +113,9 @@ def get_llm_model(provider: str, **kwargs):
         else:
             base_url = kwargs.get("base_url")
         if not kwargs.get("api_key", ""):
-            api_key = os.getenv("AZURE_OPENAI_API_KEY", "")
+            api_key = SecretStr(os.getenv("AZURE_OPENAI_API_KEY") or "")
         else:
-            api_key = kwargs.get("api_key")
+            api_key = SecretStr(kwargs.get("api_key") or "")
         return AzureChatOpenAI(
             model=kwargs.get("model_name", "gpt-4o"),
             temperature=kwargs.get("temperature", 0.0),

test_gemini_connection.py

@@ -0,0 +1,47 @@
+import google.generativeai as genai
+import os
+from dotenv import load_dotenv, find_dotenv
+
+# Force reload of environment variables
+load_dotenv(find_dotenv(), override=True)
+
+api_key = os.environ.get("GOOGLE_API_KEY")
+model_name = os.environ.get("GOOGLE_API_MODEL")
+
+if not api_key or not model_name:
+    raise ValueError("Missing required environment variables: GOOGLE_API_KEY or GOOGLE_API_MODEL")
+
+print(f"Using model: {model_name}")
+genai.configure(api_key=api_key, transport="rest")
+
+# List all available models
+print("\nAvailable models:")
+for m in genai.list_models():
+    print(f"- {m.name}")
+
+# Check that the model exists in the client
+found_model = False
+for m in genai.list_models():
+    model_id = m.name.replace("models/", "")
+    if model_id == model_name:
+        found_model = True
+        print(f"\nFound model: {m.name}")
+        break
+
+if not found_model:
+    print("\nAvailable model IDs:")
+    for m in genai.list_models():
+        print(f"- {m.name.replace('models/', '')}")
+
+assert found_model, f"Model not found: {model_name}"
+
+# Load the model
+model = genai.GenerativeModel(model_name)
+
+# Perform a simple generation task
+try:
+    response = model.generate_content("Hello, I'm testing the Gemini API connection. Please respond with a short greeting.")
+    print(f"\nResponse: {response.text}")
+except Exception as e:
+    print(f"\nError generating content: {e}")
+    raise 

tests/test_browser_cli.py

@@ -11,6 +11,7 @@
 import asyncio
 import os
 from cli.browser_use_cli import initialize_browser, run_browser_task, close_browser, main, _global_browser, _global_browser_context
+from src.utils.utils import model_names  # Import model names from utils
 
 # Configure logging for tests
 logging.basicConfig(level=logging.INFO)
@@ -164,6 +165,7 @@ async def test_model_switching(self):
                 raise
 
         # Test Gemini
+        os.environ["GOOGLE_API_MODEL"] = model_names["gemini"][0]  # Set model via environment
         result2 = await run_browser_task(
             "what do you see on the page?",
             model="gemini",
@@ -175,6 +177,9 @@ async def test_model_switching(self):
 
     async def test_vision_capability(self):
         """Test vision capabilities"""
+        # Set Gemini model via environment
+        os.environ["GOOGLE_API_MODEL"] = model_names["gemini"][0]
+        
         # Without vision
         result1 = await run_browser_task(
             "what do you see on example.com?",