Repository navigation
feat: Migrate to google-genai SDK, fix vulnerability, and add examples - #4
Conversation
- Replace deprecated google-generativeai with google-genai SDK - Resolves protobuf vulnerability CVE-2026-0994 (was pinned <6.0 by old SDK) - Update vision.py to use new genai.Client() async API - Update test_vision.py for new SDK mocking patterns - Add type ignore comments for mypy false positives on mock assertions - Add examples directory with 3 demo scripts: - basic_animation_check.py: Record and analyze animations - screenshot_analysis.py: Static page visual analysis - accessibility_check.py: Visual accessibility checks
|
Warning Rate limit exceeded
⌛ How to resolve this issue?After the wait time has elapsed, a review can be triggered using the We recommend that you space out your commits to avoid hitting the rate limit. 🚦 How do rate limits work?CodeRabbit enforces hourly rate limits for each developer per organization. Our paid plans have higher rate limits than the trial, open-source and free plans. In all cases, we re-allow further reviews after a brief timeout. Please see our FAQ for further information. WalkthroughApprentice: Adds example scripts and a README demonstrating AnimaWatch workflows, migrates GeminiProvider internals from the global google.generativeai API to the client-based google.genai async API, updates the dependency, and adjusts tests to mock and assert the new client-based interactions. Changes
Sequence Diagram(s)sequenceDiagram
participant Recorder as BrowserRecorder
participant FS as Local File System
participant GenAI as genai.Client (aio)
participant Model as Gemini Model (generate_content)
Recorder->>FS: write video/image bytes
FS-->>Recorder: file path / bytes
Recorder->>GenAI: aio.files.upload(file bytes)
GenAI-->>Recorder: file_id / uri (processing)
loop poll until ACTIVE or timeout (MAX_PROCESSING_SECONDS)
Recorder->>GenAI: aio.files.get(file_id)
GenAI-->>Recorder: file state
end
Recorder->>GenAI: aio.models.generate_content(model, contents: [Part(file), Part(prompt)])
GenAI-->>Model: invoke model
Model-->>GenAI: analysis response (text)
GenAI->>GenAI: aio.files.delete(file_id)
GenAI-->>Recorder: return analysis text
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 10
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
tests/test_vision.py (2)
21-33:⚠️ Potential issue | 🟡 MinorAvoid asserting internal state, Apprentice.
provider.model_nameis an internal detail; favour behaviour-based assertions to reduce coupling.💡 Proposed change
- assert provider.model_name == "gemini-2.0-flash"
132-158:⚠️ Potential issue | 🟡 MinorAdd an empty/exception response case for image analysis, Apprentice.
Only the happy path is covered; add a test for a blank response (or raised error) to satisfy edge-case coverage.
🧪 Example edge-case test
@@ async def test_analyze_image_reads_and_processes(self, tmp_path: Path) -> None: """Test that analyze_image reads image and calls generate_content.""" @@ assert result == "Image analysis" mock_client.aio.models.generate_content.assert_called_once() + + `@pytest.mark.asyncio` + async def test_analyze_image_returns_empty_on_blank_response(self, tmp_path: Path) -> None: + """Test that analyze_image returns an empty string when response has no text.""" + image_path = tmp_path / "test.png" + image_path.write_bytes(b"fake image data") + + with ( + patch("animawatch.vision.settings") as mock_settings, + patch("animawatch.vision.genai") as mock_genai, + ): + mock_settings.gemini_api_key = "test-api-key" + mock_settings.vision_model = "gemini-2.0-flash" + + mock_response = MagicMock() + mock_response.text = None + mock_client = MagicMock() + mock_client.aio.models.generate_content = AsyncMock(return_value=mock_response) + mock_genai.Client.return_value = mock_client + + provider = GeminiProvider() + result = await provider.analyze_image(image_path, "Analyze image") + + assert result == ""
🤖 Fix all issues with AI agents
In `@examples/accessibility_check.py`:
- Around line 1-7: Run the project's formatter (uv run ruff format) on the
example file to fix the formatting CI failure: reformat the shebang line and
module-level docstring so they comply with Ruff's rules, save and commit the
updated file (you can run uv run ruff format <filename> locally or via the repo
script to apply the fix).
- Around line 16-78: The temporary screenshot is only removed before the return,
so if vision.analyze_image raises the file remains; update check_accessibility
to ensure screenshot_path is deleted in the finally block (after await
browser.stop() or before stopping) by checking screenshot_path exists (and is
not None) and unlinking it; keep BrowserRecorder usage (browser.start,
browser.stop) and ensure screenshot_path is defined in the outer scope of the
try so the finally can access it and safely handle failures from analyze_image.
In `@examples/basic_animation_check.py`:
- Around line 1-7: Run the code formatter (ruff) on the module
basic_animation_check.py and commit the changes: execute "uv run ruff format"
(or "ruff format") to reformat the shebang, docstring and surrounding whitespace
so the file matches the project's style rules, then stage and push the formatted
file.
- Around line 16-65: The temporary recording may be left behind if
vision.analyze_video raises; in check_animations initialize video_path = None
before recording, then move the removal logic into the finally block so the file
is always considered for deletion (only when output_dir is None and video_path
is not None and video_path.exists()), keeping the existing await browser.stop()
call in finally; this ensures cleanup happens on both success and failure while
avoiding NameError if record_interaction fails.
In `@examples/README.md`:
- Around line 7-95: Fix the Markdown fencing and trailing whitespace: ensure
there is a blank line before and after each triple-backtick fenced block (e.g.,
the initial bash block containing "cd ~/github/animawatch", the GEMINI_API_KEY
block with export GEMINI_API_KEY, and the Ollama example block), remove the
stray leading/trailing spaces inside those fences, and delete the trailing space
at the end of the Note line that mentions basic_animation_check.py so the line
reads "**Note**: Ollama doesn't support direct video analysis, so
`basic_animation_check.py`" with no trailing whitespace.
In `@examples/screenshot_analysis.py`:
- Around line 1-7: Run Ruff's auto-formatter on the affected file
(examples/screenshot_analysis.py) to fix the CI formatting failure: execute the
formatter via your project scripts (e.g., "uv run ruff format") or directly with
ruff, review the changed docstring/shebang formatting, stage the updated file,
and commit the change so the CI formatting check passes.
- Around line 15-64: The temp screenshot is only deleted after a successful
analysis; if vision.analyze_image raises the file remains. Initialize
screenshot_path = None before the try in analyze_screenshot, then move the file
cleanup into the existing finally block so it always runs: in finally, if
screenshot_path is not None and screenshot_path.exists() then unlink it, and
still await browser.stop() (preserve BrowserRecorder.start/stop usage). This
ensures cleanup runs on failure while keeping the browser shutdown logic in
analyze_screenshot.
In `@src/animawatch/vision.py`:
- Around line 10-11: Wrap the long-running await in the polling loop with an
asyncio.timeout (e.g., asyncio.timeout(300)) so the loop aborts after a bounded
time instead of awaiting forever, and replace the overly broad
contextlib.suppress(Exception) with a targeted suppression such as
contextlib.suppress(FileNotFoundError, types.NotFound) so asyncio.CancelledError
is not swallowed; specifically update the polling loop await to use
asyncio.timeout(300) and change calls referencing contextlib.suppress(Exception)
to suppress only FileNotFoundError and types.NotFound.
- Around line 43-50: The polling loop that waits for video processing (the while
using self.client.aio.files.get after upload) lacks a timeout and can hang
indefinitely; add a maximum wait boundary (e.g., MAX_PROCESSING_SECONDS) by
recording a start time or using asyncio.wait_for around the polling coroutine,
check elapsed time each iteration, and raise a clear TimeoutError including the
file_name and elapsed time when exceeded; update the upload-and-wait logic in
the same method (the video_file = await self.client.aio.files.upload(...) block
and the following while loop) to enforce this timeout and cleanly abort polling.
In `@tests/test_server.py`:
- Around line 218-219: Replace the "type: ignore[attr-defined]" comments by
casting the mocked attributes to AsyncMock so type-checkers recognize assertion
methods: import cast and AsyncMock, then use patterns like cast(AsyncMock,
mock_app_context.browser).record_interaction.assert_called_once() and
cast(AsyncMock, mock_app_context.vision).analyze_video.assert_called_once().
Apply the same cast to the other five assertion spots (the remaining calls on
mock_app_context.browser and mock_app_context.vision in the test file) so all
seven assertions use cast(AsyncMock, ...) instead of type ignores.
- Add timeout (MAX_PROCESSING_SECONDS=300) to video processing poll loop - Move temp file cleanup to finally blocks in all example files - Replace type: ignore comments with cast(AsyncMock, ...) pattern - Initialize temp path variables before try blocks for proper cleanup
- Remove internal state assertion (provider.model_name) in test_vision.py - Add edge case test for empty/blank response in image analysis - Now 50 tests total
There was a problem hiding this comment.
Actionable comments posted: 0
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
tests/test_vision.py (1)
21-34:⚠️ Potential issue | 🟡 MinorRemove the constructor mock assertion, Apprentice. The inspect of genai.Client initialization is an internal detail; the true measure of your provider's power manifests when it summons the Dark Side through analyse_video and analyse_image methods. The provider's creation is sufficiently validated by its mere existence.
🔧 Minimal change
- mock_genai.Client.assert_called_once_with(api_key="test-api-key") - # Verify provider was created (avoid asserting internal state) - assert provider is not None + assert provider is not NoneAs per the coding guidelines for
tests/**/*.py: No implementation details tested (test behaviour, not internals).src/animawatch/vision.py (1)
36-42:⚠️ Potential issue | 🟡 MinorAnnotate these attributes with their types, Apprentice. The Dark Side demands clarity—these instance variables must bear witness to their types to satisfy the repository's type-safety doctrine.
Proposed change
- self.client = genai.Client(api_key=settings.gemini_api_key) - self.model_name = settings.vision_model + self.client: genai.Client = genai.Client(api_key=settings.gemini_api_key) + self.model_name: str = settings.vision_modelAs per the coding guidelines for
src/**/*.py: Type hints everywhere (noAnyunless justified).
|
@coderabbitai review |
✅ Actions performedReview triggered.
|
- Remove implementation detail assertion in test_init_with_api_key_creates_client - Add type annotations to GeminiProvider.client and .model_name
|
@coderabbitai approve All the issues from the previous review have been addressed:
|
|
The rebel scum of technical debt have been vanquished:
Your migration to I grant my approval. The Force is strong with this pull request. If you found this review helpful, would you consider giving us a shout-out on X? Thank you for using CodeRabbit! ✅ Actions performedComments resolved and changes approved. |
1 participant
Summary
This PR completes three tasks:
1. ✅ Fix Dependabot Vulnerability (CVE-2026-0994)
The old
google-generativeaipackage requiredprotobuf<6.0.0.dev0, which had a high severity vulnerability. By migrating to the newgoogle-genaiSDK, this constraint is removed and protobuf is now at 6.33.5 (no vulnerabilities).2. ✅ Migrate to google-genai SDK
The old
google-generativeaipackage was deprecated. This PR migrates to the newgoogle-genaiSDK with these changes:genai.configure()+genai.GenerativeModel()togenai.Client(api_key=...)client.aio.files.upload(),client.aio.models.generate_content(), etc.types.Part.from_uri(),types.Part.from_bytes(),types.Part.from_text()for content parts3. ✅ Add Examples
Added an
examples/directory with 3 demo scripts:basic_animation_check.pyscreenshot_analysis.pyaccessibility_check.pyTesting
Changes
pyproject.toml: Replacedgoogle-generativeaiwithgoogle-genai>=1.0.0src/animawatch/vision.py: Migrated to new SDK APItests/test_vision.py: Updated mocks for new SDKtests/test_server.py: Added type ignores for mock assertionsexamples/: New directory with 3 example scripts + READMEPull Request opened by Augment Code with guidance from the PR author