@deepsweet/mdn 0.2.0 → 0.3.0

package/README.md

@@ -55,7 +55,6 @@ The `stdio` server will spawn [llama.cpp](https://github.com/ggml-org/llama.cpp)
 | Env variable               | Default value                                                   | Description                                                                                                       |
 |----------------------------|-----------------------------------------------------------------|-------------------------------------------------------------------------------------------------------------------|
 | `MDN_DATASET_PATH`         | HuggingFace cache                                               | Custom dataset directory path                                                                                     |
-| `MDN_DATASET_LOCALE`       | `en-us`                                                         | Dataset language, currently `en-us` only                                                                          |
 | `MDN_MODEL_PATH`           | HuggingFace cache                                               | Custom model file path                                                                                            |
 | `MDN_MODEL_TTL`            | `1800`                                                          | For how long llama.cpp with embedding model should be kept loaded in memory, in seconds; `0` to prevent unloading |
 | `MDN_QUERY_DESCRIPTION`    | `Natural language query for hybrid vector and full-text search` | Custom search query description in case your LLM does a poor job asking the MCP tool                              |
@@ -64,7 +63,6 @@ The `stdio` server will spawn [llama.cpp](https://github.com/ggml-org/llama.cpp)
 ## To do
 
 - [ ] figure out a better query description so that LLM doesn't over-generate keywords
-- [ ] add more dataset [translations](https://github.com/mdn/translated-content/tree/main/files/)
 - [ ] automatically update and upload the dataset artifacts monthly with GitHub Actions
 
 ## License

package/dist/index.js

@@ -1,10 +1,28 @@
 #!/usr/bin/env node
 
+// src/huggingface.ts
+import fs from "node:fs/promises";
+import path from "path";
+import {
+  downloadFileToCacheDir,
+  getHFHubCachePath,
+  getRepoFolderName,
+  scanCachedRepo,
+  snapshotDownload
+} from "@huggingface/hub";
+
+// src/const.ts
+var DATASET_REPO = "deepsweet/mdn";
+var MODEL_REPO = "deepsweet/bge-m3-GGUF-Q4_K_M";
+var MODEL_FILE = "bge-m3-GGUF-Q4_K_M.gguf";
+var MODEL_MAX_TOKENS = 8192;
+var TABLE_NAME = "mdn";
+var TABLE_FILENAME = `${TABLE_NAME}.lance`;
+
 // src/env.ts
 import { z } from "zod";
 var env = z.object({
   MDN_DATASET_PATH: z.string().optional(),
-  MDN_DATASET_LOCALE: z.enum(["en-us"]).default("en-us"),
   MDN_MODEL_PATH: z.string().optional(),
   MDN_MODEL_TTL: z.number().default(1800),
   MDN_QUERY_DESCRIPTION: z.string().default("Natural language query for hybrid vector and full-text search"),
@@ -12,19 +30,6 @@ var env = z.object({
 }).parse(process.env);
 
 // src/huggingface.ts
-import fs from "node:fs/promises";
-import path from "path";
-import { getHFHubCachePath, getRepoFolderName, scanCachedRepo, snapshotDownload } from "@huggingface/hub";
-
-// src/utils.ts
-var getTableName = (locale) => {
-  return `mdn-${locale}`;
-};
-
-// src/huggingface.ts
-var DATASET_REPO = "deepsweet/mdn";
-var MODEL_REPO = "deepsweet/bge-m3-GGUF-Q4_K_M";
-var MODEL_FILE = "bge-m3-GGUF-Q4_K_M.gguf";
 var replaceSymlinksWithHardlinks = async (dir) => {
   const entries = await fs.readdir(dir, { withFileTypes: true });
   for (const entry of entries) {
@@ -57,18 +62,17 @@ var getLatestCachedRepoRevision = async (name, type) => {
   });
   return latestRevision.path;
 };
-var downloadDataset = async (locale) => {
-  const tableName = getTableName(locale);
+var downloadDataset = async () => {
   const dirPath = await snapshotDownload({
     repo: `datasets/${DATASET_REPO}`,
-    path: `data/${tableName}.lance`
+    path: `data/${TABLE_FILENAME}`
   });
   const dataPath = path.join(dirPath, "data");
   await replaceSymlinksWithHardlinks(dataPath);
 };
 var getDatasetPath = async () => {
-  if (process.env.MDN_DATASET_PATH != null) {
-    return process.env.MDN_DATASET_PATH;
+  if (env.MDN_DATASET_PATH != null) {
+    return env.MDN_DATASET_PATH;
   }
   const latestRevisionPath = await getLatestCachedRepoRevision(DATASET_REPO, "dataset");
   const datasetPath = path.join(latestRevisionPath, "data");
@@ -80,8 +84,8 @@ var downloadModel = async () => {
   });
 };
 var getModelPath = async () => {
-  if (process.env.MDN_MODEL_PATH != null) {
-    return process.env.MDN_MODEL_PATH;
+  if (env.MDN_MODEL_PATH != null) {
+    return env.MDN_MODEL_PATH;
   }
   const latestRevisionPath = await getLatestCachedRepoRevision(MODEL_REPO, "model");
   const modelPath = path.join(latestRevisionPath, MODEL_FILE);
@@ -95,15 +99,16 @@ import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js"
 import { z as z2 } from "zod";
 
 // src/llama.ts
+import os from "node:os";
 import { getLlama, LlamaLogLevel } from "node-llama-cpp";
-var MAX_TOKENS = 8192;
 var getLlamaContext = async (modelPath) => {
+  const threads = Math.floor(os.availableParallelism() / 2);
   const llama = await getLlama({ logLevel: LlamaLogLevel.error });
   const model = await llama.loadModel({ modelPath });
   const context = await model.createEmbeddingContext({
-    contextSize: MAX_TOKENS,
-    batchSize: MAX_TOKENS,
-    threads: 0
+    contextSize: MODEL_MAX_TOKENS,
+    batchSize: MODEL_MAX_TOKENS,
+    threads
   });
   context.onDispose.createOnceListener(() => {
     model.dispose().then(() => llama.dispose()).catch(console.error);
@@ -124,7 +129,7 @@ var vectorize = async (context, text) => {
 // src/query.ts
 var queryHybrid = async (llamaContext, table, reranker, text) => {
   const vector = await vectorize(llamaContext, text);
-  const results = await table.query().nearestTo(vector).fullTextSearch(text).rerank(reranker).limit(env.MDN_SEARCH_RESULTS_LIMIT).toArray();
+  const results = await table.query().nearestTo(vector).column("vector").fullTextSearch(text, { columns: "text" }).rerank(reranker).limit(env.MDN_SEARCH_RESULTS_LIMIT).toArray();
   return results;
 };
 var createReranker = async () => {
@@ -134,7 +139,7 @@ var createReranker = async () => {
 
 // package.json
 var name = "@deepsweet/mdn";
-var version = "0.2.0";
+var version = "0.3.0";
 
 // src/server.ts
 var startMcpServer = async () => {
@@ -142,8 +147,7 @@ var startMcpServer = async () => {
   const db = await lancedb2.connect(datasetPath);
   const reranker = await createReranker();
   const server = new McpServer({ name, version });
-  const tableName = getTableName(env.MDN_DATASET_LOCALE);
-  const table = await db.openTable(tableName);
+  const table = await db.openTable(TABLE_NAME);
   const modelPath = await getModelPath();
   const llamaTtl = env.MDN_MODEL_TTL * 1000;
   let llamaContext = null;
@@ -202,8 +206,7 @@ var startMcpServer = async () => {
 // src/index.ts
 switch (process.argv[2]) {
   case "download": {
-    const locale = process.argv[3] ?? env.MDN_DATASET_LOCALE;
-    await downloadDataset(locale);
+    await downloadDataset();
     await downloadModel();
     break;
   }
@@ -212,7 +215,7 @@ switch (process.argv[2]) {
     break;
   }
   default: {
-    console.error('Unknown or missing command, use "download" or "server"');
+    console.error('Unknown command, use "download" or "server"');
     process.exit(1);
   }
 }

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@deepsweet/mdn",
-  "version": "0.2.0",
+  "version": "0.3.0",
   "publishConfig": {
     "access": "public"
   },
@@ -27,13 +27,18 @@
     "eslint": "^10.1.0",
     "gray-matter": "^4.0.3",
     "marked": "^17.0.5",
+    "p-all": "^5.0.1",
     "rimraf": "^6.1.3",
     "typescript": "^6.0.2"
   },
   "scripts": {
+    "download": "bun scripts/download.ts",
     "chunk": "bun scripts/chunk.ts",
     "ingest": "bun scripts/ingest.ts",
+    "update": "bun scripts/update.ts",
     "query": "bun scripts/query.ts",
+    "upload": "bun scripts/upload.ts",
+    "test": "bun scripts/test.ts",
     "check": "tsc --noEmit && eslint --cache scripts/ src/",
     "dist": "bun build --format esm --target node --packages external --outdir dist/ src/index.ts"
   },