This is an automated email from the ASF dual-hosted git repository.
xtsong pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/flink-agents.git
The following commit(s) were added to refs/heads/main by this push:
new d8f2eea [Feature][integration] Add Azure chat model integration in
java (#290)
d8f2eea is described below
commit d8f2eea5213f49ec463fc2f3be439d234a8e4b26
Author: tsaiggo <[email protected]>
AuthorDate: Wed Nov 5 09:40:51 2025 +0800
[Feature][integration] Add Azure chat model integration in java (#290)
---
e2e-test/integration-test/pom.xml | 5 +
.../agents/integration/test/AgentWithAzureAI.java | 126 ++++++++++++
.../integration/test/AgentWithAzureAIExample.java | 71 +++++++
.../chat-models/azureai}/pom.xml | 45 ++---
.../azureai/AzureAIChatModelConnection.java | 223 +++++++++++++++++++++
.../chatmodels/azureai/AzureAIChatModelSetup.java | 66 ++++++
integrations/chat-models/pom.xml | 1 +
7 files changed, 507 insertions(+), 30 deletions(-)
diff --git a/e2e-test/integration-test/pom.xml
b/e2e-test/integration-test/pom.xml
index c912248..1551595 100644
--- a/e2e-test/integration-test/pom.xml
+++ b/e2e-test/integration-test/pom.xml
@@ -59,6 +59,11 @@ under the License.
<artifactId>flink-clients</artifactId>
<version>${flink.version}</version>
</dependency>
+ <dependency>
+ <groupId>org.apache.flink</groupId>
+
<artifactId>flink-agents-integrations-chat-models-azureai</artifactId>
+ <version>${project.version}</version>
+ </dependency>
<dependency>
<groupId>org.apache.flink</groupId>
<artifactId>flink-agents-integrations-chat-models-ollama</artifactId>
diff --git
a/e2e-test/integration-test/src/main/java/org/apache/flink/agents/integration/test/AgentWithAzureAI.java
b/e2e-test/integration-test/src/main/java/org/apache/flink/agents/integration/test/AgentWithAzureAI.java
new file mode 100644
index 0000000..e16f7e0
--- /dev/null
+++
b/e2e-test/integration-test/src/main/java/org/apache/flink/agents/integration/test/AgentWithAzureAI.java
@@ -0,0 +1,126 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.flink.agents.integration.test;
+
+import org.apache.flink.agents.api.Agent;
+import org.apache.flink.agents.api.InputEvent;
+import org.apache.flink.agents.api.OutputEvent;
+import org.apache.flink.agents.api.annotation.*;
+import org.apache.flink.agents.api.chat.messages.ChatMessage;
+import org.apache.flink.agents.api.chat.messages.MessageRole;
+import org.apache.flink.agents.api.context.RunnerContext;
+import org.apache.flink.agents.api.event.ChatRequestEvent;
+import org.apache.flink.agents.api.event.ChatResponseEvent;
+import org.apache.flink.agents.api.resource.ResourceDescriptor;
+import
org.apache.flink.agents.integrations.chatmodels.azureai.AzureAIChatModelConnection;
+import
org.apache.flink.agents.integrations.chatmodels.azureai.AzureAIChatModelSetup;
+
+import java.util.Collections;
+
+public class AgentWithAzureAI extends Agent {
+
+ private static final String AZURE_ENDPOINT = "";
+ private static final String AZURE_API_KEY = "";
+
+ public static boolean callingRealMode() {
+ if (AZURE_ENDPOINT != null
+ && !AZURE_ENDPOINT.isEmpty()
+ && AZURE_API_KEY != null
+ && !AZURE_API_KEY.isEmpty()) {
+ return true;
+ } else {
+ return false;
+ }
+ }
+
+ @ChatModelConnection
+ public static ResourceDescriptor azureAIChatModelConnection() {
+ return
ResourceDescriptor.Builder.newBuilder(AzureAIChatModelConnection.class.getName())
+ .addInitialArgument("endpoint", AZURE_ENDPOINT)
+ .addInitialArgument("apiKey", AZURE_API_KEY)
+ .build();
+ }
+
+ @ChatModelSetup
+ public static ResourceDescriptor azureAIChatModel() {
+ System.out.println(
+ "Calling real Azure AI service. Make sure the endpoint and
apiKey are correct.");
+ return
ResourceDescriptor.Builder.newBuilder(AzureAIChatModelSetup.class.getName())
+ .addInitialArgument("connection", "azureAIChatModelConnection")
+ .addInitialArgument("model", "gpt-4o")
+ .build();
+ }
+
+ @Tool(description = "Converts temperature between Celsius and Fahrenheit")
+ public static double convertTemperature(
+ @ToolParam(name = "value", description = "Temperature value to
convert") Double value,
+ @ToolParam(
+ name = "fromUnit",
+ description = "Source unit ('C' for Celsius or 'F'
for Fahrenheit)")
+ String fromUnit,
+ @ToolParam(
+ name = "toUnit",
+ description = "Target unit ('C' for Celsius or 'F'
for Fahrenheit)")
+ String toUnit) {
+
+ fromUnit = fromUnit.toUpperCase();
+ toUnit = toUnit.toUpperCase();
+
+ if (fromUnit.equals(toUnit)) {
+ return value;
+ }
+
+ if (fromUnit.equals("C") && toUnit.equals("F")) {
+ return (value * 9 / 5) + 32;
+ } else if (fromUnit.equals("F") && toUnit.equals("C")) {
+ return (value - 32) * 5 / 9;
+ } else {
+ throw new IllegalArgumentException("Invalid temperature units. Use
'C' or 'F'");
+ }
+ }
+
+ @Tool(description = "Calculates Body Mass Index (BMI)")
+ public static double calculateBMI(
+ @ToolParam(name = "weightKg", description = "Weight in kilograms")
Double weightKg,
+ @ToolParam(name = "heightM", description = "Height in meters")
Double heightM) {
+
+ if (weightKg <= 0 || heightM <= 0) {
+ throw new IllegalArgumentException("Weight and height must be
positive values");
+ }
+ return weightKg / (heightM * heightM);
+ }
+
+ @Tool(description = "Create a random number")
+ public static double createRandomNumber() {
+ return Math.random();
+ }
+
+ @Action(listenEvents = {InputEvent.class})
+ public static void process(InputEvent event, RunnerContext ctx) throws
Exception {
+ ctx.sendEvent(
+ new ChatRequestEvent(
+ "azureAIChatModel",
+ Collections.singletonList(
+ new ChatMessage(MessageRole.USER, (String)
event.getInput()))));
+ }
+
+ @Action(listenEvents = {ChatResponseEvent.class})
+ public static void processChatResponse(ChatResponseEvent event,
RunnerContext ctx) {
+ ctx.sendEvent(new OutputEvent(event.getResponse().getContent()));
+ }
+}
diff --git
a/e2e-test/integration-test/src/main/java/org/apache/flink/agents/integration/test/AgentWithAzureAIExample.java
b/e2e-test/integration-test/src/main/java/org/apache/flink/agents/integration/test/AgentWithAzureAIExample.java
new file mode 100644
index 0000000..bb56a21
--- /dev/null
+++
b/e2e-test/integration-test/src/main/java/org/apache/flink/agents/integration/test/AgentWithAzureAIExample.java
@@ -0,0 +1,71 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.flink.agents.integration.test;
+
+import org.apache.flink.agents.api.AgentsExecutionEnvironment;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.streaming.api.datastream.DataStream;
+import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+
+public class AgentWithAzureAIExample {
+ /** Runs the example pipeline. */
+ public static void main(String[] args) throws Exception {
+ if (!AgentWithAzureAI.callingRealMode()) {
+ // print warning information
+ System.err.println(
+ "Please set the AZURE_ENDPOINT and AZURE_API_KEY in the
AgentWithAzureAI class to run this example in real mode.");
+ System.err.println("Falling back to mock mode.");
+ AgentWithResourceExample.main(args);
+ return;
+ }
+ // Create the execution environment
+ StreamExecutionEnvironment env =
StreamExecutionEnvironment.getExecutionEnvironment();
+ env.setParallelism(1);
+
+ // Use prompts that trigger different tool calls in the agent
+ DataStream<String> inputStream =
+ env.fromData(
+ "Convert 25 degrees Celsius to Fahrenheit",
+ "What is 98.6 Fahrenheit in Celsius?",
+ "Change 32 degrees Celsius to Fahrenheit",
+ "If it's 75 degrees Fahrenheit, what would that be in
Celsius?",
+ "Convert room temperature of 20C to F",
+ "Calculate BMI for someone who is 1.75 meters tall and
weighs 70 kg",
+ "What's the BMI for a person weighing 85 kg with
height 1.80 meters?",
+ "Can you tell me the BMI if I'm 1.65m tall and weigh
60kg?",
+ "Find BMI for 75kg weight and 1.78m height",
+ "Create me a random number please");
+
+ // Create agents execution environment
+ AgentsExecutionEnvironment agentsEnv =
+ AgentsExecutionEnvironment.getExecutionEnvironment(env);
+
+ // Apply agent to the DataStream and use the prompt itself as the key
+ DataStream<Object> outputStream =
+ agentsEnv
+ .fromDataStream(inputStream, (KeySelector<String,
String>) value -> value)
+ .apply(new AgentWithAzureAI())
+ .toDataStream();
+
+ // Print the results
+ outputStream.print();
+
+ // Execute the pipeline
+ agentsEnv.execute();
+ }
+}
diff --git a/e2e-test/integration-test/pom.xml
b/integrations/chat-models/azureai/pom.xml
similarity index 54%
copy from e2e-test/integration-test/pom.xml
copy to integrations/chat-models/azureai/pom.xml
index c912248..d3cbd05 100644
--- a/e2e-test/integration-test/pom.xml
+++ b/integrations/chat-models/azureai/pom.xml
@@ -19,14 +19,17 @@ under the License.
xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance";
xsi:schemaLocation="http://maven.apache.org/POM/4.0.0
http://maven.apache.org/xsd/maven-4.0.0.xsd";>
<modelVersion>4.0.0</modelVersion>
+
<parent>
<groupId>org.apache.flink</groupId>
- <artifactId>flink-agents-e2e-tests</artifactId>
+ <artifactId>flink-agents-integrations-chat-models</artifactId>
<version>0.2-SNAPSHOT</version>
+ <relativePath>../pom.xml</relativePath>
</parent>
- <artifactId>flink-agents-integration-tests</artifactId>
- <name>Flink Agents : E2E Tests: Integration</name>
+ <artifactId>flink-agents-integrations-chat-models-azureai</artifactId>
+ <name>Flink Agents : Integrations: Chat Models: Azure AI</name>
+ <packaging>jar</packaging>
<dependencies>
<dependency>
@@ -36,38 +39,20 @@ under the License.
</dependency>
<dependency>
<groupId>org.apache.flink</groupId>
- <artifactId>flink-agents-runtime</artifactId>
+ <artifactId>flink-agents-plan</artifactId>
<version>${project.version}</version>
</dependency>
+
<dependency>
- <groupId>org.apache.flink</groupId>
- <artifactId>flink-streaming-java</artifactId>
- <version>${flink.version}</version>
- </dependency>
- <dependency>
- <groupId>org.apache.flink</groupId>
- <artifactId>flink-table-api-java-bridge</artifactId>
- <version>${flink.version}</version>
- </dependency>
- <dependency>
- <groupId>org.apache.flink</groupId>
- <artifactId>flink-table-planner_2.12</artifactId>
- <version>${flink.version}</version>
- </dependency>
- <dependency>
- <groupId>org.apache.flink</groupId>
- <artifactId>flink-clients</artifactId>
- <version>${flink.version}</version>
- </dependency>
- <dependency>
- <groupId>org.apache.flink</groupId>
-
<artifactId>flink-agents-integrations-chat-models-ollama</artifactId>
- <version>${project.version}</version>
+ <groupId>com.azure</groupId>
+ <artifactId>azure-ai-inference</artifactId>
+ <version>1.0.0-beta.5</version>
</dependency>
+
<dependency>
- <groupId>org.apache.flink</groupId>
-
<artifactId>flink-agents-integrations-embedding-models-ollama</artifactId>
- <version>${project.version}</version>
+ <groupId>com.azure</groupId>
+ <artifactId>azure-identity</artifactId>
+ <version>1.13.3</version>
</dependency>
</dependencies>
diff --git
a/integrations/chat-models/azureai/src/main/java/org/apache/flink/agents/integrations/chatmodels/azureai/AzureAIChatModelConnection.java
b/integrations/chat-models/azureai/src/main/java/org/apache/flink/agents/integrations/chatmodels/azureai/AzureAIChatModelConnection.java
new file mode 100644
index 0000000..182c4c1
--- /dev/null
+++
b/integrations/chat-models/azureai/src/main/java/org/apache/flink/agents/integrations/chatmodels/azureai/AzureAIChatModelConnection.java
@@ -0,0 +1,223 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.flink.agents.integrations.chatmodels.azureai;
+
+import com.azure.ai.inference.ChatCompletionsClient;
+import com.azure.ai.inference.ChatCompletionsClientBuilder;
+import com.azure.ai.inference.models.*;
+import com.azure.core.credential.AzureKeyCredential;
+import com.azure.core.util.BinaryData;
+import com.fasterxml.jackson.core.type.TypeReference;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.nimbusds.jose.shaded.gson.Gson;
+import org.apache.flink.agents.api.chat.messages.ChatMessage;
+import org.apache.flink.agents.api.chat.messages.MessageRole;
+import org.apache.flink.agents.api.chat.model.BaseChatModelConnection;
+import org.apache.flink.agents.api.resource.Resource;
+import org.apache.flink.agents.api.resource.ResourceDescriptor;
+import org.apache.flink.agents.api.resource.ResourceType;
+import org.apache.flink.agents.api.tools.Tool;
+
+import java.util.*;
+import java.util.function.BiFunction;
+import java.util.stream.Collectors;
+
+/**
+ * A chat model integration for Azure AI Chat Completions service.
+ *
+ * <p>This implementation adapts the generic Flink Agents chat model interface
to the Azure AI Chat
+ * Completions API.
+ *
+ * <p>See also {@link BaseChatModelConnection} for the common resource
abstractions and lifecycle.
+ *
+ * <p>Example usage:
+ *
+ * <pre>{@code
+ * public class MyAgent extends Agent {
+ * // Register the chat model connection via @ChatModelConnection metadata.
+ * @ChatModelConnection
+ * public static ResourceDesc azureAI() {
+ * return
ResourceDescriptor.Builder.newBuilder(AzureAIChatModelConnection.class.getName())
+ * .addInitialArgument("endpoint", "<your-azure-ai-endpoint>")
+ * .addInitialArgument("apiKey", "<your-azure-ai-api-key>")
+ * .build();
+ * }
+ * }
+ * }</pre>
+ */
+public class AzureAIChatModelConnection extends BaseChatModelConnection {
+
+ private final Gson gson = new Gson();
+
+ private final ChatCompletionsClient client;
+
+ /**
+ * Creates a new AzureAI chat model connection.
+ *
+ * @param descriptor a resource descriptor contains the initial parameters
+ * @param getResource a function to resolve resources (e.g., tools) by
name and type
+ * @throws IllegalArgumentException if endpoint is null or empty
+ */
+ public AzureAIChatModelConnection(
+ ResourceDescriptor descriptor, BiFunction<String, ResourceType,
Resource> getResource) {
+ super(descriptor, getResource);
+
+ String endpoint = descriptor.getArgument("endpoint");
+ String apiKey = descriptor.getArgument("apiKey");
+ if (endpoint == null || endpoint.isEmpty()) {
+ throw new IllegalArgumentException("endpoint should not be null or
empty.");
+ }
+ this.client =
+ new ChatCompletionsClientBuilder()
+ .credential(new AzureKeyCredential(apiKey))
+ .endpoint(endpoint)
+ .buildClient();
+ }
+
+ private List<ChatCompletionsToolDefinition>
convertToAzureAITools(List<Tool> tools) {
+ final ObjectMapper mapper = new ObjectMapper();
+ final List<ChatCompletionsToolDefinition> azureAITools = new
ArrayList<>();
+ try {
+ for (Tool tool : tools) {
+ final Map<String, Object> schema =
+ mapper.readValue(
+ tool.getMetadata().getInputSchema(), new
TypeReference<>() {});
+
+ final FunctionDefinition functionDef =
+ new FunctionDefinition(tool.getName())
+ .setDescription(tool.getDescription())
+ .setParameters(BinaryData.fromObject(schema));
+
+ azureAITools.add(new
ChatCompletionsFunctionToolDefinition(functionDef));
+ }
+ return azureAITools;
+ } catch (Exception e) {
+ throw new RuntimeException(e);
+ }
+ }
+
+ private ChatRequestMessage convertToChatRequestMessage(ChatMessage
message) {
+ final String content = message.getContent();
+ final MessageRole role = message.getRole();
+ final List<Map<String, Object>> toolCalls = message.getToolCalls();
+ final Map<String, Object> extraArgs = message.getExtraArgs();
+ switch (role) {
+ case SYSTEM:
+ return new ChatRequestSystemMessage(content);
+ case USER:
+ return new ChatRequestUserMessage(content);
+ case ASSISTANT:
+ final List<ChatCompletionsToolCall> azureToolCalls =
+ toolCalls != null
+ ? transformToAzureToolCalls(toolCalls)
+ : Collections.emptyList();
+ return new
ChatRequestAssistantMessage(content).setToolCalls(azureToolCalls);
+ case TOOL:
+ String toolCallId =
+ extraArgs != null &&
extraArgs.containsKey("externalId")
+ ? extraArgs.get("externalId").toString()
+ : null;
+ return new
ChatRequestToolMessage(toolCallId).setContent(content);
+ default:
+ throw new IllegalArgumentException("Unsupported role: " +
role);
+ }
+ }
+
+ // the structure of toolCalls should be like the returned value of
Method:convertToAgentsTools
+ private List<ChatCompletionsToolCall> transformToAzureToolCalls(
+ List<Map<String, Object>> toolCalls) {
+ final List<ChatCompletionsToolCall> azureToolCalls = new ArrayList<>();
+ for (Map<String, Object> call : toolCalls) {
+ final String id = (String) call.get("id");
+ final String type = (String) call.get("type");
+
+ if ("function".equals(type)) {
+ final Map<String, Object> functionCall = (Map<String, Object>)
call.get("function");
+ final String functionName = (String) functionCall.get("name");
+ final Map<String, Object> functionArguments =
+ (Map<String, Object>) functionCall.get("arguments");
+ final String functionArgumentsJson =
gson.toJson(functionArguments);
+ ChatCompletionsFunctionToolCall function =
+ new ChatCompletionsFunctionToolCall(
+ id, new FunctionCall(functionName,
functionArgumentsJson));
+ azureToolCalls.add(function);
+ }
+ }
+ return azureToolCalls;
+ }
+
+ @Override
+ public ChatMessage chat(
+ List<ChatMessage> messages, List<Tool> tools, Map<String, Object>
arguments) {
+ try {
+ final List<ChatCompletionsToolDefinition> azureTools =
convertToAzureAITools(tools);
+ final List<ChatRequestMessage> chatMessages =
+ messages.stream()
+ .map(this::convertToChatRequestMessage)
+ .collect(Collectors.toList());
+
+ ChatCompletionsOptions options =
+ new ChatCompletionsOptions(chatMessages)
+ .setModel((String) arguments.get("model"))
+ .setTools(azureTools);
+
+ ChatCompletions completions = client.complete(options);
+ ChatChoice choice = completions.getChoices().get(0);
+ ChatResponseMessage responseMessage = choice.getMessage();
+
+ ChatMessage chatMessage =
ChatMessage.assistant(responseMessage.getContent());
+
+ List<ChatCompletionsToolCall> toolCalls =
responseMessage.getToolCalls();
+ if (toolCalls != null && !toolCalls.isEmpty()) {
+ List<Map<String, Object>> convertedToolCalls =
convertToAgentsTools(toolCalls);
+ chatMessage.setToolCalls(convertedToolCalls);
+ }
+
+ return chatMessage;
+ } catch (Exception e) {
+ throw new RuntimeException(e);
+ }
+ }
+
+ private List<Map<String, Object>> convertToAgentsTools(
+ List<ChatCompletionsToolCall> azureToolCalls) {
+ final List<Map<String, Object>> toolCalls = new
ArrayList<>(azureToolCalls.size());
+ for (ChatCompletionsToolCall toolCall : azureToolCalls) {
+ if (toolCall != null) {
+ final Map<String, Object> call =
+ Map.of(
+ // todo: I don't think the magic name is a
good idea here, need to
+ // unify later (maybe we can consider
standardizing tool call
+ // structure across different LLM integrations)
+ "id", toolCall.getId(),
+ "original_id", toolCall.getId(),
+ "type", toolCall.getType(),
+ "function",
+ Map.of(
+ "name",
toolCall.getFunction().getName(),
+ "arguments",
+ gson.fromJson(
+
toolCall.getFunction()
+
.getArguments(),
+ Map.class)));
+ toolCalls.add(call);
+ }
+ }
+ return toolCalls;
+ }
+}
diff --git
a/integrations/chat-models/azureai/src/main/java/org/apache/flink/agents/integrations/chatmodels/azureai/AzureAIChatModelSetup.java
b/integrations/chat-models/azureai/src/main/java/org/apache/flink/agents/integrations/chatmodels/azureai/AzureAIChatModelSetup.java
new file mode 100644
index 0000000..465cbc8
--- /dev/null
+++
b/integrations/chat-models/azureai/src/main/java/org/apache/flink/agents/integrations/chatmodels/azureai/AzureAIChatModelSetup.java
@@ -0,0 +1,66 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.flink.agents.integrations.chatmodels.azureai;
+
+import org.apache.flink.agents.api.chat.model.BaseChatModelSetup;
+import org.apache.flink.agents.api.resource.Resource;
+import org.apache.flink.agents.api.resource.ResourceDescriptor;
+import org.apache.flink.agents.api.resource.ResourceType;
+
+/**
+ * A chat model integration for Azure AI Chat Completions service.
+ *
+ * <p>This implementation adapts the generic Flink Agents chat model interface
to the Azure AI Chat
+ * Completions API.
+ *
+ * <p>See also {@link BaseChatModelSetup} for the common resource abstractions
and lifecycle.
+ *
+ * <p>Example usage:
+ *
+ * <pre>{@code
+ * public class MyAgent extends Agent {
+ * // Register the chat model setup via @ChatModelSetup metadata.
+ * @ChatModelSetup
+ * public static ResourceDesc azureAI() {
+ * return
ResourceDescriptor.Builder.newBuilder(AzureAIChatModelSetup.class.getName())
+ * .addInitialArgument("model", "<your-azure-ai-model-name>")
+ * .addInitialArgument("prompt", "<your-prompt-template>")
+ * .addInitialArgument("tools", "<your-tool-list>")
+ * .addInitialArgument("key", "<your-azure-ai-key>")
+ * .addInitialArgument("endpoint", "<your-azure-ai-endpoint>")
+ * .build();
+ * }
+ * }
+ * }</pre>
+ */
+public class AzureAIChatModelSetup extends BaseChatModelSetup {
+
+ public AzureAIChatModelSetup(
+ ResourceDescriptor descriptor,
+ java.util.function.BiFunction<String, ResourceType, Resource>
getResource) {
+ super(descriptor, getResource);
+ }
+
+ // For any other specific parameters, please refer to
ChatCompletionsOptions
+ @Override
+ public java.util.Map<String, Object> getParameters() {
+ java.util.Map<String, Object> params = new java.util.HashMap<>();
+ params.put("model", model);
+ return params;
+ }
+}
diff --git a/integrations/chat-models/pom.xml b/integrations/chat-models/pom.xml
index bca4be1..dafeae9 100644
--- a/integrations/chat-models/pom.xml
+++ b/integrations/chat-models/pom.xml
@@ -31,6 +31,7 @@ under the License.
<packaging>pom</packaging>
<modules>
+ <module>azureai</module>
<module>ollama</module>
</modules>
