handleDB.js

import mongoose from "mongoose";
import chalk from "chalk";
import Chat from "../models/Chat.js";
import { mapChatMessagesToStoredMessages } from "@langchain/core/messages";


/**
 * opens a connection to the database
 *
 * @return  {object}          database connection
 */
export const dbConnection = async () => {
  // skip if testing
  console.log(chalk.blue('connecting to db'));
  // else try to connect
  try {
    // prevent deprication warning
    mongoose.set('strictQuery', true);
    // connect
    mongoose.connect(process.env.MONGO_URL);
    // get connection
    const dbConnection = mongoose.connection;

    dbConnection.once("open", (_) => {
      console.log(chalk.green('DB connected successfully'));
    });

    dbConnection.on("error", (error) => {
      console.error(chalk.red(error));
    });
    return dbConnection;
  } catch (error) {
    console.error(error);
    process.exit(1);
  }
};


/**
 * opens a connection to the database
 * 
 * @param   {object}  res  responde object of express
 * @param   {object}  model  model to create record in, avail in /models
 * @param   {object}  data   data to create record with
 * 
 * @return  {object}          document
 */
export const createRecord = async (model, data) => {
  try {
    return await model.create(data);
  } catch (error) {
    throw error;
  }
};


/**
 * find one record
 *
 * @param   {mongoose model}  model         [required] model to search the record in
 * @param   {object}          searchObject  [optional] search object as filter, i.e. {email: 'a@b.c', name: 'John'}
 * @param   {string}          fieldHandler  [optional] additional fields to return or drop, i.e. '+password'
 *
 * @return  {object}              found document
 */
export const findOneRecord = async (model, searchObject = {}, fieldHandler = '') => {
  performance.mark('findOneRecord:start');
  try {
    let foundRecord;
    if (!fieldHandler) {
      foundRecord = await model.findOne(searchObject);
    } else {
      foundRecord = await model.findOne(searchObject).select(fieldHandler);
    }
    performance.mark('findOneRecord:end');
    return foundRecord;
  } catch (error) {
    throw error;
  }
};

/**
 * find record by ID
 *
 * @param   {mongoose model}  model         [required] model to search the record in
 * @param   {string}          ID            [required] ID to search the record with
 * @param   {string}          fieldHandler  [optional] additional fields to return or drop, i.e. '+password'
 *
 * @return  {object}              found document
 */
export const findRecordByID = async (model, id, fieldHandler = '') => {
  try {
    let foundRecord;
    if (!fieldHandler) {
      foundRecord = await model.findById(id);
    } else {
      foundRecord = await model.findById(id).select(fieldHandler);
    }
    return foundRecord;
  } catch (error) {
    throw error;
  }
};

/**
 * find multiple records
 *
 * @param   {mongoose model}  model         [required] model to search the record in
 * @param   {object}          searchObject  [optional] search object as filter, i.e. {email: 'a@b.c', name: 'John'}
 * @param   {string}          fieldHandler  [optional] additional fields to return or drop, i.e. '+password'
 *
 * @return  {object}              found documents
 */
export const findRecords = async (model, searchObject = {}, fieldHandler = '') => {
  try {
    let foundRecords;
    if (!fieldHandler) {
      foundRecords = await model.find(searchObject);
    } else {
      foundRecords = await model.find(searchObject).select(fieldHandler);
    }
    return foundRecords;
  } catch (error) {
    throw error;
  }
};

/**
 * update one record
 *
 * @param   {mongoose document}   newData  a mongoose document with new data
 *
 * @return  {mongoose document}             updated document
 */
export const updateOneRecord = async (newData) => {
  try {
    // drop password if empty to prevent setting empty
    if (!newData.password) delete newData.password;
    return await newData.save({ new: true });
  } catch (error) {
    throw error;
  }
};

/**
 * Find a document by id and update it
 *
 * @param   {mongoose model}  model         [required] model to search the record in
 * @param   {string}          id            [required] id of the searched record
 * @param   {object}          data          [required] data to update the record with
 *
 * @return  {object}         the edited document
 */
export const findByIdAndUpdate = async (model, id, data) => {
  try {
    performance.mark('findByIdAndUpdate:start');
    const result = model.findByIdAndUpdate(id, data, { returnDocument: 'after' });
    performance.mark('findByIdAndUpdate:end');
    return result;
  } catch (error) {
    throw error;
  }
};


/**
 * Find a document by id and update it
 *
 * @param   {mongoose model}  model         [required] model to search the record in
 * @param   {object}          searchObject  [required] search object as filter, i.e. {email: 'a@b.c', name: 'John'}
 * @param   {object}          data          [required] data to update the record with
 *
 * @return  {object}         the edited document
 */
export const findOneAndUpdate = async (model, searchObject, data) => {
  try {
    return model.updateOne(searchObject, data);
  } catch (error) {
    throw error;
  }
};

/**
 * extend a chat record with a new input/response-pair
 *
 * @param   {string}  chatId    [required] chatId to extend
 * @param   {array}   messages  [required] array of message objects to extend the chat with
 *
 * @return  {object}            the edited document
 */
export const extendChat = async (chatId, messages) => {
  performance.mark('extendChat:start');
  // exit if no chatId
  if (!chatId) {
    console.error('No chatId provided');
    return;
  }
  try {
    // fetch chat record
    const record = await findRecordByID(Chat, chatId);

    // push new message into chat history
    const serializedMessages = mapChatMessagesToStoredMessages(messages);
    record.chatHistory ? record.chatHistory.push(...serializedMessages) : record.chatHistory = serializedMessages;
    performance.mark('extendChat:end');
    // save & return chat
    return await findByIdAndUpdate(Chat, chatId, { chatHistory: record.chatHistory });

  } catch (error) {
    throw error;
  }
};