UN-3008 [FEAT] Add word-level confidence support - #1672
Conversation
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
Summary by CodeRabbitRelease Notes
✏️ Tip: You can customize this high-level summary in your review settings. WalkthroughAdds word-level confidence support: new constants, model fields + migrations, threading enable_word_confidence and word_confidence_data through indexing/extraction/prompt-service/SDK, persisting word confidence, and propagating wordConfidenceData to frontend renderers. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
participant ToolCfg as Tool Config
participant Backend as PromptStudio Backend
participant Registry as Registry Export
participant Service as Prompt Service
participant LLM as LLM / SDK
participant Frontend as Frontend Renderer
ToolCfg->>Backend: set enable_word_confidence
Backend->>Registry: include enable_word_confidence in export
Registry->>Service: tool_settings.enable_word_confidence
Service->>Service: compute word_confidence_postamble if enabled
alt enabled & enable_highlight true
Service->>LLM: run prompt with postamble
else
Service->>LLM: run prompt without postamble
end
LLM->>LLM: post-process (original_text preserved)
alt returns word confidence
LLM-->>Service: metadata includes WORD_CONFIDENCE_DATA
end
Service-->>Backend: persist response + word_confidence_data
Backend-->>Frontend: include wordConfidenceData in output payload
Frontend->>Frontend: render using wordConfidenceData
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes
Pre-merge checks and finishing touches❌ Failed checks (1 warning, 1 inconclusive)
✅ Passed checks (1 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
for more information, see https://pre-commit.ci
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
unstract/sdk1/src/unstract/sdk1/llm.py (2)
477-477: Update the type hint to match the new signature.The type hint for
post_process_fnstill reflects the old 2-parameter signature, but the actual invocation at lines 500-502 passes 3 parameters includingoriginal_text. This type mismatch will cause type-checking errors.Apply this diff to fix the type hint:
def _post_process_response( self, response_text: str, extract_json: bool, - post_process_fn: Callable[[LLMResponseCompat, bool], dict[str, object]] | None, + post_process_fn: Callable[[LLMResponseCompat, bool, str], dict[str, object]] | None, ) -> tuple[str, dict[str, object]]:
228-234: Update the type casting to reflect the new signature.The type cast at lines 229-234 still uses the old 2-parameter signature for
post_process_fn. This should be updated to match the new 3-parameter signature.Apply this diff:
extract_json: bool = cast("bool", kwargs.get("extract_json", False)) post_process_fn: ( - Callable[[LLMResponseCompat, bool], dict[str, object]] | None + Callable[[LLMResponseCompat, bool, str], dict[str, object]] | None ) = cast( - "Callable[[LLMResponseCompat, bool], dict[str, object]] | None", + "Callable[[LLMResponseCompat, bool, str], dict[str, object]] | None", kwargs.get("process_text", None), )
🧹 Nitpick comments (2)
prompt-service/src/unstract/prompt_service/services/answer_prompt.py (1)
387-387: Consider removing unused parameter.Static analysis correctly identifies that
enable_word_confidenceis accepted but never used in thehandle_jsonfunction body. The word confidence data is already handled inrun_completion(lines 261-264).Unless this parameter is reserved for future functionality, consider removing it to keep the API surface clean.
If the parameter is not needed, apply this diff:
def handle_json( answer: str, structured_output: dict[str, Any], output: dict[str, Any], log_events_id: str, tool_id: str, doc_name: str, llm: LLM, enable_highlight: bool = False, - enable_word_confidence: bool = False, execution_source: str = ExecutionSource.IDE.value, metadata: dict[str, Any] | None = None, file_path: str = "", ) -> None:And update the caller in
prompt-service/src/unstract/prompt_service/controllers/answer_prompt.pylines 505-507:AnswerPromptService.handle_json( answer=answer, structured_output=structured_output, output=output, log_events_id=log_events_id, tool_id=tool_id, doc_name=doc_name, llm=llm, enable_highlight=tool_settings.get(PSKeys.ENABLE_HIGHLIGHT, False), - enable_word_confidence=tool_settings.get( - PSKeys.ENABLE_WORD_CONFIDENCE, False - ), execution_source=execution_source, metadata=metadata, file_path=file_path, )backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py (1)
1313-1321: Add model-level validation for the enable_word_confidence dependency.Verification confirms the concern is valid. The
CustomToolmodel defines bothenable_highlightandenable_word_confidencefields (lines 138-143 in models.py) without any constraint validation. While the prompt service enforces the dependency downstream at answer_prompt.py lines 124-126 (silently resettingenable_word_confidencetoFalsewhen highlighting is disabled), there is no model-level or serializer-level validation to prevent the invalid state from being persisted to the database.Recommended approach: Add a
clean()method to theCustomToolmodel or avalidate_enable_word_confidencemethod to theCustomToolSerializerto enforce: ifenable_word_confidence=True, thenenable_highlightmust also beTrue.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
Cache: Disabled due to Reviews > Disable Cache setting
Knowledge base: Disabled due to Reviews -> Disable Knowledge Base setting
📒 Files selected for processing (13)
backend/prompt_studio/prompt_studio_core_v2/constants.py(3 hunks)backend/prompt_studio/prompt_studio_core_v2/migrations/0005_customtool_enable_word_confidence.py(1 hunks)backend/prompt_studio/prompt_studio_core_v2/models.py(1 hunks)backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py(7 hunks)backend/prompt_studio/prompt_studio_registry_v2/constants.py(1 hunks)backend/prompt_studio/prompt_studio_registry_v2/prompt_studio_registry_helper.py(2 hunks)prompt-service/src/unstract/prompt_service/constants.py(2 hunks)prompt-service/src/unstract/prompt_service/controllers/answer_prompt.py(1 hunks)prompt-service/src/unstract/prompt_service/controllers/extraction.py(2 hunks)prompt-service/src/unstract/prompt_service/controllers/indexing.py(2 hunks)prompt-service/src/unstract/prompt_service/dto.py(1 hunks)prompt-service/src/unstract/prompt_service/services/answer_prompt.py(11 hunks)unstract/sdk1/src/unstract/sdk1/llm.py(2 hunks)
🧰 Additional context used
🪛 Ruff (0.14.5)
prompt-service/src/unstract/prompt_service/services/answer_prompt.py
387-387: Unused static method argument: enable_word_confidence
(ARG004)
backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py
1240-1240: Undefined name doc_id
(F821)
backend/prompt_studio/prompt_studio_core_v2/migrations/0005_customtool_enable_word_confidence.py
7-9: Mutable class attributes should be annotated with typing.ClassVar
(RUF012)
11-20: Mutable class attributes should be annotated with typing.ClassVar
(RUF012)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: build
🔇 Additional comments (17)
backend/prompt_studio/prompt_studio_registry_v2/constants.py (1)
101-101: LGTM!The constant addition follows the existing pattern and is properly placed adjacent to the related
ENABLE_HIGHLIGHTconstant.prompt-service/src/unstract/prompt_service/dto.py (1)
38-38: LGTM!The field addition is clean and follows the existing pattern. The default value of
Falseis appropriate, and the placement afterenable_highlightaligns with the feature dependency.prompt-service/src/unstract/prompt_service/controllers/indexing.py (2)
63-63: LGTM with validation caveat.The extraction and forwarding of
enable_word_confidencefollows the established pattern. However, similar to the extraction controller, ensure that dependency validation (requiringenable_highlight) is properly enforced and communicated back through this API endpoint.
83-88: LGTM!The multiline formatting of
ProcessingOptionsinitialization improves readability, and all required fields including the newenable_word_confidenceare properly passed.backend/prompt_studio/prompt_studio_core_v2/migrations/0005_customtool_enable_word_confidence.py (1)
1-20: LGTM! Migration structure is correct.The migration properly adds the
enable_word_confidencefield with appropriate defaults and documentation. Thedb_commentdocuments the dependency onenable_highlight, though enforcement should be at the model or serializer level.Note: The static analysis hints about
ClassVarare false positives—Django migration class attributes are intentionally not annotated withClassVar.backend/prompt_studio/prompt_studio_registry_v2/prompt_studio_registry_helper.py (2)
81-86: LGTM!The property definition follows the established pattern for
enable_highlightand properly documents the dependency in the description. The spec structure is consistent with other boolean flags.
292-292: LGTM!The export of
enable_word_confidenceto tool settings is consistent with theenable_highlightpattern on line 291, ensuring the feature flag is properly propagated through the registry.backend/prompt_studio/prompt_studio_core_v2/constants.py (1)
96-96: LGTM! Constants follow existing patterns.The new word confidence constants are properly structured and consistent with existing naming conventions. The default value of
Falseis appropriate for a new opt-in feature.Also applies to: 102-102, 171-171, 196-196
prompt-service/src/unstract/prompt_service/controllers/answer_prompt.py (1)
505-507: LGTM! Consistent parameter passing.The
enable_word_confidenceparameter is retrieved fromtool_settingswith an appropriate default and passed through tohandle_json, following the same pattern asenable_highlight.prompt-service/src/unstract/prompt_service/constants.py (1)
66-66: LGTM! Well-structured constant additions.The new constants for word confidence functionality are properly placed and follow the existing naming conventions across
PromptServiceConstantsandIndexingConstants.Also applies to: 74-74, 78-78, 194-194
backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py (4)
415-416: LGTM! Properly passing word confidence parameters.The
enable_word_confidenceanddoc_idparameters are correctly passed todynamic_extractor. Thedoc_idis properly calculated at lines 396-406 before this call.
875-876: LGTM! Parameters correctly threaded through.The word confidence parameters are appropriately passed to
dynamic_extractor, withdoc_idproperly calculated at lines 849-859.
979-979: LGTM! Tool settings properly configured.The
enable_word_confidenceflag andword_confidence_postambleare correctly added totool_settings, following the same pattern as other configuration settings.Also applies to: 983-985
1259-1259: LGTM! Single-pass settings configured correctly.The word confidence settings are properly added to
tool_settingsfor single-pass extraction mode, consistent with the pattern used in regular extraction.Also applies to: 1264-1266
prompt-service/src/unstract/prompt_service/services/answer_prompt.py (3)
118-126: Good dependency enforcement.The implementation correctly enforces that
enable_word_confidencerequiresenable_highlightto be enabled (lines 124-126). This prevents an invalid configuration state at the service layer.
191-192: LGTM! Postamble construction is correct.The word confidence postamble is properly appended to the platform postamble when enabled, maintaining the correct order of prompt components.
238-238: LGTM! Word confidence data flow is properly implemented.The
enable_word_confidenceflag is correctly:
- Passed to the highlight data plugin (line 238)
- Retrieved from the completion response (line 248)
- Conditionally stored in metadata when enabled (lines 261-264)
Also applies to: 248-248, 261-264
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
frontend/src/components/custom-tools/document-manager/DocumentManager.jsx (1)
383-399: Consider refactoring confidence calculation to a helper function.The new word confidence handling logic is correct and properly validates the object format. However, the confidence display logic now handles four different formats (number, object, nested array, empty array), which increases cognitive complexity within the JSX.
For improved maintainability, consider extracting this logic into a dedicated helper function:
const calculateConfidenceScore = (confidence) => { // Handle numeric confidence 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" && !isNaN(v) && isFinite(v))) { const avg = values.reduce((sum, val) => sum + val, 0) / values.length; return avg.toFixed(2); } } // Handle old nested array format if (confidence?.[0]?.[0]?.confidence) { return confidence[0][0].confidence; } // Handle old empty array format if (Array.isArray(confidence?.[0]) && confidence[0].length === 0) { return "1"; } return "NA"; };Then use it in the JSX:
<Tag color="rgb(45, 183, 245)"> Confidence Score: {calculateConfidenceScore(selectedHighlight?.confidence)} </Tag>Optional enhancement: The helper function above also adds validation for
NaNandInfinityvalues when computing the average, providing additional robustness.</review_comment_end>
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
Cache: Disabled due to Reviews > Disable Cache setting
Knowledge base: Disabled due to Reviews -> Disable Knowledge Base setting
📒 Files selected for processing (4)
frontend/src/components/custom-tools/document-manager/DocumentManager.jsx(1 hunks)frontend/src/components/custom-tools/prompt-card/DisplayPromptResult.jsx(10 hunks)frontend/src/components/custom-tools/prompt-card/PromptOutput.jsx(4 hunks)frontend/src/hooks/usePromptOutput.js(1 hunks)
🧰 Additional context used
🧬 Code graph analysis (2)
frontend/src/components/custom-tools/prompt-card/PromptOutput.jsx (1)
frontend/src/components/custom-tools/list-of-tools/ListOfTools.jsx (1)
promptDetails(34-34)
frontend/src/components/custom-tools/document-manager/DocumentManager.jsx (2)
frontend/src/components/custom-tools/prompt-card/DisplayPromptResult.jsx (1)
confidence(268-268)backend/workflow_manager/endpoint_v2/constants.py (1)
values(88-89)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: build
🔇 Additional comments (3)
frontend/src/hooks/usePromptOutput.js (1)
124-124: LGTM!The addition of
wordConfidenceDatafollows the same pattern as existing fields and properly uses optional chaining to access the API response data.frontend/src/components/custom-tools/prompt-card/PromptOutput.jsx (1)
22-22: Good catch! The missing useState import is now fixed.The
useStatehook is used at line 74, so this import was required. This change fixes a bug where the component would fail due to the missing import.frontend/src/components/custom-tools/prompt-card/DisplayPromptResult.jsx (1)
25-25: LGTM! Excellent implementation of word confidence support.The changes properly thread
wordConfidenceDatathrough all rendering paths with appropriate fallback logic:
- Fallback pattern consistently prefers word-level confidence over document-level confidence
- Data is correctly propagated through recursive
renderJsoncalls for objects and arraysTextResultcomponent properly handles the fallback logic- PropTypes are updated to include the new parameter
- Backward compatibility is maintained via optional chaining
The implementation ensures that when word-level confidence data is available, it takes precedence, while gracefully falling back to the existing confidence data when it's not.
Also applies to: 83-101, 110-225, 230-254, 262-285, 293-293, 307-307
…ext 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)
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (1)
unstract/sdk1/src/unstract/sdk1/llm.py (1)
229-233: Unifypost_process_fntype annotation with 3‑parameter call signatureYou updated the
post_process_fntype incomplete()to accept(LLMResponseCompat, bool, str), and_post_process_responsenow callspost_process_fn(response_compat, extract_json, original_text), but the_post_process_responseparameter annotation still declares a 2‑arg callable. This will trip static type checkers and is inconsistent with the actual usage.Recommend updating
_post_process_response’s signature to match:- def _post_process_response( - self, - response_text: str, - extract_json: bool, - post_process_fn: Callable[[LLMResponseCompat, bool], dict[str, object]] | None, - ) -> tuple[str, dict[str, object]]: + def _post_process_response( + self, + response_text: str, + extract_json: bool, + post_process_fn: Callable[ + [LLMResponseCompat, bool, str], dict[str, object] + ] | None, + ) -> tuple[str, dict[str, object]]:Also re‑confirm that all callbacks passed via
process_text(e.g., highlight/word‑confidence plugins) have been updated to accept the thirdoriginal_textparameter, otherwise they’ll raise aTypeErrorat runtime. This is the same concern raised in the earlier review, now that the callsite is wired up.Also applies to: 473-502
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
Cache: Disabled due to Reviews > Disable Cache setting
Knowledge base: Disabled due to Reviews -> Disable Knowledge Base setting
📒 Files selected for processing (3)
backend/prompt_studio/prompt_studio_registry_v2/prompt_studio_registry_helper.py(2 hunks)backend/workflow_manager/workflow_v2/constants.py(1 hunks)unstract/sdk1/src/unstract/sdk1/llm.py(3 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- backend/prompt_studio/prompt_studio_registry_v2/prompt_studio_registry_helper.py
🧰 Additional context used
🧬 Code graph analysis (1)
unstract/sdk1/src/unstract/sdk1/llm.py (1)
unstract/sdk1/src/unstract/sdk1/utils/common.py (1)
LLMResponseCompat(120-138)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: build
🔇 Additional comments (1)
unstract/sdk1/src/unstract/sdk1/llm.py (1)
481-483: Capturingoriginal_textbefore JSON extraction looks correctStoring
original_text = response_textbefore theextract_jsonmanipulation and passing bothresponse_compat(possibly trimmed) andoriginal_textintopost_process_fnis a clean way to support word‑level confidence / highlighting without breaking existing JSON extraction behavior.Also applies to: 499-505
- 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)
Test ResultsSummary
Runner Tests - Full Report
SDK1 Tests - Full Report
|
|
* 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) * UN-3008 Sync with main - remove unwanted params * UN-3008 Word confidence fixes --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
3 participants
Summary
Implements word-level confidence scoring feature that extends the existing document highlighting functionality. This allows users to track confidence scores at the word level during extraction, providing more granular feedback about extraction quality.
Key Features:
enable_word_confidencetoggle in CustomTool modelenable_highlightto be enabled)Technical Changes:
enable_word_confidencefieldTest Plan