91e75e620b
CI: cua-driver distro-compat matrix / debian:12 (glibc 2.36) (push) Has been cancelled
CI: SPDX Headers / Check SPDX headers (warn-only) (push) Has been cancelled
CD: Docs MCP Server / build (linux/amd64) (push) Has been cancelled
CD: Docs MCP Server / build (linux/arm64) (push) Has been cancelled
CD: Docs MCP Server / merge (push) Has been cancelled
CI: cua-driver distro-compat matrix / Resolve release version (push) Has been cancelled
CI: cua-driver distro-compat matrix / fedora:41 (glibc 2.40) (push) Has been cancelled
CI: cua-driver distro-compat matrix / rockylinux:9 (glibc 2.34) (push) Has been cancelled
CI: cua-driver distro-compat matrix / ubuntu:22.04 (glibc 2.35) (push) Has been cancelled
CI: cua-driver distro-compat matrix / ubuntu:24.04 (glibc 2.39) (push) Has been cancelled
CI: cua-driver distro-compat matrix / Distro compat summary (push) Has been cancelled
CI: Rust Linux unit / Rust Linux unit and compile (push) Has been cancelled
CI: Rust Windows unit / Rust Windows unit and compile (push) Has been cancelled
CI: Nix Linux Rust source / Nix / compositor build (push) Has been cancelled
CI: Nix Linux Rust source / Nix / driver package (push) Has been cancelled
CI: Nix Linux Rust source / Nix / Rust unit tests (push) Has been cancelled
494 lines
18 KiB
Python
494 lines
18 KiB
Python
"""
|
|
Moondream3+ composed-grounded agent loop implementation.
|
|
Grounding is handled by a local Moondream3 preview model via Transformers.
|
|
Thinking is delegated to the trailing LLM in the composed model string: "moondream3+<thinking_model>".
|
|
|
|
Differences from composed_grounded:
|
|
- Provides a singleton Moondream3 client outside the class.
|
|
- predict_click uses model.point(image, instruction, settings={"max_objects": 1}) and returns pixel coordinates.
|
|
- If the last image was a screenshot (or we take one), run model.detect(image, "all form ui") to get bboxes, then
|
|
run model.caption on each cropped bbox to label it. Overlay labels on the screenshot and emit via _on_screenshot.
|
|
- Add a user message listing all detected form UI names so the thinker can reference them.
|
|
- If the thinking model doesn't support vision, filter out image content before calling litellm.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import base64
|
|
import io
|
|
import uuid
|
|
from typing import Any, Dict, List, Optional, Tuple
|
|
|
|
import litellm
|
|
from PIL import Image, ImageDraw, ImageFont
|
|
|
|
from ..decorators import register_agent
|
|
from ..loops.base import AsyncAgentConfig
|
|
from ..responses import (
|
|
convert_completion_messages_to_responses_items,
|
|
convert_computer_calls_desc2xy,
|
|
convert_computer_calls_xy2desc,
|
|
convert_responses_items_to_completion_messages,
|
|
get_all_element_descriptions,
|
|
)
|
|
from ..types import AgentCapability
|
|
|
|
_MOONDREAM_SINGLETON = None
|
|
|
|
|
|
def get_moondream_model() -> Any:
|
|
"""Get a singleton instance of the Moondream3 preview model."""
|
|
global _MOONDREAM_SINGLETON
|
|
if _MOONDREAM_SINGLETON is None:
|
|
try:
|
|
import torch
|
|
from transformers import AutoModelForCausalLM
|
|
|
|
_MOONDREAM_SINGLETON = AutoModelForCausalLM.from_pretrained(
|
|
"moondream/moondream3-preview",
|
|
trust_remote_code=True,
|
|
torch_dtype=torch.bfloat16,
|
|
device_map="cuda",
|
|
)
|
|
except ImportError as e:
|
|
raise RuntimeError(
|
|
"moondream3 requires torch and transformers. Install with: pip install cua-agent[moondream3]"
|
|
) from e
|
|
return _MOONDREAM_SINGLETON
|
|
|
|
|
|
def _decode_image_b64(image_b64: str) -> Image.Image:
|
|
data = base64.b64decode(image_b64)
|
|
return Image.open(io.BytesIO(data)).convert("RGB")
|
|
|
|
|
|
def _image_to_b64(img: Image.Image) -> str:
|
|
buf = io.BytesIO()
|
|
img.save(buf, format="PNG")
|
|
return base64.b64encode(buf.getvalue()).decode("utf-8")
|
|
|
|
|
|
def _supports_vision(model: str) -> bool:
|
|
"""Heuristic vision support detection for thinking model."""
|
|
m = model.lower()
|
|
vision_markers = [
|
|
"gpt-4o",
|
|
"gpt-4.1",
|
|
"o1",
|
|
"o3",
|
|
"claude-3",
|
|
"claude-3.5",
|
|
"sonnet",
|
|
"haiku",
|
|
"opus",
|
|
"gemini-1.5",
|
|
"llava",
|
|
]
|
|
return any(v in m for v in vision_markers)
|
|
|
|
|
|
def _filter_images_from_completion_messages(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
|
filtered: List[Dict[str, Any]] = []
|
|
for msg in messages:
|
|
msg_copy = {**msg}
|
|
content = msg_copy.get("content")
|
|
if isinstance(content, list):
|
|
msg_copy["content"] = [c for c in content if c.get("type") != "image_url"]
|
|
filtered.append(msg_copy)
|
|
return filtered
|
|
|
|
|
|
def _annotate_detect_and_label_ui(base_img: Image.Image, model_md) -> Tuple[str, List[str]]:
|
|
"""Detect UI elements with Moondream, caption each, draw labels with backgrounds.
|
|
|
|
Args:
|
|
base_img: PIL image of the screenshot (RGB or RGBA). Will be copied/converted internally.
|
|
model_md: Moondream model instance with .detect() and .query() methods.
|
|
|
|
Returns:
|
|
A tuple of (annotated_image_base64_png, detected_names)
|
|
"""
|
|
# Ensure RGBA for semi-transparent fills
|
|
if base_img.mode != "RGBA":
|
|
base_img = base_img.convert("RGBA")
|
|
W, H = base_img.width, base_img.height
|
|
|
|
# Detect objects
|
|
try:
|
|
detect_result = model_md.detect(base_img, "all ui elements")
|
|
objects = detect_result.get("objects", []) if isinstance(detect_result, dict) else []
|
|
except Exception:
|
|
objects = []
|
|
|
|
draw = ImageDraw.Draw(base_img)
|
|
try:
|
|
font = ImageFont.load_default()
|
|
except Exception:
|
|
font = None
|
|
|
|
detected_names: List[str] = []
|
|
|
|
for i, obj in enumerate(objects):
|
|
try:
|
|
# Clamp normalized coords and crop
|
|
x_min = max(0.0, min(1.0, float(obj.get("x_min", 0.0))))
|
|
y_min = max(0.0, min(1.0, float(obj.get("y_min", 0.0))))
|
|
x_max = max(0.0, min(1.0, float(obj.get("x_max", 0.0))))
|
|
y_max = max(0.0, min(1.0, float(obj.get("y_max", 0.0))))
|
|
left, top, right, bottom = (
|
|
int(x_min * W),
|
|
int(y_min * H),
|
|
int(x_max * W),
|
|
int(y_max * H),
|
|
)
|
|
left, top = max(0, left), max(0, top)
|
|
right, bottom = min(W - 1, right), min(H - 1, bottom)
|
|
crop = base_img.crop((left, top, right, bottom))
|
|
|
|
# Prompted short caption
|
|
try:
|
|
result = model_md.query(crop, "Caption this UI element in few words.")
|
|
caption_text = (result or {}).get("answer", "")
|
|
except Exception:
|
|
caption_text = ""
|
|
|
|
name = (caption_text or "").strip() or f"element_{i+1}"
|
|
detected_names.append(name)
|
|
|
|
# Draw bbox
|
|
draw.rectangle([left, top, right, bottom], outline=(255, 215, 0, 255), width=2)
|
|
|
|
# Label background with padding and rounded corners
|
|
label = f"{i+1}. {name}"
|
|
padding = 3
|
|
if font:
|
|
text_bbox = draw.textbbox((0, 0), label, font=font)
|
|
else:
|
|
text_bbox = draw.textbbox((0, 0), label)
|
|
text_w = text_bbox[2] - text_bbox[0]
|
|
text_h = text_bbox[3] - text_bbox[1]
|
|
|
|
tx = left + 3
|
|
ty = top - (text_h + 2 * padding + 4)
|
|
if ty < 0:
|
|
ty = top + 3
|
|
|
|
bg_left = tx - padding
|
|
bg_top = ty - padding
|
|
bg_right = tx + text_w + padding
|
|
bg_bottom = ty + text_h + padding
|
|
try:
|
|
draw.rounded_rectangle(
|
|
[bg_left, bg_top, bg_right, bg_bottom],
|
|
radius=4,
|
|
fill=(0, 0, 0, 160),
|
|
outline=(255, 215, 0, 200),
|
|
width=1,
|
|
)
|
|
except Exception:
|
|
draw.rectangle(
|
|
[bg_left, bg_top, bg_right, bg_bottom],
|
|
fill=(0, 0, 0, 160),
|
|
outline=(255, 215, 0, 200),
|
|
width=1,
|
|
)
|
|
|
|
text_fill = (255, 255, 255, 255)
|
|
if font:
|
|
draw.text((tx, ty), label, fill=text_fill, font=font)
|
|
else:
|
|
draw.text((tx, ty), label, fill=text_fill)
|
|
except Exception:
|
|
continue
|
|
|
|
# Encode PNG base64
|
|
annotated = base_img
|
|
if annotated.mode not in ("RGBA", "RGB"):
|
|
annotated = annotated.convert("RGBA")
|
|
annotated_b64 = _image_to_b64(annotated)
|
|
return annotated_b64, detected_names
|
|
|
|
|
|
GROUNDED_COMPUTER_TOOL_SCHEMA = {
|
|
"type": "function",
|
|
"function": {
|
|
"name": "computer",
|
|
"description": (
|
|
"Control a computer by taking screenshots and interacting with UI elements. "
|
|
"The screenshot action will include a list of detected form UI element names when available. "
|
|
"Use element descriptions to locate and interact with UI elements on the screen."
|
|
),
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"action": {
|
|
"type": "string",
|
|
"enum": [
|
|
"screenshot",
|
|
"click",
|
|
"double_click",
|
|
"drag",
|
|
"type",
|
|
"keypress",
|
|
"scroll",
|
|
"move",
|
|
"wait",
|
|
"get_current_url",
|
|
"get_dimensions",
|
|
"get_environment",
|
|
],
|
|
"description": "The action to perform (required for all actions)",
|
|
},
|
|
"element_description": {
|
|
"type": "string",
|
|
"description": "Description of the element to interact with (required for click/double_click/move/scroll)",
|
|
},
|
|
"start_element_description": {
|
|
"type": "string",
|
|
"description": "Description of the element to start dragging from (required for drag)",
|
|
},
|
|
"end_element_description": {
|
|
"type": "string",
|
|
"description": "Description of the element to drag to (required for drag)",
|
|
},
|
|
"text": {
|
|
"type": "string",
|
|
"description": "The text to type (required for type)",
|
|
},
|
|
"keys": {
|
|
"type": "array",
|
|
"items": {"type": "string"},
|
|
"description": "Key(s) to press (required for keypress)",
|
|
},
|
|
"button": {
|
|
"type": "string",
|
|
"enum": ["left", "right", "wheel", "back", "forward"],
|
|
"description": "The mouse button to use for click/double_click",
|
|
},
|
|
"scroll_x": {
|
|
"type": "integer",
|
|
"description": "Horizontal scroll amount (required for scroll)",
|
|
},
|
|
"scroll_y": {
|
|
"type": "integer",
|
|
"description": "Vertical scroll amount (required for scroll)",
|
|
},
|
|
},
|
|
"required": ["action"],
|
|
},
|
|
},
|
|
}
|
|
|
|
|
|
@register_agent(r"moondream3\+.*", priority=2)
|
|
class Moondream3PlusConfig(AsyncAgentConfig):
|
|
def __init__(self):
|
|
self.desc2xy: Dict[str, Tuple[float, float]] = {}
|
|
|
|
async def predict_step(
|
|
self,
|
|
messages: List[Dict[str, Any]],
|
|
model: str,
|
|
tools: Optional[List[Dict[str, Any]]] = None,
|
|
max_retries: Optional[int] = None,
|
|
stream: bool = False,
|
|
computer_handler=None,
|
|
use_prompt_caching: Optional[bool] = False,
|
|
_on_api_start=None,
|
|
_on_api_end=None,
|
|
_on_usage=None,
|
|
_on_screenshot=None,
|
|
**kwargs,
|
|
) -> Dict[str, Any]:
|
|
# Parse composed model: moondream3+<thinking_model>
|
|
if "+" not in model:
|
|
raise ValueError(f"Composed model must be 'moondream3+<thinking_model>', got: {model}")
|
|
_, thinking_model = model.split("+", 1)
|
|
|
|
pre_output_items: List[Dict[str, Any]] = []
|
|
|
|
# Acquire last screenshot; if missing, take one
|
|
last_image_b64: Optional[str] = None
|
|
for message in reversed(messages):
|
|
if (
|
|
isinstance(message, dict)
|
|
and message.get("type") == "computer_call_output"
|
|
and isinstance(message.get("output"), dict)
|
|
and message["output"].get("type") == "input_image"
|
|
):
|
|
image_url = message["output"].get("image_url", "")
|
|
if image_url.startswith("data:image/png;base64,"):
|
|
last_image_b64 = image_url.split(",", 1)[1]
|
|
break
|
|
|
|
if last_image_b64 is None and computer_handler is not None:
|
|
# Take a screenshot
|
|
screenshot_b64 = await computer_handler.screenshot() # type: ignore
|
|
if screenshot_b64:
|
|
call_id = uuid.uuid4().hex
|
|
pre_output_items += [
|
|
{
|
|
"type": "message",
|
|
"role": "assistant",
|
|
"content": [
|
|
{
|
|
"type": "output_text",
|
|
"text": "Taking a screenshot to analyze the current screen.",
|
|
}
|
|
],
|
|
},
|
|
{
|
|
"type": "computer_call",
|
|
"call_id": call_id,
|
|
"status": "completed",
|
|
"action": {"type": "screenshot"},
|
|
},
|
|
{
|
|
"type": "computer_call_output",
|
|
"call_id": call_id,
|
|
"output": {
|
|
"type": "input_image",
|
|
"image_url": f"data:image/png;base64,{screenshot_b64}",
|
|
},
|
|
},
|
|
]
|
|
last_image_b64 = screenshot_b64
|
|
if _on_screenshot:
|
|
await _on_screenshot(screenshot_b64)
|
|
|
|
# If we have a last screenshot, run Moondream detection and labeling
|
|
detected_names: List[str] = []
|
|
if last_image_b64 is not None:
|
|
base_img = _decode_image_b64(last_image_b64)
|
|
model_md = get_moondream_model()
|
|
annotated_b64, detected_names = _annotate_detect_and_label_ui(base_img, model_md)
|
|
if _on_screenshot:
|
|
await _on_screenshot(annotated_b64, "annotated_form_ui")
|
|
|
|
# Also push a user message listing all detected names
|
|
if detected_names:
|
|
names_text = "\n".join(f"- {n}" for n in detected_names)
|
|
pre_output_items.append(
|
|
{
|
|
"type": "message",
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "input_text", "text": "Detected form UI elements on screen:"},
|
|
{"type": "input_text", "text": names_text},
|
|
{
|
|
"type": "input_text",
|
|
"text": "Please continue with the next action needed to perform your task.",
|
|
},
|
|
],
|
|
}
|
|
)
|
|
|
|
tool_schemas = []
|
|
for schema in tools or []:
|
|
if schema.get("type") == "computer":
|
|
tool_schemas.append(GROUNDED_COMPUTER_TOOL_SCHEMA)
|
|
else:
|
|
tool_schemas.append(schema)
|
|
|
|
# Step 1: Convert computer calls from xy to descriptions
|
|
input_messages = messages + pre_output_items
|
|
messages_with_descriptions = convert_computer_calls_xy2desc(input_messages, self.desc2xy)
|
|
|
|
# Step 2: Convert responses items to completion messages
|
|
completion_messages = convert_responses_items_to_completion_messages(
|
|
messages_with_descriptions,
|
|
allow_images_in_tool_results=False,
|
|
)
|
|
|
|
# Optionally filter images if model lacks vision
|
|
if not _supports_vision(thinking_model):
|
|
completion_messages = _filter_images_from_completion_messages(completion_messages)
|
|
|
|
# Step 3: Call thinking model with litellm.acompletion
|
|
api_kwargs = {
|
|
"model": thinking_model,
|
|
"messages": completion_messages,
|
|
"tools": tool_schemas,
|
|
"max_retries": max_retries,
|
|
"stream": stream,
|
|
**kwargs,
|
|
}
|
|
if use_prompt_caching:
|
|
api_kwargs["use_prompt_caching"] = use_prompt_caching
|
|
|
|
if _on_api_start:
|
|
await _on_api_start(api_kwargs)
|
|
|
|
response = await litellm.acompletion(**api_kwargs)
|
|
|
|
if _on_api_end:
|
|
await _on_api_end(api_kwargs, response)
|
|
|
|
usage = {
|
|
**response.usage.model_dump(), # type: ignore
|
|
"response_cost": response._hidden_params.get("response_cost", 0.0),
|
|
}
|
|
if _on_usage:
|
|
await _on_usage(usage)
|
|
|
|
# Step 4: Convert completion messages back to responses items format
|
|
response_dict = response.model_dump() # type: ignore
|
|
choice_messages = [choice["message"] for choice in response_dict["choices"]]
|
|
thinking_output_items: List[Dict[str, Any]] = []
|
|
for choice_message in choice_messages:
|
|
thinking_output_items.extend(
|
|
convert_completion_messages_to_responses_items([choice_message])
|
|
)
|
|
|
|
# Step 5: Use Moondream to get coordinates for each description
|
|
element_descriptions = get_all_element_descriptions(thinking_output_items)
|
|
if element_descriptions and last_image_b64:
|
|
for desc in element_descriptions:
|
|
for _ in range(3): # try 3 times
|
|
coords = await self.predict_click(
|
|
model=model,
|
|
image_b64=last_image_b64,
|
|
instruction=desc,
|
|
)
|
|
if coords:
|
|
self.desc2xy[desc] = coords
|
|
break
|
|
|
|
# Step 6: Convert computer calls from descriptions back to xy coordinates
|
|
final_output_items = convert_computer_calls_desc2xy(thinking_output_items, self.desc2xy)
|
|
|
|
# Step 7: Return output and usage
|
|
return {"output": pre_output_items + final_output_items, "usage": usage}
|
|
|
|
async def predict_click(
|
|
self,
|
|
model: str,
|
|
image_b64: str,
|
|
instruction: str,
|
|
**kwargs,
|
|
) -> Optional[Tuple[float, float]]:
|
|
"""Predict click coordinates using Moondream3's point API.
|
|
|
|
Returns pixel coordinates (x, y) as floats.
|
|
"""
|
|
img = _decode_image_b64(image_b64)
|
|
W, H = img.width, img.height
|
|
model_md = get_moondream_model()
|
|
try:
|
|
result = model_md.point(img, instruction, settings={"max_objects": 1})
|
|
except Exception:
|
|
return None
|
|
|
|
try:
|
|
pt = (result or {}).get("points", [])[0]
|
|
x_norm = float(pt.get("x", 0.0))
|
|
y_norm = float(pt.get("y", 0.0))
|
|
x_px = max(0.0, min(float(W - 1), x_norm * W))
|
|
y_px = max(0.0, min(float(H - 1), y_norm * H))
|
|
return (x_px, y_px)
|
|
except Exception:
|
|
return None
|
|
|
|
def get_capabilities(self) -> List[AgentCapability]:
|
|
return ["click", "step"]
|