[FIX] Fixes for v0.81.0 (#592)

* Fixes for v0.81.0

Signed-off-by: Deepak <89829542+Deepak-Kesavan@users.noreply.github.com>

* Minor refactor

Signed-off-by: Deepak <89829542+Deepak-Kesavan@users.noreply.github.com>

* Improvements in the code

Signed-off-by: Deepak <89829542+Deepak-Kesavan@users.noreply.github.com>

---------

Signed-off-by: Deepak <89829542+Deepak-Kesavan@users.noreply.github.com>

Author: Deepak-Kesavan

backend/api/api_deployment_views.py

@@ -74,28 +74,28 @@ def get(
         self, request: Request, org_name: str, api_name: str, api: APIDeployment
     ) -> Response:
         execution_id = request.query_params.get("execution_id")
-        include_metadata = request.query_params.get("include_metadata", False)
+        include_metadata = (
+            request.query_params.get(ApiExecution.INCLUDE_METADATA, "false").lower()
+            == "true"
+        )
         if not execution_id:
             raise InvalidAPIRequest("execution_id shouldn't be empty")
         response: ExecutionResponse = DeploymentHelper.get_execution_status(
             execution_id=execution_id
         )
-        if response.execution_status != CeleryTaskState.SUCCESS.value:
-            return Response(
-                {
-                    "status": response.execution_status,
-                    "message": response.result,
-                },
-                status=status.HTTP_422_UNPROCESSABLE_ENTITY,
-            )
-        if (
-            response.execution_status == CeleryTaskState.SUCCESS.value
-            and not include_metadata
-        ):
-            response.remove_result_metadata_keys()
+        response_status = status.HTTP_422_UNPROCESSABLE_ENTITY
+        if response.execution_status == CeleryTaskState.COMPLETED.value:
+            response_status = status.HTTP_200_OK
+            if include_metadata:
+                response.remove_result_metadata_keys(keys_to_remove=["highlight_data"])
+            else:
+                response.remove_result_metadata_keys()
         return Response(
-            {"status": response.execution_status, "message": response.result},
-            status=status.HTTP_200_OK,
+            data={
+                "status": response.execution_status,
+                "message": response.result,
+            },
+            status=response_status,
         )
 
 

backend/api/deployment_helper.py

@@ -178,9 +178,10 @@ def execute_workflow(
             result.status_api = DeploymentHelper.construct_status_endpoint(
                 api_endpoint=api.api_endpoint, execution_id=execution_id
             )
-            if not include_metadata:
+            if include_metadata:
+                result.remove_result_metadata_keys(keys_to_remove=["highlight_data"])
+            else:
                 result.remove_result_metadata_keys()
-            cls._send_notification(api=api, result=result)
         except Exception as error:
             DestinationConnector.delete_api_storage_dir(
                 workflow_id=workflow_id, execution_id=execution_id

backend/api/postman_collection/dto.py

@@ -136,7 +136,10 @@ def get_api_endpoint(self) -> str:
 
     def _get_status_api_request(self) -> RequestItem:
         header_list = [HeaderItem(key="Authorization", value=f"Bearer {self.api_key}")]
-        status_query_param = {"execution_id": CollectionKey.STATUS_EXEC_ID_DEFAULT}
+        status_query_param = {
+            "execution_id": CollectionKey.STATUS_EXEC_ID_DEFAULT,
+            ApiExecution.INCLUDE_METADATA: "False",
+        }
         status_query_str = urlencode(status_query_param)
         abs_api_endpoint = urljoin(settings.WEB_APP_ORIGIN_URL, self.api_endpoint)
         status_url = urljoin(abs_api_endpoint, "?" + status_query_str)

backend/prompt_studio/prompt_studio_core/constants.py

@@ -89,6 +89,7 @@ class ToolStudioPromptKeys:
     METADATA = "metadata"
     INCLUDE_METADATA = "include_metadata"
     PLATFORM_POSTAMBLE = "platform_postamble"
+    SUMMARIZE_AS_SOURCE = "summarize_as_source"
 
 
 class FileViewTypes:

backend/prompt_studio/prompt_studio_core/prompt_studio_helper.py

@@ -29,7 +29,6 @@
     IndexingAPIError,
     NoPromptsFound,
     PermissionError,
-    ToolNotValid,
 )
 from prompt_studio.prompt_studio_core.models import CustomTool
 from prompt_studio.prompt_studio_core.prompt_ide_base_tool import PromptIdeBaseTool
@@ -324,10 +323,6 @@ def index_document(
             )
             file_path = str(Path(file_path) / file_name)
 
-        if not tool:
-            logger.error(f"No tool instance found for the ID {tool_id}")
-            raise ToolNotValid()
-
         logger.info(f"[{tool_id}] Indexing started for doc: {file_name}")
         PromptStudioHelper._publish_log(
             {"tool_id": tool_id, "run_id": run_id, "doc_name": file_name},
@@ -751,6 +746,7 @@ def _fetch_response(
         tool_settings[TSPKeys.SINGLE_PASS_EXTRACTION_MODE] = (
             tool.single_pass_extraction_mode
         )
+        tool_settings[TSPKeys.SUMMARIZE_AS_SOURCE] = tool.summarize_as_source
         tool_settings[TSPKeys.PREAMBLE] = tool.preamble
         tool_settings[TSPKeys.POSTAMBLE] = tool.postamble
         tool_settings[TSPKeys.GRAMMAR] = grammar_list

backend/utils/enums.py

@@ -9,3 +9,4 @@ class CeleryTaskState(Enum):
     REVOKED = "REVOKED"
     STARTED = "STARTED"
     SUCCESS = "SUCCESS"
+    COMPLETED = "COMPLETED"

frontend/src/components/deployments/display-code/DisplayCode.jsx

@@ -8,17 +8,16 @@ import CodeSnippet from "./CodeSnippet.jsx";
 import "./DisplayCode.css";
 
 const DisplayCode = ({ isDialogOpen, setDialogOpen, url }) => {
-  const handleCloseDialog = () => {
-    setDialogOpen(false);
-  };
   const [copied, setCopied] = useState(false);
   const [language, setLanguage] = useState("python");
   const [code, setCode] = useState("");
   const [activeTabKey, setActiveTabKey] = useState("POST");
+  const params =
+    "?execution_id=REPLACE_WITH_EXECUTION_ID&include_metadata=False";
 
-  useEffect(() => {
-    generateCode();
-  }, [url, language, activeTabKey]);
+  const handleCloseDialog = () => {
+    setDialogOpen(false);
+  };
 
   const generateCode = () => {
     if (language === "python") {
@@ -51,7 +50,7 @@ const DisplayCode = ({ isDialogOpen, setDialogOpen, url }) => {
         files=[('files',('file',open(filepath,'rb'),'application/octet-stream'))]
         response = requests.request("POST", api_url, headers=headers, data=payload, files=files)
       {{else}}
-        api_url = '{{url}}?execution_id=REPLACE_WITH_EXECUTION_ID'
+        api_url = '{{url}}{{{params}}}'
         headers = {
           'Authorization': 'Bearer REPLACE_WITH_API_KEY'
         }
@@ -62,7 +61,11 @@ const DisplayCode = ({ isDialogOpen, setDialogOpen, url }) => {
     code = trimIndent(code);
     const template = Handlebars.compile(code);
 
-    const pythonCode = template({ url, isPost: activeTabKey === "POST" });
+    const pythonCode = template({
+      url,
+      params,
+      isPost: activeTabKey === "POST",
+    });
     setCode(pythonCode);
   };
 
@@ -74,14 +77,15 @@ const DisplayCode = ({ isDialogOpen, setDialogOpen, url }) => {
     --form 'timeout=300' \\
     --form 'include_metadata=false'
     {{else}}
-    curl --location '{{url}}?execution_id=REPLACE_WITH_EXECUTION_ID' \\
+    curl --location '{{url}}{{{params}}}' \\
     --header 'Authorization: Bearer REPLACE_WITH_API_KEY'
     {{/if}}
   `;
     code = trimIndent(code);
     const template = Handlebars.compile(code);
     const curlCode = template({
       url,
+      params,
       isPost: activeTabKey === "POST",
       pathToFile: "/path/to/file",
     });
@@ -100,7 +104,7 @@ const DisplayCode = ({ isDialogOpen, setDialogOpen, url }) => {
     fetch("{{url}}", requestOptions)
     {{else}}
     var requestOptions = { method: 'GET', redirect: 'follow', headers: myHeaders};
-    fetch("{{url}}?execution_id=REPLACE_WITH_EXECUTION_ID", requestOptions)
+    fetch("{{url}}{{{params}}}", requestOptions)
     {{/if}}
     .then(response => response.text())
     .then(result => console.log(result))
@@ -109,7 +113,7 @@ const DisplayCode = ({ isDialogOpen, setDialogOpen, url }) => {
     code = trimIndent(code);
     const template = Handlebars.compile(code);
 
-    const jsCode = template({ url, isPost: activeTabKey === "POST" });
+    const jsCode = template({ url, params, isPost: activeTabKey === "POST" });
     setCode(jsCode);
   };
 
@@ -155,6 +159,10 @@ const DisplayCode = ({ isDialogOpen, setDialogOpen, url }) => {
     },
   ];
 
+  useEffect(() => {
+    generateCode();
+  }, [url, language, activeTabKey]);
+
   return (
     <Modal
       title="Code Snippets"

prompt-service/src/unstract/prompt_service/constants.py

@@ -62,6 +62,7 @@ class PromptServiceContants:
     PLATFORM_POSTAMBLE = "platform_postamble"
     EXTRACT_EPILOGUE = "extract-epilogue"
     CLEAN_CONTEXT = "clean-context"
+    SUMMARIZE_AS_SOURCE = "summarize_as_source"
 
 
 class LogLevel(Enum):

prompt-service/src/unstract/prompt_service/helper.py

@@ -212,13 +212,16 @@ def construct_and_run_prompt(
     prompt: str,
     metadata: dict[str, Any],
 ) -> str:
+    platform_postamble = tool_settings.get(PSKeys.PLATFORM_POSTAMBLE, "")
+    if tool_settings.get(PSKeys.SUMMARIZE_AS_SOURCE):
+        platform_postamble = ""
     prompt = construct_prompt(
         preamble=tool_settings.get(PSKeys.PREAMBLE, ""),
         prompt=output[prompt],
         postamble=tool_settings.get(PSKeys.POSTAMBLE, ""),
         grammar_list=tool_settings.get(PSKeys.GRAMMAR, []),
         context=context,
-        platform_postamble=tool_settings.get(PSKeys.PLATFORM_POSTAMBLE, ""),
+        platform_postamble=platform_postamble,
     )
     return run_completion(
         llm=llm,

tools/structure/src/main.py

@@ -63,6 +63,7 @@ def run(
         tool_settings[SettingsKeys.ENABLE_SINGLE_PASS_EXTRACTION] = (
             single_pass_extraction_mode
         )
+        tool_settings[SettingsKeys.SUMMARIZE_AS_SOURCE] = summarize_as_source
 
         prompt_service_resp = None
         _, file_name = os.path.split(input_file)