android-midscene-automation 0.2.3 → 1.1.0

package/artemis-main/artemis/agents/video_analyzer/chunk_conversation.py

@@ -0,0 +1,491 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Native Gemini sub-agent conversation and submit_answer handling for chunks.
+
+Extracted from ``video_analyzer.py`` as a pure structural split.  Patchable
+collaborators are looked up late through the facade module so
+``mock.patch("artemis.agents.video_analyzer.video_analyzer.<name>")`` targets
+keep working.
+"""
+
+import asyncio
+import glob
+from pathlib import Path
+import re
+
+from google.genai import types
+
+from artemis.agents.video_analyzer import gemini_files
+from artemis.agents.video_analyzer import video_analyzer as _va
+from artemis.agents.video_analyzer.reliability import SubAgentAnswerExhausted
+from artemis.constants import SAFETY_SETTINGS_BLOCK_NONE
+from artemis.data_engine.trace import CURRENT_TRACE_ID, TraceSpan
+from artemis.llm.google import normalize_usage
+from third_party.mobile_use.utils.logger import get_logger
+
+logger = get_logger(__name__)
+
+#: Turns the chunk sub-agent gets to produce an accepted ``submit_answer``
+#: (shared by the static and agentic paths).
+SUB_AGENT_MAX_TURNS = 2
+
+_SUBMIT_NOW_WARNING = "[WARNING] You must call the submit_answer tool now, time is running out."
+_WRONG_TOOL_ERROR = "Tool not recognized. Please use submit_answer."
+_NO_TOOL_CALL_ERROR = "The model replied with text only and never called submit_answer."
+_MISSING_SUMMARY_ERROR = (
+    "submit_answer was called without a 'summary'. Call it again with a non-empty"
+    " 'summary' (and 'analysis') describing what you observed."
+)
+
+
+def _tool_name(name: str | None) -> str:
+    name = name or ""
+    return name.split(":")[-1] if ":" in name else name
+
+
+def _clean_text(value) -> str:
+    return str(value or "").strip().replace("\n", " ")
+
+
+def _validate_answer(
+    summary: str, timeline_events: list, start_time: float, end_time: float | None
+) -> str | None:
+    """Return a validation error for submit_answer, or None if valid."""
+    if not summary:
+        return _MISSING_SUMMARY_ERROR
+    return _validate_timeline_events(timeline_events, start_time, end_time)
+
+
+class _ChunkMedia:
+    """Mutable holder for per-attempt media artifacts shared with failure handling."""
+
+    __slots__ = ("path", "compressed_path", "actual_start", "prompt_with_context")
+
+    def __init__(self):
+        self.path: Path | None = None
+        self.compressed_path: Path | None = None
+        self.actual_start: float | None = None
+        self.prompt_with_context: str | None = None
+
+
+async def _run_native_chunk_conversation(
+    analyzer,
+    media: _ChunkMedia,
+    current_start: float,
+    current_end: float,
+    start_time: float,
+    end_time: float | None,
+    specific_query: str,
+    lease_owner: str,
+) -> str:
+    """Uploads the chunk and drives the native Gemini sub-agent to a committed answer."""
+    file = None
+    try:
+        with TraceSpan(name="upload_video_to_gemini") as span:
+            file = await analyzer.upload_and_poll_file(media.compressed_path)
+            span.result = f"Uploaded {file.name}"
+        logger.info(
+            f"Invoking Gemini for sub-agent task with model {analyzer.model_name} (streaming)..."
+        )
+        trace_id = CURRENT_TRACE_ID.get()
+
+        if analyzer.ctx.data_engine and trace_id:
+            analyzer.ctx.data_engine.record_trace(
+                type="llm_call",
+                name="video_sub_agent",
+                payload={
+                    "contents": [
+                        f"file://{file.name}",
+                        media.prompt_with_context,
+                    ],
+                    "system_instruction": analyzer.sub_system_prompt,
+                },
+                parent_trace_id=trace_id,
+            )
+
+        sub_agent_contents = [
+            types.Content(
+                role="user",
+                parts=[
+                    types.Part.from_uri(file_uri=file.uri, mime_type=file.mime_type)
+                    if hasattr(file, "uri")
+                    else file,
+                    types.Part.from_text(text=media.prompt_with_context),
+                ],
+            )
+        ]
+
+        final_summary, final_analysis = await _drive_sub_agent_loop(
+            analyzer,
+            sub_agent_contents,
+            trace_id,
+            media,
+            start_time,
+            end_time,
+            specific_query,
+        )
+
+        end_str = f"{current_end:.1f}s" if current_end is not None else "unknown"
+        chunk_result = (
+            f"[from {current_start:.1f}s to {end_str}] Summary:"
+            f" {final_summary} Analysis: {final_analysis}"
+        )
+        analyzer.blackboard.complete_segment(
+            current_start,
+            current_end,
+            specific_query,
+            lease_owner,
+            final_summary,
+            final_analysis,
+        )
+        return chunk_result
+    finally:
+        if file:
+            await gemini_files.delete_cloud_file(
+                analyzer.client, file.name, analyzer.cloud_files_to_cleanup
+            )
+
+
+async def _drive_sub_agent_loop(
+    analyzer,
+    sub_agent_contents: list,
+    trace_id,
+    media: _ChunkMedia,
+    start_time: float,
+    end_time: float | None,
+    specific_query: str,
+) -> tuple[str, str]:
+    """Return a validated answer and persist its timeline events.
+
+    Raise ``SubAgentAnswerExhausted`` if no answer is accepted within the turn limit.
+    """
+    rejection = _NO_TOOL_CALL_ERROR
+
+    for turn in range(1, SUB_AGENT_MAX_TURNS + 1):
+        if turn == SUB_AGENT_MAX_TURNS - 1:
+            sub_agent_contents.append(
+                types.Content(
+                    role="user",
+                    parts=[types.Part.from_text(text=_SUBMIT_NOW_WARNING)],
+                )
+            )
+
+        full_text, function_calls, original_parts = await _stream_sub_agent_turn(
+            analyzer, sub_agent_contents, trace_id
+        )
+
+        if function_calls:
+            answer, rejection = await _handle_sub_agent_calls(
+                analyzer,
+                sub_agent_contents,
+                function_calls,
+                original_parts,
+                full_text,
+                media,
+                start_time,
+                end_time,
+                specific_query,
+            )
+            if answer is not None:
+                return answer
+        else:
+            rejection = _NO_TOOL_CALL_ERROR
+            sub_agent_contents.append(
+                types.Content(
+                    role="model",
+                    parts=[types.Part.from_text(text=full_text)],
+                )
+            )
+
+    raise SubAgentAnswerExhausted(SUB_AGENT_MAX_TURNS, rejection)
+
+
+async def _stream_sub_agent_turn(
+    analyzer, sub_agent_contents: list, trace_id
+) -> tuple[str, list, list]:
+    """Starts one streamed Gemini turn and consumes it fully."""
+    with TraceSpan(name="gemini_stream_content_main") as span:
+        stream = await asyncio.wait_for(
+            analyzer.client.aio.models.generate_content_stream(
+                model=analyzer.model_name,
+                contents=sub_agent_contents,
+                config=types.GenerateContentConfig(
+                    system_instruction=analyzer.sub_system_prompt,
+                    tools=[types.Tool(function_declarations=[analyzer.submit_answer_tool])],
+                    safety_settings=SAFETY_SETTINGS_BLOCK_NONE,
+                ),
+            ),
+            timeout=60,
+        )
+
+    full_text, function_calls, original_parts = await asyncio.wait_for(
+        _read_sub_agent_stream(analyzer, span, stream, trace_id),
+        timeout=analyzer.model_call_timeout_seconds,
+    )
+    return full_text, function_calls, original_parts
+
+
+async def _read_sub_agent_stream(analyzer, span, target_stream, current_trace_id):
+    """Consumes a streamed response, mirroring text/thoughts into the trace."""
+    text = ""
+    function_calls = []
+    accumulated_parts = []
+    async for chunk in target_stream:
+        usage = normalize_usage(getattr(chunk, "usage_metadata", None))
+        if usage:
+            span.payload["usage_metadata"] = usage
+
+        chunk_text = chunk.text or ""
+        if chunk_text:
+            text += chunk_text
+            if analyzer.ctx.data_engine and current_trace_id:
+                analyzer.ctx.data_engine.stream_output(current_trace_id, chunk_text)
+        if chunk.function_calls:
+            function_calls.extend(chunk.function_calls)
+
+        candidates = getattr(chunk, "candidates", None)
+        if candidates and isinstance(candidates, list) and len(candidates) > 0:
+            content = getattr(candidates[0], "content", None)
+            parts = getattr(content, "parts", None) if content else None
+            if parts:
+                for part in parts:
+                    accumulated_parts.append(part)
+                    if getattr(part, "thought", False) and part.text:
+                        if analyzer.ctx.data_engine and current_trace_id:
+                            analyzer.ctx.data_engine.stream_output(
+                                current_trace_id,
+                                part.text,
+                                is_thinking=True,
+                            )
+        else:
+            if chunk.function_calls:
+                for fc in chunk.function_calls:
+                    accumulated_parts.append(types.Part(function_call=fc))
+
+    return text, function_calls, accumulated_parts
+
+
+async def _handle_sub_agent_calls(
+    analyzer,
+    sub_agent_contents: list,
+    function_calls: list,
+    original_parts: list,
+    full_text: str,
+    media: _ChunkMedia,
+    start_time: float,
+    end_time: float | None,
+    specific_query: str,
+) -> tuple[tuple[str, str] | None, str | None]:
+    """Return ``(answer, rejection)`` after processing one turn's tool calls.
+
+    Persist events for accepted answers; append rejected calls to the conversation
+    with an error for the model to correct.
+    """
+    for fc in function_calls:
+        if _tool_name(fc.name) != "submit_answer":
+            continue
+        args = dict(fc.args or {})
+        timeline_events = args.get("timeline_events", [])
+        if not isinstance(timeline_events, list):
+            timeline_events = []
+        summary = _clean_text(args.get("summary"))
+        analysis = _clean_text(args.get("analysis"))
+
+        rejection = _validate_answer(summary, timeline_events, start_time, end_time)
+        if rejection is not None:
+            _append_rejection(
+                sub_agent_contents, original_parts, function_calls, full_text, fc.name, rejection
+            )
+            return None, rejection
+
+        await _persist_timeline_events(
+            analyzer, timeline_events, media, start_time, end_time, specific_query
+        )
+        return (summary, analysis or "No analysis provided."), None
+
+    _append_rejection(
+        sub_agent_contents,
+        original_parts,
+        function_calls,
+        full_text,
+        function_calls[0].name,
+        _WRONG_TOOL_ERROR,
+    )
+    return None, _WRONG_TOOL_ERROR
+
+
+def _validate_timeline_events(
+    timeline_events: list, start_time: float, end_time: float | None
+) -> str | None:
+    """Returns the first validation error message for the submitted events, if any."""
+    for event in timeline_events:
+        if (
+            "confidence_score" not in event
+            or not isinstance(
+                event["confidence_score"],
+                (int, float),
+            )
+            or isinstance(event["confidence_score"], bool)
+            or not (0.0 <= float(event["confidence_score"]) <= 1.0)
+        ):
+            return (
+                "Error: 'confidence_score' is"
+                " missing or out of bounds in"
+                " one of the timeline_events."
+                " It must be a float between"
+                " 0.0 and 1.0 inclusive. Fix"
+                " the payload and call"
+                " submit_answer again."
+            )
+
+        if (
+            "verification_timestamp_secs" not in event
+            or not isinstance(
+                event["verification_timestamp_secs"],
+                (int, float),
+            )
+            or isinstance(event["verification_timestamp_secs"], bool)
+        ):
+            return (
+                "Error:"
+                " 'verification_timestamp_secs'"
+                " is missing or invalid in one"
+                " of the timeline_events. Fix"
+                " the payload and call"
+                " submit_answer again."
+            )
+
+        verification_ts = float(
+            event.get(
+                "verification_timestamp_secs",
+                start_time,
+            )
+        )
+        confidence = float(event.get("confidence_score", 0.0))  # noqa: F841
+
+        if verification_ts < start_time or (end_time is not None and verification_ts > end_time):
+            return (
+                "Error:"
+                " 'verification_timestamp_secs'"
+                f" {verification_ts} must fall"
+                " within the start_time"
+                f" ({start_time}) and end_time"
+                f" ({end_time}) boundaries. Fix"
+                " the payload and call"
+                " submit_answer again."
+            )
+    return None
+
+
+def _append_rejection(
+    sub_agent_contents: list,
+    original_parts: list,
+    function_calls: list,
+    full_text: str,
+    response_name: str,
+    error_msg: str,
+) -> None:
+    """Appends the model turn plus an error function response to the conversation."""
+    sub_agent_contents.append(
+        types.Content(
+            role="model",
+            parts=original_parts
+            if original_parts
+            else [types.Part(function_call=fc) for fc in function_calls]
+            + ([types.Part.from_text(text=full_text)] if full_text else []),
+        )
+    )
+    sub_agent_contents.append(
+        types.Content(
+            role="user",
+            parts=[
+                types.Part.from_function_response(
+                    name=response_name,
+                    response={"error": (error_msg)},
+                )
+            ],
+        )
+    )
+
+
+async def _persist_timeline_events(
+    analyzer,
+    timeline_events: list,
+    media: _ChunkMedia,
+    start_time: float,
+    end_time: float | None,
+    specific_query: str,
+) -> None:
+    """Captures verification screenshots and records validated events on the blackboard."""
+    for event in timeline_events:
+        verification_ts = float(
+            event.get(
+                "verification_timestamp_secs",
+                start_time,
+            )
+        )
+        confidence = float(event.get("confidence_score", 0.0))
+        relative_offset = max(0.0, verification_ts - media.actual_start)
+
+        out_dir = _va.get_temp_dir("video_analyzer")
+        out_dir.mkdir(parents=True, exist_ok=True)
+        existing_imgs = glob.glob(f"{out_dir}/img_*.jpg")
+        max_idx = -1
+        for img in existing_imgs:
+            match = re.search(r"img_(\d+)\.jpg$", img)
+            if match:
+                max_idx = max(max_idx, int(match.group(1)))
+        next_idx = max_idx + 1
+        screenshot_output_path = out_dir / f"img_{next_idx}.jpg"
+        while screenshot_output_path.exists():
+            next_idx += 1
+            screenshot_output_path = out_dir / f"img_{next_idx}.jpg"
+        screenshot_output_path.touch()
+        analyzer.local_files_to_cleanup.add(screenshot_output_path)
+        ffmpeg_cmd = [
+            "ffmpeg",
+            "-ss",
+            str(relative_offset),
+            "-i",
+            str(media.path),
+            "-vframes",
+            "1",
+            "-q:v",
+            "2",
+            "-y",
+            str(screenshot_output_path),
+        ]
+        try:
+            proc = await asyncio.create_subprocess_exec(
+                *ffmpeg_cmd,
+                stdout=asyncio.subprocess.PIPE,
+                stderr=asyncio.subprocess.PIPE,
+            )
+            await proc.communicate()
+        except Exception as e:
+            logger.warning(f"Failed to extract frame: {e}")
+
+        screenshot_path_str = str(screenshot_output_path)
+        entry = {
+            "start": event.get("start_time", start_time),
+            "end": event.get("end_time", end_time),
+            "target": specific_query,
+            "summary": event.get("transcription", ""),
+            "confidence_score": confidence,
+        }
+        if Path(screenshot_output_path).exists():
+            entry["screenshot"] = str(screenshot_path_str)
+
+        analyzer._record_blackboard_entry(entry)