[OpenAI] Replace NonAzureOpenAIKeyCredential by KeyCredential (Azure#35754)

Author: mssfang

sdk/openai/azure-ai-openai/CHANGELOG.md

@@ -8,6 +8,8 @@
 
 ### Breaking Changes
 
+- Replaced usage of class `NonAzureOpenAIKeyCredential` by Azure Core class `KeyCredential`.
+
 ### Bugs Fixed
 
 ### Other Changes

sdk/openai/azure-ai-openai/README.md

@@ -35,12 +35,12 @@ If you want to see the full code for these snippets check out our [samples folde
 
 ### Adding the package to your product
 
-[//]: # ({x-version-update-start;com.azure:azure-ai-openai;current})
+[//]: # ({x-version-update-start;com.azure:azure-ai-openai;dependency})
 ```xml
 <dependency>
     <groupId>com.azure</groupId>
     <artifactId>azure-ai-openai</artifactId>
-    <version>1.0.0-beta.2</version>
+    <version>1.0.0-beta.3</version>
 </dependency>
 ```
 [//]: # ({x-version-update-end})
@@ -77,14 +77,14 @@ The SDK also supports operating against the public non-Azure OpenAI. The respons
 
 ```java readme-sample-createNonAzureOpenAISyncClientApiKey
 OpenAIClient client = new OpenAIClientBuilder()
-    .credential(new NonAzureOpenAIKeyCredential("{openai-secret-key}"))
+    .credential(new KeyCredential("{openai-secret-key}"))
     .buildClient();
 ```
 or 
 
 ```java readme-sample-createNonAzureOpenAIAsyncClientApiKey
 OpenAIAsyncClient client = new OpenAIClientBuilder()
-    .credential(new NonAzureOpenAIKeyCredential("{openai-secret-key}"))
+    .credential(new KeyCredential("{openai-secret-key}"))
     .buildAsyncClient();
 ```
 

sdk/openai/azure-ai-openai/src/main/java/com/azure/ai/openai/OpenAIClientBuilder.java

@@ -3,10 +3,10 @@
 // Code generated by Microsoft (R) AutoRest Code Generator.
 package com.azure.ai.openai;
 
+import static com.azure.ai.openai.implementation.NonAzureOpenAIClientImpl.OPEN_AI_ENDPOINT;
+
 import com.azure.ai.openai.implementation.NonAzureOpenAIClientImpl;
-import com.azure.ai.openai.implementation.NonAzureOpenAIKeyCredentialPolicy;
 import com.azure.ai.openai.implementation.OpenAIClientImpl;
-import com.azure.ai.openai.models.NonAzureOpenAIKeyCredential;
 import com.azure.core.annotation.Generated;
 import com.azure.core.annotation.ServiceClientBuilder;
 import com.azure.core.client.traits.AzureKeyCredentialTrait;
@@ -15,6 +15,7 @@
 import com.azure.core.client.traits.HttpTrait;
 import com.azure.core.client.traits.TokenCredentialTrait;
 import com.azure.core.credential.AzureKeyCredential;
+import com.azure.core.credential.KeyCredential;
 import com.azure.core.credential.TokenCredential;
 import com.azure.core.http.HttpClient;
 import com.azure.core.http.HttpHeaders;
@@ -24,13 +25,13 @@
 import com.azure.core.http.policy.AddDatePolicy;
 import com.azure.core.http.policy.AddHeadersFromContextPolicy;
 import com.azure.core.http.policy.AddHeadersPolicy;
-import com.azure.core.http.policy.AzureKeyCredentialPolicy;
 import com.azure.core.http.policy.BearerTokenAuthenticationPolicy;
 import com.azure.core.http.policy.CookiePolicy;
 import com.azure.core.http.policy.HttpLogOptions;
 import com.azure.core.http.policy.HttpLoggingPolicy;
 import com.azure.core.http.policy.HttpPipelinePolicy;
 import com.azure.core.http.policy.HttpPolicyProviders;
+import com.azure.core.http.policy.KeyCredentialPolicy;
 import com.azure.core.http.policy.RequestIdPolicy;
 import com.azure.core.http.policy.RetryOptions;
 import com.azure.core.http.policy.RetryPolicy;
@@ -182,23 +183,22 @@ public OpenAIClientBuilder credential(TokenCredential tokenCredential) {
     @Generated private AzureKeyCredential azureKeyCredential;
 
     /** {@inheritDoc}. */
-    @Generated
     @Override
     public OpenAIClientBuilder credential(AzureKeyCredential azureKeyCredential) {
-        this.azureKeyCredential = azureKeyCredential;
-        return this;
+        return this.credential((KeyCredential) azureKeyCredential);
     }
 
-    private NonAzureOpenAIKeyCredential nonAzureOpenAIKeyCredential;
+    /** The KeyCredential used for OpenAi authentication. It could be either of Azure or Non-Azure OpenAI API key. */
+    private KeyCredential keyCredential;
 
     /**
-     * The NonAzureOpenAiKeyCredential used for public OpenAi authentication.
+     * The KeyCredential used for OpenAi authentication. It could be either of Azure or Non-Azure OpenAI API key.
      *
-     * @param nonAzureOpenAIKeyCredential The credential for non-azure public OpenAI authenticaton.
+     * @param keyCredential The credential for OpenAI authentication.
      * @return the object itself.
      */
-    public OpenAIClientBuilder credential(NonAzureOpenAIKeyCredential nonAzureOpenAIKeyCredential) {
-        this.nonAzureOpenAIKeyCredential = nonAzureOpenAIKeyCredential;
+    public OpenAIClientBuilder credential(KeyCredential keyCredential) {
+        this.keyCredential = keyCredential;
         return this;
     }
 
@@ -268,7 +268,6 @@ private OpenAIClientImpl buildInnerClient() {
         return client;
     }
 
-    @Generated
     private HttpPipeline createHttpPipeline() {
         Configuration buildConfiguration =
                 (configuration == null) ? Configuration.getGlobalConfiguration() : configuration;
@@ -292,8 +291,12 @@ private HttpPipeline createHttpPipeline() {
         HttpPolicyProviders.addBeforeRetryPolicies(policies);
         policies.add(ClientBuilderUtil.validateAndGetRetryPolicy(retryPolicy, retryOptions, new RetryPolicy()));
         policies.add(new AddDatePolicy());
-        if (azureKeyCredential != null) {
-            policies.add(new AzureKeyCredentialPolicy("api-key", azureKeyCredential));
+        policies.add(new CookiePolicy());
+        if (keyCredential != null) {
+            policies.add(
+                    useNonAzureOpenAIService()
+                            ? new KeyCredentialPolicy("Authorization", keyCredential, "Bearer")
+                            : new KeyCredentialPolicy("api-key", keyCredential));
         }
         if (tokenCredential != null) {
             policies.add(new BearerTokenAuthenticationPolicy(tokenCredential, DEFAULT_SCOPES));
@@ -313,63 +316,21 @@ private HttpPipeline createHttpPipeline() {
     }
 
     private NonAzureOpenAIClientImpl buildInnerNonAzureOpenAIClient() {
-        HttpPipeline localPipeline = (pipeline != null) ? pipeline : createHttpPipelineNonAzureOpenAI();
+        HttpPipeline localPipeline = (pipeline != null) ? pipeline : createHttpPipeline();
         NonAzureOpenAIClientImpl client =
                 new NonAzureOpenAIClientImpl(localPipeline, JacksonAdapter.createDefaultSerializerAdapter());
         return client;
     }
 
-    private HttpPipeline createHttpPipelineNonAzureOpenAI() {
-        Configuration buildConfiguration =
-                (configuration == null) ? Configuration.getGlobalConfiguration() : configuration;
-        HttpLogOptions localHttpLogOptions = this.httpLogOptions == null ? new HttpLogOptions() : this.httpLogOptions;
-        ClientOptions localClientOptions = this.clientOptions == null ? new ClientOptions() : this.clientOptions;
-        List<HttpPipelinePolicy> policies = new ArrayList<>();
-        String clientName = PROPERTIES.getOrDefault(SDK_NAME, "UnknownName");
-        String clientVersion = PROPERTIES.getOrDefault(SDK_VERSION, "UnknownVersion");
-        String applicationId = CoreUtils.getApplicationId(localClientOptions, localHttpLogOptions);
-        policies.add(new UserAgentPolicy(applicationId, clientName, clientVersion, buildConfiguration));
-        policies.add(new RequestIdPolicy());
-        policies.add(new AddHeadersFromContextPolicy());
-        HttpHeaders headers = new HttpHeaders();
-        localClientOptions.getHeaders().forEach(header -> headers.set(header.getName(), header.getValue()));
-        if (headers.getSize() > 0) {
-            policies.add(new AddHeadersPolicy(headers));
-        }
-        this.pipelinePolicies.stream()
-                .filter(p -> p.getPipelinePosition() == HttpPipelinePosition.PER_CALL)
-                .forEach(p -> policies.add(p));
-        HttpPolicyProviders.addBeforeRetryPolicies(policies);
-        policies.add(ClientBuilderUtil.validateAndGetRetryPolicy(retryPolicy, retryOptions, new RetryPolicy()));
-        policies.add(new AddDatePolicy());
-        policies.add(new CookiePolicy());
-        if (nonAzureOpenAIKeyCredential != null) {
-            policies.add(new NonAzureOpenAIKeyCredentialPolicy("Authorization", nonAzureOpenAIKeyCredential, "Bearer"));
-        }
-        this.pipelinePolicies.stream()
-                .filter(p -> p.getPipelinePosition() == HttpPipelinePosition.PER_RETRY)
-                .forEach(p -> policies.add(p));
-        HttpPolicyProviders.addAfterRetryPolicies(policies);
-        policies.add(new HttpLoggingPolicy(httpLogOptions));
-        HttpPipeline httpPipeline =
-                new HttpPipelineBuilder()
-                        .policies(policies.toArray(new HttpPipelinePolicy[0]))
-                        .httpClient(httpClient)
-                        .clientOptions(localClientOptions)
-                        .build();
-        return httpPipeline;
-    }
-
     /**
      * Builds an instance of OpenAIAsyncClient class.
      *
      * @return an instance of OpenAIAsyncClient.
      */
     public OpenAIAsyncClient buildAsyncClient() {
-        if (nonAzureOpenAIKeyCredential != null) {
-            return new OpenAIAsyncClient(buildInnerNonAzureOpenAIClient());
-        }
-        return new OpenAIAsyncClient(buildInnerClient());
+        return useNonAzureOpenAIService()
+                ? new OpenAIAsyncClient(buildInnerNonAzureOpenAIClient())
+                : new OpenAIAsyncClient(buildInnerClient());
     }
 
     /**
@@ -378,11 +339,18 @@ public OpenAIAsyncClient buildAsyncClient() {
      * @return an instance of OpenAIClient.
      */
     public OpenAIClient buildClient() {
-        if (nonAzureOpenAIKeyCredential != null) {
-            return new OpenAIClient(buildInnerNonAzureOpenAIClient());
-        }
-        return new OpenAIClient(buildInnerClient());
+        return useNonAzureOpenAIService()
+                ? new OpenAIClient(buildInnerNonAzureOpenAIClient())
+                : new OpenAIClient(buildInnerClient());
     }
 
     private static final ClientLogger LOGGER = new ClientLogger(OpenAIClientBuilder.class);
+
+    /**
+     * OpenAI service can be used by either not setting the endpoint or by setting the endpoint to start with
+     * "https://api.openai.com/"
+     */
+    private boolean useNonAzureOpenAIService() {
+        return endpoint == null || endpoint.startsWith(OPEN_AI_ENDPOINT);
+    }
 }

sdk/openai/azure-ai-openai/src/main/java/com/azure/ai/openai/implementation/NonAzureOpenAIKeyCredentialPolicy.java

@@ -7,7 +7,7 @@
 import com.azure.ai.openai.OpenAIClientBuilder;
 import com.azure.ai.openai.models.Completions;
 import com.azure.ai.openai.models.CompletionsOptions;
-import com.azure.ai.openai.models.NonAzureOpenAIKeyCredential;
+import com.azure.core.credential.KeyCredential;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -25,7 +25,7 @@ public static void main(String[] args) {
         String apiKey = "{non-azure-open-ai-api-key}";
 
         OpenAIClientBuilder builder = new OpenAIClientBuilder()
-            .credential(new NonAzureOpenAIKeyCredential(apiKey));
+            .credential(new KeyCredential(apiKey));
 
         OpenAIClient client = builder.buildClient();