feat: Streaming azure openai

packages/foundation-models/src/azure-openai/azure-openai-chat-client.test.ts

@@ -7,7 +7,7 @@ import {
 } from '../../../../test-util/mock-http.js';
 import { AzureOpenAiChatClient } from './azure-openai-chat-client.js';
 import { apiVersion } from './model-types.js';
-import type { AzureOpenAiCreateChatCompletionResponse } from './client/inference/schema';
+import type { AzureOpenAiCreateChatCompletionResponse } from './client/inference/schema/index.js';
 
 describe('Azure OpenAI chat client', () => {
   const chatCompletionEndpoint = {

packages/foundation-models/src/azure-openai/azure-openai-chat-client.ts

@@ -6,6 +6,10 @@ import {
 } from '@sap-ai-sdk/ai-api/internal.js';
 import { apiVersion, type AzureOpenAiChatModel } from './model-types.js';
 import { AzureOpenAiChatCompletionResponse } from './azure-openai-chat-completion-response.js';
+import { AzureOpenAiChatCompletionStreamResponse } from './azure-openai-chat-completion-stream-response.js';
+import { AzureOpenAiChatCompletionStream } from './azure-openai-chat-completion-stream.js';
+import type { AzureOpenAiChatCompletionStreamChunkResponse } from './azure-openai-chat-completion-stream-chunk-response.js';
+import type { HttpResponse } from '@sap-cloud-sdk/http-client';
 import type { AzureOpenAiCreateChatCompletionRequest } from './client/inference/schema/index.js';
 
 /**
@@ -28,12 +32,60 @@ export class AzureOpenAiChatClient {
     data: AzureOpenAiCreateChatCompletionRequest,
     requestConfig?: CustomRequestConfig
   ): Promise<AzureOpenAiChatCompletionResponse> {
+    const response = await this.executeRequest(data, requestConfig);
+    return new AzureOpenAiChatCompletionResponse(response);
+  }
+
+  /**
+   * Creates a completion stream for the chat messages.
+   * @param data - The input parameters for the chat completion.
+   * @param requestConfig - The request configuration.
+   * @returns The completion stream.
+   */
+  async stream(
+    data: AzureOpenAiCreateChatCompletionRequest,
+    requestConfig?: CustomRequestConfig
+  ): Promise<
+    AzureOpenAiChatCompletionStreamResponse<AzureOpenAiChatCompletionStreamChunkResponse>
+  > {
+    const response =
+      new AzureOpenAiChatCompletionStreamResponse<AzureOpenAiChatCompletionStreamChunkResponse>();
+    response.stream = (await this.createStream(data, requestConfig))
+      .pipe(AzureOpenAiChatCompletionStream.processChunk)
+      .pipe(AzureOpenAiChatCompletionStream.processFinishReason, response)
+      .pipe(AzureOpenAiChatCompletionStream.processTokenUsage, response);
+    return response;
+  }
+
+  /**
+   * Creates a completion stream of the delta content for the chat messages.
+   * @param data - The input parameters for the chat completion.
+   * @param requestConfig - The request configuration.
+   * @returns The completion stream of the delta content.
+   */
+  async streamContent(
+    data: AzureOpenAiCreateChatCompletionRequest,
+    requestConfig?: CustomRequestConfig
+  ): Promise<AzureOpenAiChatCompletionStreamResponse<string>> {
+    const response = new AzureOpenAiChatCompletionStreamResponse<string>();
+    response.stream = (await this.createStream(data, requestConfig))
+      .pipe(AzureOpenAiChatCompletionStream.processChunk)
+      .pipe(AzureOpenAiChatCompletionStream.processFinishReason, response)
+      .pipe(AzureOpenAiChatCompletionStream.processTokenUsage, response)
+      .pipe(AzureOpenAiChatCompletionStream.processContent, response);
+    return response;
+  }
+
+  private async executeRequest(
+    data: AzureOpenAiCreateChatCompletionRequest,
+    requestConfig?: CustomRequestConfig
+  ): Promise<HttpResponse> {
     const deploymentId = await getDeploymentId(
       this.modelDeployment,
       'azure-openai'
     );
     const resourceGroup = getResourceGroup(this.modelDeployment);
-    const response = await executeRequest(
+    return executeRequest(
       {
         url: `/inference/deployments/${deploymentId}/chat/completions`,
         apiVersion,
@@ -42,6 +94,25 @@ export class AzureOpenAiChatClient {
       data,
       requestConfig
     );
-    return new AzureOpenAiChatCompletionResponse(response);
+  }
+
+  private async createStream(
+    data: AzureOpenAiCreateChatCompletionRequest,
+    requestConfig?: CustomRequestConfig
+  ): Promise<AzureOpenAiChatCompletionStream<any>> {
+    const response = await this.executeRequest(
+      {
+        ...data,
+        stream: true,
+        stream_options: {
+          include_usage: true
+        }
+      },
+      {
+        ...requestConfig,
+        responseType: 'stream'
+      }
+    );
+    return AzureOpenAiChatCompletionStream.create(response);
   }
 }

packages/foundation-models/src/azure-openai/azure-openai-chat-completion-stream-chunk-response.test.ts

@@ -0,0 +1,77 @@
+import { parseMockResponse } from '../../../../test-util/mock-http.js';
+import { AzureOpenAiChatCompletionStreamChunkResponse } from './azure-openai-chat-completion-stream-chunk-response.js';
+
+describe('OpenAI chat completion stream chunk response', () => {
+  let mockResponses: {
+    tokenUsage: any;
+    finishReason: any;
+    deltaContent: any;
+  };
+  let azureOpenAiChatCompletionStreamChunkResponses: {
+    tokenUsage: AzureOpenAiChatCompletionStreamChunkResponse;
+    finishReason: AzureOpenAiChatCompletionStreamChunkResponse;
+    deltaContent: AzureOpenAiChatCompletionStreamChunkResponse;
+  };
+
+  beforeAll(async () => {
+    mockResponses = {
+      tokenUsage: await parseMockResponse<any>(
+        'foundation-models',
+        'azure-openai-chat-completion-stream-chunk-response-token-usage.json'
+      ),
+      finishReason: await parseMockResponse<any>(
+        'foundation-models',
+        'azure-openai-chat-completion-stream-chunk-response-finish-reason.json'
+      ),
+      deltaContent: await parseMockResponse<any>(
+        'foundation-models',
+        'azure-openai-chat-completion-stream-chunk-response-delta-content.json'
+      )
+    };
+    azureOpenAiChatCompletionStreamChunkResponses = {
+      tokenUsage: new AzureOpenAiChatCompletionStreamChunkResponse(
+        mockResponses.tokenUsage
+      ),
+      finishReason: new AzureOpenAiChatCompletionStreamChunkResponse(
+        mockResponses.finishReason
+      ),
+      deltaContent: new AzureOpenAiChatCompletionStreamChunkResponse(
+        mockResponses.deltaContent
+      )
+    };
+  });
+
+  it('should return the chat completion stream chunk response', () => {
+    expect(
+      azureOpenAiChatCompletionStreamChunkResponses.tokenUsage.data
+    ).toStrictEqual(mockResponses.tokenUsage);
+    expect(
+      azureOpenAiChatCompletionStreamChunkResponses.finishReason.data
+    ).toStrictEqual(mockResponses.finishReason);
+    expect(
+      azureOpenAiChatCompletionStreamChunkResponses.deltaContent.data
+    ).toStrictEqual(mockResponses.deltaContent);
+  });
+
+  it('should get token usage', () => {
+    expect(
+      azureOpenAiChatCompletionStreamChunkResponses.tokenUsage.getTokenUsage()
+    ).toMatchObject({
+      completion_tokens: expect.any(Number),
+      prompt_tokens: expect.any(Number),
+      total_tokens: expect.any(Number)
+    });
+  });
+
+  it('should return finish reason', () => {
+    expect(
+      azureOpenAiChatCompletionStreamChunkResponses.finishReason.getFinishReason()
+    ).toBe('stop');
+  });
+
+  it('should return delta content with default index 0', () => {
+    expect(
+      azureOpenAiChatCompletionStreamChunkResponses.deltaContent.getDeltaContent()
+    ).toBe(' is');
+  });
+});

packages/foundation-models/src/azure-openai/azure-openai-chat-completion-stream-chunk-response.ts

@@ -0,0 +1,37 @@
+/**
+ * Azure OpenAI chat completion stream chunk response.
+ */
+export class AzureOpenAiChatCompletionStreamChunkResponse {
+  constructor(public readonly data: any) {
+    // TODO: Change `any` to `CreateChatCompletionStreamResponse` once the preview spec becomes stable.
+    this.data = data;
+  }
+
+  /**
+   * Usage of tokens in the chunk response.
+   * @returns Token usage.
+   */
+  getTokenUsage(): this['data']['usage'] {
+    return this.data.usage;
+  }
+
+  /**
+   * Reason for stopping the completion stream chunk.
+   * @param choiceIndex - The index of the choice to parse.
+   * @returns The finish reason.
+   */
+  getFinishReason(
+    choiceIndex = 0
+  ): this['data']['choices'][0]['finish_reason'] {
+    return this.data.choices[choiceIndex]?.finish_reason;
+  }
+
+  /**
+   * Parses the chunk response and returns the delta content.
+   * @param choiceIndex - The index of the choice to parse.
+   * @returns The message delta content.
+   */
+  getDeltaContent(choiceIndex = 0): string | undefined | null {
+    return this.data.choices[choiceIndex]?.delta?.content;
+  }
+}

packages/foundation-models/src/azure-openai/azure-openai-chat-completion-stream-response.ts

@@ -0,0 +1,44 @@
+import type { AzureOpenAiCompletionUsage } from './client/inference/schema/index.js';
+import type { AzureOpenAiChatCompletionStream } from './azure-openai-chat-completion-stream.js';
+
+/**
+ * Azure OpenAI chat completion stream response.
+ */
+export class AzureOpenAiChatCompletionStreamResponse<T> {
+  private _usage: AzureOpenAiCompletionUsage | undefined;
+  private _finishReason: string | undefined;
+  private _stream: AzureOpenAiChatCompletionStream<T> | undefined;
+
+  public get usage(): AzureOpenAiCompletionUsage {
+    if (!this._usage) {
+      throw new Error('Response stream is undefined.');
+    }
+    return this._usage;
+  }
+
+  public set usage(usage: AzureOpenAiCompletionUsage) {
+    this._usage = usage;
+  }
+
+  public get finishReason(): string {
+    if (!this._finishReason) {
+      throw new Error('Response finish reason is undefined.');
+    }
+    return this._finishReason;
+  }
+
+  public set finishReason(finishReason: string) {
+    this._finishReason = finishReason;
+  }
+
+  public get stream(): AzureOpenAiChatCompletionStream<T> {
+    if (!this._stream) {
+      throw new Error('Response stream is undefined.');
+    }
+    return this._stream;
+  }
+
+  public set stream(stream: AzureOpenAiChatCompletionStream<T>) {
+    this._stream = stream;
+  }
+}

packages/foundation-models/src/azure-openai/azure-openai-chat-completion-stream.test.ts

@@ -0,0 +1,39 @@
+import { parseFileToString } from "../../../../test-util/mock-http.js";
+import { AzureOpenAiChatCompletionStream } from "./azure-openai-chat-completion-stream.js";
+import { LineDecoder } from "./azure-openai-line-decoder.js";
+import { SSEDecoder } from "./azure-openai-sse-decoder.js";
+
+describe('OpenAI chat completion stream', () => {
+  let sseChunks: string[];
+  let originalChatCompletionStream: AzureOpenAiChatCompletionStream<any>;
+
+  beforeEach(async () => {
+    const rawChunksString = await parseFileToString(
+      'foundation-models',
+      'azure-openai-chat-completion-stream-chunks.txt'
+    );
+    const lineDecoder = new LineDecoder();
+    const sseDecoder = new SSEDecoder();
+    const rawLines: string[] = lineDecoder.decode(Buffer.from(rawChunksString, 'utf-8'));
+
+    sseChunks = rawLines
+      .map((chunk) => sseDecoder.decode(chunk))
+      .filter((sse) => sse !== null)
+      .filter((sse) => !sse.data.startsWith('[DONE]'))
+      .map((sse) => JSON.parse(sse.data));
+
+    async function *iterator(): AsyncGenerator<any> {
+      for (let sseChunk of sseChunks) {
+          yield sseChunk;
+      }
+    }
+    originalChatCompletionStream = new AzureOpenAiChatCompletionStream(iterator);
+  });
+
+  it('should wrap the raw chunk', async () => {
+    for await (const chunk of AzureOpenAiChatCompletionStream.processChunk(originalChatCompletionStream)) {
+      expect(chunk).toBeDefined();
+      console.log(chunk.getDeltaContent());
+    }
+  });
+});