code-puppy 0.0.214__py3-none-any.whl → 0.0.366__py3-none-any.whl

code_puppy/tools/browser/browser_screenshot.py

@@ -1,19 +1,21 @@
-"""Screenshot and visual analysis tool with VQA capabilities."""
+"""Screenshot tool for browser automation.
 
-import asyncio
+Captures screenshots and returns them via ToolReturn with BinaryContent
+so multimodal models can directly see and analyze - no separate VQA agent needed.
+"""
+
+import time
 from datetime import datetime
 from pathlib import Path
 from tempfile import gettempdir, mkdtemp
-from typing import Any, Dict, Optional
+from typing import Any, Dict, Optional, Union
 
-from pydantic import BaseModel
-from pydantic_ai import RunContext
+from pydantic_ai import BinaryContent, RunContext, ToolReturn
 
-from code_puppy.messaging import emit_error, emit_info
+from code_puppy.messaging import emit_error, emit_info, emit_success
 from code_puppy.tools.common import generate_group_id
 
-from .camoufox_manager import get_camoufox_manager
-from .vqa_agent import run_vqa_analysis
+from .browser_manager import get_session_browser_manager
 
 _TEMP_SCREENSHOT_ROOT = Path(
     mkdtemp(prefix="code_puppy_screenshots_", dir=gettempdir())
@@ -21,21 +23,11 @@ _TEMP_SCREENSHOT_ROOT = Path(
 
 
 def _build_screenshot_path(timestamp: str) -> Path:
-    """Return the target path for a screenshot using a shared temp directory."""
+    """Return the target path for a screenshot."""
     filename = f"screenshot_{timestamp}.png"
     return _TEMP_SCREENSHOT_ROOT / filename
 
 
-class ScreenshotResult(BaseModel):
-    """Result from screenshot operation."""
-
-    success: bool
-    screenshot_path: Optional[str] = None
-    screenshot_data: Optional[bytes] = None
-    timestamp: Optional[str] = None
-    error: Optional[str] = None
-
-
 async def _capture_screenshot(
     page,
     full_page: bool = False,
@@ -45,41 +37,37 @@ async def _capture_screenshot(
 ) -> Dict[str, Any]:
     """Internal screenshot capture function."""
     try:
-        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
 
         # Take screenshot
         if element_selector:
-            # Screenshot specific element
             element = await page.locator(element_selector).first
             if not await element.is_visible():
                 return {
                     "success": False,
                     "error": f"Element '{element_selector}' is not visible",
                 }
-            screenshot_data = await element.screenshot()
+            screenshot_bytes = await element.screenshot()
         else:
-            # Screenshot page or full page
-            screenshot_data = await page.screenshot(full_page=full_page)
+            screenshot_bytes = await page.screenshot(full_page=full_page)
 
-        result = {
+        result: Dict[str, Any] = {
             "success": True,
-            "screenshot_data": screenshot_data,
+            "screenshot_bytes": screenshot_bytes,
             "timestamp": timestamp,
         }
 
         if save_screenshot:
             screenshot_path = _build_screenshot_path(timestamp)
             screenshot_path.parent.mkdir(parents=True, exist_ok=True)
-
             with open(screenshot_path, "wb") as f:
-                f.write(screenshot_data)
-
+                f.write(screenshot_bytes)
             result["screenshot_path"] = str(screenshot_path)
-            message = f"[green]Screenshot saved: {screenshot_path}[/green]"
+
             if group_id:
-                emit_info(message, message_group=group_id)
-            else:
-                emit_info(message)
+                emit_success(
+                    f"Screenshot saved: {screenshot_path}", message_group=group_id
+                )
 
         return result
 
@@ -87,46 +75,42 @@ async def _capture_screenshot(
         return {"success": False, "error": str(e)}
 
 
-async def take_screenshot_and_analyze(
-    question: str,
+async def take_screenshot(
     full_page: bool = False,
     element_selector: Optional[str] = None,
     save_screenshot: bool = True,
-) -> Dict[str, Any]:
-    """
-    Take a screenshot and analyze it using visual understanding.
+) -> Union[ToolReturn, Dict[str, Any]]:
+    """Take a screenshot of the browser page.
+
+    Returns a ToolReturn with BinaryContent so multimodal models can
+    directly see and analyze the screenshot.
 
     Args:
-        question: The specific question to ask about the screenshot
-        full_page: Whether to capture the full page or just viewport
-        element_selector: Optional selector to screenshot just a specific element
-        save_screenshot: Whether to save the screenshot to disk
+        full_page: Whether to capture full page or just viewport.
+        element_selector: Optional selector to screenshot specific element.
+        save_screenshot: Whether to save the screenshot to disk.
 
     Returns:
-        Dict containing analysis results and screenshot info
+        ToolReturn containing:
+            - return_value: Success message with screenshot path
+            - content: List with description and BinaryContent image
+            - metadata: Screenshot details (path, target, timestamp)
+        Or Dict with error info if failed.
     """
     target = element_selector or ("full_page" if full_page else "viewport")
-    group_id = generate_group_id(
-        "browser_screenshot_analyze", f"{question[:50]}_{target}"
-    )
-    emit_info(
-        f"[bold white on blue] BROWSER SCREENSHOT ANALYZE [/bold white on blue] 📷 question='{question[:100]}{'...' if len(question) > 100 else ''}' target={target}",
-        message_group=group_id,
-    )
+    group_id = generate_group_id("browser_screenshot", target)
+    emit_info(f"BROWSER SCREENSHOT 📷 target={target}", message_group=group_id)
+
     try:
-        # Get the current browser page
-        browser_manager = get_camoufox_manager()
+        browser_manager = get_session_browser_manager()
         page = await browser_manager.get_current_page()
 
         if not page:
-            return {
-                "success": False,
-                "error": "No active browser page available. Please navigate to a webpage first.",
-                "question": question,
-            }
+            error_msg = "No active browser page. Navigate to a webpage first."
+            emit_error(error_msg, message_group=group_id)
+            return {"success": False, "error": error_msg}
 
-        # Take screenshot
-        screenshot_result = await _capture_screenshot(
+        result = await _capture_screenshot(
             page,
             full_page=full_page,
             element_selector=element_selector,
@@ -134,110 +118,62 @@ async def take_screenshot_and_analyze(
             group_id=group_id,
         )
 
-        if not screenshot_result["success"]:
-            error_message = screenshot_result.get("error", "Screenshot failed")
-            emit_error(
-                f"[red]Screenshot capture failed: {error_message}[/red]",
-                message_group=group_id,
-            )
-            return {
-                "success": False,
-                "error": error_message,
-                "question": question,
-            }
-
-        screenshot_bytes = screenshot_result.get("screenshot_data")
-        if not screenshot_bytes:
-            emit_error(
-                "[red]Screenshot captured but pixel data missing; cannot run visual analysis.[/red]",
-                message_group=group_id,
-            )
-            return {
-                "success": False,
-                "error": "Screenshot captured but no image bytes available for analysis.",
-                "question": question,
-            }
-
-        try:
-            vqa_result = await asyncio.to_thread(
-                run_vqa_analysis,
-                question,
-                screenshot_bytes,
-            )
-        except Exception as exc:
-            emit_error(
-                f"[red]Visual question answering failed: {exc}[/red]",
-                message_group=group_id,
-            )
-            return {
-                "success": False,
-                "error": f"Visual analysis failed: {exc}",
-                "question": question,
-                "screenshot_info": {
-                    "path": screenshot_result.get("screenshot_path"),
-                    "timestamp": screenshot_result.get("timestamp"),
-                    "full_page": full_page,
-                    "element_selector": element_selector,
-                },
-            }
-
-        emit_info(
-            f"[green]Visual analysis answer: {vqa_result.answer}[/green]",
-            message_group=group_id,
-        )
-        emit_info(
-            f"[dim]Observations: {vqa_result.observations}[/dim]",
-            message_group=group_id,
-        )
-
-        return {
-            "success": True,
-            "question": question,
-            "answer": vqa_result.answer,
-            "confidence": vqa_result.confidence,
-            "observations": vqa_result.observations,
-            "screenshot_info": {
-                "path": screenshot_result.get("screenshot_path"),
-                "size": len(screenshot_bytes),
-                "timestamp": screenshot_result.get("timestamp"),
+        if not result["success"]:
+            emit_error(result.get("error", "Screenshot failed"), message_group=group_id)
+            return {"success": False, "error": result.get("error")}
+
+        screenshot_path = result.get("screenshot_path", "(not saved)")
+
+        # Return as ToolReturn with BinaryContent so the model can SEE the image!
+        return ToolReturn(
+            return_value=f"Screenshot captured successfully. Saved to: {screenshot_path}",
+            content=[
+                f"Here's the browser screenshot ({target}):",
+                BinaryContent(
+                    data=result["screenshot_bytes"],
+                    media_type="image/png",
+                ),
+                "Please analyze what you see and describe any relevant details.",
+            ],
+            metadata={
+                "success": True,
+                "screenshot_path": screenshot_path,
+                "target": target,
                 "full_page": full_page,
                 "element_selector": element_selector,
+                "timestamp": time.time(),
             },
-        }
+        )
 
     except Exception as e:
-        emit_info(
-            f"[red]Screenshot analysis failed: {str(e)}[/red]", message_group=group_id
-        )
-        return {"success": False, "error": str(e), "question": question}
+        error_msg = f"Screenshot failed: {str(e)}"
+        emit_error(error_msg, message_group=group_id)
+        return {"success": False, "error": error_msg}
 
 
 def register_take_screenshot_and_analyze(agent):
-    """Register the screenshot analysis tool."""
+    """Register the screenshot tool."""
 
     @agent.tool
     async def browser_screenshot_analyze(
         context: RunContext,
-        question: str,
         full_page: bool = False,
         element_selector: Optional[str] = None,
-        save_screenshot: bool = True,
-    ) -> Dict[str, Any]:
+    ) -> Union[ToolReturn, Dict[str, Any]]:
         """
-        Take a screenshot and analyze it to answer a specific question.
+        Take a screenshot of the browser page.
+
+        Returns the screenshot via ToolReturn with BinaryContent that you can
+        see directly. Use this to see what's displayed in the browser.
 
         Args:
-            question: The specific question to ask about the screenshot
-            full_page: Whether to capture the full page or just viewport
-            element_selector: Optional CSS/XPath selector to screenshot specific element
-            save_screenshot: Whether to save the screenshot to disk
+            full_page: Capture full page (True) or just viewport (False).
+            element_selector: Optional CSS selector to screenshot specific element.
 
         Returns:
-            Dict with analysis results including answer, confidence, and observations
+            ToolReturn with the screenshot image you can analyze, or error dict.
         """
-        return await take_screenshot_and_analyze(
-            question=question,
+        return await take_screenshot(
             full_page=full_page,
             element_selector=element_selector,
-            save_screenshot=save_screenshot,
         )

code_puppy/tools/browser/browser_scripts.py

@@ -4,10 +4,10 @@ from typing import Any, Dict, Optional
 
 from pydantic_ai import RunContext
 
-from code_puppy.messaging import emit_info
+from code_puppy.messaging import emit_error, emit_info, emit_success
 from code_puppy.tools.common import generate_group_id
 
-from .camoufox_manager import get_camoufox_manager
+from .browser_manager import get_session_browser_manager
 
 
 async def execute_javascript(
@@ -17,29 +17,27 @@ async def execute_javascript(
     """Execute JavaScript code in the browser context."""
     group_id = generate_group_id("browser_execute_js", script[:100])
     emit_info(
-        f"[bold white on blue] BROWSER EXECUTE JS [/bold white on blue] 📜 script='{script[:100]}{'...' if len(script) > 100 else ''}'",
+        f"BROWSER EXECUTE JS 📜 script='{script[:100]}{'...' if len(script) > 100 else ''}'",
         message_group=group_id,
     )
     try:
-        browser_manager = get_camoufox_manager()
+        browser_manager = get_session_browser_manager()
         page = await browser_manager.get_current_page()
 
         if not page:
             return {"success": False, "error": "No active browser page available"}
 
         # Execute JavaScript
-        result = await page.evaluate(script, timeout=timeout)
+        # Note: page.evaluate() does NOT accept a timeout parameter
+        # The timeout arg to this function is kept for API compatibility but unused
+        result = await page.evaluate(script)
 
-        emit_info(
-            "[green]JavaScript executed successfully[/green]", message_group=group_id
-        )
+        emit_success("JavaScript executed successfully", message_group=group_id)
 
         return {"success": True, "script": script, "result": result}
 
     except Exception as e:
-        emit_info(
-            f"[red]JavaScript execution failed: {str(e)}[/red]", message_group=group_id
-        )
+        emit_error(f"JavaScript execution failed: {str(e)}", message_group=group_id)
         return {"success": False, "error": str(e), "script": script}
 
 
@@ -52,11 +50,11 @@ async def scroll_page(
     target = element_selector or "page"
     group_id = generate_group_id("browser_scroll", f"{direction}_{amount}_{target}")
     emit_info(
-        f"[bold white on blue] BROWSER SCROLL [/bold white on blue] 📋 direction={direction} amount={amount} target='{target}'",
+        f"BROWSER SCROLL 📋 direction={direction} amount={amount} target='{target}'",
         message_group=group_id,
     )
     try:
-        browser_manager = get_camoufox_manager()
+        browser_manager = get_session_browser_manager()
         page = await browser_manager.get_current_page()
 
         if not page:
@@ -64,7 +62,7 @@ async def scroll_page(
 
         if element_selector:
             # Scroll specific element
-            element = page.locator(element_selector)
+            element = page.locator(element_selector).first
             await element.scroll_into_view_if_needed()
 
             # Get element's current scroll position and dimensions
@@ -120,9 +118,7 @@ async def scroll_page(
             })
         """)
 
-        emit_info(
-            f"[green]Scrolled {target} {direction}[/green]", message_group=group_id
-        )
+        emit_success(f"Scrolled {target} {direction}", message_group=group_id)
 
         return {
             "success": True,
@@ -148,26 +144,24 @@ async def scroll_to_element(
     """Scroll to bring an element into view."""
     group_id = generate_group_id("browser_scroll_to_element", selector[:100])
     emit_info(
-        f"[bold white on blue] BROWSER SCROLL TO ELEMENT [/bold white on blue] 🎯 selector='{selector}'",
+        f"BROWSER SCROLL TO ELEMENT 🎯 selector='{selector}'",
         message_group=group_id,
     )
     try:
-        browser_manager = get_camoufox_manager()
+        browser_manager = get_session_browser_manager()
         page = await browser_manager.get_current_page()
 
         if not page:
             return {"success": False, "error": "No active browser page available"}
 
-        element = page.locator(selector)
+        element = page.locator(selector).first
         await element.wait_for(state="attached", timeout=timeout)
         await element.scroll_into_view_if_needed()
 
         # Check if element is now visible
         is_visible = await element.is_visible()
 
-        emit_info(
-            f"[green]Scrolled to element: {selector}[/green]", message_group=group_id
-        )
+        emit_success(f"Scrolled to element: {selector}", message_group=group_id)
 
         return {"success": True, "selector": selector, "visible": is_visible}
 
@@ -182,11 +176,11 @@ async def set_viewport_size(
     """Set the viewport size."""
     group_id = generate_group_id("browser_set_viewport", f"{width}x{height}")
     emit_info(
-        f"[bold white on blue] BROWSER SET VIEWPORT [/bold white on blue] 🖥️ size={width}x{height}",
+        f"BROWSER SET VIEWPORT 🖥️ size={width}x{height}",
         message_group=group_id,
     )
     try:
-        browser_manager = get_camoufox_manager()
+        browser_manager = get_session_browser_manager()
         page = await browser_manager.get_current_page()
 
         if not page:
@@ -194,8 +188,8 @@ async def set_viewport_size(
 
         await page.set_viewport_size({"width": width, "height": height})
 
-        emit_info(
-            f"[green]Set viewport size to {width}x{height}[/green]",
+        emit_success(
+            f"Set viewport size to {width}x{height}",
             message_group=group_id,
         )
 
@@ -213,22 +207,20 @@ async def wait_for_element(
     """Wait for an element to reach a specific state."""
     group_id = generate_group_id("browser_wait_for_element", f"{selector[:50]}_{state}")
     emit_info(
-        f"[bold white on blue] BROWSER WAIT FOR ELEMENT [/bold white on blue] ⏱️ selector='{selector}' state={state} timeout={timeout}ms",
+        f"BROWSER WAIT FOR ELEMENT ⏱️ selector='{selector}' state={state} timeout={timeout}ms",
         message_group=group_id,
     )
     try:
-        browser_manager = get_camoufox_manager()
+        browser_manager = get_session_browser_manager()
         page = await browser_manager.get_current_page()
 
         if not page:
             return {"success": False, "error": "No active browser page available"}
 
-        element = page.locator(selector)
+        element = page.locator(selector).first
         await element.wait_for(state=state, timeout=timeout)
 
-        emit_info(
-            f"[green]Element {selector} is now {state}[/green]", message_group=group_id
-        )
+        emit_success(f"Element {selector} is now {state}", message_group=group_id)
 
         return {"success": True, "selector": selector, "state": state}
 
@@ -246,17 +238,17 @@ async def highlight_element(
         "browser_highlight_element", f"{selector[:50]}_{color}"
     )
     emit_info(
-        f"[bold white on blue] BROWSER HIGHLIGHT ELEMENT [/bold white on blue] 🔦 selector='{selector}' color={color}",
+        f"BROWSER HIGHLIGHT ELEMENT 🔦 selector='{selector}' color={color}",
         message_group=group_id,
     )
     try:
-        browser_manager = get_camoufox_manager()
+        browser_manager = get_session_browser_manager()
         page = await browser_manager.get_current_page()
 
         if not page:
             return {"success": False, "error": "No active browser page available"}
 
-        element = page.locator(selector)
+        element = page.locator(selector).first
         await element.wait_for(state="visible", timeout=timeout)
 
         # Add highlight style
@@ -271,9 +263,7 @@ async def highlight_element(
 
         await element.evaluate(highlight_script)
 
-        emit_info(
-            f"[green]Highlighted element: {selector}[/green]", message_group=group_id
-        )
+        emit_success(f"Highlighted element: {selector}", message_group=group_id)
 
         return {"success": True, "selector": selector, "color": color}
 
@@ -285,11 +275,11 @@ async def clear_highlights() -> Dict[str, Any]:
     """Clear all element highlights."""
     group_id = generate_group_id("browser_clear_highlights")
     emit_info(
-        "[bold white on blue] BROWSER CLEAR HIGHLIGHTS [/bold white on blue] 🧹",
+        "BROWSER CLEAR HIGHLIGHTS 🧹",
         message_group=group_id,
     )
     try:
-        browser_manager = get_camoufox_manager()
+        browser_manager = get_session_browser_manager()
         page = await browser_manager.get_current_page()
 
         if not page:
@@ -311,7 +301,7 @@ async def clear_highlights() -> Dict[str, Any]:
 
         count = await page.evaluate(clear_script)
 
-        emit_info(f"[green]Cleared {count} highlights[/green]", message_group=group_id)
+        emit_success(f"Cleared {count} highlights", message_group=group_id)
 
         return {"success": True, "cleared_count": count}