Merge pull request #70 from ComfyAssets/release/1.1.1

Release 1.1.1

README.md

@@ -1067,7 +1067,7 @@ MIT License - see [LICENSE](LICENSE) file for details.
 - **Test Coverage**: 100% (470+ comprehensive tests)
 - **Python Version**: 3.8+
 - **ComfyUI Compatibility**: Latest
-- **Dependencies**: Minimal (PyTorch, NumPy, Pillow, google-generativeai for Gemini)
+- **Dependencies**: Minimal (PyTorch, NumPy, Pillow, google-genai for Gemini)
 
 ## 🙏 Attribution
 

examples/documentation/gemini_prompt.md

@@ -21,7 +21,7 @@ Get your free Gemini API key from [Google AI Studio](https://makersuite.google.c
 ### 2. Install Dependencies
 
 ```bash
-pip install google-generativeai
+pip install google-genai
 ```
 
 ### 3. Configure API Key
@@ -204,6 +204,6 @@ Errors are displayed in the prompt output for easy debugging.
 - Verify API key is valid
 - Image might be too large (resize if needed)
 
-**Import error for google-generativeai**:
-- Run `pip install google-generativeai` in your ComfyUI environment
+**Import error for google-genai**:
+- Run `pip install google-genai` in your ComfyUI environment
 - Restart ComfyUI after installation

kikotools/tools/gemini_prompt/logic.py

@@ -115,30 +115,24 @@ def analyze_image_with_gemini(
     else:
         system_prompt = PROMPT_TEMPLATES.get(prompt_type, PROMPT_TEMPLATES["flux"])
 
-    # Here we would normally make the API call to Gemini
-    # For now, we'll import the google-generativeai library
     try:
-        import google.generativeai as genai
+        from google import genai
     except ImportError:
         return (
             "",
-            "google-generativeai library not installed. Please run: pip install google-generativeai",
+            "google-genai library not installed. Please run: pip install google-genai",
         )
 
     try:
-        # Configure Gemini
-        genai.configure(api_key=api_key)
+        client = genai.Client(api_key=api_key)
 
-        # Create model
-        model = genai.GenerativeModel(model_name)
-
-        # Generate content
-        response = model.generate_content(
-            [
+        response = client.models.generate_content(
+            model=model_name,
+            contents=[
                 system_prompt,
                 pil_image,
                 "Analyze this image and generate an appropriate prompt according to the instructions.",
-            ]
+            ],
         )
 
         # Extract text from response

kikotools/tools/gemini_prompt/models.py

@@ -59,10 +59,10 @@ def _fetch_models_from_api(
         Tuple of (model_names_list, model_descriptions_dict)
     """
     try:
-        import google.generativeai as genai
+        from google import genai
     except ImportError:
         if not silent:
-            logger.error("google-generativeai not installed")
+            logger.error("google-genai not installed")
         return [], {}
 
     # Get API key
@@ -77,15 +77,15 @@ def _fetch_models_from_api(
         return [], {}
 
     try:
-        genai.configure(api_key=api_key)
+        client = genai.Client(api_key=api_key)
 
         models = []
         descriptions = {}
 
         # Fetch all models
-        for model in genai.list_models():
+        for model in client.models.list():
             # Only include models that support generateContent
-            if "generateContent" in model.supported_generation_methods:
+            if "generateContent" in (model.supported_actions or []):
                 # Remove "models/" prefix from name
                 model_name = model.name.replace("models/", "")
                 models.append(model_name)

kikotools/tools/gemini_prompt/node.py

@@ -63,7 +63,7 @@ Supports multiple prompt formats:
 - Video: Motion and temporal descriptions for video generation
 
 Requires Gemini API key (set GEMINI_API_KEY env var or provide in node).
-Install: pip install google-generativeai
+Install: pip install google-genai
 """
 
     def generate_prompt(

pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "setuptools.build_meta"
 [project]
 name = "kikotools"
 description = "Simple tools for ComfyUI"
-version = "1.1.0"
+version = "1.1.1"
 license = {text = "MIT"}
 dependencies = []
 

requirements.txt

@@ -1,4 +1,4 @@
 # Runtime dependencies for ComfyUI-KikoTools
 
 # Gemini API integration (optional - only needed for Gemini Prompt node)
-google-generativeai
+google-genai

tests/unit/tools/test_gemini_prompt.py

@@ -187,10 +187,10 @@ class TestGeminiLogic:
         assert validate_prompt_type("") is False
         assert validate_prompt_type(None) is False
 
-    @pytest.mark.skip(reason="Requires google-generativeai library")
+    @pytest.mark.skip(reason="Requires google-genai library")
     def test_analyze_image_with_gemini_success(self):
         """Test successful image analysis with Gemini."""
-        pass  # Skipped as it requires google-generativeai
+        pass  # Skipped as it requires google-genai
 
     def test_analyze_image_no_api_key(self):
         """Test analysis without API key."""
@@ -204,10 +204,10 @@ class TestGeminiLogic:
         assert result == ""
         assert "API key not found" in error
 
-    @pytest.mark.skip(reason="Requires google-generativeai library")
+    @pytest.mark.skip(reason="Requires google-genai library")
     def test_analyze_image_with_custom_prompt(self):
         """Test analysis with custom prompt."""
-        pass  # Skipped as it requires google-generativeai
+        pass  # Skipped as it requires google-genai
 
 
 class TestPromptTemplates:

tests/unit/tools/test_gemini_prompt_extra.py

@@ -24,17 +24,22 @@ def cache_file(tmp_path, monkeypatch):
 
 def fake_genai(model_specs):
     genai = MagicMock()
-    genai.list_models.return_value = [
+    genai.Client.return_value.models.list.return_value = [
         SimpleNamespace(
             name=f"models/{name}",
             display_name=display,
-            supported_generation_methods=methods,
+            supported_actions=methods,
         )
         for name, display, methods in model_specs
     ]
     return genai
 
 
+def sdk_modules(genai):
+    """sys.modules patch that makes `from google import genai` return genai."""
+    return {"google": MagicMock(genai=genai), "google.genai": genai}
+
+
 class TestSortModels:
     def test_newer_versions_come_first(self):
         result = models._sort_models(
@@ -111,27 +116,25 @@ class TestCache:
 
 class TestFetchModelsFromApi:
     def test_missing_sdk_returns_empty(self):
-        with patch.dict(sys.modules, {"google.generativeai": None, "google": None}):
+        with patch.dict(sys.modules, {"google.genai": None, "google": None}):
             assert models._fetch_models_from_api("k") == ([], {})
 
     def test_no_api_key_returns_empty(self):
         genai = fake_genai([])
-        google = MagicMock(generativeai=genai)
         with (
-            patch.dict(sys.modules, {"google": google, "google.generativeai": genai}),
+            patch.dict(sys.modules, sdk_modules(genai)),
             patch.object(logic, "get_api_key", return_value=None),
         ):
             assert models._fetch_models_from_api(None, silent=True) == ([], {})
 
     def test_api_key_taken_from_logic_when_not_given(self):
         genai = fake_genai([("gemini-2.5-pro", "Pro", ["generateContent"])])
-        google = MagicMock(generativeai=genai)
         with (
-            patch.dict(sys.modules, {"google": google, "google.generativeai": genai}),
+            patch.dict(sys.modules, sdk_modules(genai)),
             patch.object(logic, "get_api_key", return_value="envkey"),
         ):
             names, _ = models._fetch_models_from_api(None)
-        genai.configure.assert_called_once_with(api_key="envkey")
+        genai.Client.assert_called_once_with(api_key="envkey")
         assert names == ["gemini-2.5-pro"]
 
     def test_filters_to_generate_content_and_strips_prefix(self):
@@ -142,17 +145,15 @@ class TestFetchModelsFromApi:
                 ("gemini-2.5-flash", "Flash", ["generateContent", "countTokens"]),
             ]
         )
-        google = MagicMock(generativeai=genai)
-        with patch.dict(sys.modules, {"google": google, "google.generativeai": genai}):
+        with patch.dict(sys.modules, sdk_modules(genai)):
             names, desc = models._fetch_models_from_api("k")
         assert names == ["gemini-2.5-flash", "gemini-1.5-pro"]
         assert desc == {"gemini-1.5-pro": "Pro 1.5", "gemini-2.5-flash": "Flash"}
 
     def test_api_error_returns_empty(self):
         genai = MagicMock()
-        genai.list_models.side_effect = RuntimeError("boom")
-        google = MagicMock(generativeai=genai)
-        with patch.dict(sys.modules, {"google": google, "google.generativeai": genai}):
+        genai.Client.return_value.models.list.side_effect = RuntimeError("boom")
+        with patch.dict(sys.modules, sdk_modules(genai)):
             assert models._fetch_models_from_api("k") == ([], {})
 
 
@@ -205,20 +206,19 @@ class TestLogicExtra:
 
     def test_missing_sdk_reported(self):
         img = np.zeros((1, 4, 4, 3), dtype=np.float32)
-        with patch.dict(sys.modules, {"google.generativeai": None, "google": None}):
+        with patch.dict(sys.modules, {"google.genai": None, "google": None}):
             out, err = logic.analyze_image_with_gemini(img, "flux", api_key="k")
         assert out == "" and "not installed" in err
 
     def _run(self, response=None, error=None, **kwargs):
         genai = MagicMock()
-        model = genai.GenerativeModel.return_value
+        client_models = genai.Client.return_value.models
         if error:
-            model.generate_content.side_effect = error
+            client_models.generate_content.side_effect = error
         else:
-            model.generate_content.return_value = response
-        google = MagicMock(generativeai=genai)
+            client_models.generate_content.return_value = response
         img = np.zeros((1, 4, 4, 3), dtype=np.float32)
-        with patch.dict(sys.modules, {"google": google, "google.generativeai": genai}):
+        with patch.dict(sys.modules, sdk_modules(genai)):
             result = logic.analyze_image_with_gemini(img, "flux", api_key="k", **kwargs)
         return result, genai
 
@@ -227,24 +227,28 @@ class TestLogicExtra:
             SimpleNamespace(text="  hello \n"), model_name="gemini-x"
         )
         assert (out, err) == ("hello", None)
-        genai.configure.assert_called_once_with(api_key="k")
-        genai.GenerativeModel.assert_called_once_with("gemini-x")
+        genai.Client.assert_called_once_with(api_key="k")
+        call = genai.Client.return_value.models.generate_content.call_args
+        assert call.kwargs["model"] == "gemini-x"
 
     def test_custom_prompt_is_sent_as_system_prompt(self):
         (_, _), genai = self._run(SimpleNamespace(text="x"), custom_prompt="BE TERSE")
-        sent = genai.GenerativeModel.return_value.generate_content.call_args[0][0]
+        sent = genai.Client.return_value.models.generate_content.call_args.kwargs[
+            "contents"
+        ]
         assert sent[0] == "BE TERSE"
 
     def test_unknown_prompt_type_falls_back_to_flux_template(self):
         genai = MagicMock()
-        genai.GenerativeModel.return_value.generate_content.return_value = (
+        genai.Client.return_value.models.generate_content.return_value = (
             SimpleNamespace(text="x")
         )
-        google = MagicMock(generativeai=genai)
         img = np.zeros((1, 4, 4, 3), dtype=np.float32)
-        with patch.dict(sys.modules, {"google": google, "google.generativeai": genai}):
+        with patch.dict(sys.modules, sdk_modules(genai)):
             logic.analyze_image_with_gemini(img, "nonsense", api_key="k")
-        sent = genai.GenerativeModel.return_value.generate_content.call_args[0][0]
+        sent = genai.Client.return_value.models.generate_content.call_args.kwargs[
+            "contents"
+        ]
         assert sent[0] == logic.PROMPT_TEMPLATES["flux"]
 
     def test_empty_response(self):