@nomyx/assistant/src/ai/providers/AzureProvider/streaming.ts

A powerful assistant library and cli for your AI projects. works with Vertex AI (Claude and Gemini)

import axios from 'axios';
import { ChatMessage, ChatOptions, ProviderResponse } from '../../../types/chat';
import { Tool } from '../../../types/tool';
import { AzureConfig } from './config';
import { AzureChatRequestBody, AzureChatResponse } from './types';
import { convertToolSchema, convertMessages, convertToolCall } from './tools';

export async function* streamChat(
  config: AzureConfig,
  messages: ChatMessage[],
  options: ChatOptions,
  tools?: Tool[]
): AsyncIterableIterator<ProviderResponse> {
  const url = config.getApiUrl('chat/completions');
  
  const requestBody: AzureChatRequestBody = {
    messages: convertMessages(messages),
    max_tokens: options.maxTokens,
    temperature: options.temperature,
    stream: true,
  };

  if (tools && tools.length > 0) {
    requestBody.tools = tools.map(tool => convertToolSchema(tool.getSchema()));
    requestBody.tool_choice = 'auto';
  }

  const response = await axios.post<NodeJS.ReadableStream>(url, requestBody, {
    headers: {
      'Content-Type': 'application/json',
      'api-key': config.apiKey,
    },
    responseType: 'stream',
  });

  for await (const chunk of response.data) {
    const lines = chunk.toString().split('\n').filter((line: string) => line.trim() !== '');
    for (const line of lines) {
      if (line.includes('[DONE]')) return;
      if (line.startsWith('data: ')) {
        const data = JSON.parse(line.slice(6)) as AzureChatResponse;
        yield createStreamingProviderResponse(data);
      }
    }
  }
}

function createStreamingProviderResponse(data: AzureChatResponse): ProviderResponse {
  const choice = data.choices && data.choices.length > 0 ? data.choices[0] : null;
  if (!choice) {
    throw new Error('Unexpected response format from Azure OpenAI API');
  }

  const response: ProviderResponse = {
    content: choice.delta?.content || '',
    toolCalls: [],
    usage: { promptTokens: 0, completionTokens: 0, totalTokens: 0 },
  };

  if (choice.delta?.function_call) {
    response.toolCalls = [convertToolCall(choice.delta.function_call)];
  }

  return response;
}