Plan-Execute-Reflect Agent Tracing #3964

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Open

chriswlai wants to merge 15 commits into opensearch-project:feature/agent-tracing from chriswlai:feature/pertracer

chriswlai commented Jul 7, 2025 •

edited

Loading

Description

Adds tracing to the Plan-Execute-Reflect agent. Uses custom index mapping for storage. Basic connection with Conversational agent. Also re-structures the tracing class to be more flexible for non-agent tracing.

nLJ1Ri8m33slNv5Z4OKls04qCGa9q2JGxY6r85kRLeuxiJ7-VMvQGcdfj6c7zHBPVizvzbCpiIIHYZG9vDOIIaL29T9QCQt3vMB31w0u1eA_aQX3SaUTMouUGLA1C3Docq2y1Y9jTY9DRkO3HniAT_SwboOqCeL8I2BKsuB03ch4yQpgz0qu4hb4CAvuWmWcTZJhogSWg0Pi3JOKEw0g4mv-hEztegwLL3bjNa0vr4Dc20xa

Related Issues

Resolves #3971

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.

chriswlai requested review from b4sjoo, dhrubo-os, mingshl, jngz-es, model-collapse, rbhavna, ylwu-amzn, zane-neo, Zhangxunmt, austintlee, HenryL27 and xinyual as code owners

July 7, 2025 07:57

chriswlai requested a deployment to ml-commons-cicd-env-require-approval

July 7, 2025 07:59

— with

GitHub Actions Waiting

chriswlai requested a deployment to ml-commons-cicd-env-require-approval

July 7, 2025 07:59

— with

GitHub Actions Waiting

chriswlai requested a deployment to ml-commons-cicd-env-require-approval

July 7, 2025 07:59

— with

GitHub Actions Waiting

chriswlai requested a deployment to ml-commons-cicd-env-require-approval

July 7, 2025 07:59

— with

GitHub Actions Waiting

chriswlai requested a deployment to ml-commons-cicd-env-require-approval

July 7, 2025 17:28

— with

GitHub Actions Waiting

chriswlai requested a deployment to ml-commons-cicd-env-require-approval

July 7, 2025 17:28

— with

GitHub Actions Waiting

chriswlai requested a deployment to ml-commons-cicd-env-require-approval

July 7, 2025 17:28

— with

GitHub Actions Waiting

chriswlai requested a deployment to ml-commons-cicd-env-require-approval

July 7, 2025 17:28

— with

GitHub Actions Waiting

chriswlai had a problem deploying to ml-commons-cicd-env-require-approval

July 14, 2025 18:16

— with

GitHub Actions Error

chriswlai had a problem deploying to ml-commons-cicd-env-require-approval

July 14, 2025 18:16

— with

GitHub Actions Error

chriswlai had a problem deploying to ml-commons-cicd-env-require-approval

July 14, 2025 18:16

— with

GitHub Actions Failure

chriswlai had a problem deploying to ml-commons-cicd-env-require-approval

July 14, 2025 18:16

— with

GitHub Actions Failure

chriswlai requested a deployment to ml-commons-cicd-env-require-approval

July 15, 2025 19:06

— with

GitHub Actions Waiting

chriswlai requested a deployment to ml-commons-cicd-env-require-approval

July 15, 2025 19:06

— with

GitHub Actions Waiting

chriswlai requested a deployment to ml-commons-cicd-env-require-approval

July 15, 2025 19:06

— with

GitHub Actions Waiting

chriswlai requested a deployment to ml-commons-cicd-env-require-approval

July 15, 2025 19:06

— with

GitHub Actions Waiting

chriswlai added 2 commits

July 15, 2025 12:27


          adding agent tracing to mlplugin

b04f2c7

Signed-off-by: chrislai <chrlaii@amazon.com>


          add tests

4cec2f4

Signed-off-by: chrislai <chrlaii@amazon.com>

chriswlai requested a deployment to ml-commons-cicd-env-require-approval

July 22, 2025 18:34

— with

GitHub Actions Waiting

chriswlai requested a deployment to ml-commons-cicd-env-require-approval

July 22, 2025 18:34

— with

GitHub Actions Waiting

mingshl reviewed

View reviewed changes

ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/AgentUtils.java Outdated

+                                                  if ("aws.bedrock".equalsIgnoreCase(provider)) {
+                                                      // Bedrock/Claude format: input_tokens, output_tokens (or inputTokens, outputTokens)
+                                                      if (usage.containsKey("input_tokens")) {

Collaborator

mingshl Jul 22, 2025

wondering if you check on contains "input_tokens", but if the usage has "input_tokens_meta_data" but no "input_tokens", this will through key not found exception or getting null while toString(), can you verify this?

Author

chriswlai Jul 23, 2025

addressed by adding fallback

ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/AgentUtils.java Outdated

+                                  if (tensors != null && !tensors.isEmpty()) {
+                                      var tensor = tensors.get(0);
+                                      // Try result
+                                      if (tensor.getResult() != null) {

Collaborator

mingshl Jul 22, 2025

did you test it with local model output, even though most of the users should use remote model, but what if a users accidentally passed over a local model output, what would happen here?

Author

chriswlai Jul 23, 2025

yes, very good point. addressed as well by adding fallback for local model


          reconfigure classes and write coverage tests

2b14138

Signed-off-by: chrislai <chrlaii@amazon.com>

chriswlai requested a deployment to ml-commons-cicd-env-require-approval

July 23, 2025 20:38

— with

GitHub Actions Waiting

chriswlai requested a deployment to ml-commons-cicd-env-require-approval

July 23, 2025 20:38

— with

GitHub Actions Waiting

chriswlai requested a deployment to ml-commons-cicd-env-require-approval

July 23, 2025 20:38

— with

GitHub Actions Waiting

chriswlai requested a deployment to ml-commons-cicd-env-require-approval

July 23, 2025 20:38

— with

GitHub Actions Waiting

pyek-bot reviewed

View reviewed changes

common/src/main/resources/index-mappings/ml_agent_trace.json

+                  "version": 1,
+                  "template": {
+                    "mappings": {
+                      "date_detection": false,

Contributor

pyek-bot Jul 23, 2025

Curious, what is this?

Author

chriswlai Jul 23, 2025

Disables automatic date detection in OpenSearch to ensure no field gets automatically mapped to date if we don't want it to

Contributor

pyek-bot Jul 23, 2025

what if a cluster with this index name already exists? do you think we should give a setting to configure this index?

pyek-bot reviewed

View reviewed changes

common/src/main/resources/index-mappings/ml_agent_trace.json

+                            }
+                          }
+                        },
+                        "serviceName": {

Contributor

pyek-bot Jul 23, 2025

snake case for consistency

Author

chriswlai Jul 23, 2025

Unfortunately it conflicts with OTel format, so must keep the standard of camel case

pyek-bot reviewed

View reviewed changes

.../main/java/org/opensearch/ml/engine/algorithms/agent/MLPlanExecuteAndReflectAgentRunner.java Outdated

Comment on lines 250 to 251

		Map<String, String> agentAttributes = MLAgentTracer.createAgentTaskAttributes(mlAgent.getName(), apiParams.get(QUESTION_FIELD));
		Span agentTaskSpan = MLAgentTracer.getInstance().startSpan(MLAgentTracer.AGENT_TASK_PER_SPAN, agentAttributes, null);

Contributor

pyek-bot Jul 23, 2025

since attributes depend on the type of span, can we have methods like startTaskSpan? What do you think?

Author

chriswlai Jul 25, 2025

Good point, fixed

pyek-bot reviewed

View reviewed changes

.../main/java/org/opensearch/ml/engine/algorithms/agent/MLPlanExecuteAndReflectAgentRunner.java Outdated

Comment on lines 487 to 492

+                                  Double inputTokens = planResultInfo.usage != null && planResultInfo.usage.get("inputTokens") instanceof Number
+                                      ? ((Number) planResultInfo.usage.get("inputTokens")).doubleValue()
+                                      : null;
+                                  Double outputTokens = planResultInfo.usage != null && planResultInfo.usage.get("outputTokens") instanceof Number
+                                      ? ((Number) planResultInfo.usage.get("outputTokens")).doubleValue()
+                                      : null;

Contributor

pyek-bot Jul 23, 2025

move attribute creation/setting into methods, keep each method small and reusable

Author

chriswlai Jul 25, 2025

addressed

pyek-bot reviewed

View reviewed changes

.../main/java/org/opensearch/ml/engine/algorithms/agent/MLPlanExecuteAndReflectAgentRunner.java Outdated

Comment on lines 540 to 542

+                                      phaseInputTokens.set(0.0);
+                                      phaseOutputTokens.set(0.0);
+                                      phaseTotalTokens.set(0.0);

Contributor

pyek-bot Jul 23, 2025

can we use a "reset" method?

Author

chriswlai Jul 25, 2025

yep, good point, addressed

pyek-bot reviewed

View reviewed changes

.../main/java/org/opensearch/ml/engine/algorithms/agent/MLPlanExecuteAndReflectAgentRunner.java Outdated

Comment on lines 836 to 838

+                          // if (!allParams.containsKey(LLM_RESPONSE_FILTER) || allParams.get(LLM_RESPONSE_FILTER).isEmpty()) {
+                          // throw new IllegalArgumentException("llm_response_filter not found. Please provide the path to the model output.");
+                          // }

Contributor

pyek-bot Jul 23, 2025

why has this been commented out?

Author

chriswlai Jul 23, 2025

good catch

pyek-bot reviewed

View reviewed changes

.../main/java/org/opensearch/ml/engine/algorithms/agent/MLPlanExecuteAndReflectAgentRunner.java Outdated

+                                      Map<String, String> spanContextMap = new HashMap<>();
+                                      MLAgentTracer.getInstance().injectSpanContext(executeStepSpan, spanContextMap);
+                                      reactParams.putAll(spanContextMap);
+                                      log.info("[AGENT_TRACE] PER Agent - Injected parent SpanContext: {}", spanContextMap);

Contributor

pyek-bot Jul 23, 2025

use debug logs

Author

chriswlai Jul 25, 2025

addressed


          reconfig + tests

cbe45a3

Signed-off-by: chrislai <chrlaii@amazon.com>

chriswlai temporarily deployed to ml-commons-cicd-env-require-approval

July 24, 2025 00:38

— with

GitHub Actions Inactive

chriswlai had a problem deploying to ml-commons-cicd-env-require-approval

July 24, 2025 00:38

— with

GitHub Actions Failure

chriswlai temporarily deployed to ml-commons-cicd-env-require-approval

July 24, 2025 00:38

— with

GitHub Actions Inactive

chriswlai had a problem deploying to ml-commons-cicd-env-require-approval

July 24, 2025 00:38

— with

GitHub Actions Error

chriswlai requested a deployment to ml-commons-cicd-env-require-approval

July 24, 2025 02:04

— with

GitHub Actions Waiting

chriswlai requested a deployment to ml-commons-cicd-env-require-approval

July 24, 2025 02:04

— with

GitHub Actions Waiting

mingshl reviewed

View reviewed changes

ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/tracing/MLTracer.java

    
                              newSpan.addAttribute("thread.name", Thread.currentThread().getName());

                          } catch (Exception e) {

                              log.warn("Failed to create root span, falling back to normal span creation", e);

Collaborator

mingshl Jul 24, 2025

what is the consequence if it failed to create root span and become a normal span? is it we lost the relationship of a root span?

Author

chriswlai Jul 24, 2025

Yes, but the opposite actually. If we fail to create a root span, the span will still be created except it will show to have a parent span, but that span doesn't exist.

...lgorithms/src/test/java/org/opensearch/ml/engine/algorithms/agent/tracing/MLTracerTests.java

		@@ -0,0 +1,182 @@
		package org.opensearch.ml.engine.algorithms.agent.tracing;

		import static org.junit.Assert.*;

Collaborator

mingshl Jul 24, 2025

avoid import *

Author

chriswlai Jul 25, 2025

yes, addressed

...lgorithms/src/test/java/org/opensearch/ml/engine/algorithms/agent/tracing/MLTracerTests.java

+                  private Tracer mockTracer;
+                  private MLFeatureEnabledSetting mockFeatureSetting;
+                  @Before

Collaborator

mingshl Jul 24, 2025

add proper java doc for you class and test method

Author

chriswlai Jul 25, 2025

right, thanks. done

...st/java/org/opensearch/ml/engine/algorithms/agent/tracing/MLAgentTracerStaticUtilsTests.java Outdated

		@@ -0,0 +1,180 @@
		package org.opensearch.ml.engine.algorithms.agent.tracing;

		import static org.junit.Assert.*;

Collaborator

mingshl Jul 24, 2025

avoid *

Author

chriswlai Jul 25, 2025

addressed

...st/java/org/opensearch/ml/engine/algorithms/agent/tracing/MLAgentTracerStaticUtilsTests.java Outdated

+              package org.opensearch.ml.engine.algorithms.agent.tracing;
+              import static org.junit.Assert.*;
+              import static org.mockito.Mockito.*;

Collaborator

mingshl Jul 24, 2025

also here

Author

chriswlai Jul 25, 2025

addressed

...st/java/org/opensearch/ml/engine/algorithms/agent/tracing/MLAgentTracerStaticUtilsTests.java

+              import org.opensearch.ml.common.output.model.ModelTensors;
+              import org.opensearch.telemetry.tracing.Span;
+              public class MLAgentTracerStaticUtilsTests {

Collaborator

mingshl Jul 24, 2025

java doc

Author

chriswlai Jul 25, 2025

done

...lgorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/tracing/MLAgentTracer.java Outdated

+                                                  @SuppressWarnings("unchecked")
+                                                  Map<String, Object> usage = (Map<String, Object>) usageObj;
+                                                  if ("aws.bedrock".equalsIgnoreCase(provider)) {

Collaborator

mingshl Jul 24, 2025

you are using equalsIgnoreCase for "aws.bedrock" but you are not getting "aws.bedrock".

I think "aws.bedrock" would be better to be "containsKey"

but the rest of the key that you are actually getting, for example, "input_tokens" and "prompt_tokens", before you are getting the keys, it makes more sense to check equalsIgnoreCase

Author

chriswlai Jul 25, 2025

yes, very good point. fixed


          address comments

f9f699a

Signed-off-by: chrislai <chrlaii@amazon.com>

chriswlai temporarily deployed to ml-commons-cicd-env-require-approval

July 25, 2025 16:48

— with

GitHub Actions Inactive

chriswlai temporarily deployed to ml-commons-cicd-env-require-approval

July 25, 2025 16:48

— with

GitHub Actions Inactive

chriswlai temporarily deployed to ml-commons-cicd-env-require-approval

July 25, 2025 16:48

— with

GitHub Actions Inactive

chriswlai temporarily deployed to ml-commons-cicd-env-require-approval

July 25, 2025 16:48

— with

GitHub Actions Inactive

chriswlai temporarily deployed to ml-commons-cicd-env-require-approval

July 25, 2025 17:53

— with

GitHub Actions Inactive

chriswlai deployed to ml-commons-cicd-env-require-approval

July 25, 2025 17:53

— with

GitHub Actions Active

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

mingshl mingshl left review comments

pyek-bot pyek-bot left review comments

b4sjoo Awaiting requested review from b4sjoo b4sjoo is a code owner

dhrubo-os Awaiting requested review from dhrubo-os dhrubo-os is a code owner

jngz-es Awaiting requested review from jngz-es jngz-es is a code owner

model-collapse Awaiting requested review from model-collapse model-collapse is a code owner

rbhavna Awaiting requested review from rbhavna rbhavna is a code owner

ylwu-amzn Awaiting requested review from ylwu-amzn ylwu-amzn is a code owner

zane-neo Awaiting requested review from zane-neo zane-neo is a code owner

Zhangxunmt Awaiting requested review from Zhangxunmt Zhangxunmt is a code owner

austintlee Awaiting requested review from austintlee austintlee is a code owner

HenryL27 Awaiting requested review from HenryL27 HenryL27 is a code owner

xinyual Awaiting requested review from xinyual xinyual is a code owner

Labels

None yet