droidrun 0.2.0__py3-none-any.whl → 0.3.0__py3-none-any.whl

droidrun/agent/planner/prompts.py

@@ -6,7 +6,7 @@ separated from the workflow logic for better maintainability.
 """
 
 # System prompt for the PlannerAgent that explains its role and capabilities
-DEFAULT_PLANNER_SYSTEM_PROMPT = """You are an Android Task Planner. Your job is to create short, functional plans (1-5 steps) to achieve a user's goal on an Android device.
+DEFAULT_PLANNER_SYSTEM_PROMPT = """You are an Android Task Planner. Your job is to create short, functional plans (1-5 steps) to achieve a user's goal on an Android device, and assign each task to the most appropriate specialized agent.
 
 **Inputs You Receive:**
 1.  **User's Overall Goal.**
@@ -20,34 +20,34 @@ DEFAULT_PLANNER_SYSTEM_PROMPT = """You are an Android Task Planner. Your job is
     * For failed tasks, the detailed reasons for failure.
     * This history persists across all planning cycles and is never lost, even when creating new tasks.
 
+**Available Specialized Agents:**
+You have access to specialized agents, each optimized for specific types of tasks:
+{agents}
+
 **Your Task:**
-Given the goal, current state, and task history, devise the **next 1-5 functional steps**. Focus on what to achieve, not how. Planning fewer steps at a time improves accuracy, as the state can change.
+Given the goal, current state, and task history, devise the **next 1-5 functional steps** and assign each to the most appropriate specialized agent.
+Focus on what to achieve, not how. Planning fewer steps at a time improves accuracy, as the state can change.
 
 **Step Format:**
-Each step must be a functional goal. A **precondition** describing the expected starting screen/state for that step is highly recommended for clarity, especially for steps after the first in your 1-5 step plan. Each task string can start with "Precondition: ... Goal: ...". If a specific precondition isn't critical for the first step in your current plan segment, you can use "Precondition: None. Goal: ..." or simply state the goal if the context is implicitly clear from the first step of a new sequence.
-
-**Executor Agent Capabilities:**
-The plan you create will be executed by another agent. This executor can:
-*   `swipe(direction: str, distance_percentage: int)`
-*   `input_text(text: str, element_hint: Optional[str] = None)`
-*   `press_key(keycode: int)` (Common: 3=HOME, 4=BACK)
-*   `tap_by_coordinates(x: int, y: int)` (This is a fallback; prefer functional goals)
-*   `start_app(package_name: str)`
-*   `remember(info: str)
-`
-*   (The executor will use the UI JSON to find elements for your functional goals like "Tap 'Settings button'" or "Enter text into 'Username field'").
+Each step must be a functional goal.
+A **precondition** describing the expected starting screen/state for that step is highly recommended for clarity, especially for steps after the first in your 1-5 step plan.
+Each task string can start with "Precondition: ... Goal: ...".
+If a specific precondition isn't critical for the first step in your current plan segment, you can use "Precondition: None. Goal: ..." or simply state the goal if the context is implicitly clear from the first step of a new sequence.
 
 **Your Output:**
-*   Use the `set_tasks` tool to provide your 1-5 step plan as a list of strings.
-*   **After your planned steps are executed, you will be invoked again with the new device state.** You will then:
+*   Use the `set_tasks_with_agents` tool to provide your 1-5 step plan with agent assignments.
+*   Each task should be assigned to a specialized agent using it's name.
+
+*   **After your planned steps are executed, you will be invoked again with the new device state.**
+You will then:
     1.  Assess if the **overall user goal** is complete.
     2.  If complete, call the `complete_goal(message: str)` tool.
-    3.  If not complete, generate the next 1-5 steps using `set_tasks`.
+    3.  If not complete, generate the next 1-5 steps using `set_tasks_with_agents`.
 
 **Memory Persistence:**
 *   You maintain a COMPLETE memory of ALL tasks across the entire session:
     * Every task that was completed or failed is preserved in your context.
-    * Previously completed steps are never lost when calling `set_tasks()` for new steps.
+    * Previously completed steps are never lost when calling `set_tasks_with_agents()` for new steps.
     * You will see all historical tasks each time you're called.
     * Use this accumulated knowledge to build progressively on successful steps.
     * When you see discovered information (e.g., dates, locations), use it explicitly in future tasks.
@@ -57,67 +57,47 @@ The plan you create will be executed by another agent. This executor can:
 *   **NO Low-Level Actions:** Do NOT specify swipes, taps on coordinates, or element IDs in your plan.
 *   **Short Plans (1-5 steps):** Plan only the immediate next actions.
 *   **Learn From History:** If a task failed previously, try a different approach.
-*   **Use Tools:** Your response *must* be a Python code block calling `set_tasks` or `complete_goal`.
+*   **Use Tools:** Your response *must* be a Python code block calling `set_tasks_with_agents` or `complete_goal`.
+*   **Smart Agent Assignment:** Choose the most appropriate agent for each task type.
 
 **Available Planning Tools:**
-*   `set_tasks(tasks: List[str])`: Defines the sequence of tasks. Each element in the list is a string representing a single task.
+*   `set_tasks_with_agents(task_assignments: List[Dict[str, str]])`: Defines the sequence of tasks with agent assignments. Each element should be a dictionary with 'task' and 'agent' keys.
 *   `complete_goal(message: str)`: Call this when the overall user goal has been achieved. The message can summarize the completion.
 
 ---
 
 **Example Interaction Flow:**
 
-**User Goal:** Turn on Wi-Fi.
+**User Goal:** Open Gmail and compose a new email.
 
 **(Round 1) Planner Input:**
-*   Goal: "Turn on Wi-Fi"
+*   Goal: "Open Gmail and compose a new email"
 *   Current State: Screenshot of Home screen, UI JSON.
 *   Task History: None (first planning cycle)
 
 **Planner Thought Process (Round 1):**
-Need to open settings first, then go to Network settings. This is the first plan.
-1. Task 1: "Precondition: None. Goal: Open the Settings app."
-2. Task 2: "Precondition: Settings main screen is open. Goal: Navigate to 'Network & internet' settings."
+Need to first open Gmail app, then navigate to compose. The first task is app launching, the second is UI navigation.
 
 **Planner Output (Round 1):**
 ```python
-set_tasks(tasks=[
-    "Precondition: None. Goal: Open the Settings app.",
-    "Precondition: Settings main screen is open. Goal: Navigate to 'Network & internet' settings."
+set_tasks_with_agents([
+    {{'task': 'Precondition: None. Goal: Open the Gmail app.', 'agent': <Specialized_Agent>}},
+    {{'task': 'Precondition: Gmail app is open and loaded. Goal: Navigate to compose new email.', 'agent': <Specialized Agents>}}
 ])
 ```
 
-**(After Executor performs these steps...)**
+**(After specialized agents perform these steps...)**
 
 **(Round 2) Planner Input:**
-*   Goal: "Turn on Wi-Fi"
-*   Current State: Screenshot of "Network & internet" screen, UI JSON showing "Wi-Fi" option.
-*   Task History: Shows ALL previously completed tasks, including "Open the Settings app" and "Navigate to 'Network & internet' settings"
-
-**Planner Thought Process (Round 2):**
-Now on "Network & internet". Need to tap Wi-Fi, then enable it. I can see from history that we've already opened Settings and navigated to Network & internet.
-1. Task 1: "Precondition: 'Network & internet' screen is open. Goal: Tap the 'Wi-Fi' option."
-2. Task 2: "Precondition: Wi-Fi settings screen is open. Goal: Enable the Wi-Fi toggle if it's off."
+*   Goal: "Open Gmail and compose a new email"
+*   Current State: Screenshot of Gmail compose screen, UI JSON showing compose interface.
+*   Task History: Shows completed tasks with their assigned agents
 
 **Planner Output (Round 2):**
 ```python
-set_tasks(tasks=[
-    "Precondition: 'Network & internet' screen is open. Goal: Tap the 'Wi-Fi' option.",
-    "Precondition: Wi-Fi settings screen is open. Goal: Enable the Wi-Fi toggle if it's off."
-])
+complete_goal(message="Gmail has been opened and compose email screen is ready for use.")
 ```
-
-**(After Executor performs these steps...)**
-
-**(Round 3) Planner Input:**
-*   Goal: "Turn on Wi-Fi"
-*   Current State: Screenshot of Wi-Fi screen, UI JSON showing Wi-Fi is now ON.
-*   Task History: Shows ALL previous tasks completed successfully (all 4 tasks from previous rounds)
-
-**Planner Output (Round 3):**
-```python
-complete_goal(message="Wi-Fi has been successfully enabled.")
-```"""
+"""
 
 # User prompt template that simply states the goal
 DEFAULT_PLANNER_USER_PROMPT = """Goal: {goal}"""

droidrun/agent/utils/__init__.py

@@ -0,0 +1,3 @@
+"""
+Utility modules for DroidRun agents.
+"""

droidrun/agent/utils/async_utils.py

@@ -1,7 +1,4 @@
 import asyncio
-import nest_asyncio
-nest_asyncio_applied = False
-
 
 def async_to_sync(func):
     """
@@ -15,42 +12,6 @@ def async_to_sync(func):
     """
 
     def wrapper(*args, **kwargs):
-        global nest_asyncio_applied # Declare modification of global at the start of the scope
-        coro = func(*args, **kwargs)
-        try:
-            # Try to get the running event loop.
-            loop = asyncio.get_running_loop()
-
-            # If the loop is running, apply nest_asyncio if available and needed.
-            # Removed global declaration from here
-            if nest_asyncio and not nest_asyncio_applied:
-                nest_asyncio.apply()
-                nest_asyncio_applied = True
-            # Run the coroutine to completion within the running loop.
-            # This requires nest_asyncio to work correctly in nested scenarios.
-            # Changed from ensure_future to run_until_complete to make it truly synchronous.
-            return loop.run_until_complete(coro)
-
-        except RuntimeError:
-            # No running event loop found.
-            try:
-                # Check if there's a loop policy and a current event loop set, even if not running.
-                loop = asyncio.get_event_loop_policy().get_event_loop()
-                if loop.is_running():
-                     # This case should ideally be caught by get_running_loop(),
-                     # but as a fallback, handle similarly if loop is running.
-                     # Removed global declaration from here
-                     if nest_asyncio and not nest_asyncio_applied:
-                         nest_asyncio.apply()
-                         nest_asyncio_applied = True
-                     return loop.run_until_complete(coro)
-                else:
-                    # Loop exists but is not running, run until complete.
-                    return loop.run_until_complete(coro)
-            except RuntimeError:
-                 # If get_event_loop() also fails (no loop set at all for this thread),
-                 # use asyncio.run() which creates a new loop.
-                 return asyncio.run(coro)
-
+        return asyncio.run(func(*args, **kwargs))
 
     return wrapper

droidrun/agent/utils/chat_utils.py

@@ -1,14 +1,18 @@
 import base64
+import re
+import inspect
+
+
 import json
 import logging
-from typing import List, TYPE_CHECKING
+from typing import List, TYPE_CHECKING, Optional, Tuple
+from droidrun.agent.context import Reflection
 from llama_index.core.base.llms.types import ChatMessage, ImageBlock, TextBlock
 
 if TYPE_CHECKING:
-    from ...tools import Tools
+    from droidrun.tools import Tools
 
 logger = logging.getLogger("droidrun")
-logging.basicConfig(level=logging.INFO)
 
 def message_copy(message: ChatMessage, deep = True) -> ChatMessage:
     if deep:
@@ -23,29 +27,98 @@ def message_copy(message: ChatMessage, deep = True) -> ChatMessage:
 
     return copied_message
 
-async def add_ui_text_block(tools: 'Tools', chat_history: List[ChatMessage], retry = 5, copy = True) -> List[ChatMessage]:
+async def add_reflection_summary(reflection: Reflection, chat_history: List[ChatMessage]) -> List[ChatMessage]:
+    """Add reflection summary and advice to help the planner understand what went wrong and what to do differently."""
+    
+    reflection_text = "\n### The last task failed. You have additional information about what happenend. \nThe Reflection from Previous Attempt:\n"
+    
+    if reflection.summary:
+        reflection_text += f"**What happened:** {reflection.summary}\n\n"
+    
+    if reflection.advice:
+        reflection_text += f"**Recommended approach for this retry:** {reflection.advice}\n"
+    
+    reflection_block = TextBlock(text=reflection_text)
+    
+    # Copy chat_history and append reflection block to the last message
+    chat_history = chat_history.copy()
+    chat_history[-1] = message_copy(chat_history[-1])
+    chat_history[-1].blocks.append(reflection_block)
+    
+    return chat_history
+
+def _format_ui_elements(ui_data, level=0) -> str:
+    """Format UI elements in natural language: index. className: resourceId, text - bounds"""
+    if not ui_data:
+        return ""
+    
+    formatted_lines = []
+    indent = "  " * level  # Indentation for nested elements
+    
+    # Handle both list and single element
+    elements = ui_data if isinstance(ui_data, list) else [ui_data]
+    
+    for element in elements:
+        if not isinstance(element, dict):
+            continue
+            
+        # Extract element properties
+        index = element.get('index', '')
+        class_name = element.get('className', '')
+        resource_id = element.get('resourceId', '')
+        text = element.get('text', '')
+        bounds = element.get('bounds', '')
+        children = element.get('children', [])
+        
+        
+        # Format the line: index. className: resourceId, text - bounds
+        line_parts = []
+        if index != '':
+            line_parts.append(f"{index}.")
+        if class_name:
+            line_parts.append(class_name + ":")
+        
+        details = []
+        if resource_id:
+            details.append(f'"{resource_id}"')
+        if text:
+            details.append(f'"{text}"')
+        if details:
+            line_parts.append(", ".join(details))
+        
+        if bounds:
+            line_parts.append(f"- ({bounds})")
+        
+        formatted_line = f"{indent}{' '.join(line_parts)}"
+        formatted_lines.append(formatted_line)
+        
+        # Recursively format children with increased indentation
+        if children:
+            child_formatted = _format_ui_elements(children, level + 1)
+            if child_formatted:
+                formatted_lines.append(child_formatted)
+    
+    return "\n".join(formatted_lines)
+
+async def add_ui_text_block(ui_state: str, chat_history: List[ChatMessage], copy = True) -> List[ChatMessage]:
     """Add UI elements to the chat history without modifying the original."""
-    ui_elements = None
-    for i in range(retry):
+    if ui_state:
+        # Parse the JSON and format it in natural language
         try:
-            ui_elements = await tools.get_clickables()
-            if ui_elements:
-                break
-        except Exception as e:
-            if i < 4:
-                logger.warning(f"  - Error getting UI elements: {e}. Retrying...")
-            else:
-                logger.error(f"  - Error getting UI elements: {e}. No UI elements will be sent.")
-    if ui_elements:
-        ui_block = TextBlock(text="\nCurrent Clickable UI elements from the device using the custom TopViewService:\n```json\n" + json.dumps(ui_elements) + "\n```\n")
+            ui_data = json.loads(ui_state) if isinstance(ui_state, str) else ui_state
+            formatted_ui = _format_ui_elements(ui_data)
+            ui_block = TextBlock(text=f"\nCurrent Clickable UI elements from the device in the schema 'index. className: resourceId, text - bounds(x1,y1,x2,y2)':\n{formatted_ui}\n")
+        except (json.JSONDecodeError, TypeError):
+            # Fallback to original format if parsing fails
+            ui_block = TextBlock(text="\nCurrent Clickable UI elements from the device using the custom TopViewService:\n```json\n" + json.dumps(ui_state) + "\n```\n")
+        
         if copy:
             chat_history = chat_history.copy()
             chat_history[-1] = message_copy(chat_history[-1])
         chat_history[-1].blocks.append(ui_block)
     return chat_history
 
-async def add_screenshot_image_block(tools: 'Tools', chat_history: List[ChatMessage], retry: int = 5, copy = True) -> None:
-    screenshot = await take_screenshot(tools, retry)
+async def add_screenshot_image_block(screenshot, chat_history: List[ChatMessage], copy = True) -> None:
     if screenshot:
         image_block = ImageBlock(image=base64.b64encode(screenshot))
         if copy:
@@ -55,38 +128,185 @@ async def add_screenshot_image_block(tools: 'Tools', chat_history: List[ChatMess
     return chat_history
 
 
-async def take_screenshot(tools: 'Tools', retry: int = 5) -> None:
-    """Take a screenshot and return the image."""
-    # Retry taking screenshot
-    tools.last_screenshot = None
-    for i in range(retry):
-        try:
-            await tools.take_screenshot()
-            if tools.last_screenshot:
-                break
-        except Exception as e:
-            if i < 4:
-                logger.warning(f"  - Error taking screenshot: {e}. Retrying...")
-            else:
-                logger.error(f"  - Error taking screenshot: {e}. No screenshot will be sent.")
-                return None
-    screenshot = tools.last_screenshot
-    tools.last_screenshot = None # Reset last screenshot after taking it
-    return screenshot
-
-async def add_screenshot(chat_history: List[ChatMessage], screenshot, copy = True) -> List[ChatMessage]:
-    """Add a screenshot to the chat history."""
-    image_block = ImageBlock(image=base64.b64encode(screenshot))
-    if copy:
-        chat_history[-1] = message_copy(chat_history[-1])
-        chat_history = chat_history.copy()  # Create a copy of chat history to avoid modifying the original
-    chat_history[-1].blocks.append(image_block)
+async def add_phone_state_block(phone_state, chat_history: List[ChatMessage]) -> List[ChatMessage]:
+    
+    # Format the phone state data nicely
+    if isinstance(phone_state, dict) and 'error' not in phone_state:
+        current_app = phone_state.get('currentApp', 'Unknown')
+        package_name = phone_state.get('packageName', 'Unknown')
+        keyboard_visible = phone_state.get('keyboardVisible', False)
+        focused_element = phone_state.get('focusedElement')
+        
+        # Format the focused element
+        if focused_element:
+            element_text = focused_element.get('text', 'No text')
+            element_class = focused_element.get('className', 'Unknown')
+            element_bounds = focused_element.get('bounds', 'Unknown')
+            element_type = focused_element.get('type', 'unknown')
+            element_resource_id = focused_element.get('resourceId', '')
+            
+            # Build focused element description
+            focused_desc = f"'{element_text}' ({element_class})"
+            if element_resource_id:
+                focused_desc += f" | ID: {element_resource_id}"
+            focused_desc += f" | Bounds: {element_bounds} | Type: {element_type}"
+        else:
+            focused_desc = "None"
+        
+        phone_state_text = f"""
+**Current Phone State:**
+• **App:** {current_app} ({package_name})
+• **Keyboard:** {'Visible' if keyboard_visible else 'Hidden'}
+• **Focused Element:** {focused_desc}
+        """
+    else:
+        # Handle error cases or malformed data
+        if isinstance(phone_state, dict) and 'error' in phone_state:
+            phone_state_text = f"\n📱 **Phone State Error:** {phone_state.get('message', 'Unknown error')}\n"
+        else:
+            phone_state_text = f"\n📱 **Phone State:** {phone_state}\n"
+    
+    ui_block = TextBlock(text=phone_state_text)
+    chat_history = chat_history.copy()
+    chat_history[-1] = message_copy(chat_history[-1])
+    chat_history[-1].blocks.append(ui_block)
     return chat_history
 
-async def add_phone_state_block(tools: 'Tools', chat_history: List[ChatMessage]) -> List[ChatMessage]:
-    phone_state = await tools.get_phone_state()
-    ui_block = TextBlock(text=f"\nCurrent Phone state: {phone_state}\n```\n")
+async def add_packages_block(packages, chat_history: List[ChatMessage]) -> List[ChatMessage]:
+    
+    ui_block = TextBlock(text=f"\nInstalled packages: {packages}\n```\n")
     chat_history = chat_history.copy()
     chat_history[-1] = message_copy(chat_history[-1])
     chat_history[-1].blocks.append(ui_block)
-    return chat_history
+    return chat_history
+
+async def add_memory_block(memory: List[str], chat_history: List[ChatMessage]) -> List[ChatMessage]:
+    memory_block = "\n### Remembered Information:\n"
+    for idx, item in enumerate(memory, 1):
+        memory_block += f"{idx}. {item}\n"
+    
+    for i, msg in enumerate(chat_history):
+        if msg.role == "user":
+            if isinstance(msg.content, str):
+                updated_content = f"{memory_block}\n\n{msg.content}"
+                chat_history[i] = ChatMessage(role="user", content=updated_content)
+            elif isinstance(msg.content, list):
+                memory_text_block = TextBlock(text=memory_block)
+                content_blocks = [memory_text_block] + msg.content
+                chat_history[i] = ChatMessage(role="user", content=content_blocks)
+            break
+    return chat_history
+
+async def get_reflection_block(reflections: List[Reflection]) -> ChatMessage:
+    reflection_block = "\n### You also have additional Knowledge to help you guide your current task from previous expierences:\n"
+    for reflection in reflections:
+        reflection_block += f"**{reflection.advice}\n"
+    
+    return ChatMessage(role="user", content=reflection_block)
+        
+async def add_task_history_block(completed_tasks: list[dict], failed_tasks: list[dict], chat_history: List[ChatMessage]) -> List[ChatMessage]: 
+    task_history = ""
+
+    # Combine all tasks and show in chronological order
+    all_tasks = completed_tasks + failed_tasks
+    
+    if all_tasks:
+        task_history += "Task History (chronological order):\n"
+        for i, task in enumerate(all_tasks, 1):
+            if hasattr(task, 'description'):
+                status_indicator = "[success]" if hasattr(task, 'status') and task.status == "completed" else "[failed]"
+                task_history += f"{i}. {status_indicator} {task.description}\n"
+            elif isinstance(task, dict):
+                # For backward compatibility with dict format
+                task_description = task.get('description', str(task))
+                status_indicator = "[success]" if task in completed_tasks else "[failed]"
+                task_history += f"{i}. {status_indicator} {task_description}\n"
+            else:
+                status_indicator = "[success]" if task in completed_tasks else "[failed]"
+                task_history += f"{i}. {status_indicator} {task}\n"
+
+    
+    task_block = TextBlock(text=f"{task_history}")
+
+    chat_history = chat_history.copy()
+    chat_history[-1] = message_copy(chat_history[-1])
+    chat_history[-1].blocks.append(task_block)
+    return chat_history
+
+def parse_tool_descriptions(tool_list) -> str:
+    """Parses the available tools and their descriptions for the system prompt."""
+    logger.info("🛠️  Parsing tool descriptions...")
+    tool_descriptions = []
+    
+    for tool in tool_list.values():
+        assert callable(tool), f"Tool {tool} is not callable."
+        tool_name = tool.__name__
+        tool_signature = inspect.signature(tool)
+        tool_docstring = tool.__doc__ or "No description available."
+        formatted_signature = f"def {tool_name}{tool_signature}:\n    \"\"\"{tool_docstring}\"\"\"\n..."
+        tool_descriptions.append(formatted_signature)
+        logger.debug(f"  - Parsed tool: {tool_name}")
+    descriptions = "\n".join(tool_descriptions)
+    logger.info(f"🔩 Found {len(tool_descriptions)} tools.")
+    return descriptions
+
+
+def parse_persona_description(personas) -> str:
+    """Parses the available agent personas and their descriptions for the system prompt."""
+    logger.debug("👥 Parsing agent persona descriptions for Planner Agent...")
+    
+    if not personas:
+        logger.warning("No agent personas provided to Planner Agent")
+        return "No specialized agents available."
+    
+    persona_descriptions = []
+    for persona in personas:
+        # Format each persona with name, description, and expertise areas
+        expertise_list = ", ".join(persona.expertise_areas) if persona.expertise_areas else "General tasks"
+        formatted_persona = f"- **{persona.name}**: {persona.description}\n  Expertise: {expertise_list}"
+        persona_descriptions.append(formatted_persona)
+        logger.debug(f"  - Parsed persona: {persona.name}")
+    
+    # Join all persona descriptions into a single string
+    descriptions = "\n".join(persona_descriptions)
+    logger.debug(f"👤 Found {len(persona_descriptions)} agent personas.")
+    return descriptions
+
+
+def extract_code_and_thought(response_text: str) -> Tuple[Optional[str], str]:
+    """
+    Extracts code from Markdown blocks (```python ... ```) and the surrounding text (thought),
+    handling indented code blocks.
+
+    Returns:
+        Tuple[Optional[code_string], thought_string]
+    """
+    logger.debug("✂️ Extracting code and thought from response...")
+    code_pattern = r"^\s*```python\s*\n(.*?)\n^\s*```\s*?$"
+    code_matches = list(re.finditer(code_pattern, response_text, re.DOTALL | re.MULTILINE))
+
+    if not code_matches:
+        logger.debug("  - No code block found. Entire response is thought.")
+        return None, response_text.strip()
+
+    extracted_code_parts = []
+    for match in code_matches:
+            code_content = match.group(1)
+            extracted_code_parts.append(code_content)
+
+    extracted_code = "\n\n".join(extracted_code_parts)
+
+
+    thought_parts = []
+    last_end = 0
+    for match in code_matches:
+        start, end = match.span(0)
+        thought_parts.append(response_text[last_end:start])
+        last_end = end
+    thought_parts.append(response_text[last_end:])
+
+    thought_text = "".join(thought_parts).strip()
+    thought_preview = (thought_text[:100] + '...') if len(thought_text) > 100 else thought_text
+    logger.debug(f"  - Extracted thought: {thought_preview}")
+
+    return extracted_code, thought_text