minitap-mobile-use 2.0.0__py3-none-any.whl → 2.0.1__py3-none-any.whl

minitap/mobile_use/agents/planner/planner.md

@@ -25,7 +25,7 @@ You work like an agile tech lead: defining the key milestones without locking in
 
 ### Output
 
-You must output a **list of strings**, each representing a clear subgoal.
+You must output a **list of subgoals (description + optional subgoal ID)**, each representing a clear subgoal.
 Each subgoal should be:
 
 - Focused on **realistic mobile interactions**
@@ -33,21 +33,23 @@ Each subgoal should be:
 - Sequential (later steps may depend on earlier ones)
 - Don't use loop-like formulation unless necessary (e.g. don't say "repeat this X times", instead reuse the same steps X times as subgoals)
 
+If you're replaning and need to keep a previous subgoal, you **must keep the same subgoal ID**.
+
 ### Examples
 
 #### **Initial Goal**: "Open WhatsApp and send 'I’m running late' to Alice"
 
 **Plan**:
 
-- Open the WhatsApp app
-- Locate or search for Alice
-- Open the conversation with Alice
-- Type the message "I’m running late"
-- Send the message
+- Open the WhatsApp app (ID: None -> will be generated as a UUID like bc3c362d-f498-4f1a-991e-4a2d1f8c1226)
+- Locate or search for Alice (ID: None)
+- Open the conversation with Alice (ID: None)
+- Type the message "I’m running late" (ID: None)
+- Send the message (ID: None)
 
 #### **Replanning Example**
 
-**Original Plan**: same as above
+**Original Plan**: same as above with IDs set
 **Agent Thoughts**:
 
 - Couldn’t find Alice in recent chats
@@ -56,9 +58,8 @@ Each subgoal should be:
 
 **New Plan**:
 
-- Unlock the phone if needed
-- Open WhatsApp
-- Tap the search bar
-- Search for "Alice"
-- Select the correct chat
-- Type and send "I’m running late"
+- Open WhatsApp (ID: bc3c362d-f498-4f1a-991e-4a2d1f8c1226)
+- Tap the search bar (ID: None)
+- Search for "Alice" (ID: None)
+- Select the correct chat (ID: None)
+- Type and send "I’m running late" (ID: None)

minitap/mobile_use/agents/planner/planner.py

@@ -1,4 +1,5 @@
 from pathlib import Path
+import uuid
 
 from jinja2 import Template
 from langchain_core.messages import HumanMessage, SystemMessage
@@ -47,7 +48,8 @@ class PlannerNode:
 
         subgoals_plan = [
             Subgoal(
-                description=subgoal,
+                id=subgoal.id or str(uuid.uuid4()),
+                description=subgoal.description,
                 status=SubgoalStatus.NOT_STARTED,
                 completion_reason=None,
             )
@@ -61,4 +63,5 @@ class PlannerNode:
             update={
                 "subgoal_plan": subgoals_plan,
             },
+            agent="planner",
         )

minitap/mobile_use/agents/planner/types.py

@@ -5,8 +5,13 @@ from pydantic import BaseModel
 from typing_extensions import Annotated
 
 
+class PlannerSubgoalOutput(BaseModel):
+    id: Annotated[Optional[str], "If not provided, it will be generated"] = None
+    description: str
+
+
 class PlannerOutput(BaseModel):
-    subgoals: list[str]
+    subgoals: list[PlannerSubgoalOutput]
 
 
 class SubgoalStatus(Enum):
@@ -17,6 +22,7 @@ class SubgoalStatus(Enum):
 
 
 class Subgoal(BaseModel):
+    id: Annotated[str, "Unique identifier of the subgoal"]
     description: Annotated[str, "Description of the subgoal"]
     completion_reason: Annotated[
         Optional[str], "Reason why the subgoal was completed (failure or success)"
@@ -35,7 +41,7 @@ class Subgoal(BaseModel):
             case SubgoalStatus.NOT_STARTED:
                 status_emoji = "(not started yet)"
 
-        output = f"- {self.description} : {status_emoji}."
+        output = f"- [ID:{self.id}]: {self.description} : {status_emoji}."
         if self.completion_reason:
             output += f" Completion reason: {self.completion_reason}"
         return output

minitap/mobile_use/agents/planner/utils.py

@@ -5,6 +5,10 @@ def get_current_subgoal(subgoals: list[Subgoal]) -> Subgoal | None:
     return next((s for s in subgoals if s.status == SubgoalStatus.PENDING), None)
 
 
+def get_subgoals_by_ids(subgoals: list[Subgoal], ids: list[str]) -> list[Subgoal]:
+    return [s for s in subgoals if s.id in ids]
+
+
 def get_next_subgoal(subgoals: list[Subgoal]) -> Subgoal | None:
     return next((s for s in subgoals if s.status == SubgoalStatus.NOT_STARTED), None)
 
@@ -21,6 +25,13 @@ def complete_current_subgoal(subgoals: list[Subgoal]) -> list[Subgoal]:
     return subgoals
 
 
+def complete_subgoals_by_ids(subgoals: list[Subgoal], ids: list[str]) -> list[Subgoal]:
+    for subgoal in subgoals:
+        if subgoal.id in ids:
+            subgoal.status = SubgoalStatus.SUCCESS
+    return subgoals
+
+
 def fail_current_subgoal(subgoals: list[Subgoal]) -> list[Subgoal]:
     current_subgoal = get_current_subgoal(subgoals)
     if not current_subgoal:

minitap/mobile_use/clients/device_hardware_client.py

@@ -12,6 +12,9 @@ class DeviceHardwareClient:
         url = urljoin(self.base_url, f"/api/{path.lstrip('/')}")
         return self.session.get(url, **kwargs)
 
+    def get_rich_hierarchy(self) -> list[dict]:
+        return self.get("last-view-hierarchy").json().get("children", [])
+
     def post(self, path: str, **kwargs):
         url = urljoin(self.base_url, f"/api/{path.lstrip('/')}")
         return self.session.post(url, **kwargs)

minitap/mobile_use/config.py

@@ -22,6 +22,8 @@ class Settings(BaseSettings):
     XAI_API_KEY: Optional[SecretStr] = None
     OPEN_ROUTER_API_KEY: Optional[SecretStr] = None
 
+    OPENAI_BASE_URL: Optional[str] = None
+
     DEVICE_SCREEN_API_BASE_URL: Optional[str] = None
     DEVICE_HARDWARE_BRIDGE_BASE_URL: Optional[str] = None
     ADB_HOST: Optional[str] = None

minitap/mobile_use/constants.py

@@ -1,2 +1,3 @@
 RECURSION_LIMIT = 400
 MAX_MESSAGES_IN_HISTORY = 25
+EXECUTOR_MESSAGES_KEY = "executor_messages"

minitap/mobile_use/controllers/mobile_command_controller.py

@@ -9,6 +9,7 @@ from requests import JSONDecodeError
 
 from minitap.mobile_use.clients.device_hardware_client import DeviceHardwareClient
 from minitap.mobile_use.clients.screen_api_client import ScreenApiClient
+from minitap.mobile_use.config import initialize_llm_config
 from minitap.mobile_use.context import DeviceContext, DevicePlatform, MobileUseContext
 from minitap.mobile_use.utils.errors import ControllerErrors
 from minitap.mobile_use.utils.logger import get_logger
@@ -331,12 +332,10 @@ def run_flow_with_wait_for_animation_to_end(
 
 
 if __name__ == "__main__":
-    # long press, erase
-    # input_text(text="test")
-    # erase_text()
     ctx = MobileUseContext(
+        llm_config=initialize_llm_config(),
         device=DeviceContext(
-            host_platform="LINUX",
+            host_platform="WINDOWS",
             mobile_platform=DevicePlatform.ANDROID,
             device_id="emulator-5554",
             device_width=1080,
@@ -347,7 +346,6 @@ if __name__ == "__main__":
     )
     screen_data = get_screen_data(ctx.screen_api_client)
     from minitap.mobile_use.graph.state import State
-    from minitap.mobile_use.tools.mobile.erase_text import get_erase_text_tool
 
     dummy_state = State(
         latest_ui_hierarchy=screen_data.elements,
@@ -358,20 +356,21 @@ if __name__ == "__main__":
         focused_app_info=None,
         device_date="",
         structured_decisions=None,
-        executor_retrigger=False,
-        executor_failed=False,
+        complete_subgoals_by_ids=[],
         executor_messages=[],
         cortex_last_thought="",
         agents_thoughts=[],
     )
 
-    # invoke erase_text tool
-    input_resource_id = "com.google.android.settings.intelligence:id/open_search_view_edit_text"
-    command_output: Command = get_erase_text_tool(ctx=ctx).invoke(
+    from minitap.mobile_use.tools.mobile.input_text import get_input_text_tool
+
+    input_resource_id = "com.google.android.apps.nexuslauncher:id/search_container_hotseat"
+    command_output: Command = get_input_text_tool(ctx=ctx).invoke(
         {
             "tool_call_id": uuid.uuid4().hex,
             "agent_thought": "",
-            "input_text_resource_id": input_resource_id,
+            "text_input_resource_id": input_resource_id,
+            "text": "Hello World",
             "state": dummy_state,
             "executor_metadata": None,
         }

minitap/mobile_use/graph/graph.py

@@ -6,13 +6,10 @@ from langchain_core.messages import (
 from langgraph.constants import END, START
 from langgraph.graph import StateGraph
 from langgraph.graph.state import CompiledStateGraph
-from langgraph.prebuilt import ToolNode
 from minitap.mobile_use.agents.contextor.contextor import ContextorNode
 from minitap.mobile_use.agents.cortex.cortex import CortexNode
 from minitap.mobile_use.agents.executor.executor import ExecutorNode
-from minitap.mobile_use.agents.executor.executor_context_cleaner import (
-    executor_context_cleaner_node,
-)
+from minitap.mobile_use.agents.executor.tool_node import ExecutorToolNode
 from minitap.mobile_use.agents.orchestrator.orchestrator import OrchestratorNode
 from minitap.mobile_use.agents.planner.planner import PlannerNode
 from minitap.mobile_use.agents.planner.utils import (
@@ -21,6 +18,7 @@ from minitap.mobile_use.agents.planner.utils import (
     one_of_them_is_failure,
 )
 from minitap.mobile_use.agents.summarizer.summarizer import SummarizerNode
+from minitap.mobile_use.constants import EXECUTOR_MESSAGES_KEY
 from minitap.mobile_use.context import MobileUseContext
 from minitap.mobile_use.graph.state import State
 from minitap.mobile_use.tools.index import EXECUTOR_WRAPPERS_TOOLS, get_tools_from_wrappers
@@ -53,7 +51,7 @@ def post_cortex_gate(
     state: State,
 ) -> Literal["continue", "end_subgoal"]:
     logger.info("Starting post_cortex_gate")
-    if not state.structured_decisions:
+    if len(state.complete_subgoals_by_ids) > 0:
         return "end_subgoal"
     return "continue"
 
@@ -62,7 +60,7 @@ def post_executor_gate(
     state: State,
 ) -> Literal["invoke_tools", "skip"]:
     logger.info("Starting post_executor_gate")
-    messages = state.messages
+    messages = state.executor_messages
     if not messages:
         return "skip"
     last_message = messages[-1]
@@ -77,17 +75,6 @@ def post_executor_gate(
     return "skip"
 
 
-def post_executor_tools_gate(
-    state: State,
-) -> Literal["continue", "failed", "done"]:
-    logger.info("Starting post_executor_tools_gate")
-    if state.executor_failed:
-        return "failed"
-    if state.executor_retrigger:
-        return "continue"
-    return "done"
-
-
 async def get_graph(ctx: MobileUseContext) -> CompiledStateGraph:
     graph_builder = StateGraph(State)
 
@@ -100,12 +87,12 @@ async def get_graph(ctx: MobileUseContext) -> CompiledStateGraph:
     graph_builder.add_node("cortex", CortexNode(ctx))
 
     graph_builder.add_node("executor", ExecutorNode(ctx))
-    executor_tool_node = ToolNode(
-        get_tools_from_wrappers(ctx=ctx, wrappers=EXECUTOR_WRAPPERS_TOOLS)
+    executor_tool_node = ExecutorToolNode(
+        tools=get_tools_from_wrappers(ctx=ctx, wrappers=EXECUTOR_WRAPPERS_TOOLS),
+        messages_key=EXECUTOR_MESSAGES_KEY,
     )
     graph_builder.add_node("executor_tools", executor_tool_node)
 
-    graph_builder.add_node("executor_context_cleaner", executor_context_cleaner_node)
     graph_builder.add_node("summarizer", SummarizerNode(ctx))
 
     # Linking nodes
@@ -132,18 +119,9 @@ async def get_graph(ctx: MobileUseContext) -> CompiledStateGraph:
     graph_builder.add_conditional_edges(
         "executor",
         post_executor_gate,
-        {"invoke_tools": "executor_tools", "skip": "executor_context_cleaner"},
-    )
-    graph_builder.add_conditional_edges(
-        "executor_tools",
-        post_executor_tools_gate,
-        {
-            "continue": "executor",
-            "done": "executor_context_cleaner",
-            "failed": "executor_context_cleaner",
-        },
+        {"invoke_tools": "executor_tools", "skip": "summarizer"},
     )
-    graph_builder.add_edge("executor_context_cleaner", "summarizer")
+    graph_builder.add_edge("executor_tools", "summarizer")
     graph_builder.add_edge("summarizer", "contextor")
 
     return graph_builder.compile()

minitap/mobile_use/graph/state.py

@@ -4,6 +4,7 @@ from langgraph.prebuilt.chat_agent_executor import AgentStatePydantic
 from typing_extensions import Annotated, Optional
 
 from minitap.mobile_use.agents.planner.types import Subgoal
+from minitap.mobile_use.config import AgentNode
 from minitap.mobile_use.utils.logger import get_logger
 from minitap.mobile_use.utils.recorder import record_interaction
 from minitap.mobile_use.context import MobileUseContext
@@ -36,10 +37,13 @@ class State(AgentStatePydantic):
         "Structured decisions made by the cortex, for the executor to follow",
         take_last,
     ]
+    complete_subgoals_by_ids: Annotated[
+        list[str],
+        "List of subgoal IDs to complete",
+        take_last,
+    ]
 
     # executor related keys
-    executor_retrigger: Annotated[Optional[bool], "Whether the executor must be retriggered"]
-    executor_failed: Annotated[bool, "Whether a tool call made by the executor failed"]
     executor_messages: Annotated[list[AnyMessage], "Sequential Executor messages", add_messages]
     cortex_last_thought: Annotated[Optional[str], "Last thought of the cortex for the executor"]
 
@@ -47,11 +51,18 @@ class State(AgentStatePydantic):
     agents_thoughts: Annotated[
         list[str],
         "All thoughts and reasons that led to actions (why a tool was called, expected outcomes..)",
+        take_last,
     ]
 
-    def sanitize_update(self, ctx: MobileUseContext, update: dict):
+    def sanitize_update(
+        self,
+        ctx: MobileUseContext,
+        update: dict,
+        agent: Optional[AgentNode] = None,
+    ):
         """
         Sanitizes the state update to ensure it is valid and apply side effect logic where required.
+        The agent is required if the update contains the "agents_thoughts" key.
         """
         updated_agents_thoughts: Optional[str | list[str]] = update.get("agents_thoughts", None)
         if updated_agents_thoughts is not None:
@@ -59,15 +70,24 @@ class State(AgentStatePydantic):
                 updated_agents_thoughts = [updated_agents_thoughts]
             elif not isinstance(updated_agents_thoughts, list):
                 raise ValueError("agents_thoughts must be a str or list[str]")
+            if agent is None:
+                raise ValueError("Agent is required when updating the 'agents_thoughts' key")
             update["agents_thoughts"] = _add_agent_thoughts(
                 ctx=ctx,
                 old=self.agents_thoughts,
                 new=updated_agents_thoughts,
+                agent=agent,
             )
         return update
 
 
-def _add_agent_thoughts(ctx: MobileUseContext, old: list[str], new: list[str]) -> list[str]:
+def _add_agent_thoughts(
+    ctx: MobileUseContext,
+    old: list[str],
+    new: list[str],
+    agent: AgentNode,
+) -> list[str]:
+    named_thoughts = [f"[{agent}] {thought}" for thought in new]
     if ctx.execution_setup:
-        record_interaction(ctx, response=AIMessage(content=str(new)))
-    return old + new
+        record_interaction(ctx, response=AIMessage(content=str(named_thoughts)))
+    return old + named_thoughts

minitap/mobile_use/main.py

@@ -1,9 +1,10 @@
 import asyncio
 import os
-from adbutils import AdbClient
 from typing import Optional
 
 import typer
+from adbutils import AdbClient
+from langchain.callbacks.base import Callbacks
 from rich.console import Console
 from typing_extensions import Annotated
 
@@ -26,6 +27,7 @@ async def run_automation(
     test_name: Optional[str] = None,
     traces_output_path_str: str = "traces",
     output_description: Optional[str] = None,
+    graph_config_callbacks: Callbacks = [],
 ):
     llm_config = initialize_llm_config()
     agent_profile = AgentProfile(name="default", llm_config=llm_config)
@@ -37,11 +39,13 @@ async def run_automation(
         config.with_hw_bridge_base_url(url=settings.DEVICE_HARDWARE_BRIDGE_BASE_URL)
     if settings.DEVICE_SCREEN_API_BASE_URL:
         config.with_screen_api_base_url(url=settings.DEVICE_SCREEN_API_BASE_URL)
+    if graph_config_callbacks:
+        config.with_graph_config_callbacks(graph_config_callbacks)
 
     agent = Agent(config=config.build())
     agent.init(
         retry_count=int(os.getenv("MOBILE_USE_HEALTH_RETRIES", 5)),
-        retry_wait_seconds=int(os.getenv("MOBILE_USE_HEALTH_DELAY", 5)),
+        retry_wait_seconds=int(os.getenv("MOBILE_USE_HEALTH_DELAY", 2)),
     )
 
     task = agent.new_task(goal)

minitap/mobile_use/sdk/agent.py

@@ -1,64 +1,63 @@
 import asyncio
-from datetime import datetime
-from pathlib import Path
 import sys
 import tempfile
 import time
+import uuid
+from datetime import datetime
+from pathlib import Path
 from types import NoneType
 from typing import Optional, TypeVar, overload
-import uuid
+
 from adbutils import AdbClient
 from langchain_core.messages import AIMessage
 from pydantic import BaseModel
-from minitap.mobile_use.agents.outputter.outputter import outputter
 
+from minitap.mobile_use.agents.outputter.outputter import outputter
+from minitap.mobile_use.clients.device_hardware_client import DeviceHardwareClient
+from minitap.mobile_use.clients.screen_api_client import ScreenApiClient
 from minitap.mobile_use.config import OutputConfig, record_events
-from minitap.mobile_use.graph.graph import get_graph
-from minitap.mobile_use.graph.state import State
-from minitap.mobile_use.sdk.builders.agent_config_builder import get_default_agent_config
-from minitap.mobile_use.sdk.builders.task_request_builder import TaskRequestBuilder
-from minitap.mobile_use.sdk.constants import (
-    DEFAULT_HW_BRIDGE_BASE_URL,
-    DEFAULT_SCREEN_API_BASE_URL,
-)
-from minitap.mobile_use.sdk.types.agent import AgentConfig
 from minitap.mobile_use.context import (
     DeviceContext,
     DevicePlatform,
     ExecutionSetup,
     MobileUseContext,
 )
-from minitap.mobile_use.clients.device_hardware_client import DeviceHardwareClient
-from minitap.mobile_use.clients.screen_api_client import ScreenApiClient
 from minitap.mobile_use.controllers.mobile_command_controller import (
     ScreenDataResponse,
     get_screen_data,
 )
 from minitap.mobile_use.controllers.platform_specific_commands_controller import get_first_device
-
-from minitap.mobile_use.servers.stop_servers import stop_servers
-from minitap.mobile_use.servers.device_hardware_bridge import BridgeStatus
-from minitap.mobile_use.servers.start_servers import (
-    start_device_hardware_bridge,
-    start_device_screen_api,
+from minitap.mobile_use.graph.graph import get_graph
+from minitap.mobile_use.graph.state import State
+from minitap.mobile_use.sdk.builders.agent_config_builder import get_default_agent_config
+from minitap.mobile_use.sdk.builders.task_request_builder import TaskRequestBuilder
+from minitap.mobile_use.sdk.constants import (
+    DEFAULT_HW_BRIDGE_BASE_URL,
+    DEFAULT_SCREEN_API_BASE_URL,
 )
-from minitap.mobile_use.utils.logger import get_logger
+from minitap.mobile_use.sdk.types.agent import AgentConfig
 from minitap.mobile_use.sdk.types.exceptions import (
+    AgentNotInitializedError,
     AgentProfileNotFoundError,
     AgentTaskRequestError,
     DeviceNotFoundError,
     ServerStartupError,
-    AgentNotInitializedError,
 )
 from minitap.mobile_use.sdk.types.task import AgentProfile, Task, TaskRequest, TaskStatus
+from minitap.mobile_use.servers.device_hardware_bridge import BridgeStatus
+from minitap.mobile_use.servers.start_servers import (
+    start_device_hardware_bridge,
+    start_device_screen_api,
+)
+from minitap.mobile_use.servers.stop_servers import stop_servers
+from minitap.mobile_use.utils.logger import get_logger
 from minitap.mobile_use.utils.media import (
     create_gif_from_trace_folder,
     create_steps_json_from_trace_folder,
     remove_images_from_trace_folder,
     remove_steps_json_from_trace_folder,
 )
-from minitap.mobile_use.utils.recorder import log_agent_thoughts
-
+from minitap.mobile_use.utils.recorder import log_agent_thought
 
 logger = get_logger(__name__)
 
@@ -127,7 +126,10 @@ class Agent:
                     f"Server start failed, attempting restart "
                     f"{restart_attempt}/{server_restart_attempts}"
                 )
-                time.sleep(3)
+                stop_servers(
+                    should_stop_screen_api=self._is_default_screen_api,
+                    should_stop_hw_bridge=self._is_default_hw_bridge,
+                )
             else:
                 error_msg = "Mobile-use servers failed to start after all restart attempts."
                 logger.error(error_msg)
@@ -261,17 +263,31 @@ class Agent:
                 input=graph_input,
                 config={
                     "recursion_limit": task.request.max_steps,
+                    "callbacks": self._config.graph_config_callbacks,
                 },
-                stream_mode=["messages", "custom", "values"],
+                stream_mode=["messages", "custom", "updates", "values"],
             ):
-                stream_mode, content = chunk
+                stream_mode, payload = chunk
                 if stream_mode == "values":
-                    last_state_snapshot = content  # type: ignore
+                    last_state_snapshot = payload  # type: ignore
                     last_state = State(**last_state_snapshot)  # type: ignore
-                    log_agent_thoughts(
-                        agents_thoughts=last_state.agents_thoughts,
-                        output_path=task.request.thoughts_output_path,
-                    )
+                    if task.request.thoughts_output_path:
+                        record_events(
+                            output_path=task.request.thoughts_output_path,
+                            events=last_state.agents_thoughts,
+                        )
+
+                if stream_mode == "updates":
+                    for key, value in payload.items():  # type: ignore
+                        if value and "agents_thoughts" in value:
+                            new_thoughts = value["agents_thoughts"]
+                            last_item = new_thoughts[-1] if new_thoughts else None
+                            if last_item:
+                                log_agent_thought(
+                                    prefix=key,
+                                    agent_thought=last_item,
+                                )
+
             if not last_state:
                 err = f"[{task_name}] No result received from graph"
                 logger.warning(err)
@@ -302,12 +318,12 @@ class Agent:
             self._finalize_tracing(task=task, context=context)
         return output
 
-    def clean(self):
-        if not self._initialized:
+    def clean(self, force: bool = False):
+        if not self._initialized and not force:
             return
         screen_api_ok, hw_bridge_ok = stop_servers(
-            device_screen_api=not self._is_default_screen_api,
-            device_hardware_bridge=not self._is_default_hw_bridge,
+            should_stop_screen_api=self._is_default_screen_api,
+            should_stop_hw_bridge=self._is_default_hw_bridge,
         )
         if not screen_api_ok:
             logger.warning("Failed to stop Device Screen API.")
@@ -402,10 +418,9 @@ class Agent:
             focused_app_info=None,
             device_date=None,
             structured_decisions=None,
+            complete_subgoals_by_ids=[],
             agents_thoughts=[],
             remaining_steps=task.request.max_steps,
-            executor_retrigger=False,
-            executor_failed=False,
             executor_messages=[],
             cortex_last_thought=None,
         )

minitap/mobile_use/sdk/builders/agent_config_builder.py

@@ -2,19 +2,20 @@
 Builder for AgentConfig objects using a fluent interface.
 """
 
-from typing import Dict, Optional, List
 import copy
+from typing import Dict, List, Optional
+
+from langchain_core.callbacks.base import Callbacks
 
 from minitap.mobile_use.config import get_default_llm_config
+from minitap.mobile_use.context import DevicePlatform
 from minitap.mobile_use.sdk.constants import (
     DEFAULT_HW_BRIDGE_BASE_URL,
     DEFAULT_PROFILE_NAME,
     DEFAULT_SCREEN_API_BASE_URL,
 )
-from minitap.mobile_use.sdk.types.agent import ApiBaseUrl, AgentConfig, ServerConfig
-from minitap.mobile_use.sdk.types.agent import AgentProfile
+from minitap.mobile_use.sdk.types.agent import AgentConfig, AgentProfile, ApiBaseUrl, ServerConfig
 from minitap.mobile_use.sdk.types.task import TaskRequestCommon
-from minitap.mobile_use.context import DevicePlatform
 
 
 class AgentConfigBuilder:
@@ -44,6 +45,7 @@ class AgentConfigBuilder:
         self._device_id: Optional[str] = None
         self._device_platform: Optional[DevicePlatform] = None
         self._servers: ServerConfig = get_default_servers()
+        self._graph_config_callbacks: Callbacks = None
 
     def add_profile(self, profile: AgentProfile) -> "AgentConfigBuilder":
         """
@@ -151,6 +153,16 @@ class AgentConfigBuilder:
         self._servers = copy.deepcopy(servers)
         return self
 
+    def with_graph_config_callbacks(self, callbacks: Callbacks) -> "AgentConfigBuilder":
+        """
+        Set the graph config callbacks.
+
+        Args:
+            callbacks: The graph config callbacks to use
+        """
+        self._graph_config_callbacks = callbacks
+        return self
+
     def build(self) -> AgentConfig:
         """
         Build the mobile-use AgentConfig object.
@@ -197,6 +209,7 @@ class AgentConfigBuilder:
             device_id=self._device_id,
             device_platform=self._device_platform,
             servers=self._servers,
+            graph_config_callbacks=self._graph_config_callbacks,
         )
 
 

minitap/mobile_use/sdk/types/agent.py

@@ -1,5 +1,7 @@
 from typing import Dict, Literal, Optional
 from urllib.parse import urlparse
+
+from langchain_core.callbacks.base import Callbacks
 from pydantic import BaseModel
 
 from minitap.mobile_use.context import DevicePlatform
@@ -71,3 +73,6 @@ class AgentConfig(BaseModel):
     device_id: Optional[str] = None
     device_platform: Optional[DevicePlatform] = None
     servers: ServerConfig
+    graph_config_callbacks: Callbacks = None
+
+    model_config = {"arbitrary_types_allowed": True}