forked from run-llama/chat-llamaindex
-
Notifications
You must be signed in to change notification settings - Fork 0
/
generate.mjs
86 lines (77 loc) · 2.45 KB
/
generate.mjs
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
import {
serviceContextFromDefaults,
storageContextFromDefaults,
SimpleDirectoryReader,
VectorStoreIndex,
} from "llamaindex";
import {
DATASOURCES_CACHE_DIR,
DATASOURCES_DIR,
DATASOURCES_CHUNK_SIZE,
DATASOURCES_CHUNK_OVERLAP,
} from "./constants.mjs";
import { exit } from "process";
import dotenv from "dotenv";
import path from "path";
import fs from "fs";
async function getRuntime(func) {
const start = Date.now();
await func();
const end = Date.now();
return end - start;
}
async function generateDatasource(serviceContext, datasource) {
console.log(`Generating storage context for datasource '${datasource}'...`);
// Split documents, create embeddings and store them in the storage context
const ms = await getRuntime(async () => {
const storageContext = await storageContextFromDefaults({
persistDir: `${DATASOURCES_CACHE_DIR}/${datasource}`,
});
const documents = await new SimpleDirectoryReader().loadData({
directoryPath: `${DATASOURCES_DIR}/${datasource}`,
});
await VectorStoreIndex.fromDocuments(documents, {
storageContext,
serviceContext,
});
});
console.log(
`Storage context for datasource '${datasource}' successfully generated in ${
ms / 1000
}s.`,
);
}
async function ensureEnv(fileName) {
try {
const __dirname = path.dirname(new URL(import.meta.url).pathname);
const envFileContent = await fs.promises.readFile(
path.join(__dirname, "..", fileName),
);
const envConfig = dotenv.parse(envFileContent);
if (envConfig && envConfig.OPENAI_API_KEY) {
process.env.OPENAI_API_KEY = envConfig.OPENAI_API_KEY;
} else {
throw new Error(`OPENAI_API_KEY not found in '${fileName}'`);
}
} catch (e) {
console.log(`Error getting OPENAI_API_KEY from ${fileName}: ${e.message}`);
exit(1);
}
console.log(`Using OPENAI_API_KEY=${process.env.OPENAI_API_KEY}`);
}
const datasource = process.argv[2];
if (!datasource) {
console.log("Error: You must provide a datasource as the parameter.");
console.log("Usage: pnpm run generate <datasource>");
exit(1);
}
(async () => {
// get OPENAI_API_KEY from Next.JS's .env.development.local
await ensureEnv(".env.development.local");
const serviceContext = serviceContextFromDefaults({
chunkSize: DATASOURCES_CHUNK_SIZE,
chunkOverlap: DATASOURCES_CHUNK_OVERLAP,
});
await generateDatasource(serviceContext, datasource);
console.log("Finished generating datasource.");
})();