-
Notifications
You must be signed in to change notification settings - Fork 0
/
core.js
93 lines (76 loc) · 2.19 KB
/
core.js
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
const { VectorAlgorithms, SchemaFieldTypes } = require("redis")
const { OpenAI } = require("openai")
const {
OPENAI_KEY,
INDEX_NAME = "idx:doc"
} = process.env
const openai = new OpenAI({ apiKey: OPENAI_KEY })
const generateEmbeddingBuffer = async (text) => {
const { data } = await openai.embeddings.create({
model: "text-embedding-ada-002",
input: text
})
const vector = data[0].embedding
const buffer = Buffer.alloc(vector.length * 4);
for (let i = 0; i < vector.length; i++) {
buffer.writeFloatLE(vector[i], i * 4);
}
return buffer
}
const addEmbedding = async (redisClient, id, text) => {
const buffer = await generateEmbeddingBuffer(text)
await redisClient.hSet(`doc:${id}`, { text, embedding: buffer })
}
const search = async(redisClient, text, knn) => {
const buffer = await generateEmbeddingBuffer(text)
const data = await redisClient.ft.search(INDEX_NAME, `*=>[KNN ${knn} @embedding $vector AS score]`, {
PARAMS: {
vector: buffer
},
SORTBY: "score",
DIALECT: 2,
RETURN: ["score", "text"]
})
data.documents.forEach(doc => {
doc.value.score = Math.floor((1 - doc.value.score) * 100) / 100
})
return data
}
const chat = async (prompt, redisClient) => {
let data = await search(redisClient, prompt, 5)
let chatCompletion = await openai.chat.completions.create({
model: "gpt-4",
messages: [
{ role: "system", content: "You're a generic assistant of an event called socrates-IT" },
{ role: "system", content: `Your dataset: ${JSON.stringify(data)}` },
{ role: "user", content: prompt },
],
})
if (!chatCompletion?.choices?.length) {
throw new Error("No chat completion response")
}
return chatCompletion.choices[0].message.content
}
const createRedisStackIndex = async(redisClient) => await redisClient.ft.create(INDEX_NAME, {
text: {
type: SchemaFieldTypes.TEXT,
SORTABLE: false
},
embedding: {
type: SchemaFieldTypes.VECTOR,
ALGORITHM: VectorAlgorithms.HNSW,
TYPE: "FLOAT32",
DISTANCE_METRIC: "COSINE",
DIM: 1536
}
}, {
ON: "HASH",
PREFIX: ["doc:"]
})
module.exports = {
createRedisStackIndex,
generateEmbeddingBuffer,
addEmbedding,
search,
chat
}