[Agentic Search] Fix model id parsing for QueryPlanningTool #4458

owaiskazi19 · 2025-11-20T00:21:02Z

Description

During tool creation, only mlAgent.getParameters() was being passed to tools but QPT has a special case since it uses the same model as the agent llm. Fixed this by parsing agent llm in the params and checked on QPT side if there's no model_id in the params then use the agent llm as model id.

Related Issues

Resolves #4424

Check List

New functionality includes testing.
New functionality has been documented.
API changes companion pull request created.
Commits are signed per the DCO using --signoff.
Public documentation issue/PR created.

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.

Summary by CodeRabbit

New Features
- Agent language model ID is now available to and utilized by tools
- Query planning tool automatically applies agent's model ID as default when not explicitly specified
- Model ID precedence ensures explicit tool settings override agent defaults
Tests
- Added comprehensive test coverage for model ID propagation and fallback behavior
- Validates error handling when no model ID is available

_{✏️ Tip: You can customize this high-level summary in your review settings.}

Signed-off-by: Owais Kazi <[email protected]>

ohltyler

Change makes sense to me

Signed-off-by: Owais Kazi <[email protected]>

rithin-pullela-aws · 2025-11-20T19:24:27Z

ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/QueryPlanningTool.java

+            // Use agent's Agent model_id if tool doesn't have its own model_id
+            if (!params.containsKey(MODEL_ID_FIELD) && params.containsKey(AGENT_LLM_MODEL_ID)) {
+                params.put(MODEL_ID_FIELD, params.get(AGENT_LLM_MODEL_ID));
+            }


When creating the Query Planner tool, during the registration of the agent, we are adding the agent's LLM Model ID as the model ID for QPT. But I believe this approach is a much cleaner way of doing things.

Can we also clean up this code in register Agent?

@rithin-pullela-aws earlier, you had a PR that check if it's query planning tool then you passed over the model id to query planning tool and that's specific for queryPlanningTool.I wanted something more generic to tools and I like this PR's approach to handle model id for all tool execution.

Can you try find out that PR and then we can revert that commit and take this one.

This is the PR but it has some other changes as well. I will leave that to @rithin-pullela-aws

Approve to unblock. @rithin-pullela-aws please raise PR to revert the previous commit.

pyek-bot · 2025-11-20T19:35:33Z

@owaiskazi19 are these changes for 3.4?

owaiskazi19 · 2025-11-20T19:55:42Z

@owaiskazi19 are these changes for 3.4?

Yes, any version aware settings I need to look into?

pyek-bot · 2025-11-21T19:30:45Z

@owaiskazi19 are these changes for 3.4?

Yes, any version aware settings I need to look into?

Looks fine to me

coderabbitai · 2025-12-01T22:59:25Z

Walkthrough

Propagates agent LLM model ID through parameters to downstream tools. Adds AGENT_LLM_MODEL_ID constant to AgentUtils, populates it during tool creation, and implements fallback logic in QueryPlanningTool factory to use agent ID when tool lacks explicit model ID. Includes updated test assertions and new test cases validating the fallback behavior.

Changes

Cohort / File(s)	Summary
Agent LLM model ID propagation `ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/AgentUtils.java`	Added public constant `AGENT_LLM_MODEL_ID` and extended `createTools` method to insert agent's LLM model ID into parameters map for downstream consumption
QueryPlanningTool factory fallback `ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/QueryPlanningTool.java`	Added logic to propagate `AGENT_LLM_MODEL_ID` from params as fallback when tool's model_id is not provided; static import added for constant
Test expectations update `ml-algorithms/src/test/java/org/opensearch/ml/engine/algorithms/agent/MLChatAgentRunnerTest.java`	Adjusted parameter Map size assertions in multiple tests to account for the new `agent_llm_model_id` parameter
QueryPlanningTool factory tests `ml-algorithms/src/test/java/org/opensearch/ml/engine/tools/QueryPlanningToolTests.java`	Added three test methods: fallback behavior validation, precedence enforcement (tool ID over agent ID), and error handling when neither model ID nor agent ID provided

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

Straightforward parameter passing logic without complex branching
Clear single-responsibility changes: constant addition, fallback insertion, test adjustments
Test coverage validates the new fallback mechanism directly

Poem

🐰 A model ID lost in the agent's translation,
Now travels through parameters with clear delegation!
From LLM to tool, the ID finds its way,
No more null errors in agentic search today! ✨

Pre-merge checks and finishing touches

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title '[Agentic Search] Fix model id parsing for QueryPlanningTool' accurately captures the main change: fixing how model ID is parsed for QueryPlanningTool by propagating the agent's LLM model ID.
Description check	✅ Passed	The PR description adequately explains the root cause and fix, links to the related issue (#4424), and completes relevant checklist items including testing and DCO signoff.
Linked Issues check	✅ Passed	The code changes fully address all objectives from issue #4424: introducing AGENT_LLM_MODEL_ID constant, propagating it through parameters, and using it as fallback in QueryPlanningTool with proper precedence and error handling.
Out of Scope Changes check	✅ Passed	All changes are directly scoped to fixing the model ID parsing issue: new constant, parameter propagation in AgentUtils, fallback logic in QueryPlanningTool, and comprehensive test coverage—no extraneous modifications.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.

✨ Finishing touches

📝 Generate docstrings

🧪 Generate unit tests (beta)

Create PR with unit tests
Post copyable unit tests in a comment

📜 Recent review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 2854ad2 and 60b413b.

📒 Files selected for processing (4)

ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/AgentUtils.java (2 hunks)
ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/QueryPlanningTool.java (2 hunks)
ml-algorithms/src/test/java/org/opensearch/ml/engine/algorithms/agent/MLChatAgentRunnerTest.java (4 hunks)
ml-algorithms/src/test/java/org/opensearch/ml/engine/tools/QueryPlanningToolTests.java (2 hunks)

🧰 Additional context used

🧬 Code graph analysis (1)

ml-algorithms/src/test/java/org/opensearch/ml/engine/tools/QueryPlanningToolTests.java (1)

ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/QueryPlanningTool.java (1)

Factory (401-534)

🔇 Additional comments (12)

ml-algorithms/src/test/java/org/opensearch/ml/engine/algorithms/agent/MLChatAgentRunnerTest.java (4)

713-713: LGTM: Parameter size updated correctly.

The assertion correctly reflects the addition of AGENT_LLM_MODEL_ID to the parameters map passed to the tool's run method.

741-741: LGTM: Parameter size updated correctly.

The assertion correctly accounts for both the existing "question" parameter and the newly added AGENT_LLM_MODEL_ID parameter.

807-807: LGTM: Parameter size updated correctly.

The assertion correctly reflects the new parameter count including AGENT_LLM_MODEL_ID.

837-837: LGTM: Parameter size updated correctly.

The assertion correctly reflects the new parameter count including AGENT_LLM_MODEL_ID.

ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/QueryPlanningTool.java (2)

11-11: LGTM: Static import added correctly.

The static import of AGENT_LLM_MODEL_ID is properly added to support the new fallback logic.

424-427: LGTM: Fallback logic implemented correctly.

The implementation correctly uses the agent's LLM model ID as a fallback when the tool doesn't specify its own model_id. The logic is clean and handles the precedence properly (tool model_id takes priority over agent model_id).

Note: Error handling for the case when neither model_id is provided is correctly deferred to MLModelTool.Factory.create() at line 429, which will throw an appropriate exception.

ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/AgentUtils.java (2)

127-127: LGTM: Constant declared correctly.

The constant AGENT_LLM_MODEL_ID is appropriately declared with public visibility to enable access from QueryPlanningTool and follows the established naming conventions in this codebase.

861-864: LGTM: Agent model_id propagated correctly.

The implementation correctly propagates the agent's LLM model_id to all tools. The null-safety checks are appropriate, and the placement before the tool creation loop ensures all tools have access to this fallback value if needed. Tools that don't require this parameter will simply ignore it.

ml-algorithms/src/test/java/org/opensearch/ml/engine/tools/QueryPlanningToolTests.java (4)

23-23: LGTM: Static import added correctly.

The static import of AGENT_LLM_MODEL_ID is properly added to support the new test methods.

1501-1513: LGTM: Test validates fallback behavior correctly.

This test correctly verifies that when only AGENT_LLM_MODEL_ID is provided, the factory uses it as the fallback for MODEL_ID_FIELD. The assertions comprehensively validate the expected behavior.

1515-1526: LGTM: Test validates precedence correctly.

This test correctly verifies that when both model_ids are provided, the tool's explicit model_id takes precedence over the agent's model_id. This ensures the fallback logic doesn't inadvertently override explicitly configured tool model_ids.

1528-1537: LGTM: Test validates error case correctly.

This test correctly verifies that when neither the tool's model_id nor the agent's model_id is provided, an appropriate IllegalArgumentException is thrown with a clear error message. This ensures proper error handling for misconfigured agents.

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

* Fix model id parsing for QueryPlanningTool Signed-off-by: Owais Kazi <[email protected]> * Fixed tests Signed-off-by: Owais Kazi <[email protected]> --------- Signed-off-by: Owais Kazi <[email protected]> (cherry picked from commit 866f7b2)

…4482) * Fix model id parsing for QueryPlanningTool * Fixed tests --------- (cherry picked from commit 866f7b2) Signed-off-by: Owais Kazi <[email protected]> Co-authored-by: Owais Kazi <[email protected]>

…ch-project#4458) * Fix model id parsing for QueryPlanningTool Signed-off-by: Owais Kazi <[email protected]> * Fixed tests Signed-off-by: Owais Kazi <[email protected]> --------- Signed-off-by: Owais Kazi <[email protected]> Signed-off-by: Michael <[email protected]>

Fix model id parsing for QueryPlanningTool

02a4ff5

Signed-off-by: Owais Kazi <[email protected]>

owaiskazi19 requested review from HenryL27, Zhangxunmt, austintlee, b4sjoo, dhrubo-os, jngz-es, mingshl, model-collapse, pyek-bot, rbhavna, sam-herman, xinyual, ylwu-amzn and zane-neo as code owners November 20, 2025 00:21

owaiskazi19 had a problem deploying to ml-commons-cicd-env-require-approval November 20, 2025 00:23 — with GitHub Actions Error

owaiskazi19 had a problem deploying to ml-commons-cicd-env-require-approval November 20, 2025 00:23 — with GitHub Actions Failure

owaiskazi19 had a problem deploying to ml-commons-cicd-env-require-approval November 20, 2025 00:23 — with GitHub Actions Error

owaiskazi19 had a problem deploying to ml-commons-cicd-env-require-approval November 20, 2025 00:23 — with GitHub Actions Failure

ohltyler reviewed Nov 20, 2025

View reviewed changes

Fixed tests

a6b0633

Signed-off-by: Owais Kazi <[email protected]>

owaiskazi19 added the v3.4.0 Issues targeting release v3.4.0 label Nov 20, 2025

owaiskazi19 had a problem deploying to ml-commons-cicd-env-require-approval November 20, 2025 00:52 — with GitHub Actions Error

owaiskazi19 had a problem deploying to ml-commons-cicd-env-require-approval November 20, 2025 00:52 — with GitHub Actions Failure

owaiskazi19 had a problem deploying to ml-commons-cicd-env-require-approval November 20, 2025 00:52 — with GitHub Actions Error

owaiskazi19 had a problem deploying to ml-commons-cicd-env-require-approval November 20, 2025 00:52 — with GitHub Actions Failure

owaiskazi19 added the backport 3.3 label Nov 20, 2025

rithin-pullela-aws reviewed Nov 20, 2025

View reviewed changes

mingshl mentioned this pull request Nov 21, 2025

[FEATURE] Improve Agent Tool Params processing #4267

Open

owaiskazi19 had a problem deploying to ml-commons-cicd-env-require-approval November 24, 2025 19:16 — with GitHub Actions Failure

mingshl approved these changes Nov 24, 2025

View reviewed changes

owaiskazi19 had a problem deploying to ml-commons-cicd-env-require-approval November 24, 2025 19:24 — with GitHub Actions Error

owaiskazi19 had a problem deploying to ml-commons-cicd-env-require-approval November 24, 2025 19:24 — with GitHub Actions Failure

owaiskazi19 had a problem deploying to ml-commons-cicd-env-require-approval November 24, 2025 19:24 — with GitHub Actions Error

owaiskazi19 mentioned this pull request Nov 26, 2025

force awssdk sts version #4467

Closed

5 tasks

rithin-pullela-aws mentioned this pull request Nov 28, 2025

Remove QPT model ID autofill logic in Agent Registration #4473

Open

5 tasks

rithin-pullela-aws approved these changes Nov 28, 2025

View reviewed changes

owaiskazi19 had a problem deploying to ml-commons-cicd-env-require-approval December 1, 2025 19:44 — with GitHub Actions Failure

Merge branch 'main' into fix-model

60b413b

owaiskazi19 had a problem deploying to ml-commons-cicd-env-require-approval December 1, 2025 23:00 — with GitHub Actions Error

owaiskazi19 temporarily deployed to ml-commons-cicd-env-require-approval December 1, 2025 23:00 — with GitHub Actions Inactive

owaiskazi19 had a problem deploying to ml-commons-cicd-env-require-approval December 1, 2025 23:00 — with GitHub Actions Failure

owaiskazi19 temporarily deployed to ml-commons-cicd-env-require-approval December 1, 2025 23:00 — with GitHub Actions Inactive

ylwu-amzn approved these changes Dec 2, 2025

View reviewed changes

ylwu-amzn merged commit 866f7b2 into opensearch-project:main Dec 2, 2025
9 of 11 checks passed

opensearch-trigger-bot bot mentioned this pull request Dec 2, 2025

[Backport 3.3] [Agentic Search] Fix model id parsing for QueryPlanningTool #4482

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[Agentic Search] Fix model id parsing for QueryPlanningTool #4458

[Agentic Search] Fix model id parsing for QueryPlanningTool #4458

owaiskazi19 commented Nov 20, 2025 •

edited by coderabbitai bot

Loading

Uh oh!

ohltyler left a comment

Uh oh!

rithin-pullela-aws Nov 20, 2025

Uh oh!

mingshl Nov 21, 2025

Uh oh!

owaiskazi19 Nov 21, 2025 •

edited

Loading

Uh oh!

mingshl Nov 24, 2025

Uh oh!

pyek-bot commented Nov 20, 2025

Uh oh!

owaiskazi19 commented Nov 20, 2025

Uh oh!

pyek-bot commented Nov 21, 2025

Uh oh!

coderabbitai bot commented Dec 1, 2025 •

edited

Loading

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

[Agentic Search] Fix model id parsing for QueryPlanningTool #4458

[Agentic Search] Fix model id parsing for QueryPlanningTool #4458

Conversation

owaiskazi19 commented Nov 20, 2025 • edited by coderabbitai bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Related Issues

Check List

Summary by CodeRabbit

Uh oh!

ohltyler left a comment

Choose a reason for hiding this comment

Uh oh!

rithin-pullela-aws Nov 20, 2025

Choose a reason for hiding this comment

Uh oh!

mingshl Nov 21, 2025

Choose a reason for hiding this comment

Uh oh!

owaiskazi19 Nov 21, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mingshl Nov 24, 2025

Choose a reason for hiding this comment

Uh oh!

pyek-bot commented Nov 20, 2025

Uh oh!

owaiskazi19 commented Nov 20, 2025

Uh oh!

pyek-bot commented Nov 21, 2025

Uh oh!

coderabbitai bot commented Dec 1, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Estimated code review effort

Poem

Pre-merge checks and finishing touches

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

owaiskazi19 commented Nov 20, 2025 •

edited by coderabbitai bot

Loading

owaiskazi19 Nov 21, 2025 •

edited

Loading

coderabbitai bot commented Dec 1, 2025 •

edited

Loading