535 lines
17 KiB
Python
535 lines
17 KiB
Python
# 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)
|