Merge branch 'main' into UN-1582-ABS-CBN-Add-Oracle-connector

Zipstack · Dec 10, 2024 · fcea33e · fcea33e
2 parents 3116b30 + 1d8bfcb
commit fcea33e
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 17 deletions.
diff --git 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
@@ -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
@@ -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)