llms-py 3.0.0b2__py3-none-any.whl → 3.0.0b4__py3-none-any.whl

llms/providers/nvidia.py

@@ -0,0 +1,107 @@
+import json
+import time
+
+import aiohttp
+
+
+def install(ctx):
+    from llms.main import GeneratorBase
+
+    class NvidiaGenAi(GeneratorBase):
+        sdk = "nvidia/image"
+
+        def __init__(self, **kwargs):
+            super().__init__(**kwargs)
+            self.width = int(kwargs.get("width", 1024))
+            self.height = int(kwargs.get("height", 1024))
+            self.cfg_scale = float(kwargs.get("cfg_scale", 3))
+            self.steps = int(kwargs.get("steps", 20))
+            self.mode = kwargs.get("mode", "base")
+            self.gen_url = kwargs.get("api", "https://ai.api.nvidia.com/v1/genai")
+
+        def to_response(self, response, chat, started_at):
+            if "artifacts" in response:
+                for artifact in response["artifacts"]:
+                    base64 = artifact.get("base64")
+                    seed = artifact.get("seed")
+                    filename = f"{seed}.png"
+                    if "model" in chat:
+                        last_model = "/" in chat["model"] and chat["model"].split("/")[-1] or chat["model"]
+                        filename = f"{last_model}_{seed}.png"
+
+                    image_info = {
+                        "seed": seed,
+                    }
+                    relative_url, info = ctx.save_image_to_cache(base64, filename, image_info)
+                    return {
+                        "choices": [
+                            {
+                                "message": {
+                                    "role": "assistant",
+                                    "content": self.default_content,
+                                    "images": [
+                                        {
+                                            "type": "image_url",
+                                            "image_url": {
+                                                "url": relative_url,
+                                            },
+                                        }
+                                    ],
+                                }
+                            }
+                        ]
+                    }
+            raise Exception("No artifacts in response")
+
+        async def chat(self, chat, provider=None):
+            headers = self.get_headers(provider, chat)
+            if provider is not None:
+                chat["model"] = provider.provider_model(chat["model"]) or chat["model"]
+
+            prompt = ctx.last_user_prompt(chat)
+
+            gen_request = {
+                "prompt": prompt,
+            }
+            modalities = chat.get("modalities", ["text"])
+            if "image" in modalities:
+                image_config = chat.get("image_config", {})
+                aspect_ratio = image_config.get("aspect_ratio")
+                if aspect_ratio:
+                    dimension = ctx.app.aspect_ratios.get(aspect_ratio)
+                    if dimension:
+                        width, height = dimension.split("×")
+                        gen_request["width"] = int(width)
+                        gen_request["height"] = int(height)
+                    else:
+                        gen_request["width"] = self.width
+                        gen_request["height"] = self.height
+
+                gen_request["mode"] = self.mode
+                gen_request["cfg_scale"] = self.cfg_scale
+                gen_request["steps"] = self.steps
+
+            gen_url = f"{self.gen_url}/{chat['model']}"
+            ctx.log(f"POST {gen_url}")
+            ctx.log(self.gen_summary(gen_request))
+            # remove metadata if any (conflicts with some providers, e.g. Z.ai)
+            gen_request.pop("metadata", None)
+            started_at = time.time()
+
+            if ctx.MOCK:
+                ctx.log("Mocking NvidiaGenAi")
+                text = ctx.text_from_file(f"{ctx.MOCK_DIR}/nvidia-image.json")
+                return self.to_response(json.loads(text), chat, started_at)
+            else:
+                async with aiohttp.ClientSession() as session, session.post(
+                    gen_url,
+                    headers=headers,
+                    data=json.dumps(gen_request),
+                    timeout=aiohttp.ClientTimeout(total=120),
+                ) as response:
+                    return self.to_response(await self.response_json(response), chat, started_at)
+
+    ctx.add_provider(NvidiaGenAi)
+
+
+__install__ = install

llms/providers/openai.py

@@ -0,0 +1,159 @@
+import base64
+import json
+import mimetypes
+import time
+
+import aiohttp
+
+
+def install(ctx):
+    from llms.main import GeneratorBase, OpenAiCompatible
+
+    class OpenAiProvider(OpenAiCompatible):
+        sdk = "@ai-sdk/openai"
+
+        def __init__(self, **kwargs):
+            if "api" not in kwargs:
+                kwargs["api"] = "https://api.openai.com/v1"
+            super().__init__(**kwargs)
+            self.modalities["image"] = OpenAiGenerator(**kwargs)
+
+    # https://platform.openai.com/docs/api-reference/images
+    class OpenAiGenerator(GeneratorBase):
+        sdk = "openai/image"
+
+        def __init__(self, **kwargs):
+            super().__init__(**kwargs)
+            self.api = "https://api.openai.com/v1/images/generations"
+            self.map_image_models = kwargs.get(
+                "map_image_models",
+                {
+                    "gpt-5.1-codex-mini": "gpt-image-1-mini",
+                },
+            )
+
+        def aspect_ratio_to_size(self, aspect_ratio, model):
+            w, h = aspect_ratio.split(":")
+            width = int(w)
+            height = int(h)
+            if model == "dall-e-2":
+                return "1024x1024"
+            if model == "dall-e-3":
+                if width > height:
+                    return "1792x1024"
+                elif height > width:
+                    return "1024x1792"
+            if width > height:
+                return "1536x1024"
+            elif height > width:
+                return "1024x1536"
+            return "1024x1024"
+
+        async def to_response(self, response, chat, started_at):
+            # go through all image responses and save them to cache
+            # Try to extract and save images
+            images = []
+            if "data" in response:
+                for i, item in enumerate(response["data"]):
+                    image_url = item.get("url")
+                    b64_json = item.get("b64_json")
+
+                    ext = "png"
+                    image_data = None
+
+                    if b64_json:
+                        image_data = base64.b64decode(b64_json)
+                    elif image_url:
+                        ctx.log(f"GET {image_url}")
+                        async with aiohttp.ClientSession() as session, await session.get(image_url) as res:
+                            if res.status == 200:
+                                image_data = await res.read()
+                                content_type = res.headers.get("Content-Type")
+                                if content_type:
+                                    ext = mimetypes.guess_extension(content_type)
+                                    if ext:
+                                        ext = ext.lstrip(".")  # remove leading dot
+                                    # Fallback if guess_extension returns None or if we want to be safe
+                                    if not ext:
+                                        ext = "png"
+                            else:
+                                raise Exception(f"Failed to download image: {res.status}")
+
+                    if image_data:
+                        relative_url, info = ctx.save_image_to_cache(
+                            image_data,
+                            f"{chat['model']}-{i}.{ext}",
+                            {
+                                "model": chat["model"],
+                                "prompt": ctx.last_user_prompt(chat),
+                            },
+                        )
+                        images.append(
+                            {
+                                "type": "image_url",
+                                "image_url": {
+                                    "url": relative_url,
+                                },
+                            }
+                        )
+                    else:
+                        raise Exception("No image data found")
+
+                return {
+                    "choices": [
+                        {
+                            "message": {
+                                "role": "assistant",
+                                "content": self.default_content,
+                                "images": images,
+                            }
+                        }
+                    ]
+                }
+            if "error" in response:
+                raise Exception(response["error"]["message"])
+
+            ctx.log(json.dumps(response, indent=2))
+            raise Exception("No 'data' field in response.")
+
+        async def chat(self, chat, provider=None):
+            headers = self.get_headers(provider, chat)
+
+            if chat["model"] in self.map_image_models:
+                chat["model"] = self.map_image_models[chat["model"]]
+
+            aspect_ratio = "1:1"
+            if "image_config" in chat and "aspect_ratio" in chat["image_config"]:
+                aspect_ratio = chat["image_config"].get("aspect_ratio", "1:1")
+            payload = {
+                "model": chat["model"],
+                "prompt": ctx.last_user_prompt(chat),
+                "size": self.aspect_ratio_to_size(aspect_ratio, chat["model"]),
+            }
+            if provider is not None:
+                chat["model"] = provider.provider_model(chat["model"]) or chat["model"]
+
+            started_at = time.time()
+            if ctx.MOCK:
+                print("Mocking OpenAiGenerator")
+                text = ctx.text_from_file(f"{ctx.MOCK_DIR}/openai-image.json")
+                return await self.to_response(json.loads(text), chat, started_at)
+            else:
+                ctx.log(f"POST {self.api}")
+                # _log(json.dumps(headers, indent=2))
+                ctx.log(json.dumps(payload, indent=2))
+                async with aiohttp.ClientSession() as session, session.post(
+                    self.api, headers=headers, json=payload
+                ) as response:
+                    text = await response.text()
+                    ctx.log(text[:1024] + (len(text) > 1024 and "..." or ""))
+                    if response.status < 300:
+                        return ctx.log_json(await self.to_response(json.loads(text), chat, started_at))
+                    else:
+                        raise Exception(f"Failed to generate image {response.status}")
+
+    ctx.add_provider(OpenAiProvider)
+    ctx.add_provider(OpenAiGenerator)
+
+
+__install__ = install

llms/providers/openrouter.py

@@ -0,0 +1,70 @@
+import json
+import time
+
+import aiohttp
+
+
+def install(ctx):
+    from llms.main import GeneratorBase
+
+    # https://openrouter.ai/docs/guides/overview/multimodal/image-generation
+    class OpenRouterGenerator(GeneratorBase):
+        sdk = "openrouter/image"
+
+        def __init__(self, **kwargs):
+            super().__init__(**kwargs)
+
+        def to_response(self, response, chat, started_at):
+            # go through all image responses and save them to cache
+            for choice in response["choices"]:
+                if "message" in choice and "images" in choice["message"]:
+                    for image in choice["message"]["images"]:
+                        if choice["message"]["content"] == "":
+                            choice["message"]["content"] = self.default_content
+                        if "image_url" in image:
+                            data_uri = image["image_url"]["url"]
+                            if data_uri.startswith("data:"):
+                                parts = data_uri.split(",", 1)
+                                ext = parts[0].split(";")[0].split("/")[1]
+                                base64_data = parts[1]
+                                model = chat["model"].split("/")[-1]
+                                filename = f"{model}-{choice['index']}.{ext}"
+                                info = {
+                                    "model": model,
+                                    "prompt": ctx.last_user_prompt(chat),
+                                }
+                                relative_url, info = ctx.save_image_to_cache(base64_data, filename, info)
+                                image["image_url"]["url"] = relative_url
+
+            return response
+
+        async def chat(self, chat, provider=None):
+            headers = self.get_headers(provider, chat)
+            if provider is not None:
+                chat["model"] = provider.provider_model(chat["model"]) or chat["model"]
+
+            started_at = time.time()
+            if ctx.MOCK:
+                print("Mocking OpenRouterGenerator")
+                text = ctx.text_from_file(f"{ctx.MOCK_DIR}/openrouter-image.json")
+                return ctx.log_json(self.to_response(json.loads(text), chat, started_at))
+            else:
+                chat_url = provider.chat_url
+                chat = await self.process_chat(chat, provider_id=self.id)
+                ctx.log(f"POST {chat_url}")
+                ctx.log(provider.chat_summary(chat))
+                # remove metadata if any (conflicts with some providers, e.g. Z.ai)
+                chat.pop("metadata", None)
+
+                async with aiohttp.ClientSession() as session, session.post(
+                    chat_url,
+                    headers=headers,
+                    data=json.dumps(chat),
+                    timeout=aiohttp.ClientTimeout(total=300),
+                ) as response:
+                    return ctx.log_json(self.to_response(await self.response_json(response), chat, started_at))
+
+    ctx.add_provider(OpenRouterGenerator)
+
+
+__install__ = install

llms/providers-extra.json

@@ -0,0 +1,356 @@
+{
+    "openrouter": {
+        "models": {
+            "google/gemini-2.5-flash-image": {
+                "name": "Gemini 2.5 Flash Image (Nano Banana)",
+                "modalities": {
+                    "input": [
+                        "text"
+                    ],
+                    "output": [
+                        "image"
+                    ]
+                },
+                "cost": {
+                    "input": 0.30,
+                    "output": 2.50
+                }
+            },
+            "google/gemini-2.5-flash-image-preview": {
+                "name": "Gemini 2.5 Flash Image Preview (Nano Banana)",
+                "modalities": {
+                    "input": [
+                        "text"
+                    ],
+                    "output": [
+                        "image"
+                    ]
+                },
+                "cost": {
+                    "input": 0.30,
+                    "output": 2.50
+                }
+            },
+            "google/gemini-3-pro-image-preview": {
+                "name": "Google: Nano Banana Pro Preview (Gemini 3 Pro)",
+                "modalities": {
+                    "input": [
+                        "text"
+                    ],
+                    "output": [
+                        "image"
+                    ]
+                },
+                "cost": {
+                    "input": 2.0,
+                    "output": 12.0
+                }
+            },
+            "sourceful/riverflow-v2-fast-preview": {
+                "name": "Sourceful: Riverflow V2 Fast Preview",
+                "modalities": {
+                    "input": [
+                        "text"
+                    ],
+                    "output": [
+                        "image"
+                    ]
+                },
+                "cost": {
+                    "input": 0,
+                    "output": 7.19
+                }
+            },
+            "sourceful/riverflow-v2-standard-preview": {
+                "name": "Sourceful: Riverflow V2 Standard Preview",
+                "modalities": {
+                    "input": [
+                        "text"
+                    ],
+                    "output": [
+                        "image"
+                    ]
+                },
+                "cost": {
+                    "input": 0,
+                    "output": 8.38
+                }
+            },
+            "sourceful/riverflow-v2-max-preview": {
+                "name": "Sourceful: Riverflow V2 Max Preview",
+                "modalities": {
+                    "input": [
+                        "text"
+                    ],
+                    "output": [
+                        "image"
+                    ]
+                },
+                "cost": {
+                    "input": 0,
+                    "output": 17.96
+                }
+            },
+            "openai/gpt-5-image-mini": {
+                "name": "OpenAI: GPT-5 Image Mini",
+                "modalities": {
+                    "input": [
+                        "text"
+                    ],
+                    "output": [
+                        "image"
+                    ]
+                },
+                "cost": {
+                    "input": 2.50,
+                    "output": 2
+                }
+            },
+            "openai/gpt-5-image": {
+                "name": "OpenAI: GPT-5 Image",
+                "modalities": {
+                    "input": [
+                        "text"
+                    ],
+                    "output": [
+                        "image"
+                    ]
+                },
+                "cost": {
+                    "input": 10.0,
+                    "output": 10.0
+                }
+            },
+            "black-forest-labs/flux.2-pro": {
+                "name": "Black Forest Labs: FLUX.2 Pro",
+                "modalities": {
+                    "input": [
+                        "text"
+                    ],
+                    "output": [
+                        "image"
+                    ]
+                },
+                "cost": {
+                    "input": 3.66,
+                    "output": 3.66
+                }
+            },
+            "black-forest-labs/flux.2-max": {
+                "name": "Black Forest Labs: FLUX.2 Max",
+                "modalities": {
+                    "input": [
+                        "text"
+                    ],
+                    "output": [
+                        "image"
+                    ]
+                },
+                "cost": {
+                    "input": 7.32,
+                    "output": 7.32
+                }
+            },
+            "black-forest-labs/flux.2-flex": {
+                "name": "Black Forest Labs: FLUX.2 Flex",
+                "modalities": {
+                    "input": [
+                        "text"
+                    ],
+                    "output": [
+                        "image"
+                    ]
+                },
+                "cost": {
+                    "input": 14.64,
+                    "output": 14.64
+                }
+            }
+        }
+    },
+    "openai": {
+        "models": {
+            "chatgpt-image-latest": {
+                "name": "ChatGPT Image Latest",
+                "modalities": {
+                    "input": [
+                        "text"
+                    ],
+                    "output": [
+                        "image"
+                    ]
+                },
+                "cost": {
+                    "input": 5.0,
+                    "output": 32.0
+                }
+            },
+            "gpt-image-1.5": {
+                "name": "GPT Image 1.5",
+                "modalities": {
+                    "input": [
+                        "text"
+                    ],
+                    "output": [
+                        "image"
+                    ]
+                },
+                "cost": {
+                    "input": 5.0,
+                    "output": 32.0
+                }
+            },
+            "gpt-image-1": {
+                "name": "GPT Image 1",
+                "modalities": {
+                    "input": [
+                        "text"
+                    ],
+                    "output": [
+                        "image"
+                    ]
+                },
+                "cost": {
+                    "input": 5.0,
+                    "output": 40.0
+                }
+            },
+            "gpt-image-1-mini": {
+                "name": "GPT Image 1 Mini",
+                "modalities": {
+                    "input": [
+                        "text"
+                    ],
+                    "output": [
+                        "image"
+                    ]
+                },
+                "cost": {
+                    "input": 5.0,
+                    "output": 8.0
+                }
+            }
+        }
+    },
+    "chutes": {
+        "models": {
+            "chutes-z-image-turbo": {
+                "name": "Z Image Turbo",
+                "modalities": {
+                    "input": [
+                        "text"
+                    ],
+                    "output": [
+                        "image"
+                    ]
+                }
+            },
+            "qwen-image": {
+                "modalities": {
+                    "input": [
+                        "text"
+                    ],
+                    "output": [
+                        "image"
+                    ]
+                }
+            },
+            "chutes-qwen-image-edit-2509": {
+                "name": "Qwen Image Edit 2509",
+                "modalities": {
+                    "input": [
+                        "text",
+                        "image"
+                    ],
+                    "output": [
+                        "image"
+                    ]
+                }
+            },
+            "chutes-hidream": {
+                "name": "Hidream",
+                "modalities": {
+                    "input": [
+                        "text"
+                    ],
+                    "output": [
+                        "image"
+                    ]
+                }
+            },
+            "chutes-hunyuan-image-3": {
+                "name": "Hunyuan Image 3",
+                "modalities": {
+                    "input": [
+                        "text"
+                    ],
+                    "output": [
+                        "image"
+                    ]
+                }
+            },
+            "FLUX.1-schnell": {
+                "name": "FLUX.1 Schnell",
+                "modalities": {
+                    "input": [
+                        "text"
+                    ],
+                    "output": [
+                        "image"
+                    ]
+                }
+            },
+            "chroma": {
+                "modalities": {
+                    "input": [
+                        "text"
+                    ],
+                    "output": [
+                        "image"
+                    ]
+                }
+            },
+            "JuggernautXL-Ragnarok": {
+                "name": "JuggernautXL Ragnarok",
+                "modalities": {
+                    "input": [
+                        "text"
+                    ],
+                    "output": [
+                        "image"
+                    ]
+                }
+            },
+            "Animij": {
+                "modalities": {
+                    "input": [
+                        "text"
+                    ],
+                    "output": [
+                        "image"
+                    ]
+                }
+            },
+            "Illustrij": {
+                "modalities": {
+                    "input": [
+                        "text"
+                    ],
+                    "output": [
+                        "image"
+                    ]
+                }
+            },
+            "iLustMix": {
+                "modalities": {
+                    "input": [
+                        "text"
+                    ],
+                    "output": [
+                        "image"
+                    ]
+                }
+            }
+        }
+    }
+}