Azure OpenAI support (#912)

Co-authored-by: Max Leiter <max.leiter@vercel.com>

.changeset/pink-yaks-lie.md

@@ -0,0 +1,5 @@
+---
+'ai': patch
+---
+
+OpenAIStream: Add support for the Azure OpenAI client library

docs/pages/docs/guides.mdx

@@ -4,6 +4,7 @@ Vercel AI SDK is compatible with many popular AI and model providers. This secti
 with models and services from these providers:
 
 - [OpenAI](./guides/providers/openai)
+- [Azure OpenAI](./guides/providers/azure-openai)
 - [AWS Bedrock](./guides/providers/aws-bedrock)
 - [Anthropic](./guides/providers/anthropic)
 - [Cohere](./guides/providers/cohere)

docs/pages/docs/guides/providers/azure-openai.mdx

@@ -0,0 +1,146 @@
+---
+title: Azure OpenAI
+---
+
+import { Steps, Callout } from 'nextra-theme-docs';
+
+# Azure OpenAI
+
+Vercel AI SDK provides a set of utilities to make it easy to use the [Azure OpenAI client library](https://learn.microsoft.com/en-us/javascript/api/overview/azure/openai-readme?view=azure-node-preview). In this guide, we'll walk through how to use the utilities to create a chat bot.
+
+## Guide: Chat Bot
+
+<Steps>
+
+### Create a Next.js app
+
+Create a Next.js application and install `ai` and `@azure/openai`, the Vercel AI SDK and Azure OpenAI client respectively:
+
+```sh
+pnpm dlx create-next-app my-ai-app
+cd my-ai-app
+pnpm install ai @azure/openai
+```
+
+### Add your Azure OpenAI API Key to `.env`
+
+Create a `.env` file in your project root and add your Azure OpenAI API Key:
+
+```env filename=".env"
+AZURE_OPENAI_API_KEY=xxxxxxxxx
+```
+
+### Create a Route Handler
+
+Create a Next.js Route Handler that uses the Edge Runtime that we'll use to generate a chat completion via Azure OpenAI that we'll then stream back to our Next.js.
+
+For this example, we'll create a route handler at `app/api/chat/route.ts` that accepts a `POST` request with a `messages` array of strings:
+
+```tsx filename="app/api/chat/route.ts" showLineNumbers
+import { OpenAIClient, AzureKeyCredential } from '@azure/openai';
+import { OpenAIStream, StreamingTextResponse } from 'ai';
+
+// Create an OpenAI API client (that's edge friendly!)
+const client = new OpenAIClient(
+  'https://YOUR-AZURE-OPENAI-ENDPOINT',
+  new AzureKeyCredential(process.env.AZURE_OPENAI_API_KEY!),
+);
+
+// IMPORTANT! Set the runtime to edge
+export const runtime = 'edge';
+
+export async function POST(req: Request) {
+  const { messages } = await req.json();
+
+  // Ask Azure OpenAI for a streaming chat completion given the prompt
+  const response = await client.streamChatCompletions(
+    'YOUR_DEPLOYED_INSTANCE_NAME',
+    messages,
+  });
+
+  // Convert the response into a friendly text-stream
+  const stream = OpenAIStream(response);
+  // Respond with the stream
+  return new StreamingTextResponse(stream);
+}
+```
+
+<Callout>
+  Vercel AI SDK provides 2 utility helpers to make the above seamless: First, we
+  pass the streaming `response` we receive from Azure OpenAI library to
+  [`OpenAIStream`](/docs/api-reference/openai-stream). This method
+  decodes/extracts the text tokens in the response and then re-encodes them
+  properly for simple consumption. We can then pass that new stream directly to
+  [`StreamingTextResponse`](/docs/api-reference/streaming-text-response). This
+  is another utility class that extends the normal Node/Edge Runtime `Response`
+  class with the default headers you probably want (hint: `'Content-Type':
+  'text/plain; charset=utf-8'` is already set for you).
+</Callout>
+
+### Wire up the UI
+
+Create a Client component with a form that we'll use to gather the prompt from the user and then stream back the completion from.
+By default, the [`useChat`](/docs/api-reference/use-chat) hook will use the `POST` Route Handler we created above (it defaults to `/api/chat`). You can override this by passing a `api` prop to `useChat({ api: '...'})`.
+
+```tsx filename="app/page.tsx" showLineNumbers
+'use client';
+
+import { useChat } from 'ai/react';
+
+export default function Chat() {
+  const { messages, input, handleInputChange, handleSubmit } = useChat();
+  return (
+    <div className="flex flex-col w-full max-w-md py-24 mx-auto stretch">
+      {messages.map(m => (
+        <div key={m.id} className="whitespace-pre-wrap">
+          {m.role === 'user' ? 'User: ' : 'AI: '}
+          {m.content}
+        </div>
+      ))}
+
+      <form onSubmit={handleSubmit}>
+        <input
+          className="fixed bottom-0 w-full max-w-md p-2 mb-8 border border-gray-300 rounded shadow-xl"
+          value={input}
+          placeholder="Say something..."
+          onChange={handleInputChange}
+        />
+      </form>
+    </div>
+  );
+}
+```
+
+</Steps>
+
+## Guide: Save to Database After Completion
+
+It’s common to want to save the result of a completion to a database after streaming it back to the user. The `OpenAIStream` adapter accepts a couple of optional callbacks that can be used to do this.
+
+```tsx filename="app/api/completion/route.ts" showLineNumbers
+export async function POST(req: Request) {
+  // ...
+
+  // Convert the response into a friendly text-stream
+  const stream = OpenAIStream(response, {
+    onStart: async () => {
+      // This callback is called when the stream starts
+      // You can use this to save the prompt to your database
+      await savePromptToDatabase(prompt);
+    },
+    onToken: async (token: string) => {
+      // This callback is called for each token in the stream
+      // You can use this to debug the stream or save the tokens to your database
+      console.log(token);
+    },
+    onCompletion: async (completion: string) => {
+      // This callback is called when the stream completes
+      // You can use this to save the final completion to your database
+      await saveCompletionToDatabase(completion);
+    },
+  });
+
+  // Respond with the stream
+  return new StreamingTextResponse(stream);
+}
+```

docs/pages/docs/guides/providers/openai.mdx

@@ -265,38 +265,6 @@ export async function POST(req: Request) {
 }
 ```
 
-## Guide: Use with Azure OpenAI Service
-
-You can pass custom options to the `Configuration` from the OpenAI package to connect to the an Azure instance.
-See the [OpenAI client repository](https://github.com/openai/openai-node/blob/v4/examples/azure.ts) for a more complete example.
-
-```tsx filename="app/api/completion/route.ts"
-import OpenAI from 'openai';
-
-const resource = '<your resource name>';
-const model = '<your model>';
-
-const apiKey = process.env.AZURE_OPENAI_API_KEY;
-if (!apiKey) {
-  throw new Error('AZURE_OPENAI_API_KEY is missing from the environment.');
-}
-
-// Azure OpenAI requires a custom baseURL, api-version query param, and api-key header.
-const openai = new OpenAI({
-  apiKey,
-  baseURL: `https://${resource}.openai.azure.com/openai/deployments/${model}`,
-  defaultQuery: { 'api-version': '2023-06-01-preview' },
-  defaultHeaders: { 'api-key': apiKey },
-});
-```
-
-<Callout>
-  Note: Before the release of `openai@4`, we previously recommended using the
-  `openai-edge` library because of it's compatibility with Vercel Edge Runtime.
-  The OpenAI SDK now supports Edge Runtime out of the box, so we recommend using
-  the official `openai` library instead.
-</Callout>
-
 ## Guide: Using Images with GPT 4 Vision and useChat
 
 You can use the extra `data` property that is part of `handleSubmit` to send additional data

docs/pages/docs/index.mdx

@@ -186,6 +186,7 @@ export function A({ href, children }) {
 {<div className="grid lg:grid-cols-2 grid-cols-1 gap-4 mt-8">
 
   <IntegrationCard href="/docs/guides/providers/openai" title="OpenAI" description="OpenAI is an AI research and deployment company. The Vercel AI SDK provides a simple way to use OpenAI in your frontend web applications." />
+  <IntegrationCard href="/docs/guides/providers/azure-openai" title="Azure OpenAI" description="Azure OpenAI Service offers industry-leading coding and language AI models that you can fine-tune to your specific needs for a variety of use cases." />
   <IntegrationCard href="/docs/guides/providers/langchain" title="LangChain" description="LangChain is an open source prompt engineering framework for developing applications powered by language models. The Vercel AI SDK provides a simple way to use LangChain in your frontend web applications." />
   <IntegrationCard href="/docs/guides/providers/mistral" title="Mistral" description="Mistral is an AI research and deployment company. The Vercel AI SDK provides a simple way to use Mistral models in your frontend web applications." />
   <IntegrationCard href="/docs/guides/providers/aws-bedrock" title="AWS Bedrock" description="Bedrock is a fully managed service that offers a choice of high-performing foundation models from leading AI companies through AWS. The Vercel AI SDK provides a simple way to use AWS Bedrock models in your frontend web applications." />

packages/core/streams/azure-openai-types.ts

@@ -0,0 +1,39 @@
+export declare interface AzureChatCompletions {
+  id: string;
+  created: Date;
+  choices: AzureChatChoice[];
+  systemFingerprint?: string;
+  usage?: AzureCompletionsUsage;
+  promptFilterResults: any[]; // marker
+}
+
+export declare interface AzureChatChoice {
+  message?: AzureChatResponseMessage;
+  index: number;
+  finishReason: string | null;
+  delta?: AzureChatResponseMessage;
+}
+
+export declare interface AzureChatResponseMessage {
+  role: string;
+  content: string | null;
+  toolCalls: AzureChatCompletionsFunctionToolCall[];
+  functionCall?: AzureFunctionCall;
+}
+
+export declare interface AzureCompletionsUsage {
+  completionTokens: number;
+  promptTokens: number;
+  totalTokens: number;
+}
+
+export declare interface AzureFunctionCall {
+  name: string;
+  arguments: string;
+}
+
+export declare interface AzureChatCompletionsFunctionToolCall {
+  type: 'function';
+  function: AzureFunctionCall;
+  id: string;
+}

packages/core/streams/openai-stream.test.ts

@@ -11,6 +11,7 @@ import {
 } from '../tests/snapshots/openai-chat';
 import { createClient, readAllChunks } from '../tests/utils/mock-client';
 import { DEFAULT_TEST_URL, createMockServer } from '../tests/utils/mock-server';
+import { azureOpenaiChatCompletionChunks } from '../tests/snapshots/azure-openai';
 
 const FUNCTION_CALL_TEST_URL = DEFAULT_TEST_URL + 'mock-func-call';
 const TOOL_CALL_TEST_URL = DEFAULT_TEST_URL + 'mock-tool-call';
@@ -353,4 +354,46 @@ describe('OpenAIStream', () => {
       });
     });
   });
+
+  describe('Azure SDK', () => {
+    async function* asyncIterableFromArray(array: any[]) {
+      for (const item of array) {
+        // You can also perform any asynchronous operations here if needed
+        yield item;
+      }
+    }
+
+    describe('StreamData prototcol', () => {
+      it('should send text', async () => {
+        const data = new experimental_StreamData();
+
+        const stream = OpenAIStream(
+          asyncIterableFromArray(azureOpenaiChatCompletionChunks),
+          {
+            onFinal() {
+              data.close();
+            },
+            experimental_streamData: true,
+          },
+        );
+
+        const response = new StreamingTextResponse(stream, {}, data);
+
+        const client = createClient(response);
+        const chunks = await client.readAll();
+
+        expect(chunks).toEqual([
+          '0:"Hello"\n',
+          '0:"!"\n',
+          '0:" How"\n',
+          '0:" can"\n',
+          '0:" I"\n',
+          '0:" assist"\n',
+          '0:" you"\n',
+          '0:" today"\n',
+          '0:"?"\n',
+        ]);
+      });
+    });
+  });
 });

packages/core/streams/openai-stream.ts

@@ -16,6 +16,7 @@ import {
   createCallbacksTransformer,
   ToolCallPayload,
 } from './ai-stream';
+import { AzureChatCompletions } from './azure-openai-types';
 import { createStreamDataTransformer } from './stream-data';
 
 export type OpenAIStreamCallbacks = AIStreamCallbacksAndOptions & {
@@ -277,8 +278,38 @@ function parseOpenAIStream(): (data: string) => string | void {
  */
 async function* streamable(stream: AsyncIterableOpenAIStreamReturnTypes) {
   const extract = chunkToText();
-  for await (const chunk of stream) {
+
+  for await (let chunk of stream) {
+    // convert chunk if it is an Azure chat completion. Azure does not expose all
+    // properties in the interfaces, and also uses camelCase instead of snake_case
+    if ('promptFilterResults' in chunk) {
+      chunk = {
+        id: chunk.id,
+        created: chunk.created.getDate(),
+        object: (chunk as any).object, // not exposed by Azure API
+        model: (chunk as any).model, // not exposed by Azure API
+        choices: chunk.choices.map(choice => ({
+          delta: {
+            content: choice.delta?.content,
+            function_call: choice.delta?.functionCall,
+            role: choice.delta?.role as any,
+            tool_calls: choice.delta?.toolCalls?.length
+              ? choice.delta?.toolCalls?.map((toolCall, index) => ({
+                  index,
+                  id: toolCall.id,
+                  function: toolCall.function,
+                  type: toolCall.type,
+                }))
+              : undefined,
+          },
+          finish_reason: choice.finishReason as any,
+          index: choice.index,
+        })),
+      } satisfies ChatCompletionChunk;
+    }
+
     const text = extract(chunk);
+
     if (text) yield text;
   }
 }
@@ -350,7 +381,8 @@ const __internal__OpenAIFnMessagesSymbol = Symbol(
 
 type AsyncIterableOpenAIStreamReturnTypes =
   | AsyncIterable<ChatCompletionChunk>
-  | AsyncIterable<Completion>;
+  | AsyncIterable<Completion>
+  | AsyncIterable<AzureChatCompletions>;
 
 type ExtractType<T> = T extends AsyncIterable<infer U> ? U : never;