UN-3008 [FEAT] Add word-level confidence support (#1672)

* UN-3008 [FEAT] Add word-level confidence support

Add word-level confidence feature that extends the existing highlight functionality. This feature allows tracking confidence scores at the word level during extraction.

Key changes:
- Add enable_word_confidence field to CustomTool model
- Add word_confidence_postamble support for custom prompts
- Pass word_confidence flag through extraction and indexing pipelines
- Update SDK to preserve original text for post-processing
- Add dependency check to ensure word confidence requires highlight to be enabled

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* UN-3008 Confidence fixes

* UN-3008 Confidence fixes in backend

* UN-3008 Fix type annotation for post_process_fn to include original_text parameter

- Updated post_process_fn type signature from Callable[[LLMResponseCompat, bool], ...]
  to Callable[[LLMResponseCompat, bool, str], ...] to match the actual call at line 500-502
- Addresses review comment: #1672 (comment)
- The highlight_data plugin's run() method already accepts the third parameter (original_text: str)

* UN-3008 Minor fixes

* UN-3008 Minor fixes

* UN-3008 Remove unused WORD_CONFIDENCE_DATA constant from ResultKeys

- Removed WORD_CONFIDENCE_DATA constant as it's not used in main repo's workflow manager
- The constant is only needed in unstract-cloud repo which has the rule engine
- Prompt studio code uses the string directly, which is appropriate
- Addresses review comment: #1672 (comment)

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: Deepak-Kesavan

backend/prompt_studio/prompt_studio_core_v2/constants.py

@@ -93,11 +93,13 @@ class ToolStudioPromptKeys:
     TXT_EXTENTION = ".txt"
     TABLE = "table"
     PLATFORM_POSTAMBLE = "platform_postamble"
+    WORD_CONFIDENCE_POSTAMBLE = "word_confidence_postamble"
     SUMMARIZE_AS_SOURCE = "summarize_as_source"
     VARIABLE_MAP = "variable_map"
     RECORD = "record"
     FILE_PATH = "file_path"
     ENABLE_HIGHLIGHT = "enable_highlight"
+    ENABLE_WORD_CONFIDENCE = "enable_word_confidence"
     REQUIRED = "required"
     EXECUTION_SOURCE = "execution_source"
     LINE_ITEM = "line-item"
@@ -166,6 +168,7 @@ class IndexingConstants:
     FILE_HASH = "file_hash"
     OUTPUT_FILE_PATH = "output_file_path"
     ENABLE_HIGHLIGHT = "enable_highlight"
+    ENABLE_WORD_CONFIDENCE = "enable_word_confidence"
     USAGE_KWARGS = "usage_kwargs"
     PROCESS_TEXT = "process_text"
     EXTRACTED_TEXT = "extracted_text"
@@ -190,6 +193,7 @@ class DefaultValues:
     DEFAULT_EXCLUDE_FAILED = True
     DEFAULT_ENABLE_CHALLENGE = False
     DEFAULT_ENABLE_HIGHLIGHT = False
+    DEFAULT_ENABLE_WORD_CONFIDENCE = False
     DEFAULT_SUMMARIZE_AS_SOURCE = False
     DEFAULT_SUMMARIZE_CONTEXT = False
     DEFAULT_SINGLE_PASS_EXTRACTION_MODE = False

backend/prompt_studio/prompt_studio_core_v2/migrations/0005_customtool_enable_word_confidence.py

@@ -0,0 +1,20 @@
+# Generated by Django 4.2.1 on 2025-11-18 05:07
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("prompt_studio_core_v2", "0004_add_shared_to_org_to_custom_tool"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="customtool",
+            name="enable_word_confidence",
+            field=models.BooleanField(
+                db_comment="Flag to enable or disable word-level confidence (depends on enable_highlight)",
+                default=False,
+            ),
+        ),
+    ]

backend/prompt_studio/prompt_studio_core_v2/models.py

@@ -138,6 +138,10 @@ class CustomTool(DefaultOrganizationMixin, BaseModel):
     enable_highlight = models.BooleanField(
         db_comment="Flag to enable or disable document highlighting", default=False
     )
+    enable_word_confidence = models.BooleanField(
+        db_comment="Flag to enable or disable word-level confidence (depends on enable_highlight)",
+        default=False,
+    )
 
     # Introduced field to establish M2M relation between users and custom_tool.
     # This will introduce intermediary table which relates both the models.

backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py

@@ -412,6 +412,8 @@ def index_document(
             document_id=document_id,
             run_id=run_id,
             enable_highlight=tool.enable_highlight,
+            enable_word_confidence=tool.enable_word_confidence,
+            doc_id=doc_id,
         )
         if tool.summarize_context:
             summarize_file_path = PromptStudioHelper.summarize(
@@ -972,9 +974,13 @@ def _fetch_response(
         tool_settings[TSPKeys.POSTAMBLE] = tool.postamble
         tool_settings[TSPKeys.GRAMMAR] = grammar_list
         tool_settings[TSPKeys.ENABLE_HIGHLIGHT] = tool.enable_highlight
+        tool_settings[TSPKeys.ENABLE_WORD_CONFIDENCE] = tool.enable_word_confidence
         tool_settings[TSPKeys.PLATFORM_POSTAMBLE] = getattr(
             settings, TSPKeys.PLATFORM_POSTAMBLE.upper(), ""
         )
+        tool_settings[TSPKeys.WORD_CONFIDENCE_POSTAMBLE] = getattr(
+            settings, TSPKeys.WORD_CONFIDENCE_POSTAMBLE.upper(), ""
+        )
         file_hash = fs_instance.get_hash_from_file(path=doc_path)
 
         payload = {
@@ -1246,10 +1252,14 @@ def _fetch_single_pass_response(
         tool_settings[TSPKeys.CHUNK_OVERLAP] = default_profile.chunk_overlap
         tool_settings[TSPKeys.ENABLE_CHALLENGE] = tool.enable_challenge
         tool_settings[TSPKeys.ENABLE_HIGHLIGHT] = tool.enable_highlight
+        tool_settings[TSPKeys.ENABLE_WORD_CONFIDENCE] = tool.enable_word_confidence
         tool_settings[TSPKeys.CHALLENGE_LLM] = challenge_llm
         tool_settings[TSPKeys.PLATFORM_POSTAMBLE] = getattr(
             settings, TSPKeys.PLATFORM_POSTAMBLE.upper(), ""
         )
+        tool_settings[TSPKeys.WORD_CONFIDENCE_POSTAMBLE] = getattr(
+            settings, TSPKeys.WORD_CONFIDENCE_POSTAMBLE.upper(), ""
+        )
         tool_settings[TSPKeys.SUMMARIZE_AS_SOURCE] = tool.summarize_as_source
         for prompt in prompts:
             if not prompt.prompt:

backend/prompt_studio/prompt_studio_output_manager_v2/migrations/0004_promptstudiooutputmanager_word_confidence_data.py

@@ -0,0 +1,24 @@
+# Generated migration for word_confidence_data field
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        (
+            "prompt_studio_output_manager_v2",
+            "0003_promptstudiooutputmanager_confidence_data",
+        ),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="promptstudiooutputmanager",
+            name="word_confidence_data",
+            field=models.JSONField(
+                blank=True,
+                db_comment="Field to store word-level confidence data",
+                null=True,
+            ),
+        ),
+    ]

backend/prompt_studio/prompt_studio_output_manager_v2/models.py

@@ -37,6 +37,12 @@ class PromptStudioOutputManager(BaseModel):
         null=True,
         blank=True,
     )
+    word_confidence_data = models.JSONField(
+        db_comment="Field to store word-level confidence data",
+        editable=True,
+        null=True,
+        blank=True,
+    )
     eval_metrics = models.JSONField(
         db_column="eval_metrics",
         null=False,

backend/prompt_studio/prompt_studio_output_manager_v2/output_manager_helper.py

@@ -63,6 +63,7 @@ def update_or_create_prompt_output(
             challenge_data: dict[str, Any] | None,
             highlight_data: dict[str, Any] | None,
             confidence_data: dict[str, Any] | None,
+            word_confidence_data: dict[str, Any] | None,
         ) -> PromptStudioOutputManager:
             """Handles creating or updating a single prompt output and returns
             the instance.
@@ -81,6 +82,7 @@ def update_or_create_prompt_output(
                         "challenge_data": challenge_data,
                         "highlight_data": highlight_data,
                         "confidence_data": confidence_data,
+                        "word_confidence_data": word_confidence_data,
                     },
                 )
 
@@ -103,6 +105,7 @@ def update_or_create_prompt_output(
                     "challenge_data": challenge_data,
                     "highlight_data": highlight_data,
                     "confidence_data": confidence_data,
+                    "word_confidence_data": word_confidence_data,
                 }
                 PromptStudioOutputManager.objects.filter(
                     document_manager=document_manager,
@@ -126,6 +129,7 @@ def update_or_create_prompt_output(
         challenge_data = metadata.get("challenge_data")
         highlight_data = metadata.get("highlight_data")
         confidence_data = metadata.get("confidence_data")
+        word_confidence_data = metadata.get("word_confidence_data")
 
         if not prompts:
             return serialized_data
@@ -147,6 +151,8 @@ def update_or_create_prompt_output(
                     highlight_data = highlight_data.get(prompt.prompt_key)
                 if confidence_data:
                     confidence_data = confidence_data.get(prompt.prompt_key)
+                if word_confidence_data:
+                    word_confidence_data = word_confidence_data.get(prompt.prompt_key)
                 if challenge_data:
                     challenge_data = challenge_data.get(prompt.prompt_key)
 
@@ -171,6 +177,7 @@ def update_or_create_prompt_output(
                 challenge_data=challenge_data,
                 highlight_data=highlight_data,
                 confidence_data=confidence_data,
+                word_confidence_data=word_confidence_data,
             )
 
             # Serialize the instance

backend/prompt_studio/prompt_studio_registry_v2/constants.py

@@ -98,6 +98,7 @@ class JsonSchemaKey:
     SUMMARIZE_PROMPT = "summarize_prompt"
     SUMMARIZE_AS_SOURCE = "summarize_as_source"
     ENABLE_HIGHLIGHT = "enable_highlight"
+    ENABLE_WORD_CONFIDENCE = "enable_word_confidence"
     PLATFORM_POSTAMBLE = "platform_postamble"
     REQUIRED = "required"
     # Webhook postprocessing settings

backend/prompt_studio/prompt_studio_registry_v2/prompt_studio_registry_helper.py

@@ -78,6 +78,12 @@ def frame_spec(tool: CustomTool) -> Spec:
                 "default": False,
                 "description": "Enables highlight",
             },
+            "enable_word_confidence": {
+                "type": "boolean",
+                "title": "Enable word confidence",
+                "default": False,
+                "description": "Enables word-level confidence (depends on enable_highlight)",
+            },
         }
 
         spec = Spec(
@@ -283,9 +289,13 @@ def frame_export_json(
             tool.single_pass_extraction_mode
         )
         tool_settings[JsonSchemaKey.ENABLE_HIGHLIGHT] = tool.enable_highlight
+        tool_settings[JsonSchemaKey.ENABLE_WORD_CONFIDENCE] = tool.enable_word_confidence
         tool_settings[JsonSchemaKey.PLATFORM_POSTAMBLE] = getattr(
             settings, JsonSchemaKey.PLATFORM_POSTAMBLE.upper(), ""
         )
+        tool_settings[JsonSchemaKey.WORD_CONFIDENCE_POSTAMBLE] = getattr(
+            settings, JsonSchemaKey.WORD_CONFIDENCE_POSTAMBLE.upper(), ""
+        )
 
         for prompt in prompts:
             if prompt.prompt_type == JsonSchemaKey.NOTES or not prompt.active:

frontend/src/components/custom-tools/document-manager/DocumentManager.jsx

@@ -380,6 +380,23 @@ function DocumentManager({ generateIndex, handleUpdateTool, handleDocChange }) {
                     if (typeof confidence === "number") {
                       return confidence.toFixed(2);
                     }
+                    // Handle word confidence format: object with line numbers as keys
+                    if (
+                      confidence &&
+                      typeof confidence === "object" &&
+                      !Array.isArray(confidence)
+                    ) {
+                      const values = Object.values(confidence);
+                      if (
+                        values.length > 0 &&
+                        values.every((v) => typeof v === "number")
+                      ) {
+                        const avg =
+                          values.reduce((sum, val) => sum + val, 0) /
+                          values.length;
+                        return avg.toFixed(2);
+                      }
+                    }
                     // Handle old format: nested array structure
                     if (confidence?.[0]?.[0]?.confidence) {
                       return confidence[0][0].confidence;