ccs-llmconnector 1.0.4__py3-none-any.whl → 1.0.6__py3-none-any.whl

llmconnector/gemini_client.py

@@ -1,224 +1,299 @@
-"""Thin wrapper around the Google Gemini API via the google-genai SDK."""
-
-from __future__ import annotations
-
-import base64
-import mimetypes
-from pathlib import Path
-import logging
-from typing import Optional, Sequence, Union
-from urllib.error import URLError
-from urllib.request import urlopen
-
-from google import genai
-from google.genai import types
-
-ImageInput = Union[str, Path]
-logger = logging.getLogger(__name__)
-
-
-class GeminiClient:
-    """Convenience wrapper around the Google Gemini SDK."""
-
-    def generate_response(
-        self,
-        *,
-        api_key: str,
-        prompt: str,
-        model: str,
-        max_tokens: int = 32000,
-        reasoning_effort: Optional[str] = None,
-        images: Optional[Sequence[ImageInput]] = None,
-    ) -> str:
-        """Generate a response from the specified Gemini model.
-
-        Args:
-            api_key: API key used to authenticate with the Gemini API.
-            prompt: Natural-language instruction or query for the model.
-            model: Identifier of the Gemini model to target (for example, ``"gemini-2.5-flash"``).
-            max_tokens: Cap for tokens across the entire exchange, defaults to 32000.
-            reasoning_effort: Included for API parity; currently unused by the Gemini SDK.
-            images: Optional collection of image references (local paths, URLs, or data URLs).
-
-        Returns:
-            The text output produced by the model.
-
-        Raises:
-            ValueError: If required arguments are missing or the request payload is empty.
-            URLError: If an image URL cannot be retrieved.
-            google.genai.errors.APIError: If the underlying Gemini request fails.
-        """
-        if not api_key:
-            raise ValueError("api_key must be provided.")
-        if not prompt and not images:
-            raise ValueError("At least one of prompt or images must be provided.")
-        if not model:
-            raise ValueError("model must be provided.")
-
-        parts: list[types.Part] = []
-        if prompt:
-            parts.append(types.Part.from_text(text=prompt))
-
-        if images:
-            for image in images:
-                parts.append(self._to_image_part(image))
-
-        if not parts:
-            raise ValueError("No content provided for response generation.")
-
-        content = types.Content(role="user", parts=parts)
-
-        config = types.GenerateContentConfig(max_output_tokens=max_tokens)
-        # reasoning_effort is accepted for compatibility but not currently applied because the
-        # Gemini SDK does not expose an equivalent configuration parameter.
-
-        client = genai.Client(api_key=api_key)
-        try:
-            try:
-                response = client.models.generate_content(
-                    model=model,
-                    contents=[content],
-                    config=config,
-                )
-            except Exception as exc:
-                logger.exception("Gemini generate_content failed: %s", exc)
-                raise
-        finally:
-            closer = getattr(client, "close", None)
-            if callable(closer):
-                try:
-                    closer()
-                except Exception:
-                    pass
-
-        if response.text:
-            result_text = response.text
-            logger.info(
-                "Gemini generate_content succeeded: model=%s images=%d text_len=%d",
-                model,
-                len(images or []),
-                len(result_text or ""),
-            )
-            return result_text
-
-        candidate_texts: list[str] = []
-        for candidate in getattr(response, "candidates", []) or []:
-            content_obj = getattr(candidate, "content", None)
-            if not content_obj:
-                continue
-            for part in getattr(content_obj, "parts", []) or []:
-                text = getattr(part, "text", None)
-                if text:
-                    candidate_texts.append(text)
-
-        if candidate_texts:
-            result_text = "\n".join(candidate_texts)
-            logger.info(
-                "Gemini generate_content succeeded (candidates): model=%s images=%d text_len=%d",
-                model,
-                len(images or []),
-                len(result_text or ""),
-            )
-            return result_text
-
-        # Treat successful calls without textual content as a successful, empty response
-        # rather than raising. This aligns with callers that handle empty outputs gracefully.
-        logger.info(
-            "Gemini generate_content succeeded with no text: model=%s images=%d",
-            model,
-            len(images or []),
-        )
-        return ""
-
-    def list_models(self, *, api_key: str) -> list[dict[str, Optional[str]]]:
-        """Return the models available to the authenticated Gemini account."""
-        if not api_key:
-            raise ValueError("api_key must be provided.")
-
-        models: list[dict[str, Optional[str]]] = []
-        client = genai.Client(api_key=api_key)
-        try:
-            try:
-                iterator = client.models.list()
-            except Exception as exc:
-                logger.exception("Gemini list models failed: %s", exc)
-                raise
-            for model in iterator:
-                model_id = getattr(model, "name", None)
-                if model_id is None and isinstance(model, dict):
-                    model_id = model.get("name")
-                if not model_id:
-                    continue
-
-                # Normalize IDs like "models/<id>" -> "<id>"
-                if isinstance(model_id, str) and model_id.startswith("models/"):
-                    model_id = model_id.split("/", 1)[1]
-
-                display_name = getattr(model, "display_name", None)
-                if display_name is None and isinstance(model, dict):
-                    display_name = model.get("display_name")
-
-                models.append({"id": model_id, "display_name": display_name})
-        finally:
-            closer = getattr(client, "close", None)
-            if callable(closer):
-                try:
-                    closer()
-                except Exception:
-                    pass
-
-        logger.info("Gemini list_models succeeded: count=%d", len(models))
-        return models
-
-    @staticmethod
-    def _to_image_part(image: ImageInput) -> types.Part:
-        """Convert an image reference into a Gemini SDK part."""
-        if isinstance(image, Path):
-            return _part_from_path(image)
-
-        if image.startswith("data:"):
-            return _part_from_data_url(image)
-
-        if image.startswith(("http://", "https://")):
-            return _part_from_url(image)
-
-        return _part_from_path(Path(image))
-
-
-def _part_from_path(path: Path) -> types.Part:
-    """Create an image part from a local filesystem path."""
-    expanded = path.expanduser()
-    data = expanded.read_bytes()
-    mime_type = mimetypes.guess_type(expanded.name)[0] or "application/octet-stream"
-    return types.Part.from_bytes(data=data, mime_type=mime_type)
-
-
-def _part_from_url(url: str) -> types.Part:
-    """Create an image part by downloading content from a URL."""
-    with urlopen(url) as response:
-        data = response.read()
-        mime_type = response.info().get_content_type()
-
-    if not mime_type or mime_type == "application/octet-stream":
-        mime_type = mimetypes.guess_type(url)[0] or "application/octet-stream"
-
-    return types.Part.from_bytes(data=data, mime_type=mime_type)
-
-
-def _part_from_data_url(data_url: str) -> types.Part:
-    """Create an image part from a data URL."""
-    header, encoded = data_url.split(",", 1)
-    metadata = header[len("data:") :]
-    mime_type = "application/octet-stream"
-
-    if ";" in metadata:
-        mime_type, _, metadata = metadata.partition(";")
-    elif metadata:
-        mime_type = metadata
-
-    if "base64" in metadata:
-        data = base64.b64decode(encoded)
-    else:
-        data = encoded.encode("utf-8")
-
-    return types.Part.from_bytes(data=data, mime_type=mime_type or "application/octet-stream")
+"""Thin wrapper around the Google Gemini API via the google-genai SDK."""
+
+from __future__ import annotations
+
+import base64
+import mimetypes
+from pathlib import Path
+import logging
+from typing import Optional, Sequence, Union
+from urllib.error import URLError
+from urllib.request import urlopen
+
+from google import genai
+from google.genai import types
+
+ImageInput = Union[str, Path]
+logger = logging.getLogger(__name__)
+
+
+class GeminiClient:
+    """Convenience wrapper around the Google Gemini SDK."""
+
+    def generate_response(
+        self,
+        *,
+        api_key: str,
+        prompt: str,
+        model: str,
+        max_tokens: int = 32000,
+        reasoning_effort: Optional[str] = None,
+        images: Optional[Sequence[ImageInput]] = None,
+    ) -> str:
+        """Generate a response from the specified Gemini model.
+
+        Args:
+            api_key: API key used to authenticate with the Gemini API.
+            prompt: Natural-language instruction or query for the model.
+            model: Identifier of the Gemini model to target (for example, ``"gemini-2.5-flash"``).
+            max_tokens: Cap for tokens across the entire exchange, defaults to 32000.
+            reasoning_effort: Included for API parity; currently unused by the Gemini SDK.
+            images: Optional collection of image references (local paths, URLs, or data URLs).
+
+        Returns:
+            The text output produced by the model.
+
+        Raises:
+            ValueError: If required arguments are missing or the request payload is empty.
+            URLError: If an image URL cannot be retrieved.
+            google.genai.errors.APIError: If the underlying Gemini request fails.
+        """
+        if not api_key:
+            raise ValueError("api_key must be provided.")
+        if not prompt and not images:
+            raise ValueError("At least one of prompt or images must be provided.")
+        if not model:
+            raise ValueError("model must be provided.")
+
+        parts: list[types.Part] = []
+        if prompt:
+            parts.append(types.Part.from_text(text=prompt))
+
+        if images:
+            for image in images:
+                parts.append(self._to_image_part(image))
+
+        if not parts:
+            raise ValueError("No content provided for response generation.")
+
+        content = types.Content(role="user", parts=parts)
+
+        config = types.GenerateContentConfig(max_output_tokens=max_tokens)
+        # reasoning_effort is accepted for compatibility but not currently applied because the
+        # Gemini SDK does not expose an equivalent configuration parameter.
+
+        client = genai.Client(api_key=api_key)
+        try:
+            try:
+                response = client.models.generate_content(
+                    model=model,
+                    contents=[content],
+                    config=config,
+                )
+            except Exception as exc:
+                logger.exception("Gemini generate_content failed: %s", exc)
+                raise
+        finally:
+            closer = getattr(client, "close", None)
+            if callable(closer):
+                try:
+                    closer()
+                except Exception:
+                    pass
+
+        if response.text:
+            result_text = response.text
+            logger.info(
+                "Gemini generate_content succeeded: model=%s images=%d text_len=%d",
+                model,
+                len(images or []),
+                len(result_text or ""),
+            )
+            return result_text
+
+        candidate_texts: list[str] = []
+        for candidate in getattr(response, "candidates", []) or []:
+            content_obj = getattr(candidate, "content", None)
+            if not content_obj:
+                continue
+            for part in getattr(content_obj, "parts", []) or []:
+                text = getattr(part, "text", None)
+                if text:
+                    candidate_texts.append(text)
+
+        if candidate_texts:
+            result_text = "\n".join(candidate_texts)
+            logger.info(
+                "Gemini generate_content succeeded (candidates): model=%s images=%d text_len=%d",
+                model,
+                len(images or []),
+                len(result_text or ""),
+            )
+            return result_text
+
+        # Treat successful calls without textual content as a successful, empty response
+        # rather than raising. This aligns with callers that handle empty outputs gracefully.
+        logger.info(
+            "Gemini generate_content succeeded with no text: model=%s images=%d",
+            model,
+            len(images or []),
+        )
+        return ""
+
+    def generate_image(
+        self,
+        *,
+        api_key: str,
+        prompt: str,
+        model: str,
+        image_size: Optional[str] = None,
+        aspect_ratio: Optional[str] = None,
+        image: Optional[ImageInput] = None,
+    ) -> bytes:
+        """Generate an image using Gemini 3 Pro Image.
+
+        Args:
+            api_key: API key used to authenticate with the Gemini API.
+            prompt: Text prompt for image generation.
+            model: Identifier of the Gemini model to target (e.g., "gemini-3-pro-image-preview").
+            image_size: Size of the generated image (e.g., "2K", "4K"). Defaults to "2K".
+            aspect_ratio: Aspect ratio of the generated image (e.g., "16:9", "4:3").
+            image: Optional input image for editing tasks.
+
+        Returns:
+            The generated image data as bytes.
+
+        Raises:
+            ValueError: If required arguments are missing or no image is returned.
+            google.genai.errors.APIError: If the underlying Gemini request fails.
+        """
+        if not api_key:
+            raise ValueError("api_key must be provided.")
+        if not prompt:
+            raise ValueError("prompt must be provided.")
+        if not model:
+            raise ValueError("model must be provided.")
+
+        client = genai.Client(api_key=api_key)
+        
+        config = types.GenerateContentConfig(
+            tools=[{"google_search": {}}],
+            image_config=types.ImageConfig(
+                image_size=image_size or "2K",
+                aspect_ratio=aspect_ratio,
+            )
+        )
+
+        contents = [prompt]
+        if image:
+            contents.append(self._to_image_part(image))
+
+        try:
+            try:
+                response = client.models.generate_content(
+                    model=model,
+                    contents=contents,
+                    config=config,
+                )
+            except Exception as exc:
+                logger.exception("Gemini generate_image failed: %s", exc)
+                raise
+        finally:
+            closer = getattr(client, "close", None)
+            if callable(closer):
+                try:
+                    closer()
+                except Exception:
+                    pass
+
+        if not response.parts:
+             raise ValueError("No content returned from Gemini.")
+
+        for part in response.parts:
+            if part.inline_data:
+                return part.inline_data.data
+        
+        raise ValueError("No image data found in response.")
+
+    def list_models(self, *, api_key: str) -> list[dict[str, Optional[str]]]:
+        """Return the models available to the authenticated Gemini account."""
+        if not api_key:
+            raise ValueError("api_key must be provided.")
+
+        models: list[dict[str, Optional[str]]] = []
+        client = genai.Client(api_key=api_key)
+        try:
+            try:
+                iterator = client.models.list()
+            except Exception as exc:
+                logger.exception("Gemini list models failed: %s", exc)
+                raise
+            for model in iterator:
+                model_id = getattr(model, "name", None)
+                if model_id is None and isinstance(model, dict):
+                    model_id = model.get("name")
+                if not model_id:
+                    continue
+
+                # Normalize IDs like "models/<id>" -> "<id>"
+                if isinstance(model_id, str) and model_id.startswith("models/"):
+                    model_id = model_id.split("/", 1)[1]
+
+                display_name = getattr(model, "display_name", None)
+                if display_name is None and isinstance(model, dict):
+                    display_name = model.get("display_name")
+
+                models.append({"id": model_id, "display_name": display_name})
+        finally:
+            closer = getattr(client, "close", None)
+            if callable(closer):
+                try:
+                    closer()
+                except Exception:
+                    pass
+
+        logger.info("Gemini list_models succeeded: count=%d", len(models))
+        return models
+
+    @staticmethod
+    def _to_image_part(image: ImageInput) -> types.Part:
+        """Convert an image reference into a Gemini SDK part."""
+        if isinstance(image, Path):
+            return _part_from_path(image)
+
+        if image.startswith("data:"):
+            return _part_from_data_url(image)
+
+        if image.startswith(("http://", "https://")):
+            return _part_from_url(image)
+
+        return _part_from_path(Path(image))
+
+
+def _part_from_path(path: Path) -> types.Part:
+    """Create an image part from a local filesystem path."""
+    expanded = path.expanduser()
+    data = expanded.read_bytes()
+    mime_type = mimetypes.guess_type(expanded.name)[0] or "application/octet-stream"
+    return types.Part.from_bytes(data=data, mime_type=mime_type)
+
+
+def _part_from_url(url: str) -> types.Part:
+    """Create an image part by downloading content from a URL."""
+    with urlopen(url) as response:
+        data = response.read()
+        mime_type = response.info().get_content_type()
+
+    if not mime_type or mime_type == "application/octet-stream":
+        mime_type = mimetypes.guess_type(url)[0] or "application/octet-stream"
+
+    return types.Part.from_bytes(data=data, mime_type=mime_type)
+
+
+def _part_from_data_url(data_url: str) -> types.Part:
+    """Create an image part from a data URL."""
+    header, encoded = data_url.split(",", 1)
+    metadata = header[len("data:") :]
+    mime_type = "application/octet-stream"
+
+    if ";" in metadata:
+        mime_type, _, metadata = metadata.partition(";")
+    elif metadata:
+        mime_type = metadata
+
+    if "base64" in metadata:
+        data = base64.b64decode(encoded)
+    else:
+        data = encoded.encode("utf-8")
+
+    return types.Part.from_bytes(data=data, mime_type=mime_type or "application/octet-stream")