browser use upgrade

models.py

@@ -19,7 +19,7 @@ import litellm
 import openai
 
 from python.helpers import dotenv
-from python.helpers import settings
+from python.helpers import settings, dirty_json
 from python.helpers.dotenv import load_dotenv
 from python.helpers.providers import get_provider_config
 from python.helpers.rate_limiter import RateLimiter
@@ -545,7 +545,27 @@ class LiteLLMChatWrapper(SimpleChatModel):
                 await asyncio.sleep(retry_delay_s)
 
 
-class BrowserCompatibleChatWrapper(LiteLLMChatWrapper):
+class AsyncAIChatReplacement:
+    class _Completions:
+        def __init__(self, wrapper):
+            self._wrapper = wrapper
+
+        async def create(self, *args, **kwargs):
+            # call the async _acall method on the wrapper
+            return await self._wrapper._acall(*args, **kwargs)
+
+    class _Chat:
+        def __init__(self, wrapper):
+            self.completions = AsyncAIChatReplacement._Completions(wrapper)
+
+    def __init__(self, wrapper, *args, **kwargs):
+        self._wrapper = wrapper
+        self.chat = AsyncAIChatReplacement._Chat(wrapper)
+
+
+from browser_use.llm import ChatOllama, ChatOpenRouter, ChatGoogle, ChatAnthropic, ChatGroq, ChatOpenAI
+
+class BrowserCompatibleChatWrapper(ChatOpenRouter):
     """
     A wrapper for browser agent that can filter/sanitize messages
     before sending them to the LLM.
@@ -553,32 +573,61 @@ class BrowserCompatibleChatWrapper(LiteLLMChatWrapper):
 
     def __init__(self, *args, **kwargs):
         turn_off_logging()
-        super().__init__(*args, **kwargs)
+        # Create the underlying LiteLLM wrapper
+        self._wrapper = LiteLLMChatWrapper(*args, **kwargs)
         # Browser-use may expect a 'model' attribute
-        self.model = self.model_name
+        self.model = self._wrapper.model_name
+        self.kwargs = self._wrapper.kwargs
 
-    def _call(
-        self,
-        messages: List[BaseMessage],
-        stop: Optional[List[str]] = None,
-        run_manager: Optional[CallbackManagerForLLMRun] = None,
-        **kwargs: Any,
-    ) -> str:
-        turn_off_logging()
-        result = super()._call(messages, stop, run_manager, **kwargs)
-        return result
+    @property
+    def model_name(self) -> str:
+        return self._wrapper.model_name
 
-    async def _astream(
+    @property
+    def provider(self) -> str:
+        return self._wrapper.provider
+
+    def get_client(self, *args, **kwargs):  # type: ignore
+        return AsyncAIChatReplacement(self, *args, **kwargs)
+
+    async def _acall(
         self,
         messages: List[BaseMessage],
         stop: Optional[List[str]] = None,
-        run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
+        run_manager: Optional[CallbackManagerForLLMRun] = None,
         **kwargs: Any,
-    ) -> AsyncIterator[ChatGenerationChunk]:
-        turn_off_logging()
-        async for chunk in super()._astream(messages, stop, run_manager, **kwargs):
-            yield chunk
+    ):
+        # Apply rate limiting if configured
+        apply_rate_limiter_sync(self._wrapper.a0_model_conf, str(messages))
 
+        # Call the model
+        try:
+            model = kwargs.pop("model", None)
+            kwrgs = {**self._wrapper.kwargs, **kwargs}
+
+            # hack from browser-use to fix json schema for gemini
+            if "response_format" in kwrgs and "json_schema" in kwrgs["response_format"] and model.startswith("gemini/"):
+                kwrgs["response_format"]["json_schema"] = ChatGoogle("")._fix_gemini_schema(self._wrapper.kwargs)
+
+            resp = await acompletion(
+                model=self._wrapper.model_name,
+                messages=messages,
+                stop=stop,
+                **kwrgs,
+            )
+        except Exception as e:
+            raise e
+
+        # another hack for browser-use post process invalid jsons
+        try:
+            if "response_format" in kwrgs and "json_schema" in kwrgs["response_format"] or "json_object" in kwrgs["response_format"]:
+                if resp.choices[0].message.content is not None and not resp.choices[0].message.content.startswith("{"): # type: ignore
+                    js = dirty_json.parse(resp.choices[0].message.content) # type: ignore
+                    resp.choices[0].message.content = dirty_json.stringify(js) # type: ignore
+        except Exception as e:
+            pass
+
+        return resp
 
 class LiteLLMEmbeddingWrapper(Embeddings):
     model_name: str

python/tools/browser_agent.py

@@ -148,6 +148,7 @@ class State:
                 ),
                 controller=controller,
                 enable_memory=False,  # Disable memory to avoid state conflicts
+                llm_timeout=3000, # TODO rem
                 sensitive_data=cast(dict[str, str | dict[str, str]] | None, secrets_dict or {}),  # Pass secrets
             )
         except Exception as e:
@@ -387,7 +388,7 @@ class BrowserAgent(Tool):
 def get_use_agent_log(use_agent: browser_use.Agent | None):
     result = ["🚦 Starting task"]
     if use_agent:
-        action_results = use_agent.state.history.action_results() or []
+        action_results = use_agent.history.action_results() or []
         short_log = []
         for item in action_results:
             # final results

requirements.txt

@@ -1,6 +1,6 @@
 a2wsgi==1.10.8
 ansio==0.0.1
-browser-use==0.2.5
+browser-use==0.5.11
 docker==7.1.0
 duckduckgo-search==6.1.12
 faiss-cpu==1.11.0
@@ -19,7 +19,7 @@ langchain-unstructured[all-docs]==0.1.6
 openai-whisper==20240930
 lxml_html_clean==0.3.1
 markdown==3.7
-mcp==1.12.4
+mcp==1.13.1
 newspaper3k==0.2.8
 paramiko==3.5.0
 playwright==1.52.0

run_ui.py

@@ -1,3 +1,4 @@
+import asyncio
 from datetime import timedelta
 import os
 import secrets
@@ -8,6 +9,7 @@ import struct
 from functools import wraps
 import threading
 from flask import Flask, request, Response, session, redirect, url_for, render_template_string
+from werkzeug.wrappers.response import Response as BaseResponse
 import initialize
 from python.helpers import files, git, mcp_server, fasta2a_server
 from python.helpers.files import get_abs_path
@@ -217,7 +219,7 @@ def run():
         name = handler.__module__.split(".")[-1]
         instance = handler(app, lock)
 
-        async def handler_wrap():
+        async def handler_wrap() -> BaseResponse:
             return await instance.handle_request(request=request)
 
         if handler.requires_loopback():