From d9d6fb2e420eec4d3c214d82c757df64e5763115 Mon Sep 17 00:00:00 2001 From: STJ Date: Fri, 13 Mar 2026 14:38:41 -0700 Subject: [PATCH] metadata + cleanup some actions --- containers/Dockerfile | 2 +- .../tool_components/browser_renderer.py | 5 +- strix/llm/llm.py | 1 + strix/tools/browser/browser_actions.py | 116 +++++++++++------- .../tools/browser/browser_actions_schema.xml | 66 +--------- strix/tools/browser/browser_manager.py | 4 +- strix/tools/browser/litellm/chat.py | 15 +-- 7 files changed, 87 insertions(+), 122 deletions(-) diff --git a/containers/Dockerfile b/containers/Dockerfile index d385c7e7..fe7a87a8 100644 --- a/containers/Dockerfile +++ b/containers/Dockerfile @@ -33,7 +33,7 @@ RUN apt-get update && \ net-tools dnsutils whois \ jq parallel ripgrep grep \ less man-db procps htop \ - iproute2 iputils-ping netcat-traditional socat \ + iproute2 iputils-ping netcat-traditional \ nmap ncat ndiff \ sqlmap nuclei subfinder naabu ffuf \ nodejs npm pipx \ diff --git a/strix/interface/tool_components/browser_renderer.py b/strix/interface/tool_components/browser_renderer.py index 799ce1f7..d66f9352 100644 --- a/strix/interface/tool_components/browser_renderer.py +++ b/strix/interface/tool_components/browser_renderer.py @@ -387,9 +387,10 @@ class BrowserRenderer(BaseToolRenderer): @classmethod def _build_fallback(cls, action: str, status: str) -> Text: + if not action: + return Text() text = Text("@ ", style=cls.DIM) - if action: - text.append(action, style=cls.DIM) + text.append(action, style=cls.DIM) text.append_text(cls._status_mark(status)) return text diff --git a/strix/llm/llm.py b/strix/llm/llm.py index 4f624956..3840ce86 100644 --- a/strix/llm/llm.py +++ b/strix/llm/llm.py @@ -247,6 +247,7 @@ class LLM: "messages": messages, "timeout": self.config.timeout, "stream_options": {"include_usage": True}, + "metadata": {"litellm_session_id": self.agent_id}, } if self.config.api_key: diff --git a/strix/tools/browser/browser_actions.py b/strix/tools/browser/browser_actions.py index b51894b6..6b6a09e2 100644 --- a/strix/tools/browser/browser_actions.py +++ b/strix/tools/browser/browser_actions.py @@ -2,6 +2,7 @@ import asyncio import json import logging import re +from functools import partial from typing import Any, Literal from browser_use import Agent @@ -26,7 +27,6 @@ BrowserUseLocalAction = Literal[ "launch", "run", "close_browser", - "search", "navigate", "go_back", "wait", @@ -46,10 +46,6 @@ BrowserUseLocalAction = Literal[ "evaluate", "switch", "close_tab", - "write_file", - "read_file", - "replace_file", - "read_long_content", "done", ] @@ -69,7 +65,7 @@ def _is_ws_error(exc: BaseException) -> bool: return any(kw in msg for kw in _WS_ERRORS) -def _build_llm() -> Any: +def _build_llm(metadata: dict[str, Any] | None = None) -> Any: from strix.config.config import resolve_llm_config from .litellm.chat import ChatLiteLLM @@ -82,6 +78,7 @@ def _build_llm() -> Any: model=model, api_key=api_key, api_base=api_base, + metadata=metadata, ) @@ -135,8 +132,9 @@ async def _run_browser_agent( session: BrowserSession, task: str, return_fields: list[str] | None, + metadata: dict[str, Any] | None = None, ) -> dict[str, Any]: - llm = _build_llm() + llm = _build_llm(metadata=metadata) agent: Any = Agent( task=task, llm=llm, @@ -148,35 +146,26 @@ async def _run_browser_agent( async def log_step(step: Any) -> None: logger.info("Agent step completed: %s", step) - try: - result = await agent.run(on_step_end=log_step) - - if hasattr(result, "is_successful") and not result.is_successful(): - final_result = ( - result.final_result() if callable(result.final_result) else result.final_result - ) - return {"error": final_result or "Agent failed", "is_running": False} + result = await agent.run(on_step_end=log_step) + if hasattr(result, "is_successful") and not result.is_successful(): final_result = ( - result.final_result() - if hasattr(result, "final_result") and callable(result.final_result) - else getattr(result, "final_result", str(result)) + result.final_result() if callable(result.final_result) else result.final_result ) - out = {"message": "Task completed", "result": final_result, "is_running": False} + return {"error": final_result or "Agent failed", "is_running": False} - if return_fields: - fields = {f: getattr(result, f, None) for f in return_fields} - out["fields"] = fields + final_result = ( + result.final_result() + if hasattr(result, "final_result") and callable(result.final_result) + else getattr(result, "final_result", str(result)) + ) + out = {"message": "Task completed", "result": final_result, "is_running": False} - return out - finally: - for name in ("close", "stop"): - if fn := getattr(agent, name, None): - try: - if asyncio.iscoroutine(coro := fn()): - await asyncio.wait_for(coro, timeout=5) - except Exception: # noqa: BLE001,S110 - pass + if return_fields: + fields = {f: getattr(result, f, None) for f in return_fields} + out["fields"] = fields + + return out def _fix_json_in_xml(kws: dict[str, Any]) -> dict[str, Any]: @@ -198,11 +187,12 @@ async def _run_browser_tool( session: BrowserSession, action: str, params: dict[str, Any], + metadata: dict[str, Any] | None = None, ) -> Any: if not session.browser.is_cdp_connected: await session.browser.start() - llm = _build_llm() + llm = _build_llm(metadata=metadata) if session.local: from pathlib import Path @@ -234,6 +224,43 @@ async def _run_browser_tool( ) +async def populate_response(session: BrowserSession, response: dict[str, Any]) -> dict[str, Any]: + try: + screenshot = await session.browser.take_screenshot() + except Exception as e: # noqa: BLE001 + screenshot = None + response["screenshot_error"] = str(e) + + try: + title = await session.browser.get_current_page_title() + url = await session.browser.get_current_page_url() + all_tabs = await session.browser.get_tabs() + except Exception as e: # noqa: BLE001 + # TODO: Consider a fallback way to retrieve the url? + url = f"URL retrieval failed: {e}" + title = "Title retrieval failed with same error" + all_tabs = [] + + try: + # this can fail if the browser disconnects during the tool completion + vp = getattr(session.browser.browser_profile, "viewport", None) + viewport = { + "width": vp.width if vp else None, + "height": vp.height if vp else None, + } + except Exception as e: # noqa: BLE001 + viewport = {"error": f"Viewport retrieval failed: {e}"} + + return { + **response, + "screenshot": screenshot, + "url": url, + "title": title, + "viewport": viewport, + "tabs": all_tabs, + } + + @register_tool(sandbox_execution=False) async def browser_actions( action: BrowserUseLocalAction, @@ -249,6 +276,8 @@ async def browser_actions( agent_id = get_current_agent_id() + metadata = {"litellm_session_id": agent_id} + if action == "launch": has_sandbox = agent_state and getattr(agent_state, "sandbox_info", None) if not has_sandbox: @@ -285,18 +314,19 @@ async def browser_actions( if not task: return {"error": "task required for run action", "is_running": False} - return await _execute_task( - session, - lambda: _run_browser_agent(session, task, return_fields), - task, - ) + runner = partial(_run_browser_agent, session, task, return_fields, metadata) + desc = task + else: + params = _fix_json_in_xml(kwargs) + runner = partial(_run_browser_tool, session, action, params, metadata) + desc = f"{action}({list(kwargs.keys())[:3]})" - params = _fix_json_in_xml(kwargs) - return await _execute_task( - session, - lambda: _run_browser_tool(session, action, params), - f"{action}({list(params.keys())[:3]})", - ) + task_output = await _execute_task(session, runner, desc) + + if "error" in task_output: + return task_output + + return await populate_response(session, task_output) except Exception as error: logger.exception("browser_actions error: %s", action) diff --git a/strix/tools/browser/browser_actions_schema.xml b/strix/tools/browser/browser_actions_schema.xml index 19e94d7f..1a7fc650 100644 --- a/strix/tools/browser/browser_actions_schema.xml +++ b/strix/tools/browser/browser_actions_schema.xml @@ -25,7 +25,6 @@ Optionally pass 'return_fields' to select which history data to include in the response. **Navigation & Browser Control:** - - search: Search queries (DuckDuckGo, Google, Bing). Requires 'query', optional 'engine'. - navigate: Navigate to URLs. Requires 'url', optional 'new_tab'. - go_back: Go back in browser history. - wait: Wait for specified seconds. Optional 'seconds'. @@ -58,12 +57,6 @@ - switch: Switch between tabs. Requires 'tab_id'. - close_tab: Close browser tab. Requires 'tab_id'. - **File Operations:** [Requires to be running in non-sandbox environment] - - write_file: Write content to files. Requires 'file_name' and 'content', optional 'append', 'trailing_newline', 'leading_newline'. - - read_file: Read file contents. Requires 'file_name'. - - replace_file: Replace text in files. Requires 'file_name', 'old_str', 'new_str'. - - read_long_content: Read long content intelligently. Requires 'goal', optional 'source', 'context'. - **Task Completion:** - done: Complete the task. Requires 'text', optional 'success', 'files_to_display'. @@ -76,10 +69,7 @@ directory to use, e.g. "Default", "Profile 1". If not specified, auto-selected. - Search query. Required for 'search' and 'extract' actions. - - - Search engine for 'search' action. Options: "google", "duckduckgo", "bing". + Query string. Required for 'extract' action. URL to navigate to. Required for 'navigate' action. @@ -168,33 +158,6 @@ Tab ID for 'switch' and 'close_tab' actions. Last 4 chars of target_id from browser state. - - File content for 'write_file' action. - - - Append to file instead of overwriting for 'write_file' action. Default: false. - - - Add trailing newline for 'write_file' action. - - - Add leading newline for 'write_file' action. - - - String to replace for 'replace_file' action. - - - Replacement string for 'replace_file' action. - - - Information goal for 'read_long_content' action. - - - Content source for 'read_long_content' action. Options: "page", file path. - - - Additional context for 'read_long_content' action. - Task success status for 'done' action. @@ -259,12 +222,6 @@ Go to https://example.com/login, fill in username "admin" and password "secret", then click the login button - # Search with DuckDuckGo - - search - Python web scraping - - # Navigate to a URL navigate @@ -387,27 +344,6 @@ a3f2 - # Write content to file - - write_file - results.txt - Extracted data here - - - # Read a file - - read_file - data.json - - - # Replace text in file - - replace_file - config.txt - old_value - new_value - - # Complete the task done diff --git a/strix/tools/browser/browser_manager.py b/strix/tools/browser/browser_manager.py index 3983ccc2..04cf2da1 100644 --- a/strix/tools/browser/browser_manager.py +++ b/strix/tools/browser/browser_manager.py @@ -5,6 +5,8 @@ import logging import time from typing import Any +from browser_use import Browser + logger = logging.getLogger(__name__) @@ -79,7 +81,7 @@ class BrowserSession: def __init__( self, - browser: Any, + browser: Browser, cdp_url: str, ws_url: str, *, diff --git a/strix/tools/browser/litellm/chat.py b/strix/tools/browser/litellm/chat.py index 03f4e6a9..31a27594 100644 --- a/strix/tools/browser/litellm/chat.py +++ b/strix/tools/browser/litellm/chat.py @@ -43,17 +43,10 @@ class ChatLiteLLM(BaseChatModel): temperature: float | None = 0.0 max_tokens: int | None = 4096 max_retries: int = 3 + metadata: dict[str, Any] | None = None - _provider_name: str = field( - default="", - init=False, - repr=False, - ) - _clean_model: str = field( - default="", - init=False, - repr=False, - ) + _provider_name: str = field(default="", init=False, repr=False) + _clean_model: str = field(default="", init=False, repr=False) def __post_init__(self) -> None: """Resolve provider info from the model string via litellm.""" @@ -163,6 +156,8 @@ class ChatLiteLLM(BaseChatModel): params["api_key"] = self.api_key if self.api_base: params["api_base"] = self.api_base + if self.metadata: + params["metadata"] = self.metadata if output_format is not None: schema = SchemaOptimizer.create_optimized_json_schema(output_format)