prompture 1.3.3.dev1__tar.gz → 1.4.1.dev1__tar.gz

{prompture-1.3.3.dev1 → prompture-1.4.1.dev1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: prompture
-Version: 1.3.3.dev1
+Version: 1.4.1.dev1
 Summary: Ask LLMs to return structured JSON and run cross-model tests. API-first.
 Author-email: Juan Denis <juan@vene.co>
 License-Expression: MIT

{prompture-1.3.3.dev1 → prompture-1.4.1.dev1}/prompture/_version.py

@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
 commit_id: str | None
 __commit_id__: str | None
 
-__version__ = version = '1.3.3.dev1'
-__version_tuple__ = version_tuple = (1, 3, 3, 'dev1')
+__version__ = version = '1.4.1.dev1'
+__version_tuple__ = version_tuple = (1, 4, 1, 'dev1')
 
 __commit_id__ = commit_id = None

prompture-1.4.1.dev1/prompture/drivers/vertex_img_gen_driver.py

@@ -0,0 +1,135 @@
+"""Google Vertex AI image generation via Gemini image models.
+
+Unlike the Imagen driver (``generate_images``), the ``gemini-*-flash-image``
+models emit images through ``generate_content`` with ``response_modalities``
+including ``IMAGE`` — so this is a distinct driver. Uses the ``google-genai``
+SDK with ``vertexai=True``.
+
+Auth (same as the Vertex LLM driver): a Vertex Express API key
+(``GOOGLE_VERTEX_API_KEY``) for the express path, or project + ADC
+(``GOOGLE_VERTEX_PROJECT_ID`` / ``GOOGLE_VERTEX_LOCATION``).
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+from typing import Any
+
+from ..infra.cost_mixin import ImageCostMixin
+from ..media.image import image_from_bytes
+from .img_gen_base import ImageGenDriver
+
+logger = logging.getLogger(__name__)
+
+
+class VertexImageGenDriver(ImageCostMixin, ImageGenDriver):
+    """Image generation via Vertex AI Gemini image models (generate_content)."""
+
+    supports_multiple = False  # generate_content returns a single image
+    supports_size_variants = True
+    # Aspect ratios (Gemini image gen is aspect-ratio based, not pixel sizes).
+    supported_sizes = ["1:1", "16:9", "9:16", "4:3", "3:4"]
+    max_images = 1
+
+    # Per-image USD approximations; refine as Google publishes rates.
+    IMAGE_PRICING = {
+        "gemini-2.5-flash-image": {"default": 0.04},
+        "gemini-3.1-flash-image": {"default": 0.04},
+    }
+
+    KNOWN_MODELS = ["gemini-2.5-flash-image", "gemini-3.1-flash-image"]
+
+    def __init__(
+        self,
+        api_key: str | None = None,
+        project_id: str | None = None,
+        location: str | None = None,
+        model: str = "gemini-2.5-flash-image",
+    ):
+        self.api_key = api_key or os.getenv("GOOGLE_VERTEX_API_KEY")
+        self.project_id = project_id or os.getenv("GOOGLE_VERTEX_PROJECT_ID")
+        self.location = location or os.getenv("GOOGLE_VERTEX_LOCATION", "global")
+        self.model = model
+        self._client: Any = None
+
+    def _get_client(self) -> Any:
+        if self._client is None:
+            try:
+                from google import genai
+            except ImportError as exc:
+                raise RuntimeError("google-genai package is not installed") from exc
+
+            kwargs: dict[str, Any] = {"vertexai": True}
+            if self.api_key:
+                # Express mode: api key is mutually exclusive with project/location.
+                kwargs["api_key"] = self.api_key
+            elif self.project_id:
+                kwargs["project"] = self.project_id
+                kwargs["location"] = self.location
+            else:
+                raise RuntimeError(
+                    "Vertex image generation requires GOOGLE_VERTEX_API_KEY "
+                    "or GOOGLE_VERTEX_PROJECT_ID"
+                )
+            self._client = genai.Client(**kwargs)
+        return self._client
+
+    def generate_image(self, prompt: str, options: dict[str, Any]) -> dict[str, Any]:
+        """Generate an image via Vertex Gemini ``generate_content``.
+
+        Options: ``model``, ``aspect_ratio`` (or ``size``), ``image_size``
+        (e.g. "1K"/"2K"). Gemini image gen produces one image per call.
+        """
+        client = self._get_client()
+        from google.genai import types
+
+        model = options.get("model", self.model)
+        aspect = options.get("aspect_ratio") or options.get("size") or "1:1"
+        # If a pixel size slipped through (e.g. "1024x1024"), don't send it as an
+        # aspect ratio — fall back to square.
+        if "x" in str(aspect) and ":" not in str(aspect):
+            aspect = "1:1"
+        image_size = options.get("image_size", "1K")
+
+        config = types.GenerateContentConfig(
+            response_modalities=["TEXT", "IMAGE"],
+            image_config=types.ImageConfig(aspect_ratio=aspect, image_size=image_size),
+        )
+
+        try:
+            resp = client.models.generate_content(model=model, contents=prompt, config=config)
+        except Exception as exc:
+            name = type(exc).__name__
+            if "Blocked" in name or "blocked" in str(exc).lower() or "safety" in str(exc).lower():
+                raise ValueError(f"Prompt blocked by safety filter: {exc}") from exc
+            raise
+
+        images = []
+        for cand in (getattr(resp, "candidates", None) or []):
+            content = getattr(cand, "content", None)
+            for part in (getattr(content, "parts", None) or []):
+                inline = getattr(part, "inline_data", None)
+                data = getattr(inline, "data", None) if inline is not None else None
+                if data:
+                    mime = getattr(inline, "mime_type", None) or "image/png"
+                    images.append(image_from_bytes(data, media_type=mime))
+
+        if not images:
+            raise ValueError(
+                "Vertex returned no image — the model may have refused or replied with text only."
+            )
+
+        cost = self._calculate_image_cost("vertex_ai", model, n=len(images))
+
+        return {
+            "images": images,
+            "meta": {
+                "image_count": len(images),
+                "size": aspect,
+                "revised_prompt": None,
+                "cost": round(cost, 6),
+                "model_name": f"vertex_ai/{model}",
+                "raw_response": {},
+            },
+        }

{prompture-1.3.3.dev1 → prompture-1.4.1.dev1}/prompture/infra/discovery.py

@@ -1325,6 +1325,15 @@ def get_available_image_gen_models(
         for model_id in BFLImageGenDriver.KNOWN_MODELS:
             available.add(f"bfl/{model_id}")
 
+    # Vertex AI Gemini image models (generate_content with image output).
+    if _cfg_value(env, "google_vertex_api_key", "GOOGLE_VERTEX_API_KEY") or _cfg_value(
+        env, "google_vertex_project_id", "GOOGLE_VERTEX_PROJECT_ID"
+    ):
+        from ..drivers.vertex_img_gen_driver import VertexImageGenDriver
+
+        for model_id in VertexImageGenDriver.KNOWN_MODELS:
+            available.add(f"vertex_ai/{model_id}")
+
     # Dynamic discovery: check modalities_output from models.dev capabilities
     # for any models that the pricing dicts don't know about yet.
     from .model_rates import get_model_capabilities
@@ -1347,6 +1356,9 @@ _IMG_SIZE_PARAM: dict[str, str] = {
     "stability": "aspect_ratio",
     "ideogram": "aspect_ratio",
     "runway": "ratio",
+    "vertex_ai": "aspect_ratio",
+    "google_vertexai": "aspect_ratio",
+    "vertex": "aspect_ratio",
 }
 
 # Preferred representative ratio per orientation, used when a model offers it.
@@ -1371,6 +1383,9 @@ def _img_gen_driver_class(provider: str):
         if provider in ("grok", "xai"):
             from ..drivers.grok_img_gen_driver import GrokImageGenDriver
             return GrokImageGenDriver
+        if provider in ("vertex_ai", "google_vertexai", "vertex", "vertexai"):
+            from ..drivers.vertex_img_gen_driver import VertexImageGenDriver
+            return VertexImageGenDriver
         if provider == "stability":
             from ..drivers.stability_img_gen_driver import StabilityImageGenDriver
             return StabilityImageGenDriver

{prompture-1.3.3.dev1 → prompture-1.4.1.dev1}/prompture/plugins/builtins/google_plugin.py

@@ -63,6 +63,13 @@ class GooglePlugin(ProviderPlugin):
             "location": "google_vertex_location",
             "access_token": "google_vertex_access_token",  # nosec B105
         }
+        # Image gen uses google-genai with vertexai=True; access_token isn't a
+        # constructor arg for the image driver, so use a trimmed kwarg map.
+        vertex_img_kw = {
+            "api_key": "google_vertex_api_key",
+            "project_id": "google_vertex_project_id",
+            "location": "google_vertex_location",
+        }
         vertex_desc = ProviderDescriptor(
             name="google_vertexai",
             **_llm(
@@ -73,6 +80,11 @@ class GooglePlugin(ProviderPlugin):
                 vertex_kw,
                 "google_vertex_model",
             ),
+            img_gen_sync=DriverSpec(
+                "vertex_img_gen_driver.VertexImageGenDriver",
+                vertex_img_kw,
+                "gemini-2.5-flash-image",
+            ),
             display_name="Google Vertex AI",
             is_configured_fn=_google_vertexai_is_configured,
             list_models_kwargs=[
@@ -87,6 +99,9 @@ class GooglePlugin(ProviderPlugin):
             google_desc,
             vertex_desc,
             build_alias("gemini", google_desc, {"llm", "img_gen"}),
-            build_alias("vertex", vertex_desc, {"llm"}),
-            build_alias("vertexai", vertex_desc, {"llm"}),
+            build_alias("vertex", vertex_desc, {"llm", "img_gen"}),
+            build_alias("vertexai", vertex_desc, {"llm", "img_gen"}),
+            # The dashboard (and OpenAI-style "provider/model" ids) use the
+            # "vertex_ai" prefix — alias it for both LLM and image gen.
+            build_alias("vertex_ai", vertex_desc, {"llm", "img_gen"}),
         ]

{prompture-1.3.3.dev1 → prompture-1.4.1.dev1}/prompture.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: prompture
-Version: 1.3.3.dev1
+Version: 1.4.1.dev1
 Summary: Ask LLMs to return structured JSON and run cross-model tests. API-first.
 Author-email: Juan Denis <juan@vene.co>
 License-Expression: MIT

{prompture-1.3.3.dev1 → prompture-1.4.1.dev1}/prompture.egg-info/SOURCES.txt

@@ -255,6 +255,7 @@ prompture/drivers/runway_video_gen_driver.py
 prompture/drivers/stability_img_gen_driver.py
 prompture/drivers/stt_base.py
 prompture/drivers/tts_base.py
+prompture/drivers/vertex_img_gen_driver.py
 prompture/drivers/video_gen_base.py
 prompture/drivers/video_gen_registry.py
 prompture/drivers/vision_helpers.py