Simple vector db embeddings.

vector-db.js

import readline from 'readline'
import chalk from 'chalk'
import OpenAI from 'openai'
import { LocalIndex } from 'vectra'
import { fileURLToPath } from 'url'
import path, { dirname } from 'path'

const __filename = fileURLToPath(import.meta.url)
const __dirname = dirname(__filename)

const openAi = new OpenAI({
  apiKey: 'OPENAI_API_KEY',
})

const index = new LocalIndex(path.join(__dirname, '..', 'index'))

const getVector = async (input) => {
  try {
    const response = await openAi.embeddings.create({
      model: 'text-embedding-ada-002',
      input,
    })

    return response.data[0].embedding
  } catch (e) {
    console.log('_ERROR_', e)
  }

  return null
}

const createIndex = async () => {
  if (!await index.isIndexCreated()) {
    await index.createIndex()
  }

  return index
}

const addItem = async (text) => {
  await index.insertItem({
    vector: await getVector(text),
    metadata: { text },
  })
}

const query = async (text) => {
  const vector = await getVector(text)
  const results = await index.queryItems(vector, 3)
  if (results.length > 0) {
    for (const result of results) {
      console.log(`[${result.score}] ${result.item.metadata.text}`)
    }
  } else {
    console.log(`No results found.`)
  }

  return results
}

const rl = readline.createInterface({
  input: process.stdin,
  output: process.stdout,
  prompt: chalk.greenBright('chatbot> ')
})

console.log(chalk.bold.bgBlue.white(' Welcome to the Chatbot '))
console.log(chalk.gray('Type "store" to save a document. Type "query" to find documents. Type "help" for more commands.'))

await createIndex()

rl.prompt()

rl.on('line', async (line) => {
  const input = line.trim().split(' ')

  switch (input[0]) {
    case 'store':
      if (input.length < 2) {
        console.log(chalk.red('Please provide a document to store.'))
      } else {
        const doc = input.slice(1).join(' ')

        const results = await addItem(doc)

        console.log(results)
        console.log(chalk.green('Document stored!'));
      }
      break

    case 'query':
      const queryText = input.slice(1).join(' ')

      const results = await query(queryText)

      if (results && results.length > 0) {
        console.log(chalk.cyan('Matching documents:'))
        results.forEach(doc => {
          console.log(chalk.cyan(`- ${doc.item.metadata.text}`))
        })
      } else {
        console.log(chalk.red('No matching documents found.'))
      }
      break

    case 'help':
      console.log(chalk.cyan('Available commands: '))
      console.log(chalk.cyan('- store [document]: To store a document.'))
      console.log(chalk.cyan('- query [text]: To query the stored documents.'))
      break

    case 'exit':
      console.log(chalk.yellow('Goodbye!'))
      process.exit(0)
      break

    default:
      console.log(chalk.red('Unknown command. Type "help" for available commands.'))
  }

  rl.prompt()
}).on('close', () => {
  console.log(chalk.yellow('Goodbye!'))
  process.exit(0)
})