clean up

Author: iamNarcisse

.gitignore

@@ -6,4 +6,7 @@ jigsawstack.egg-info/
 .env
 build/  
 dist/ 
-/demo
+/demo
+
+test.py
+test_web.py

jigsawstack/prompt_engine.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, List, Union, cast
+from typing import Any, Dict, List, Union, cast, Generator
 from typing_extensions import NotRequired, TypedDict
 from .request import Request, RequestConfig
 from typing import List, Union
@@ -15,13 +15,13 @@ class PromptEngineRunParams(TypedDict):
     prompt: str
     inputs: NotRequired[List[object]]
     return_prompt: Union[str, List[object], Dict[str, str]]
-    # grok_key: NotRequired[str]
     input_values: NotRequired[Dict[str, str]]
 
 
 class PromptEngineExecuteParams(TypedDict):
     id: str
     input_values: object
+    stream: Union[bool, None] = False
 
 
 class PromptEngineRunResponse(TypedDict):
@@ -138,9 +138,22 @@ def run_prompt_direct(
         ).perform_with_content()
         return resp
 
-    def run(self, params: PromptEngineExecuteParams) -> PromptEngineRunResponse:
+    def run(
+        self, params: PromptEngineExecuteParams
+    ) -> Union[PromptEngineRunResponse, Generator[Any, None, None]]:
         id = params.get("id")
         path = f"/prompt_engine/{id}"
+        stream = params.get("stream")
+
+        if stream:
+            resp = Request(
+                config=self.config,
+                path=path,
+                params=cast(Dict[Any, Any], params),
+                verb="post",
+            ).perform_with_content_streaming()
+            return resp
+
         resp = Request(
             config=self.config,
             path=path,

jigsawstack/request.py

@@ -1,7 +1,8 @@
-from typing import Any, Dict, Generic, List, Union, cast, TypedDict
+from typing import Any, Dict, Generic, List, Union, cast, TypedDict, Generator
 import requests
 from typing_extensions import Literal, TypeVar
 from .exceptions import NoContentError, raise_for_code_and_type
+import json
 
 RequestVerb = Literal["get", "post", "put", "patch", "delete"]
 
@@ -24,6 +25,7 @@ def __init__(
         verb: RequestVerb,
         headers: Dict[str, str] = {"Content-Type": "application/json"},
         data: Union[bytes, None] = None,
+        stream: Union[bool, None] = False,
     ):
         self.path = path
         self.params = params
@@ -33,6 +35,7 @@ def __init__(
         self.data = data
         self.headers = headers
         self.disable_request_logging = config.get("disable_request_logging")
+        self.stream = stream
 
     def perform(self) -> Union[T, None]:
         """Is the main function that makes the HTTP request
@@ -152,6 +155,75 @@ def __get_headers(self) -> Dict[Any, Any]:
 
         return _headers
 
+    def perform_streaming(self) -> Generator[Union[T, str], None, None]:
+        """Is the main function that makes the HTTP request
+        to the JigsawStack API. It uses the path, params, and verb attributes
+        to make the request.
+
+        Returns:
+            Generator[bytes, None, None]: A generator of bytes
+
+        Raises:
+            requests.HTTPError: If the request fails
+        """
+        resp = self.make_request(url=f"{self.api_url}{self.path}")
+
+        # delete calls do not return a body
+        if resp.text == "":
+            return 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"]:
+            raise_for_code_and_type(
+                code=500,
+                message="Failed to parse JigsawStack API response. Please try again.",
+            )
+
+        if resp.status_code != 200:
+            error = resp.json()
+            raise_for_code_and_type(
+                code=resp.status_code,
+                message=error.get("message"),
+                err=error.get("error"),
+            )
+
+        def try_parse_data(chunk: bytes) -> Union[T, str]:
+            if not chunk:
+                return chunk
+            # Decode bytes to text
+            text = chunk.decode("utf-8")
+
+            try:
+                # Try to parse as JSON
+                return json.loads(text)
+            except json.JSONDecodeError:
+                # Return as text if not valid JSON
+                return text
+
+        # Yield content in chunks
+        def chunk_generator():
+            for chunk in resp.iter_content(chunk_size=1024):  # 1KB chunks
+                if chunk:  # Filter out keep-alive new chunks
+                    yield try_parse_data(chunk)
+
+        return chunk_generator()
+
+    def perform_with_content_streaming(self) -> T:
+        """
+        Perform an HTTP request and return the response content as a streaming response.
+
+        Returns:
+            T: The content of the response
+
+        Raises:
+            NoContentError: If the response content is `None`.
+        """
+        resp = self.perform_streaming()
+        if resp is None:
+            raise NoContentError()
+        return resp
+
     def make_request(self, url: str) -> requests.Response:
         """make_request is a helper function that makes the actual
         HTTP request to the JigsawStack API.
@@ -183,6 +255,7 @@ def make_request(self, url: str) -> requests.Response:
                 json=params,
                 headers=headers,
                 data=data,
+                stream=self.stream,
             )
         except requests.HTTPError as e:
             raise e

tests/test_prompt_engine.py

@@ -3,16 +3,20 @@
 from jigsawstack.exceptions import JigsawStackError
 import jigsawstack
 import pytest
+
 # flake8: noq
 
 jigsaw = jigsawstack.JigsawStack()
 
+
 @pytest.mark.skip(reason="Skipping TestWebAPI class for now")
 class TestPromptEngine(unittest.TestCase):
 
     def test_get_prompt_engine_response_success(self) -> None:
         try:
-            result = jigsaw.prompt_engine.get("b08921b8-0b30-409e-8257-06fa1620c7e6")
+            result = jigsaw.prompt_engine.run(
+                {"id": "b08921b8-0b30-409e-8257-06fa1620c7e6", "stream": True}
+            )
             assert result["success"] == True
         except JigsawStackError as e:
-            assert e.message == "Failed to parse API response. Please try again."
+            assert e.message == "Failed to parse API response. Please try again."

tests/test_search.py

@@ -3,6 +3,7 @@
 from jigsawstack.exceptions import JigsawStackError
 import jigsawstack
 import pytest
+
 # flake8: noq
 
 jigsaw = jigsawstack.JigsawStack()
@@ -12,21 +13,17 @@
 class TestSearchAPI(unittest.TestCase):
 
     def test_search_suggestion_response_success(self) -> None:
-        params = {
-            "query": "Time Square New Yor"
-        }
+        params = {"query": "Time Square New Yor"}
         try:
             result = jigsaw.search.suggestion(params)
             assert result["success"] == True
         except JigsawStackError as e:
             assert e.message == "Failed to parse API response. Please try again."
 
     def test_ai_search_response_success(self) -> None:
-        params = {
-            "query": "Time Square New Yor"
-        }
+        params = {"query": "Time Square New Yor"}
         try:
             result = jigsaw.search.ai_search(params)
             assert result["success"] == True
         except JigsawStackError as e:
-            assert e.message == "Failed to parse API response. Please try again."
+            assert e.message == "Failed to parse API response. Please try again."

tests/test_web.py

@@ -4,25 +4,28 @@
 from jigsawstack import JigsawStack
 
 import pytest
+
 # flake8: noqa
 
 client = JigsawStack()
+
+
 @pytest.mark.skip(reason="Skipping TestWebAPI class for now")
 class TestWebAPI(unittest.TestCase):
     def test_ai_scrape_success_response(self) -> None:
         params = {
-         "url": "https://supabase.com/pricing",
-         "element_prompts": ["Plan title", "Plan price"],
+            "url": "https://supabase.com/pricing",
+            "element_prompts": ["Plan title", "Plan price"],
         }
         try:
-            result =client.file.upload(params)
+            result = client.file.upload(params)
             assert result["success"] == True
         except JigsawStackError as e:
             assert e.message == "Failed to parse API response. Please try again."
 
     def test_scrape_success_response(self) -> None:
         params = {
-         "url": "https://supabase.com/pricing",
+            "url": "https://supabase.com/pricing",
         }
         try:
             result = client.web.scrape(params)
@@ -33,10 +36,10 @@ def test_scrape_success_response(self) -> None:
     def test_dns_success_response(self) -> None:
 
         params = {
-                "url": "https://supabase.com/pricing",
+            "url": "https://supabase.com/pricing",
         }
         try:
             result = client.web.dns(params)
             assert result["success"] == True
         except JigsawStackError as e:
-            assert e.message == "Failed to parse API response. Please try again."
+            assert e.message == "Failed to parse API response. Please try again."