refactor(video): split video analyzer into engine modules

artemis/agents/video_analyzer/universal_engine.py

@@ -0,0 +1,501 @@
+# 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.
+
+"""Universal (OpenAI-compatible / LangChain) execution paths for the video analyzer.
+
+Extracted from ``video_analyzer.py`` as a pure structural split.  Patchable
+collaborators (``extract_keyframes_from_video``, ``extract_audio_from_video``,
+semaphores, retry helper, ...) are looked up late through the facade module so
+``mock.patch("artemis.agents.video_analyzer.video_analyzer.<name>")`` targets
+keep working.
+"""
+
+import asyncio
+import base64
+import glob
+from pathlib import Path
+import re
+from typing import Any
+
+from langchain_core.messages import (
+    BaseMessage,
+    HumanMessage,
+    SystemMessage,
+    ToolMessage,
+)
+
+from artemis.agents.video_analyzer import video_analyzer as _va
+from artemis.agents.video_analyzer.reliability import classify_video_failure
+from artemis.data_engine.trace import CURRENT_TRACE_ID, TraceSpan
+from artemis.utils.logger import get_logger
+
+logger = get_logger(__name__)
+
+
+async def exec_single_chunk_universal(
+    analyzer,
+    compressed_path: Path,
+    raw_path: Path,
+    start_time: float,
+    end_time: float | None,
+    actual_start: float,
+    prompt_with_context: str,
+    specific_query: str,
+    force_fallback: bool = False,
+) -> str:
+    """Executes sub-agent chunk analysis via keyframe extraction and LangChain ChatModel."""
+    keyframes = _gather_universal_frames(
+        analyzer, raw_path, start_time, end_time, actual_start
+    )
+
+    user_blocks, audio_block = await _build_universal_user_blocks(
+        analyzer, keyframes, actual_start, prompt_with_context, raw_path
+    )
+
+    messages = [
+        SystemMessage(content=analyzer.sub_system_prompt),
+        HumanMessage(content=user_blocks),
+    ]
+
+    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_universal",
+            payload={
+                "num_keyframes": len(keyframes),
+                "prompt": prompt_with_context,
+                "system_instruction": analyzer.sub_system_prompt,
+            },
+            parent_trace_id=trace_id,
+        )
+
+    response = await _invoke_universal_with_visual_retry(
+        analyzer, messages, user_blocks, audio_block, force_fallback
+    )
+
+    final_summary, final_analysis, timeline_events = _extract_submit_payload(response)
+
+    await _persist_universal_events(
+        analyzer,
+        timeline_events,
+        start_time,
+        end_time,
+        actual_start,
+        raw_path,
+        specific_query,
+        final_summary,
+    )
+
+    end_str = f"{end_time:.1f}s" if end_time is not None else "unknown"
+    return (
+        f"[from {start_time:.1f}s to {end_str}] Summary: {final_summary}"
+        f" Analysis: {final_analysis}"
+    )
+
+
+def _gather_universal_frames(
+    analyzer,
+    raw_path: Path,
+    start_time: float,
+    end_time: float | None,
+    actual_start: float,
+) -> list:
+    """Extracts uniform keyframes plus dense action-window frames, merged by timestamp."""
+    # Universal providers consume stills rather than the uploaded clip. Use
+    # the raw extracted segment so slowdown/transcode PTS never skew labels.
+    keyframes = _va.extract_keyframes_from_video(
+        raw_path,
+        fps=1.0,
+        max_frames=30,
+        max_dimension=1080,
+    )
+    dense_offsets = analyzer._dense_action_offsets(
+        start_time,
+        end_time if end_time is not None else start_time,
+        actual_start,
+    )
+    targeted_frames = _va.extract_frames_at_timestamps(
+        raw_path,
+        dense_offsets,
+        max_frames=analyzer.max_dense_action_frames,
+        max_dimension=1080,
+    )
+    merged_frames = {round(timestamp, 3): data for timestamp, data in keyframes}
+    merged_frames.update(
+        {round(timestamp, 3): data for timestamp, data in targeted_frames}
+    )
+    return sorted(merged_frames.items())
+
+
+async def _build_universal_user_blocks(
+    analyzer,
+    keyframes: list,
+    actual_start: float,
+    prompt_with_context: str,
+    raw_path: Path,
+) -> tuple[list[dict[str, Any]], dict[str, Any] | None]:
+    """Builds the multimodal user content blocks (text, frames, optional audio)."""
+    user_blocks: list[dict[str, Any]] = [{"type": "text", "text": prompt_with_context}]
+
+    for ts_sec, frame_bytes in keyframes:
+        b64_str = base64.b64encode(frame_bytes).decode("utf-8")
+        abs_time = actual_start + ts_sec
+        user_blocks.append(
+            {
+                "type": "text",
+                "text": f"--- Video Keyframe at {abs_time:.1f}s ---",
+            }
+        )
+        user_blocks.append(
+            {
+                "type": "image_url",
+                "image_url": {"url": f"data:image/jpeg;base64,{b64_str}"},
+            }
+        )
+
+    audio_block = None
+    try:
+        async with _va.TRANSCODE_SEMAPHORE:
+            audio_path = await _va.extract_audio_from_video(raw_path)
+        analyzer.local_files_to_cleanup.add(audio_path)
+        analyzer.local_dirs_to_cleanup.add(audio_path.parent)
+        audio_block = analyzer._audio_content_block(audio_path)
+        user_blocks.append(audio_block)
+    except Exception as audio_error:
+        logger.info(f"No usable audio track for universal video chunk: {audio_error}")
+
+    return user_blocks, audio_block
+
+
+async def _invoke_universal_with_visual_retry(
+    analyzer,
+    messages: list[BaseMessage],
+    user_blocks: list[dict[str, Any]],
+    audio_block: dict[str, Any] | None,
+    force_fallback: bool,
+):
+    """Invokes the universal model, retrying visually when audio input is rejected."""
+    with TraceSpan(name="universal_sub_agent_call") as span:
+        try:
+            response = await analyzer._invoke_universal_model(
+                messages,
+                [_va.UNIVERSAL_SUBMIT_ANSWER_TOOL],
+                timeout=analyzer.model_call_timeout_seconds,
+                label="Universal video sub-agent",
+                force_fallback=force_fallback,
+            )
+        except Exception as multimodal_error:
+            failure = classify_video_failure(multimodal_error)
+            if audio_block is None or failure.category.value not in {
+                "bad_request",
+                "media_processing",
+            }:
+                raise
+            logger.warning(
+                "Configured video model rejected audio input; retrying this chunk visually"
+            )
+            visual_blocks = [block for block in user_blocks if block is not audio_block]
+            messages = [
+                SystemMessage(content=analyzer.sub_system_prompt),
+                HumanMessage(content=visual_blocks),
+            ]
+            response = await analyzer._invoke_universal_model(
+                messages,
+                [_va.UNIVERSAL_SUBMIT_ANSWER_TOOL],
+                timeout=analyzer.model_call_timeout_seconds,
+                label="Universal visual-only sub-agent",
+                force_fallback=force_fallback,
+            )
+        span.result = (
+            f"Received response (tool_calls={len(getattr(response, 'tool_calls', []))})"
+        )
+    return response
+
+
+def _extract_submit_payload(response) -> tuple[str, str, list]:
+    """Pulls summary/analysis/timeline events from a universal model response."""
+    final_summary = "No summary provided."
+    final_analysis = "No analysis provided."
+    timeline_events = []
+
+    if getattr(response, "tool_calls", None):
+        for tc in response.tool_calls:
+            if tc.get("name") == "submit_answer":
+                args = tc.get("args", {})
+                final_summary = (
+                    str(args.get("summary", "")).strip().replace("\n", " ") or final_summary
+                )
+                final_analysis = (
+                    str(args.get("analysis", "")).strip().replace("\n", " ") or final_analysis
+                )
+                timeline_events = args.get("timeline_events", [])
+                break
+    elif getattr(response, "content", None):
+        text_content = str(response.content).strip()
+        final_summary = text_content[:200].replace("\n", " ")
+        final_analysis = text_content
+
+    return final_summary, final_analysis, timeline_events
+
+
+async def _persist_universal_events(
+    analyzer,
+    timeline_events: list,
+    start_time: float,
+    end_time: float | None,
+    actual_start: float,
+    raw_path: Path,
+    specific_query: str,
+    final_summary: str,
+) -> None:
+    """Processes timeline events & saves screenshots for the blackboard ledger."""
+    for event in timeline_events:
+        if not isinstance(event, dict):
+            continue
+        try:
+            v_ts = float(event.get("verification_timestamp_secs", start_time))
+            conf = float(event.get("confidence_score", 0.8))
+            event_start = float(event.get("start_time", start_time))
+            event_end = float(
+                event.get(
+                    "end_time",
+                    end_time if end_time is not None else event_start,
+                )
+            )
+        except (TypeError, ValueError):
+            logger.warning(f"Ignoring malformed universal video event: {event}")
+            continue
+        interval_end = end_time if end_time is not None else max(event_end, v_ts)
+        if (
+            not 0.0 <= conf <= 1.0
+            or event_end < event_start
+            or event_start < start_time
+            or event_end > interval_end
+            or v_ts < start_time
+            or v_ts > interval_end
+        ):
+            logger.warning(f"Ignoring out-of-range universal video event: {event}")
+            continue
+
+        screenshot_path_str = None
+        try:
+            rel_offset = max(0.0, v_ts - 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:
+                m = re.search(r"img_(\d+)\.jpg$", img)
+                if m:
+                    max_idx = max(max_idx, int(m.group(1)))
+            next_idx = max_idx + 1
+            shot_path = out_dir / f"img_{next_idx}.jpg"
+
+            ffmpeg_cmd = [
+                "ffmpeg",
+                "-ss",
+                str(rel_offset),
+                "-i",
+                str(raw_path),
+                "-vframes",
+                "1",
+                "-q:v",
+                "2",
+                "-y",
+                str(shot_path),
+            ]
+            proc = await asyncio.create_subprocess_exec(
+                *ffmpeg_cmd,
+                stdout=asyncio.subprocess.PIPE,
+                stderr=asyncio.subprocess.PIPE,
+            )
+            await proc.communicate()
+            if shot_path.exists():
+                analyzer.local_files_to_cleanup.add(shot_path)
+                screenshot_path_str = str(shot_path)
+        except Exception as fe:
+            logger.warning(f"Failed extracting keyframe for event {v_ts}s: {fe}")
+
+        entry = {
+            "start": event_start,
+            "end": event_end,
+            "target": specific_query,
+            "summary": event.get("transcription", final_summary),
+            "confidence_score": conf,
+        }
+        if screenshot_path_str:
+            entry["screenshot"] = screenshot_path_str
+        analyzer._record_blackboard_entry(entry)
+
+
+async def exec_analyze_audio_universal(
+    analyzer,
+    audio_path: Path,
+    start_time: float,
+    end_time: float | None,
+    actual_start: float,
+    prompt_with_context: str,
+    specific_query: str,
+    force_fallback: bool = False,
+) -> str:
+    """Executes audio analysis via Universal LangChain ChatModel."""
+    messages = [
+        SystemMessage(content=analyzer.audio_system_prompt),
+        HumanMessage(
+            content=[
+                {"type": "text", "text": prompt_with_context},
+                analyzer._audio_content_block(audio_path),
+            ]
+        ),
+    ]
+    response = await analyzer._invoke_universal_model(
+        messages,
+        [_va.UNIVERSAL_SUBMIT_ANSWER_TOOL],
+        timeout=analyzer.model_call_timeout_seconds,
+        label="Universal audio sub-agent",
+        force_fallback=force_fallback,
+    )
+    final_summary = "No audio summary provided."
+    final_analysis = "No audio analysis provided."
+
+    if getattr(response, "tool_calls", None):
+        for tc in response.tool_calls:
+            if tc.get("name") == "submit_answer":
+                args = tc.get("args", {})
+                final_summary = (
+                    str(args.get("summary", "")).strip().replace("\n", " ") or final_summary
+                )
+                final_analysis = (
+                    str(args.get("analysis", "")).strip().replace("\n", " ") or final_analysis
+                )
+                break
+    elif getattr(response, "content", None):
+        final_summary = str(response.content)[:200].replace("\n", " ")
+        final_analysis = str(response.content)
+
+    end_val = end_time if end_time is not None else "unknown"
+    analyzer._record_blackboard_entry(
+        {
+            "start": start_time,
+            "end": end_val,
+            "target": specific_query,
+            "summary": final_summary,
+            "confidence_score": 0.8,
+        },
+        modality="audio",
+    )
+
+    end_str = f"{end_time:.1f}s" if end_time is not None else "unknown"
+    return (
+        f"[from {start_time:.1f}s to {end_str}] Audio Summary: {final_summary}"
+        f" Analysis: {final_analysis}"
+    )
+
+
+async def run_universal(
+    analyzer,
+    time_description: str,
+    purpose: str,
+    system_prompt: str,
+    *,
+    force_fallback: bool = False,
+) -> tuple[str, str]:
+    """Runs VideoAnalyzer main reasoning loop using Universal LangChain BaseChatModel."""
+    messages: list[BaseMessage] = [
+        SystemMessage(content=system_prompt),
+        HumanMessage(
+            content=(
+                f"Time description: {time_description}\n"
+                f"Purpose: {purpose}\n"
+                f"Persistent video blackboard: {analyzer.blackboard_ledger}\n"
+                "Reuse sufficient prior evidence and do not request an identical "
+                "interval/query analysis again.\n"
+                "Please coordinate video metadata extraction and sub-agent analysis."
+            )
+        ),
+    ]
+
+    max_iterations = 6
+    iterations = 0
+    agent_outcome = ""
+    status = "success"
+
+    try:
+        while iterations < max_iterations:
+            iterations += 1
+            response = await _va._invoke_with_retry(
+                lambda: analyzer._invoke_universal_model(
+                    messages,
+                    _va.UNIVERSAL_MAIN_TOOLS,
+                    timeout=300,
+                    label="Universal video coordinator",
+                    force_fallback=force_fallback,
+                ),
+                "Universal video coordinator",
+            )
+            messages.append(response)
+
+            if not getattr(response, "tool_calls", None):
+                agent_outcome = str(response.content)
+                break
+
+            for tc in response.tool_calls:
+                t_name = tc["name"]
+                t_args = tc.get("args", {})
+                t_id = tc.get("id", f"call_{iterations}")
+                logger.info(f"Executing universal tool '{t_name}' with args {t_args}")
+
+                res = await _execute_universal_tool(analyzer, t_name, t_args)
+
+                messages.append(ToolMessage(tool_call_id=t_id, name=t_name, content=str(res)))
+
+        if not agent_outcome:
+            agent_outcome = "Video analyzer completed all turns without error."
+
+    except Exception as e:
+        logger.error(f"Universal video analyzer failed: {e}")
+        agent_outcome = f"Video analysis failed: {e}"
+        status = "error"
+
+    analyzer._record_reliability_metrics()
+    return agent_outcome, status
+
+
+async def _execute_universal_tool(analyzer, t_name: str, t_args: dict) -> str:
+    """Dispatches one universal coordinator tool call, converting errors to text."""
+    try:
+        if t_name == "extract_segment_metadata":
+            res = await analyzer.exec_extract_segment_metadata(
+                t_args.get("start_time"), t_args.get("end_time")
+            )
+        elif t_name == "spawn_sub_agent":
+            res = await analyzer.exec_spawn_sub_agent(
+                t_args.get("start_time"),
+                t_args.get("end_time"),
+                t_args.get("specific_query"),
+            )
+        elif t_name == "analyze_audio_only":
+            res = await analyzer.exec_analyze_audio_only(
+                t_args.get("start_time"),
+                t_args.get("end_time"),
+                t_args.get("specific_query"),
+            )
+        else:
+            res = f"Error: Tool '{t_name}' is not recognized."
+    except Exception as tool_err:
+        logger.error(f"Error executing tool '{t_name}': {tool_err}")
+        res = f"Error executing tool '{t_name}': {tool_err}"
+    return res