Update package versions and add support for retrieving a specific num…

n4ze3m · Apr 26, 2024 · efc922d · efc922d
1 parent e95ee8a
commit efc922d
Show file tree

Hide file tree

Showing 9 changed files with 86 additions and 39 deletions.
diff --git a/app/ui/package.json b/app/ui/package.json
@@ -1,7 +1,7 @@
 {
   "name": "app",
   "private": true,
-  "version": "1.8.1",
+  "version": "1.8.2",
   "type": "module",
   "scripts": {
     "dev": "vite",

diff --git a/app/ui/src/@types/bot.ts b/app/ui/src/@types/bot.ts
@@ -6,6 +6,7 @@ export type BotSettings = {
   public_id: string;
   temperature: number;
   embedding: string;
+  noOfDocumentsToRetrieve: number;
   qaPrompt: string;
   questionGeneratorPrompt: string;
   streaming: boolean;

diff --git a/app/ui/src/components/Bot/Settings/SettingsCard.tsx b/app/ui/src/components/Bot/Settings/SettingsCard.tsx
@@ -1,4 +1,12 @@
-import { Form, Input, notification, Select, Slider, Switch } from "antd";
+import {
+  Form,
+  Input,
+  InputNumber,
+  notification,
+  Select,
+  Slider,
+  Switch,
+} from "antd";
 import { useNavigate, useParams } from "react-router-dom";
 import api from "../../../services/api";
 import { useMutation, useQueryClient } from "@tanstack/react-query";
@@ -155,6 +163,7 @@ export const SettingsCard: React.FC<BotSettings> = ({
             bot_protect: data.bot_protect,
             use_rag: data.use_rag,
             bot_model_api_key: data.bot_model_api_key,
+            noOfDocumentsToRetrieve: data.noOfDocumentsToRetrieve,
           }}
           form={form}
           requiredMark={false}
@@ -260,6 +269,24 @@ export const SettingsCard: React.FC<BotSettings> = ({
                   />
                 </Form.Item>
 
+                <Form.Item
+                  name="noOfDocumentsToRetrieve"
+                  label="Number of documents to retrieve"
+                  rules={[
+                    {
+                      required: true,
+                      message:
+                        "Please input a number of documents to retrieve!",
+                    },
+                  ]}
+                >
+                  <InputNumber
+                    min={0}
+                    style={{ width: "100%" }}
+                    placeholder="Enter number of documents to retrieve"
+                  />
+                </Form.Item>
+
                 <Form.Item
                   label={"Question Answering Prompt (System Prompt)"}
                   name="qaPrompt"

diff --git a/package.json b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "dialoqbase",
-  "version": "1.8.1",
+  "version": "1.8.2",
   "description": "Create chatbots with ease",
   "scripts": {
     "ui:dev": "pnpm run --filter ui dev",

diff --git a/server/prisma/migrations/q_26/migration.sql b/server/prisma/migrations/q_26/migration.sql
@@ -0,0 +1,2 @@
+-- AlterTable
+ALTER TABLE "Bot" ADD COLUMN     "noOfDocumentsToRetrieve" INTEGER DEFAULT 4;
diff --git a/server/prisma/schema.prisma b/server/prisma/schema.prisma
@@ -18,6 +18,7 @@ model Bot {
   description                 String?
   createdAt                   DateTime         @default(now())
   temperature                 Float            @default(0.7)
+  noOfDocumentsToRetrieve     Int?             @default(4)
   model                       String           @default("gpt-3.5-turbo")
   provider                    String           @default("openai")
   embedding                   String           @default("openai")

diff --git a/server/src/schema/api/v1/bot/bot/index.ts b/server/src/schema/api/v1/bot/bot/index.ts
@@ -159,6 +159,9 @@ export const updateBotByIdSchema: FastifySchema = {
       bot_model_api_key: {
         type: "string",
       },
+      noOfDocumentsToRetrieve: {
+        type: "number",
+      }
     },
   },
 };

diff --git a/server/src/utils/hybrid.ts b/server/src/utils/hybrid.ts
@@ -40,47 +40,53 @@ export class DialoqbaseHybridRetrival extends BaseRetriever {
   protected async similaritySearch(
     query: string,
     k: number,
-    _callbacks?: Callbacks,
+    _callbacks?: Callbacks
   ): Promise<SearchResult[]> {
     try {
+      const embeddedQuery = await this.embeddings.embedQuery(query);
 
-    const embeddedQuery = await this.embeddings.embedQuery(query);
+      const vector = `[${embeddedQuery.join(",")}]`;
+      const bot_id = this.botId;
 
-    const vector = `[${embeddedQuery.join(",")}]`;
-    const bot_id = this.botId;
-
-    const data = await prisma.$queryRaw`
+      const data = await prisma.$queryRaw`
      SELECT * FROM "similarity_search_v2"(query_embedding := ${vector}::vector, botId := ${bot_id}::text,match_count := ${k}::int)
     `;
 
-    const result: [Document, number, number][] = (
-      data as SearchEmbeddingsResponse[]
-    ).map((resp) => [
-      new Document({
-        metadata: resp.metadata,
-        pageContent: resp.content,
-      }),
-      resp.similarity * 10,
-      resp.id,
-    ]);
-
-
-    return result;
-} catch (e) {
-    console.log(e)
-    return []
-}
+      const result: [Document, number, number][] = (
+        data as SearchEmbeddingsResponse[]
+      ).map((resp) => [
+        new Document({
+          metadata: resp.metadata,
+          pageContent: resp.content,
+        }),
+        resp.similarity * 10,
+        resp.id,
+      ]);
+
+      return result;
+    } catch (e) {
+      console.log(e);
+      return [];
+    }
   }
 
   protected async keywordSearch(
     query: string,
-    k: number,
+    k: number
   ): Promise<SearchResult[]> {
     const query_text = query;
     const bot_id = this.botId;
 
+    const botInfo = await prisma.bot.findFirst({
+      where: {
+        id: bot_id,
+      },
+    });
+
+    const match_count = botInfo?.noOfDocumentsToRetrieve || k;
+
     const data = await prisma.$queryRaw`
-     SELECT * FROM "kw_match_documents"(query_text := ${query_text}::text, bot_id := ${bot_id}::text,match_count := ${k}::int)
+     SELECT * FROM "kw_match_documents"(query_text := ${query_text}::text, bot_id := ${bot_id}::text,match_count := ${match_count}::int)
     `;
 
     const result: [Document, number, number][] = (
@@ -104,12 +110,12 @@ export class DialoqbaseHybridRetrival extends BaseRetriever {
     query: string,
     similarityK: number,
     keywordK: number,
-    callbacks?: Callbacks,
+    callbacks?: Callbacks
   ): Promise<SearchResult[]> {
     const similarity_search = this.similaritySearch(
       query,
       similarityK,
-      callbacks,
+      callbacks
     );
 
     const keyword_search = this.keywordSearch(query, keywordK);
@@ -136,13 +142,13 @@ export class DialoqbaseHybridRetrival extends BaseRetriever {
 
   async _getRelevantDocuments(
     query: string,
-    runManager?: CallbackManagerForRetrieverRun,
+    runManager?: CallbackManagerForRetrieverRun
   ): Promise<Document[]> {
     const searchResults = await this.hybridSearch(
       query,
       this.similarityK,
       this.keywordK,
-      runManager?.getChild("hybrid_search"),
+      runManager?.getChild("hybrid_search")
     );
 
     return searchResults.map(([doc]) => doc);

diff --git a/server/src/utils/store.ts b/server/src/utils/store.ts
@@ -40,8 +40,7 @@ export class DialoqbaseVectorStore extends VectorStore {
         if (row?.embedding) {
           const vector = `[${row.embedding.join(",")}]`;
           const content = row?.content.replace(/\x00/g, "").trim();
-          await prisma
-            .$executeRaw`INSERT INTO "BotDocument" ("content", "embedding", "metadata", "botId", "sourceId") VALUES (${content}, ${vector}::vector, ${row.metadata}, ${row.botId}, ${row.sourceId})`;
+          await prisma.$executeRaw`INSERT INTO "BotDocument" ("content", "embedding", "metadata", "botId", "sourceId") VALUES (${content}, ${vector}::vector, ${row.metadata}, ${row.botId}, ${row.sourceId})`;
         }
       });
     } catch (e) {
@@ -57,7 +56,7 @@ export class DialoqbaseVectorStore extends VectorStore {
   static async fromDocuments(
     docs: Document[],
     embeddings: Embeddings,
-    dbConfig: DialoqbaseLibArgs,
+    dbConfig: DialoqbaseLibArgs
   ) {
     const instance = new this(embeddings, dbConfig);
     await instance.addDocuments(docs);
@@ -68,7 +67,7 @@ export class DialoqbaseVectorStore extends VectorStore {
     texts: string[],
     metadatas: object[] | object,
     embeddings: Embeddings,
-    dbConfig: DialoqbaseLibArgs,
+    dbConfig: DialoqbaseLibArgs
   ) {
     const docs = [];
     for (let i = 0; i < texts.length; i += 1) {
@@ -84,7 +83,7 @@ export class DialoqbaseVectorStore extends VectorStore {
 
   static async fromExistingIndex(
     embeddings: Embeddings,
-    dbConfig: DialoqbaseLibArgs,
+    dbConfig: DialoqbaseLibArgs
   ) {
     const instance = new this(embeddings, dbConfig);
     return instance;
@@ -93,14 +92,22 @@ export class DialoqbaseVectorStore extends VectorStore {
   async similaritySearchVectorWithScore(
     query: number[],
     k: number,
-    filter?: this["FilterType"] | undefined,
+    filter?: this["FilterType"] | undefined
   ): Promise<[Document<Record<string, any>>, number][]> {
     console.log(this.botId);
     const vector = `[${query.join(",")}]`;
     const bot_id = this.botId;
 
+    const botInfo = await prisma.bot.findFirst({
+      where: {
+        id: bot_id,
+      },
+    });
+
+    const match_count = botInfo?.noOfDocumentsToRetrieve || k;
+
     const data = await prisma.$queryRaw`
-     SELECT * FROM "similarity_search_v2"(query_embedding := ${vector}::vector, botId := ${bot_id}::text,match_count := ${k}::int)
+     SELECT * FROM "similarity_search_v2"(query_embedding := ${vector}::vector, botId := ${bot_id}::text,match_count := ${match_count}::int)
     `;
 
     const result: [Document, number][] = (