Improve semantic fact extraction prompt and add JSON enforcement #4282

dhrubo-os · 2025-10-08T05:24:57Z

Summary

This PR improves the semantic fact extraction system in the agentic memory feature by replacing the verbose XML-heavy prompt with a universal semantic fact extraction agent and adding JSON enforcement mechanisms for reliable LLM output.

Changes Made

1. Updated Semantic Facts Extraction Prompt

File: common/src/main/java/org/opensearch/ml/common/memorycontainer/MemoryContainerConstants.java
Replaced XML-heavy prompt with a cleaner, universal semantic fact extraction agent
New prompt handles multiple contexts: personal facts, technical investigations, and RCA
Improved clarity with SCOPE, STYLE & RULES, and OUTPUT sections
Better handling of both user facts and assistant conclusions/results

2. Added JSON Enforcement Mechanism

File: common/src/main/java/org/opensearch/ml/common/memorycontainer/MemoryContainerConstants.java
Added JSON_ENFORCEMENT_MESSAGE constant to guarantee proper LLM output format
Message: "Respond NOW with ONE LINE of valid JSON ONLY exactly as {"facts":["fact1","fact2",...]}. No extra text, no code fences, no newlines or tabs, and no spaces after commas or colons."

3. Updated Service Logic

File: plugin/src/main/java/org/opensearch/ml/action/memorycontainer/memory/MemoryProcessingService.java
Modified extractFactsFromConversation method to automatically append JSON enforcement message to all fact extraction requests
Ensures consistent JSON output across all LLM interactions

4. Comprehensive Testing

File: plugin/src/test/java/org/opensearch/ml/action/memorycontainer/memory/MemoryProcessingServiceTests.java
Added unit test testExtractFactsFromConversation_JsonEnforcementMessageAppended to verify JSON enforcement
Fixed compilation issues with proper imports and parameter access
All existing tests continue to pass

Benefits

Universal Prompt: Single prompt handles semantic facts, user preferences, and technical investigations
Improved LLM Compliance: JSON enforcement message significantly improves structured output reliability
Better Maintainability: Cleaner prompt structure without verbose XML formatting
Comprehensive Coverage: Handles both user facts and assistant conclusions/results for technical investigations

Testing

✅ All memory processing service tests pass
✅ JSON enforcement test specifically validates the new functionality
✅ Main compilation succeeds
✅ Code formatting applied with Spotless

Backward Compatibility

This change is backward compatible as it only modifies the internal prompt structure and adds enforcement messaging. The API and external behavior remain unchanged.

Related Issues

This addresses improvements to the agentic memory feature's fact extraction reliability and prompt maintainability.

codecov · 2025-10-08T06:31:03Z

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 80.15%. Comparing base (a3cb4fe) to head (9cb73cd).
⚠️ Report is 1 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff            @@
##               main    #4282   +/-   ##
=========================================
  Coverage     80.15%   80.15%           
  Complexity    10167    10167           
=========================================
  Files           854      854           
  Lines         44181    44183    +2     
  Branches       5112     5112           
=========================================
+ Hits          35413    35417    +4     
+ Misses         6614     6613    -1     
+ Partials       2154     2153    -1

Flag	Coverage Δ
ml-commons	`80.15% <100.00%> (+<0.01%)`	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

common/src/main/java/org/opensearch/ml/common/memorycontainer/MemoryContainerConstants.java

- Replace XML-heavy prompt with universal semantic fact extraction agent - Add JSON_ENFORCEMENT_MESSAGE constant for reliable LLM output format - Update MemoryProcessingService to append enforcement message to all fact extraction requests - Add comprehensive unit tests to verify JSON enforcement functionality - Support multiple contexts: personal facts, technical investigations, and RCA - Improve fact extraction reliability and maintainability Signed-off-by: Dhrubo Saha <dhrubo@amazon.com>

- Convert SEMANTIC_FACTS_EXTRACTION_PROMPT to use Java text blocks - Convert JSON_ENFORCEMENT_MESSAGE to use text blocks - Improves code readability and makes prompts easier to edit and visualize - Addresses PR feedback about using triple quotes

Signed-off-by: Dhrubo Saha <dhrubo@amazon.com>

pyek-bot · 2025-10-08T21:53:17Z

@dhrubo-os should USER_PREFERENCE_FACTS_EXTRACTION_PROMPT also be changed to the new structure? It's plain xml at the moment

rithin-pullela-aws · 2025-10-08T21:56:15Z

...n/src/main/java/org/opensearch/ml/action/memorycontainer/memory/MemoryProcessingService.java

+            MessageInput enforcementMessage = getMessageInput(JSON_ENFORCEMENT_MESSAGE);
+            enforcementMessage.toXContent(messagesBuilder, ToXContent.EMPTY_PARAMS);


Did you find adding the enforcement message in the User prompt instead of the system prompt works better?

Because in my experience they seem to adhere better in System prompt.
If we add too many messages in user prompt, the LLM might fail to adhere to the output format. The length of system prompt is in our control but the length of user prompt varies, I am not sure if this would adversely impact the output format

The general guidance from OpenAI and others is to put output format instructions in System Prompt.

If we do not see any significant improvements, I suggest adding it in the system prompt, to keep the code clean and readable.

+1

even in the user prompt when i ask LLM to not adhere to this format, it responds by saying it cannot and has been explicitly requested to follow a certain format. This must be because system prompt is given higher priority.

I noticed significant improvements with adding this JSON_ENFORCEMENT_MESSAGE. By using this I had 100% success to extract facts with our desired facts with 30+ iterations (different types of memories). So I'll insist to keep this.

Btw, I got this suggestion from LLMs like claude & openai :)

sure, that sounds good.

But recommend to try system_prompt as an iteration in the future, it should perform better! I had some issues with PER responses when it was in user_prompt, moving to system_prompt significantly helped. As you have tested, I don't want to block you! Thanks!

dhrubo-os · 2025-10-08T21:57:16Z

@dhrubo-os should USER_PREFERENCE_FACTS_EXTRACTION_PROMPT also be changed to the new structure? It's plain xml at the moment

I'm still investigating what else to improve in the user preference. After performing my analysis I'll raise another separate PR focusing that.

* Improve semantic fact extraction prompt and add JSON enforcement - Replace XML-heavy prompt with universal semantic fact extraction agent - Add JSON_ENFORCEMENT_MESSAGE constant for reliable LLM output format - Update MemoryProcessingService to append enforcement message to all fact extraction requests - Add comprehensive unit tests to verify JSON enforcement functionality - Support multiple contexts: personal facts, technical investigations, and RCA - Improve fact extraction reliability and maintainability Signed-off-by: Dhrubo Saha <dhrubo@amazon.com> * Use triple quotes (text blocks) for better prompt readability - Convert SEMANTIC_FACTS_EXTRACTION_PROMPT to use Java text blocks - Convert JSON_ENFORCEMENT_MESSAGE to use text blocks - Improves code readability and makes prompts easier to edit and visualize - Addresses PR feedback about using triple quotes * updated prompt based on comment Signed-off-by: Dhrubo Saha <dhrubo@amazon.com> --------- Signed-off-by: Dhrubo Saha <dhrubo@amazon.com> (cherry picked from commit cac8ddb)

…) (#4286) * Improve semantic fact extraction prompt and add JSON enforcement - Replace XML-heavy prompt with universal semantic fact extraction agent - Add JSON_ENFORCEMENT_MESSAGE constant for reliable LLM output format - Update MemoryProcessingService to append enforcement message to all fact extraction requests - Add comprehensive unit tests to verify JSON enforcement functionality - Support multiple contexts: personal facts, technical investigations, and RCA - Improve fact extraction reliability and maintainability * Use triple quotes (text blocks) for better prompt readability - Convert SEMANTIC_FACTS_EXTRACTION_PROMPT to use Java text blocks - Convert JSON_ENFORCEMENT_MESSAGE to use text blocks - Improves code readability and makes prompts easier to edit and visualize - Addresses PR feedback about using triple quotes * updated prompt based on comment --------- (cherry picked from commit cac8ddb) Signed-off-by: Dhrubo Saha <dhrubo@amazon.com> Co-authored-by: Dhrubo Saha <dhrubo@amazon.com>

…nsearch-project#4282) * Improve semantic fact extraction prompt and add JSON enforcement - Replace XML-heavy prompt with universal semantic fact extraction agent - Add JSON_ENFORCEMENT_MESSAGE constant for reliable LLM output format - Update MemoryProcessingService to append enforcement message to all fact extraction requests - Add comprehensive unit tests to verify JSON enforcement functionality - Support multiple contexts: personal facts, technical investigations, and RCA - Improve fact extraction reliability and maintainability Signed-off-by: Dhrubo Saha <dhrubo@amazon.com> * Use triple quotes (text blocks) for better prompt readability - Convert SEMANTIC_FACTS_EXTRACTION_PROMPT to use Java text blocks - Convert JSON_ENFORCEMENT_MESSAGE to use text blocks - Improves code readability and makes prompts easier to edit and visualize - Addresses PR feedback about using triple quotes * updated prompt based on comment Signed-off-by: Dhrubo Saha <dhrubo@amazon.com> --------- Signed-off-by: Dhrubo Saha <dhrubo@amazon.com> Signed-off-by: Anuj Soni <sonianuj287@gmail.com>

dhrubo-os requested review from HenryL27, Zhangxunmt, austintlee, b4sjoo, jngz-es, mingshl, model-collapse, pyek-bot, rbhavna, sam-herman, xinyual, ylwu-amzn and zane-neo as code owners October 8, 2025 05:24

dhrubo-os force-pushed the improve-semantic-fact-extraction-prompt branch from c8d4f1b to e79b701 Compare October 8, 2025 05:26

dhrubo-os temporarily deployed to ml-commons-cicd-env October 8, 2025 05:28 — with GitHub Actions Inactive

dhrubo-os had a problem deploying to ml-commons-cicd-env October 8, 2025 06:31 — with GitHub Actions Failure

dhrubo-os had a problem deploying to ml-commons-cicd-env October 8, 2025 06:31 — with GitHub Actions Error

pyek-bot reviewed Oct 8, 2025

View reviewed changes

common/src/main/java/org/opensearch/ml/common/memorycontainer/MemoryContainerConstants.java Show resolved Hide resolved

pyek-bot reviewed Oct 8, 2025

View reviewed changes

common/src/main/java/org/opensearch/ml/common/memorycontainer/MemoryContainerConstants.java Outdated Show resolved Hide resolved

b4sjoo approved these changes Oct 8, 2025

View reviewed changes

dhrubo-os temporarily deployed to ml-commons-cicd-env October 8, 2025 21:00 — with GitHub Actions Inactive

dhrubo-os added 2 commits October 8, 2025 14:02

updated prompt based on comment

9cb73cd

Signed-off-by: Dhrubo Saha <dhrubo@amazon.com>

b4sjoo force-pushed the improve-semantic-fact-extraction-prompt branch from 9d55e94 to 9cb73cd Compare October 8, 2025 21:02

b4sjoo temporarily deployed to ml-commons-cicd-env October 8, 2025 21:04 — with GitHub Actions Inactive

b4sjoo had a problem deploying to ml-commons-cicd-env October 8, 2025 21:04 — with GitHub Actions Failure

b4sjoo had a problem deploying to ml-commons-cicd-env October 8, 2025 21:04 — with GitHub Actions Error

rithin-pullela-aws reviewed Oct 8, 2025

View reviewed changes

pyek-bot approved these changes Oct 8, 2025

View reviewed changes

b4sjoo had a problem deploying to ml-commons-cicd-env October 8, 2025 22:07 — with GitHub Actions Error

b4sjoo had a problem deploying to ml-commons-cicd-env October 8, 2025 22:07 — with GitHub Actions Failure

dhrubo-os merged commit cac8ddb into opensearch-project:main Oct 8, 2025
9 of 13 checks passed

dhrubo-os added the backport 3.3 label Oct 8, 2025

opensearch-trigger-bot bot mentioned this pull request Oct 8, 2025

[Backport 3.3] Improve semantic fact extraction prompt and add JSON enforcement #4286

Merged

peterzhuamazon mentioned this pull request Oct 9, 2025

[RELEASE] Release version 3.3.0 opensearch-project/opensearch-build#5693

Closed

66 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Improve semantic fact extraction prompt and add JSON enforcement #4282

Improve semantic fact extraction prompt and add JSON enforcement #4282

Uh oh!

dhrubo-os commented Oct 8, 2025

Uh oh!

codecov bot commented Oct 8, 2025 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

pyek-bot commented Oct 8, 2025

Uh oh!

rithin-pullela-aws Oct 8, 2025

Uh oh!

pyek-bot Oct 8, 2025

Uh oh!

dhrubo-os Oct 8, 2025

Uh oh!

dhrubo-os Oct 8, 2025

Uh oh!

pyek-bot Oct 8, 2025

Uh oh!

dhrubo-os commented Oct 8, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

		MessageInput enforcementMessage = getMessageInput(JSON_ENFORCEMENT_MESSAGE);
		enforcementMessage.toXContent(messagesBuilder, ToXContent.EMPTY_PARAMS);

Improve semantic fact extraction prompt and add JSON enforcement #4282

Improve semantic fact extraction prompt and add JSON enforcement #4282

Uh oh!

Conversation

dhrubo-os commented Oct 8, 2025

Summary

Changes Made

1. Updated Semantic Facts Extraction Prompt

2. Added JSON Enforcement Mechanism

3. Updated Service Logic

4. Comprehensive Testing

Benefits

Testing

Backward Compatibility

Related Issues

Uh oh!

codecov bot commented Oct 8, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Uh oh!

Uh oh!

pyek-bot commented Oct 8, 2025

Uh oh!

rithin-pullela-aws Oct 8, 2025

Choose a reason for hiding this comment

Uh oh!

pyek-bot Oct 8, 2025

Choose a reason for hiding this comment

Uh oh!

dhrubo-os Oct 8, 2025

Choose a reason for hiding this comment

Uh oh!

dhrubo-os Oct 8, 2025

Choose a reason for hiding this comment

Uh oh!

pyek-bot Oct 8, 2025

Choose a reason for hiding this comment

Uh oh!

dhrubo-os commented Oct 8, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

codecov bot commented Oct 8, 2025 •

edited

Loading