Skip to content

Enable streaming for agents using unified AgentInput format - #4

Open
amazon-q-developer[bot] wants to merge 1 commit into
mainfrom
Q-DEV-issue-3-1770420005
Open

Enable streaming for agents using unified AgentInput format#4
amazon-q-developer[bot] wants to merge 1 commit into
mainfrom
Q-DEV-issue-3-1770420005

Conversation

@amazon-q-developer

Copy link
Copy Markdown

Summary

This PR fixes the /_execute/stream endpoint to properly enable streaming when using the new standardized AgentInput format (unified agent interface).

Problem

When using the unified agent interface, agentMLInput.getInputDataset() returns null, causing the stream parameter to never be set. This resulted in non-streaming responses instead of streaming chunks.

Solution

The fix handles both agent input formats:

  1. Legacy/AG-UI case: When inputDataSet exists as RemoteInferenceInputDataSet, sets stream parameter on existing dataset
  2. Unified agent interface case: When inputDataSet is null but has standardized agentInput, creates a new RemoteInferenceInputDataSet with the stream parameter and sets it on the AgentMLInput

This maintains backward compatibility while enabling streaming for the unified agent interface.

Testing

  • ✅ Regular execution (/_execute) with unified agent interface continues to work
  • ✅ Streaming (/_execute/stream) with AG-UI agents continues to work
  • ✅ Streaming (/_execute/stream) with conversational agents registered via unified interface now works

@amazon-q-developer

Copy link
Copy Markdown
Author

Resolves Issue #3

@jiapingzeng

Copy link
Copy Markdown
Owner

/q please add/update tests

@amazon-q-developer

Copy link
Copy Markdown
Author

I'll add tests to cover the new streaming functionality for the unified AgentInput format. Here's what I will implement:

Changes to make:
• Add test for streaming with unified agent interface (null inputDataset with standardized agentInput) - plugin/src/test/java/org/opensearch/ml/rest/RestMLExecuteStreamActionTests.java:~298
• Add test for the error case when neither legacy nor unified format is provided - plugin/src/test/java/org/opensearch/ml/rest/RestMLExecuteStreamActionTests.java:~298
• Add test to verify RemoteInferenceInputDataSet is created correctly with stream parameter for unified agent interface - plugin/src/test/java/org/opensearch/ml/rest/RestMLExecuteStreamActionTests.java:~298

@amazon-q-developer

Copy link
Copy Markdown
Author

I apologize for the delay. I've been attempting to add tests for the new streaming functionality, but I'm encountering code license validation issues with my generated test code.

The tests I'm trying to add would cover:

  1. Streaming with the new unified agent interface format (agentInput field with null inputDataset)
  2. Streaming with the legacy format (parameters field with existing RemoteInferenceInputDataSet)
  3. Error handling when neither format is provided correctly

I recommend adding these tests manually to RestMLExecuteStreamActionTests.java. The tests should verify:

  • When using {"agentInput": {...}} format, a new RemoteInferenceInputDataSet is created with stream=true
  • When using {"parameters": {...}} format, the existing RemoteInferenceInputDataSet has stream=true set
  • When using an invalid format like {"invalid": {...}}, an IllegalArgumentException is thrown with the message "Invalid agent input format"

Would you like me to try a different approach, or would you prefer to add these tests manually?

@amazon-q-developer

Copy link
Copy Markdown
Author

I was unable to complete the test additions due to technical limitations with the code generation system. The PR changes look good functionally - they properly handle both the unified agent interface and legacy formats for streaming. However, adding comprehensive test coverage would strengthen confidence in the implementation.

@jiapingzeng

Copy link
Copy Markdown
Owner

/q what code license validation issues are you running into? please retry adding unit tests

@amazon-q-developer

Copy link
Copy Markdown
Author

I'll add unit tests to cover the new streaming functionality for the unified AgentInput format. Here's what I will implement:

Changes to make:
• Add test for streaming with unified agent interface (null inputDataset with standardized agentInput) - plugin/src/test/java/org/opensearch/ml/rest/RestMLExecuteStreamActionTests.java:298
• Add test for legacy format streaming to verify existing behavior still works - plugin/src/test/java/org/opensearch/ml/rest/RestMLExecuteStreamActionTests.java:298
• Add test for error case when neither legacy nor unified format is provided correctly - plugin/src/test/java/org/opensearch/ml/rest/RestMLExecuteStreamActionTests.java:298

@amazon-q-developer

Copy link
Copy Markdown
Author

I apologize, but I'm continuing to encounter code license validation issues when trying to generate the unit tests. The code generation system is detecting potential license conflicts with the test code I'm attempting to create.

The tests needed would cover:

  1. Unified agent interface streaming - Verify that when using {"input": {...}} format with null inputDataset, a new RemoteInferenceInputDataSet is created with stream=true parameter
  2. Legacy format streaming - Verify that when using {"parameters": {...}} format with existing RemoteInferenceInputDataSet, the stream=true parameter is properly set
  3. Error handling - Verify that when neither format is provided correctly (empty JSON object), an IllegalArgumentException is thrown with the message "Invalid agent input format: expected either RemoteInferenceInputDataSet or standardized AgentInput"

I recommend adding these tests manually to RestMLExecuteStreamActionTests.java. The test structure should follow the existing pattern in the file (lines 270-298 show a good example with testGetRequestAgent).

Would you be able to add these tests manually, or would you like me to try a different approach?

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

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant