OpenSpace/openspace/agents/visual_analyzer.py

from __future__ import annotations

import json
from typing import TYPE_CHECKING, Any, Dict, List, Optional

from openspace.grounding.core.types import ToolResult
from openspace.platforms.screenshot import ScreenshotClient
from openspace.prompts import GroundingAgentPrompts
from openspace.utils.logging import Logger

if TYPE_CHECKING:
    from openspace.llm import LLMClient

logger = Logger.get_logger(__name__)


class VisualAnalyzer:
    """Handles screenshot capture and LLM-based visual analysis of tool results."""

    def __init__(
        self,
        llm_client: Optional[LLMClient] = None,
        visual_analysis_model: Optional[str] = None,
        visual_analysis_timeout: float = 30.0,
    ) -> None:
        self._llm_client = llm_client
        self._visual_analysis_model = visual_analysis_model
        self._visual_analysis_timeout = visual_analysis_timeout

    async def analyze_tool_result(
        self,
        result: ToolResult,
        tool_name: str,
        tool_call: Dict,
        backend: str,
        task_description: str = "",
    ) -> ToolResult:
        """Callback for LLMClient to handle visual analysis after tool execution."""
        skip_visual_analysis = False
        try:
            arguments = tool_call.function.arguments
            if isinstance(arguments, str):
                args = json.loads(arguments.strip() or "{}")
            else:
                args = arguments

            if isinstance(args, dict) and args.get("skip_visual_analysis"):
                skip_visual_analysis = True
                logger.info(f"Visual analysis skipped for {tool_name} (meta-parameter set by LLM)")
        except Exception as e:
            logger.debug(f"Could not parse tool arguments: {e}")

        if skip_visual_analysis:
            return result

        if backend != "gui":
            return result

        metadata = getattr(result, "metadata", None)
        has_screenshots = metadata and (
            metadata.get("screenshot") or metadata.get("screenshots")
        )

        if not has_screenshots:
            try:
                logger.info(f"No visual data from {tool_name}, capturing screenshot...")
                screenshot_client = ScreenshotClient()
                screenshot_bytes = await screenshot_client.capture()

                if screenshot_bytes:
                    if metadata is None:
                        result.metadata = {}
                        metadata = result.metadata
                    metadata["screenshot"] = screenshot_bytes
                    has_screenshots = True
                    logger.info("Screenshot captured for visual analysis")
                else:
                    logger.warning("Failed to capture screenshot")
            except Exception as e:
                logger.warning(f"Error capturing screenshot: {e}")

        if not has_screenshots:
            logger.debug(f"No visual data available for {tool_name}")
            return result

        return await self._enhance_result(result, tool_name, task_description)

    async def _enhance_result(
        self,
        result: ToolResult,
        tool_name: str,
        task_description: str = "",
    ) -> ToolResult:
        """Enhance tool result with LLM-based visual analysis."""
        import asyncio
        import base64
        import litellm

        try:
            metadata = getattr(result, "metadata", None)
            if not metadata:
                return result

            screenshots_bytes: List[bytes] = []
            if metadata.get("screenshots"):
                screenshots_list = metadata["screenshots"]
                if isinstance(screenshots_list, list):
                    screenshots_bytes = [s for s in screenshots_list if s]
            elif metadata.get("screenshot"):
                screenshots_bytes = [metadata["screenshot"]]

            if not screenshots_bytes:
                return result

            selected_screenshots = self._select_key_screenshots(
                screenshots_bytes, max_count=3
            )

            visual_b64_list = []
            for visual_data in selected_screenshots:
                if isinstance(visual_data, bytes):
                    visual_b64_list.append(
                        base64.b64encode(visual_data).decode("utf-8")
                    )
                else:
                    visual_b64_list.append(visual_data)

            num_screenshots = len(visual_b64_list)

            prompt = GroundingAgentPrompts.visual_analysis(
                tool_name=tool_name,
                num_screenshots=num_screenshots,
                task_description=task_description,
            )

            content: List[Dict[str, Any]] = [{"type": "text", "text": prompt}]
            for visual_b64 in visual_b64_list:
                content.append(
                    {
                        "type": "image_url",
                        "image_url": {"url": f"data:image/png;base64,{visual_b64}"},
                    }
                )

            visual_model = self._visual_analysis_model or (
                self._llm_client.model
                if self._llm_client
                else "openrouter/anthropic/claude-sonnet-4.5"
            )
            _llm_extra: Dict[str, Any] = {}
            if self._llm_client and visual_model == self._llm_client.model:
                _llm_extra = (
                    getattr(self._llm_client, "litellm_kwargs", {}) or {}
                )
            elif self._visual_analysis_model:
                try:
                    from openspace.host_detection import build_llm_kwargs

                    visual_model, _llm_extra = build_llm_kwargs(visual_model)
                except Exception as e:
                    logger.debug(
                        f"Failed to resolve dedicated visual model credentials: {e}"
                    )
                    _llm_extra = {}

            response = await asyncio.wait_for(
                litellm.acompletion(
                    model=visual_model,
                    messages=[{"role": "user", "content": content}],
                    timeout=self._visual_analysis_timeout,
                    **_llm_extra,
                ),
                timeout=self._visual_analysis_timeout + 5,
            )

            analysis = response.choices[0].message.content.strip()

            original_content = result.content or "(no text output)"
            enhanced_content = (
                f"{original_content}\n\n**Visual content**: {analysis}"
            )

            enhanced_result = ToolResult(
                status=result.status,
                content=enhanced_content,
                error=result.error,
                metadata={
                    **metadata,
                    "visual_analyzed": True,
                    "visual_analysis": analysis,
                },
                execution_time=result.execution_time,
            )

            logger.info(
                f"Enhanced {tool_name} result with visual analysis "
                f"({num_screenshots} screenshot(s))"
            )
            return enhanced_result

        except asyncio.TimeoutError:
            logger.warning(
                f"Visual analysis timed out for {tool_name}, returning original result"
            )
            return result
        except Exception as e:
            logger.warning(
                f"Failed to analyze visual content for {tool_name}: {e}"
            )
            return result

    @staticmethod
    def _select_key_screenshots(
        screenshots: List[bytes],
        max_count: int = 3,
    ) -> List[bytes]:
        """Select key screenshots from a sequence, preferring first/last/evenly-spaced."""
        if len(screenshots) <= max_count:
            return screenshots

        selected_indices: set[int] = set()

        selected_indices.add(len(screenshots) - 1)

        if max_count >= 2:
            selected_indices.add(0)

        remaining_slots = max_count - len(selected_indices)
        if remaining_slots > 0:
            available_indices = [
                i
                for i in range(1, len(screenshots) - 1)
                if i not in selected_indices
            ]

            if available_indices:
                step = max(1, len(available_indices) // (remaining_slots + 1))
                for i in range(remaining_slots):
                    idx = min((i + 1) * step, len(available_indices) - 1)
                    if idx < len(available_indices):
                        selected_indices.add(available_indices[idx])

        selected = [screenshots[i] for i in sorted(selected_indices)]

        logger.debug(
            f"Selected {len(selected)} screenshots at indices "
            f"{sorted(selected_indices)} from total of {len(screenshots)}"
        )

        return selected