Tests / Test apps.device-host-agent.tests.test_mcp_token.test_load_or_create_concurrent_calls_do_not_corrupt failed
195 lines
8.1 KiB
Python
195 lines
8.1 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
from dataclasses import dataclass
|
|
from typing import Any, Callable
|
|
|
|
from runtime.planner_config import PlannerConfig
|
|
from runtime.tool_specs import ACTION_TOOL_SPECS, ToolSpec
|
|
|
|
READ_TOOL_SPECS = [
|
|
ToolSpec("take_screenshot", "Capture the current device screen.", {"type": "object", "properties": {"device_id": {"type": "string"}}}),
|
|
ToolSpec("describe_screen", "Inspect the current screen and UI.", {"type": "object", "properties": {"device_id": {"type": "string"}}}),
|
|
ToolSpec("find_text", "Find visible text on the current screen.", {"type": "object", "required": ["query"], "properties": {"query": {"type": "string"}, "device_id": {"type": "string"}}}),
|
|
ToolSpec("get_ui_tree", "Read the current accessibility/UI tree.", {"type": "object", "properties": {"device_id": {"type": "string"}, "include_app_info": {"type": "boolean"}}}),
|
|
ToolSpec("list_devices", "List configured devices.", {"type": "object", "properties": {}}),
|
|
ToolSpec("device_status", "Read one device status.", {"type": "object", "required": ["device_id"], "properties": {"device_id": {"type": "string"}}}),
|
|
]
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ChatResult:
|
|
content: str
|
|
tool_calls: int
|
|
|
|
|
|
class ConversationAgent:
|
|
"""Small OpenAI-compatible agent loop backed by Host Agent tools."""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
config: PlannerConfig,
|
|
tools: dict[str, Callable[..., Any]],
|
|
max_rounds: int = 8,
|
|
event_logger: Callable[[dict[str, Any]], None] | None = None,
|
|
) -> None:
|
|
self.config = config
|
|
self.tools = tools
|
|
self.max_rounds = max_rounds
|
|
self.event_logger = event_logger
|
|
self._client: Any | None = None
|
|
|
|
def chat(self, messages: list[dict[str, Any]]) -> ChatResult:
|
|
if self.config.provider not in {"openai", "openai_compatible"}:
|
|
raise RuntimeError(
|
|
"conversation chat currently requires an OpenAI-compatible provider"
|
|
)
|
|
client = self._get_client()
|
|
history: list[dict[str, Any]] = [
|
|
{
|
|
"role": "system",
|
|
"content": (
|
|
"You are a mobile device assistant. Reply naturally when no "
|
|
"device action is needed. Call a tool when the user asks you "
|
|
"to inspect or operate a device. Never claim an action was "
|
|
"completed unless the tool result confirms it."
|
|
),
|
|
},
|
|
*[_normalize_message(message) for message in messages],
|
|
]
|
|
calls = 0
|
|
for _ in range(self.max_rounds):
|
|
response = client.chat.completions.create(
|
|
model=self.config.resolved_model(),
|
|
messages=history,
|
|
tools=[_openai_tool(spec) for spec in [*ACTION_TOOL_SPECS, *READ_TOOL_SPECS]],
|
|
tool_choice="auto",
|
|
parallel_tool_calls=False,
|
|
timeout=self.config.timeout,
|
|
)
|
|
message = response.choices[0].message
|
|
text = getattr(message, "content", None)
|
|
thinking = getattr(message, "reasoning_content", None) or getattr(message, "reasoning", None)
|
|
tool_calls = getattr(message, "tool_calls", None) or []
|
|
self._log({"type": "llm_response", "content": text, "thinking": thinking,
|
|
"tool_calls": [{"id": c.id, "name": c.function.name, "arguments": c.function.arguments} for c in tool_calls]})
|
|
if not tool_calls:
|
|
self._log({"type": "final_reply", "content": str(text or "")})
|
|
return ChatResult(content=str(text or ""), tool_calls=calls)
|
|
history.append(_message_dict(message))
|
|
for tool_call in tool_calls:
|
|
name = tool_call.function.name
|
|
arguments: dict[str, Any] = {}
|
|
try:
|
|
decoded_arguments = json.loads(tool_call.function.arguments or "{}")
|
|
if not isinstance(decoded_arguments, dict):
|
|
raise ValueError("tool arguments must be an object")
|
|
arguments = decoded_arguments
|
|
arguments.pop("purpose", None)
|
|
arguments.pop("expected_outcome", None)
|
|
result = self.tools[name](**arguments)
|
|
except Exception as exc:
|
|
result = {"ok": False, "error": str(exc)}
|
|
calls += 1
|
|
self._log({"type": "tool_result", "tool_call_id": tool_call.id,
|
|
"tool_name": name, "arguments": arguments, "result": result})
|
|
history.append(
|
|
{
|
|
"role": "tool",
|
|
"tool_call_id": tool_call.id,
|
|
"content": json.dumps(result, ensure_ascii=True, default=str),
|
|
}
|
|
)
|
|
raise RuntimeError("conversation exceeded the maximum tool-call rounds")
|
|
|
|
def _log(self, event: dict[str, Any]) -> None:
|
|
if self.event_logger is not None:
|
|
try:
|
|
self.event_logger(event)
|
|
except Exception:
|
|
pass
|
|
|
|
def _get_client(self) -> Any:
|
|
if self._client is None:
|
|
from openai import OpenAI
|
|
|
|
kwargs: dict[str, Any] = {}
|
|
if self.config.api_key:
|
|
kwargs["api_key"] = self.config.api_key
|
|
if self.config.base_url:
|
|
kwargs["base_url"] = self.config.base_url
|
|
self._client = OpenAI(**kwargs)
|
|
return self._client
|
|
|
|
|
|
def _openai_tool(spec: ToolSpec) -> dict[str, Any]:
|
|
parameters = dict(spec.parameters)
|
|
properties = dict(parameters.get("properties") or {})
|
|
properties.setdefault(
|
|
"device_id",
|
|
{"type": "string", "description": "Target device ID when needed."},
|
|
)
|
|
parameters["properties"] = properties
|
|
return {
|
|
"type": "function",
|
|
"function": {
|
|
"name": spec.name,
|
|
"description": spec.description,
|
|
"parameters": parameters,
|
|
},
|
|
}
|
|
|
|
|
|
def _message_dict(message: Any) -> dict[str, Any]:
|
|
result: dict[str, Any] = {"role": "assistant"}
|
|
content = getattr(message, "content", None)
|
|
if content is not None:
|
|
result["content"] = content
|
|
calls = getattr(message, "tool_calls", None) or []
|
|
result["tool_calls"] = [
|
|
{
|
|
"id": call.id,
|
|
"type": "function",
|
|
"function": {
|
|
"name": call.function.name,
|
|
"arguments": call.function.arguments,
|
|
},
|
|
}
|
|
for call in calls
|
|
]
|
|
return result
|
|
|
|
|
|
def _normalize_message(message: dict[str, Any]) -> dict[str, Any]:
|
|
"""Accept OpenAI text/image content blocks and a compact image_base64 form."""
|
|
role = message.get("role")
|
|
content = message.get("content")
|
|
if isinstance(content, str):
|
|
return {"role": role, "content": content}
|
|
if isinstance(content, list):
|
|
blocks: list[dict[str, Any]] = []
|
|
for block in content:
|
|
if not isinstance(block, dict):
|
|
continue
|
|
if block.get("type") == "text" and isinstance(block.get("text"), str):
|
|
blocks.append({"type": "text", "text": block["text"]})
|
|
elif block.get("type") == "image_url":
|
|
image_url = block.get("image_url")
|
|
if isinstance(image_url, str):
|
|
image_url = {"url": image_url}
|
|
if isinstance(image_url, dict) and isinstance(image_url.get("url"), str):
|
|
blocks.append({"type": "image_url", "image_url": {"url": image_url["url"]}})
|
|
if blocks:
|
|
return {"role": role, "content": blocks}
|
|
image = message.get("image_base64")
|
|
if isinstance(image, str) and image:
|
|
mime = str(message.get("mime_type") or "image/png")
|
|
text = message.get("text")
|
|
blocks = []
|
|
if isinstance(text, str) and text:
|
|
blocks.append({"type": "text", "text": text})
|
|
blocks.append({"type": "image_url", "image_url": {"url": f"data:{mime};base64,{image}"}})
|
|
return {"role": role, "content": blocks}
|
|
raise ValueError("message content must be text, image blocks, or image_base64")
|