Add test for chat service && pull out document search into an individual service (#270)

LittleLittleCloud · IEvangelist · web-flow · commit 9543b548feca · 2024-02-02T12:57:34.000-06:00
## Purpose
&lt;!-- Describe the intention of the changes being proposed. What problem
does it solve or functionality does it add? --&gt;
* ...

## Does this introduce a breaking change?
&lt;!-- Mark one with an "x". --&gt;
```
[ ] Yes
[ ] No
```

## Pull Request Type
What kind of change does this Pull Request introduce?

&lt;!-- Please check the one that applies to this PR using "x". --&gt;
```
[ ] Bugfix
[ ] Feature
[ ] Code style update (formatting, local variables)
[ ] Refactoring (no functional changes, no api changes)
[ ] Documentation content changes
[ ] Other... Please describe:
```

## How to Test
*  Get the code

```
git clone [repo-address]
cd [repo-name]
git checkout [branch-name]
npm install
```

* Test the code
&lt;!-- Add steps to run the tests suite and/or manually test --&gt;
```
```

## What to Check
Verify that the following are valid
* ...

## Other Information
&lt;!-- Add any other helpful information that may be needed here. --&gt;

---------

Co-authored-by: David Pine &lt;david.pine@microsoft.com&gt;
diff --git a/app/backend/Extensions/SearchClientExtensions.cs b/app/backend/Extensions/SearchClientExtensions.cs
@@ -41,7 +41,7 @@ internal static async Task<SupportingContentRecord[]> QueryDocumentsAsync(
                 Size = top,
             };
 
-        if (embedding != null && overrides?.RetrievalMode != "Text")
+        if (embedding != null && overrides?.RetrievalMode != RetrievalMode.Text)
         {
             var k = useSemanticRanker ? 50 : top;
             var vectorQuery = new VectorizedQuery(embedding)
diff --git a/app/backend/Extensions/ServiceCollectionExtensions.cs b/app/backend/Extensions/ServiceCollectionExtensions.cs
@@ -27,18 +27,19 @@ internal static IServiceCollection AddAzureServices(this IServiceCollection serv
             return sp.GetRequiredService<BlobServiceClient>().GetBlobContainerClient(azureStorageContainer);
         });
 
-        services.AddSingleton<SearchClient>(sp =>
+        services.AddSingleton<IDocumentService, AzureDocumentService>(sp =>
         {
             var config = sp.GetRequiredService<IConfiguration>();
-            var (azureSearchServiceEndpoint, azureSearchIndex) =
-                (config["AzureSearchServiceEndpoint"], config["AzureSearchIndex"]);
-
+            var azureSearchServiceEndpoint = config["AzureSearchServiceEndpoint"];
             ArgumentNullException.ThrowIfNullOrEmpty(azureSearchServiceEndpoint);
 
+            var azureSearchIndex = config["AzureSearchIndex"];
+            ArgumentNullException.ThrowIfNullOrEmpty(azureSearchIndex);
+
             var searchClient = new SearchClient(
-                new Uri(azureSearchServiceEndpoint), azureSearchIndex, s_azureCredential);
+                               new Uri(azureSearchServiceEndpoint), azureSearchIndex, s_azureCredential);
 
-            return searchClient;
+            return new AzureDocumentService(searchClient);
         });
 
         services.AddSingleton<DocumentAnalysisClient>(sp =>
diff --git a/app/backend/Services/ReadRetrieveReadChatService.cs b/app/backend/Services/ReadRetrieveReadChatService.cs
@@ -4,12 +4,12 @@ namespace MinimalApi.Services;
 
 public class ReadRetrieveReadChatService
 {
-    private readonly SearchClient _searchClient;
+    private readonly IDocumentService _searchClient;
     private readonly IKernel _kernel;
     private readonly IConfiguration _configuration;
 
     public ReadRetrieveReadChatService(
-        SearchClient searchClient,
+        IDocumentService searchClient,
         OpenAIClient client,
         IConfiguration configuration)
     {
@@ -45,15 +45,15 @@ public async Task<ApproachResponse> ReplyAsync(
         var question = history.LastOrDefault()?.User is { } userQuestion
             ? userQuestion
             : throw new InvalidOperationException("Use question is null");
-        if (overrides?.RetrievalMode != "Text" && embedding is not null)
+        if (overrides?.RetrievalMode != RetrievalMode.Text && embedding is not null)
         {
             embeddings = (await embedding.GenerateEmbeddingAsync(question, cancellationToken: cancellationToken)).ToArray();
         }
 
         // step 1
         // use llm to get query if retrieval mode is not vector
         string? query = null;
-        if (overrides?.RetrievalMode != "Vector")
+        if (overrides?.RetrievalMode != RetrievalMode.Vector)
         {
             var getQueryChat = chat.CreateNewChat(@"You are a helpful AI assistant, generate search query for followup question.
 Make your respond simple and precise. Return the query only, do not return any other text.
diff --git a/app/shared/Shared/Models/RequestOverrides.cs b/app/shared/Shared/Models/RequestOverrides.cs
@@ -1,12 +1,11 @@
 ﻿// Copyright (c) Microsoft. All rights reserved.
 
 namespace Shared.Models;
-
 public record RequestOverrides
 {
     public bool SemanticRanker { get; set; } = false;
 
-    public string RetrievalMode { get; set; } = "Vector"; // available option: Text, Vector, Hybrid
+    public RetrievalMode RetrievalMode { get; set; } = RetrievalMode.Vector; // available option: Text, Vector, Hybrid
 
     public bool? SemanticCaptions { get; set; }
     public string? ExcludeCategory { get; set; }
diff --git a/app/shared/Shared/Models/RetrievalMode.cs b/app/shared/Shared/Models/RetrievalMode.cs
@@ -0,0 +1,24 @@
+﻿// Copyright (c) Microsoft. All rights reserved.
+
+namespace Shared.Models;
+
+/// <summary>
+/// retrieval mode for azure search service
+/// </summary>
+public enum RetrievalMode
+{
+    /// <summary>
+    /// Text-only model, where only query will be used to retrieve the results
+    /// </summary>
+    Text = 0,
+
+    /// <summary>
+    /// Vector-only model, where only embeddings will be used to retrieve the results
+    /// </summary>
+    Vector,
+
+    /// <summary>
+    /// Text + Vector model, where both query and embeddings will be used to retrieve the results
+    /// </summary>
+    Hybrid,
+}
diff --git a/app/shared/Shared/Services/AzureComputerVisionService.cs b/app/shared/Shared/Services/AzureComputerVisionService.cs
@@ -2,30 +2,30 @@
 
 using System.Net.Http.Headers;
 using System.Text;
+using System.Text.Json;
+using Azure.Core;
 
-namespace MinimalApi.Services;
-
-public class AzureComputerVisionService(IHttpClientFactory httpClientFactory, string endPoint, string apiKey)
+public class AzureComputerVisionService(HttpClient client, string endPoint, TokenCredential tokenCredential)
 {
     // add virtual keyword to make it mockable
-    public virtual async Task<ImageEmbeddingResponse> VectorizeImageAsync(string imagePathOrUrl, CancellationToken ct = default)
+    public async Task<ImageEmbeddingResponse> VectorizeImageAsync(string imagePathOrUrl, CancellationToken ct = default)
     {
         var api = $"{endPoint}/computervision/retrieval:vectorizeImage?api-version=2023-02-01-preview&modelVersion=latest";
+        var token = await tokenCredential.GetTokenAsync(new TokenRequestContext(new[] { "https://cognitiveservices.azure.com/.default" }), ct);
         // first try to read as local file
         if (File.Exists(imagePathOrUrl))
         {
             using var request = new HttpRequestMessage(HttpMethod.Post, api);
 
             // set authorization header
-            request.Headers.Add("Ocp-Apim-Subscription-Key", apiKey);
+            request.Headers.Add("Authorization", $"Bearer {token.Token}");
 
             // set body
             var bytes = await File.ReadAllBytesAsync(imagePathOrUrl, ct);
             request.Content = new ByteArrayContent(bytes);
             request.Content.Headers.ContentType = new MediaTypeHeaderValue("image/*");
 
             // send request
-            using var client = httpClientFactory.CreateClient();
             using var response = await client.SendAsync(request, ct);
             response.EnsureSuccessStatusCode();
 
@@ -44,14 +44,13 @@ public virtual async Task<ImageEmbeddingResponse> VectorizeImageAsync(string ima
             request.Headers.Accept.Add(new MediaTypeWithQualityHeaderValue("application/json"));
 
             // set authorization header
-            request.Headers.Add("Ocp-Apim-Subscription-Key", apiKey);
+            request.Headers.Add("Authorization", $"Bearer {token.Token}");
 
             // set body
             var body = new { url = imagePathOrUrl };
             request.Content = new StringContent(JsonSerializer.Serialize(body), Encoding.UTF8, "application/json");
 
             // send request
-            using var client = httpClientFactory.CreateClient();
             using var response = await client.SendAsync(request, ct);
             response.EnsureSuccessStatusCode();
 
@@ -67,13 +66,14 @@ public virtual async Task<ImageEmbeddingResponse> VectorizeTextAsync(string text
     {
         var api = $"{endPoint}/computervision/retrieval:vectorizeText?api-version=2023-02-01-preview&modelVersion=latest";
 
+        var token = await tokenCredential.GetTokenAsync(new TokenRequestContext(new[] { "https://cognitiveservices.azure.com/.default" }), ct);
         using var request = new HttpRequestMessage(HttpMethod.Post, api);
 
         // set content type to application/json
         request.Headers.Accept.Add(new MediaTypeWithQualityHeaderValue("application/json"));
 
         // set authorization header
-        request.Headers.Add("Ocp-Apim-Subscription-Key", apiKey);
+        request.Headers.Add("Authorization", $"Bearer {token.Token}");
 
         // set body
         var body = new { text };
diff --git a/app/shared/Shared/Services/AzureDocumentService.cs b/app/shared/Shared/Services/AzureDocumentService.cs
@@ -0,0 +1,123 @@
+﻿using Azure;
+using Azure.Search.Documents;
+using Azure.Search.Documents.Models;
+using Shared.Models;
+
+public interface IDocumentService
+{
+    Task<SupportingContentRecord[]> QueryDocumentsAsync(
+               string? query = null,
+               float[]? embedding = null,
+               RequestOverrides? overrides = null,
+               CancellationToken cancellationToken = default);
+}
+
+public class AzureDocumentService(SearchClient searchClient) : IDocumentService
+{
+    public async Task<SupportingContentRecord[]> QueryDocumentsAsync(
+        string? query = null,
+        float[]? embedding = null,
+        RequestOverrides? overrides = null,
+        CancellationToken cancellationToken = default)
+    {
+        if (query is null && embedding is null)
+        {
+            throw new ArgumentException("Either query or embedding must be provided");
+        }
+
+        var documentContents = string.Empty;
+        var top = overrides?.Top ?? 3;
+        var exclude_category = overrides?.ExcludeCategory;
+        var filter = exclude_category == null ? string.Empty : $"category ne '{exclude_category}'";
+        var useSemanticRanker = overrides?.SemanticRanker ?? false;
+        var useSemanticCaptions = overrides?.SemanticCaptions ?? false;
+
+        SearchOptions searchOptions = useSemanticRanker
+            ? new SearchOptions
+            {
+                Filter = filter,
+                QueryType = SearchQueryType.Semantic,
+                SemanticSearch = new()
+                {
+                    SemanticConfigurationName = "default",
+                    QueryCaption = new(useSemanticCaptions
+                        ? QueryCaptionType.Extractive
+                        : QueryCaptionType.None),
+                },
+                // TODO: Find if these options are assignable
+                //QueryLanguage = "en-us",
+                //QuerySpeller = "lexicon",
+                Size = top,
+            }
+            : new SearchOptions
+            {
+                Filter = filter,
+                Size = top,
+            };
+
+        if (embedding != null && overrides?.RetrievalMode != RetrievalMode.Text)
+        {
+            var k = useSemanticRanker ? 50 : top;
+            var vectorQuery = new VectorizedQuery(embedding)
+            {
+                // if semantic ranker is enabled, we need to set the rank to a large number to get more
+                // candidates for semantic reranking
+                KNearestNeighborsCount = useSemanticRanker ? 50 : top,
+            };
+            vectorQuery.Fields.Add("embedding");
+            searchOptions.VectorSearch = new();
+            searchOptions.VectorSearch.Queries.Add(vectorQuery);
+        }
+
+        var searchResultResponse = await searchClient.SearchAsync<SearchDocument>(
+            query, searchOptions, cancellationToken);
+        if (searchResultResponse.Value is null)
+        {
+            throw new InvalidOperationException("fail to get search result");
+        }
+
+        SearchResults<SearchDocument> searchResult = searchResultResponse.Value;
+
+        // Assemble sources here.
+        // Example output for each SearchDocument:
+        // {
+        //   "@search.score": 11.65396,
+        //   "id": "Northwind_Standard_Benefits_Details_pdf-60",
+        //   "content": "x-ray, lab, or imaging service, you will likely be responsible for paying a copayment or coinsurance. The exact amount you will be required to pay will depend on the type of service you receive. You can use the Northwind app or website to look up the cost of a particular service before you receive it.\nIn some cases, the Northwind Standard plan may exclude certain diagnostic x-ray, lab, and imaging services. For example, the plan does not cover any services related to cosmetic treatments or procedures. Additionally, the plan does not cover any services for which no diagnosis is provided.\nIt’s important to note that the Northwind Standard plan does not cover any services related to emergency care. This includes diagnostic x-ray, lab, and imaging services that are needed to diagnose an emergency condition. If you have an emergency condition, you will need to seek care at an emergency room or urgent care facility.\nFinally, if you receive diagnostic x-ray, lab, or imaging services from an out-of-network provider, you may be required to pay the full cost of the service. To ensure that you are receiving services from an in-network provider, you can use the Northwind provider search ",
+        //   "category": null,
+        //   "sourcepage": "Northwind_Standard_Benefits_Details-24.pdf",
+        //   "sourcefile": "Northwind_Standard_Benefits_Details.pdf"
+        // }
+        var sb = new List<SupportingContentRecord>();
+        foreach (var doc in searchResult.GetResults())
+        {
+            doc.Document.TryGetValue("sourcepage", out var sourcePageValue);
+            string? contentValue;
+            try
+            {
+                if (useSemanticCaptions)
+                {
+                    var docs = doc.SemanticSearch.Captions.Select(c => c.Text);
+                    contentValue = string.Join(" . ", docs);
+                }
+                else
+                {
+                    doc.Document.TryGetValue("content", out var value);
+                    contentValue = (string)value;
+                }
+            }
+            catch (ArgumentNullException)
+            {
+                contentValue = null;
+            }
+
+            if (sourcePageValue is string sourcePage && contentValue is string content)
+            {
+                content = content.Replace('\r', ' ').Replace('\n', ' ');
+                sb.Add(new SupportingContentRecord(sourcePage, content));
+            }
+        }
+
+        return [.. sb];
+    }
+}
diff --git a/app/shared/Shared/Shared.csproj b/app/shared/Shared/Shared.csproj
@@ -1,4 +1,4 @@
-<Project Sdk="Microsoft.NET.Sdk">
+﻿<Project Sdk="Microsoft.NET.Sdk">
 
   <PropertyGroup>
     <TargetFramework>net8.0</TargetFramework>
@@ -7,4 +7,15 @@
     <LangVersion>preview</LangVersion>
   </PropertyGroup>
 
+  <ItemGroup>
+    <PackageReference Include="Azure.AI.FormRecognizer" />
+    <PackageReference Include="Azure.AI.OpenAI" />
+    <PackageReference Include="Azure.Extensions.AspNetCore.Configuration.Secrets" />
+    <PackageReference Include="Azure.Identity" />
+    <PackageReference Include="Azure.Search.Documents" />
+    <PackageReference Include="Azure.Storage.Blobs" />
+    <PackageReference Include="Microsoft.ApplicationInsights.AspNetCore" />
+    <PackageReference Include="PdfSharpCore" />
+  </ItemGroup>
+
 </Project>
diff --git a/app/tests/MinimalApi.Tests/Attribute/EnvironmentSpecificFactAttribute.cs b/app/tests/MinimalApi.Tests/Attribute/EnvironmentSpecificFactAttribute.cs
@@ -29,11 +29,11 @@ protected EnvironmentSpecificFactAttribute(string skipMessage)
     protected abstract bool IsEnvironmentSupported();
 }
 
-public sealed class ApiKeyFactAttribute : EnvironmentSpecificFactAttribute
+public sealed class EnvironmentVariablesFactAttribute : EnvironmentSpecificFactAttribute
 {
     private readonly string[] _envVariableNames;
     
-    public ApiKeyFactAttribute(params string[] envVariableNames) : base($"{string.Join(", ", envVariableNames)} is not found in env")
+    public EnvironmentVariablesFactAttribute(params string[] envVariableNames) : base($"{string.Join(", ", envVariableNames)} is not found in env")
     {
         _envVariableNames = envVariableNames;
     }
diff --git a/app/tests/MinimalApi.Tests/AzureComputerVisionServiceTest.cs b/app/tests/MinimalApi.Tests/AzureComputerVisionServiceTest.cs
@@ -1,5 +1,7 @@
 ﻿// Copyright (c) Microsoft. All rights reserved.
 
+using Azure.Core;
+using Azure.Identity;
 using FluentAssertions;
 using MinimalApi.Services;
 using NSubstitute;
@@ -8,16 +10,14 @@ namespace MinimalApi.Tests;
 
 public class AzureComputerVisionServiceTest
 {
-    [ApiKeyFact("AZURE_COMPUTER_VISION_API_KEY", "AZURE_COMPUTER_VISION_ENDPOINT")]
+    [EnvironmentVariablesFact("AZURE_COMPUTER_VISION_ENDPOINT")]
     public async Task VectorizeImageTestAsync()
     {
         var endpoint = Environment.GetEnvironmentVariable("AZURE_COMPUTER_VISION_ENDPOINT") ?? throw new InvalidOperationException();
-        var apiKey = Environment.GetEnvironmentVariable("AZURE_COMPUTER_VISION_API_KEY") ?? throw new InvalidOperationException();
-        var httpClientFactory = Substitute.For<IHttpClientFactory>();
-        httpClientFactory.CreateClient().ReturnsForAnyArgs(x => new HttpClient());
-        var service = new AzureComputerVisionService(httpClientFactory, endpoint, apiKey);
+        using var httpClient = new HttpClient();
         var imageUrl = @"https://learn.microsoft.com/azure/ai-services/computer-vision/media/quickstarts/presentation.png";
 
+        var service = new AzureComputerVisionService(httpClient, endpoint, new DefaultAzureCredential());
         var result = await service.VectorizeImageAsync(imageUrl);
 
         result.modelVersion.Should().NotBeNullOrEmpty();
@@ -45,14 +45,12 @@ public async Task VectorizeImageTestAsync()
         }
     }
 
-    [ApiKeyFact("AZURE_COMPUTER_VISION_API_KEY", "AZURE_COMPUTER_VISION_ENDPOINT")]
+    [EnvironmentVariablesFact("AZURE_COMPUTER_VISION_ENDPOINT")]
     public async Task VectorizeTextTestAsync()
     {
         var endpoint = Environment.GetEnvironmentVariable("AZURE_COMPUTER_VISION_ENDPOINT") ?? throw new InvalidOperationException();
-        var apiKey = Environment.GetEnvironmentVariable("AZURE_COMPUTER_VISION_API_KEY") ?? throw new InvalidOperationException();
-        var httpClientFactory = Substitute.For<IHttpClientFactory>();
-        httpClientFactory.CreateClient().ReturnsForAnyArgs(x => new HttpClient());
-        var service = new AzureComputerVisionService(httpClientFactory, endpoint, apiKey);
+        using var httpClient = new HttpClient();
+        var service = new AzureComputerVisionService(httpClient, endpoint, new DefaultAzureCredential());
         var text = "Hello world";
         var result = await service.VectorizeTextAsync(text);
 
diff --git a/app/tests/MinimalApi.Tests/AzureDocumentSearchServiceTest.cs b/app/tests/MinimalApi.Tests/AzureDocumentSearchServiceTest.cs
diff --git a/app/tests/MinimalApi.Tests/ReadRetrieveReadChatServiceTest.cs b/app/tests/MinimalApi.Tests/ReadRetrieveReadChatServiceTest.cs

Original file line number	Diff line number	Diff line change
`@@ -41,7 +41,7 @@ internal static async Task<SupportingContentRecord[]> QueryDocumentsAsync(`
`41`	`41`	`Size = top,`
`42`	`42`	`};`
`43`	`43`
`44`		`- if (embedding != null && overrides?.RetrievalMode != "Text")`
	`44`	`+ if (embedding != null && overrides?.RetrievalMode != RetrievalMode.Text)`
`45`	`45`	`{`
`46`	`46`	`var k = useSemanticRanker ? 50 : top;`
`47`	`47`	`var vectorQuery = new VectorizedQuery(embedding)`
Original file line number	Diff line number	Diff line change
`@@ -4,12 +4,12 @@ namespace MinimalApi.Services;`
`4`	`4`
`5`	`5`	`public class ReadRetrieveReadChatService`
`6`	`6`	`{`
`7`		`- private readonly SearchClient _searchClient;`
	`7`	`+ private readonly IDocumentService _searchClient;`
`8`	`8`	`private readonly IKernel _kernel;`
`9`	`9`	`private readonly IConfiguration _configuration;`
`10`	`10`
`11`	`11`	`public ReadRetrieveReadChatService(`
`12`		`- SearchClient searchClient,`
	`12`	`+ IDocumentService searchClient,`
`13`	`13`	`OpenAIClient client,`
`14`	`14`	`IConfiguration configuration)`
`15`	`15`	`{`
`@@ -45,15 +45,15 @@ public async Task<ApproachResponse> ReplyAsync(`
`45`	`45`	`var question = history.LastOrDefault()?.User is { } userQuestion`
`46`	`46`	`? userQuestion`
`47`	`47`	`: throw new InvalidOperationException("Use question is null");`
`48`		`- if (overrides?.RetrievalMode != "Text" && embedding is not null)`
	`48`	`+ if (overrides?.RetrievalMode != RetrievalMode.Text && embedding is not null)`
`49`	`49`	`{`
`50`	`50`	`embeddings = (await embedding.GenerateEmbeddingAsync(question, cancellationToken: cancellationToken)).ToArray();`
`51`	`51`	`}`
`52`	`52`
`53`	`53`	`// step 1`
`54`	`54`	`// use llm to get query if retrieval mode is not vector`
`55`	`55`	`string? query = null;`
`56`		`- if (overrides?.RetrievalMode != "Vector")`
	`56`	`+ if (overrides?.RetrievalMode != RetrievalMode.Vector)`
`57`	`57`	`{`
`58`	`58`	`var getQueryChat = chat.CreateNewChat(@"You are a helpful AI assistant, generate search query for followup question.`
`59`	`59`	`Make your respond simple and precise. Return the query only, do not return any other text.`
Original file line number	Diff line number	Diff line change
`@@ -29,11 +29,11 @@ protected EnvironmentSpecificFactAttribute(string skipMessage)`
`29`	`29`	`protected abstract bool IsEnvironmentSupported();`
`30`	`30`	`}`
`31`	`31`
`32`		`-public sealed class ApiKeyFactAttribute : EnvironmentSpecificFactAttribute`
	`32`	`+public sealed class EnvironmentVariablesFactAttribute : EnvironmentSpecificFactAttribute`
`33`	`33`	`{`
`34`	`34`	`private readonly string[] _envVariableNames;`
`35`	`35`
`36`		`- public ApiKeyFactAttribute(params string[] envVariableNames) : base($"{string.Join(", ", envVariableNames)} is not found in env")`
	`36`	`+ public EnvironmentVariablesFactAttribute(params string[] envVariableNames) : base($"{string.Join(", ", envVariableNames)} is not found in env")`
`37`	`37`	`{`
`38`	`38`	`_envVariableNames = envVariableNames;`
`39`	`39`	`}`