Enhance RagTool to choose neural sparse query type #140

mingshl · 2024-01-22T21:45:59Z

Description

semantic search is based on text embedding models, and neural sparse search is based on sparse embedding models. We are now enhancing RagTool with a parameter query_type, to enable both options.

now that users can set parameter query_type=neural to run neural query ,
or set parameter query_type=neural_sparse to run neural_sparse query

There are two phases of RAGTool, first is through retrieval phase, to query, and second phase is content generation through inference mode.

The newly added parameter enable_Content_Generation, the default is true, When it sets to false, it will return search retrieval result directly. When it sets to true, the RAGTool will continue trigger inference model to generate content.

Also fixing the build failure related to JDK 21

Check List

New functionality includes testing.
- All tests pass
New functionality has been documented.
- New functionality has javadoc added
Commits are signed per the DCO using --signoff

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
For more information on following Developer Certificate of Origin and signing off your commits, please check here.

Signed-off-by: Mingshi Liu <[email protected]>

mingshl · 2024-01-22T21:50:55Z

known failure on CI GitHub Actions to run on JDK 21.0.1 due to opensearch-project/flow-framework#426

src/main/java/org/opensearch/agent/tools/RAGTool.java

zhichao-aws

Thanks for your work! But I think it still need some changes.
Besides the comments for this PR, Since RAGTool doesn't call getQueryBody, and doesn't use the logic in AbstractRetrieverTool, I think we don't need to extend the AbstractRetrieverTool. Now the retrieval part has been done in RAGTool's component, i.e. VectorDBTool and NeuralSparseTool. It doesn't make sense to extend the AbstractRetrieverTool for RAGTool itself.

mingshl · 2024-01-23T22:45:11Z

Thanks for your work! But I think it still need some changes. Besides the comments for this PR, Since RAGTool doesn't call getQueryBody, and doesn't use the logic in AbstractRetrieverTool, I think we don't need to extend the AbstractRetrieverTool. Now the retrieval part has been done in RAGTool's component, i.e. VectorDBTool and NeuralSparseTool. It doesn't make sense to extend the AbstractRetrieverTool for RAGTool itself.

I am more leaning towards keeping RAGTool extending from AbstractRetrieverTool, because it uses similar logic of creating a tool and run the parameters, it's also retrieving data through the sub tools(VectorDBTool and NeuralSparseTool). From the purpose of RAGTool, it retrieves data from subtool and use model for inference, so it should count towards a RetrieverTool as well.

Signed-off-by: Mingshi Liu <[email protected]>

codecov · 2024-01-23T22:51:57Z

Codecov Report

Attention: 5 lines in your changes are missing coverage. Please review.

Comparison is base (e278d8d) 80.83% compared to head (44b8e23) 81.59%.
Report is 4 commits behind head on main.

Files	Patch %	Lines
.../main/java/org/opensearch/agent/tools/RAGTool.java	88.09%	2 Missing and 3 partials ⚠️

Additional details and impacted files

@@             Coverage Diff              @@
##               main     #140      +/-   ##
============================================
+ Coverage     80.83%   81.59%   +0.75%     
- Complexity      195      199       +4     
============================================
  Files            13       13              
  Lines          1002     1027      +25     
  Branches        133      133              
============================================
+ Hits            810      838      +28     
+ Misses          141      139       -2     
+ Partials         51       50       -1

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

mingshl · 2024-01-23T22:56:51Z

fixing the JDK issue in the commit Work around JDK 21.0.2 bug impacting scaling executors

src/main/java/org/opensearch/agent/tools/RAGTool.java

zhichao-aws · 2024-01-24T02:29:32Z

Thanks for your work! But I think it still need some changes. Besides the comments for this PR, Since RAGTool doesn't call getQueryBody, and doesn't use the logic in AbstractRetrieverTool, I think we don't need to extend the AbstractRetrieverTool. Now the retrieval part has been done in RAGTool's component, i.e. VectorDBTool and NeuralSparseTool. It doesn't make sense to extend the AbstractRetrieverTool for RAGTool itself.

I am more leaning towards keeping RAGTool extending from AbstractRetrieverTool, because it uses similar logic of creating a tool and run the parameters, it's also retrieving data through the sub tools(VectorDBTool and NeuralSparseTool). From the purpose of RAGTool, it retrieves data from subtool and use model for inference, so it should count towards a RetrieverTool as well.

The core logic of AbstractRetrieverTool is getQueryBody and run, but we don't use getQueryBody, and override run. It doesn't make sense to extend an abstract class but doesn't use it's logic at all. If we just want to mark the retrieval purpose, maybe we can make them implements a new interface. Extending the AbstractRetrieverTool will confuse other contributors.

There is also another implementation if we do want to extend AbstractRetrieverTool. In getQueryBody, we generate neural or neural_sparse query, and in run, we call super.run() to do search.

zhichao-aws · 2024-01-24T02:33:24Z

I also notice we still don't have integ test for RAGTool now. We can use MockHttpServer to mock an LLM endpoint to do integ test for RAGTool. You can check PPLToolIT for reference (

skills/src/test/java/org/opensearch/integTest/PPLToolIT.java

Line 28 in 8074492

List<PromptHandler> promptHandlers() {

).

I don't think this is a blocker for this PR, but we need to add integ test in follow up PRs.

Signed-off-by: Mingshi Liu <[email protected]>

zhichao-aws · 2024-01-25T12:40:21Z

src/main/java/org/opensearch/agent/tools/RAGTool.java

                        .queryTool(neuralSparseSearchTool)
                        .build();
                case "neural":
                    VectorDBTool.Factory.getInstance().init(client, xContentRegistry);
-                    params.put(VectorDBTool.MODEL_ID_FIELD, EMBEDDING_MODEL_ID_FIELD);
+                    params.put(VectorDBTool.MODEL_ID_FIELD, embeddingModelId);


We should also do this for NeuralSparseTool

We don't need to call VectorDBTool.Factory.getInstance().init(client, xContentRegistry) or NeuralSparseSearchTool.Factory.getInstance().init(client, xContentRegistry) here. This has been done during the initialization of skills plugin

added the initiation to the test to simulate the initiations of tools in ToolPlugins.java

mapped the inference_model_id to model_id for NeuralSparseTool

zhichao-aws · 2024-01-25T12:56:32Z

src/main/java/org/opensearch/agent/tools/RAGTool.java

+    public String getVersion() {
+        return null;
+    }
+


We can also remove the extends logic in below lines. public static class Factory extends AbstractRetrieverTool.Factory<RAGTool>

used Factory implements Tool.Factory

Signed-off-by: Mingshi Liu <[email protected]>

zhichao-aws

We can get PR this merged first, but may need to alter the prompt handling logic of RAGTool in the future based on e2e test result.

Need to fix the flaky test and add integration test for RAGTool in follow up PR.

opensearch-trigger-bot · 2024-01-26T03:24:34Z

The backport to 2.x failed:

The process '/usr/bin/git' failed with exit code 128

To backport manually, run these commands in your terminal:

# Navigate to the root of your repository
cd $(git rev-parse --show-toplevel)
# Fetch latest updates from GitHub
git fetch
# Create a new working tree
git worktree add ../.worktrees/skills/backport-2.x 2.x
# Navigate to the new working tree
pushd ../.worktrees/skills/backport-2.x
# Create a new branch
git switch --create backport/backport-140-to-2.x
# Cherry-pick the merged commit of this pull request and resolve the conflicts
git cherry-pick -x --mainline 1 d7945d55dc28f95126891de8d8ca4ac2ea26e280
# Push it to GitHub
git push --set-upstream origin backport/backport-140-to-2.x
# Go back to the original working tree
popd
# Delete the working tree
git worktree remove ../.worktrees/skills/backport-2.x

Then, create a pull request where the base branch is 2.x and the compare/head branch is backport/backport-140-to-2.x.

* enhance RagTool to choose neural sparse query type Signed-off-by: Mingshi Liu <[email protected]> * Work around JDK 21.0.2 bug impacting scaling executors Signed-off-by: Mingshi Liu <[email protected]> * Modify RAGTool factory create to initiate subtools Signed-off-by: Mingshi Liu <[email protected]> * Add enableContentGeneration to RAGTool Signed-off-by: Mingshi Liu <[email protected]> * Map embedding_model_id to model_id for NeuralSparseTool Signed-off-by: Mingshi Liu <[email protected]> * Map embedding_model_id to model_id for NeuralSparseTool Signed-off-by: Mingshi Liu <[email protected]> --------- Signed-off-by: Mingshi Liu <[email protected]> (cherry picked from commit d7945d5) Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

* enhance RagTool to choose neural sparse query type * Work around JDK 21.0.2 bug impacting scaling executors * Modify RAGTool factory create to initiate subtools * Add enableContentGeneration to RAGTool * Map embedding_model_id to model_id for NeuralSparseTool * Map embedding_model_id to model_id for NeuralSparseTool --------- (cherry picked from commit d7945d5) Signed-off-by: Mingshi Liu <[email protected]> Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

…t#140) (opensearch-project#152) * enhance RagTool to choose neural sparse query type * Work around JDK 21.0.2 bug impacting scaling executors * Modify RAGTool factory create to initiate subtools * Add enableContentGeneration to RAGTool * Map embedding_model_id to model_id for NeuralSparseTool * Map embedding_model_id to model_id for NeuralSparseTool --------- (cherry picked from commit d7945d5) Signed-off-by: Mingshi Liu <[email protected]> Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> Signed-off-by: yuye-aws <[email protected]>

mingshl added 2 commits January 22, 2024 13:25

enhance RagTool to choose neural sparse query type

af2ef84

Signed-off-by: Mingshi Liu <[email protected]>

Work around JDK 21.0.2 bug impacting scaling executors

b316179

Signed-off-by: Mingshi Liu <[email protected]>

mingshl requested review from b4sjoo, dhrubo-os, jngz-es, model-collapse, rbhavna, ylwu-amzn, zane-neo, Zhangxunmt, xinyual, amitgalitz, jackiehanyang, owaiskazi19, ohltyler, joshpalis, dbwiddis, kaituo, lezzago, eirsep, sbcd90, zhichao-aws and joshuali925 as code owners January 22, 2024 21:46

opensearch-trigger-bot bot added the backport 2.x label Jan 22, 2024

zhichao-aws reviewed Jan 23, 2024

View reviewed changes

src/main/java/org/opensearch/agent/tools/RAGTool.java Outdated Show resolved Hide resolved

zhichao-aws reviewed Jan 23, 2024

View reviewed changes

src/main/java/org/opensearch/agent/tools/RAGTool.java Outdated Show resolved Hide resolved

zhichao-aws reviewed Jan 23, 2024

View reviewed changes

src/main/java/org/opensearch/agent/tools/RAGTool.java Show resolved Hide resolved

zhichao-aws reviewed Jan 23, 2024

View reviewed changes

Modify RAGTool factory create to initiate subtools

eaf309a

Signed-off-by: Mingshi Liu <[email protected]>

zhichao-aws reviewed Jan 24, 2024

View reviewed changes

src/main/java/org/opensearch/agent/tools/RAGTool.java Outdated Show resolved Hide resolved

Add enableContentGeneration to RAGTool

b12fd8a

Signed-off-by: Mingshi Liu <[email protected]>

zhichao-aws reviewed Jan 25, 2024

View reviewed changes

mingshl added 2 commits January 25, 2024 09:49

Map embedding_model_id to model_id for NeuralSparseTool

431b9f9

Signed-off-by: Mingshi Liu <[email protected]>

Map embedding_model_id to model_id for NeuralSparseTool

44b8e23

Signed-off-by: Mingshi Liu <[email protected]>

mingshl requested a review from zhichao-aws January 26, 2024 01:07

zhichao-aws approved these changes Jan 26, 2024

View reviewed changes

zhichao-aws merged commit d7945d5 into opensearch-project:main Jan 26, 2024
11 of 13 checks passed

opensearch-trigger-bot bot added the failed backport label Jan 26, 2024

mingshl added backport 2.x and removed failed backport backport 2.x labels Jan 26, 2024

opensearch-trigger-bot bot mentioned this pull request Jan 26, 2024

[Backport 2.x] Enhance RagTool to choose neural sparse query type #152

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Enhance RagTool to choose neural sparse query type #140

Enhance RagTool to choose neural sparse query type #140

mingshl commented Jan 22, 2024 •

edited

Loading

mingshl commented Jan 22, 2024 •

edited

Loading

zhichao-aws left a comment

mingshl commented Jan 23, 2024

codecov bot commented Jan 23, 2024 •

edited

Loading

mingshl commented Jan 23, 2024

zhichao-aws commented Jan 24, 2024

zhichao-aws commented Jan 24, 2024

zhichao-aws Jan 25, 2024

zhichao-aws Jan 25, 2024

mingshl Jan 25, 2024

zhichao-aws Jan 25, 2024

mingshl Jan 25, 2024

zhichao-aws left a comment

opensearch-trigger-bot bot commented Jan 26, 2024

Enhance RagTool to choose neural sparse query type #140

Enhance RagTool to choose neural sparse query type #140

Conversation

mingshl commented Jan 22, 2024 • edited Loading

Description

Check List

mingshl commented Jan 22, 2024 • edited Loading

zhichao-aws left a comment

Choose a reason for hiding this comment

mingshl commented Jan 23, 2024

codecov bot commented Jan 23, 2024 • edited Loading

Codecov Report

mingshl commented Jan 23, 2024

zhichao-aws commented Jan 24, 2024

zhichao-aws commented Jan 24, 2024

zhichao-aws Jan 25, 2024

Choose a reason for hiding this comment

zhichao-aws Jan 25, 2024

Choose a reason for hiding this comment

mingshl Jan 25, 2024

Choose a reason for hiding this comment

zhichao-aws Jan 25, 2024

Choose a reason for hiding this comment

mingshl Jan 25, 2024

Choose a reason for hiding this comment

zhichao-aws left a comment

Choose a reason for hiding this comment

opensearch-trigger-bot bot commented Jan 26, 2024

mingshl commented Jan 22, 2024 •

edited

Loading

mingshl commented Jan 22, 2024 •

edited

Loading

codecov bot commented Jan 23, 2024 •

edited

Loading