From 1d8bfcbd59fdfd7a904c2a84486488eb29f092b9 Mon Sep 17 00:00:00 2001
From: vishnuszipstack <117254672+vishnuszipstack@users.noreply.github.com>
Date: Tue, 10 Dec 2024 10:50:05 +0530
Subject: [PATCH] Fix/confidence data to prompt service (#877)

* added constant and helper update in promptservice

* removed unused epilogue
---
 .../src/unstract/prompt_service/constants.py      |  1 +
 .../src/unstract/prompt_service/helper.py         | 15 +++++++++++++--
 tools/structure/src/main.py                       | 15 ---------------
 3 files changed, 14 insertions(+), 17 deletions(-)

diff --git a/prompt-service/src/unstract/prompt_service/constants.py b/prompt-service/src/unstract/prompt_service/constants.py
index e905eec1e..a654ce5c2 100644
--- a/prompt-service/src/unstract/prompt_service/constants.py
+++ b/prompt-service/src/unstract/prompt_service/constants.py
@@ -71,6 +71,7 @@ class PromptServiceContants:
     ENABLE_HIGHLIGHT = "enable_highlight"
     FILE_PATH = "file_path"
     HIGHLIGHT_DATA = "highlight_data"
+    CONFIDENCE_DATA = "confidence_data"
 
 
 class RunLevel(Enum):
diff --git a/prompt-service/src/unstract/prompt_service/helper.py b/prompt-service/src/unstract/prompt_service/helper.py
index cf4ef63cb..62a24f2d1 100644
--- a/prompt-service/src/unstract/prompt_service/helper.py
+++ b/prompt-service/src/unstract/prompt_service/helper.py
@@ -296,8 +296,19 @@ def run_completion(
         )
         answer: str = completion[PSKeys.RESPONSE].text
         highlight_data = completion.get(PSKeys.HIGHLIGHT_DATA)
-        if all([metadata, highlight_data, prompt_key]):
-            metadata.setdefault(PSKeys.HIGHLIGHT_DATA, {})[prompt_key] = highlight_data
+        confidence_data = completion.get(PSKeys.CONFIDENCE_DATA)
+
+        if metadata is not None and prompt_key:
+            if highlight_data:
+                metadata.setdefault(PSKeys.HIGHLIGHT_DATA, {})[
+                    prompt_key
+                ] = highlight_data
+
+            if confidence_data:
+                metadata.setdefault(PSKeys.CONFIDENCE_DATA, {})[
+                    prompt_key
+                ] = confidence_data
+
         return answer
     # TODO: Catch and handle specific exception here
     except SdkRateLimitError as e:
diff --git a/tools/structure/src/main.py b/tools/structure/src/main.py
index d5fc55c09..c4adade11 100644
--- a/tools/structure/src/main.py
+++ b/tools/structure/src/main.py
@@ -239,21 +239,6 @@ def run(
 
         if not summarize_as_source:
             metadata = structured_output_dict[SettingsKeys.METADATA]
-            epilogue = metadata.pop(SettingsKeys.EPILOGUE, None)
-            if epilogue:
-                try:
-                    from helper import (  # type: ignore [attr-defined]
-                        get_confidence_data,
-                    )
-
-                    metadata[SettingsKeys.CONFIDENCE_DATA] = get_confidence_data(
-                        epilogue, tool_data_dir
-                    )
-                except ImportError:
-                    self.stream_log(
-                        f"Confidence data is not added. {PAID_FEATURE_MSG}",
-                        level=LogLevel.WARN,
-                    )
             # Update the dictionary with modified metadata
             structured_output_dict[SettingsKeys.METADATA] = metadata
             structured_output = json.dumps(structured_output_dict)