Merge branch 'main' into feat/remote-storage-prompt-studio

Zipstack · Nov 13, 2024 · eb0e1aa · eb0e1aa
2 parents 7381cf2 + 6a64d2b
commit eb0e1aa
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 1 deletion.
diff --git a/tools/structure/src/constants.py b/tools/structure/src/constants.py
@@ -72,4 +72,5 @@ class SettingsKeys:
     METADATA = "metadata"
     EPILOGUE = "epilogue"
     HIGHLIGHT_DATA = "highlight_data"
+    CONFIDENCE_DATA = "confidence_data"
     EXECUTION_RUN_DATA_FOLDER = "EXECUTION_RUN_DATA_FOLDER"
diff --git a/tools/structure/src/main.py b/tools/structure/src/main.py
@@ -216,10 +216,16 @@ def run(
             epilogue = metadata.pop(SettingsKeys.EPILOGUE, None)
             if epilogue:
                 try:
-                    from helper import transform_dict  # type: ignore [attr-defined]
+                    from helper import (  # type: ignore [attr-defined]
+                        get_confidence_data,
+                        transform_dict,
+                    )
 
                     highlight_data = transform_dict(epilogue, tool_data_dir)
                     metadata[SettingsKeys.HIGHLIGHT_DATA] = highlight_data
+                    metadata[SettingsKeys.CONFIDENCE_DATA] = get_confidence_data(
+                        epilogue, tool_data_dir
+                    )
                 except ImportError:
                     self.stream_log(
                         f"Highlight metadata is not added. {PAID_FEATURE_MSG}",