llms-py 2.0.9__py3-none-any.whl → 3.0.10__py3-none-any.whl

llms/extensions/providers/openai.py

@@ -0,0 +1,154 @@
+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, context=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 = ctx.chat_to_aspect_ratio(chat) or "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, context=context))
+                    else:
+                        raise Exception(f"Failed to generate image {response.status}")
+
+    ctx.add_provider(OpenAiProvider)
+    ctx.add_provider(OpenAiGenerator)
+
+
+__install__ = install_openai

llms/extensions/providers/openrouter.py

@@ -0,0 +1,74 @@
+import json
+import time
+
+import aiohttp
+
+
+def install_openrouter(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
+            cost = None
+            if "usage" in response and "cost" in response["usage"]:
+                cost = response["usage"]["cost"]
+            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}"
+                                relative_url, info = ctx.save_image_to_cache(
+                                    base64_data, filename, ctx.to_file_info(chat, {"cost": cost})
+                                )
+                                image["image_url"]["url"] = relative_url
+
+            return response
+
+        async def chat(self, chat, provider=None, context=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
+                # remove tools
+                chat.pop("tools", None)
+                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)
+                metadata = 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:
+                    if metadata:
+                        chat["metadata"] = metadata
+                    return ctx.log_json(
+                        self.to_response(await self.response_json(response), chat, started_at, context=context)
+                    )
+
+    ctx.add_provider(OpenRouterGenerator)

llms/extensions/providers/zai.py

@@ -0,0 +1,182 @@
+import json
+import time
+from typing import Optional
+
+import aiohttp
+
+
+def install_zai(ctx):
+    from llms.main import GeneratorBase
+
+    # https://docs.z.ai/guides/image/glm-image
+    class ZaiGenerator(GeneratorBase):
+        sdk = "zai/image"
+
+        def __init__(self, **kwargs):
+            super().__init__(**kwargs)
+            self.aspect_ratios = {
+                "1:1": "1280×1280",
+                "2:3": "1056×1568",
+                "3:2": "1568×1056",
+                "3:4": "1088×1472",
+                "4:3": "1472×1088",
+                "4:5": "1088×1472",
+                "5:4": "1472×1088",
+                "9:16": "960×1728",
+                "16:9": "1728×960",
+                "21:9": "1728×960",
+            }
+            self.model: str = kwargs.get("model", "glm-image")
+            self.n: Optional[int] = kwargs.get("n")
+            self.quality: Optional[str] = kwargs.get("quality")
+            self.response_format: Optional[str] = kwargs.get("response_format")
+            self.size: Optional[str] = kwargs.get("size")
+            self.style: Optional[str] = kwargs.get("style")
+            self.sensitive_word_check: Optional[str] = kwargs.get("sensitive_word_check")
+            self.user: Optional[str] = kwargs.get("user")
+            self.request_id: Optional[str] = kwargs.get("request_id")
+            self.user_id: Optional[str] = kwargs.get("user_id")
+            self.extra_headers: Optional[dict] = kwargs.get("extra_headers")
+            self.extra_body: Optional[dict] = kwargs.get("extra_body")
+            self.disable_strict_validation: Optional[bool] = kwargs.get("disable_strict_validation")
+            self.timeout: Optional[float] = float(kwargs.get("timeout") or 300)
+            self.watermark_enabled: Optional[bool] = kwargs.get("watermark_enabled")
+
+        async def chat(self, chat, provider=None, context=None):
+            headers = {"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"}
+            if self.extra_headers:
+                headers.update(self.extra_headers)
+
+            chat_url = "https://api.z.ai/api/paas/v4/images/generations"
+            if provider is not None:
+                headers["Authorization"] = f"Bearer {provider.api_key}"
+                chat["model"] = provider.provider_model(chat["model"]) or chat["model"]
+                chat_url = provider.api + "/images/generations"
+
+            body = {}
+            attrs = [
+                "model",
+                "n",
+                "quality",
+                "response_format",
+                "size",
+                "style",
+                "sensitive_word_check",
+                "user",
+                "request_id",
+                "user_id",
+                "disable_strict_validation",
+                "watermark_enabled",
+            ]
+            for attr in attrs:
+                if hasattr(self, attr) and getattr(self, attr) is not None:
+                    body[attr] = getattr(self, attr)
+
+            if self.extra_body:
+                body.update(self.extra_body)
+
+            if "model" in chat:
+                body["model"] = chat["model"]
+
+            body["prompt"] = ctx.last_user_prompt(chat)
+
+            aspect_ratio = ctx.chat_to_aspect_ratio(chat) or "1:1"
+            size = self.aspect_ratios.get(aspect_ratio, "1280x1280").replace("×", "x")
+            body["size"] = size
+
+            username = ctx.context_to_username(context)
+            if username:
+                body["user"] = username
+
+            ctx.dbg(f"ZaiProvider.chat: {chat_url}")
+            ctx.dbg(json.dumps(body, indent=2))
+            started_at = time.time()
+            async with aiohttp.ClientSession() as session, session.post(
+                chat_url,
+                headers=headers,
+                data=json.dumps(body),
+                timeout=aiohttp.ClientTimeout(total=self.timeout),
+            ) as response:
+                # Example Response
+                # {
+                #   "created": 1768451303,
+                #   "data": [
+                #     {
+                #       "url": "https://mfile.z.ai/1768451374203-b334959408a643a8a6c74eb104746dcb.png?ufileattname=202601151228236805d575507d4570_watermark.png"
+                #     }
+                #   ],
+                #   "id": "202601151228236805d575507d4570",
+                #   "request_id": "202601151228236805d575507d4570",
+                #   "usage": {
+                #     "tokens": 0,
+                #     "price": 0,
+                #     "cost": 0.0,
+                #     "duration": 71
+                #   },
+                #   "timestamp": 1768451374519,
+                #   "model": "GLM-Image"
+                # }
+
+                response_json = await self.response_json(response)
+                duration = int(time.time() - started_at)
+                usage = response_json.get("usage", {})
+                if context is not None:
+                    context["providerResponse"] = response_json
+                    if "cost" in usage:
+                        context["cost"] = usage.get("cost")
+
+                images = []
+                for image in response_json.get("data", []):
+                    url = image.get("url")
+                    if not url:
+                        continue
+                    # download url with aiohttp
+                    async with session.get(url) as image_response:
+                        headers = image_response.headers
+                        # get filename from Content-Disposition
+                        # attachment; filename="202601151228236805d575507d4570_watermark.png"
+                        mime_type = headers.get("Content-Type") or "image/png"
+                        disposition = headers.get("Content-Disposition")
+                        if disposition:
+                            start = disposition.index('filename="') + len('filename="')
+                            end = disposition.index('"', start)
+                            filename = disposition[start:end]
+                        else:
+                            ext = mime_type.split("/")[1]
+                            filename = f"{body['model'].lower()}-{response_json.get('id', int(started_at))}.{ext}"
+                        image_bytes = await image_response.read()
+
+                        info = {
+                            "prompt": body["prompt"],
+                            "type": mime_type,
+                            "width": int(size.split("x")[0]),
+                            "height": int(size.split("x")[1]),
+                            "duration": duration,
+                        }
+                        info.update(usage)
+                        cache_url, info = ctx.save_image_to_cache(
+                            image_bytes, filename, image_info=info, ignore_info=True
+                        )
+
+                    images.append(
+                        {
+                            "type": "image_url",
+                            "image_url": {
+                                "url": cache_url,
+                            },
+                        }
+                    )
+
+                chat_response = {
+                    "choices": [{"message": {"role": "assistant", "content": self.default_content, "images": images}}],
+                    "created": int(time.time()),
+                    "usage": {
+                        "prompt_tokens": 0,
+                        "completion_tokens": 1_000_000,  # Price per image is 0.015, so 1M token is 0.015
+                    },
+                }
+                if "cost" in usage:
+                    chat_response["cost"] = usage["cost"]
+                return ctx.log_json(chat_response)
+
+    ctx.add_provider(ZaiGenerator)

llms/extensions/system_prompts/README.md

@@ -0,0 +1,22 @@
+# System Prompts Extension
+
+This extension configures AI requests with a library of **over 200+** awesome curated system prompts that can be selected from the UI.
+
+## Custom System Prompts
+
+You can also maintain your own library of system prompts which can be maintained for all anonymous users at:
+`~/.llms/user/default/system-prompts.json`
+
+Or for signed in users at:
+`~/.llms/user/<github-user>/system-prompts.json`
+
+The JSON file should contain an array of Prompt objects, e.g:
+
+```json
+[
+    {
+        "name": "Helpful Assistant",
+        "prompt": "You are a helpful assistant."
+    }
+]
+```

llms/extensions/system_prompts/__init__.py

@@ -0,0 +1,45 @@
+import json
+import os
+
+from aiohttp import web
+
+default_prompts = [
+    {"name": "Helpful Assistant", "prompt": "You are a helpful assistant."},
+]
+
+
+# runs after providers are configured but before server is run
+def install(ctx):
+    # helper to get user or default prompts
+    def get_user_prompts(request):
+        candidate_paths = []
+        # check if user is signed in
+        username = ctx.get_username(request)
+        if username:
+            # if signed in (Github OAuth), return the prompts for this user if exists
+            candidate_paths.append(os.path.join(ctx.get_user_path(username), "system_prompts", "prompts.json"))
+        # return default prompts for all users if exists
+        candidate_paths.append(os.path.join(ctx.get_user_path(), "system_prompts", "prompts.json"))
+        # otherwise return the default prompts from this repo
+        candidate_paths.append(os.path.join(ctx.path, "ui", "prompts.json"))
+
+        # iterate all candidate paths and when exists return its json
+        for path in candidate_paths:
+            if os.path.exists(path):
+                with open(path, encoding="utf-8") as f:
+                    txt = f.read()
+                    return json.loads(txt)
+        return default_prompts
+
+    # API Handler to get prompts
+    async def get_prompts(request):
+        prompts_json = get_user_prompts(request)
+        return web.json_response(prompts_json)
+
+    ctx.add_get("prompts.json", get_prompts)
+
+
+# register install extension handler
+__install__ = install
+
+__order__ = -10