update prompt engine

Author: iamNarcisse

jigsawstack/prompt_engine.py

@@ -16,6 +16,8 @@ class PromptEngineRunParams(TypedDict):
     inputs: NotRequired[List[object]]
     return_prompt: Union[str, List[object], Dict[str, str]]
     input_values: NotRequired[Dict[str, str]]
+    stream: Union[bool, None] = False
+    use_internet: Union[bool, None] = False
 
 
 class PromptEngineExecuteParams(TypedDict):
@@ -31,9 +33,10 @@ class PromptEngineRunResponse(TypedDict):
 
 class PromptEngineCreateParams(TypedDict):
     prompt: str
-    # grok_key: NotRequired[str]
     inputs: NotRequired[List[object]]
     return_prompt: Union[str, List[object], Dict[str, str]]
+    use_internet: Union[bool, None] = False
+    optimize_prompt: Union[bool, None] = False
 
 
 class PromptEngineCreateResponse(TypedDict):
@@ -128,8 +131,18 @@ def delete(self, id: str) -> PromptEngineDeleteResponse:
 
     def run_prompt_direct(
         self, params: PromptEngineRunParams
-    ) -> PromptEngineRunResponse:
+    ) -> Union[PromptEngineRunResponse, Generator[Any, None, None]]:
         path = "/prompt_engine/run"
+        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

@@ -172,14 +172,6 @@ def perform_streaming(self) -> Generator[Union[T, str], None, None]:
         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(
@@ -209,7 +201,7 @@ def chunk_generator():
 
         return chunk_generator()
 
-    def perform_with_content_streaming(self) -> T:
+    def perform_with_content_streaming(self) -> Generator[Union[T, str], None, None]:
         """
         Perform an HTTP request and return the response content as a streaming response.
 

setup.py

@@ -6,7 +6,7 @@
 
 setup(
     name="jigsawstack",
-    version="0.1.17",
+    version="0.1.18",
     description="JigsawStack Python SDK",
     long_description=open("README.md", encoding="utf8").read(),
     long_description_content_type="text/markdown",