droidrun 0.5.0.dev5__tar.gz → 0.5.0.dev7__tar.gz

{droidrun-0.5.0.dev5 → droidrun-0.5.0.dev7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: droidrun
-Version: 0.5.0.dev5
+Version: 0.5.0.dev7
 Summary: A framework for controlling Android devices through LLM agents
 Project-URL: Homepage, https://github.com/droidrun/droidrun
 Project-URL: Bug Tracker, https://github.com/droidrun/droidrun/issues

{droidrun-0.5.0.dev5 → droidrun-0.5.0.dev7}/droidrun/agent/codeact/codeact_agent.py

@@ -1,8 +1,10 @@
 import asyncio
+import copy
 import inspect
 import logging
 from typing import TYPE_CHECKING, Optional, Type
 
+from llama_index.core.base.llms.types import ChatMessage, ImageBlock, TextBlock
 from llama_index.core.llms.llm import LLM
 from llama_index.core.workflow import Context, StartEvent, StopEvent, Workflow, step
 from opentelemetry import trace
@@ -21,7 +23,6 @@ from droidrun.agent.usage import get_usage_from_response
 from droidrun.agent.utils.chat_utils import (
     extract_code_and_thought,
     limit_history,
-    to_chat_messages,
 )
 from droidrun.agent.utils.executer import ExecuterState, SimpleCodeExecutor
 from droidrun.agent.utils.inference import acall_with_retries
@@ -52,7 +53,7 @@ class CodeActAgent(Workflow):
     Agent that generates and executes Python code using atomic actions.
 
     Uses ReAct cycle: Thought -> Code -> Observation -> repeat until complete().
-    Messages stored as list[dict], converted to ChatMessage only for LLM calls.
+    Messages stored as list[ChatMessage] to preserve thinking tokens across turns.
     """
 
     def __init__(
@@ -90,7 +91,7 @@ class CodeActAgent(Workflow):
         self.prompt_resolver = prompt_resolver or PromptResolver()
         self.tracing_config = tracing_config
 
-        self.system_prompt: dict | None = None
+        self.system_prompt: ChatMessage | None = None
         self.code_exec_counter = 0
         self.remembered_info: list[str] | None = None
 
@@ -150,7 +151,7 @@ class CodeActAgent(Workflow):
 
         logger.debug("CodeActAgent initialized.")
 
-    async def _build_system_prompt(self) -> dict:
+    async def _build_system_prompt(self) -> ChatMessage:
         """Build system prompt message."""
         # Build template context with available tools for conditional examples
         template_context = {
@@ -178,9 +179,9 @@ class CodeActAgent(Workflow):
                 str(PathResolver.resolve(prompt_path, must_exist=True)),
                 template_context,
             )
-        return {"role": "system", "content": [{"text": system_text}]}
+        return ChatMessage(role="system", content=system_text)
 
-    async def _build_user_prompt(self, goal: str) -> dict:
+    async def _build_user_prompt(self, goal: str) -> ChatMessage:
         """Build initial user prompt message."""
         custom_user_prompt = self.prompt_resolver.get_prompt("fast_agent_user")
         if custom_user_prompt:
@@ -207,7 +208,7 @@ class CodeActAgent(Workflow):
                     ),
                 },
             )
-        return {"role": "user", "content": [{"text": user_text}]}
+        return ChatMessage(role="user", content=user_text)
 
     @step
     async def prepare_chat(self, ctx: Context, ev: StartEvent) -> CodeActInputEvent:
@@ -240,8 +241,8 @@ class CodeActAgent(Workflow):
             for idx, item in enumerate(remembered_info, 1):
                 memory_text += f"{idx}. {item}\n"
             # Append to first user message
-            self.shared_state.message_history[0]["content"].append(
-                {"text": memory_text}
+            self.shared_state.message_history[0].blocks.append(
+                TextBlock(text=memory_text)
             )
 
         return CodeActInputEvent()
@@ -296,7 +297,10 @@ class CodeActAgent(Workflow):
             ui_state = await self.state_provider.get_state()
             self.action_ctx.ui = ui_state
 
-            # Update shared state
+            # Update shared state (previous ← current, current ← new)
+            self.shared_state.previous_formatted_device_state = (
+                self.shared_state.formatted_device_state
+            )
             self.shared_state.formatted_device_state = ui_state.formatted_text
             self.shared_state.focused_text = ui_state.focused_text
             self.shared_state.a11y_tree = ui_state.elements
@@ -311,11 +315,6 @@ class CodeActAgent(Workflow):
             # Stream formatted state for trajectory
             ctx.write_event_to_stream(RecordUIStateEvent(ui_state=ui_state.elements))
 
-            # Add device state to last user message
-            self.shared_state.message_history[-1]["content"].append(
-                {"text": f"\n{ui_state.formatted_text}\n"}
-            )
-
         except DeviceDisconnectedError:
             raise
         except Exception as e:
@@ -323,27 +322,51 @@ class CodeActAgent(Workflow):
             if self.debug:
                 logger.error("State retrieval error details:", exc_info=True)
 
-        # Add screenshot to message if vision enabled
-        if self.vision and screenshot:
-            self.shared_state.message_history[-1]["content"].append(
-                {"image": screenshot}
-            )
-
-        # Limit history and prepare for LLM
+        # Limit history and build ephemeral copy for LLM
         limited_history = limit_history(
             self.shared_state.message_history,
             LLM_HISTORY_LIMIT * 2,
             preserve_first=True,
         )
+        messages_to_send = [self.system_prompt] + copy.deepcopy(limited_history)
+
+        # Inject device state and screenshot into the copy (not the original)
+        user_indices = [
+            i for i, msg in enumerate(messages_to_send) if msg.role == "user"
+        ]
+        if user_indices:
+            last_user_idx = user_indices[-1]
+
+            # Current device state → last user message
+            current_state = self.shared_state.formatted_device_state.strip()
+            if current_state:
+                messages_to_send[last_user_idx].blocks.append(
+                    TextBlock(
+                        text=f"\n<device_state>\n{current_state}\n</device_state>\n"
+                    )
+                )
+
+            # Screenshot → last user message
+            if self.vision and screenshot:
+                messages_to_send[last_user_idx].blocks.append(
+                    ImageBlock(image=screenshot)
+                )
 
-        # Build final messages: system + history
-        messages_to_send = [self.system_prompt] + limited_history
-        chat_messages = to_chat_messages(messages_to_send)
+            # Previous device state → second-to-last user message
+            if len(user_indices) >= 2:
+                second_last_idx = user_indices[-2]
+                prev_state = self.shared_state.previous_formatted_device_state.strip()
+                if prev_state:
+                    messages_to_send[second_last_idx].blocks.append(
+                        TextBlock(
+                            text=f"\n<previous_device_state>\n{prev_state}\n</previous_device_state>\n"
+                        )
+                    )
 
         # Call LLM
         logger.info("CodeAct response:", extra={"color": "yellow"})
         response = await acall_with_retries(
-            self.llm, chat_messages, stream=self.agent_config.streaming
+            self.llm, messages_to_send, stream=self.agent_config.streaming
         )
 
         if response is None:
@@ -360,11 +383,9 @@ class CodeActAgent(Workflow):
         except Exception as e:
             logger.warning(f"Could not get usage: {e}")
 
-        # Store assistant response
+        # Store assistant response (preserves ThinkingBlock, additional_kwargs, etc.)
+        self.shared_state.message_history.append(response.message)
         response_text = response.message.content
-        self.shared_state.message_history.append(
-            {"role": "assistant", "content": [{"text": response_text}]}
-        )
 
         # Extract thought and code
         code, thought = extract_code_and_thought(response_text)
@@ -391,7 +412,7 @@ class CodeActAgent(Workflow):
                 "Now, describe the next step you will take to address the original goal."
             )
             self.shared_state.message_history.append(
-                {"role": "user", "content": [{"text": no_thoughts_text}]}
+                ChatMessage(role="user", content=no_thoughts_text)
             )
         else:
             logger.debug(f"Reasoning: {ev.thought}")
@@ -408,7 +429,7 @@ class CodeActAgent(Workflow):
                 "function within a <python></python> code block."
             )
             self.shared_state.message_history.append(
-                {"role": "user", "content": [{"text": no_code_text}]}
+                ChatMessage(role="user", content=no_code_text)
             )
             return CodeActInputEvent()
 
@@ -482,7 +503,7 @@ class CodeActAgent(Workflow):
         # Add execution output as user message
         observation_text = f"Execution Result:\n<result>\n{output}\n</result>"
         self.shared_state.message_history.append(
-            {"role": "user", "content": [{"text": observation_text}]}
+            ChatMessage(role="user", content=observation_text)
         )
 
         return CodeActInputEvent()

{droidrun-0.5.0.dev5 → droidrun-0.5.0.dev7}/droidrun/agent/codeact/tools_agent.py

@@ -9,9 +9,12 @@ compatibility with DroidAgent's execute_task() method.
 """
 
 import asyncio
+import copy
 import logging
+import os
 from typing import TYPE_CHECKING, Optional, Type
 
+from llama_index.core.base.llms.types import ChatMessage, ImageBlock, TextBlock
 from llama_index.core.llms.llm import LLM
 from llama_index.core.workflow import Context, StartEvent, StopEvent, Workflow, step
 from opentelemetry import trace
@@ -35,7 +38,7 @@ from droidrun.agent.codeact.xml_parser import (
 from droidrun.agent.common.constants import LLM_HISTORY_LIMIT
 from droidrun.agent.common.events import RecordUIStateEvent, ScreenshotEvent
 from droidrun.agent.usage import get_usage_from_response
-from droidrun.agent.utils.chat_utils import limit_history, to_chat_messages
+from droidrun.agent.utils.chat_utils import limit_history
 from droidrun.agent.utils.inference import acall_with_retries
 from droidrun.agent.utils.prompt_resolver import PromptResolver
 from droidrun.agent.utils.tracing_setup import record_langfuse_screenshot
@@ -56,7 +59,7 @@ class FastAgent(Workflow):
     """Agent that uses XML tool-calling instead of code generation.
 
     Uses ReAct cycle: Thought -> Tool Call -> Observation -> repeat until complete().
-    Messages stored as list[dict], converted to ChatMessage only for LLM calls.
+    Messages stored as list[ChatMessage] to preserve thinking tokens across turns.
     """
 
     def __init__(
@@ -88,12 +91,15 @@ class FastAgent(Workflow):
         self.action_ctx = action_ctx
         self.state_provider = state_provider
         self.save_trajectory = save_trajectory
+        self._stream_screenshots = os.environ.get(
+            "DROIDRUN_STREAM_SCREENSHOTS", ""
+        ).lower() in ("1", "true")
         self.shared_state = shared_state
         self.output_model = output_model
         self.prompt_resolver = prompt_resolver or PromptResolver()
         self.tracing_config = tracing_config
 
-        self.system_prompt: dict | None = None
+        self.system_prompt: ChatMessage | None = None
         self.tool_call_counter = 0
         self.remembered_info: list[str] | None = None
 
@@ -108,7 +114,7 @@ class FastAgent(Workflow):
 
         logger.debug("FastAgent initialized.")
 
-    async def _build_system_prompt(self) -> dict:
+    async def _build_system_prompt(self) -> ChatMessage:
         """Build system prompt message."""
         template_context = {
             "tool_descriptions": self.tool_descriptions,
@@ -133,9 +139,9 @@ class FastAgent(Workflow):
                 self.agent_config.get_fast_agent_system_prompt_path(),
                 template_context,
             )
-        return {"role": "system", "content": [{"text": system_text}]}
+        return ChatMessage(role="system", content=system_text)
 
-    async def _build_user_prompt(self, goal: str) -> dict:
+    async def _build_user_prompt(self, goal: str) -> ChatMessage:
         """Build initial user prompt message."""
         custom_user_prompt = self.prompt_resolver.get_prompt("fast_agent_user")
         if custom_user_prompt:
@@ -158,7 +164,7 @@ class FastAgent(Workflow):
                     ),
                 },
             )
-        return {"role": "user", "content": [{"text": user_text}]}
+        return ChatMessage(role="user", content=user_text)
 
     @step
     async def prepare_chat(self, ctx: Context, ev: StartEvent) -> FastAgentInputEvent:
@@ -190,8 +196,8 @@ class FastAgent(Workflow):
             memory_text = "\n### Remembered Information:\n"
             for idx, item in enumerate(remembered_info, 1):
                 memory_text += f"{idx}. {item}\n"
-            self.shared_state.message_history[0]["content"].append(
-                {"text": memory_text}
+            self.shared_state.message_history[0].blocks.append(
+                TextBlock(text=memory_text)
             )
 
         return FastAgentInputEvent()
@@ -218,7 +224,7 @@ class FastAgent(Workflow):
 
         # Capture screenshot if needed
         screenshot = None
-        if self.vision or self.save_trajectory != "none":
+        if self.vision or self._stream_screenshots or self.save_trajectory != "none":
             try:
                 screenshot = await self.action_ctx.driver.screenshot()
 
@@ -246,7 +252,10 @@ class FastAgent(Workflow):
             ui_state = await self.state_provider.get_state()
             self.action_ctx.ui = ui_state
 
-            # Update shared state
+            # Update shared state (previous ← current, current ← new)
+            self.shared_state.previous_formatted_device_state = (
+                self.shared_state.formatted_device_state
+            )
             self.shared_state.formatted_device_state = ui_state.formatted_text
             self.shared_state.focused_text = ui_state.focused_text
             self.shared_state.a11y_tree = ui_state.elements
@@ -261,11 +270,6 @@ class FastAgent(Workflow):
             # Stream formatted state for trajectory
             ctx.write_event_to_stream(RecordUIStateEvent(ui_state=ui_state.elements))
 
-            # Add device state to last user message
-            self.shared_state.message_history[-1]["content"].append(
-                {"text": f"\n{ui_state.formatted_text}\n"}
-            )
-
         except DeviceDisconnectedError:
             raise
         except Exception as e:
@@ -273,27 +277,51 @@ class FastAgent(Workflow):
             if self.debug:
                 logger.error("State retrieval error details:", exc_info=True)
 
-        # Add screenshot to message if vision enabled
-        if self.vision and screenshot:
-            self.shared_state.message_history[-1]["content"].append(
-                {"image": screenshot}
-            )
-
-        # Limit history and prepare for LLM
+        # Limit history and build ephemeral copy for LLM
         limited_history = limit_history(
             self.shared_state.message_history,
             LLM_HISTORY_LIMIT * 2,
             preserve_first=True,
         )
+        messages_to_send = [self.system_prompt] + copy.deepcopy(limited_history)
+
+        # Inject device state and screenshot into the copy (not the original)
+        user_indices = [
+            i for i, msg in enumerate(messages_to_send) if msg.role == "user"
+        ]
+        if user_indices:
+            last_user_idx = user_indices[-1]
+
+            # Current device state → last user message
+            current_state = self.shared_state.formatted_device_state.strip()
+            if current_state:
+                messages_to_send[last_user_idx].blocks.append(
+                    TextBlock(
+                        text=f"\n<device_state>\n{current_state}\n</device_state>\n"
+                    )
+                )
+
+            # Screenshot → last user message
+            if self.vision and screenshot:
+                messages_to_send[last_user_idx].blocks.append(
+                    ImageBlock(image=screenshot)
+                )
 
-        # Build final messages: system + history
-        messages_to_send = [self.system_prompt] + limited_history
-        chat_messages = to_chat_messages(messages_to_send)
+            # Previous device state → second-to-last user message
+            if len(user_indices) >= 2:
+                second_last_idx = user_indices[-2]
+                prev_state = self.shared_state.previous_formatted_device_state.strip()
+                if prev_state:
+                    messages_to_send[second_last_idx].blocks.append(
+                        TextBlock(
+                            text=f"\n<previous_device_state>\n{prev_state}\n</previous_device_state>\n"
+                        )
+                    )
 
         # Call LLM
         logger.info("FastAgent response:", extra={"color": "yellow"})
         response = await acall_with_retries(
-            self.llm, chat_messages, stream=self.agent_config.streaming
+            self.llm, messages_to_send, stream=self.agent_config.streaming
         )
 
         if response is None:
@@ -310,11 +338,9 @@ class FastAgent(Workflow):
         except Exception as e:
             logger.warning(f"Could not get usage: {e}")
 
-        # Store assistant response
+        # Store assistant response (preserves ThinkingBlock, additional_kwargs, etc.)
+        self.shared_state.message_history.append(response.message)
         response_text = response.message.content
-        self.shared_state.message_history.append(
-            {"role": "assistant", "content": [{"text": response_text}]}
-        )
 
         # Parse tool calls from response
         thought, tool_calls = parse_tool_calls(response_text, self.param_types)
@@ -360,7 +386,7 @@ class FastAgent(Workflow):
                 "Now, describe the next step you will take to address the original goal."
             )
             self.shared_state.message_history.append(
-                {"role": "user", "content": [{"text": no_thoughts_text}]}
+                ChatMessage(role="user", content=no_thoughts_text)
             )
         else:
             logger.debug(f"Reasoning: {ev.thought}")
@@ -382,7 +408,7 @@ class FastAgent(Workflow):
                 "</function_calls>"
             )
             self.shared_state.message_history.append(
-                {"role": "user", "content": [{"text": no_tools_text}]}
+                ChatMessage(role="user", content=no_tools_text)
             )
             return FastAgentInputEvent()
 
@@ -469,7 +495,7 @@ class FastAgent(Workflow):
 
         # Add results as user message
         self.shared_state.message_history.append(
-            {"role": "user", "content": [{"text": output}]}
+            ChatMessage(role="user", content=output)
         )
 
         return FastAgentInputEvent()

{droidrun-0.5.0.dev5 → droidrun-0.5.0.dev7}/droidrun/agent/droid/droid_agent.py

@@ -8,6 +8,7 @@ Architecture:
 """
 
 import logging
+import os
 import traceback
 from typing import TYPE_CHECKING, Awaitable, Type, Union
 
@@ -208,6 +209,10 @@ class DroidAgent(Workflow):
         # Check if using external agent - skip LLM loading
         self._using_external_agent = self.config.agent.name != "droidrun"
 
+        self._stream_screenshots = os.environ.get(
+            "DROIDRUN_STREAM_SCREENSHOTS", ""
+        ).lower() in ("1", "true")
+
         self.timeout = timeout
 
         # Store user custom tools
@@ -1039,17 +1044,21 @@ class DroidAgent(Workflow):
                 if self.config.logging.debug:
                     logger.error(traceback.format_exc())
 
-        # Capture final screenshot before saving trajectory
-        if self.config.logging.save_trajectory != "none":
+        # Capture final screenshot (independent of trajectory persistence)
+        vision_any = (
+            self.config.agent.manager.vision
+            or self.config.agent.executor.vision
+            or self.config.agent.fast_agent.vision
+        )
+        if (
+            vision_any
+            or self._stream_screenshots
+            or self.config.logging.save_trajectory != "none"
+        ):
             try:
                 screenshot = await self.action_ctx.driver.screenshot()
                 if screenshot:
                     ctx.write_event_to_stream(ScreenshotEvent(screenshot=screenshot))
-                    vision_any = (
-                        self.config.agent.manager.vision
-                        or self.config.agent.executor.vision
-                        or self.config.agent.fast_agent.vision
-                    )
                     parent_span = trace.get_current_span()
                     record_langfuse_screenshot(
                         screenshot,
@@ -1061,6 +1070,8 @@ class DroidAgent(Workflow):
             except Exception as e:
                 logger.warning(f"Failed to capture final screenshot: {e}")
 
+        # Save trajectory to disk
+        if self.config.logging.save_trajectory != "none":
             # Populate macro data from RecordingDriver log
             if isinstance(self.driver, RecordingDriver):
                 self.trajectory.macro = list(self.driver.log)

{droidrun-0.5.0.dev5 → droidrun-0.5.0.dev7}/droidrun/agent/droid/state.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 from typing import Dict, List, Optional
 
+from llama_index.core.base.llms.types import ChatMessage
 from pydantic import BaseModel, ConfigDict, Field
 
 from droidrun.telemetry import PackageVisitEvent, capture
@@ -86,9 +87,9 @@ class DroidAgentState(BaseModel):
     success: Optional[bool] = None
 
     # ========================================================================
-    # Message History (for stateful agents - list of dicts)
+    # Message History (for stateful agents - preserves ChatMessage blocks)
     # ========================================================================
-    message_history: List[Dict] = Field(default_factory=list)
+    message_history: List[ChatMessage] = Field(default_factory=list)
 
     # ========================================================================
     # Error Handling

{droidrun-0.5.0.dev5 → droidrun-0.5.0.dev7}/droidrun/agent/executor/executor_agent.py

@@ -14,18 +14,18 @@ import json
 import logging
 from typing import TYPE_CHECKING, Optional
 
+from llama_index.core.base.llms.types import ChatMessage, ImageBlock, TextBlock
 from llama_index.core.llms.llm import LLM
 from llama_index.core.workflow import Context, StartEvent, StopEvent, Workflow, step
 
 from droidrun.agent.executor.events import (
     ExecutorActionEvent,
+    ExecutorActionResultEvent,
     ExecutorContextEvent,
     ExecutorResponseEvent,
-    ExecutorActionResultEvent,
 )
 from droidrun.agent.executor.prompts import parse_executor_response
 from droidrun.agent.usage import get_usage_from_response
-from droidrun.agent.utils.chat_utils import to_chat_messages
 from droidrun.agent.utils.inference import acall_with_retries
 from droidrun.agent.utils.prompt_resolver import PromptResolver
 from droidrun.config_manager.config_manager import AgentConfig
@@ -44,7 +44,7 @@ class ExecutorAgent(Workflow):
     Action execution agent that performs specific actions.
 
     Single-turn agent: receives subgoal, selects action, executes it.
-    Uses dict messages, converts to ChatMessage at LLM call time.
+    Uses ChatMessage objects directly for LLM calls.
     """
 
     # Flow-control tools hidden from executor's LLM prompt
@@ -123,14 +123,14 @@ class ExecutorAgent(Workflow):
                 variables,
             )
 
-        # Build message as dict
-        messages = [{"role": "user", "content": [{"text": prompt_text}]}]
+        # Build message
+        messages = [ChatMessage(role="user", blocks=[TextBlock(text=prompt_text)])]
 
         # Add screenshot if vision enabled
         if self.vision:
             screenshot = self.shared_state.screenshot
             if screenshot is not None:
-                messages[0]["content"].append({"image": screenshot})
+                messages[0].blocks.append(ImageBlock(image=screenshot))
                 logger.debug("📸 Using screenshot for Executor")
             else:
                 logger.warning("⚠️ Vision enabled but no screenshot available")
@@ -149,13 +149,10 @@ class ExecutorAgent(Workflow):
         # Get messages from context
         messages = await ctx.store.get("executor_messages")
 
-        # Convert to ChatMessage and call LLM
-        chat_messages = to_chat_messages(messages)
-
         try:
             logger.info("Executor response:", extra={"color": "green"})
             response = await acall_with_retries(
-                self.llm, chat_messages, stream=self.agent_config.streaming
+                self.llm, messages, stream=self.agent_config.streaming
             )
             response_text = str(response)
         except ValueError as e: