From 5989ddbe28092f3fbbeadfaabdfc6b2b819f1afd Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 26 Mar 2026 12:49:12 +0000 Subject: [PATCH 1/9] refactor: unify retry logic and consolidate pricing registry Task 1: Refactor ChatModel.__call__() and AnthropicChatModel.__call__() to use shared retry helpers from llm_utils.py instead of inline loops. Task 2: Create unified get_pricing(model_name, provider) dispatcher in tracking.py. Rename provider-specific functions to private helpers with public aliases for backward compatibility. Task 3: Remove deprecated ChatModelArgs class (no external consumers). Co-Authored-By: Lukas Burger --- src/agentlab/llm/chat_api.py | 143 +++++++++++------------------------ src/agentlab/llm/tracking.py | 81 +++++++++++++++----- 2 files changed, 109 insertions(+), 115 deletions(-) diff --git a/src/agentlab/llm/chat_api.py b/src/agentlab/llm/chat_api.py index f497894eb..b01fa0195 100644 --- a/src/agentlab/llm/chat_api.py +++ b/src/agentlab/llm/chat_api.py @@ -14,7 +14,12 @@ import agentlab.llm.tracking as tracking from agentlab.llm.base_api import AbstractChatModel, BaseModelArgs from agentlab.llm.huggingface_utils import HFBaseChatModel -from agentlab.llm.llm_utils import AIMessage, Discussion +from agentlab.llm.llm_utils import ( + AIMessage, + Discussion, + call_anthropic_api_with_retries, + call_openai_api_with_retries, +) def make_system_message(content: str) -> dict: @@ -160,39 +165,6 @@ def make_model(self): raise ValueError(f"Backend {self.backend} is not supported") -@dataclass -class ChatModelArgs(BaseModelArgs): - """Object added for backward compatibility with the old ChatModelArgs.""" - - model_path: str = None - model_url: str = None - model_size: str = None - training_total_tokens: int = None - hf_hosted: bool = False - is_model_operational: str = False - sliding_window: bool = False - n_retry_server: int = 4 - infer_tokens_length: bool = False - vision_support: bool = False - shard_support: bool = True - extra_tgi_args: dict = None - tgi_image: str = None - info: dict = None - - def __post_init__(self): - import warnings - - warnings.simplefilter("always", DeprecationWarning) - warnings.warn( - "ChatModelArgs is deprecated and used only for xray. Use one of the specific model args classes instead.", - DeprecationWarning, - ) - warnings.simplefilter("default", DeprecationWarning) - - def make_model(self): - pass - - def _extract_wait_time(error_message, min_retry_wait_time=60): """Extract the wait time from an OpenAI RateLimitError message.""" match = re.search(r"try again in (\d+(\.\d+)?)s", error_message) @@ -277,42 +249,21 @@ def __init__( ) def __call__(self, messages: list[dict], n_samples: int = 1, temperature: float = None) -> dict: - # Initialize retry tracking attributes - self.retries = 0 - self.success = False - self.error_types = [] - - completion = None - e = None - for itr in range(self.max_retry): - self.retries += 1 - temperature = temperature if temperature is not None else self.temperature - try: - completion = self.client.chat.completions.create( - model=self.model_name, - messages=messages, - n=n_samples, - temperature=temperature, - max_completion_tokens=self.max_tokens, - logprobs=self.log_probs, - ) + temperature = temperature if temperature is not None else self.temperature + api_params = { + "model": self.model_name, + "messages": messages, + "n": n_samples, + "temperature": temperature, + "max_completion_tokens": self.max_tokens, + "logprobs": self.log_probs, + } - if completion.usage is None: - raise OpenRouterError( - "The completion object does not contain usage information. This is likely a bug in the OpenRouter API." - ) - - self.success = True - break - except openai.OpenAIError as e: - error_type = handle_error(e, itr, self.min_retry_wait_time, self.max_retry) - self.error_types.append(error_type) - - if not completion: - raise RetryError( - f"Failed to get a response from the API after {self.max_retry} retries\n" - f"Last error: {error_type}" - ) + completion = call_openai_api_with_retries( + self.client.chat.completions.create, + api_params, + max_retries=self.max_retry, + ) input_tokens = completion.usage.prompt_tokens output_tokens = completion.usage.completion_tokens @@ -516,35 +467,31 @@ def __call__(self, messages: list[dict], n_samples: int = 1, temperature: float temperature = temperature if temperature is not None else self.temperature - for attempt in range(self.max_retry): - try: - kwargs = { - "model": self.model_name, - "messages": anthropic_messages, - "max_tokens": self.max_tokens, - "temperature": temperature, - } - - if system_message: - kwargs["system"] = system_message - - response = self.client.messages.create(**kwargs) - - # Track usage if available - if hasattr(tracking.TRACKER, "instance"): - tracking.TRACKER.instance( - response.usage.input_tokens, - response.usage.output_tokens, - 0, # cost calculation would need pricing info - ) - - return AIMessage(response.content[0].text) - - except Exception as e: - if attempt == self.max_retry - 1: - raise e - logging.warning(f"Anthropic API error (attempt {attempt + 1}): {e}") - time.sleep(60) # Simple retry delay + api_params = { + "model": self.model_name, + "messages": anthropic_messages, + "max_tokens": self.max_tokens, + "temperature": temperature, + } + + if system_message: + api_params["system"] = system_message + + response = call_anthropic_api_with_retries( + self.client.messages.create, + api_params, + max_retries=self.max_retry, + ) + + # Track usage if available + if hasattr(tracking.TRACKER, "instance"): + tracking.TRACKER.instance( + response.usage.input_tokens, + response.usage.output_tokens, + 0, # cost calculation would need pricing info + ) + + return AIMessage(response.content[0].text) @dataclass diff --git a/src/agentlab/llm/tracking.py b/src/agentlab/llm/tracking.py index 7b4b64f65..05c69e822 100644 --- a/src/agentlab/llm/tracking.py +++ b/src/agentlab/llm/tracking.py @@ -88,8 +88,66 @@ def wrapper(self, obs): return wrapper -@cache +def get_pricing(model_name: str, provider: str) -> dict: + """Unified pricing lookup that dispatches to the appropriate provider-specific function. + + Args: + model_name: The model name to look up pricing for. + provider: One of 'openai', 'anthropic', 'openrouter', 'litellm'. + + Returns: + A dict mapping model names to {"prompt": float, "completion": float} pricing info. + """ + pricing_fn_map = { + "openai": _get_pricing_openai, + "anthropic": _get_pricing_anthropic, + "openrouter": _get_pricing_openrouter, + "litellm": partial(_get_pricing_litellm, model_name=model_name), + } + pricing_fn = pricing_fn_map.get(provider) + if pricing_fn is None: + logging.warning( + f"Unsupported provider: {provider}. Supported providers are: {list(pricing_fn_map.keys())}" + ) + return {} + return pricing_fn() + + +# Keep the public aliases for backward compatibility def get_pricing_openrouter(): + """Returns a dictionary of model pricing for OpenRouter models. + + .. deprecated:: Use get_pricing(model_name, 'openrouter') instead. + """ + return _get_pricing_openrouter() + + +def get_pricing_openai(): + """Returns a dictionary of model pricing for OpenAI models. + + .. deprecated:: Use get_pricing(model_name, 'openai') instead. + """ + return _get_pricing_openai() + + +def get_pricing_anthropic(): + """Returns a dictionary of model pricing for Anthropic models. + + .. deprecated:: Use get_pricing(model_name, 'anthropic') instead. + """ + return _get_pricing_anthropic() + + +def get_pricing_litellm(model_name): + """Returns a dictionary of model pricing for a LiteLLM model. + + .. deprecated:: Use get_pricing(model_name, 'litellm') instead. + """ + return _get_pricing_litellm(model_name) + + +@cache +def _get_pricing_openrouter(): """Returns a dictionary of model pricing for OpenRouter models.""" api_key = os.getenv("OPENROUTER_API_KEY") assert api_key, "OpenRouter API key is required" @@ -108,7 +166,7 @@ def get_pricing_openrouter(): } -def get_pricing_openai(): +def _get_pricing_openai(): """Returns a dictionary of model pricing for OpenAI models.""" try: cost_dict = openai_info.MODEL_COST_PER_1K_TOKENS @@ -138,7 +196,7 @@ def _remove_version_suffix(model_name): return re.sub(r"anthropic.", "", no_version) -def get_pricing_anthropic(): +def _get_pricing_anthropic(): """Returns a dictionary of model pricing for Anthropic models.""" try: input_cost_dict = bedrock_anthropic_callback.MODEL_COST_PER_1K_INPUT_TOKENS @@ -163,7 +221,7 @@ def get_pricing_anthropic(): return res -def get_pricing_litellm(model_name): +def _get_pricing_litellm(model_name): """Returns a dictionary of model pricing for a LiteLLM model.""" try: info = get_model_info(model_name) @@ -217,19 +275,8 @@ def fetch_pricing_information_from_provider(self) -> Optional[dict]: Returns: Optional[dict]: A dict mapping model names to pricing info, or None if not found. """ - pricing_fn_map = { - "openai": get_pricing_openai, - "anthropic": get_pricing_anthropic, - "openrouter": get_pricing_openrouter, - "litellm": partial(get_pricing_litellm, self.model_name), - } - pricing_fn = pricing_fn_map.get(self._pricing_api, None) - if pricing_fn is None: - logging.warning( - f"Unsupported provider: {self._pricing_api}. Supported providers are: {list(pricing_fn_map.keys())}" - ) - return None - return pricing_fn() + result = get_pricing(self.model_name, self._pricing_api) + return result if result else None def set_pricing_attributes(self) -> None: """Set the pricing attributes for the model based on the provider.""" From c6345915afa22c08bd68ee2d4137d546a4c28fd4 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 26 Mar 2026 12:49:19 +0000 Subject: [PATCH 2/9] refactor: split dynamic_prompting.py into prompting/ subpackage Task 4: Create src/agentlab/agents/prompting/ package with: - flags.py: Flags, ObsFlags, ActionFlags - base.py: PromptElement, Shrinkable, Trunkater, fit_tokens - elements.py: all concrete prompt element classes - __init__.py: re-exports for backward compatibility Update dynamic_prompting.py to be a thin re-export shim. Co-Authored-By: Lukas Burger --- src/agentlab/agents/dynamic_prompting.py | 907 +--------------------- src/agentlab/agents/prompting/__init__.py | 56 ++ src/agentlab/agents/prompting/base.py | 181 +++++ src/agentlab/agents/prompting/elements.py | 514 ++++++++++++ src/agentlab/agents/prompting/flags.py | 86 ++ 5 files changed, 869 insertions(+), 875 deletions(-) create mode 100644 src/agentlab/agents/prompting/__init__.py create mode 100644 src/agentlab/agents/prompting/base.py create mode 100644 src/agentlab/agents/prompting/elements.py create mode 100644 src/agentlab/agents/prompting/flags.py diff --git a/src/agentlab/agents/dynamic_prompting.py b/src/agentlab/agents/dynamic_prompting.py index 92ad25b99..dbd105e49 100644 --- a/src/agentlab/agents/dynamic_prompting.py +++ b/src/agentlab/agents/dynamic_prompting.py @@ -1,876 +1,33 @@ -import abc -import logging -import platform -import time -from copy import copy, deepcopy -from dataclasses import asdict, dataclass -from textwrap import dedent -from typing import Literal -from warnings import warn - -import bgym -from bgym import HighLevelActionSetArgs -from browsergym.core.action.base import AbstractActionSet -from browsergym.utils.obs import flatten_axtree_to_str, flatten_dom_to_str, overlay_som, prune_html - -from agentlab.llm.llm_utils import ( - BaseMessage, - ParseError, - count_tokens, - extract_code_blocks, - image_to_jpg_base64_url, - parse_html_tags_raise, +"""Dynamic prompting module for AgentLab agents. + +This module re-exports all symbols from the prompting subpackage for backward compatibility. +New code should import directly from agentlab.agents.prompting instead. +""" + +# Re-export everything from the prompting subpackage for backward compatibility +from agentlab.agents.prompting import ( # noqa: F401 + HTML, + AXTree, + ActionFlags, + ActionPrompt, + BeCautious, + ChatInstructions, + Error, + Flags, + FocusedElement, + GoalInstructions, + Hints, + History, + HistoryStep, + MacNote, + Observation, + ObsFlags, + PromptElement, + Shrinkable, + SystemPrompt, + Tabs, + Think, + Trunkater, + fit_tokens, + make_obs_preprocessor, ) - - -class Flags: - """Base class for flags. Mostly for backward compatibility.""" - - def copy(self): - return deepcopy(self) - - def asdict(self): - """Helper for JSON serializable requirement.""" - return asdict(self) - - @classmethod - def from_dict(self, flags_dict): - """Helper for JSON serializable requirement.""" - if isinstance(flags_dict, ObsFlags): - return flags_dict - - if not isinstance(flags_dict, dict): - raise ValueError(f"Unregcognized type for flags_dict of type {type(flags_dict)}.") - return ObsFlags(**flags_dict) - - -@dataclass -class ObsFlags(Flags): - """ - A class to represent various flags used to control features in an application. - - Attributes: - use_html (bool): Use the HTML in the prompt. - use_ax_tree (bool): Use the accessibility tree in the prompt. - use_focused_element (bool): Provide the ID of the focused element. - use_error_logs (bool): Expose the previous error in the prompt. - use_history (bool): Enable history of previous steps in the prompt. - use_past_error_logs (bool): If use_history is True, expose all previous errors in the history. - use_action_history (bool): If use_history is True, include the actions in the history. - use_think_history (bool): If use_history is True, include all previous chains of thoughts in the history. - use_diff (bool): Add a diff of the current and previous HTML to the prompt. - html_type (str): Type of HTML to use in the prompt, may depend on preprocessing of observation. - use_screenshot (bool): Add a screenshot of the page to the prompt, following OpenAI's API. This will be automatically disabled if the model does not have vision capabilities. - use_som (bool): Add a set of marks to the screenshot. - extract_visible_tag (bool): Add a "visible" tag to visible elements in the AXTree. - extract_clickable_tag (bool): Add a "clickable" tag to clickable elements in the AXTree. - extract_coords (Literal['False', 'center', 'box']): Add the coordinates of the elements. - filter_visible_elements_only (bool): Only show visible elements in the AXTree. - """ - - use_html: bool = True - use_ax_tree: bool = False - use_tabs: bool = False - use_focused_element: bool = False - use_error_logs: bool = False - use_history: bool = False - use_past_error_logs: bool = False - use_action_history: bool = False - use_think_history: bool = False - use_diff: bool = False # - html_type: str = "pruned_html" - use_screenshot: bool = True - use_som: bool = False - extract_visible_tag: bool = False - extract_clickable_tag: bool = False - extract_coords: Literal["False", "center", "box"] = "False" - filter_visible_elements_only: bool = False - # low sets the token count of each image to 65 (85?) - # high sets the token count of each image to 2*65 (2*85?) times the amount of 512x512px patches - # auto chooses between low and high based on image size (openai default) - openai_vision_detail: Literal["low", "high", "auto"] = "auto" - filter_with_bid_only: bool = False - filter_som_only: bool = False - - -@dataclass -class ActionFlags(Flags): - action_set: HighLevelActionSetArgs = None # should be set by the set_benchmark method - long_description: bool = True - individual_examples: bool = False - - # for backward compatibility - multi_actions: bool = None - is_strict: bool = None - - -class PromptElement: - """Base class for all prompt elements. Prompt elements can be hidden.""" - - _prompt = "" - _abstract_ex = "" - _concrete_ex = "" - - def __init__(self, visible: bool = True) -> None: - """Prompt element that can be hidden. - - Args: - visible : bool, optional - Whether the prompt element should be visible, by default True. Can - be a callable that returns a bool. This is useful when a specific - flag changes during a shrink iteration. - """ - self._visible = visible - - @property - def prompt(self) -> str | BaseMessage: - """Avoid overriding this method. Override _prompt instead.""" - if self.is_visible: - return self._prompt - else: - return "" - - @property - def abstract_ex(self): - """Useful when this prompt element is requesting an answer from the llm. - Provide an abstract example of the answer here. See Memory for an - example. - - Avoid overriding this method. Override _abstract_ex instead - - Returns: - str: The abstract example - """ - if self.is_visible: - return self._abstract_ex - else: - return "" - - @property - def concrete_ex(self): - """Useful when this prompt element is requesting an answer from the llm. - Provide a concrete example of the answer here. See Memory for an - example. - - Avoid overriding this method. Override _concrete_ex instead - - Returns: - str: The concrete example - """ - if self.is_visible: - return self._concrete_ex - else: - return "" - - @property - def is_visible(self): - """Handle the case where visible is a callable.""" - visible = self._visible - if callable(visible): - visible = visible() - return visible - - def _parse_answer(self, text_answer): - """Override to actually extract elements from the answer.""" - return {} - - def parse_answer(self, text_answer) -> dict: - if self.is_visible: - return self._parse_answer(text_answer) - else: - return {} - - -class Shrinkable(PromptElement, abc.ABC): - @abc.abstractmethod - def shrink(self) -> None: - """Implement shrinking of this prompt element. - - You need to recursively call all shrinkable elements that are part of - this prompt. You can also implement a shriking startegy for this prompt. - Shrinking is can be called multiple times to progressively shrink the - prompt until it fits max_tokens. Default max shrink iterations is 20. - """ - pass - - -class Trunkater(Shrinkable): - """Shrinkable element that truncates the prompt element from the bottom - after a certain number of iterations.""" - - def __init__(self, visible, shrink_speed=0.3, start_trunkate_iteration=10): - super().__init__(visible=visible) - self.shrink_speed = shrink_speed - self.start_trunkate_iteration = start_trunkate_iteration - self.shrink_calls = 0 - self.deleted_lines = 0 - - def shrink(self) -> None: - if self.is_visible and self.shrink_calls >= self.start_trunkate_iteration: - # remove the fraction of _prompt - lines = self._prompt.splitlines() - new_line_count = int(len(lines) * (1 - self.shrink_speed)) - self.deleted_lines += len(lines) - new_line_count - self._prompt = "\n".join(lines[:new_line_count]) - self._prompt += f"\n... Deleted {self.deleted_lines} lines to reduce prompt size." - - self.shrink_calls += 1 - - -def fit_tokens( - shrinkable: Shrinkable, - max_prompt_tokens=None, - max_iterations=20, - model_name="openai/gpt-4", - additional_prompts=[""], -): - """Shrink a prompt element until it fits `max_prompt_tokens`. - - Args: - shrinkable (Shrinkable): The prompt element to shrink. - max_prompt_tokens (int): The maximum number of tokens allowed. - max_iterations (int, optional): The maximum number of shrink iterations, by default 20. - model_name (str, optional): The name of the model used when tokenizing. - additional_prompts (str or List[str], optional): Additional prompts to account for when shrinking, by default [""]. - - Returns: - str: the prompt after shrinking. - - Raises: - ValueError: Unrecognized type for prompt - """ - - if max_prompt_tokens is None: - return shrinkable.prompt - - if isinstance(additional_prompts, str): - additional_prompts = [additional_prompts] - - for prompt in additional_prompts: - max_prompt_tokens -= count_tokens(prompt, model=model_name) + 1 # +1 because why not ? - - for _ in range(max_iterations): - prompt = shrinkable.prompt - if isinstance(prompt, str): - prompt_str = prompt - elif isinstance(prompt, list): - # warn deprecated - warn( - "Using list of prompts is deprecated. Use a Discussion object instead.", - DeprecationWarning, - ) - prompt_str = "\n".join([p["text"] for p in prompt if p["type"] == "text"]) - elif isinstance(prompt, BaseMessage): - prompt_str = prompt.__str__(warn_if_image=False) - else: - raise ValueError(f"Unrecognized type for prompt: {type(prompt)}") - n_token = count_tokens(prompt_str, model=model_name) - if n_token <= max_prompt_tokens: - return prompt - shrinkable.shrink() - - logging.info( - dedent( - f"""\ - After {max_iterations} shrink iterations, the prompt is still - {count_tokens(prompt_str)} tokens (greater than {max_prompt_tokens}). Returning the prompt as is.""" - ) - ) - return prompt - - -class HTML(Trunkater): - def __init__(self, html, visible_elements_only: bool, visible: bool = True, prefix="") -> None: - super().__init__(visible=visible, start_trunkate_iteration=5) - if visible_elements_only: - visible_elements_note = """\ -Note: only elements that are visible in the viewport are presented. You might need to scroll the page, or open tabs or menus to see more. - -""" - else: - visible_elements_note = "" - self._prompt = f"\n{prefix}HTML:\n{visible_elements_note}{html}\n" - - -class AXTree(Trunkater): - def __init__( - self, - ax_tree, - visible_elements_only: bool, - visible: bool = True, - coord_type=None, - visible_tag=True, - prefix="", - ) -> None: - super().__init__(visible=visible, start_trunkate_iteration=10) - bid_info = """\ -Note: [bid] is the unique alpha-numeric identifier at the beginning of lines for each element in the AXTree. Always use bid to refer to elements in your actions. - -""" - if coord_type == "center": - coord_note = """\ -Note: center coordinates are provided in parenthesis and are relative to the top left corner of the page. - -""" - elif coord_type == "box": - coord_note = """\ -Note: bounding box of each object are provided in parenthesis and are relative to the top left corner of the page. - -""" - else: - coord_note = "" - if visible_elements_only: - visible_elements_note = """\ -Note: only elements that are visible in the viewport are presented. You might need to scroll the page, or open tabs or menus to see more. - -""" - else: - visible_elements_note = "" - - if visible_tag: - vsible_tag_note = """\ -Note: You can only interact with visible elements. If the "visible" tag is not -present, the element is not visible on the page. - -""" - else: - vsible_tag_note = "" - self._prompt = f"\n{prefix}AXTree:\n{bid_info}{coord_note}{visible_elements_note}{vsible_tag_note}{ax_tree}\n" - - -class Error(PromptElement): - def __init__(self, error: str, visible: bool = True, prefix="", limit_logs=True) -> None: - logs_separator = "Call log:" - if limit_logs and logs_separator in error: - error, logs = error.split(logs_separator) - logs = "\n".join(logs.split("\n")[:10]) - error = error + f"\n{logs_separator}\n{logs}" - - super().__init__(visible=visible) - self._prompt = f"\n{prefix}Error from previous action:\n{error}\n" - - -class FocusedElement(PromptElement): - def __init__(self, bid, visible: bool = True, prefix="") -> None: - super().__init__(visible=visible) - self._prompt = f""" -{prefix}Focused element: -""" - if bid: - self._prompt += f"""\ -bid={repr(bid)} -""" - else: - self._prompt += f"""\ -None -""" - - -class Tabs(PromptElement): - def __init__(self, obs, visible: bool = True, prefix="") -> None: - super().__init__(visible=visible) - self.obs = obs - self.prefix = prefix - - @property - def _prompt(self) -> str: - # by implementing this as a property, it's only coputed if visible - prompt_pieces = [f"\n{self.prefix}Currently open tabs:"] - for page_index, (page_url, page_title) in enumerate( - zip(self.obs["open_pages_urls"], self.obs["open_pages_titles"]) - ): - active_or_not = " (active tab)" if page_index == self.obs["active_page_index"] else "" - prompt_piece = f"""\ -Tab {page_index}{active_or_not}: - Title: {page_title} - URL: {page_url} -""" - prompt_pieces.append(prompt_piece) - return "\n".join(prompt_pieces) - - -class Observation(Shrinkable): - """Observation of the current step. - - Contains the html, the accessibility tree and the error logs. - """ - - def __init__(self, obs, flags: ObsFlags) -> None: - super().__init__() - self.flags = flags - self.obs = obs - - self.tabs = Tabs( - obs, - visible=lambda: flags.use_tabs, - prefix="## ", - ) - - self.html = HTML( - obs[flags.html_type], - visible_elements_only=flags.filter_visible_elements_only, - visible=lambda: flags.use_html, - prefix="## ", - ) - self.ax_tree = AXTree( - obs["axtree_txt"], - visible_elements_only=flags.filter_visible_elements_only, - visible=lambda: flags.use_ax_tree, - coord_type=flags.extract_coords, - visible_tag=flags.extract_visible_tag, - prefix="## ", - ) - self.error = Error( - obs["last_action_error"], - visible=lambda: flags.use_error_logs and obs["last_action_error"], - prefix="## ", - ) - self.focused_element = FocusedElement( - obs["focused_element_bid"], - visible=flags.use_focused_element, - prefix="## ", - ) - - def shrink(self): - self.ax_tree.shrink() - self.html.shrink() - - @property - def _prompt(self) -> str: - return f""" -# Observation of current step: -{self.tabs.prompt}{self.html.prompt}{self.ax_tree.prompt}{self.focused_element.prompt}{self.error.prompt} - -""" - - def add_screenshot(self, prompt: BaseMessage) -> BaseMessage: - if self.flags.use_screenshot: - if self.flags.use_som: - screenshot = self.obs["screenshot_som"] - prompt.add_text( - "\n## Screenshot:\nHere is a screenshot of the page, it is annotated with bounding boxes and corresponding bids:" - ) - else: - screenshot = self.obs["screenshot"] - prompt.add_text("\n## Screenshot:\nHere is a screenshot of the page:") - img_url = image_to_jpg_base64_url(screenshot) - prompt.add_image(img_url, detail=self.flags.openai_vision_detail) - return prompt - - -class MacNote(PromptElement): - def __init__(self) -> None: - super().__init__(visible=platform.system() == "Darwin") - self._prompt = ( - "\nNote: you are on mac so you should use Meta instead of Control for Control+C etc.\n" - ) - - -class BeCautious(PromptElement): - def __init__(self, visible: bool = True) -> None: - super().__init__(visible=visible) - self._prompt = f"""\ -\nBe very cautious. Avoid submitting anything before verifying the effect of your -actions. Take the time to explore the effect of safe actions first. For example -you can fill a few elements of a form, but don't click submit before verifying -that everything was filled correctly.\n""" - - -class GoalInstructions(PromptElement): - def __init__(self, goal_object, visible: bool = True, extra_instructions=None) -> None: - super().__init__(visible) - self._prompt = [ - dict( - type="text", - text=f"""\ -# Instructions -Review the current state of the page and all other information to find the best -possible next action to accomplish your goal. Your answer will be interpreted -and executed by a program, make sure to follow the formatting instructions. - -## Goal: -""", - ) - ] - - self._prompt += goal_object - - if extra_instructions: - self._prompt += [ - dict( - type="text", - text=f""" - -## Extra instructions: - -{extra_instructions} -""", - ) - ] - - -class ChatInstructions(PromptElement): - def __init__(self, chat_messages, visible: bool = True, extra_instructions=None) -> None: - super().__init__(visible) - self._prompt = f"""\ -# Instructions - -You are a UI Assistant, your goal is to help the user perform tasks using a web browser. You can -communicate with the user via a chat, in which the user gives you instructions and in which you -can send back messages. You have access to a web browser that both you and the user can see, -and with which only you can interact via specific commands. - -Review the instructions from the user, the current state of the page and all other information -to find the best possible next action to accomplish your goal. Your answer will be interpreted -and executed by a program, make sure to follow the formatting instructions. - -## Chat messages: - -""" - self._prompt += "\n".join( - [ - f"""\ - - [{msg['role']}] UTC Time: {time.asctime(time.gmtime(msg['timestamp']))} - Local Time: {time.asctime(time.localtime(msg['timestamp']))} - {msg['message']}""" - for msg in chat_messages - ] - ) - - if extra_instructions: - self._prompt += f""" - -## Extra instructions: - -{extra_instructions} -""" - - -class Hints(PromptElement): - """Not super useful and stale.""" - - # NOTE: are these hints still relevant? - _prompt = """\ -Note: -* Some tasks may be game like and may require to interact with the mouse position -in x, y coordinates. -* Some text field might have auto completion. To see it, you have to type a few -characters and wait until next step. -* If you have to cut and paste, don't forget to select the text first. -* Coordinate inside an SVG are relative to it's top left corner. -* Make sure to use bid to identify elements when using commands. -* Interacting with combobox, dropdowns and auto-complete fields can be tricky, -sometimes you need to use select_option, while other times you need to use fill -or click and wait for the reaction of the page. -""" - - -class SystemPrompt(PromptElement): - _prompt = """\ -You are an agent trying to solve a web task based on the content of the page and -user instructions. You can interact with the page and explore, and send messages to the user. Each time you -submit an action it will be sent to the browser and you will receive a new page.""" - - -class ActionPrompt(PromptElement): - - _concrete_ex = """ - -click('a324') - -""" - - def __init__(self, action_set: AbstractActionSet, action_flags: ActionFlags) -> None: - super().__init__() - self.action_set = action_set - self.action_flags = action_flags - action_set_generic_info = """\ -Note: This action set allows you to interact with your environment. Most of them -are python function executing playwright code. The primary way of referring to -elements in the page is through bid which are specified in your observations. - -""" - action_description = action_set.describe( - with_long_description=action_flags.long_description, - with_examples=action_flags.individual_examples, - ) - self._prompt = ( - f"# Action space:\n{action_set_generic_info}{action_description}{MacNote().prompt}\n" - ) - self._abstract_ex = f""" - -{self.action_set.example_action(abstract=True)} - -""" - - # self._concrete_ex = f""" - # - # {self.action_set.example_action(abstract=False)} - # - # """ - - def _parse_answer(self, text_answer): - try: - ans_dict = parse_html_tags_raise(text_answer, keys=["action"], merge_multiple=True) - except ParseError as e: - if self.action_flags.is_strict: - raise e - else: - # try to extract code blocks - blocks = extract_code_blocks(text_answer) - if len(blocks) == 0: - raise e - else: - code = "\n".join([block for _, block in blocks]) - ans_dict = {"action": code, "parse_error": str(e)} - - try: - if ans_dict["action"] == "None": - # Used by reproducibility agent for backward compatibility of - # traces missing LLM's response in chat messages. - ans_dict["action"] = None - else: - # just check if action can be mapped to python code but keep action as is - # the environment will be responsible for mapping it to python - self.action_set.to_python_code(ans_dict["action"]) - except Exception as e: - raise ParseError( - f"Error while parsing action\n: {e}\n" - "Make sure your answer is restricted to the allowed actions." - ) - - return ans_dict - - -# def make_action_set(action_flags: ActionFlags) -> AbstractActionSet: - -# if action_flags.action_set == "python": -# action_set = PythonActionSet(strict=action_flags.is_strict) -# if action_flags.demo_mode != "off": -# warn( -# f'Action_set "python" is incompatible with demo_mode={repr(action_flags.demo_mode)}.' -# ) -# return action_set - -# action_set = HighLevelActionSet( -# subsets=list(set(["chat"] + ["infeas"] + action_flags.action_set.split("+"))), -# multiaction=action_flags.multi_actions, -# strict=action_flags.is_strict, -# demo_mode=action_flags.demo_mode, -# ) - -# return action_set - - -class Think(PromptElement): - _prompt = "" - - _abstract_ex = """ - -Think step by step. If you need to make calculations such as coordinates, write them here. Describe the effect -that your previous action had on the current content of the page. - -""" - _concrete_ex = """ - -From previous action I tried to set the value of year to "2022", -using select_option, but it doesn't appear to be in the form. It may be a -dynamic dropdown, I will try using click with the bid "a324" and look at the -response from the page. - -""" - - def _parse_answer(self, text_answer): - try: - return parse_html_tags_raise(text_answer, keys=["think"], merge_multiple=True) - except ParseError as e: - return {"think": text_answer, "parse_error": str(e)} - - -# def diff(previous, new): -# """Return a string showing the difference between original and new. - -# If the difference is above diff_threshold, return the diff string.""" - -# if previous == new: -# return "Identical", [] - -# if len(previous) == 0 or previous is None: -# return "previous is empty", [] - -# diff_gen = difflib.ndiff(previous.splitlines(), new.splitlines()) - -# diff_lines = [] -# plus_count = 0 -# minus_count = 0 -# for line in diff_gen: -# if line.strip().startswith("+"): -# diff_lines.append(line) -# plus_count += 1 -# elif line.strip().startswith("-"): -# diff_lines.append(line) -# minus_count += 1 -# else: -# continue - -# header = f"{plus_count} lines added and {minus_count} lines removed:" - -# return header, diff_lines - - -# class Diff(Shrinkable): -# def __init__( -# self, previous, new, prefix="", max_line_diff=20, shrink_speed=2, visible=True -# ) -> None: -# super().__init__(visible=visible) -# self.previous = previous -# self.new = new -# self.max_line_diff = max_line_diff -# self.shrink_speed = shrink_speed -# self.prefix = prefix - -# def shrink(self): -# self.max_line_diff -= self.shrink_speed -# self.max_line_diff = max(1, self.max_line_diff) - -# @property -# def _prompt(self) -> str: -# header, diff_lines = diff(self.previous, self.new) - -# diff_str = "\n".join(diff_lines[: self.max_line_diff]) -# if len(diff_lines) > self.max_line_diff: -# original_count = len(diff_lines) -# diff_str = f"{diff_str}\nDiff truncated, {original_count - self.max_line_diff} changes now shown." -# return f"{self.prefix}{header}\n{diff_str}\n" - - -class HistoryStep(Shrinkable): - def __init__( - self, previous_obs, current_obs, action, memory, thought, flags: ObsFlags, shrink_speed=1 - ) -> None: - super().__init__() - # self.html_diff = Diff( - # previous_obs[flags.html_type], - # current_obs[flags.html_type], - # prefix="\n### HTML diff:\n", - # shrink_speed=shrink_speed, - # visible=lambda: flags.use_html and flags.use_diff, - # ) - # self.ax_tree_diff = Diff( - # previous_obs["axtree_txt"], - # current_obs["axtree_txt"], - # prefix=f"\n### Accessibility tree diff:\n", - # shrink_speed=shrink_speed, - # visible=lambda: flags.use_ax_tree and flags.use_diff, - # ) - self.error = Error( - current_obs["last_action_error"], - visible=( - lambda: flags.use_error_logs - and current_obs["last_action_error"] - and flags.use_past_error_logs - ), - prefix="### ", - ) - self.shrink_speed = shrink_speed - self.action = action - self.memory = memory - self.thought = thought - self.flags = flags - - def shrink(self): - super().shrink() - # self.html_diff.shrink() - # self.ax_tree_diff.shrink() - - @property - def _prompt(self) -> str: - prompt = "" - - if self.flags.use_think_history: - prompt += f"\n\n{self.thought}\n\n" - - if self.flags.use_action_history: - prompt += f"\n\n{self.action}\n\n" - - # prompt += f"{self.error.prompt}{self.html_diff.prompt}{self.ax_tree_diff.prompt}" - prompt += f"{self.error.prompt}" - - if self.memory is not None: - prompt += f"\n\n{self.memory}\n\n" - - return prompt - - -class History(Shrinkable): - def __init__( - self, history_obs, actions, memories, thoughts, flags: ObsFlags, shrink_speed=1 - ) -> None: - if memories is None: - memories = [None] * len(actions) - super().__init__(visible=lambda: flags.use_history) - assert len(history_obs) == len(actions) + 1 - assert len(history_obs) == len(memories) + 1 - - self.shrink_speed = shrink_speed - self.history_steps: list[HistoryStep] = [] - - for i in range(1, len(history_obs)): - self.history_steps.append( - HistoryStep( - history_obs[i - 1], - history_obs[i], - actions[i - 1], - memories[i - 1], - thoughts[i - 1], - flags, - ) - ) - - def shrink(self): - """Shrink individual steps""" - # TODO set the shrink speed of older steps to be higher - super().shrink() - for step in self.history_steps: - step.shrink() - - @property - def _prompt(self): - prompts = ["# History of interaction with the task:\n"] - for i, step in enumerate(self.history_steps): - prompts.append(f"## step {i}") - prompts.append(step.prompt) - return "\n".join(prompts) + "\n" - - -def make_obs_preprocessor(flags: ObsFlags): - def obs_mapping(obs: dict): - obs = copy(obs) - obs["dom_txt"] = flatten_dom_to_str( - obs["dom_object"], - extra_properties=obs["extra_element_properties"], - with_visible=flags.extract_visible_tag, - with_clickable=flags.extract_clickable_tag, - with_center_coords=flags.extract_coords == "center", - with_bounding_box_coords=flags.extract_coords == "box", - filter_visible_only=flags.filter_visible_elements_only, - filter_with_bid_only=flags.filter_with_bid_only, - filter_som_only=flags.filter_som_only, - ) - obs["axtree_txt"] = flatten_axtree_to_str( - obs["axtree_object"], - extra_properties=obs["extra_element_properties"], - with_visible=flags.extract_visible_tag, - with_clickable=flags.extract_clickable_tag, - with_center_coords=flags.extract_coords == "center", - with_bounding_box_coords=flags.extract_coords == "box", - filter_visible_only=flags.filter_visible_elements_only, - filter_with_bid_only=flags.filter_with_bid_only, - filter_som_only=flags.filter_som_only, - ) - obs["pruned_html"] = prune_html(obs["dom_txt"]) - obs["screenshot_som"] = overlay_som( - obs["screenshot"], extra_properties=obs["extra_element_properties"] - ) - - return obs - - return obs_mapping diff --git a/src/agentlab/agents/prompting/__init__.py b/src/agentlab/agents/prompting/__init__.py new file mode 100644 index 000000000..457db918b --- /dev/null +++ b/src/agentlab/agents/prompting/__init__.py @@ -0,0 +1,56 @@ +"""Prompting subpackage for AgentLab agents. + +This package provides modular prompt construction components split into: +- flags: Configuration flags (Flags, ObsFlags, ActionFlags) +- base: Base prompt element classes (PromptElement, Shrinkable, Trunkater, fit_tokens) +- elements: Concrete prompt element implementations +""" + +from .base import PromptElement, Shrinkable, Trunkater, fit_tokens +from .elements import ( + HTML, + AXTree, + ActionPrompt, + BeCautious, + ChatInstructions, + Error, + FocusedElement, + GoalInstructions, + Hints, + History, + HistoryStep, + MacNote, + Observation, + SystemPrompt, + Tabs, + Think, + make_obs_preprocessor, +) +from .flags import ActionFlags, Flags, ObsFlags + +__all__ = [ + "Flags", + "ObsFlags", + "ActionFlags", + "PromptElement", + "Shrinkable", + "Trunkater", + "fit_tokens", + "HTML", + "AXTree", + "Error", + "FocusedElement", + "Tabs", + "Observation", + "MacNote", + "BeCautious", + "GoalInstructions", + "ChatInstructions", + "Hints", + "SystemPrompt", + "ActionPrompt", + "Think", + "HistoryStep", + "History", + "make_obs_preprocessor", +] diff --git a/src/agentlab/agents/prompting/base.py b/src/agentlab/agents/prompting/base.py new file mode 100644 index 000000000..5b9800e06 --- /dev/null +++ b/src/agentlab/agents/prompting/base.py @@ -0,0 +1,181 @@ +import abc +import logging +from textwrap import dedent +from warnings import warn + +from agentlab.llm.llm_utils import BaseMessage, count_tokens + + +class PromptElement: + """Base class for all prompt elements. Prompt elements can be hidden.""" + + _prompt = "" + _abstract_ex = "" + _concrete_ex = "" + + def __init__(self, visible: bool = True) -> None: + """Prompt element that can be hidden. + + Args: + visible : bool, optional + Whether the prompt element should be visible, by default True. Can + be a callable that returns a bool. This is useful when a specific + flag changes during a shrink iteration. + """ + self._visible = visible + + @property + def prompt(self) -> str | BaseMessage: + """Avoid overriding this method. Override _prompt instead.""" + if self.is_visible: + return self._prompt + else: + return "" + + @property + def abstract_ex(self): + """Useful when this prompt element is requesting an answer from the llm. + Provide an abstract example of the answer here. See Memory for an + example. + + Avoid overriding this method. Override _abstract_ex instead + + Returns: + str: The abstract example + """ + if self.is_visible: + return self._abstract_ex + else: + return "" + + @property + def concrete_ex(self): + """Useful when this prompt element is requesting an answer from the llm. + Provide a concrete example of the answer here. See Memory for an + example. + + Avoid overriding this method. Override _concrete_ex instead + + Returns: + str: The concrete example + """ + if self.is_visible: + return self._concrete_ex + else: + return "" + + @property + def is_visible(self): + """Handle the case where visible is a callable.""" + visible = self._visible + if callable(visible): + visible = visible() + return visible + + def _parse_answer(self, text_answer): + """Override to actually extract elements from the answer.""" + return {} + + def parse_answer(self, text_answer) -> dict: + if self.is_visible: + return self._parse_answer(text_answer) + else: + return {} + + +class Shrinkable(PromptElement, abc.ABC): + @abc.abstractmethod + def shrink(self) -> None: + """Implement shrinking of this prompt element. + + You need to recursively call all shrinkable elements that are part of + this prompt. You can also implement a shriking startegy for this prompt. + Shrinking is can be called multiple times to progressively shrink the + prompt until it fits max_tokens. Default max shrink iterations is 20. + """ + pass + + +class Trunkater(Shrinkable): + """Shrinkable element that truncates the prompt element from the bottom + after a certain number of iterations.""" + + def __init__(self, visible, shrink_speed=0.3, start_trunkate_iteration=10): + super().__init__(visible=visible) + self.shrink_speed = shrink_speed + self.start_trunkate_iteration = start_trunkate_iteration + self.shrink_calls = 0 + self.deleted_lines = 0 + + def shrink(self) -> None: + if self.is_visible and self.shrink_calls >= self.start_trunkate_iteration: + # remove the fraction of _prompt + lines = self._prompt.splitlines() + new_line_count = int(len(lines) * (1 - self.shrink_speed)) + self.deleted_lines += len(lines) - new_line_count + self._prompt = "\n".join(lines[:new_line_count]) + self._prompt += f"\n... Deleted {self.deleted_lines} lines to reduce prompt size." + + self.shrink_calls += 1 + + +def fit_tokens( + shrinkable: Shrinkable, + max_prompt_tokens=None, + max_iterations=20, + model_name="openai/gpt-4", + additional_prompts=[""], +): + """Shrink a prompt element until it fits `max_prompt_tokens`. + + Args: + shrinkable (Shrinkable): The prompt element to shrink. + max_prompt_tokens (int): The maximum number of tokens allowed. + max_iterations (int, optional): The maximum number of shrink iterations, by default 20. + model_name (str, optional): The name of the model used when tokenizing. + additional_prompts (str or List[str], optional): Additional prompts to account for when shrinking, by default [""]. + + Returns: + str: the prompt after shrinking. + + Raises: + ValueError: Unrecognized type for prompt + """ + + if max_prompt_tokens is None: + return shrinkable.prompt + + if isinstance(additional_prompts, str): + additional_prompts = [additional_prompts] + + for prompt in additional_prompts: + max_prompt_tokens -= count_tokens(prompt, model=model_name) + 1 # +1 because why not ? + + for _ in range(max_iterations): + prompt = shrinkable.prompt + if isinstance(prompt, str): + prompt_str = prompt + elif isinstance(prompt, list): + # warn deprecated + warn( + "Using list of prompts is deprecated. Use a Discussion object instead.", + DeprecationWarning, + ) + prompt_str = "\n".join([p["text"] for p in prompt if p["type"] == "text"]) + elif isinstance(prompt, BaseMessage): + prompt_str = prompt.__str__(warn_if_image=False) + else: + raise ValueError(f"Unrecognized type for prompt: {type(prompt)}") + n_token = count_tokens(prompt_str, model=model_name) + if n_token <= max_prompt_tokens: + return prompt + shrinkable.shrink() + + logging.info( + dedent( + f"""\ + After {max_iterations} shrink iterations, the prompt is still + {count_tokens(prompt_str)} tokens (greater than {max_prompt_tokens}). Returning the prompt as is.""" + ) + ) + return prompt diff --git a/src/agentlab/agents/prompting/elements.py b/src/agentlab/agents/prompting/elements.py new file mode 100644 index 000000000..a7a3f6f81 --- /dev/null +++ b/src/agentlab/agents/prompting/elements.py @@ -0,0 +1,514 @@ +import platform +import time +from copy import copy + +from browsergym.core.action.base import AbstractActionSet +from browsergym.utils.obs import flatten_axtree_to_str, flatten_dom_to_str, overlay_som, prune_html + +from agentlab.llm.llm_utils import ( + ParseError, + extract_code_blocks, + image_to_jpg_base64_url, + parse_html_tags_raise, + BaseMessage, +) + +from .base import PromptElement, Shrinkable, Trunkater +from .flags import ActionFlags, ObsFlags + + +class HTML(Trunkater): + def __init__(self, html, visible_elements_only: bool, visible: bool = True, prefix="") -> None: + super().__init__(visible=visible, start_trunkate_iteration=5) + if visible_elements_only: + visible_elements_note = """\ +Note: only elements that are visible in the viewport are presented. You might need to scroll the page, or open tabs or menus to see more. + +""" + else: + visible_elements_note = "" + self._prompt = f"\n{prefix}HTML:\n{visible_elements_note}{html}\n" + + +class AXTree(Trunkater): + def __init__( + self, + ax_tree, + visible_elements_only: bool, + visible: bool = True, + coord_type=None, + visible_tag=True, + prefix="", + ) -> None: + super().__init__(visible=visible, start_trunkate_iteration=10) + bid_info = """\ +Note: [bid] is the unique alpha-numeric identifier at the beginning of lines for each element in the AXTree. Always use bid to refer to elements in your actions. + +""" + if coord_type == "center": + coord_note = """\ +Note: center coordinates are provided in parenthesis and are relative to the top left corner of the page. + +""" + elif coord_type == "box": + coord_note = """\ +Note: bounding box of each object are provided in parenthesis and are relative to the top left corner of the page. + +""" + else: + coord_note = "" + if visible_elements_only: + visible_elements_note = """\ +Note: only elements that are visible in the viewport are presented. You might need to scroll the page, or open tabs or menus to see more. + +""" + else: + visible_elements_note = "" + + if visible_tag: + vsible_tag_note = """\ +Note: You can only interact with visible elements. If the "visible" tag is not +present, the element is not visible on the page. + +""" + else: + vsible_tag_note = "" + self._prompt = f"\n{prefix}AXTree:\n{bid_info}{coord_note}{visible_elements_note}{vsible_tag_note}{ax_tree}\n" + + +class Error(PromptElement): + def __init__(self, error: str, visible: bool = True, prefix="", limit_logs=True) -> None: + logs_separator = "Call log:" + if limit_logs and logs_separator in error: + error, logs = error.split(logs_separator) + logs = "\n".join(logs.split("\n")[:10]) + error = error + f"\n{logs_separator}\n{logs}" + + super().__init__(visible=visible) + self._prompt = f"\n{prefix}Error from previous action:\n{error}\n" + + +class FocusedElement(PromptElement): + def __init__(self, bid, visible: bool = True, prefix="") -> None: + super().__init__(visible=visible) + self._prompt = f""" +{prefix}Focused element: +""" + if bid: + self._prompt += f"""\ +bid={repr(bid)} +""" + else: + self._prompt += f"""\ +None +""" + + +class Tabs(PromptElement): + def __init__(self, obs, visible: bool = True, prefix="") -> None: + super().__init__(visible=visible) + self.obs = obs + self.prefix = prefix + + @property + def _prompt(self) -> str: + # by implementing this as a property, it's only coputed if visible + prompt_pieces = [f"\n{self.prefix}Currently open tabs:"] + for page_index, (page_url, page_title) in enumerate( + zip(self.obs["open_pages_urls"], self.obs["open_pages_titles"]) + ): + active_or_not = " (active tab)" if page_index == self.obs["active_page_index"] else "" + prompt_piece = f"""\ +Tab {page_index}{active_or_not}: + Title: {page_title} + URL: {page_url} +""" + prompt_pieces.append(prompt_piece) + return "\n".join(prompt_pieces) + + +class Observation(Shrinkable): + """Observation of the current step. + + Contains the html, the accessibility tree and the error logs. + """ + + def __init__(self, obs, flags: ObsFlags) -> None: + super().__init__() + self.flags = flags + self.obs = obs + + self.tabs = Tabs( + obs, + visible=lambda: flags.use_tabs, + prefix="## ", + ) + + self.html = HTML( + obs[flags.html_type], + visible_elements_only=flags.filter_visible_elements_only, + visible=lambda: flags.use_html, + prefix="## ", + ) + self.ax_tree = AXTree( + obs["axtree_txt"], + visible_elements_only=flags.filter_visible_elements_only, + visible=lambda: flags.use_ax_tree, + coord_type=flags.extract_coords, + visible_tag=flags.extract_visible_tag, + prefix="## ", + ) + self.error = Error( + obs["last_action_error"], + visible=lambda: flags.use_error_logs and obs["last_action_error"], + prefix="## ", + ) + self.focused_element = FocusedElement( + obs["focused_element_bid"], + visible=flags.use_focused_element, + prefix="## ", + ) + + def shrink(self): + self.ax_tree.shrink() + self.html.shrink() + + @property + def _prompt(self) -> str: + return f""" +# Observation of current step: +{self.tabs.prompt}{self.html.prompt}{self.ax_tree.prompt}{self.focused_element.prompt}{self.error.prompt} + +""" + + def add_screenshot(self, prompt: BaseMessage) -> BaseMessage: + if self.flags.use_screenshot: + if self.flags.use_som: + screenshot = self.obs["screenshot_som"] + prompt.add_text( + "\n## Screenshot:\nHere is a screenshot of the page, it is annotated with bounding boxes and corresponding bids:" + ) + else: + screenshot = self.obs["screenshot"] + prompt.add_text("\n## Screenshot:\nHere is a screenshot of the page:") + img_url = image_to_jpg_base64_url(screenshot) + prompt.add_image(img_url, detail=self.flags.openai_vision_detail) + return prompt + + +class MacNote(PromptElement): + def __init__(self) -> None: + super().__init__(visible=platform.system() == "Darwin") + self._prompt = ( + "\nNote: you are on mac so you should use Meta instead of Control for Control+C etc.\n" + ) + + +class BeCautious(PromptElement): + def __init__(self, visible: bool = True) -> None: + super().__init__(visible=visible) + self._prompt = f"""\ +\nBe very cautious. Avoid submitting anything before verifying the effect of your +actions. Take the time to explore the effect of safe actions first. For example +you can fill a few elements of a form, but don't click submit before verifying +that everything was filled correctly.\n""" + + +class GoalInstructions(PromptElement): + def __init__(self, goal_object, visible: bool = True, extra_instructions=None) -> None: + super().__init__(visible) + self._prompt = [ + dict( + type="text", + text=f"""\ +# Instructions +Review the current state of the page and all other information to find the best +possible next action to accomplish your goal. Your answer will be interpreted +and executed by a program, make sure to follow the formatting instructions. + +## Goal: +""", + ) + ] + + self._prompt += goal_object + + if extra_instructions: + self._prompt += [ + dict( + type="text", + text=f""" + +## Extra instructions: + +{extra_instructions} +""", + ) + ] + + +class ChatInstructions(PromptElement): + def __init__(self, chat_messages, visible: bool = True, extra_instructions=None) -> None: + super().__init__(visible) + self._prompt = f"""\ +# Instructions + +You are a UI Assistant, your goal is to help the user perform tasks using a web browser. You can +communicate with the user via a chat, in which the user gives you instructions and in which you +can send back messages. You have access to a web browser that both you and the user can see, +and with which only you can interact via specific commands. + +Review the instructions from the user, the current state of the page and all other information +to find the best possible next action to accomplish your goal. Your answer will be interpreted +and executed by a program, make sure to follow the formatting instructions. + +## Chat messages: + +""" + self._prompt += "\n".join( + [ + f"""\ + - [{msg['role']}] UTC Time: {time.asctime(time.gmtime(msg['timestamp']))} - Local Time: {time.asctime(time.localtime(msg['timestamp']))} - {msg['message']}""" + for msg in chat_messages + ] + ) + + if extra_instructions: + self._prompt += f""" + +## Extra instructions: + +{extra_instructions} +""" + + +class Hints(PromptElement): + """Not super useful and stale.""" + + # NOTE: are these hints still relevant? + _prompt = """\ +Note: +* Some tasks may be game like and may require to interact with the mouse position +in x, y coordinates. +* Some text field might have auto completion. To see it, you have to type a few +characters and wait until next step. +* If you have to cut and paste, don't forget to select the text first. +* Coordinate inside an SVG are relative to it's top left corner. +* Make sure to use bid to identify elements when using commands. +* Interacting with combobox, dropdowns and auto-complete fields can be tricky, +sometimes you need to use select_option, while other times you need to use fill +or click and wait for the reaction of the page. +""" + + +class SystemPrompt(PromptElement): + _prompt = """\ +You are an agent trying to solve a web task based on the content of the page and +user instructions. You can interact with the page and explore, and send messages to the user. Each time you +submit an action it will be sent to the browser and you will receive a new page.""" + + +class ActionPrompt(PromptElement): + + _concrete_ex = """ + +click('a324') + +""" + + def __init__(self, action_set: AbstractActionSet, action_flags: ActionFlags) -> None: + super().__init__() + self.action_set = action_set + self.action_flags = action_flags + action_set_generic_info = """\ +Note: This action set allows you to interact with your environment. Most of them +are python function executing playwright code. The primary way of referring to +elements in the page is through bid which are specified in your observations. + +""" + action_description = action_set.describe( + with_long_description=action_flags.long_description, + with_examples=action_flags.individual_examples, + ) + self._prompt = ( + f"# Action space:\n{action_set_generic_info}{action_description}{MacNote().prompt}\n" + ) + self._abstract_ex = f""" + +{self.action_set.example_action(abstract=True)} + +""" + + def _parse_answer(self, text_answer): + try: + ans_dict = parse_html_tags_raise(text_answer, keys=["action"], merge_multiple=True) + except ParseError as e: + if self.action_flags.is_strict: + raise e + else: + # try to extract code blocks + blocks = extract_code_blocks(text_answer) + if len(blocks) == 0: + raise e + else: + code = "\n".join([block for _, block in blocks]) + ans_dict = {"action": code, "parse_error": str(e)} + + try: + if ans_dict["action"] == "None": + # Used by reproducibility agent for backward compatibility of + # traces missing LLM's response in chat messages. + ans_dict["action"] = None + else: + # just check if action can be mapped to python code but keep action as is + # the environment will be responsible for mapping it to python + self.action_set.to_python_code(ans_dict["action"]) + except Exception as e: + raise ParseError( + f"Error while parsing action\n: {e}\n" + "Make sure your answer is restricted to the allowed actions." + ) + + return ans_dict + + +class Think(PromptElement): + _prompt = "" + + _abstract_ex = """ + +Think step by step. If you need to make calculations such as coordinates, write them here. Describe the effect +that your previous action had on the current content of the page. + +""" + _concrete_ex = """ + +From previous action I tried to set the value of year to "2022", +using select_option, but it doesn't appear to be in the form. It may be a +dynamic dropdown, I will try using click with the bid "a324" and look at the +response from the page. + +""" + + def _parse_answer(self, text_answer): + try: + return parse_html_tags_raise(text_answer, keys=["think"], merge_multiple=True) + except ParseError as e: + return {"think": text_answer, "parse_error": str(e)} + + +class HistoryStep(Shrinkable): + def __init__( + self, previous_obs, current_obs, action, memory, thought, flags: ObsFlags, shrink_speed=1 + ) -> None: + super().__init__() + self.error = Error( + current_obs["last_action_error"], + visible=( + lambda: flags.use_error_logs + and current_obs["last_action_error"] + and flags.use_past_error_logs + ), + prefix="### ", + ) + self.shrink_speed = shrink_speed + self.action = action + self.memory = memory + self.thought = thought + self.flags = flags + + def shrink(self): + super().shrink() + + @property + def _prompt(self) -> str: + prompt = "" + + if self.flags.use_think_history: + prompt += f"\n\n{self.thought}\n\n" + + if self.flags.use_action_history: + prompt += f"\n\n{self.action}\n\n" + + prompt += f"{self.error.prompt}" + + if self.memory is not None: + prompt += f"\n\n{self.memory}\n\n" + + return prompt + + +class History(Shrinkable): + def __init__( + self, history_obs, actions, memories, thoughts, flags: ObsFlags, shrink_speed=1 + ) -> None: + if memories is None: + memories = [None] * len(actions) + super().__init__(visible=lambda: flags.use_history) + assert len(history_obs) == len(actions) + 1 + assert len(history_obs) == len(memories) + 1 + + self.shrink_speed = shrink_speed + self.history_steps: list[HistoryStep] = [] + + for i in range(1, len(history_obs)): + self.history_steps.append( + HistoryStep( + history_obs[i - 1], + history_obs[i], + actions[i - 1], + memories[i - 1], + thoughts[i - 1], + flags, + ) + ) + + def shrink(self): + """Shrink individual steps""" + # TODO set the shrink speed of older steps to be higher + super().shrink() + for step in self.history_steps: + step.shrink() + + @property + def _prompt(self): + prompts = ["# History of interaction with the task:\n"] + for i, step in enumerate(self.history_steps): + prompts.append(f"## step {i}") + prompts.append(step.prompt) + return "\n".join(prompts) + "\n" + + +def make_obs_preprocessor(flags: ObsFlags): + def obs_mapping(obs: dict): + obs = copy(obs) + obs["dom_txt"] = flatten_dom_to_str( + obs["dom_object"], + extra_properties=obs["extra_element_properties"], + with_visible=flags.extract_visible_tag, + with_clickable=flags.extract_clickable_tag, + with_center_coords=flags.extract_coords == "center", + with_bounding_box_coords=flags.extract_coords == "box", + filter_visible_only=flags.filter_visible_elements_only, + filter_with_bid_only=flags.filter_with_bid_only, + filter_som_only=flags.filter_som_only, + ) + obs["axtree_txt"] = flatten_axtree_to_str( + obs["axtree_object"], + extra_properties=obs["extra_element_properties"], + with_visible=flags.extract_visible_tag, + with_clickable=flags.extract_clickable_tag, + with_center_coords=flags.extract_coords == "center", + with_bounding_box_coords=flags.extract_coords == "box", + filter_visible_only=flags.filter_visible_elements_only, + filter_with_bid_only=flags.filter_with_bid_only, + filter_som_only=flags.filter_som_only, + ) + obs["pruned_html"] = prune_html(obs["dom_txt"]) + obs["screenshot_som"] = overlay_som( + obs["screenshot"], extra_properties=obs["extra_element_properties"] + ) + + return obs + + return obs_mapping diff --git a/src/agentlab/agents/prompting/flags.py b/src/agentlab/agents/prompting/flags.py new file mode 100644 index 000000000..aac27ed62 --- /dev/null +++ b/src/agentlab/agents/prompting/flags.py @@ -0,0 +1,86 @@ +from copy import deepcopy +from dataclasses import asdict, dataclass +from typing import Literal + +from bgym import HighLevelActionSetArgs + + +class Flags: + """Base class for flags. Mostly for backward compatibility.""" + + def copy(self): + return deepcopy(self) + + def asdict(self): + """Helper for JSON serializable requirement.""" + return asdict(self) + + @classmethod + def from_dict(self, flags_dict): + """Helper for JSON serializable requirement.""" + if isinstance(flags_dict, ObsFlags): + return flags_dict + + if not isinstance(flags_dict, dict): + raise ValueError(f"Unregcognized type for flags_dict of type {type(flags_dict)}.") + return ObsFlags(**flags_dict) + + +@dataclass +class ObsFlags(Flags): + """ + A class to represent various flags used to control features in an application. + + Attributes: + use_html (bool): Use the HTML in the prompt. + use_ax_tree (bool): Use the accessibility tree in the prompt. + use_focused_element (bool): Provide the ID of the focused element. + use_error_logs (bool): Expose the previous error in the prompt. + use_history (bool): Enable history of previous steps in the prompt. + use_past_error_logs (bool): If use_history is True, expose all previous errors in the history. + use_action_history (bool): If use_history is True, include the actions in the history. + use_think_history (bool): If use_history is True, include all previous chains of thoughts in the history. + use_diff (bool): Add a diff of the current and previous HTML to the prompt. + html_type (str): Type of HTML to use in the prompt, may depend on preprocessing of observation. + use_screenshot (bool): Add a screenshot of the page to the prompt, following OpenAI's API. This will be automatically disabled if the model does not have vision capabilities. + use_som (bool): Add a set of marks to the screenshot. + extract_visible_tag (bool): Add a "visible" tag to visible elements in the AXTree. + extract_clickable_tag (bool): Add a "clickable" tag to clickable elements in the AXTree. + extract_coords (Literal['False', 'center', 'box']): Add the coordinates of the elements. + filter_visible_elements_only (bool): Only show visible elements in the AXTree. + """ + + use_html: bool = True + use_ax_tree: bool = False + use_tabs: bool = False + use_focused_element: bool = False + use_error_logs: bool = False + use_history: bool = False + use_past_error_logs: bool = False + use_action_history: bool = False + use_think_history: bool = False + use_diff: bool = False # + html_type: str = "pruned_html" + use_screenshot: bool = True + use_som: bool = False + extract_visible_tag: bool = False + extract_clickable_tag: bool = False + extract_coords: Literal["False", "center", "box"] = "False" + filter_visible_elements_only: bool = False + # low sets the token count of each image to 65 (85?) + # high sets the token count of each image to 2*65 (2*85?) times the amount of 512x512px patches + # auto chooses between low and high based on image size (openai default) + openai_vision_detail: Literal["low", "high", "auto"] = "auto" + filter_with_bid_only: bool = False + filter_som_only: bool = False + + +@dataclass +class ActionFlags(Flags): + action_set: HighLevelActionSetArgs = None # should be set by the set_benchmark method + long_description: bool = True + individual_examples: bool = False + + # for backward compatibility + multi_actions: bool = None + is_strict: bool = None From 4662a4d17f0fa2e5f0af6a23517895ea49e2c247 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 26 Mar 2026 12:49:29 +0000 Subject: [PATCH 3/9] refactor: clean up benchmark hack and legacy env var fallback Task 5: Replace inline AssistantBench hack in EnvArgs.make_env() with a generic pre_env_setup_hooks callback list. Benchmarks can now register hooks instead of hard-coding task_name checks. Task 6: Add deprecation warning when legacy UI_COPILOT_RESULTS_DIR env var is used. Backward compatibility preserved with DeprecationWarning. Co-Authored-By: Lukas Burger --- src/agentlab/experiments/exp_utils.py | 12 +++++++++++- src/agentlab/experiments/loop.py | 14 ++++++++------ 2 files changed, 19 insertions(+), 7 deletions(-) diff --git a/src/agentlab/experiments/exp_utils.py b/src/agentlab/experiments/exp_utils.py index 27b909b2e..83f724c5e 100644 --- a/src/agentlab/experiments/exp_utils.py +++ b/src/agentlab/experiments/exp_utils.py @@ -16,7 +16,17 @@ # TODO move this to a more appropriate place RESULTS_DIR = os.environ.get("AGENTLAB_EXP_ROOT", None) if RESULTS_DIR is None: - RESULTS_DIR = os.environ.get("UI_COPILOT_RESULTS_DIR", None) + _legacy_dir = os.environ.get("UI_COPILOT_RESULTS_DIR", None) + if _legacy_dir is not None: + import warnings + + warnings.warn( + "UI_COPILOT_RESULTS_DIR is deprecated. Use AGENTLAB_EXP_ROOT instead. " + "Support for UI_COPILOT_RESULTS_DIR will be removed in a future version.", + DeprecationWarning, + stacklevel=1, + ) + RESULTS_DIR = _legacy_dir if RESULTS_DIR is None: logging.info("$AGENTLAB_EXP_ROOT is not defined, Using $HOME/agentlab_results.") RESULTS_DIR = Path.home() / "agentlab_results" diff --git a/src/agentlab/experiments/loop.py b/src/agentlab/experiments/loop.py index de4b976a6..f7a519723 100644 --- a/src/agentlab/experiments/loop.py +++ b/src/agentlab/experiments/loop.py @@ -49,6 +49,10 @@ class EnvArgs(DataClassJsonMixin): storage_state: Optional[str | Path | dict] = None task_kwargs: Optional[dict] = None # use default value from BrowserGym pre_observation_delay: float = None # seconds, wait for JS events to be fired + pre_env_setup_hooks: Optional[list] = field( + default=None, + metadata={"dataclasses_json": {"exclude": lambda x: True}}, + ) def make_env( self, action_mapping, exp_dir, exp_task_kwargs: dict = {}, use_raw_page_output=True @@ -81,12 +85,10 @@ def make_env( if exp_task_kwargs: extra_kwargs["task_kwargs"] = extra_kwargs.get("task_kwargs", {}) | exp_task_kwargs - # assistantbench hack, write the task output (agent prediction) to a file in the experiment's directory - # TODO: find a better way to deal with this - if self.task_name.startswith("assistantbench.test"): - extra_kwargs["task_kwargs"] = extra_kwargs.get("task_kwargs", {}) | { - "output_file": exp_dir / "assistantbench-prediction.json" - } + # Run any registered pre-env-setup hooks (e.g. benchmark-specific config) + if self.pre_env_setup_hooks: + for hook in self.pre_env_setup_hooks: + extra_kwargs = hook(self, extra_kwargs, exp_dir) return gym.make( _get_env_name(self.task_name), From 98718e67bfced949f74bf1af1c4cd90231385f30 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 26 Mar 2026 12:49:36 +0000 Subject: [PATCH 4/9] refactor: extract shared test fixtures to tests/helpers/mock_llms.py Task 7: Move CheatMiniWoBLLM_ParseRetry, CheatMiniWoBLLMArgs_ParseRetry, CheatLLM_LLMError, and CheatLLMArgs_LLMError from test_agent.py into a shared tests/helpers/mock_llms.py module for reuse across test files. Co-Authored-By: Lukas Burger --- tests/agents/test_agent.py | 104 +++---------------------------------- tests/helpers/__init__.py | 1 + tests/helpers/mock_llms.py | 103 ++++++++++++++++++++++++++++++++++++ 3 files changed, 112 insertions(+), 96 deletions(-) create mode 100644 tests/helpers/__init__.py create mode 100644 tests/helpers/mock_llms.py diff --git a/tests/agents/test_agent.py b/tests/agents/test_agent.py index 7fcaafa28..2e98d0e6f 100644 --- a/tests/agents/test_agent.py +++ b/tests/agents/test_agent.py @@ -1,17 +1,19 @@ -import re import tempfile -from dataclasses import dataclass from pathlib import Path -from openai import OpenAIError - from agentlab.agents.generic_agent.agent_configs import FLAGS_GPT_3_5 from agentlab.agents.generic_agent.generic_agent import GenericAgentArgs from agentlab.analyze import inspect_results from agentlab.experiments import launch_exp from agentlab.experiments.loop import EnvArgs, ExpArgs -from agentlab.llm.chat_api import BaseModelArgs, CheatMiniWoBLLMArgs -from agentlab.llm.llm_utils import Discussion +from agentlab.llm.chat_api import CheatMiniWoBLLMArgs + +from tests.helpers.mock_llms import ( + CheatLLM_LLMError, + CheatLLMArgs_LLMError, + CheatMiniWoBLLM_ParseRetry, + CheatMiniWoBLLMArgs_ParseRetry, +) def test_generic_agent(): @@ -45,96 +47,6 @@ def test_generic_agent(): assert result_record[key].iloc[0] == target_val -@dataclass -class CheatMiniWoBLLM_ParseRetry: - """For unit-testing purposes only. It only work with miniwob.click-test task.""" - - n_retry: int - retry_count: int = 0 - - def __call__(self, messages) -> str: - if self.retry_count < self.n_retry: - self.retry_count += 1 - return dict(role="assistant", content="I'm retrying") - - if isinstance(messages, Discussion): - prompt = messages.to_string() - else: - prompt = messages[1].get("content", "") - match = re.search(r"^\s*\[(\d+)\].*button", prompt, re.MULTILINE | re.IGNORECASE) - - if match: - bid = match.group(1) - action = f'click("{bid}")' - else: - raise Exception("Can't find the button's bid") - - answer = f"""I'm clicking the button as requested. - -{action} - -""" - return dict(role="assistant", content=answer) - - def get_stats(self): - return {} - - -@dataclass -class CheatMiniWoBLLMArgs_ParseRetry(BaseModelArgs): - n_retry: int = 2 - model_name: str = "test/cheat_miniwob_click_test_parse_retry" - - def make_model(self): - return CheatMiniWoBLLM_ParseRetry(n_retry=self.n_retry) - - -@dataclass -class CheatLLM_LLMError: - """For unit-testing purposes only. Fails to call LLM""" - - n_retry: int = 0 - success: bool = False - - def __call__(self, messages) -> str: - if self.success: - if isinstance(messages, Discussion): - prompt = messages.to_string() - else: - prompt = messages[1].get("content", "") - match = re.search(r"^\s*\[(\d+)\].*button", prompt, re.MULTILINE | re.IGNORECASE) - - if match: - bid = match.group(1) - action = f'click("{bid}")' - else: - raise Exception("Can't find the button's bid") - - answer = f"""I'm clicking the button as requested. - - {action} - - """ - return dict(role="assistant", content=answer) - raise OpenAIError("LLM failed to respond") - - def get_stats(self): - return {"n_llm_retry": self.n_retry, "n_llm_busted_retry": int(not self.success)} - - -@dataclass -class CheatLLMArgs_LLMError(BaseModelArgs): - n_retry: int = 2 - success: bool = False - model_name: str = "test/cheat_miniwob_click_test_parse_retry" - - def make_model(self): - return CheatLLM_LLMError( - n_retry=self.n_retry, - success=self.success, - ) - - def test_generic_agent_parse_retry(): exp_args = ExpArgs( agent_args=GenericAgentArgs( diff --git a/tests/helpers/__init__.py b/tests/helpers/__init__.py new file mode 100644 index 000000000..8b1378917 --- /dev/null +++ b/tests/helpers/__init__.py @@ -0,0 +1 @@ + diff --git a/tests/helpers/mock_llms.py b/tests/helpers/mock_llms.py new file mode 100644 index 000000000..b1a55649f --- /dev/null +++ b/tests/helpers/mock_llms.py @@ -0,0 +1,103 @@ +"""Shared mock LLM classes for testing. + +These mock LLMs simulate various behaviors (parse retries, LLM errors) for +unit-testing agent behavior without making real API calls. +""" + +import re +from dataclasses import dataclass + +from openai import OpenAIError + +from agentlab.llm.chat_api import BaseModelArgs +from agentlab.llm.llm_utils import Discussion + + +@dataclass +class CheatMiniWoBLLM_ParseRetry: + """For unit-testing purposes only. It only work with miniwob.click-test task.""" + + n_retry: int + retry_count: int = 0 + + def __call__(self, messages) -> str: + if self.retry_count < self.n_retry: + self.retry_count += 1 + return dict(role="assistant", content="I'm retrying") + + if isinstance(messages, Discussion): + prompt = messages.to_string() + else: + prompt = messages[1].get("content", "") + match = re.search(r"^\s*\[(\d+)\].*button", prompt, re.MULTILINE | re.IGNORECASE) + + if match: + bid = match.group(1) + action = f'click("{bid}")' + else: + raise Exception("Can't find the button's bid") + + answer = f"""I'm clicking the button as requested. + +{action} + +""" + return dict(role="assistant", content=answer) + + def get_stats(self): + return {} + + +@dataclass +class CheatMiniWoBLLMArgs_ParseRetry(BaseModelArgs): + n_retry: int = 2 + model_name: str = "test/cheat_miniwob_click_test_parse_retry" + + def make_model(self): + return CheatMiniWoBLLM_ParseRetry(n_retry=self.n_retry) + + +@dataclass +class CheatLLM_LLMError: + """For unit-testing purposes only. Fails to call LLM""" + + n_retry: int = 0 + success: bool = False + + def __call__(self, messages) -> str: + if self.success: + if isinstance(messages, Discussion): + prompt = messages.to_string() + else: + prompt = messages[1].get("content", "") + match = re.search(r"^\s*\[(\d+)\].*button", prompt, re.MULTILINE | re.IGNORECASE) + + if match: + bid = match.group(1) + action = f'click("{bid}")' + else: + raise Exception("Can't find the button's bid") + + answer = f"""I'm clicking the button as requested. + + {action} + + """ + return dict(role="assistant", content=answer) + raise OpenAIError("LLM failed to respond") + + def get_stats(self): + return {"n_llm_retry": self.n_retry, "n_llm_busted_retry": int(not self.success)} + + +@dataclass +class CheatLLMArgs_LLMError(BaseModelArgs): + n_retry: int = 2 + success: bool = False + model_name: str = "test/cheat_miniwob_click_test_parse_retry" + + def make_model(self): + return CheatLLM_LLMError( + n_retry=self.n_retry, + success=self.success, + ) From c48c37ae3691bc5469d898df8102614e4cbccf7c Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 26 Mar 2026 12:49:42 +0000 Subject: [PATCH 5/9] refactor: fix Block.make() cloning with copy.deepcopy Task 8: Replace asdict(self) with copy.deepcopy(self) in Block.make() to safely handle non-serializable fields. Update docstring accordingly. Co-Authored-By: Lukas Burger --- .../agents/tool_use_agent/tool_use_agent.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/src/agentlab/agents/tool_use_agent/tool_use_agent.py b/src/agentlab/agents/tool_use_agent/tool_use_agent.py index 375c829eb..54cc220f5 100644 --- a/src/agentlab/agents/tool_use_agent/tool_use_agent.py +++ b/src/agentlab/agents/tool_use_agent/tool_use_agent.py @@ -6,8 +6,8 @@ import time from abc import ABC, abstractmethod from collections import defaultdict -from copy import copy -from dataclasses import asdict, dataclass, field +from copy import copy, deepcopy +from dataclasses import dataclass, field from pathlib import Path from typing import Any, Literal @@ -51,14 +51,16 @@ def _init(self): pass def make(self) -> "Block": - """Returns a copy so the init can start adding some stuff to `self` without changing the - original datatclass that should only contain a config. - The aim is avoid having 2 calss definition for each block, e.g. Block and BlockArgs. + """Returns a deep copy so the init can start adding some stuff to `self` without + changing the original dataclass that should only contain a config. + The aim is to avoid having 2 class definitions for each block, e.g. Block and BlockArgs. + + Uses copy.deepcopy instead of asdict to safely handle non-serializable fields. Returns: Block: A copy of the current block instance with initialization applied. """ - block = self.__class__(**asdict(self)) + block = deepcopy(self) block._init() return block From bdc0ebe6ebe40827ffc6216452c447fb97e43a25 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 26 Mar 2026 12:52:55 +0000 Subject: [PATCH 6/9] fix: add missing docstring annotations for darglint compliance Co-Authored-By: Lukas Burger --- src/agentlab/llm/tracking.py | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/src/agentlab/llm/tracking.py b/src/agentlab/llm/tracking.py index 05c69e822..26593de1d 100644 --- a/src/agentlab/llm/tracking.py +++ b/src/agentlab/llm/tracking.py @@ -96,7 +96,7 @@ def get_pricing(model_name: str, provider: str) -> dict: provider: One of 'openai', 'anthropic', 'openrouter', 'litellm'. Returns: - A dict mapping model names to {"prompt": float, "completion": float} pricing info. + dict: Pricing information from the specified provider. """ pricing_fn_map = { "openai": _get_pricing_openai, @@ -118,6 +118,9 @@ def get_pricing_openrouter(): """Returns a dictionary of model pricing for OpenRouter models. .. deprecated:: Use get_pricing(model_name, 'openrouter') instead. + + Returns: + dict: Pricing information for OpenRouter models. """ return _get_pricing_openrouter() @@ -126,6 +129,9 @@ def get_pricing_openai(): """Returns a dictionary of model pricing for OpenAI models. .. deprecated:: Use get_pricing(model_name, 'openai') instead. + + Returns: + dict: Pricing information for OpenAI models. """ return _get_pricing_openai() @@ -134,6 +140,9 @@ def get_pricing_anthropic(): """Returns a dictionary of model pricing for Anthropic models. .. deprecated:: Use get_pricing(model_name, 'anthropic') instead. + + Returns: + dict: Pricing information for Anthropic models. """ return _get_pricing_anthropic() @@ -142,6 +151,12 @@ def get_pricing_litellm(model_name): """Returns a dictionary of model pricing for a LiteLLM model. .. deprecated:: Use get_pricing(model_name, 'litellm') instead. + + Args: + model_name: The model name to look up pricing for. + + Returns: + dict: Pricing information for the specified LiteLLM model. """ return _get_pricing_litellm(model_name) From 6e89e2df1e315740b4b3205159efbbca0c8d5f26 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 26 Mar 2026 13:10:42 +0000 Subject: [PATCH 7/9] fix: initialize self.retries in ChatModel.__init__ and preserve assistantbench hack - ChatModel.get_stats() referenced self.retries which was no longer set after the retry logic was moved to shared helpers. Initialize it in __init__. - Re-add the assistantbench.test output_file hack alongside the new pre_env_setup_hooks mechanism. Add TODO to migrate it to a hook. Co-Authored-By: Lukas Burger --- src/agentlab/experiments/loop.py | 7 +++++++ src/agentlab/llm/chat_api.py | 3 +++ 2 files changed, 10 insertions(+) diff --git a/src/agentlab/experiments/loop.py b/src/agentlab/experiments/loop.py index f7a519723..595492056 100644 --- a/src/agentlab/experiments/loop.py +++ b/src/agentlab/experiments/loop.py @@ -90,6 +90,13 @@ def make_env( for hook in self.pre_env_setup_hooks: extra_kwargs = hook(self, extra_kwargs, exp_dir) + # assistantbench hack, write the task output (agent prediction) to a file in the experiment's directory + # TODO: migrate this to a pre_env_setup_hook registered in the AssistantBench benchmark module + if self.task_name.startswith("assistantbench.test"): + extra_kwargs["task_kwargs"] = extra_kwargs.get("task_kwargs", {}) | { + "output_file": exp_dir / "assistantbench-prediction.json" + } + return gym.make( _get_env_name(self.task_name), disable_env_checker=True, diff --git a/src/agentlab/llm/chat_api.py b/src/agentlab/llm/chat_api.py index b01fa0195..ebc232b33 100644 --- a/src/agentlab/llm/chat_api.py +++ b/src/agentlab/llm/chat_api.py @@ -247,6 +247,9 @@ def __init__( api_key=api_key, **client_args, ) + # Initialize retry tracking (updated by shared retry helper is not possible, + # so we keep a default value for backward compatibility with get_stats()) + self.retries = 0 def __call__(self, messages: list[dict], n_samples: int = 1, temperature: float = None) -> dict: temperature = temperature if temperature is not None else self.temperature From 8419d59c846c1e008aa26460ff5b81a17dc3f9ef Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 26 Mar 2026 13:18:14 +0000 Subject: [PATCH 8/9] fix: add usage-None validation for OpenRouter and expose retry stats - is_openai_response_valid now checks for usage=None (known OpenRouter issue) and returns False to trigger a retry instead of crashing with AttributeError on completion.usage.prompt_tokens. - generic_call_api_with_retries accepts an optional stats dict that gets populated with the attempt count on success. - call_openai_api_with_retries passes stats through to the generic helper. - ChatModel.__call__ uses the stats dict to update self.retries so get_stats() reports the actual retry count. Co-Authored-By: Lukas Burger --- src/agentlab/llm/chat_api.py | 5 +++-- src/agentlab/llm/llm_utils.py | 18 +++++++++++++++--- 2 files changed, 18 insertions(+), 5 deletions(-) diff --git a/src/agentlab/llm/chat_api.py b/src/agentlab/llm/chat_api.py index ebc232b33..26ae71e91 100644 --- a/src/agentlab/llm/chat_api.py +++ b/src/agentlab/llm/chat_api.py @@ -247,8 +247,6 @@ def __init__( api_key=api_key, **client_args, ) - # Initialize retry tracking (updated by shared retry helper is not possible, - # so we keep a default value for backward compatibility with get_stats()) self.retries = 0 def __call__(self, messages: list[dict], n_samples: int = 1, temperature: float = None) -> dict: @@ -262,11 +260,14 @@ def __call__(self, messages: list[dict], n_samples: int = 1, temperature: float "logprobs": self.log_probs, } + retry_stats = {} completion = call_openai_api_with_retries( self.client.chat.completions.create, api_params, max_retries=self.max_retry, + stats=retry_stats, ) + self.retries = retry_stats.get("attempts", 1) input_tokens = completion.usage.prompt_tokens output_tokens = completion.usage.completion_tokens diff --git a/src/agentlab/llm/llm_utils.py b/src/agentlab/llm/llm_utils.py index 10013b72a..ae37729e7 100644 --- a/src/agentlab/llm/llm_utils.py +++ b/src/agentlab/llm/llm_utils.py @@ -109,6 +109,7 @@ def generic_call_api_with_retries( max_retries=10, initial_retry_delay_seconds=20, max_retry_delay_seconds=60 * 5, + stats=None, ): """ Makes an API call with retries for transient failures, rate limiting, @@ -125,6 +126,8 @@ def generic_call_api_with_retries( max_retries: Maximum number of retry attempts. initial_retry_delay_seconds: Initial delay between retries in seconds. max_retry_delay_seconds: Maximum delay between retries in seconds. + stats: Optional dict that will be populated with retry statistics + (e.g. ``{"attempts": int}``) when provided. Returns: The API response if successful. @@ -221,6 +224,8 @@ def _handle_api_error(e, attempt): if is_response_valid_fn(response): logging.info(f"[Attempt {attempt}/{max_retries}] API call succeeded.") + if stats is not None: + stats["attempts"] = attempt return response else: if _handle_invalid_response_content(attempt): @@ -251,7 +256,7 @@ def _handle_api_error(e, attempt): raise RuntimeError(f"API call failed after {max_retries} retries.") -def call_openai_api_with_retries(client_function, api_params, max_retries=10): +def call_openai_api_with_retries(client_function, api_params, max_retries=10, stats=None): """ Makes an OpenAI API call with retries for transient failures, rate limiting, and invalid or error-containing responses. @@ -261,6 +266,7 @@ def call_openai_api_with_retries(client_function, api_params, max_retries=10): client_function: The OpenAI API client function to call. api_params: Parameters to pass to the client function. max_retries: Maximum number of retry attempts. + stats: Optional dict populated with retry statistics. Returns: The OpenAI API response if successful. @@ -272,6 +278,13 @@ def is_openai_response_valid(response): logging.warning(f"OpenAI API response contains an error attribute: {response.error}") return False # Treat as invalid for retry purposes if hasattr(response, "choices") and response.choices: # Chat Completion API + # Validate that usage info is present (OpenRouter sometimes returns usage=None) + if hasattr(response, "usage") and response.usage is None: + logging.warning( + "OpenAI API response has choices but usage is None. " + "This is likely an OpenRouter bug. Retrying." + ) + return False return True if hasattr(response, "output") and response.output: # Response API return True @@ -289,8 +302,7 @@ def get_openai_status_code(exception): api_error_exceptions=(openai.APIError,), # openai.RateLimitError is caught first get_status_code_fn=get_openai_status_code, max_retries=max_retries, - # You can also pass initial_retry_delay_seconds and max_retry_delay_seconds - # if you want to customize them from their defaults in the generic function. + stats=stats, ) From 3ec705da51d5fdda58b9bce1730fb45b5a009c2c Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 26 Mar 2026 13:27:01 +0000 Subject: [PATCH 9/9] fix: add isinstance guard to Anthropic tracker and pass min_retry_wait_time through - AnthropicChatModel tracker guard now checks isinstance(TRACKER.instance, LLMTracker) consistent with ChatModel, preventing TypeError when TRACKER.instance is None. - call_openai_api_with_retries now accepts initial_retry_delay_seconds and forwards it to generic_call_api_with_retries. - ChatModel.__call__ passes self.min_retry_wait_time as initial_retry_delay_seconds so the parameter is no longer silently ignored. Co-Authored-By: Lukas Burger --- src/agentlab/llm/chat_api.py | 5 ++++- src/agentlab/llm/llm_utils.py | 6 +++++- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/src/agentlab/llm/chat_api.py b/src/agentlab/llm/chat_api.py index 26ae71e91..fd623ebd5 100644 --- a/src/agentlab/llm/chat_api.py +++ b/src/agentlab/llm/chat_api.py @@ -266,6 +266,7 @@ def __call__(self, messages: list[dict], n_samples: int = 1, temperature: float api_params, max_retries=self.max_retry, stats=retry_stats, + initial_retry_delay_seconds=self.min_retry_wait_time, ) self.retries = retry_stats.get("attempts", 1) @@ -488,7 +489,9 @@ def __call__(self, messages: list[dict], n_samples: int = 1, temperature: float ) # Track usage if available - if hasattr(tracking.TRACKER, "instance"): + if hasattr(tracking.TRACKER, "instance") and isinstance( + tracking.TRACKER.instance, tracking.LLMTracker + ): tracking.TRACKER.instance( response.usage.input_tokens, response.usage.output_tokens, diff --git a/src/agentlab/llm/llm_utils.py b/src/agentlab/llm/llm_utils.py index ae37729e7..d06b502d2 100644 --- a/src/agentlab/llm/llm_utils.py +++ b/src/agentlab/llm/llm_utils.py @@ -256,7 +256,9 @@ def _handle_api_error(e, attempt): raise RuntimeError(f"API call failed after {max_retries} retries.") -def call_openai_api_with_retries(client_function, api_params, max_retries=10, stats=None): +def call_openai_api_with_retries( + client_function, api_params, max_retries=10, stats=None, initial_retry_delay_seconds=20 +): """ Makes an OpenAI API call with retries for transient failures, rate limiting, and invalid or error-containing responses. @@ -267,6 +269,7 @@ def call_openai_api_with_retries(client_function, api_params, max_retries=10, st api_params: Parameters to pass to the client function. max_retries: Maximum number of retry attempts. stats: Optional dict populated with retry statistics. + initial_retry_delay_seconds: Initial delay between retries in seconds. Returns: The OpenAI API response if successful. @@ -303,6 +306,7 @@ def get_openai_status_code(exception): get_status_code_fn=get_openai_status_code, max_retries=max_retries, stats=stats, + initial_retry_delay_seconds=initial_retry_delay_seconds, )