Build a RAG command line chatbot

npm install @datastax/astra-db-ts openai

import { DataAPIClient } from "@datastax/astra-db-ts";
import OpenAI from "openai";

import { createInterface } from "node:readline/promises";
import { stdin as input, stdout as output } from "node:process";

export async function main() {
  const {
    API_ENDPOINT: endpoint,
    APPLICATION_TOKEN: applicationToken,
    OPENAI_API_KEY: openaiApiKey,
  } = process.env; (1)
  const keyspace = "default_keyspace"; (2)
  const collectionName = "quickstart_collection"; (3)

  if (!endpoint || !applicationToken || !openaiApiKey) {
    throw new Error(
      "Environment variables API_ENDPOINT, APPLICATION_TOKEN, OPENAI_API_KEY must be defined.",
    );
  }

  // Instantiate a `DataAPIClient` and get a reference to your collection
  const client = new DataAPIClient();
  const database = client.db(endpoint, { token: applicationToken, keyspace });
  const collection = database.collection(collectionName);

  // Instantiate the OpenAI client
  const openai = new OpenAI({
    apiKey: openaiApiKey,
  });

  // This list of messages will be sent to OpenAI with every query.
  // It starts with a single system prompt and grows as the chat progresses.
  const messages: OpenAI.ChatCompletionMessageParam[] = [
    {
      role: "system",
      content:
        "You are an AI assistant that can answer questions based on the the context you are given. Don't mention the context, just use it to inform your answers.",
    },
  ];

  // Use the built-in Node.js readline to implement the CLI
  const cli = createInterface({ input, output });

  // Start the chat by writing a message to the CLI
  let userInput = await cli.question(
    `Greetings! I am an AI assistant that is ready to help you with your questions. You can ask me anything you like.\nIf you want to exit, type ".exit".\n\n> `,
  );

  // Run this loop continuously until the user inputs the exit command
  while (userInput.toLowerCase() !== ".exit") {
    // If the user didn't input text, re-prompt them
    if (userInput.trim() === "") {
      userInput = await cli.question("> ");
      continue;
    }

    try {
      // Perform a vector search in your collection,
      // using the user input as the search string to vectorize.
      // Limit the search to 10 documents.
      // Use a projection to return just the $vectorize field of each document.
      const response = collection.find(
        {},
        {
          sort: { $vectorize: userInput },
          limit: 10,
          projection: { $vectorize: 1 },
        },
      );

      // Join the $vectorize fields of the returned documents into a single string
      const context = (await response.toArray())
        .map((doc) => doc.$vectorize)
        .join("\n");

      // Combine the user question with the context from vector search
      const ragMessage: OpenAI.ChatCompletionUserMessageParam = {
        role: "user",
        content: `${context}\n---\nGiven the above context, answer the following question:\n${userInput}`,
      };

      // Send the list of previous messages, plus the context augmented question
      const stream = await openai.chat.completions.create({
        model: "gpt-4o-mini",
        messages: [...messages, ragMessage],
        stream: true,
      });

      // Write OpenAI's response to the CLI as it comes in,
      // and also record it in a string
      let message = "";
      for await (const chunk of stream) {
        const delta = chunk.choices[0]?.delta?.content ?? "";
        output.write(delta);
        message += delta;
      }

      // Record the user question, without the added context, in the list of messages
      messages.push({ role: "user", content: userInput });

      // Record the OpenAI response in the list of messages
      messages.push({ role: "assistant", content: message });

      // Prompt the user for their next question
      userInput = await cli.question("\n\n> ");
    } catch (error) {
      if (error instanceof Error) {
        console.error(error.message);
      }
      userInput = await cli.question(
        "\nSomething went wrong, try asking again\n\n> ",
      );
    }
  }

  cli.close();
}

main().catch((err) => {
  console.error(err);
  process.exit(1);
});