refraction-cli/src/index.ts

AI code generation for your terminal.

import { program } from 'commander';
import chalk from 'chalk';
import { encode } from '@beskar-labs/gpt-encoder';
import pkg from '../package.json';
import { setConfig } from './lib/config';
import { getStream } from './lib/stream';
import { parseError } from './lib/helpers';
import authenticate from './lib/authenticate';

const { log } = console;

program.name('refraction').description(pkg.description).version(pkg.version);

program
  .command('generate')
  .description('Generate a Bash command using AI')
  .argument('<command>', 'command to generate')
  .option(
    '--language <language>',
    'language to generate the command in',
    'shell'
  )
  .action(async (code: string, options: { language: string }) => {
    if (!code) {
      log(chalk.red('No code provided.'));
      return;
    }

    const tokens = encode(code).length;

    if (tokens > 3000) {
      log(
        chalk.red(
          `Sorry, but the code you've selected is too long (${tokens} tokens). Please reduce the length of your code to 3000 tokens or less.`
        )
      );
    }

    try {
      await authenticate();

      const data = await getStream(options.language, code);

      data.on('readable', () => {
        const chunk: string | Buffer | null = data.read();

        if (!chunk) {
          return;
        }

        const value = typeof chunk === 'string' ? chunk : chunk.toString();
        process.stdout.write(value);
      });

      data.on('finish', () => {
        process.stdout.write('\n');
      });
    } catch (error) {
      const message = parseError(error);

      if (message === 'Request failed with status code 401') {
        log(
          chalk.red(
            'Your Refraction User ID or Team ID are incorrect. Please run `refraction login` to update them.'
          )
        );
        return;
      }

      log(chalk.red(message));
    }
  });

program
  .command('login')
  .description('Add Refraction User and Team IDs')
  .argument('<userId>', 'Your User ID')
  .argument('[teamId]', 'Your Team ID')
  .action((userId: string, teamId: string) => {
    setConfig('userId', userId);
    setConfig('teamId', teamId);

    log(chalk.magenta('Successfully saved your Refraction details!'));
  });

program.parse();