feat: generate one datasource

yxiaolu · Oct 24, 2023 · 831a544 · 831a544
1 parent 85710f9
commit 831a544
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 6 deletions.
diff --git a/scripts/constants.mjs b/scripts/constants.mjs
@@ -1,4 +1,3 @@
-export const DATASOURCES = ["redhat", "basic_law_germany", "watchos"];
 export const DATASOURCES_DIR = "./datasources";
 export const DATASOURCES_CACHE_DIR = "./cache";
 export const DATASOURCES_CHUNK_SIZE = 512;

diff --git a/scripts/generate.mjs b/scripts/generate.mjs
@@ -6,12 +6,12 @@ import {
 } from "llamaindex";
 
 import {
- DATASOURCES,
  DATASOURCES_CACHE_DIR,
  DATASOURCES_DIR,
  DATASOURCES_CHUNK_SIZE,
  DATASOURCES_CHUNK_OVERLAP,
 } from "./constants.mjs";
+import { exit } from "process";
 
 async function getRuntime(func) {
  const start = Date.now();
@@ -42,14 +42,20 @@ async function generateDatasource(serviceContext, datasource) {
  );
 }
 
+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 () => {
  const serviceContext = serviceContextFromDefaults({
  chunkSize: DATASOURCES_CHUNK_SIZE,
  chunkOverlap: DATASOURCES_CHUNK_OVERLAP,
  });
 
- for (const datasource of DATASOURCES) {
- await generateDatasource(serviceContext, datasource);
- }
- console.log("Finished generating datasources");
+ await generateDatasource(serviceContext, datasource);
+ console.log("Finished generating datasource.");
 })();