FIX: Fixed missing keys in highlight data (#861)

Fixed missing keys in highlight data Signed-off-by: Deepak K <[email protected]>
Zipstack · Jan 9, 2025 · b67768c · b67768c
1 parent 91e4c6c
commit b67768c
Showing 1 changed file with 2 additions and 5 deletions.
diff --git a/prompt-service/src/unstract/prompt_service/helper.py b/prompt-service/src/unstract/prompt_service/helper.py
@@ -327,13 +327,10 @@ def run_completion(
             extract_json=prompt_type.lower() != PSKeys.TEXT,
         )
         answer: str = completion[PSKeys.RESPONSE].text
-        highlight_data = completion.get(PSKeys.HIGHLIGHT_DATA)
+        highlight_data = completion.get(PSKeys.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
+            metadata.setdefault(PSKeys.HIGHLIGHT_DATA, {})[prompt_key] = highlight_data
 
             if confidence_data:
                 metadata.setdefault(PSKeys.CONFIDENCE_DATA, {})[