Move the overlay_action to utils

Author: amanjaiswal73892

src/agentlab/agents/agent_utils.py

@@ -1,5 +1,8 @@
+import copy
 from PIL import Image, ImageDraw
 from playwright.sync_api import Page
+from agentlab.llm.llm_utils import img_to_base_64
+from agentlab.analyze import overlay_utils
 
 
 def draw_mouse_pointer(image: Image.Image, x: int, y: int) -> Image.Image:
@@ -128,3 +131,14 @@ def zoom_webpage(page: Page, zoom_factor: float = 1.5):
 
     page.evaluate(f"document.documentElement.style.zoom='{zoom_factor*100}%'")
     return page
+
+
+def overlay_action(obs, action):
+    """Overlays actions on screenshot in-place"""
+    act_img = copy.deepcopy(obs["screenshot"])
+    act_img = Image.fromarray(act_img)
+    overlay_utils.annotate_action(act_img, action, properties=obs["extra_element_properties"])
+    return img_to_base_64(act_img)
+
+
+

src/agentlab/agents/hitl_agent/generic_human_guided_agent.py

@@ -24,7 +24,8 @@
     SystemMessage,
 )
 from agentlab.llm.tracking import cost_tracker_decorator
-
+from agentlab.agents.agent_utils import overlay_action
+from agentlab.llm.llm_utils import img_to_base_64
 
 class CandidatesGeneration(dp.PromptElement):
     # Ask for multiple alternatives; each candidate must contain <think> and <action>.
@@ -119,24 +120,6 @@ def _parse_answer(self, text_answer: str) -> Dict[str, Dict[str, str]]:
         return result
 
 
-def overlay_action(obs, action):
-    """Overlays actions on screenshot in-place"""
-    act_img = copy.deepcopy(obs["screenshot"])
-    act_img = Image.fromarray(act_img)
-    overlay_utils.annotate_action(act_img, action, properties=obs["extra_element_properties"])
-    return img_to_base_64(act_img)
-
-
-def img_to_base_64(image: Image.Image | np.ndarray) -> str:
-    """Converts a PIL Image or NumPy array to a base64-encoded string."""
-    if isinstance(image, np.ndarray):
-        image = Image.fromarray(image)
-    buffer = io.BytesIO()
-    image.save(buffer, format="PNG")
-    b64_str = base64.b64encode(buffer.getvalue()).decode("utf-8")
-    return b64_str
-
-
 @dataclass
 class MultipleProposalGenericAgentArgs(GenericAgentArgs):
 

src/agentlab/agents/hitl_agent/hitl_agent.py

@@ -1,22 +1,19 @@
-import base64
-import copy
-import io
 from dataclasses import dataclass
 from typing import Optional
 
 import bgym
-import numpy as np
 import playwright
 from browsergym.experiments.agent import Agent
-from PIL import Image
+
 
 from agentlab.agents.agent_args import AgentArgs
 from agentlab.agents.hitl_agent.base_multi_candidate_agent import MultiCandidateAgent
 from agentlab.agents.hitl_agent.hint_labelling import (
     HintLabeling,
     HintLabelingInputs,
 )
-from agentlab.analyze import overlay_utils
+from agentlab.agents.agent_utils import overlay_action
+from agentlab.llm.llm_utils import img_to_base_64
 from agentlab.llm.tracking import cost_tracker_decorator
 
 
@@ -158,24 +155,6 @@ def set_reproducibility_mode(self):
             self.subagent_args.set_reproducibility_mode()
 
 
-def overlay_action(obs, action):
-    """Overlays actions on screenshot in-place"""
-    act_img = copy.deepcopy(obs["screenshot"])
-    act_img = Image.fromarray(act_img)
-    overlay_utils.annotate_action(act_img, action, properties=obs["extra_element_properties"])
-    return img_to_base_64(act_img)
-
-
-def img_to_base_64(image: Image.Image | np.ndarray) -> str:
-    """Converts a PIL Image or NumPy array to a base64-encoded string."""
-    if isinstance(image, np.ndarray):
-        image = Image.fromarray(image)
-    buffer = io.BytesIO()
-    image.save(buffer, format="PNG")
-    b64_str = base64.b64encode(buffer.getvalue()).decode("utf-8")
-    return b64_str
-
-
 def get_base_human_in_the_loop_genericagent(llm_config):
     from agentlab.agents.generic_agent.tmlr_config import BASE_FLAGS
     from agentlab.agents.hitl_agent.hitl_agent import HumanInTheLoopAgentArgs

src/agentlab/llm/llm_utils.py

@@ -727,6 +727,16 @@ def image_to_png_base64_url(image: np.ndarray | Image.Image):
     return f"data:image/png;base64,{image_base64}"
 
 
+def img_to_base_64(image: Image.Image | np.ndarray) -> str:
+    """Converts a PIL Image or NumPy array to a base64-encoded string."""
+    if isinstance(image, np.ndarray):
+        image = Image.fromarray(image)
+    buffer = io.BytesIO()
+    image.save(buffer, format="PNG")
+    b64_str = base64.b64encode(buffer.getvalue()).decode("utf-8")
+    return b64_str
+
+
 class BaseMessage(dict):
     def __init__(self, role: str, content: Union[str, list[dict]], **kwargs):
         allowed_attrs = {"log_probs"}