@nomyx/assistant/src/ai/providers/OpenAIProvider/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, ActionContext } from '../../../types/tool';
import { OpenAIConfig } from './config';
import { convertMessages, convertToolSchema, convertToolCalls } from './utils';

export async function* streamChat(config: OpenAIConfig, messages: ChatMessage[], options: ChatOptions, tools?: Tool[]): AsyncIterableIterator<ProviderResponse> {
  const logger = config.getLogger();

  try {
    const requestBody: any = {
      model: config.getModel(),
      messages: convertMessages(messages, config.getModel()),
      max_tokens: options.maxTokens,
      temperature: options.temperature,
      stream: true,
    };

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

    console.log('Sending request to OpenAI API', { 
      model: requestBody.model, 
      messageCount: requestBody.messages.length,
      maxTokens: requestBody.max_tokens,
      temperature: requestBody.temperature,
      toolCount: tools?.length
    });

    const response = await axios.post(
      'https://api.openai.com/v1/chat/completions',
      requestBody,
      {
        headers: {
          'Authorization': `Bearer ${config.getApiKey()}`,
          'Content-Type': 'application/json',
        },
        responseType: 'stream',
      }
    );

    console.log('Received response from OpenAI API', { status: response.status });

    let buffer = '';
    let currentToolCall: any = null;

    // Create a basic ActionContext
    const actionContext: ActionContext = {
      user: { id: 'default-user', name: 'Default User' },
      session: { id: 'default-session', timestamp: new Date().toISOString() },
    };

    for await (const chunk of response.data) {
      console.log('Received chunk', { chunkSize: chunk.length });
      buffer += chunk.toString();
      let newlineIndex;
      while ((newlineIndex = buffer.indexOf('\n')) !== -1) {
        const line = buffer.slice(0, newlineIndex).trim();
        buffer = buffer.slice(newlineIndex + 1);

        if (line === 'data: [DONE]') {
          console.log('Received [DONE] signal');
          return;
        }

        if (line.startsWith('data: ')) {
          try {
            const data = JSON.parse(line.slice(6));

            const content = data.choices?.[0]?.delta?.content || '';
            const toolCalls = data.choices?.[0]?.delta?.tool_calls;
            
            console.log(content);
            console.log(toolCalls);

            if (content) {
              yield { content, toolCalls: [], usage: { promptTokens: 0, completionTokens: 0, totalTokens: 0 } };
            }

            if (toolCalls) {
              for (const toolCall of toolCalls) {
                if (!currentToolCall || toolCall.index === 0) {
                  currentToolCall = { ...toolCall, function: { ...toolCall.function } };
                } else {
                  if (toolCall.function.name) {
                    currentToolCall.function.name = toolCall.function.name;
                  }
                  if (toolCall.function.arguments) {
                    currentToolCall.function.arguments = (currentToolCall.function.arguments || '') + toolCall.function.arguments;
                  }
                }
              }

              process.stdout.write(data.choices[0].delta.content);

              if (data.choices[0].finish_reason === 'tool_calls') {
                try {
                  const fullToolCall = {
                    name: currentToolCall.function.name,
                    arguments: currentToolCall.function.arguments
                  };
                  console.log('Full tool call:', fullToolCall);
                  
                  // Execute the tool
                  const tool = tools?.find(t => t.name === fullToolCall.name);
                  if (tool) {
                    const result = await tool.execute(fullToolCall.arguments, actionContext);
                    console.log('Tool execution result:', result);
                    yield { content: `Tool ${fullToolCall.name} executed. Result: ${JSON.stringify(result)}`, toolCalls: [fullToolCall], usage: { promptTokens: 0, completionTokens: 0, totalTokens: 0 } };
                  } else {
                    console.error(`Tool ${fullToolCall.name} not found`);
                    yield { content: `Error: Tool ${fullToolCall.name} not found`, toolCalls: [fullToolCall], usage: { promptTokens: 0, completionTokens: 0, totalTokens: 0 } };
                  }
                  
                  currentToolCall = null;
                } catch (parseError) {
                //   console.error('Error parsing or executing tool call:', parseError);
                //   console.error('Raw tool call arguments:', currentToolCall.function.arguments);
                  yield { content: `Error executing tool: ${parseError instanceof Error ? parseError.message : String(parseError)}`, toolCalls: [], usage: { promptTokens: 0, completionTokens: 0, totalTokens: 0 } };
                }
              }
            }
          } catch (parseError) {
            // console.error('Error parsing JSON from OpenAI API:', { line, error: parseError });
            // Continue to the next line instead of throwing an error
            continue;
          }
        }
      }
    }
  } catch (error) {
    handleError(config, error);
    throw error;
  }
}

function handleError(config: OpenAIConfig, error: any): void {
  if (axios.isAxiosError(error)) {
    console.error('OpenAIProvider: API request failed:', {
      status: error.response?.status,
      statusText: error.response?.statusText,
      data: error.response?.data,
      headers: error.response?.headers,
      config: error.config
    });
  } else {
    console.error('OpenAIProvider: Unexpected error:', error);
  }
}