llms-py 2.0.35__py3-none-any.whl → 3.0.0__py3-none-any.whl

llms/extensions/providers/google.py

@@ -0,0 +1,378 @@
+import base64
+import io
+import json
+import time
+import wave
+
+import aiohttp
+
+# class GoogleOpenAiProvider(OpenAiCompatible):
+#     sdk = "google-openai-compatible"
+
+#     def __init__(self, api_key, **kwargs):
+#         super().__init__(api="https://generativelanguage.googleapis.com", api_key=api_key, **kwargs)
+#         self.chat_url = "https://generativelanguage.googleapis.com/v1beta/chat/completions"
+
+
+def install_google(ctx):
+    from llms.main import OpenAiCompatible
+
+    def gemini_chat_summary(gemini_chat):
+        """Summarize Gemini chat completion request for logging. Replace inline_data with size of content only"""
+        clone = json.loads(json.dumps(gemini_chat))
+        for content in clone["contents"]:
+            for part in content["parts"]:
+                if "inline_data" in part:
+                    data = part["inline_data"]["data"]
+                    part["inline_data"]["data"] = f"({len(data)})"
+        return json.dumps(clone, indent=2)
+
+    def gemini_response_summary(obj):
+        to = {}
+        for k, v in obj.items():
+            if k == "candidates":
+                candidates = []
+                for candidate in v:
+                    c = {}
+                    for ck, cv in candidate.items():
+                        if ck == "content":
+                            content = {}
+                            for content_k, content_v in cv.items():
+                                if content_k == "parts":
+                                    parts = []
+                                    for part in content_v:
+                                        p = {}
+                                        for pk, pv in part.items():
+                                            if pk == "inlineData":
+                                                p[pk] = {
+                                                    "mimeType": pv.get("mimeType"),
+                                                    "data": f"({len(pv.get('data'))})",
+                                                }
+                                            else:
+                                                p[pk] = pv
+                                        parts.append(p)
+                                    content[content_k] = parts
+                                else:
+                                    content[content_k] = content_v
+                            c[ck] = content
+                        else:
+                            c[ck] = cv
+                    candidates.append(c)
+                to[k] = candidates
+            else:
+                to[k] = v
+        return to
+
+    class GoogleProvider(OpenAiCompatible):
+        sdk = "@ai-sdk/google"
+
+        def __init__(self, **kwargs):
+            new_kwargs = {"api": "https://generativelanguage.googleapis.com", **kwargs}
+            super().__init__(**new_kwargs)
+            self.safety_settings = kwargs.get("safety_settings")
+            self.thinking_config = kwargs.get("thinking_config")
+            self.speech_config = kwargs.get("speech_config")
+            self.tools = kwargs.get("tools")
+            self.curl = kwargs.get("curl")
+            self.headers = kwargs.get("headers", {"Content-Type": "application/json"})
+            # Google fails when using Authorization header, use query string param instead
+            if "Authorization" in self.headers:
+                del self.headers["Authorization"]
+
+        async def chat(self, chat):
+            chat["model"] = self.provider_model(chat["model"]) or chat["model"]
+
+            chat = await self.process_chat(chat)
+            generation_config = {}
+
+            # Filter out system messages and convert to proper Gemini format
+            contents = []
+            system_prompt = None
+
+            async with aiohttp.ClientSession() as session:
+                for message in chat["messages"]:
+                    if message["role"] == "system":
+                        content = message["content"]
+                        if isinstance(content, list):
+                            for item in content:
+                                if "text" in item:
+                                    system_prompt = item["text"]
+                                    break
+                        elif isinstance(content, str):
+                            system_prompt = content
+                    elif "content" in message:
+                        if isinstance(message["content"], list):
+                            parts = []
+                            for item in message["content"]:
+                                if "type" in item:
+                                    if item["type"] == "image_url" and "image_url" in item:
+                                        image_url = item["image_url"]
+                                        if "url" not in image_url:
+                                            continue
+                                        url = image_url["url"]
+                                        if not url.startswith("data:"):
+                                            raise Exception("Image was not downloaded: " + url)
+                                        # Extract mime type from data uri
+                                        mimetype = url.split(";", 1)[0].split(":", 1)[1] if ";" in url else "image/png"
+                                        base64_data = url.split(",", 1)[1]
+                                        parts.append({"inline_data": {"mime_type": mimetype, "data": base64_data}})
+                                    elif item["type"] == "input_audio" and "input_audio" in item:
+                                        input_audio = item["input_audio"]
+                                        if "data" not in input_audio:
+                                            continue
+                                        data = input_audio["data"]
+                                        format = input_audio["format"]
+                                        mimetype = f"audio/{format}"
+                                        parts.append({"inline_data": {"mime_type": mimetype, "data": data}})
+                                    elif item["type"] == "file" and "file" in item:
+                                        file = item["file"]
+                                        if "file_data" not in file:
+                                            continue
+                                        data = file["file_data"]
+                                        if not data.startswith("data:"):
+                                            raise (Exception("File was not downloaded: " + data))
+                                        # Extract mime type from data uri
+                                        mimetype = (
+                                            data.split(";", 1)[0].split(":", 1)[1]
+                                            if ";" in data
+                                            else "application/octet-stream"
+                                        )
+                                        base64_data = data.split(",", 1)[1]
+                                        parts.append({"inline_data": {"mime_type": mimetype, "data": base64_data}})
+                                if "text" in item:
+                                    text = item["text"]
+                                    parts.append({"text": text})
+                            if len(parts) > 0:
+                                contents.append(
+                                    {
+                                        "role": message["role"]
+                                        if "role" in message and message["role"] == "user"
+                                        else "model",
+                                        "parts": parts,
+                                    }
+                                )
+                        else:
+                            content = message["content"]
+                            contents.append(
+                                {
+                                    "role": message["role"]
+                                    if "role" in message and message["role"] == "user"
+                                    else "model",
+                                    "parts": [{"text": content}],
+                                }
+                            )
+
+                gemini_chat = {
+                    "contents": contents,
+                }
+
+                if self.safety_settings:
+                    gemini_chat["safetySettings"] = self.safety_settings
+
+                # Add system instruction if present
+                if system_prompt is not None:
+                    gemini_chat["systemInstruction"] = {"parts": [{"text": system_prompt}]}
+
+                if "max_completion_tokens" in chat:
+                    generation_config["maxOutputTokens"] = chat["max_completion_tokens"]
+                if "stop" in chat:
+                    generation_config["stopSequences"] = [chat["stop"]]
+                if "temperature" in chat:
+                    generation_config["temperature"] = chat["temperature"]
+                if "top_p" in chat:
+                    generation_config["topP"] = chat["top_p"]
+                if "top_logprobs" in chat:
+                    generation_config["topK"] = chat["top_logprobs"]
+
+                if "thinkingConfig" in chat:
+                    generation_config["thinkingConfig"] = chat["thinkingConfig"]
+                elif self.thinking_config:
+                    generation_config["thinkingConfig"] = self.thinking_config
+
+                if len(generation_config) > 0:
+                    gemini_chat["generationConfig"] = generation_config
+
+                if "tools" in chat:
+                    # gemini_chat["tools"] = chat["tools"]
+                    ctx.log("Error: tools not supported in Gemini")
+                elif self.tools:
+                    # gemini_chat["tools"] = self.tools.copy()
+                    ctx.log("Error: tools not supported in Gemini")
+
+                if "modalities" in chat:
+                    generation_config["responseModalities"] = [modality.upper() for modality in chat["modalities"]]
+                    if "image" in chat["modalities"] and "image_config" in chat:
+                        # delete thinkingConfig
+                        del generation_config["thinkingConfig"]
+                        config_map = {
+                            "aspect_ratio": "aspectRatio",
+                            "image_size": "imageSize",
+                        }
+                        generation_config["imageConfig"] = {
+                            config_map[k]: v for k, v in chat["image_config"].items() if k in config_map
+                        }
+                    if "audio" in chat["modalities"] and self.speech_config:
+                        del generation_config["thinkingConfig"]
+                        generation_config["speechConfig"] = self.speech_config.copy()
+                        # Currently Google Audio Models only accept AUDIO
+                        generation_config["responseModalities"] = ["AUDIO"]
+
+                started_at = int(time.time() * 1000)
+                gemini_chat_url = f"https://generativelanguage.googleapis.com/v1beta/models/{chat['model']}:generateContent?key={self.api_key}"
+
+                ctx.log(f"POST {gemini_chat_url}")
+                ctx.log(gemini_chat_summary(gemini_chat))
+                started_at = time.time()
+
+                if ctx.MOCK and "modalities" in chat:
+                    print("Mocking Google Gemini Image")
+                    with open(f"{ctx.MOCK_DIR}/gemini-image.json") as f:
+                        obj = json.load(f)
+                else:
+                    try:
+                        async with session.post(
+                            gemini_chat_url,
+                            headers=self.headers,
+                            data=json.dumps(gemini_chat),
+                            timeout=aiohttp.ClientTimeout(total=120),
+                        ) as res:
+                            obj = await self.response_json(res)
+                    except Exception as e:
+                        ctx.log(f"Error: {res.status} {res.reason}: {e}")
+                        text = await res.text()
+                        try:
+                            obj = json.loads(text)
+                        except:
+                            ctx.log(text)
+                            raise e
+
+                if "error" in obj:
+                    ctx.log(f"Error: {obj['error']}")
+                    raise Exception(obj["error"]["message"])
+
+                if ctx.debug:
+                    ctx.dbg(json.dumps(gemini_response_summary(obj), indent=2))
+
+                # calculate cost per generation
+                cost = None
+                token_costs = obj.get("metadata", {}).get("pricing", "")
+                if token_costs:
+                    input_price, output_price = token_costs.split("/")
+                    input_per_token = float(input_price) / 1000000
+                    output_per_token = float(output_price) / 1000000
+                    if "usageMetadata" in obj:
+                        input_tokens = obj["usageMetadata"].get("promptTokenCount", 0)
+                        output_tokens = obj["usageMetadata"].get("candidatesTokenCount", 0)
+                        cost = (input_per_token * input_tokens) + (output_per_token * output_tokens)
+
+                response = {
+                    "id": f"chatcmpl-{started_at}",
+                    "created": started_at,
+                    "model": obj.get("modelVersion", chat["model"]),
+                }
+                choices = []
+                for i, candidate in enumerate(obj["candidates"]):
+                    role = "assistant"
+                    if "content" in candidate and "role" in candidate["content"]:
+                        role = "assistant" if candidate["content"]["role"] == "model" else candidate["content"]["role"]
+
+                    # Safely extract content from all text parts
+                    content = ""
+                    reasoning = ""
+                    images = []
+                    audios = []
+                    if "content" in candidate and "parts" in candidate["content"]:
+                        text_parts = []
+                        reasoning_parts = []
+                        for part in candidate["content"]["parts"]:
+                            if "text" in part:
+                                if "thought" in part and part["thought"]:
+                                    reasoning_parts.append(part["text"])
+                                else:
+                                    text_parts.append(part["text"])
+                            if "inlineData" in part:
+                                inline_data = part["inlineData"]
+                                mime_type = inline_data.get("mimeType", "image/png")
+                                if mime_type.startswith("image"):
+                                    ext = mime_type.split("/")[1]
+                                    base64_data = inline_data["data"]
+                                    filename = f"{chat['model'].split('/')[-1]}-{len(images)}.{ext}"
+                                    ctx.log(f"inlineData {len(base64_data)} {mime_type} {filename}")
+                                    relative_url, info = ctx.save_image_to_cache(
+                                        base64_data,
+                                        filename,
+                                        ctx.to_file_info(chat, {"cost": cost}),
+                                    )
+                                    images.append(
+                                        {
+                                            "type": "image_url",
+                                            "index": len(images),
+                                            "image_url": {
+                                                "url": relative_url,
+                                            },
+                                        }
+                                    )
+                                elif mime_type.startswith("audio"):
+                                    # mime_type audio/L16;codec=pcm;rate=24000
+                                    base64_data = inline_data["data"]
+
+                                    pcm = base64.b64decode(base64_data)
+                                    # Convert PCM to WAV
+                                    wav_io = io.BytesIO()
+                                    with wave.open(wav_io, "wb") as wf:
+                                        wf.setnchannels(1)
+                                        wf.setsampwidth(2)
+                                        wf.setframerate(24000)
+                                        wf.writeframes(pcm)
+                                    wav_data = wav_io.getvalue()
+
+                                    ext = mime_type.split("/")[1].split(";")[0]
+                                    pcm_filename = f"{chat['model'].split('/')[-1]}-{len(audios)}.{ext}"
+                                    filename = pcm_filename.replace(f".{ext}", ".wav")
+                                    ctx.log(f"inlineData {len(base64_data)} {mime_type} {filename}")
+
+                                    relative_url, info = ctx.save_bytes_to_cache(
+                                        wav_data,
+                                        filename,
+                                        ctx.to_file_info(chat, {"cost": cost}),
+                                    )
+
+                                    audios.append(
+                                        {
+                                            "type": "audio_url",
+                                            "index": len(audios),
+                                            "audio_url": {
+                                                "url": relative_url,
+                                            },
+                                        }
+                                    )
+                        content = " ".join(text_parts)
+                        reasoning = " ".join(reasoning_parts)
+
+                    choice = {
+                        "index": i,
+                        "finish_reason": candidate.get("finishReason", "stop"),
+                        "message": {
+                            "role": role,
+                            "content": content,
+                        },
+                    }
+                    if reasoning:
+                        choice["message"]["reasoning"] = reasoning
+                    if len(images) > 0:
+                        choice["message"]["images"] = images
+                    if len(audios) > 0:
+                        choice["message"]["audios"] = audios
+                    choices.append(choice)
+                response["choices"] = choices
+                if "usageMetadata" in obj:
+                    usage = obj["usageMetadata"]
+                    response["usage"] = {
+                        "completion_tokens": usage["candidatesTokenCount"],
+                        "total_tokens": usage["totalTokenCount"],
+                        "prompt_tokens": usage["promptTokenCount"],
+                    }
+
+                return ctx.log_json(self.to_response(response, chat, started_at))
+
+    ctx.add_provider(GoogleProvider)

llms/extensions/providers/nvidia.py

@@ -0,0 +1,105 @@
+import json
+import time
+
+import aiohttp
+
+
+def install_nvidia(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"
+
+                    relative_url, info = ctx.save_image_to_cache(
+                        base64,
+                        filename,
+                        ctx.to_file_info(chat, {"seed": seed}),
+                    )
+                    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)

llms/extensions/providers/openai.py

@@ -0,0 +1,156 @@
+import base64
+import json
+import mimetypes
+import time
+
+import aiohttp
+
+
+def install_openai(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}",
+                            ctx.to_file_info(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_openai