ggai/openai.js

OpenAI LLM Agent Interface

const axios = require('axios');
const { estimateTokenLength } = require('./pricing.js');
const chalk = require('chalk');

const getInitialState = () => {
  return {
    usage: {
      promptTokens: 0,
      completionTokens: 0,
      totalTokens: 0,
    }
  };
};

let state = getInitialState();
let incompleteChunk = '';

const processChunk = (chunk, streamingIndicator = false) => {
  if (chunk === '[DONE]') {
    const _state = Object.assign({}, state);
    state = getInitialState();
    return {
      state: _state, final: true
    }
  }
  if (incompleteChunk) {
    chunk = incompleteChunk + chunk;
    incompleteChunk = '';
  }
  try {
    chunk = JSON.parse(chunk);
  } catch (err) {
    incompleteChunk = chunk;
    return {
      state, final: false
    }
  }

  for (const key in chunk) {
    if (typeof chunk[key] === 'string') {
      state[key] = chunk[key];
    } else {
      if (key === 'choices') {
        let item = {};
        const choice = chunk[key][0];
        for (const key in choice) {
          if (key !== 'delta') {
            item[key] = choice[key];
          } else {
            item = { ...item, ...choice.delta }
          }
        }
        if (!state?.choices) {
          state['choices'] = [item];
        }
        for (const key in item) {
          if (key in state.choices[0] && item[key] !== null) {
            if (key === 'index') {

              if ('tool_calls' in item && !('tool_calls' in state.choices[0])) {
                state.choices[0].tool_calls = [];
              }
              if ('tool_calls' in item) {
                if (streamingIndicator) {
                  process.stdout.write(chalk.yellow('.'));
                }
                for (const toolCall of item.tool_calls) {
                  if (!state.choices[0].tool_calls[toolCall.index]) {
                    state.choices[0].tool_calls[toolCall.index] = {
                      ...toolCall,
                      index: undefined
                    };
                  } else {
                    state.choices[0].tool_calls[toolCall.index].function.arguments += toolCall.function.arguments;
                  }
                }
              }
            }
            if (key === 'content') {
              state.choices[0][key] += item[key];
              if (streamingIndicator) {
                process.stdout.write(chalk.cyan(item[key]));
              }
            }
          }
        }
      }
    }
  }

  const final = Boolean(chunk.choices[0].finish_reason);

  if (final) {
    const _state = Object.assign({}, state);
    state = getInitialState();
    return {
      state: _state, final
    }
  }

  return {
    state, final
  }
}

/**
 * Streaming request
 * Makes a request to the OpenAI API.
 * @param {string} endpoint - The API endpoint to call.
 * @param {object} data - The data to send in the request.
 * @returns {Promise<object>} - The response from the API.
 * @throws {Error} Error making OpenAI API request
 */
const streamingRequest = async (endpoint, data, streamingIndicators = false, debug = false) => {
  return new Promise(async (resolve, reject) => {

    if (!process.env.OPENAI_API_KEY) {
      return reject(new Error('OPENAI_API_KEY environment variable not set'))
    }

    if (debug) {
      console.log(chalk.bgCyan(JSON.stringify(data, null, 2)));
    }

    const config = {
      headers: {
        'Authorization': `Bearer ${process.env.OPENAI_API_KEY}`
      },
      responseType: 'stream'
    };
    try {
      const response = await axios.post(`https://api.openai.com/v1/${endpoint}`, {
        ...data,
        stream: true
      }, config);
      response.data.on('data', raw => {
        const str = raw.toString().replace(/data: /g, '');
        const chunks = str.split('\n').filter(Boolean);
        for (const chunk of chunks) {
          const { state, final } = processChunk(chunk, streamingIndicators);
          if (final) {
            const promptLength = estimateTokenLength(JSON.stringify(data));
            const completionLength = estimateTokenLength(JSON.stringify(state));
            resolve({
              ...state,
              usage: {
                prompt_tokens: promptLength,
                completion_tokens: completionLength,
                total_tokens: promptLength + completionLength,
              }
            });
          }

        }
      });
      return response.data;
    } catch (error) {
      if (error.response) {
        console.log(error.response.status);
      }
      throw new Error("Error making OpenAI API request");
    }
  });
};

const getChatCompletion = (data, streamingIndicators = true, debug = false) => {
  return streamingRequest('/chat/completions', data, streamingIndicators, debug);
}

module.exports = {
  getChatCompletion
};