FIELDENG-589 EntraId support for Vectorizers (#522)

Entra ID support for Vectorizers

Author: slorello89

src/Redis.OM.Vectorizers/AzureOpenAIVectorizer.cs

@@ -1,6 +1,9 @@
 
+using System.Net.Http.Headers;
 using System.Net.Http.Json;
 using System.Text.Json;
+using Azure.Core;
+using Azure.Identity;
 using Redis.OM.Contracts;
 using Redis.OM.Modeling;
 
@@ -11,7 +14,8 @@ namespace Redis.OM.Vectorizers;
 /// </summary>
 public class AzureOpenAIVectorizer : IVectorizer<string>
 {
-    private readonly string _apiKey;
+    private readonly TokenCredential? _tokenCredential;
+    private readonly string? _apiKey;
     private readonly string _resourceName;
     private readonly string _deploymentName;
 
@@ -28,6 +32,22 @@ public AzureOpenAIVectorizer(string apiKey, string resourceName, string deployme
         _resourceName = resourceName;
         _deploymentName = deploymentName;
         Dim = dim;
+        _tokenCredential = new DefaultAzureCredential();
+    }
+    
+    /// <summary>
+    /// Initialize vectorizer
+    /// </summary>
+    /// <param name="resourceName">The Azure Resource's name</param>
+    /// <param name="deploymentName">The Azure deployment name</param>
+    /// <param name="dim">The dimensions of the model addressed by this resource/deployment.</param>
+    public AzureOpenAIVectorizer(string resourceName, string deploymentName, int dim)
+    {
+        _apiKey = null;
+        _resourceName = resourceName;
+        _deploymentName = deploymentName;
+        Dim = dim;
+        _tokenCredential = new DefaultAzureCredential();
     }
 
     /// <inheritdoc />
@@ -37,9 +57,9 @@ public AzureOpenAIVectorizer(string apiKey, string resourceName, string deployme
     public int Dim { get; }
 
     /// <inheritdoc />
-    public byte[] Vectorize(string str) => GetFloats(str, _resourceName, _deploymentName, _apiKey).SelectMany(BitConverter.GetBytes).ToArray();
+    public byte[] Vectorize(string str) => GetFloats(str, _resourceName, _deploymentName, _apiKey, _tokenCredential).SelectMany(BitConverter.GetBytes).ToArray();
 
-    internal static float[] GetFloats(string s, string resourceName, string deploymentName, string apiKey)
+    internal static float[] GetFloats(string s, string resourceName, string deploymentName, string? apiKey, TokenCredential? azureCredentials)
     {
         var client = Configuration.Instance.Client;
         var requestContent = JsonContent.Create(new { input = s });
@@ -49,9 +69,24 @@ internal static float[] GetFloats(string s, string resourceName, string deployme
             Method = HttpMethod.Post,
             RequestUri = new Uri(
                 $"https://{resourceName}.openai.azure.com/openai/deployments/{deploymentName}/embeddings?api-version=2023-05-15"),
-            Content = requestContent,
-            Headers = { { "api-key", apiKey } }
+            Content = requestContent
         };
+        
+        if(string.IsNullOrEmpty(apiKey) && azureCredentials != null)
+        {
+            var scope = "https://cognitiveservices.azure.com/.default";
+            AccessToken token = azureCredentials.GetToken(new TokenRequestContext(new []{scope}), CancellationToken.None);
+            request.Headers.Authorization = new AuthenticationHeaderValue("Bearer", token.Token);
+        }
+        else if(!string.IsNullOrEmpty(apiKey))
+        {
+            request.Headers.Add("api-key", apiKey);
+        }
+        else
+        {
+            throw new InvalidOperationException("Either apiKey or azureCredentials must be provided.");
+        }
+        
 
         var res = client.Send(request);
         if (!res.IsSuccessStatusCode)

src/Redis.OM.Vectorizers/AzureOpenAIVectorizerAttribute.cs

@@ -1,3 +1,5 @@
+using Azure.Core;
+using Azure.Identity;
 using Redis.OM.Contracts;
 using Redis.OM.Modeling;
 
@@ -6,13 +8,19 @@ namespace Redis.OM.Vectorizers;
 /// <inheritdoc />
 public class AzureOpenAIVectorizerAttribute : VectorizerAttribute<string>
 {
+    private readonly string? _apikey;
+    private readonly TokenCredential? _tokenCredential;
+    
     /// <inheritdoc />
     public AzureOpenAIVectorizerAttribute(string deploymentName, string resourceName, int dim)
     {
         DeploymentName = deploymentName;
         ResourceName = resourceName;
         Dim = dim;
-        Vectorizer = new AzureOpenAIVectorizer(Configuration.Instance.AzureOpenAIApiKey, ResourceName, DeploymentName, Dim);
+        _apikey = Configuration.Instance.AzureOpenAIApiKey;
+        _tokenCredential = new DefaultAzureCredential();
+        Vectorizer = string.IsNullOrEmpty(Configuration.Instance.AzureOpenAIApiKey) ? new AzureOpenAIVectorizer(resourceName, deploymentName, dim) : new AzureOpenAIVectorizer(Configuration.Instance.AzureOpenAIApiKey, ResourceName, DeploymentName, dim);
+        
     }
 
     /// <summary>
@@ -42,7 +50,7 @@ public override byte[] Vectorize(object obj)
             throw new ArgumentException("Object must be a string to be embedded", nameof(obj));
         }
 
-        var floats = AzureOpenAIVectorizer.GetFloats(s, ResourceName, DeploymentName, Configuration.Instance.AzureOpenAIApiKey);
+        var floats = AzureOpenAIVectorizer.GetFloats(s, ResourceName, DeploymentName, _apikey, _tokenCredential);
         return floats.SelectMany(BitConverter.GetBytes).ToArray();
     }
 }

src/Redis.OM.Vectorizers/Redis.OM.Vectorizers.csproj

@@ -26,6 +26,7 @@
         <PackageReference Include="Microsoft.Extensions.Configuration" Version="7.0.0" />
         <PackageReference Include="Microsoft.Extensions.Configuration.Abstractions" Version="7.0.0" />
         <PackageReference Include="Microsoft.Extensions.Configuration.Binder" Version="7.0.4" />
+        <PackageReference Include="Azure.Identity" Version="1.13.2" />
     </ItemGroup>
 
     <ItemGroup>

test/Redis.OM.Unit.Tests/RediSearchTests/VectorTests/AzureOpenAIVectors.cs

@@ -0,0 +1,24 @@
+using Redis.OM.Modeling;
+using Redis.OM.Modeling.Vectors;
+using Redis.OM.Vectorizers;
+
+namespace Redis.OM.Unit.Tests;
+
+[Document(StorageType = StorageType.Json)]
+public class AzureOpenAIVectors
+{
+    [RedisIdField]
+    public string Id { get; set; }
+    
+    [Indexed]
+    [AzureOpenAIVectorizer("redisom-embedding", "redisom-openai", 1536)]
+    public Vector<string> Sentence { get; set; }
+
+    [Indexed]
+    public string Name { get; set; }
+
+    [Indexed]
+    public int Age { get; set; }
+
+    public VectorScores VectorScore { get; set; }
+}

test/Redis.OM.Unit.Tests/RediSearchTests/VectorTests/VectorFunctionalTests.cs

@@ -388,4 +388,27 @@ public void OpenAIQueryTest()
         result = collection.First(x=>x.TimeStamp > ts && x.Prompt.VectorRange(queryPrompt, .15));
         Assert.Equal("Paris", result.Response);
     }
+    
+    [SkipIfMissingEnvVar("REDIS_OM_AZURE_OAI_TOKEN")]
+    public void TestAzureOpenAIVectorizer()
+    {
+        _connection.DropIndexAndAssociatedRecords(typeof(AzureOpenAIVectors));
+        _connection.CreateIndex(typeof(AzureOpenAIVectors));
+        var collection = new RedisCollection<AzureOpenAIVectors>(_connection);
+        var sentenceVector = Vector.Of("Hello World this is Hal."); 
+        var obj = new AzureOpenAIVectors
+        {
+            Age = 45,
+            Sentence = sentenceVector,
+            Name = "Hal"
+        };
+
+        collection.Insert(obj);
+        var queryVector = Vector.Of("Hello World this is Hal.");
+        var res = collection.NearestNeighbors(x => x.Sentence, 2, queryVector).First();
+        Assert.Equal(obj.Id, res.Id);
+        Assert.True(res.VectorScore.NearestNeighborsScore < .01);
+        Assert.Equal(obj.Sentence.Value, res.Sentence.Value);
+        Assert.Equal(obj.Sentence.Embedding, res.Sentence.Embedding);
+    }
 }