Merge pull request #48 from JigsawStack/update/params

Update/params

Author: yoeven

jigsawstack/image_generation.py

@@ -6,34 +6,51 @@
 from typing import List, Union
 from ._config import ClientConfig
 
+class AdvanceConfig(TypedDict):
+    negative_prompt: NotRequired[str]
+    guidance: NotRequired[int]
+    seed: NotRequired[int]
+
 class ImageGenerationParams(TypedDict):
     prompt: Required[str]
-    """"
-    The text to generate the image from."
+    """
+    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."
+    """
+    The width of the image. The default is 512.
     """
     height: NotRequired[int]
     """
-    The height of the image. The default is 512."
+    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 number of steps to generate the image.
+    """
+    output_format: NotRequired[Literal["png", "svg"]]
+    """
+    The output format of the generated image. Can be 'png' or 'svg'.
+    """
+    advance_config: NotRequired[AdvanceConfig]
     """
     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.
+    - `negative_prompt`: The negative prompt for the image generation
+    - `guidance`: The guidance scale for the image generation
+    - `seed`: The seed for reproducible generation
+    """
+    url: NotRequired[str]
+    """
+    URL to use as image input.
+    """
+    file_store_key: NotRequired[str]
+    """
+    File store key to use as image input.
     """
 
 class ImageGenerationResponse(TypedDict):

jigsawstack/search.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, List, Union, cast
+from typing import Any, Dict, List, Union, cast, Literal
 from typing_extensions import NotRequired, TypedDict
 from .request import Request, RequestConfig
 from .async_request import AsyncRequest, AsyncRequestConfig
@@ -9,30 +9,61 @@
 class SearchResponse(TypedDict):
     success: bool
     """
-    Indicates whether the translation was successful.
+    Whether the search request was successful
     """
-    results: List[Any]
 
-    is_safe: bool
+    query: str
+    """
+    The search query that was used
+    """
 
     ai_overview: str
+    """
+    AI-generated overview/summary of the search results
+    or deep research results if enabled
+    """
 
-    spell_fixed: str
+    results: List[Any]
+    """
+    List of search result items
+    """
+
+    is_safe: bool
+    """
+    Whether the search results passed safe search filtering
+    """
+
+    spell_fixed: bool
+    """
+    Whether the query was spell-checked and fixed
+    """
 
     geo_results: List[Any]
+    """
+    List of location/geographic search results if applicable
+    """
 
     image_urls: List[str]
+    """
+    List of image URLs found in the search results
+    """
 
     links: List[str]
+    """
+    List of web page URLs found in the search results
+    """
 
 
 class SearchSuggestionsResponse(TypedDict):
     success: bool
     """
-    Indicates whether the translation was successful.
+    Whether the search suggestions request was successful
     """
 
     suggestions: List[str]
+    """
+    List of search suggestions
+    """
 
 
 class SearchSuggestionsParams(TypedDict):
@@ -41,26 +72,59 @@ class SearchSuggestionsParams(TypedDict):
     The search value. The maximum query character length is 200.
     """
 
+class DeepResearchConfig(TypedDict):
+    max_depth: NotRequired[int]
+    max_breadth: NotRequired[int]
+    max_output_tokens: NotRequired[int]
+    target_output_tokens: NotRequired[int]
 
 class SearchParams(TypedDict):
     query: str
     """
-    The search value. The maximum query character length is 200.
+    The search query string to execute
     """
-    ai_overview: NotRequired[bool] = True
+
+    spell_check: NotRequired[bool]
     """
-    Include AI powered overview in the search results. The default value is True
+    Whether to perform spell checking on the query. Defaults to True.
     """
-    safe_search: NotRequired[str] = "moderate"
+
+    safe_search: NotRequired[Literal["strict", "moderate", "off"]]
     """
-    Include offensive results in the search results. The default value is "moderate". Supported values:  moderate, strict, off
+    Safe search filtering level. Can be 'strict', 'moderate', or 'off'
     """
-    spell_check: NotRequired[bool] = True
+
+    ai_overview: NotRequired[bool]
     """
-    Spell check the search query.
+    Whether to generate an AI-powered overview of the search results. Defaults to True.
     """
 
     byo_urls: NotRequired[List[str]]
+    """
+    List of custom URLs to include in the search results
+    """
+
+    country_code: NotRequired[str]
+    """
+    Two-letter country code to localize search results (e.g. 'US', 'GB')
+    """
+
+    auto_scrape: bool
+    """
+    Whether to automatically scrape content from search result URLs
+    """
+
+    deep_research: NotRequired[bool]
+    """
+    Enable deep research mode for more comprehensive results
+    """
+
+    deep_research_config: NotRequired[DeepResearchConfig]
+    """
+    Configuration options for deep research mode
+    """
+
+
 
 
 class Search(ClientConfig):

tests/test_image_generation.py

@@ -0,0 +1,53 @@
+from unittest.mock import MagicMock
+import unittest
+from jigsawstack.exceptions import JigsawStackError
+import jigsawstack
+import pytest
+import asyncio
+import logging
+import io
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+jigsaw = jigsawstack.JigsawStack()
+async_jigsaw = jigsawstack.AsyncJigsawStack()
+
+
+def test_image_generation_response():
+    async def _test():
+        client = jigsawstack.AsyncJigsawStack()
+        try:
+            result = await client.image_generation({
+                "prompt": "A beautiful mountain landscape at sunset",
+                "aspect_ratio": "16:9"
+            })
+            # Just check if we got some data back
+            assert result is not None
+            assert len(result) > 0
+        except JigsawStackError as e:
+            pytest.fail(f"Unexpected JigsawStackError: {e}")
+
+    asyncio.run(_test())
+
+
+def test_image_generation_with_advanced_config():
+    async def _test():
+        client = jigsawstack.AsyncJigsawStack()
+        try:
+            result = await client.image_generation({
+                "prompt": "A beautiful mountain landscape at sunset",
+                "output_format": "png",
+                "advance_config": {
+                    "negative_prompt": "blurry, low quality",
+                    "guidance": 7,
+                    "seed": 42
+                }
+            })
+            # Just check if we got some data back
+            assert result is not None
+            assert len(result) > 0
+        except JigsawStackError as e:
+            pytest.fail(f"Unexpected JigsawStackError: {e}")
+
+    asyncio.run(_test())