Feat/sperate highlighting

backend/api_v2/api_deployment_views.py

@@ -87,11 +87,11 @@ def get(
         response: ExecutionResponse = DeploymentHelper.get_execution_status(
             execution_id
         )
-
         # Determine response status
         response_status = status.HTTP_422_UNPROCESSABLE_ENTITY
         if response.execution_status == CeleryTaskState.COMPLETED.value:
             response_status = status.HTTP_200_OK
+            response.remove_result_metadata_keys(["highlight_data"])
             if not include_metadata:
                 response.remove_result_metadata_keys()
             if not include_metrics:

backend/api_v2/deployment_helper.py

@@ -185,6 +185,7 @@ def execute_workflow(
             result.status_api = DeploymentHelper.construct_status_endpoint(
                 api_endpoint=api.api_endpoint, execution_id=execution_id
             )
+            result.remove_result_metadata_keys(["highlight_data"])
             if not include_metadata:
                 result.remove_result_metadata_keys()
             if not include_metrics:

backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py

@@ -385,6 +385,7 @@ def index_document(
             user_id=user_id,
             process_text=process_text,
             fs=fs_instance,
+            enable_highlight=tool.enable_highlight,
         )
 
         elapsed_time = time.time() - start_time
@@ -761,6 +762,7 @@ def _fetch_response(
             user_id=user_id,
             process_text=process_text,
             fs=fs_instance,
+            enable_highlight=tool.enable_highlight,
         )
         if index_result.get("status") == IndexingStatus.PENDING_STATUS.value:
             return {
@@ -915,6 +917,7 @@ def dynamic_indexer(
         run_id: str = None,
         process_text: Optional[Callable[[str], str]] = None,
         fs: FileStorage = FileStorage(provider=FileStorageProvider.LOCAL),
+        enable_highlight: bool = False,
     ) -> Any:
         """Used to index a file based on the passed arguments.
 
@@ -998,6 +1001,7 @@ def dynamic_indexer(
                 usage_kwargs=usage_kwargs.copy(),
                 process_text=process_text,
                 fs=fs,
+                enable_highlight=enable_highlight,
             )
 
             PromptStudioIndexHelper.handle_index_manager(
@@ -1075,6 +1079,7 @@ def _fetch_single_pass_response(
             user_id=user_id,
             process_text=process_text,
             fs=fs_instance,
+            enable_highlight=tool.enable_highlight,
         )
         if index_result.get("status") == IndexingStatus.PENDING_STATUS.value:
             return {

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

@@ -77,6 +77,7 @@ class PromptServiceContants:
     EXECUTION_SOURCE = "execution_source"
     METRICS = "metrics"
     LINE_ITEM = "line-item"
+    LINE_NUMBERS = "line_numbers"
 
 
 class RunLevel(Enum):

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

@@ -327,8 +327,10 @@ def run_completion(
         answer: str = completion[PSKeys.RESPONSE].text
         highlight_data = completion.get(PSKeys.HIGHLIGHT_DATA, [])
         confidence_data = completion.get(PSKeys.CONFIDENCE_DATA)
+        line_numbers = completion.get(PSKeys.LINE_NUMBERS, [])
         if metadata is not None and prompt_key:
             metadata.setdefault(PSKeys.HIGHLIGHT_DATA, {})[prompt_key] = highlight_data
+            metadata.setdefault(PSKeys.LINE_NUMBERS, {})[prompt_key] = line_numbers
 
             if confidence_data:
                 metadata.setdefault(PSKeys.CONFIDENCE_DATA, {})[