wenjin272 commented on code in PR #290: URL: https://github.com/apache/flink-agents/pull/290#discussion_r2480331186
########## integrations/chat-models/azureai/src/main/java/org/apache/flink/agents/integrations/chatmodels/azureai/AzureAIChatModelSetup.java: ########## @@ -0,0 +1,99 @@ +/* + * 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 { + + private final String model; + private final String prompt; + private final java.util.List<String> tools; + private final String key; + private final String endpoint; + + public AzureAIChatModelSetup( + ResourceDescriptor descriptor, + java.util.function.BiFunction<String, ResourceType, Resource> getResource) { + super(descriptor, getResource); + this.model = descriptor.getArgument("model"); + this.prompt = descriptor.getArgument("prompt"); + this.tools = (java.util.List<String>) descriptor.getArgument("tools"); + this.key = descriptor.getArgument("key"); + this.endpoint = descriptor.getArgument("endpoint"); Review Comment: `key` and `endpoint` should be the arguments of `AzureAIChatModelConnection`. Different `AzureAIChatModelSetup` can share the same `AzureAIChatModelConnection`. ########## integrations/chat-models/azureai/src/main/java/org/apache/flink/agents/integrations/chatmodels/azureai/AzureAIChatModelConnection.java: ########## @@ -0,0 +1,183 @@ +/* + * 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 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 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(); + } + + @SuppressWarnings("unchecked") Review Comment: My idea show this is a redundant suppression. ########## integrations/chat-models/azureai/src/main/java/org/apache/flink/agents/integrations/chatmodels/azureai/AzureAIChatModelSetup.java: ########## @@ -0,0 +1,99 @@ +/* + * 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 { + + private final String model; + private final String prompt; + private final java.util.List<String> tools; + private final String key; + private final String endpoint; + + public AzureAIChatModelSetup( + ResourceDescriptor descriptor, + java.util.function.BiFunction<String, ResourceType, Resource> getResource) { + super(descriptor, getResource); + this.model = descriptor.getArgument("model"); + this.prompt = descriptor.getArgument("prompt"); + this.tools = (java.util.List<String>) descriptor.getArgument("tools"); Review Comment: `prompt` and `tools` will be retrieved in constructor of base class `BaseChatModelSetup`, no need to process here again. ########## integrations/chat-models/azureai/src/main/java/org/apache/flink/agents/integrations/chatmodels/azureai/AzureAIChatModelSetup.java: ########## @@ -0,0 +1,99 @@ +/* + * 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 { + + private final String model; + private final String prompt; + private final java.util.List<String> tools; + private final String key; + private final String endpoint; + + public AzureAIChatModelSetup( + ResourceDescriptor descriptor, + java.util.function.BiFunction<String, ResourceType, Resource> getResource) { + super(descriptor, getResource); + this.model = descriptor.getArgument("model"); + this.prompt = descriptor.getArgument("prompt"); + this.tools = (java.util.List<String>) descriptor.getArgument("tools"); + this.key = descriptor.getArgument("key"); + this.endpoint = descriptor.getArgument("endpoint"); + } + + public AzureAIChatModelSetup( + String model, + String prompt, + java.util.List<String> tools, + String key, + String endpoint, + java.util.function.BiFunction<String, ResourceType, Resource> getResource) { + this( + new ResourceDescriptor( + AzureAIChatModelSetup.class.getName(), + java.util.Map.of( + "model", model, + "prompt", prompt, + "tools", tools, + "key", key, + "endpoint", endpoint)), + getResource); + } + + @Override + public java.util.Map<String, Object> getParameters() { + java.util.Map<String, Object> params = new java.util.HashMap<>(); + params.put("model", model); + params.put("prompt", prompt); + params.put("tools", tools); Review Comment: Base class `BaseChatModelSetup ` will process the `prompt` and `tools` directly, no need to pass them as parameters here. ########## integrations/chat-models/azureai/src/main/java/org/apache/flink/agents/integrations/chatmodels/azureai/AzureAIChatModelConnection.java: ########## @@ -0,0 +1,183 @@ +/* + * 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 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 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(); + } + + @SuppressWarnings("unchecked") + 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(); + switch (role) { + case SYSTEM: + return new ChatRequestSystemMessage(content); + case USER: + return new ChatRequestUserMessage(content); + case ASSISTANT: + return new ChatRequestAssistantMessage(content); + case TOOL: + return new ChatRequestToolMessage(content); Review Comment: 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); ... } ``` ########## integrations/chat-models/azureai/src/main/java/org/apache/flink/agents/integrations/chatmodels/azureai/AzureAIChatModelConnection.java: ########## @@ -0,0 +1,183 @@ +/* + * 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 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 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(); + } + + @SuppressWarnings("unchecked") + 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(); + switch (role) { + case SYSTEM: + return new ChatRequestSystemMessage(content); + case USER: + return new ChatRequestUserMessage(content); + case ASSISTANT: + return new ChatRequestAssistantMessage(content); Review Comment: Here, only the `content` of `ChatRequestAssistantMessage` is set, but no `toolCalls` are set. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
