From 3ad74d282b84a8682d75a906b335ec0fcae81264 Mon Sep 17 00:00:00 2001
From: Deepak K <89829542+Deepak-Kesavan@users.noreply.github.com>
Date: Thu, 4 Jul 2024 10:36:41 +0530
Subject: [PATCH] Added support to pass params for prompt service API calls
 (#69)

* Added support to pass params for prompt service API calls

* Changed from positional args to keyword args

* SDK version bump
---
 src/unstract/sdk/__init__.py |  2 +-
 src/unstract/sdk/prompt.py   | 57 ++++++++++++++++++++++++++----------
 2 files changed, 42 insertions(+), 17 deletions(-)

diff --git a/src/unstract/sdk/__init__.py b/src/unstract/sdk/__init__.py
index e73e016e..b1599aad 100644
--- a/src/unstract/sdk/__init__.py
+++ b/src/unstract/sdk/__init__.py
@@ -1,4 +1,4 @@
-__version__ = "0.35.0"
+__version__ = "0.36.0"
 
 
 def get_sdk_version():
diff --git a/src/unstract/sdk/prompt.py b/src/unstract/sdk/prompt.py
index 194e5dbd..d81b3fed 100644
--- a/src/unstract/sdk/prompt.py
+++ b/src/unstract/sdk/prompt.py
@@ -25,32 +25,52 @@ def __init__(
             tool (AbstractTool): Instance of AbstractTool
             prompt_host (str): Host of platform service
             prompt_host (str): Port of platform service
-
         """
         self.tool = tool
         self.base_url = SdkHelper.get_platform_base_url(prompt_host, prompt_port)
         self.bearer_token = tool.get_env_or_die(ToolEnv.PLATFORM_API_KEY)
 
-    def answer_prompt(self, payload: dict[str, Any]) -> dict[str, Any]:
-        return self._post_call("answer-prompt", payload)
-
-    def single_pass_extraction(self, payload: dict[str, Any]) -> dict[str, Any]:
-        return self._post_call("single-pass-extraction", payload)
-
-    def summarize(self, payload: dict[str, Any]) -> dict[str, Any]:
-        return self._post_call("summarize", payload)
-
-    def _post_call(self, url_path: str, payload: dict[str, Any]) -> dict[str, Any]:
+    def answer_prompt(
+        self, payload: dict[str, Any], params: Optional[dict[str, str]] = None
+    ) -> dict[str, Any]:
+        return self._post_call(
+            url_path="answer-prompt",
+            payload=payload,
+            params=params,
+        )
+
+    def single_pass_extraction(
+        self, payload: dict[str, Any], params: Optional[dict[str, str]] = None
+    ) -> dict[str, Any]:
+        return self._post_call(
+            url_path="single-pass-extraction",
+            payload=payload,
+            params=params,
+        )
+
+    def summarize(
+        self, payload: dict[str, Any], params: Optional[dict[str, str]] = None
+    ) -> dict[str, Any]:
+        return self._post_call(url_path="summarize", payload=payload, params=params)
+
+    def _post_call(
+        self,
+        url_path: str,
+        payload: dict[str, Any],
+        params: Optional[dict[str, str]] = None,
+    ) -> dict[str, Any]:
         """Invokes and communicates to prompt service to fetch response for the
         prompt.
 
         Args:
-            file_name (str): File in which the prompt is processed
-            outputs (dict): dict of all input data for the tool
-            tool_id (str): Unique ID of the tool to be processed
+            url_path (str): URL path to the service endpoint
+            payload (dict): Payload to send in the request body
+            params (dict, optional): Query parameters to include in the request
 
         Returns:
-            Sample return dict:
+            dict: Response from the prompt service
+
+            Sample Response:
             {
                 "status": "OK",
                 "error": "",
@@ -68,7 +88,12 @@ def _post_call(self, url_path: str, payload: dict[str, Any]) -> dict[str, Any]:
         headers: dict[str, str] = {"Authorization": f"Bearer {self.bearer_token}"}
         response: Response = Response()
         try:
-            response = requests.post(url, json=payload, headers=headers)
+            response = requests.post(
+                url=url,
+                json=payload,
+                headers=headers,
+                params=params,
+            )
             response.raise_for_status()
             result["status"] = "OK"
             result["structure_output"] = response.text