cleaned up CM workspace APIs

sdk/provisioning/Azure.Provisioning.CloudMachine/api/Azure.Provisioning.CloudMachine.netstandard2.0.cs

@@ -6,17 +6,15 @@ public partial class CloudMachineClient : Azure.CloudMachine.CloudMachineWorkspa
         public Azure.CloudMachine.MessagingServices Messaging { get { throw null; } }
         public Azure.CloudMachine.StorageServices Storage { get { throw null; } }
     }
-    public partial class CloudMachineWorkspace : Azure.Core.WorkspaceClient
+    public partial class CloudMachineWorkspace : Azure.Core.ClientWorkspace
     {
         public CloudMachineWorkspace(Azure.Core.TokenCredential? credential = null, Microsoft.Extensions.Configuration.IConfiguration? configuration = null) { }
         [System.ComponentModel.EditorBrowsableAttribute(System.ComponentModel.EditorBrowsableState.Never)]
-        public override Azure.Core.TokenCredential Credential { get { throw null; } }
-        [System.ComponentModel.EditorBrowsableAttribute(System.ComponentModel.EditorBrowsableState.Never)]
         public string Id { get { throw null; } }
         [System.ComponentModel.EditorBrowsableAttribute(System.ComponentModel.EditorBrowsableState.Never)]
         public override bool Equals(object? obj) { throw null; }
         [System.ComponentModel.EditorBrowsableAttribute(System.ComponentModel.EditorBrowsableState.Never)]
-        public override Azure.Core.ClientConfiguration? GetConfiguration(string clientId, string? instanceId = null) { throw null; }
+        public override Azure.Core.ClientConnectionOptions GetConnectionOptions(System.Type clientType, string? instanceId = null) { throw null; }
         [System.ComponentModel.EditorBrowsableAttribute(System.ComponentModel.EditorBrowsableState.Never)]
         public override int GetHashCode() { throw null; }
         [System.ComponentModel.EditorBrowsableAttribute(System.ComponentModel.EditorBrowsableState.Never)]
@@ -46,30 +44,34 @@ namespace Azure.Core
     public partial class ClientCache
     {
         public ClientCache() { }
-        public T Get<T>(string id, System.Func<T> value) where T : class { throw null; }
+        public T Get<T>(System.Func<T> value, string? id = null) where T : class { throw null; }
+    }
+    public enum ClientConnectionKind
+    {
+        EntraId = 0,
+        ApiKey = 1,
+        OutOfBand = 2,
     }
     [System.Runtime.InteropServices.StructLayoutAttribute(System.Runtime.InteropServices.LayoutKind.Sequential)]
-    public readonly partial struct ClientConfiguration
+    public readonly partial struct ClientConnectionOptions
     {
         private readonly object _dummy;
         private readonly int _dummyPrimitive;
-        public ClientConfiguration(string endpoint, string? apiKey = null) { throw null; }
-        public string? ApiKey { get { throw null; } }
-        public Azure.Core.CredentialType CredentialType { get { throw null; } }
-        public string Endpoint { get { throw null; } }
+        public ClientConnectionOptions(string subclientId) { throw null; }
+        public ClientConnectionOptions(System.Uri endpoint, Azure.Core.TokenCredential credential) { throw null; }
+        public ClientConnectionOptions(System.Uri endpoint, string apiKey) { throw null; }
+        public string? ApiKeyCredential { get { throw null; } }
+        public Azure.Core.ClientConnectionKind ConnectionKind { get { throw null; } }
+        public System.Uri? Endpoint { get { throw null; } }
+        public string? Id { get { throw null; } }
+        public Azure.Core.TokenCredential? TokenCredential { get { throw null; } }
     }
-    public enum CredentialType
+    public abstract partial class ClientWorkspace
     {
-        EntraId = 0,
-        ApiKey = 1,
-    }
-    public abstract partial class WorkspaceClient
-    {
-        protected WorkspaceClient() { }
-        public abstract Azure.Core.TokenCredential Credential { get; }
+        protected ClientWorkspace() { }
         [System.ComponentModel.EditorBrowsableAttribute(System.ComponentModel.EditorBrowsableState.Never)]
         public Azure.Core.ClientCache Subclients { get { throw null; } }
-        public abstract Azure.Core.ClientConfiguration? GetConfiguration(string clientId, string? instanceId = null);
+        public abstract Azure.Core.ClientConnectionOptions GetConnectionOptions(System.Type clientType, string? instanceId = null);
     }
 }
 namespace Azure.Provisioning.CloudMachine
@@ -96,7 +98,7 @@ namespace Azure.Provisioning.CloudMachine.KeyVault
 {
     public static partial class KeyVaultExtensions
     {
-        public static Azure.Security.KeyVault.Secrets.SecretClient GetKeyVaultSecretsClient(this Azure.Core.WorkspaceClient workspace) { throw null; }
+        public static Azure.Security.KeyVault.Secrets.SecretClient GetKeyVaultSecretsClient(this Azure.Core.ClientWorkspace workspace) { throw null; }
     }
     public partial class KeyVaultFeature : Azure.Provisioning.CloudMachine.CloudMachineFeature
     {
@@ -107,15 +109,15 @@ public override void AddTo(Azure.Provisioning.CloudMachine.CloudMachineInfrastru
 }
 namespace Azure.Provisioning.CloudMachine.OpenAI
 {
+    public static partial class AzureOpenAIExtensions
+    {
+        public static OpenAI.Chat.ChatClient GetOpenAIChatClient(this Azure.Core.ClientWorkspace workspace) { throw null; }
+    }
     public partial class OpenAIFeature : Azure.Provisioning.CloudMachine.CloudMachineFeature
     {
         public OpenAIFeature(string model, string modelVersion) { }
         public string Model { get { throw null; } }
         public string ModelVersion { get { throw null; } }
         public override void AddTo(Azure.Provisioning.CloudMachine.CloudMachineInfrastructure cloudMachine) { }
     }
-    public static partial class OpenAIFeatureExtensions
-    {
-        public static OpenAI.Chat.ChatClient GetOpenAIChatClient(this Azure.Core.WorkspaceClient workspace) { throw null; }
-    }
 }

sdk/provisioning/Azure.Provisioning.CloudMachine/src/AzureSdkExtensions/KeyVaultFeature.cs

@@ -64,19 +64,13 @@ public override void AddTo(CloudMachineInfrastructure infrastructure)
 
 public static class KeyVaultExtensions
 {
-    public static SecretClient GetKeyVaultSecretsClient(this WorkspaceClient workspace)
+    public static SecretClient GetKeyVaultSecretsClient(this ClientWorkspace workspace)
     {
-        ClientConfiguration? connectionMaybe = workspace.GetConfiguration(typeof(SecretClient).FullName);
-        if (connectionMaybe == null)
+        ClientConnectionOptions connection = workspace.GetConnectionOptions(typeof(SecretClient));
+        if (connection.ConnectionKind == ClientConnectionKind.EntraId)
         {
-            throw new Exception("Connection not found");
+            return new(connection.Endpoint, connection.TokenCredential);
         }
-
-        ClientConfiguration connection = connectionMaybe.Value;
-        if (connection.CredentialType == CredentialType.EntraId)
-        {
-            return new(new Uri(connection.Endpoint), workspace.Credential);
-        }
-        throw new Exception("ApiKey not supported");
+        throw new Exception("API key not supported");
     }
 }

sdk/provisioning/Azure.Provisioning.CloudMachine/src/AzureSdkExtensions/OpenAIFeature.cs

@@ -3,10 +3,7 @@
 
 using System;
 using System.ClientModel;
-using System.ClientModel.Primitives;
-using System.Diagnostics.Contracts;
 using Azure.AI.OpenAI;
-using Azure.CloudMachine;
 using Azure.Core;
 using Azure.Provisioning.Authorization;
 using Azure.Provisioning.CognitiveServices;
@@ -62,44 +59,36 @@ public override void AddTo(CloudMachineInfrastructure cloudMachine)
     }
 }
 
-public static class OpenAIFeatureExtensions
+public static class AzureOpenAIExtensions
 {
-    public static ChatClient GetOpenAIChatClient(this WorkspaceClient workspace)
+    public static ChatClient GetOpenAIChatClient(this ClientWorkspace workspace)
     {
-        string chatClientId = typeof(ChatClient).FullName;
-
-        ChatClient client = workspace.Subclients.Get(chatClientId, () =>
+        ChatClient chatClient = workspace.Subclients.Get(() =>
         {
-            string azureOpenAIClientId = typeof(AzureOpenAIClient).FullName;
-
-            AzureOpenAIClient aoia = workspace.Subclients.Get(azureOpenAIClientId, () =>
-            {
-                ClientConfiguration? connectionMaybe = workspace.GetConfiguration(typeof(AzureOpenAIClient).FullName);
-                if (connectionMaybe == null) throw new Exception("Connection not found");
+            AzureOpenAIClient aoiaClient = workspace.Subclients.Get(() => CreateAzureOpenAIClient(workspace));
+            return workspace.CreateChatClient(aoiaClient);
+        });
 
-                ClientConfiguration connection = connectionMaybe.Value;
-                Uri endpoint = new(connection.Endpoint);
-                var clientOptions = new AzureOpenAIClientOptions();
-                if (connection.CredentialType == CredentialType.EntraId)
-                {
-                    AzureOpenAIClient aoai = new(endpoint, workspace.Credential, clientOptions);
-                    return aoai;
-                }
-                else
-                {
-                    AzureOpenAIClient aoai = new(endpoint, new ApiKeyCredential(connection.ApiKey!), clientOptions);
-                    return aoai;
-                }
-            });
+        return chatClient;
+    }
 
-            string azureOpenAIChatClientId = typeof(ChatClient).FullName;
-            ClientConfiguration? connectionMaybe = workspace.GetConfiguration(azureOpenAIChatClientId);
-            if (connectionMaybe == null) throw new Exception("Connection not found");
-            var connection = connectionMaybe.Value;
-            ChatClient chat = aoia.GetChatClient(connection.Endpoint);
-            return chat;
-        });
+    private static AzureOpenAIClient CreateAzureOpenAIClient(this ClientWorkspace workspace)
+    {
+        ClientConnectionOptions connection = workspace.GetConnectionOptions(typeof(AzureOpenAIClient));
+        if (connection.ConnectionKind == ClientConnectionKind.EntraId)
+        {
+            return new(connection.Endpoint, connection.TokenCredential);
+        }
+        else
+        {
+            return  new(connection.Endpoint, new ApiKeyCredential(connection.ApiKeyCredential!));
+        }
+    }
 
-        return client;
+    private static ChatClient CreateChatClient(this ClientWorkspace workspace, AzureOpenAIClient client)
+    {
+        ClientConnectionOptions connection = workspace.GetConnectionOptions(typeof(ChatClient));
+        ChatClient chat = client.GetChatClient(connection.Id);
+        return chat;
     }
 }

sdk/provisioning/Azure.Provisioning.CloudMachine/src/Core/ClientCache.cs

@@ -9,25 +9,24 @@ namespace Azure.Core;
 // TODO: this is a very demo implementation. We need to do better
 public class ClientCache
 {
-    private readonly Dictionary<string, object> _clients = new Dictionary<string, object>();
+    private readonly Dictionary<(Type, string?), object> _clients = new Dictionary<(Type, string?), object>();
 
-    // TODO: consider uisng ICLientCreator instead of Func
-    public T Get<T>(string id, Func<T> value) where T: class
+    public T Get<T>(Func<T> value, string? id = default) where T: class
     {
+        var client = (typeof(T), id);
         lock (_clients)
         {
-            if (_clients.TryGetValue(id, out object cached))
+            if (_clients.TryGetValue(client, out object cached))
             {
-                T client = (T)cached;
-                return client;
+                return (T)cached;
             }
 
             if (_clients.Count > 100)
             {
                 GC();
             }
             T created = value();
-            _clients.Add(id, created);
+            _clients.Add(client, created);
             return created;
         }
 

sdk/provisioning/Azure.Provisioning.CloudMachine/src/Core/ClientWorkspace.cs

@@ -0,0 +1,50 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+using System;
+using System.ComponentModel;
+
+namespace Azure.Core;
+
+public abstract class ClientWorkspace
+{
+    public abstract ClientConnectionOptions GetConnectionOptions(Type clientType, string? instanceId = default);
+
+    [EditorBrowsable(EditorBrowsableState.Never)]
+    public ClientCache Subclients { get; } = new ClientCache();
+}
+
+public readonly struct ClientConnectionOptions
+{
+    public ClientConnectionOptions(Uri endpoint, string apiKey)
+    {
+        Endpoint = endpoint;
+        ApiKeyCredential = apiKey;
+        ConnectionKind = ClientConnectionKind.ApiKey;
+    }
+    public ClientConnectionOptions(Uri endpoint, TokenCredential credential)
+    {
+        Endpoint = endpoint;
+        TokenCredential = credential;
+        ConnectionKind = ClientConnectionKind.EntraId;
+    }
+    public ClientConnectionOptions(string subclientId)
+    {
+        Id = subclientId;
+        ConnectionKind = ClientConnectionKind.OutOfBand;
+    }
+
+    public ClientConnectionKind ConnectionKind { get; }
+
+    public Uri? Endpoint { get; }
+    public string? Id { get; }
+    public string? ApiKeyCredential { get; }
+    public TokenCredential? TokenCredential { get; }
+}
+
+public enum ClientConnectionKind
+{
+    EntraId,
+    ApiKey,
+    OutOfBand
+}

sdk/provisioning/Azure.Provisioning.CloudMachine/src/OFX/CloudMachineWorkspace.cs

@@ -10,17 +10,16 @@
 
 namespace Azure.CloudMachine;
 
-public class CloudMachineWorkspace : WorkspaceClient
+public class CloudMachineWorkspace : ClientWorkspace
 {
-    [EditorBrowsable(EditorBrowsableState.Never)]
-    public string Id { get; }
-
-    [EditorBrowsable(EditorBrowsableState.Never)]
-    public override TokenCredential Credential { get; } = new ChainedTokenCredential(
+    private TokenCredential Credential { get; } = new ChainedTokenCredential(
         new AzureCliCredential(),
         new AzureDeveloperCliCredential()
     );
 
+    [EditorBrowsable(EditorBrowsableState.Never)]
+    public string Id { get; }
+
     [SuppressMessage("Usage", "AZC0007:DO provide a minimal constructor that takes only the parameters required to connect to the service.", Justification = "<Pending>")]
     public CloudMachineWorkspace(TokenCredential? credential = default, IConfiguration? configuration = default)
     {
@@ -45,29 +44,25 @@ public CloudMachineWorkspace(TokenCredential? credential = default, IConfigurati
     }
 
     [EditorBrowsable(EditorBrowsableState.Never)]
-    public override ClientConfiguration? GetConfiguration(string clientId, string? instanceId = default)
+    public override ClientConnectionOptions GetConnectionOptions(Type clientType, string? instanceId = default)
     {
+        string clientId = clientType.FullName;
         switch (clientId)
         {
             case "Azure.Security.KeyVault.Secrets.SecretClient":
-                return new ClientConfiguration($"https://{this.Id}.vault.azure.net/");
+                return new ClientConnectionOptions(new($"https://{this.Id}.vault.azure.net/"), Credential);
             case "Azure.Messaging.ServiceBus.ServiceBusClient":
-                return new ClientConfiguration($"{this.Id}.servicebus.windows.net");
+                return new ClientConnectionOptions(new($"{this.Id}.servicebus.windows.net"), Credential);
             case "Azure.Messaging.ServiceBus.ServiceBusSender":
                 if (instanceId == default) instanceId = "cm_default_topic_sender";
-                return new ClientConfiguration(instanceId);
+                return new ClientConnectionOptions(instanceId);
             case "Azure.Storage.Blobs.BlobContainerClient":
                 if (instanceId == default) instanceId = "default";
-                return new ClientConfiguration($"https://{this.Id}.blob.core.windows.net/{instanceId}");
+                return new ClientConnectionOptions(new($"https://{this.Id}.blob.core.windows.net/{instanceId}"), Credential);
             case "Azure.AI.OpenAI.AzureOpenAIClient":
-                string endpoint = $"https://{this.Id}.openai.azure.com";
-                string? key = null; // Environment.GetEnvironmentVariable("openai_cm_key");
-                if (key != null)
-                    return new ClientConfiguration(endpoint, key);
-                else
-                    return new ClientConfiguration(endpoint);
+                return new ClientConnectionOptions(new($"https://{this.Id}.openai.azure.com"), Credential);
             case "OpenAI.Chat.ChatClient":
-                return new ClientConfiguration(this.Id);
+                return new ClientConnectionOptions(Id);
             default:
                 throw new Exception($"unknown client {clientId}");
         }