Fix highlight not working

Zipstack · Feb 4, 2025 · eb83572 · eb83572
1 parent 9535e2b
commit eb83572
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 3 deletions.
diff --git a/backend/prompt_studio/prompt_studio_core_v2/prompt_ide_base_tool.py b/backend/prompt_studio/prompt_studio_core_v2/prompt_ide_base_tool.py
@@ -5,6 +5,14 @@
 from unstract.sdk.constants import LogLevel
 from unstract.sdk.tool.stream import StreamMixin
 
+from backend.constants import FeatureFlag
+from unstract.flags.feature_flag import check_feature_flag_status
+
+if check_feature_flag_status(FeatureFlag.REMOTE_FILE_STORAGE):
+    from unstract.sdk.file_storage.constants import StorageType
+    from unstract.sdk.file_storage.env_helper import EnvHelper
+    from utils.file_storage.constants import FileStorageKeys
+
 
 class PromptIdeBaseTool(StreamMixin):
     def __init__(self, log_level: LogLevel = LogLevel.INFO, org_id: str = "") -> None:
@@ -17,6 +25,12 @@ def __init__(self, log_level: LogLevel = LogLevel.INFO, org_id: str = "") -> Non
         self.log_level = log_level
         self.org_id = org_id
         self.workflow_filestorage = None
+        if check_feature_flag_status(FeatureFlag.REMOTE_FILE_STORAGE):
+
+            self.workflow_filestorage = EnvHelper.get_storage(
+                storage_type=StorageType.PERMANENT,
+                env_name=FileStorageKeys.PERMANENT_REMOTE_STORAGE,
+            )
 
         super().__init__(log_level=log_level)
 

diff --git a/prompt-service/src/unstract/prompt_service/helper.py b/prompt-service/src/unstract/prompt_service/helper.py
@@ -233,6 +233,7 @@ def construct_and_run_prompt(
     prompt: str,
     metadata: dict[str, Any],
     file_path: str = "",
+    execution_source: Optional[str] = ExecutionSource.IDE.value,
 ) -> str:
     platform_postamble = tool_settings.get(PSKeys.PLATFORM_POSTAMBLE, "")
     summarize_as_source = tool_settings.get(PSKeys.SUMMARIZE_AS_SOURCE)
@@ -255,6 +256,7 @@ def construct_and_run_prompt(
         prompt_type=output.get(PSKeys.TYPE, PSKeys.TEXT),
         enable_highlight=enable_highlight,
         file_path=file_path,
+        execution_source=execution_source,
     )
 
 
@@ -312,9 +314,9 @@ def run_completion(
                         storage_type=StorageType.PERMANENT,
                         env_name=FileStorageKeys.PERMANENT_REMOTE_STORAGE,
                     )
-                if execution_source == ExecutionSource.TOOL.value:
+                elif execution_source == ExecutionSource.TOOL.value:
                     fs_instance = EnvHelper.get_storage(
-                        storage_type=StorageType.TEMPORARY,
+                        storage_type=StorageType.SHARED_TEMPORARY,
                         env_name=FileStorageKeys.TEMPORARY_REMOTE_STORAGE,
                     )
                 highlight_data = highlight_data_plugin["entrypoint_cls"](

diff --git a/prompt-service/src/unstract/prompt_service/main.py b/prompt-service/src/unstract/prompt_service/main.py
@@ -1,12 +1,13 @@
 import time
 import traceback
 from json import JSONDecodeError
-from typing import Any
+from typing import Any, Optional
 
 from flask import json, jsonify, request
 from llama_index.core.vector_stores import ExactMatchFilter, MetadataFilters
 from unstract.prompt_service.authentication_middleware import AuthenticationMiddleware
 from unstract.prompt_service.config import create_app, db
+from unstract.prompt_service.constants import ExecutionSource
 from unstract.prompt_service.constants import PromptServiceContants as PSKeys
 from unstract.prompt_service.constants import RunLevel
 from unstract.prompt_service.exceptions import APIError, ErrorResponse, NoPayloadError
@@ -330,6 +331,7 @@ def prompt_processor() -> Any:
                     prompt="promptx",
                     metadata=metadata,
                     file_path=file_path,
+                    execution_source=execution_source,
                 )
                 metadata[PSKeys.CONTEXT][output[PSKeys.NAME]] = get_cleaned_context(
                     context
@@ -360,6 +362,7 @@ def prompt_processor() -> Any:
                         vector_index=vector_index,
                         retrieval_type=retrieval_strategy,
                         metadata=metadata,
+                        execution_source=execution_source,
                     )
                     metadata[PSKeys.CONTEXT][output[PSKeys.NAME]] = get_cleaned_context(
                         context
@@ -807,6 +810,7 @@ def run_retrieval(  # type:ignore
     vector_index,
     retrieval_type: str,
     metadata: dict[str, Any],
+    execution_source: Optional[str] = ExecutionSource.IDE.value,
 ) -> tuple[str, set[str]]:
     context: set[str] = set()
     prompt = output[PSKeys.PROMPTX]
@@ -869,6 +873,7 @@ def run_retrieval(  # type:ignore
         context="\n".join(context),
         prompt="promptx",
         metadata=metadata,
+        execution_source=execution_source,
     )
 
     return (answer, context)