Files
Artemis-Eapil/artemis/agents/video_analyzer/chunk_conversation.py
T
somew1nd 67248c887e feat: improve task execution, verification and packaging
Add execution incidents, action bursts and shared Flash history.

Expose verification and Explorer settings with streamed check results.

Support graceful cancellation, recording recovery and bundled UI resources.

Add CI quality checks and clean up comments and documentation.
2026-09-02 21:45:51 -07:00

501 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(
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:
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)