Merge pull request #39 from JigsawStack/pysdk/updates

Author: yoeven

biome.json

@@ -0,0 +1,37 @@
+{
+  "$schema": "https://biomejs.dev/schemas/1.9.4/schema.json",
+  "files": {
+    "ignoreUnknown": false,
+    "ignore": []
+  },
+  "formatter": {
+    "enabled": true,
+    "useEditorconfig": true,
+    "formatWithErrors": false,
+    "indentStyle": "space",
+    "indentWidth": 2,
+    "lineEnding": "lf",
+    "lineWidth": 150,
+    "attributePosition": "auto",
+    "bracketSpacing": true
+  },
+  "organizeImports": {
+    "enabled": true
+  },
+  "linter": {
+    "enabled": false
+  },
+  "javascript": {
+    "formatter": {
+      "jsxQuoteStyle": "double",
+      "quoteProperties": "asNeeded",
+      "trailingCommas": "es5",
+      "semicolons": "always",
+      "arrowParentheses": "always",
+      "bracketSameLine": false,
+      "quoteStyle": "double",
+      "attributePosition": "auto",
+      "bracketSpacing": true
+    }
+  }
+}

jigsawstack/__init__.py

@@ -15,11 +15,13 @@
 from .prompt_engine import PromptEngine, AsyncPromptEngine
 from .embedding import Embedding, AsyncEmbedding
 from .exceptions import JigsawStackError
+from .image_generation import ImageGeneration, AsyncImageGeneration
 
 
 class JigsawStack:
     audio: Audio
     vision: Vision
+    image_generation: ImageGeneration
     file: Store
     web: Web
     search: Search
@@ -116,6 +118,11 @@ def __init__(
             api_url=api_url,
             disable_request_logging=disable_request_logging,
         ).execute
+        self.image_generation = ImageGeneration(
+            api_key=api_key,
+            api_url=api_url,
+            disable_request_logging=disable_request_logging,
+        ).image_generation
 
 
 class AsyncJigsawStack:
@@ -124,6 +131,7 @@ class AsyncJigsawStack:
     web: AsyncWeb
     audio: AsyncAudio
     vision: AsyncVision
+    image_generation: AsyncImageGeneration
     store: AsyncStore
     prompt_engine: AsyncPromptEngine
     api_key: str
@@ -227,6 +235,12 @@ def __init__(
             disable_request_logging=disable_request_logging,
         ).execute
 
+        self.image_generation = AsyncImageGeneration(
+            api_key=api_key,
+            api_url=api_url,
+            disable_request_logging=disable_request_logging,
+        ).image_generation
+
 
 # Create a global instance of the Web class
 __all__ = ["JigsawStack", "Search", "JigsawStackError", "AsyncJigsawStack"]

jigsawstack/image_generation.py

@@ -0,0 +1,103 @@
+from typing import Any, Dict, List, Union, cast
+from typing_extensions import NotRequired, TypedDict, Literal, Required
+from .request import Request, RequestConfig
+from .async_request import AsyncRequest
+
+from typing import List, Union
+from ._config import ClientConfig
+
+class ImageGenerationParams(TypedDict):
+    prompt: Required[str]
+    """"
+    The text to generate the image from."
+    """
+    aspect_ratio: NotRequired[Literal["1:1", "16:9", "21:9", "3:2", "2:3", "4:5", "5:4", "3:4", "4:3", "9:16", "9:21"]]
+    """
+    The aspect ratio of the image. The default is 1:1.
+    """
+    width: NotRequired[int]
+    """"
+    The width of the image. The default is 512."
+    """
+    height: NotRequired[int]
+    """
+    The height of the image. The default is 512."
+    """
+    steps: NotRequired[int]
+    """"
+    The number of steps to generate the image.""
+    """
+    advance_config: NotRequired[Dict[str, Union[int, str]]]
+    """
+    The advance configuration for the image generation. The default is None.
+    You can pass the following:
+    - `seed`: The seed for the image generation. The default is None.
+    - `guidance`: The guidance for the image generation. The default is None.
+    - `negative_prompt`: The negative prompt for the image generation. The default is None.
+    """
+
+class ImageGenerationResponse(TypedDict):
+    success: bool
+    """
+    Indicates whether the image generation was successful.
+    """
+    image: bytes
+    """
+    The generated image as a blob.
+    """
+
+class ImageGeneration(ClientConfig):
+    config: RequestConfig
+
+    def __init__(
+        self,
+        api_key: str,
+        api_url: str,
+        disable_request_logging: Union[bool, None] = False,
+    ):
+        super().__init__(api_key, api_url, disable_request_logging=disable_request_logging)
+        self.config = RequestConfig(
+            api_url=api_url,
+            api_key=api_key,
+            disable_request_logging=disable_request_logging,
+        )
+
+    def image_generation(self, params: ImageGenerationParams) -> ImageGenerationResponse:
+        path = "/ai/image_generation"
+        resp = Request(
+            config=self.config,
+            path=path,
+            params=cast(Dict[Any, Any], params), # type: ignore
+            verb="post",
+        ).perform()
+        return resp
+    
+class AsyncImageGeneration(ClientConfig):
+    config: RequestConfig
+
+    def __init__(
+        self,
+        api_key: str,
+        api_url: str,
+        disable_request_logging: Union[bool, None] = False,
+    ):
+        super().__init__(api_key, api_url, disable_request_logging=disable_request_logging)
+        self.config = RequestConfig(
+            api_url=api_url,
+            api_key=api_key,
+            disable_request_logging=disable_request_logging,
+        )
+
+    async def image_generation(self, params: ImageGenerationParams) -> ImageGenerationResponse:
+        path = "/ai/image_generation"
+        resp = await AsyncRequest(
+            config=self.config,
+            path=path,
+            params=cast(Dict[Any, Any], params), # type: ignore
+            verb="post",
+        ).perform()
+        return resp
+    
+
+
+    

jigsawstack/request.py

@@ -56,7 +56,9 @@ def perform(self) -> Union[T, None]:
 
         # this is a safety net, if we get here it means the JigsawStack API is having issues
         # and most likely the gateway is returning htmls
-        if "application/json" not in resp.headers["content-type"] and "audio/wav" not in resp.headers["content-type"]:
+        if "application/json" not in resp.headers["content-type"] \
+            and "audio/wav" not in resp.headers["content-type"] \
+            and "image/png" not in resp.headers["content-type"]:
             raise_for_code_and_type(
                 code=500,
                 message="Failed to parse JigsawStack API response. Please try again.",
@@ -72,9 +74,9 @@ def perform(self) -> Union[T, None]:
                 err=error.get("error"),
             )
 
-        if "audio/wav" in resp.headers["content-type"]:
+        if "audio/wav" or "image/png" in resp.headers["content-type"]:
             return cast(T, resp) # we return the response object, instead of the json
-
+        
         return cast(T, resp.json())
 
     def perform_file(self) -> Union[T, None]: