Update app.py

app.py

@@ -1,7 +1,11 @@
 import os
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
+
+# LLM wrappers
 from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain_openai import ChatOpenAI
+
 from browser_use import Agent
 
 app = FastAPI()
@@ -10,22 +14,53 @@ class Task(BaseModel):
     task: str
 
 def make_llm():
+    """
+    Return a LangChain ChatModel based on env:
+      LLM_PROVIDER = 'gemini' | 'azure' | 'azure-openai'
+    """
+    provider = os.getenv("LLM_PROVIDER", "gemini").lower()
+    
     # LangChain helper; will raise if GOOGLE_API_KEY missing
     # return ChatGoogleGenerativeAI(model="gemini-2.0-flash-exp")
     
-    # Allow override from HF Secrets, but fall back to a stable, supported model
-    model_id = os.getenv("GEMINI_MODEL", "gemini-1.5-pro-latest")
-    return ChatGoogleGenerativeAI(
-        model=model_id,
-        max_retries=3,          # built-in exponential back-off
-        temperature=0.2,
-    )
+    if provider == "gemini":
+        # Google Gemini via Generative AI API
+        model = os.getenv("GEMINI_MODEL", "gemini-2.0-flash")  # or 1.5-pro, etc.
+        return ChatGoogleGenerativeAI(
+            model=model,
+            max_retries=3,
+            timeout=60,
+            temperature=0.3,
+        )
+
+    elif provider in {"azure", "azure-openai"}:
+        # Azure OpenAI  (ChatGPT family)
+        return ChatOpenAI(
+            azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT"),
+            azure_deployment_name=os.getenv("AZURE_OPENAI_DEPLOYMENT", "gpt-35-turbo"),
+            azure_api_version=os.getenv("AZURE_OPENAI_API_VERSION", "2024-02-15-preview"),
+            api_key=os.getenv("AZURE_OPENAI_API_KEY"),        # picked up by ChatOpenAI
+            max_retries=3,
+            timeout=60,
+            temperature=float(os.getenv("AZURE_OPENAI_TEMPERATURE", "0.3")),
+        )
+
+    else:
+        raise ValueError(
+            f"Unsupported LLM_PROVIDER '{provider}'. "
+            "Use 'gemini' or 'azure-openai'."
+        )
 
 @app.post("/run")
 async def run_task(t: Task):
     try:
         llm = make_llm()
-        agent = Agent(task=t.task, llm=llm)   # Browser-Use agent
+        agent = Agent(
+            task=t.task,
+            llm=llm,
+            max_steps=int(os.getenv("AGENT_MAX_STEPS", 6)),
+            step_delay=float(os.getenv("AGENT_STEP_DELAY", 2.0)),
+        )
         result = await agent.run_async()      # async version
         return result                         # Browser-Use returns dict
     except Exception as e: