[azopenai] Adding in support for using Azure Chat Extensions with external data sources (Azure#21426)

Adding in support for chat extensions which allow you (in Azure OpenAI) to use your own data from a Cognitive Search index.

Also, updated the example files to use the `context.TODO()` consistently, and to use our new "comment + panic" in testable examples.

Fixes Azure#21373

Author: richardpark-msft

sdk/ai/azopenai/CHANGELOG.md

@@ -1,9 +1,12 @@
 # Release History
 
-## 0.1.2 (Unreleased)
+## 0.2.0 (2023-08-28)
 
 ### Features Added
 
+- ChatCompletions supports Azure OpenAI's newest feature to use Azure OpenAI with your own data. See `example_client_getchatcompletions_extensions_test.go`
+  for a working example. (PR#21426)
+
 ### Breaking Changes
 
 - ChatCompletionsOptions, CompletionsOptions, EmbeddingsOptions `DeploymentID` field renamed to `Deployment`.
@@ -13,9 +16,7 @@
 
 - EventReader, used by GetChatCompletionsStream and GetCompletionsStream for streaming results, would not return an 
   error if the underlying Body reader was closed or EOF'd before the actual DONE: token arrived. This could result in an
-  infinite loop for callers. (PR#)
-
-### Other Changes
+  infinite loop for callers. (PR#21323)
 
 ## 0.1.1 (2023-07-26)
 

sdk/ai/azopenai/assets.json

@@ -2,5 +2,5 @@
   "AssetsRepo": "Azure/azure-sdk-assets",
   "AssetsRepoPrefixPath": "go",
   "TagPrefix": "go/ai/azopenai",
-  "Tag": "go/ai/azopenai_2bf13bba09"
+  "Tag": "go/ai/azopenai_7be6ae3c15"
 }

sdk/ai/azopenai/autorest.md

@@ -100,6 +100,15 @@ directive:
   - from: openapi-document
     where: $.components.schemas["ChatChoice"].properties.finish_reason
     transform: $["$ref"] = "#/components/schemas/CompletionsFinishReason"; delete $.oneOf;
+  - from: openapi-document
+    where: $.components.schemas["AzureChatExtensionConfiguration"].properties.type
+    transform: $["$ref"] = "#/components/schemas/AzureChatExtensionType"; delete $.allOf;
+  - from: openapi-document
+    where: $.components.schemas["AzureChatExtensionConfiguration"].properties.type
+    transform: $["$ref"] = "#/components/schemas/AzureChatExtensionType"; delete $.allOf;
+  - from: openapi-document
+    where: $.components.schemas["AzureCognitiveSearchChatExtensionConfiguration"].properties.queryType
+    transform: $["$ref"] = "#/components/schemas/AzureCognitiveSearchQueryType"; delete $.allOf;
   # Fix "AutoGenerated" models
   - from: openapi-document
     where: $.components.schemas["ChatCompletions"].properties.usage
@@ -294,4 +303,45 @@ directive:
       return $
         .replace(/populate\(objectMap, "model", (c|e).Model\)/g, 'populate(objectMap, "model", &$1.Deployment)')
         .replace(/err = unpopulate\(val, "Model", &(c|e).Model\)/g, 'err = unpopulate(val, "Model", &$1.Deployment)');
+
+  # Make the Azure extensions internal - we expose these through the GetChatCompletions*() functions
+  # and just treat which endpoint we use as an implementation detail.
+  - from: client.go
+    where: $
+    transform: |
+      return $
+        .replace(/GetChatCompletionsWithAzureExtensions([ (])/g, "getChatCompletionsWithAzureExtensions$1")
+        .replace(/GetChatCompletions([ (])/g, "getChatCompletions$1");
+
+  # move the Azure extensions options into place
+  - from: models.go
+    where: $
+    transform: return $.replace(/(\/\/ The configuration entries for Azure OpenAI.+?)DataSources \[\]AzureChatExtensionConfiguration/s, "$1AzureExtensionsOptions *AzureChatExtensionOptions");
+  - from: models_serde.go
+    where: $
+    transform: |
+      return $
+        .replace(/populate\(objectMap, "dataSources", c.DataSources\)/, 'if c.AzureExtensionsOptions != nil { populate(objectMap, "dataSources", c.AzureExtensionsOptions.Extensions) }')
+        // technically not used, but let's be completionists...
+        .replace(/err = unpopulate\(val, "DataSources", &c.DataSources\)/, 'c.AzureExtensionsOptions = &AzureChatExtensionOptions{}; err = unpopulate(val, "DataSources", &c.AzureExtensionsOptions.Extensions)')
+
+  # try to fix some of the generated types.
+
+  # swap the `Parameters` and `Type` fields (Type really drives what's in Parameters)
+  - from: models.go
+    where: $
+    transform: |
+      let typeRE = /(\/\/ REQUIRED; The label for the type of an Azure chat extension.*?Type \*AzureChatExtensionType)/s;
+      let paramsRE = /(\/\/ REQUIRED; The configuration payload used for the Azure chat extension.*?Parameters any)/s;
+
+      return $
+        .replace(paramsRE, "")
+        .replace(typeRE, $.match(typeRE)[1] + "\n\n" + $.match(paramsRE)[1]);
+
+  - from: constants.go
+    where: $
+    transform: |
+      return $.replace(
+        /(AzureChatExtensionTypeAzureCognitiveSearch AzureChatExtensionType)/, 
+        "// AzureChatExtensionTypeAzureCognitiveSearch enables the use of an Azure Cognitive Search index with chat completions.\n// [AzureChatExtensionConfiguration.Parameter] should be of type [AzureCognitiveSearchChatExtensionConfiguration].\n$1");
 ```

sdk/ai/azopenai/ci.yml

@@ -55,3 +55,9 @@ stages:
         OPENAI_EMBEDDINGS_MODEL: $(OPENAI-EMBEDDINGS-MODEL)
         OPENAI_CHAT_COMPLETIONS_MODEL: $(OPENAI-CHAT-COMPLETIONS-MODEL)
         OPENAI_COMPLETIONS_MODEL: $(OPENAI-COMPLETIONS-MODEL)
+
+        # used for BYOD scenarios with ChatCompletions
+        COGNITIVE_SEARCH_API_ENDPOINT: $(COGNITIVE-SEARCH-API-ENDPOINT)
+        COGNITIVE_SEARCH_API_INDEX: $(COGNITIVE-SEARCH-API-INDEX)
+        COGNITIVE_SEARCH_API_KEY: $(COGNITIVE-SEARCH-API-KEY)
+

sdk/ai/azopenai/client.go

@@ -0,0 +1,102 @@
+//go:build go1.18
+// +build go1.18
+
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License. See License.txt in the project root for license information.
+
+package azopenai_test
+
+import (
+	"context"
+	"errors"
+	"io"
+	"testing"
+
+	"github.com/Azure/azure-sdk-for-go/sdk/ai/azopenai"
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/to"
+	"github.com/stretchr/testify/require"
+)
+
+func TestChatCompletions_extensions_bringYourOwnData(t *testing.T) {
+	client := newAzureOpenAIClientForTest(t, azureOpenAI)
+
+	resp, err := client.GetChatCompletions(context.Background(), azopenai.ChatCompletionsOptions{
+		Messages: []azopenai.ChatMessage{
+			{Content: to.Ptr("What does PR complete mean?"), Role: to.Ptr(azopenai.ChatRoleUser)},
+		},
+		MaxTokens: to.Ptr[int32](512),
+		AzureExtensionsOptions: &azopenai.AzureChatExtensionOptions{
+			Extensions: []azopenai.AzureChatExtensionConfiguration{
+				{
+					Type:       to.Ptr(azopenai.AzureChatExtensionTypeAzureCognitiveSearch),
+					Parameters: azureOpenAI.Cognitive,
+				},
+			},
+		},
+		Deployment: "gpt-4",
+	}, nil)
+	require.NoError(t, err)
+
+	// when you BYOD you get some extra content showing you metadata/info from the external
+	// data source.
+	msgContext := resp.Choices[0].Message.Context
+	require.NotEmpty(t, msgContext.Messages[0].Content)
+	require.Equal(t, azopenai.ChatRoleTool, *msgContext.Messages[0].Role)
+
+	require.NotEmpty(t, *resp.Choices[0].Message.Content)
+	require.Equal(t, azopenai.CompletionsFinishReasonStop, *resp.Choices[0].FinishReason)
+}
+
+func TestChatExtensionsStreaming_extensions_bringYourOwnData(t *testing.T) {
+	client := newAzureOpenAIClientForTest(t, azureOpenAI)
+
+	streamResp, err := client.GetChatCompletionsStream(context.Background(), azopenai.ChatCompletionsOptions{
+		Messages: []azopenai.ChatMessage{
+			{Content: to.Ptr("What does PR complete mean?"), Role: to.Ptr(azopenai.ChatRoleUser)},
+		},
+		MaxTokens: to.Ptr[int32](512),
+		AzureExtensionsOptions: &azopenai.AzureChatExtensionOptions{
+			Extensions: []azopenai.AzureChatExtensionConfiguration{
+				{
+					Type:       to.Ptr(azopenai.AzureChatExtensionTypeAzureCognitiveSearch),
+					Parameters: azureOpenAI.Cognitive,
+				},
+			},
+		},
+		Deployment: "gpt-4",
+	}, nil)
+
+	require.NoError(t, err)
+	defer streamResp.ChatCompletionsStream.Close()
+
+	text := ""
+
+	first := false
+
+	for {
+		event, err := streamResp.ChatCompletionsStream.Read()
+
+		if errors.Is(err, io.EOF) {
+			break
+		}
+
+		require.NoError(t, err)
+
+		if first {
+			// when you BYOD you get some extra content showing you metadata/info from the external
+			// data source.
+			first = false
+			msgContext := event.Choices[0].Message.Context
+			require.NotEmpty(t, msgContext.Messages[0].Content)
+			require.Equal(t, azopenai.ChatRoleTool, *msgContext.Messages[0].Role)
+		}
+
+		for _, choice := range event.Choices {
+			if choice.Delta != nil && choice.Delta.Content != nil {
+				text += *choice.Delta.Content
+			}
+		}
+	}
+
+	require.NotEmpty(t, text)
+}