autoglm-gui 1.4.0__py3-none-any.whl → 1.4.1__py3-none-any.whl

mai_agent/mai_grounding_agent.py

@@ -0,0 +1,263 @@
+# Copyright (c) 2025, Alibaba Cloud and its affiliates;
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+MAI Grounding Agent - A GUI grounding agent for locating UI elements.
+
+This module provides the MAIGroundingAgent class that uses vision-language models
+to locate UI elements based on natural language instructions.
+"""
+
+import json
+import re
+from io import BytesIO
+from typing import Any, Dict, Optional, Tuple, Union
+
+from openai import OpenAI
+from PIL import Image
+
+from prompt import MAI_MOBILE_SYS_PROMPT_GROUNDING
+from utils import pil_to_base64
+
+
+# Constants
+SCALE_FACTOR = 999
+
+
+def parse_grounding_response(text: str) -> Dict[str, Any]:
+    """
+    Parse model output text containing grounding_think and answer tags.
+
+    Args:
+        text: Raw model output containing <grounding_think> and <answer> tags.
+
+    Returns:
+        Dictionary with keys:
+            - "thinking": The model's reasoning process
+            - "coordinate": Normalized [x, y] coordinate
+
+    Raises:
+        ValueError: If parsing fails or JSON is invalid.
+    """
+    text = text.strip()
+
+    result: Dict[str, Any] = {
+        "thinking": None,
+        "coordinate": None,
+    }
+
+    # Extract thinking content
+    think_pattern = r"<grounding_think>(.*?)</grounding_think>"
+    think_match = re.search(think_pattern, text, re.DOTALL)
+    if think_match:
+        result["thinking"] = think_match.group(1).strip()
+
+    # Extract answer content
+    answer_pattern = r"<answer>(.*?)</answer>"
+    answer_match = re.search(answer_pattern, text, re.DOTALL)
+    if answer_match:
+        answer_text = answer_match.group(1).strip()
+        try:
+            answer_json = json.loads(answer_text)
+            coordinates = answer_json.get("coordinate", [])
+            if len(coordinates) == 2:
+                # Normalize coordinates from SCALE_FACTOR range to [0, 1]
+                point_x = coordinates[0] / SCALE_FACTOR
+                point_y = coordinates[1] / SCALE_FACTOR
+                result["coordinate"] = [point_x, point_y]
+            else:
+                raise ValueError(
+                    f"Invalid coordinate format: expected 2 values, got {len(coordinates)}"
+                )
+        except json.JSONDecodeError as e:
+            raise ValueError(f"Invalid JSON in answer: {e}")
+
+    return result
+
+
+class MAIGroundingAgent:
+    """
+    GUI grounding agent using vision-language models.
+
+    This agent processes a screenshot and natural language instruction to
+    locate a specific UI element and return its coordinates.
+
+    Attributes:
+        llm_base_url: Base URL for the LLM API endpoint.
+        model_name: Name of the model to use for predictions.
+        runtime_conf: Configuration dictionary for runtime parameters.
+    """
+
+    def __init__(
+        self,
+        llm_base_url: str,
+        model_name: str,
+        runtime_conf: Optional[Dict[str, Any]] = None,
+    ):
+        """
+        Initialize the MAIGroundingAgent.
+
+        Args:
+            llm_base_url: Base URL for the LLM API endpoint.
+            model_name: Name of the model to use.
+            runtime_conf: Optional configuration dictionary with keys:
+                - max_pixels: Maximum pixels for image processing
+                - min_pixels: Minimum pixels for image processing
+                - temperature: Sampling temperature (default: 0.0)
+                - top_k: Top-k sampling parameter (default: -1)
+                - top_p: Top-p sampling parameter (default: 1.0)
+                - max_tokens: Maximum tokens in response (default: 2048)
+        """
+        # Set default configuration
+        default_conf = {
+            "temperature": 0.0,
+            "top_k": -1,
+            "top_p": 1.0,
+            "max_tokens": 2048,
+        }
+        self.runtime_conf = {**default_conf, **(runtime_conf or {})}
+
+        self.llm_base_url = llm_base_url
+        self.model_name = model_name
+        self.llm = OpenAI(
+            base_url=self.llm_base_url,
+            api_key="empty",
+        )
+
+        # Extract frequently used config values
+        self.temperature = self.runtime_conf["temperature"]
+        self.top_k = self.runtime_conf["top_k"]
+        self.top_p = self.runtime_conf["top_p"]
+        self.max_tokens = self.runtime_conf["max_tokens"]
+
+    @property
+    def system_prompt(self) -> str:
+        """Return the system prompt for grounding tasks."""
+        return MAI_MOBILE_SYS_PROMPT_GROUNDING
+
+    def _build_messages(
+        self,
+        instruction: str,
+        image: Image.Image,
+    ) -> list:
+        """
+        Build the message list for the LLM API call.
+
+        Args:
+            instruction: Grounding instruction from user.
+            image: PIL Image of the screenshot.
+            magic_prompt: Whether to use the magic prompt format.
+
+        Returns:
+            List of message dictionaries for the API.
+        """
+        encoded_string = pil_to_base64(image)
+
+        messages = [
+            {
+                "role": "system",
+                "content": [
+                    {
+                        "type": "text",
+                        "text": self.system_prompt,
+                    }
+                ],
+            }
+        ]
+
+        messages.append(
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "text",
+                        "text": instruction + "\n",
+                    },
+                    {
+                        "type": "image_url",
+                        "image_url": {"url": f"data:image/png;base64,{encoded_string}"},
+                    },
+                ],
+            }
+        )
+
+        return messages
+
+    def predict(
+        self,
+        instruction: str,
+        image: Union[Image.Image, bytes],
+        **kwargs: Any,
+    ) -> Tuple[str, Dict[str, Any]]:
+        """
+        Predict the coordinate of the UI element based on the instruction.
+
+        Args:
+            instruction: Grounding instruction describing the UI element to locate.
+            image: PIL Image or bytes of the screenshot.
+            **kwargs: Additional arguments (unused).
+
+        Returns:
+            Tuple of (prediction_text, result_dict) where:
+                - prediction_text: Raw model response or error message
+                - result_dict: Dictionary containing:
+                    - "thinking": Model's reasoning process
+                    - "coordinate": Normalized [x, y] coordinate
+        """
+        # Convert bytes to PIL Image if necessary
+        if isinstance(image, bytes):
+            image = Image.open(BytesIO(image))
+
+        if image.mode != "RGB":
+            image = image.convert("RGB")
+
+        # Build messages
+        messages = self._build_messages(instruction, image)
+
+        # Make API call with retry logic
+        max_retries = 3
+        prediction = None
+        result = None
+
+        for attempt in range(max_retries):
+            try:
+                response = self.llm.chat.completions.create(
+                    model=self.model_name,
+                    messages=messages,
+                    max_tokens=self.max_tokens,
+                    temperature=self.temperature,
+                    top_p=self.top_p,
+                    frequency_penalty=0.0,
+                    presence_penalty=0.0,
+                    extra_body={"repetition_penalty": 1.0, "top_k": self.top_k},
+                    seed=42,
+                )
+                prediction = response.choices[0].message.content.strip()
+                print(f"Raw response:\n{prediction}")
+
+                # Parse response
+                result = parse_grounding_response(prediction)
+                print(f"Parsed result:\n{result}")
+                break
+
+            except Exception as e:
+                print(f"Error on attempt {attempt + 1}: {e}")
+                prediction = None
+                result = None
+
+        # Return error if all retries failed
+        if prediction is None or result is None:
+            print("Max retry attempts reached, returning error flag.")
+            return "llm client error", {"thinking": None, "coordinate": None}
+
+        return prediction, result