refactor(video): split video analyzer into engine modules
This commit is contained in:
somew1nd committed
1 parent
9cad5b690f
commit
7cc20bbddc
8 files changed
+3200
-2355
No files matched your search
@@ -0,0 +1,534 @@
|
||||
# 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 video_analyzer as _va
|
||||
from artemis.constants import SAFETY_SETTINGS_BLOCK_NONE
|
||||
from artemis.data_engine.trace import CURRENT_TRACE_ID, TraceSpan
|
||||
from artemis.utils.logger import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
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(
|
||||
"Invoking Gemini for sub-agent task with model"
|
||||
f" {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:
|
||||
logger.info(f"Cleaning up cloud file {file.name}...")
|
||||
try:
|
||||
await asyncio.wait_for(
|
||||
analyzer.client.aio.files.delete(name=file.name),
|
||||
timeout=30,
|
||||
)
|
||||
analyzer.cloud_files_to_cleanup.discard(file.name)
|
||||
except Exception as ce:
|
||||
logger.error(f"Failed to delete cloud file {file.name}: {ce}")
|
||||
|
||||
|
||||
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]:
|
||||
"""Iterates the sub-agent conversation until submit_answer is accepted."""
|
||||
sub_max_iterations = 2
|
||||
sub_iterations = 0
|
||||
final_analysis = "No analysis provided."
|
||||
final_summary = "No summary provided."
|
||||
final_full_text = ""
|
||||
|
||||
while sub_iterations < sub_max_iterations:
|
||||
sub_iterations += 1
|
||||
if sub_iterations == sub_max_iterations - 1:
|
||||
sub_agent_contents.append(
|
||||
types.Content(
|
||||
role="user",
|
||||
parts=[
|
||||
types.Part.from_text(
|
||||
text=(
|
||||
"[WARNING] You must call the"
|
||||
" submit_answer tool now, time"
|
||||
" is running out."
|
||||
)
|
||||
)
|
||||
],
|
||||
)
|
||||
)
|
||||
|
||||
full_text, function_calls, original_parts = await _stream_sub_agent_turn(
|
||||
analyzer, sub_agent_contents, trace_id
|
||||
)
|
||||
final_full_text += full_text + "\n"
|
||||
|
||||
if function_calls:
|
||||
answered, final_summary, final_analysis = await _handle_sub_agent_calls(
|
||||
analyzer,
|
||||
sub_agent_contents,
|
||||
function_calls,
|
||||
original_parts,
|
||||
full_text,
|
||||
media,
|
||||
start_time,
|
||||
end_time,
|
||||
specific_query,
|
||||
final_summary,
|
||||
final_analysis,
|
||||
)
|
||||
if answered is True:
|
||||
break
|
||||
else:
|
||||
sub_agent_contents.append(
|
||||
types.Content(
|
||||
role="model",
|
||||
parts=[types.Part.from_text(text=full_text)],
|
||||
)
|
||||
)
|
||||
|
||||
return final_summary, final_analysis
|
||||
|
||||
|
||||
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:
|
||||
if hasattr(chunk, "usage_metadata") and chunk.usage_metadata:
|
||||
span.payload["usage_metadata"] = {
|
||||
"prompt_token_count": getattr(
|
||||
chunk.usage_metadata,
|
||||
"prompt_token_count",
|
||||
0,
|
||||
),
|
||||
"candidates_token_count": getattr(
|
||||
chunk.usage_metadata,
|
||||
"candidates_token_count",
|
||||
0,
|
||||
),
|
||||
"total_token_count": getattr(
|
||||
chunk.usage_metadata,
|
||||
"total_token_count",
|
||||
0,
|
||||
),
|
||||
}
|
||||
|
||||
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,
|
||||
final_summary: str,
|
||||
final_analysis: str,
|
||||
):
|
||||
"""Processes function calls from one turn; returns (answered, summary, analysis)."""
|
||||
answered = False
|
||||
for fc in function_calls:
|
||||
if (
|
||||
fc.name.split(":")[-1] if ":" in fc.name else fc.name
|
||||
) == "submit_answer":
|
||||
args = fc.args
|
||||
timeline_events = args.get("timeline_events", [])
|
||||
final_summary = (
|
||||
args.get("summary", "No summary provided.")
|
||||
.strip()
|
||||
.replace("\n", " ")
|
||||
)
|
||||
final_analysis = (
|
||||
args.get("analysis", "No analysis provided.")
|
||||
.strip()
|
||||
.replace("\n", " ")
|
||||
)
|
||||
|
||||
error_msg = _validate_timeline_events(timeline_events, start_time, end_time)
|
||||
if error_msg is not None:
|
||||
_append_rejection(
|
||||
sub_agent_contents,
|
||||
original_parts,
|
||||
function_calls,
|
||||
full_text,
|
||||
fc.name,
|
||||
error_msg,
|
||||
)
|
||||
answered = "error"
|
||||
break
|
||||
|
||||
await _persist_timeline_events(
|
||||
analyzer, timeline_events, media, start_time, end_time, specific_query
|
||||
)
|
||||
|
||||
answered = True
|
||||
break
|
||||
|
||||
if answered is True or answered == "error":
|
||||
return answered, final_summary, final_analysis
|
||||
|
||||
_append_rejection(
|
||||
sub_agent_contents,
|
||||
original_parts,
|
||||
function_calls,
|
||||
full_text,
|
||||
function_calls[0].name,
|
||||
(
|
||||
"Tool not recognized."
|
||||
" Please use"
|
||||
" submit_answer."
|
||||
),
|
||||
)
|
||||
return answered, final_summary, final_analysis
|
||||
|
||||
|
||||
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)
|
||||
Reference in new issue
Block a user