Skip to content

Commit efc922d

Browse files
n4ze3mn4ze3m
committed
Update package versions and add support for retrieving a specific number of documents
1 parent e95ee8a commit efc922d

File tree

9 files changed

+86
-39
lines changed

9 files changed

+86
-39
lines changed

‎app/ui/package.json

+1-1
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
{
22
"name": "app",
33
"private": true,
4-
"version": "1.8.1",
4+
"version": "1.8.2",
55
"type": "module",
66
"scripts": {
77
"dev": "vite",

‎app/ui/src/@types/bot.ts

+1
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ export type BotSettings = {
66
public_id: string;
77
temperature: number;
88
embedding: string;
9+
noOfDocumentsToRetrieve: number;
910
qaPrompt: string;
1011
questionGeneratorPrompt: string;
1112
streaming: boolean;

‎app/ui/src/components/Bot/Settings/SettingsCard.tsx

+28-1
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,12 @@
1-
import { Form, Input, notification, Select, Slider, Switch } from "antd";
1+
import {
2+
Form,
3+
Input,
4+
InputNumber,
5+
notification,
6+
Select,
7+
Slider,
8+
Switch,
9+
} from "antd";
210
import { useNavigate, useParams } from "react-router-dom";
311
import api from "../../../services/api";
412
import { useMutation, useQueryClient } from "@tanstack/react-query";
@@ -155,6 +163,7 @@ export const SettingsCard: React.FC<BotSettings> = ({
155163
bot_protect: data.bot_protect,
156164
use_rag: data.use_rag,
157165
bot_model_api_key: data.bot_model_api_key,
166+
noOfDocumentsToRetrieve: data.noOfDocumentsToRetrieve,
158167
}}
159168
form={form}
160169
requiredMark={false}
@@ -260,6 +269,24 @@ export const SettingsCard: React.FC<BotSettings> = ({
260269
/>
261270
</Form.Item>
262271

272+
<Form.Item
273+
name="noOfDocumentsToRetrieve"
274+
label="Number of documents to retrieve"
275+
rules={[
276+
{
277+
required: true,
278+
message:
279+
"Please input a number of documents to retrieve!",
280+
},
281+
]}
282+
>
283+
<InputNumber
284+
min={0}
285+
style={{ width: "100%" }}
286+
placeholder="Enter number of documents to retrieve"
287+
/>
288+
</Form.Item>
289+
263290
<Form.Item
264291
label={"Question Answering Prompt (System Prompt)"}
265292
name="qaPrompt"

‎package.json

+1-1
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
{
22
"name": "dialoqbase",
3-
"version": "1.8.1",
3+
"version": "1.8.2",
44
"description": "Create chatbots with ease",
55
"scripts": {
66
"ui:dev": "pnpm run --filter ui dev",

‎server/prisma/migrations/q_26/migration.sql

+2
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
-- AlterTable
2+
ALTER TABLE "Bot" ADD COLUMN "noOfDocumentsToRetrieve" INTEGER DEFAULT 4;

‎server/prisma/schema.prisma

+1
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ model Bot {
1818
description String?
1919
createdAt DateTime @default(now())
2020
temperature Float @default(0.7)
21+
noOfDocumentsToRetrieve Int? @default(4)
2122
model String @default("gpt-3.5-turbo")
2223
provider String @default("openai")
2324
embedding String @default("openai")

‎server/src/schema/api/v1/bot/bot/index.ts

+3
Original file line numberDiff line numberDiff line change
@@ -159,6 +159,9 @@ export const updateBotByIdSchema: FastifySchema = {
159159
bot_model_api_key: {
160160
type: "string",
161161
},
162+
noOfDocumentsToRetrieve: {
163+
type: "number",
164+
}
162165
},
163166
},
164167
};

‎server/src/utils/hybrid.ts

+35-29
Original file line numberDiff line numberDiff line change
@@ -40,47 +40,53 @@ export class DialoqbaseHybridRetrival extends BaseRetriever {
4040
protected async similaritySearch(
4141
query: string,
4242
k: number,
43-
_callbacks?: Callbacks,
43+
_callbacks?: Callbacks
4444
): Promise<SearchResult[]> {
4545
try {
46+
const embeddedQuery = await this.embeddings.embedQuery(query);
4647

47-
const embeddedQuery = await this.embeddings.embedQuery(query);
48+
const vector = `[${embeddedQuery.join(",")}]`;
49+
const bot_id = this.botId;
4850

49-
const vector = `[${embeddedQuery.join(",")}]`;
50-
const bot_id = this.botId;
51-
52-
const data = await prisma.$queryRaw`
51+
const data = await prisma.$queryRaw`
5352
SELECT * FROM "similarity_search_v2"(query_embedding := ${vector}::vector, botId := ${bot_id}::text,match_count := ${k}::int)
5453
`;
5554

56-
const result: [Document, number, number][] = (
57-
data as SearchEmbeddingsResponse[]
58-
).map((resp) => [
59-
new Document({
60-
metadata: resp.metadata,
61-
pageContent: resp.content,
62-
}),
63-
resp.similarity * 10,
64-
resp.id,
65-
]);
66-
67-
68-
return result;
69-
} catch (e) {
70-
console.log(e)
71-
return []
72-
}
55+
const result: [Document, number, number][] = (
56+
data as SearchEmbeddingsResponse[]
57+
).map((resp) => [
58+
new Document({
59+
metadata: resp.metadata,
60+
pageContent: resp.content,
61+
}),
62+
resp.similarity * 10,
63+
resp.id,
64+
]);
65+
66+
return result;
67+
} catch (e) {
68+
console.log(e);
69+
return [];
70+
}
7371
}
7472

7573
protected async keywordSearch(
7674
query: string,
77-
k: number,
75+
k: number
7876
): Promise<SearchResult[]> {
7977
const query_text = query;
8078
const bot_id = this.botId;
8179

80+
const botInfo = await prisma.bot.findFirst({
81+
where: {
82+
id: bot_id,
83+
},
84+
});
85+
86+
const match_count = botInfo?.noOfDocumentsToRetrieve || k;
87+
8288
const data = await prisma.$queryRaw`
83-
SELECT * FROM "kw_match_documents"(query_text := ${query_text}::text, bot_id := ${bot_id}::text,match_count := ${k}::int)
89+
SELECT * FROM "kw_match_documents"(query_text := ${query_text}::text, bot_id := ${bot_id}::text,match_count := ${match_count}::int)
8490
`;
8591

8692
const result: [Document, number, number][] = (
@@ -104,12 +110,12 @@ export class DialoqbaseHybridRetrival extends BaseRetriever {
104110
query: string,
105111
similarityK: number,
106112
keywordK: number,
107-
callbacks?: Callbacks,
113+
callbacks?: Callbacks
108114
): Promise<SearchResult[]> {
109115
const similarity_search = this.similaritySearch(
110116
query,
111117
similarityK,
112-
callbacks,
118+
callbacks
113119
);
114120

115121
const keyword_search = this.keywordSearch(query, keywordK);
@@ -136,13 +142,13 @@ export class DialoqbaseHybridRetrival extends BaseRetriever {
136142

137143
async _getRelevantDocuments(
138144
query: string,
139-
runManager?: CallbackManagerForRetrieverRun,
145+
runManager?: CallbackManagerForRetrieverRun
140146
): Promise<Document[]> {
141147
const searchResults = await this.hybridSearch(
142148
query,
143149
this.similarityK,
144150
this.keywordK,
145-
runManager?.getChild("hybrid_search"),
151+
runManager?.getChild("hybrid_search")
146152
);
147153

148154
return searchResults.map(([doc]) => doc);

‎server/src/utils/store.ts

+14-7
Original file line numberDiff line numberDiff line change
@@ -40,8 +40,7 @@ export class DialoqbaseVectorStore extends VectorStore {
4040
if (row?.embedding) {
4141
const vector = `[${row.embedding.join(",")}]`;
4242
const content = row?.content.replace(/\x00/g, "").trim();
43-
await prisma
44-
.$executeRaw`INSERT INTO "BotDocument" ("content", "embedding", "metadata", "botId", "sourceId") VALUES (${content}, ${vector}::vector, ${row.metadata}, ${row.botId}, ${row.sourceId})`;
43+
await prisma.$executeRaw`INSERT INTO "BotDocument" ("content", "embedding", "metadata", "botId", "sourceId") VALUES (${content}, ${vector}::vector, ${row.metadata}, ${row.botId}, ${row.sourceId})`;
4544
}
4645
});
4746
} catch (e) {
@@ -57,7 +56,7 @@ export class DialoqbaseVectorStore extends VectorStore {
5756
static async fromDocuments(
5857
docs: Document[],
5958
embeddings: Embeddings,
60-
dbConfig: DialoqbaseLibArgs,
59+
dbConfig: DialoqbaseLibArgs
6160
) {
6261
const instance = new this(embeddings, dbConfig);
6362
await instance.addDocuments(docs);
@@ -68,7 +67,7 @@ export class DialoqbaseVectorStore extends VectorStore {
6867
texts: string[],
6968
metadatas: object[] | object,
7069
embeddings: Embeddings,
71-
dbConfig: DialoqbaseLibArgs,
70+
dbConfig: DialoqbaseLibArgs
7271
) {
7372
const docs = [];
7473
for (let i = 0; i < texts.length; i += 1) {
@@ -84,7 +83,7 @@ export class DialoqbaseVectorStore extends VectorStore {
8483

8584
static async fromExistingIndex(
8685
embeddings: Embeddings,
87-
dbConfig: DialoqbaseLibArgs,
86+
dbConfig: DialoqbaseLibArgs
8887
) {
8988
const instance = new this(embeddings, dbConfig);
9089
return instance;
@@ -93,14 +92,22 @@ export class DialoqbaseVectorStore extends VectorStore {
9392
async similaritySearchVectorWithScore(
9493
query: number[],
9594
k: number,
96-
filter?: this["FilterType"] | undefined,
95+
filter?: this["FilterType"] | undefined
9796
): Promise<[Document<Record<string, any>>, number][]> {
9897
console.log(this.botId);
9998
const vector = `[${query.join(",")}]`;
10099
const bot_id = this.botId;
101100

101+
const botInfo = await prisma.bot.findFirst({
102+
where: {
103+
id: bot_id,
104+
},
105+
});
106+
107+
const match_count = botInfo?.noOfDocumentsToRetrieve || k;
108+
102109
const data = await prisma.$queryRaw`
103-
SELECT * FROM "similarity_search_v2"(query_embedding := ${vector}::vector, botId := ${bot_id}::text,match_count := ${k}::int)
110+
SELECT * FROM "similarity_search_v2"(query_embedding := ${vector}::vector, botId := ${bot_id}::text,match_count := ${match_count}::int)
104111
`;
105112

106113
const result: [Document, number][] = (

0 commit comments

Comments
 (0)