@samitouri / QOSami_Agents_Zero-2 / commits / f3401076

browser-use upgrade

frdel committed Aug 26, 2025 at 10:38 UTC f3401076cb9e511f4f33be8f9bb010ce73bdef01

4 files changed +169 -55

agent.py

@@ -335,6 +335,7 @@ class Agent:

335 await self.call_extensions("before_main_llm_call", loop_data=self.loop_data)

336

337 async def reasoning_callback(chunk: str, full: str):

338 + await self.handle_intervention()

339 if chunk == full:

340 printer.print("Reasoning: ") # start of reasoning

341 # Pass chunk and full data to extensions for processing

@@ -349,6 +350,7 @@ class Agent:

350 await self.handle_reasoning_stream(stream_data["full"])

351

352 async def stream_callback(chunk: str, full: str):

353 + await self.handle_intervention()

354 # output the agent response stream

355 if chunk == full:

356 printer.print("Response: ") # start of response

@@ -804,6 +806,7 @@ class Agent:

806 )

807

808 async def handle_reasoning_stream(self, stream: str):

809 + await self.handle_intervention()

810 await self.call_extensions(

811 "reasoning_stream",

812 loop_data=self.loop_data,

@@ -811,6 +814,7 @@ class Agent:

814 )

815

816 async def handle_response_stream(self, stream: str):

817 + await self.handle_intervention()

818 try:

819 if len(stream) < 25:

820 return # no reason to try

models.py

+158 -50

@@ -23,6 +23,7 @@ from python.helpers.dotenv import load_dotenv

23 from python.helpers.providers import get_provider_config

24 from python.helpers.rate_limiter import RateLimiter

25 from python.helpers.tokens import approximate_tokens

26 +from python.helpers import dirty_json

28 from langchain_core.language_models.chat_models import SimpleChatModel

29 from langchain_core.outputs.chat_generation import ChatGenerationChunk

@@ -90,6 +91,7 @@ class ChatChunk(TypedDict):

91 rate_limiters: dict[str, RateLimiter] = {}

92 api_keys_round_robin: dict[str, int] = {}

94 +

95 def get_api_key(service: str) -> str:

96 # get api key for the service

97 key = (

@@ -116,7 +118,14 @@ def get_rate_limiter(

118 limiter.limits["output"] = output or 0

119 return limiter

120

119 -async def apply_rate_limiter(model_config: ModelConfig|None, input_text: str, rate_limiter_callback: Callable[[str, str, int, int], Awaitable[bool]] | None = None):

121 +

122 +async def apply_rate_limiter(

123 + model_config: ModelConfig | None,

124 + input_text: str,

125 + rate_limiter_callback: (

126 + Callable[[str, str, int, int], Awaitable[bool]] | None

127 + ) = None,

128 +):

129 if not model_config:

130 return

131 limiter = get_rate_limiter(

@@ -131,25 +140,41 @@ async def apply_rate_limiter(model_config: ModelConfig|None, input_text: str, ra

140 await limiter.wait(rate_limiter_callback)

141 return limiter

142

134 -def apply_rate_limiter_sync(model_config: ModelConfig|None, input_text: str, rate_limiter_callback: Callable[[str, str, int, int], Awaitable[bool]] | None = None):

143 +

144 +def apply_rate_limiter_sync(

145 + model_config: ModelConfig | None,

146 + input_text: str,

147 + rate_limiter_callback: (

148 + Callable[[str, str, int, int], Awaitable[bool]] | None

149 + ) = None,

150 +):

151 if not model_config:

152 return

153 import asyncio, nest_asyncio

154 +

155 nest_asyncio.apply()

139 - return asyncio.run(apply_rate_limiter(model_config, input_text, rate_limiter_callback))

156 + return asyncio.run(

157 + apply_rate_limiter(model_config, input_text, rate_limiter_callback)

158 + )

159

160

161 class LiteLLMChatWrapper(SimpleChatModel):

162 model_name: str

163 provider: str

164 kwargs: dict = {}

146 -

165 +

166 class Config:

167 arbitrary_types_allowed = True

168 extra = "allow" # Allow extra attributes

169 validate_assignment = False # Don't validate on assignment

170

152 - def __init__(self, model: str, provider: str, model_config: Optional[ModelConfig] = None, **kwargs: Any):

171 + def __init__(

172 + self,

173 + model: str,

174 + provider: str,

175 + model_config: Optional[ModelConfig] = None,

176 + **kwargs: Any,

177 + ):

178 model_value = f"{provider}/{model}"

179 super().__init__(model_name=model_value, provider=provider, kwargs=kwargs) # type: ignore

180 # Set A0 model config as instance attribute after parent init

@@ -158,7 +183,7 @@ class LiteLLMChatWrapper(SimpleChatModel):

183 @property

184 def _llm_type(self) -> str:

185 return "litellm-chat"

161 -

186 +

187 def _convert_messages(self, messages: List[BaseMessage]) -> List[dict]:

188 result = []

189 # Map LangChain message types to LiteLLM roles

@@ -169,7 +194,9 @@ class LiteLLMChatWrapper(SimpleChatModel):

194 "tool": "tool",

195 }

196 for m in messages:

172 - role = role_mapping.get(m.type, m.type)

197 + m_type = getattr(m, "type", getattr(m, "role", ""))

198 + role = role_mapping.get(m_type, m_type)

199 + content = getattr(m, "content", getattr(m, "text", ""))

200 message_dict = {"role": role, "content": m.content}

201

202 # Handle tool calls for AI messages

@@ -215,12 +242,12 @@ class LiteLLMChatWrapper(SimpleChatModel):

242 **kwargs: Any,

243 ) -> str:

244 import asyncio

218 -

245 +

246 msgs = self._convert_messages(messages)

220 -

247 +

248 # Apply rate limiting if configured

249 apply_rate_limiter_sync(self.a0_model_conf, str(msgs))

223 -

250 +

251 # Call the model

252 resp = completion(

253 model=self.model_name, messages=msgs, stop=stop, **{**self.kwargs, **kwargs}

@@ -238,12 +265,12 @@ class LiteLLMChatWrapper(SimpleChatModel):

265 **kwargs: Any,

266 ) -> Iterator[ChatGenerationChunk]:

267 import asyncio

241 -

268 +

269 msgs = self._convert_messages(messages)

243 -

270 +

271 # Apply rate limiting if configured

272 apply_rate_limiter_sync(self.a0_model_conf, str(msgs))

246 -

273 +

274 for chunk in completion(

275 model=self.model_name,

276 messages=msgs,

@@ -266,11 +293,10 @@ class LiteLLMChatWrapper(SimpleChatModel):

293 **kwargs: Any,

294 ) -> AsyncIterator[ChatGenerationChunk]:

295 msgs = self._convert_messages(messages)

269 -

296 +

297 # Apply rate limiting if configured

298 await apply_rate_limiter(self.a0_model_conf, str(msgs))

272 -

273 -

299 +

300 response = await acompletion(

301 model=self.model_name,

302 messages=msgs,

@@ -294,7 +320,9 @@ class LiteLLMChatWrapper(SimpleChatModel):

320 response_callback: Callable[[str, str], Awaitable[None]] | None = None,

321 reasoning_callback: Callable[[str, str], Awaitable[None]] | None = None,

322 tokens_callback: Callable[[str, int], Awaitable[None]] | None = None,

297 - rate_limiter_callback: Callable[[str, str, int, int], Awaitable[bool]] | None = None,

323 + rate_limiter_callback: (

324 + Callable[[str, str, int, int], Awaitable[bool]] | None

325 + ) = None,

326 **kwargs: Any,

327 ) -> Tuple[str, str]:

328

@@ -312,7 +340,9 @@ class LiteLLMChatWrapper(SimpleChatModel):

340 msgs_conv = self._convert_messages(messages)

341

342 # Apply rate limiting if configured

315 - limiter = await apply_rate_limiter(self.a0_model_conf, str(msgs_conv), rate_limiter_callback)

343 + limiter = await apply_rate_limiter(

344 + self.a0_model_conf, str(msgs_conv), rate_limiter_callback

345 + )

346

347 # call model

348 _completion = await acompletion(

@@ -360,7 +390,27 @@ class LiteLLMChatWrapper(SimpleChatModel):

390 return response, reasoning

391

392

363 -class BrowserCompatibleChatWrapper(LiteLLMChatWrapper):

393 +class AsyncAIChatReplacement:

394 + class _Completions:

395 + def __init__(self, wrapper):

396 + self._wrapper = wrapper

397 +

398 + async def create(self, *args, **kwargs):

399 + # call the async _acall method on the wrapper

400 + return await self._wrapper._acall(*args, **kwargs)

401 +

402 + class _Chat:

403 + def __init__(self, wrapper):

404 + self.completions = AsyncAIChatReplacement._Completions(wrapper)

405 +

406 + def __init__(self, wrapper, *args, **kwargs):

407 + self._wrapper = wrapper

408 + self.chat = AsyncAIChatReplacement._Chat(wrapper)

409 +

410 +

411 +from browser_use.llm import ChatOllama, ChatOpenRouter, ChatGoogle, ChatAnthropic, ChatGroq, ChatOpenAI

412 +

413 +class BrowserCompatibleChatWrapper(ChatOpenRouter):

414 """

415 A wrapper for browser agent that can filter/sanitize messages

416 before sending them to the LLM.

@@ -368,31 +418,61 @@ class BrowserCompatibleChatWrapper(LiteLLMChatWrapper):

418

419 def __init__(self, *args, **kwargs):

420 turn_off_logging()

371 - super().__init__(*args, **kwargs)

421 + # Create the underlying LiteLLM wrapper

422 + self._wrapper = LiteLLMChatWrapper(*args, **kwargs)

423 # Browser-use may expect a 'model' attribute

373 - self.model = self.model_name

424 + self.model = self._wrapper.model_name

425 + self.kwargs = self._wrapper.kwargs

426

375 - def _call(

427 + @property

428 + def model_name(self) -> str:

429 + return self._wrapper.model_name

430 +

431 + @property

432 + def provider(self) -> str:

433 + return self._wrapper.provider

434 +

435 + def get_client(self, *args, **kwargs): # type: ignore

436 + return AsyncAIChatReplacement(self, *args, **kwargs)

437 +

438 + async def _acall(

439 self,

440 messages: List[BaseMessage],

441 stop: Optional[List[str]] = None,

442 run_manager: Optional[CallbackManagerForLLMRun] = None,

443 **kwargs: Any,

381 - ) -> str:

382 - turn_off_logging()

383 - result = super()._call(messages, stop, run_manager, **kwargs)

384 - return result

444 + ):

445 + # Apply rate limiting if configured

446 + apply_rate_limiter_sync(self._wrapper.a0_model_conf, str(messages))

447

386 - async def _astream(

387 - self,

388 - messages: List[BaseMessage],

389 - stop: Optional[List[str]] = None,

390 - run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,

391 - **kwargs: Any,

392 - ) -> AsyncIterator[ChatGenerationChunk]:

393 - turn_off_logging()

394 - async for chunk in super()._astream(messages, stop, run_manager, **kwargs):

395 - yield chunk

448 + # Call the model

449 + try:

450 + model = kwargs.pop("model", None)

451 + kwrgs = {**self._wrapper.kwargs, **kwargs}

452 +

453 + # hack from browser-use to fix json schema for gemini

454 + if "response_format" in kwrgs and "json_schema" in kwrgs["response_format"] and model.startswith("gemini/"):

455 + kwrgs["response_format"]["json_schema"] = ChatGoogle("")._fix_gemini_schema(self._wrapper.kwargs)

456 +

457 + resp = await acompletion(

458 + model=self._wrapper.model_name,

459 + messages=messages,

460 + stop=stop,

461 + **kwrgs,

462 + )

463 + except Exception as e:

464 + raise e

465 +

466 + # another hack for browser-use post process invalid jsons

467 + try:

468 + if "response_format" in kwrgs and "json_schema" in kwrgs["response_format"] or "json_object" in kwrgs["response_format"]:

469 + if resp.choices[0].message.content is not None and not resp.choices[0].message.content.startswith("{"): # type: ignore

470 + js = dirty_json.parse(resp.choices[0].message.content) # type: ignore

471 + resp.choices[0].message.content = dirty_json.stringify(js) # type: ignore

472 + except Exception as e:

473 + pass

474 +

475 + return resp

476

477

478 class LiteLLMEmbeddingWrapper(Embeddings):

@@ -400,15 +480,21 @@ class LiteLLMEmbeddingWrapper(Embeddings):

480 kwargs: dict = {}

481 a0_model_conf: Optional[ModelConfig] = None

482

403 - def __init__(self, model: str, provider: str, model_config: Optional[ModelConfig] = None, **kwargs: Any):

483 + def __init__(

484 + self,

485 + model: str,

486 + provider: str,

487 + model_config: Optional[ModelConfig] = None,

488 + **kwargs: Any,

489 + ):

490 self.model_name = f"{provider}/{model}" if provider != "openai" else model

491 self.kwargs = kwargs

492 self.a0_model_conf = model_config

407 -

493 +

494 def embed_documents(self, texts: List[str]) -> List[List[float]]:

495 # Apply rate limiting if configured

496 apply_rate_limiter_sync(self.a0_model_conf, " ".join(texts))

411 -

497 +

498 resp = embedding(model=self.model_name, input=texts, **self.kwargs)

499 return [

500 item.get("embedding") if isinstance(item, dict) else item.embedding # type: ignore

@@ -418,7 +504,7 @@ class LiteLLMEmbeddingWrapper(Embeddings):

504 def embed_query(self, text: str) -> List[float]:

505 # Apply rate limiting if configured

506 apply_rate_limiter_sync(self.a0_model_conf, text)

421 -

507 +

508 resp = embedding(model=self.model_name, input=[text], **self.kwargs)

509 item = resp.data[0] # type: ignore

510 return item.get("embedding") if isinstance(item, dict) else item.embedding # type: ignore

@@ -427,7 +513,13 @@ class LiteLLMEmbeddingWrapper(Embeddings):

513 class LocalSentenceTransformerWrapper(Embeddings):

514 """Local wrapper for sentence-transformers models to avoid HuggingFace API calls"""

515

430 - def __init__(self, provider: str, model: str, model_config: Optional[ModelConfig] = None, **kwargs: Any):

516 + def __init__(

517 + self,

518 + provider: str,

519 + model: str,

520 + model_config: Optional[ModelConfig] = None,

521 + **kwargs: Any,

522 + ):

523 # Clean common user-input mistakes

524 model = model.strip().strip('"').strip("'")

525

@@ -449,18 +541,18 @@ class LocalSentenceTransformerWrapper(Embeddings):

541 self.model = SentenceTransformer(model, **st_kwargs)

542 self.model_name = model

543 self.a0_model_conf = model_config

452 -

544 +

545 def embed_documents(self, texts: List[str]) -> List[List[float]]:

546 # Apply rate limiting if configured

547 apply_rate_limiter_sync(self.a0_model_conf, " ".join(texts))

456 -

548 +

549 embeddings = self.model.encode(texts, convert_to_tensor=False) # type: ignore

550 return embeddings.tolist() if hasattr(embeddings, "tolist") else embeddings # type: ignore

551

552 def embed_query(self, text: str) -> List[float]:

553 # Apply rate limiting if configured

554 apply_rate_limiter_sync(self.a0_model_conf, text)

463 -

555 +

556 embedding = self.model.encode([text], convert_to_tensor=False) # type: ignore

557 result = (

558 embedding[0].tolist() if hasattr(embedding[0], "tolist") else embedding[0]

@@ -485,10 +577,17 @@ def _get_litellm_chat(

577 provider_name, model_name, kwargs = _adjust_call_args(

578 provider_name, model_name, kwargs

579 )

488 - return cls(provider=provider_name, model=model_name, model_config=model_config, **kwargs)

580 + return cls(

581 + provider=provider_name, model=model_name, model_config=model_config, **kwargs

582 + )

583

584

491 -def _get_litellm_embedding(model_name: str, provider_name: str, model_config: Optional[ModelConfig] = None, **kwargs: Any):

585 +def _get_litellm_embedding(

586 + model_name: str,

587 + provider_name: str,

588 + model_config: Optional[ModelConfig] = None,

589 + **kwargs: Any,

590 +):

591 # Check if this is a local sentence-transformers model

592 if provider_name == "huggingface" and model_name.startswith(

593 "sentence-transformers/"

@@ -498,7 +597,10 @@ def _get_litellm_embedding(model_name: str, provider_name: str, model_config: Op

597 provider_name, model_name, kwargs

598 )

599 return LocalSentenceTransformerWrapper(

501 - provider=provider_name, model=model_name, model_config=model_config, **kwargs

600 + provider=provider_name,

601 + model=model_name,

602 + model_config=model_config,

603 + **kwargs,

604 )

605

606 # use api key from kwargs or env

@@ -511,7 +613,9 @@ def _get_litellm_embedding(model_name: str, provider_name: str, model_config: Op

613 provider_name, model_name, kwargs = _adjust_call_args(

614 provider_name, model_name, kwargs

615 )

514 - return LiteLLMEmbeddingWrapper(model=model_name, provider=provider_name, model_config=model_config, **kwargs)

616 + return LiteLLMEmbeddingWrapper(

617 + model=model_name, provider=provider_name, model_config=model_config, **kwargs

618 + )

619

620

621 def _parse_chunk(chunk: Any) -> ChatChunk:

@@ -599,10 +703,14 @@ def _merge_provider_defaults(

703 return provider_name, kwargs

704

705

602 -def get_chat_model(provider: str, name: str, model_config: Optional[ModelConfig] = None, **kwargs: Any) -> LiteLLMChatWrapper:

706 +def get_chat_model(

707 + provider: str, name: str, model_config: Optional[ModelConfig] = None, **kwargs: Any

708 +) -> LiteLLMChatWrapper:

709 orig = provider.lower()

710 provider_name, kwargs = _merge_provider_defaults("chat", orig, kwargs)

605 - return _get_litellm_chat(LiteLLMChatWrapper, name, provider_name, model_config, **kwargs)

711 + return _get_litellm_chat(

712 + LiteLLMChatWrapper, name, provider_name, model_config, **kwargs

713 + )

714

715

716 def get_browser_model(

python/tools/browser_agent.py

+3 -1

@@ -48,6 +48,7 @@ class State:

48 accept_downloads=True,

49 downloads_dir=files.get_abs_path("tmp/downloads"),

50 downloads_path=files.get_abs_path("tmp/downloads"),

51 + allowed_domains=["*"],

52 executable_path=pw_binary,

53 keep_alive=True,

54 minimum_wait_page_load_time=1.0,

@@ -143,6 +144,7 @@ class State:

144 ),

145 controller=controller,

146 enable_memory=False, # Disable memory to avoid state conflicts

147 + llm_timeout=3000, # TODO rem

148 sensitive_data=cast(dict[str, str | dict[str, str]] | None, secrets_dict or {}), # Pass secrets

149 )

150 except Exception as e:

@@ -382,7 +384,7 @@ class BrowserAgent(Tool):

384 def get_use_agent_log(use_agent: browser_use.Agent | None):

385 result = ["🚦 Starting task"]

386 if use_agent:

385 - action_results = use_agent.state.history.action_results()

387 + action_results = use_agent.history.action_results() or []

388 short_log = []

389 for item in action_results:

390 # final results

requirements.txt

+4 -4

@@ -1,6 +1,6 @@

1 a2wsgi==1.10.8

2 ansio==0.0.1

3 -browser-use==0.2.5

3 +browser-use==0.5.11

4 docker==7.1.0

5 duckduckgo-search==6.1.12

6 faiss-cpu==1.11.0

@@ -19,11 +19,11 @@ langchain-unstructured[all-docs]==0.1.6

19 openai-whisper==20240930

20 lxml_html_clean==0.3.1

21 markdown==3.7

22 -mcp==1.9.0

22 +mcp==1.13.1

23 newspaper3k==0.2.8

24 paramiko==3.5.0

25 playwright==1.52.0

26 -pypdf==4.3.1

26 +pypdf==6.0.0

27 python-dotenv==1.1.0

28 pytz==2024.2

29 sentence-transformers==3.0.1

@@ -33,7 +33,7 @@ unstructured-client==0.31.0

33 webcolors==24.6.0

34 nest-asyncio==1.6.0

35 crontab==1.0.1

36 -litellm==1.76

36 +litellm==1.75.3

37 markdownify==1.1.0

38 pymupdf==1.25.3

39 pytesseract==0.3.13