diff --git a/biome.json b/biome.json new file mode 100644 index 0000000..5ad6df5 --- /dev/null +++ b/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 + } + } +} diff --git a/jigsawstack/__init__.py b/jigsawstack/__init__.py index 7931505..d75a6f0 100644 --- a/jigsawstack/__init__.py +++ b/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"] diff --git a/jigsawstack/image_generation.py b/jigsawstack/image_generation.py new file mode 100644 index 0000000..8d4c3f0 --- /dev/null +++ b/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 + + + + \ No newline at end of file diff --git a/jigsawstack/request.py b/jigsawstack/request.py index 101e937..1a6b464 100644 --- a/jigsawstack/request.py +++ b/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]: diff --git a/jigsawstack/store.py b/jigsawstack/store.py index 51a0713..16ff78c 100644 --- a/jigsawstack/store.py +++ b/jigsawstack/store.py @@ -3,87 +3,23 @@ from .request import Request, RequestConfig from .async_request import AsyncRequest, AsyncRequestConfig from ._config import ClientConfig -from typing import Any, Dict, List, cast -from typing_extensions import NotRequired, TypedDict - - -class FileDeleteResponse(TypedDict): - success: bool -class KVGetParams(TypedDict): - key: str -class KVGetResponse(TypedDict): +class FileDeleteResponse(TypedDict): success: bool - value: str -class KVAddParams(TypedDict): - key: str - value: str - encrypt: NotRequired[bool] - byo_secret: NotRequired[str] - - -class KVAddResponse(TypedDict): - success: bool - class FileUploadParams(TypedDict): overwrite: bool filename: str content_type: NotRequired[str] - -class KV(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) - self.config = RequestConfig( - api_url=api_url, - api_key=api_key, - disable_request_logging=disable_request_logging, - ) - - def add(self, params: KVAddParams) -> KVAddResponse: - path = "/store/kv" - resp = Request( - config=self.config, - path=path, - params=cast(Dict[Any, Any], params), - verb="post", - ).perform_with_content() - return resp - - def get(self, key: str) -> KVGetResponse: - path = f"/store/kv/{key}" - resp = Request(config=self.config, path=path, verb="get").perform_with_content() - return resp - - def delete(self, key: str) -> KVGetResponse: - path = f"/store/kv/{key}" - resp = Request( - config=self.config, - path=path, - params=cast(Dict[Any, Any], params={}), - verb="delete", - ).perform_with_content() - return resp - - class Store(ClientConfig): config: RequestConfig - kv: KV def __init__( self, @@ -98,8 +34,6 @@ def __init__( disable_request_logging=disable_request_logging, ) - self.kv = KV(api_key, api_url, disable_request_logging) - def upload(self, file: bytes, options=FileUploadParams) -> Any: overwrite = options.get("overwrite") filename = options.get("filename") @@ -121,7 +55,7 @@ def upload(self, file: bytes, options=FileUploadParams) -> Any: return resp def get(self, key: str) -> Any: - path = f"/store/file/{key}" + path = f"/store/file/read/{key}" resp = Request( config=self.config, path=path, @@ -131,56 +65,11 @@ def get(self, key: str) -> Any: return resp def delete(self, key: str) -> FileDeleteResponse: - path = f"/store/file/{key}" + path = f"/store/file/read/{key}" resp = Request( config=self.config, path=path, - params=cast(Dict[Any, Any], params={}), - verb="delete", - ).perform_with_content() - return resp - - -class AsyncKV(ClientConfig): - - config: AsyncRequestConfig - - 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) - self.config = AsyncRequestConfig( - api_url=api_url, - api_key=api_key, - disable_request_logging=disable_request_logging, - ) - - async def add(self, params: KVAddParams) -> KVAddResponse: - path = "/store/kv" - resp = await AsyncRequest( - config=self.config, - path=path, - params=cast(Dict[Any, Any], params), - verb="post", - ).perform_with_content() - return resp - - async def get(self, key: str) -> KVGetResponse: - path = f"/store/kv/{key}" - resp = await AsyncRequest( - config=self.config, path=path, verb="get", params={} - ).perform_with_content() - return resp - - async def delete(self, key: str) -> KVGetResponse: - path = f"/store/kv/{key}" - resp = await AsyncRequest( - config=self.config, - path=path, - params=cast(Dict[Any, Any], params={}), + params=key, verb="delete", ).perform_with_content() return resp @@ -188,7 +77,6 @@ async def delete(self, key: str) -> KVGetResponse: class AsyncStore(ClientConfig): config: AsyncRequestConfig - kv: AsyncKV def __init__( self, @@ -202,7 +90,7 @@ def __init__( api_key=api_key, disable_request_logging=disable_request_logging, ) - self.kv = AsyncKV(api_key, api_url, disable_request_logging) + async def upload(self, file: bytes, options=FileUploadParams) -> Any: overwrite = options.get("overwrite") @@ -225,7 +113,7 @@ async def upload(self, file: bytes, options=FileUploadParams) -> Any: return resp async def get(self, key: str) -> Any: - path = f"/store/file/{key}" + path = f"/store/file/read/{key}" resp = await AsyncRequest( config=self.config, path=path, @@ -235,7 +123,7 @@ async def get(self, key: str) -> Any: return resp async def delete(self, key: str) -> FileDeleteResponse: - path = f"/store/file/{key}" + path = f"/store/file/read/{key}" resp = AsyncRequest( config=self.config, path=path, diff --git a/jigsawstack/validate.py b/jigsawstack/validate.py index 80dab07..daafa06 100644 --- a/jigsawstack/validate.py +++ b/jigsawstack/validate.py @@ -96,14 +96,14 @@ def email(self, params: EmailValidationParams) -> EmailValidationResponse: ).perform_with_content() return resp - def nsfw(self, url: str) -> NSFWResponse: + def nsfw(self, params: NSFWParams) -> NSFWResponse: path = f"/validate/nsfw" resp = Request( config=self.config, path=path, params=cast( Dict[Any, Any], - params={"url": url}, + params ), verb="post", ).perform_with_content() @@ -174,14 +174,14 @@ async def email(self, params: EmailValidationParams) -> EmailValidationResponse: ).perform_with_content() return resp - async def nsfw(self, url: str) -> NSFWResponse: + async def nsfw(self, params: NSFWParams) -> NSFWResponse: path = f"/validate/nsfw" resp = await AsyncRequest( config=self.config, path=path, params=cast( Dict[Any, Any], - params={"url": url}, + params, ), verb="post", ).perform_with_content() diff --git a/tests/test_store.py b/tests/test_store.py index c4862d9..bf03d3b 100644 --- a/tests/test_store.py +++ b/tests/test_store.py @@ -11,29 +11,41 @@ @pytest.mark.skip(reason="Skipping TestWebAPI class for now") -def test_async_kv_response(): - async def _test(): - client = AsyncJigsawStack() - try: - result = await client.store.kv.add( - {"key": "hello", "value": "world", "encrypt": False} - ) - logger.info(result) - assert result["success"] == True - except JigsawStackError as e: - pytest.fail(f"Unexpected JigsawStackError: {e}") +class TestAsyncFileOperations: + """ + Test class for async file operations. + Add your file operation tests here. + """ + + def test_async_file_upload(self): + # Template for future file upload tests + pass + + def test_async_file_retrieval(self): + # Template for future file retrieval tests + pass + + def test_async_file_deletion(self): + # Template for future file deletion tests + pass - asyncio.run(_test()) - -def test_async_retriev_kv_response(): +# Example file upload test +# Uncomment and modify as needed +""" +def test_async_file_upload_example(): async def _test(): client = AsyncJigsawStack() try: - result = await client.store.kv.get("hello") + file_content = b"test file content" + result = await client.store.upload( + file_content, + {"filename": "test.txt", "overwrite": True} + ) logger.info(result) - assert result["success"] == True + assert result["success"] == True except JigsawStackError as e: pytest.fail(f"Unexpected JigsawStackError: {e}") asyncio.run(_test()) +""" \ No newline at end of file