[feature][java][chatmodels] Add Azure AI chat model integration #290

tsaiggo · 2025-10-29T13:25:48Z

Linked issue: #286

Purpose of change

This PR adds support for the Azure AI chat model as a foundational model for the agent.

Tests

I added one test to verify this change, but I ran into some issues while running the test cases:

When I tried to run the AgentWithOllamaExample, the output stream didn't produce any data. Consequently, the same thing happened when I tested my new AgentWithAzureAIExample. My assumption is that this isn't related to the chat model implementation itself, but is more likely a higher-level issue (perhaps in the core agent logic or streaming setup).
Because of this, I switched to a different test case, ReActAgentExample. I modified it locally to use my Azure AI implementation, and it worked perfectly. However, I did not commit this change, as my PR should not introduce breaking changes to existing examples.
Finally, since this PR integrates a cloud-hosted model, I cannot check in my personal endpoint or API key. Therefore, I created the new AgentWithAzureAIExample and leveraged the MockChatModel (from AgentWithResource) to add mock tests. The "happy path" for this mock test passes successfully.

API

No

Documentation

Yes, exactly. It's because we now offer more model options.

…ation classes

wenjin272 · 2025-10-31T04:00:10Z

Hi, @tsaiggo, thanks for your work. I will take a look today.

I verified the question of ollama test you mentioned, it is generated after we minor refactor the OllamaChatModelConnection in #279 but forgot to adapt the ollama integration test. I fix it in #291.

wenjin272

Hi, @tsaiggo, thanks for your contribution. I left some minor comments.

...main/java/org/apache/flink/agents/integrations/chatmodels/azureai/AzureAIChatModelSetup.java

wenjin272 · 2025-10-31T06:58:33Z

...main/java/org/apache/flink/agents/integrations/chatmodels/azureai/AzureAIChatModelSetup.java

+        this.prompt = descriptor.getArgument("prompt");
+        this.tools = (java.util.List<String>) descriptor.getArgument("tools");
+        this.key = descriptor.getArgument("key");
+        this.endpoint = descriptor.getArgument("endpoint");


key and endpoint should be the arguments of AzureAIChatModelConnection. Different AzureAIChatModelSetup can share the same AzureAIChatModelConnection.

wenjin272 · 2025-10-31T06:59:32Z

...main/java/org/apache/flink/agents/integrations/chatmodels/azureai/AzureAIChatModelSetup.java

+        java.util.Map<String, Object> params = new java.util.HashMap<>();
+        params.put("model", model);
+        params.put("prompt", prompt);
+        params.put("tools", tools);


Base class BaseChatModelSetup will process the prompt and tools directly, no need to pass them as parameters here.

wenjin272 · 2025-10-31T07:07:14Z

...java/org/apache/flink/agents/integrations/chatmodels/azureai/AzureAIChatModelConnection.java

+            case USER:
+                return new ChatRequestUserMessage(content);
+            case ASSISTANT:
+                return new ChatRequestAssistantMessage(content);


Here, only the content of ChatRequestAssistantMessage is set, but no toolCalls are set.

wenjin272 · 2025-10-31T07:16:01Z

...java/org/apache/flink/agents/integrations/chatmodels/azureai/AzureAIChatModelConnection.java

+            case ASSISTANT:
+                return new ChatRequestAssistantMessage(content);
+            case TOOL:
+                return new ChatRequestToolMessage(content);


Here, the content will be set as toolCallId of ChatRequestToolMessage, so we should pass the function call id generated by Azure llm.

ChatRequestToolMessage msg = new ChatRequestToolMessage( (String) message.getExtraArgs().get("externalId")); msg.setContent(content); return msg;

and the externalId should be set in convertToAgentsTools:

private List<Map<String, Object>> convertToAgentsTools( List<ChatCompletionsToolCall> azureToolCalls) { ... final Map<String, Object> call = Map.of( "id", functionCall.getId(), "original_id", functionCall.getId(), "type", "function", "function", Map.of( "name", functionCall.getFunction().getName(), "arguments", functionCall.getFunction().getArguments())); toolCalls.add(call); ... }

...java/org/apache/flink/agents/integrations/chatmodels/azureai/AzureAIChatModelConnection.java

…enhancing event handling

…back in example

tsaiggo · 2025-10-31T18:23:21Z

pushed new changes, please re-review this PR, thx @wenjin272
btw, there might be something wrong with the python-related test case, can you help retry the failed job? I have no access to do this

wenjin272

Hi, @tsaiggo, thanks for your update. LGTM, only one comment.

The failed test case in CI is caused by the request time out of ollama chat model. I fixed it also in pr #291, you can rebase main branch after it is merged.

wenjin272 · 2025-11-03T06:15:07Z

...main/java/org/apache/flink/agents/integrations/chatmodels/azureai/AzureAIChatModelSetup.java

+        super(descriptor, getResource);
+    }
+
+    // For any other specific parameters, please refer to ChatCompletionsOptions    @Override


The @Override annotation is commented out.

Oops, good catch! 😂 Blame it on the late-night brain fog. Will fix it

tsaiggo · 2025-11-03T14:24:48Z

Looks like the issue of python-related test case has been self-recovered.
lease let me know if you have any further comments or feedback—even nitpicks are welcome. If not, is this PR good to go? @wenjin272

wenjin272 · 2025-11-04T08:28:02Z

Hi, @tsaiggo, this pr looks good to me.

Since merging a PR requires apache committer privileges, @xintongsong could you take a look at this pr at your convenience?

xintongsong

@tsaiggo, thanks for your contribution.

The PR overall looks good. I just have one minor comment.

Additionally, I think this new support for Azure chat model needs to be documented. We can either do that in this PR, or create a separate issue for it.

xintongsong · 2025-11-04T11:59:16Z

examples/pom.xml

+            <groupId>org.apache.flink</groupId>
+            <artifactId>flink-agents-integrations-chat-models-azureai</artifactId>
+            <version>${project.version}</version>
+        </dependency>


Why do we need this dependency in examples/ ?

I prefer to keep documentation updates in a separate PR

added the azureaichatmodel to the examples module because I intended it to be our first example of a cloud-hosted model. My original plan was to add specific use cases for cloud-hosted models, but I haven't come up with any good ideas for that yet, so it's temporarily on hold. If you think this addition should be removed for now, I'm happy to take care of that. @xintongsong

Yes, I think the dependency should be introduced only when it's needed. Thanks for addressing my comments.

FYI, I just created #305 for the documentation work.

xintongsong

LGTM. Merging.

tsaiggo added 4 commits October 29, 2025 21:09

[Feature][Integrations][AzureAI] Add Azure AI chat model integration

7a1bd5e

[Refactor] Improve code readability and formatting in Azure AI integr…

248c04b

…ation classes

[Refactor] Improve code readability and formatting in Azure AI integr…

d3c753c

…ation classes

[Refactor] Improve code readability and formatting in Azure AI integr…

8f31d0e

…ation classes

wenjin272 reviewed Oct 31, 2025

View reviewed changes

tsaiggo added 3 commits November 1, 2025 01:37

[Refactor] Simplify Azure AI integration by removing unused code and …

b81b947

…enhancing event handling

[Feature] Add method to check Azure API configuration and handle fall…

0846236

…back in example

[Refactor] Improve error message formatting in Azure AI example

b1c1854

wenjin272 reviewed Nov 3, 2025

View reviewed changes

[Refactor] Fix formatting of comments in AzureAIChatModelSetup.java

48a880b

xintongsong reviewed Nov 4, 2025

View reviewed changes

[Refactor] Remove unused Azure AI dependency from pom.xml

7c1b189

xintongsong approved these changes Nov 5, 2025

View reviewed changes

xintongsong merged commit d8f2eea into apache:main Nov 5, 2025
11 checks passed

tsaiggo mentioned this pull request Nov 13, 2025

[feature][documentation] Add Azure AI chat model integration details #318

Open

Sxnan added priority/major Default priority of the PR or issue. fixVersion/0.2.0 The feature or bug should be implemented/fixed in the 0.2.0 version. doc-not-needed Your PR changes do not impact docs labels Nov 14, 2025

[feature][java][chatmodels] Add Azure AI chat model integration #290

[feature][java][chatmodels] Add Azure AI chat model integration #290

Uh oh!

Conversation

tsaiggo commented Oct 29, 2025

Purpose of change

Tests

API

Documentation

Uh oh!

wenjin272 commented Oct 31, 2025

Uh oh!

wenjin272 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

tsaiggo commented Oct 31, 2025

Uh oh!

wenjin272 left a comment • 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!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tsaiggo commented Nov 3, 2025

Uh oh!

wenjin272 commented Nov 4, 2025

Uh oh!

xintongsong left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tsaiggo Nov 4, 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!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

xintongsong left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

wenjin272 left a comment •

edited

Loading

tsaiggo Nov 4, 2025 •

edited

Loading