feat: stream uploads to the Gemini Files API with polling and cleanup

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

skills/watch/scripts/gemini.py

@@ -8,7 +8,10 @@ from __future__ import annotations
 
 import json
 import math
+import mimetypes
 import re
+import time
+from pathlib import Path
 from urllib.error import HTTPError, URLError
 from urllib.request import Request, urlopen
 
@@ -25,6 +28,7 @@ _HINTS = {
     'service': 'Gemini had a server error; retry shortly',
     'network': 'could not reach generativelanguage.googleapis.com; check network, proxy, or egress rules',
     'response': 'Gemini returned something this version cannot read',
+    'upload': 'the video could not be uploaded to or processed by the Gemini Files API',
 }
 DEFAULT_QUESTION = ('Give a thorough summary of this video: what is shown, any on-screen text or code, '
                     'and what is said, in chronological order.')
@@ -108,3 +112,49 @@ def ask(video: dict, question: str | None, *, model: str, key: str, clip=None, t
         raise _fail('response', f'no answer text (status {data.get("status")})', key)
     return {'text': '\n'.join(parts).strip(), 'model': model, 'processing': label,
             'total_tokens': tokens if isinstance(tokens, int) else None}
+
+
+def upload_file(path: Path, key: str, *, poll_seconds: float = 2.0, max_wait: float = 900.0, sleep=time.sleep) -> dict:
+    path = Path(path)
+    size = path.stat().st_size
+    if size == 0:
+        raise _fail('upload', f'{path.name} is empty', key)
+    guessed = mimetypes.guess_type(path.name)[0] or ''
+    mime = guessed if guessed.startswith('video/') else 'video/mp4'
+    _, headers, _ = _call('POST', f'{API}/upload/v1beta/files', key,
+                          data=json.dumps({'file': {'display_name': 'watch-upload'}}).encode('utf-8'),
+                          headers={'X-Goog-Upload-Protocol': 'resumable', 'X-Goog-Upload-Command': 'start',
+                                   'X-Goog-Upload-Header-Content-Length': str(size),
+                                   'X-Goog-Upload-Header-Content-Type': mime, 'Content-Type': 'application/json'})
+    session = {name.lower(): value for name, value in headers.items()}.get('x-goog-upload-url')
+    if not session:
+        raise _fail('upload', 'the Files API did not return an upload session URL', key)
+    with path.open('rb') as stream:  # Streamed in blocks; never held in memory.
+        _, _, body = _call('POST', session, key, data=stream,
+                           headers={'Content-Length': str(size), 'X-Goog-Upload-Offset': '0',
+                                    'X-Goog-Upload-Command': 'upload, finalize'}, timeout=3600.0)
+    try:
+        info = json.loads(body)['file']
+        waited = 0.0
+        while info.get('state') == 'PROCESSING':
+            if waited >= max_wait:
+                delete_file(info['name'], key)
+                raise _fail('upload', f'still processing after {int(max_wait)}s', key)
+            sleep(poll_seconds)
+            waited += poll_seconds
+            info = json.loads(_call('GET', f"{API}/v1beta/{info['name']}", key)[2])
+        if info.get('state') != 'ACTIVE':
+            delete_file(info['name'], key)
+            raise _fail('upload', f"file state {info.get('state')}", key)
+        return {'name': info['name'], 'uri': info['uri'], 'mime_type': info.get('mimeType') or mime}
+    except (ValueError, KeyError, TypeError):
+        raise _fail('upload', 'unreadable Files API response', key) from None
+
+
+def delete_file(name: str, key: str) -> str | None:
+    try:
+        _call('DELETE', f'{API}/v1beta/{name}', key)
+        return None
+    except (SystemExit, Exception) as exc:  # Cleanup must never mask the real result or error.
+        return (f'Could not delete uploaded {name} from Google ({str(exc).split(". Detail:")[0]}); '
+                'it expires on its own within 48 hours.')

tests/test_gemini.py

@@ -136,3 +136,77 @@ def test_key_echoed_by_the_server_is_redacted(calls):
     with pytest.raises(SystemExit) as caught:
         gemini.ask({'uri': 'https://youtu.be/abc'}, None, model='m', key=KEY)
     assert KEY not in str(caught.value) and '[redacted]' in str(caught.value)
+
+
+def file_state(state, name='files/abc'):
+    return FakeResponse(json.dumps({'name': name, 'state': state, 'mimeType': 'video/mp4',
+                                    'uri': f'https://generativelanguage.googleapis.com/v1beta/{name}'}).encode())
+
+
+def test_upload_streams_file_then_polls_until_active(calls, tmp_path):
+    clip = tmp_path / 'clip.mp4'
+    clip.write_bytes(b'x' * 1000)
+    calls.replies = [
+        FakeResponse(headers={'X-Goog-Upload-URL': 'https://upload.invalid/session?upload_id=1'}),
+        FakeResponse(json.dumps({'file': {'name': 'files/abc', 'state': 'PROCESSING', 'mimeType': 'video/mp4',
+                                          'uri': 'https://generativelanguage.googleapis.com/v1beta/files/abc'}}).encode()),
+        file_state('PROCESSING'), file_state('ACTIVE'),
+    ]
+    naps = []
+    result = gemini.upload_file(clip, KEY, sleep=naps.append)
+    start, body, poll = calls[0], calls[1], calls[2]
+    assert start.full_url == 'https://generativelanguage.googleapis.com/upload/v1beta/files'
+    assert start.get_header('X-goog-upload-command') == 'start'
+    assert start.get_header('X-goog-upload-header-content-length') == '1000'
+    assert start.get_header('X-goog-upload-header-content-type') == 'video/mp4'
+    assert body.full_url == 'https://upload.invalid/session?upload_id=1'
+    assert body.get_header('X-goog-upload-command') == 'upload, finalize'
+    assert body.get_header('Content-length') == '1000'
+    assert hasattr(body.data, 'read'), 'the body must be streamed, not loaded into memory'
+    assert poll.full_url.endswith('/v1beta/files/abc') and poll.get_method() == 'GET'
+    assert naps == [2.0, 2.0]
+    assert result == {'name': 'files/abc', 'mime_type': 'video/mp4',
+                      'uri': 'https://generativelanguage.googleapis.com/v1beta/files/abc'}
+
+
+def test_upload_failed_state_and_timeout(calls, tmp_path):
+    clip = tmp_path / 'clip.mkv'
+    clip.write_bytes(b'x')
+    first = [FakeResponse(headers={'x-goog-upload-url': 'https://upload.invalid/s'}),
+             FakeResponse(json.dumps({'file': {'name': 'files/abc', 'state': 'PROCESSING', 'uri': 'u'}}).encode())]
+    calls.replies = [*first, file_state('FAILED')]
+    with pytest.raises(SystemExit, match='Gemini upload:.*FAILED'):
+        gemini.upload_file(clip, KEY, sleep=lambda s: None)
+    calls.replies = [*first, *[file_state('PROCESSING') for _ in range(3)]]
+    with pytest.raises(SystemExit, match='Gemini upload:.*still processing'):
+        gemini.upload_file(clip, KEY, poll_seconds=2.0, max_wait=6.0, sleep=lambda s: None)
+
+
+def test_upload_without_session_url_or_with_empty_file(calls, tmp_path):
+    empty = tmp_path / 'empty.mp4'
+    empty.write_bytes(b'')
+    with pytest.raises(SystemExit, match='Gemini upload:.*empty'):
+        gemini.upload_file(empty, KEY)
+    clip = tmp_path / 'clip.mp4'
+    clip.write_bytes(b'x')
+    calls.replies = [FakeResponse(headers={})]
+    with pytest.raises(SystemExit, match='Gemini upload:.*session'):
+        gemini.upload_file(clip, KEY)
+
+
+def test_unknown_extension_falls_back_to_mp4(calls, tmp_path):
+    clip = tmp_path / 'clip.unknownext'
+    clip.write_bytes(b'x')
+    calls.replies = [FakeResponse(headers={'x-goog-upload-url': 'https://upload.invalid/s'}),
+                     FakeResponse(json.dumps({'file': {'name': 'files/a', 'state': 'ACTIVE', 'uri': 'u'}}).encode())]
+    assert gemini.upload_file(clip, KEY)['mime_type'] == 'video/mp4'
+    assert calls[0].get_header('X-goog-upload-header-content-type') == 'video/mp4'
+
+
+def test_delete_is_best_effort(calls):
+    calls.replies = [FakeResponse(b'{}')]
+    assert gemini.delete_file('files/abc', KEY) is None
+    assert calls[0].get_method() == 'DELETE' and calls[0].full_url.endswith('/v1beta/files/abc')
+    calls.replies = [http_error(500, b'boom')]
+    warning = gemini.delete_file('files/abc', KEY)
+    assert 'files/abc' in warning and '48 hours' in warning and KEY not in warning