feat: vector search

Author: fenil3357

controllers/insertContent.js

@@ -1,20 +1,17 @@
 import { Document } from "../models/document.js";
-import { openAI } from "../openai.js";
+import { generateEmbeddings, openAI } from "../openai.js";
 
 export const generateEmbeddingsController = async (req, res) => {
   try {
     const { content } = req?.body; // You can pass the embeddings manually in request body if you don't want to use openai
 
     // Generate embeddings of given text content using openAI's text-embedding-ada-002 embedding model
-    const embeddingResponse = await openAI.embeddings.create({
-      input: content,
-      model: 'text-embedding-ada-002'
-    });
+    const embeddings = await generateEmbeddings(content);
 
     // Save embeddings in DB
     const document = await Document.create({
       content,
-      embeddings: embeddingResponse.data[0].embedding
+      embeddings
     }, { raw: true });
 
     return res.json({

controllers/vectorSearch.js

@@ -0,0 +1,34 @@
+import { DB } from "../db.js";
+import { generateEmbeddings } from "../openai.js";
+
+export const vectorSearchController = async (req, res) => {
+  try {
+    let { query, match_count, similarity_threshold } = req?.body; // You can pass the embeddings manually in request body if you don't want to use openai
+
+    if (!match_count) match_count = 3;
+    if (!similarity_threshold) similarity_threshold = 0.01
+
+    // Generate the embeddings of given query
+    const query_embeddings = await generateEmbeddings(query);
+
+    // convert embeddings into postgres vector type
+    const query_embeddings_vector = `[${query_embeddings.join(', ')}]`
+
+    // Invoke the vector_search function to find the similar content
+    const data = await DB.query(
+      `SELECT id, content, similarity FROM vector_search(:query_embeddings, :similarity_threshold, :match_count)`,
+      {
+        replacements: {
+          query_embeddings: query_embeddings_vector,
+          similarity_threshold,
+          match_count
+        },
+      }
+    );
+
+    return res.json(data[0]);
+  } catch (error) {
+    console.log("🚀 ~ vectorSearchController ~ error:", error)
+    return res.status(500).json({ error })
+  }
+}

db.js

@@ -3,5 +3,6 @@ import * as dotenv from 'dotenv'
 dotenv.config()
 
 export const DB = new Sequelize(process.env.DB_CONNECTION_URL, {
-  dialect: 'postgres'
+  dialect: 'postgres',
+  logging: false
 });

openai.js

@@ -4,4 +4,18 @@ dotenv.config();
 
 export const openAI = new OpenAI({
   apiKey: process.env.OPENAI_API_KEY
-})
+})
+
+export const generateEmbeddings = async (content) => {
+  try {
+    const embeddingResponse = await openAI.embeddings.create({
+      input: content,
+      model: 'text-embedding-ada-002'
+    });
+
+    return embeddingResponse.data[0].embedding;
+  } catch (error) {
+    console.log("🚀 ~ generateEmbeddings ~ error:", error)
+    throw new Error(error?.message || 'Something went wrong while generating embeddings');
+  }
+}

routes.js

@@ -1,8 +1,10 @@
 import { Router } from "express";
 import { generateEmbeddingsController } from "./controllers/insertContent.js";
+import { vectorSearchController } from "./controllers/vectorSearch.js";
 
 const indexRouter = Router();
 
 indexRouter.post('/generate-embeddings', generateEmbeddingsController);
+indexRouter.post('/vector-search', vectorSearchController);
 
 export default indexRouter;

sql/3.define-vector-search-function.sql renamed to sql/3.create-vector-search-function.sql

@@ -15,7 +15,7 @@ as $$
     Documents.content,
     1 - (Documents.embeddings <=> query_embeddings) as similarity
   from Documents
-  where (1 - Documents.embeddings <=> query_embeddings) > similarity_threshold
+  where (1 - (Documents.embeddings <=> query_embeddings)) > similarity_threshold
   order by Documents.embeddings <=> query_embeddings
   limit match_count;
 $$;