astro-relatinator 1.2.0 → 2.0.1

package/README.md

@@ -10,7 +10,7 @@ Assuming you're running Astro v3 or later, you can install this integration by r
 npx astro add astro-relatinator
 ```
 
-OR 
+OR
 
 ```
 npm install astro-relatinator
@@ -18,7 +18,12 @@ npm install astro-relatinator
 
 ## Configuration
 
-The integration takes a single input argument, `paths`, which is an array of paths to the `.md` or `.mdx` files you want to use for training the model. If you're using the [Content Collections API](https://docs.astro.build/en/guides/content-collections/), you can find these at `./src/collections/<collection-name>/`.
+The integration takes four input arguments, `paths`, `schema`, `similarityMethod`, and `debug`.
+
+- `paths` is an array of paths to the `.md` or `.mdx` files you want to use for training the model. If you're using the [Content Collections API](https://docs.astro.build/en/guides/content-collections/), you can find these at `./src/collections/<collection-name>/`.
+- `schema` is an array of frontmatter fields you want to use for training the model.
+- `similarityMethod` is the algorithm to use for training the model. Can be `"tfidf"` or `"bm25"`.
+- `debug` is a boolean flag to enable debug mode. It gets VERY verbose, so use it with caution.
 
 ### Example:
 
@@ -30,31 +35,33 @@ export default defineConfig({
     relatinatorIntegration({
       paths: ["./src/content/posts"], // path yo your content files
       schema: ["title", "descriptions", "tags"], // frontmatter fields to use for training
+      similarityMethod: "bm25", // use BM25 for training
+      debug: true, // enable debug mode
     }),
     // ...
   ],
 });
 ```
 
-__Note:__ the fields you specify in the `schema` array must be present in the frontmatter of the files you're training on. If they're not, the integration will throw an error or lead to unexpected results.
+**Note:** the fields you specify in the `schema` array must be present in the frontmatter of the files you're training on. If they're not, the integration will throw an error or lead to unexpected results.
 
 ## Usage
 
-The integration hooks into the Astro build (or dev) process, and will traing the tf-idf model based on the files you've provided. You can then use the `find` function from the package to find related content.
+The integration hooks into the Astro build (or dev) process, and will train the model based on the files you've provided. You can then use the `findRelated` function from the package to find related content.
 
 ### Example:
 
 ```astro
 ---
 import { getCollection } from "astro:content";
-import { find } from "astro-relatinator";
+import { findRelated } from "astro-relatinator";
 
 const blogEntries = await getCollection("posts"); // assuming your collection is called posts
 // let's say we want to compare against the first blog entry
 const first = blogEntries[0];
 
 // Concatenate the same fields we used for training and declare how many related posts we want to find
-const related = find(`${first.frontmatter.title} ${first.frontmatter.description} ${first.frontmatter.tags.join(' ')} ${first.body}`, 3)
+const related = findRelated(`${first.frontmatter.title} ${first.frontmatter.description} ${first.frontmatter.tags.join(' ')} ${first.body}`, first.id, 3)
 
 ---
 
@@ -71,4 +78,4 @@ const related = find(`${first.frontmatter.title} ${first.frontmatter.description
 }) }
 </div>
 
-```
+```

package/integration.ts

@@ -1,50 +1,55 @@
 import type { AstroIntegration } from "astro";
 import path from "node:path";
-import { train, tfIdf } from "relatinator";
+import { TfIdfUtils, BM25Utils } from "relatinator";
 import { glob } from "glob";
 import matter from "gray-matter";
 import fs from "node:fs/promises";
-import { TfIdfDocument } from "natural";
-import { Document } from "relatinator";
+import { RelatinatorDocument } from "relatinator";
 
-const stringifyData = (data: any): string => {
-  if (typeof data === 'object') {
-    if (Array.isArray(data)) {
-      return data.map(stringifyData).join(' ');
-    } else {
-      return Object.values(data).map(stringifyData).join(' ');
-    }
-  } else {
-    return String(data);
-  }
+// Helper to write debug info to a file
+async function writeDebug(message: string) {
+  console.log("[astro-relatinator]", message);
 }
 
-async function trainModel(paths: string[], schema: string[]) {
-  const documents: Document[] = [];
+async function trainModel(
+  similarityMethod: "tfidf" | "bm25",
+  paths: string[],
+  schema: string[],
+  debug = false
+) {
+  const documents: RelatinatorDocument[] = [];
 
-  if (tfIdf.documents) {
-    if ((tfIdf.documents as unknown as TfIdfDocument[]).length > 0) {
-      (tfIdf.documents as unknown as TfIdfDocument[]) = [];
-      (tfIdf.documents as unknown as TfIdfDocument[]).length = 0;
-    }
+  if (debug) {
+    await writeDebug("[trainModel] Training TF-IDF on existing data...");
   }
 
   for (const path of paths) {
     const markdownFiles = await glob(`${path}/**/*.{md,mdx}`);
 
+    if (debug) {
+      await writeDebug(
+        `[trainModel] Found ${markdownFiles.length} files in ${path}`
+      );
+    }
+
     for (const file of markdownFiles) {
       const content = await fs.readFile(file, "utf-8");
       const parsed = matter(content);
-      const idArr = file.split("/");
-      const id = idArr[idArr.length - 1];
 
-      let documentContent = '';
+      // Use the document's title as the ID
+      const id = parsed.data.title;
+
+      let documentContent = "";
       for (const key of schema) {
         if (parsed.data[key]) {
-          documentContent += ' ' + stringifyData(parsed.data[key]);
+          documentContent += " " + stringifyData(parsed.data[key]);
         }
       }
-      documentContent += ' ' + parsed.content;
+      documentContent += " " + parsed.content;
+
+      if (debug) {
+        await writeDebug(`[trainModel] Adding document with ID: ${id}`);
+      }
 
       documents.push({
         id,
@@ -53,15 +58,44 @@ async function trainModel(paths: string[], schema: string[]) {
     }
   }
 
-  train(documents);
+  if (debug) {
+    await writeDebug(
+      `[trainModel] Training with documents: ${documents
+        .map((d) => d.id)
+        .join(", ")}`
+    );
+  }
+
+  // Train the existing instance
+  if (similarityMethod === "tfidf") {
+    TfIdfUtils.train(documents, debug);
+  } else {
+    BM25Utils.train(documents, debug);
+  }
 }
 
+const stringifyData = (data: any): string => {
+  if (typeof data === "object") {
+    if (Array.isArray(data)) {
+      return data.map(stringifyData).join(" ");
+    } else {
+      return Object.values(data).map(stringifyData).join(" ");
+    }
+  } else {
+    return `${data}`;
+  }
+};
+
 const relatinatorIntegration = ({
   paths,
   schema,
+  similarityMethod = "tfidf",
+  debug = false,
 }: {
   paths: string[];
   schema: string[];
+  similarityMethod?: "tfidf" | "bm25";
+  debug?: boolean;
 }): AstroIntegration => {
   return {
     name: "astro-relatinator",
@@ -69,8 +103,15 @@ const relatinatorIntegration = ({
       "astro:server:setup": async ({ server, logger }) => {
         logger.info("relatinator server:setup");
 
-        logger.info("Training TF-IDF on existing data...");
-        await trainModel(paths, schema);
+        logger.info(
+          `Training Relatinator using ${similarityMethod} on existing data...`
+        );
+        if (similarityMethod === "tfidf") {
+          await TfIdfUtils.getInstance();
+        } else {
+          await BM25Utils.getInstance();
+        }
+        await trainModel(similarityMethod, paths, schema, debug);
         logger.info("Initial training done.");
 
         server.watcher.on("all", async (eventName, filePath) => {
@@ -80,9 +121,14 @@ const relatinatorIntegration = ({
 
           if (path.extname(filePath).includes("md")) {
             logger.info(
-              `Markdown file changed: ${filePath}; Re-training TF-IDF...`
+              `Markdown file changed: ${filePath}; Re-training ${similarityMethod}...`
             );
-            await trainModel(paths, schema);
+            if (similarityMethod === "tfidf") {
+              await TfIdfUtils.getInstance();
+            } else {
+              await BM25Utils.getInstance();
+            }
+            await trainModel(similarityMethod, paths, schema, debug);
           } else {
             return;
           }
@@ -91,13 +137,20 @@ const relatinatorIntegration = ({
       "astro:build:setup": async ({ logger }) => {
         logger.info("relatinator build:setup");
 
-        logger.info("Training TF-IDF on existing data...");
-        await trainModel(paths, schema);
-        logger.info("Training done.");
-      }
+        logger.info(
+          `Training Relatinator using ${similarityMethod} on existing data...`
+        );
+        if (similarityMethod === "tfidf") {
+          await TfIdfUtils.getInstance();
+        } else {
+          await BM25Utils.getInstance();
+        }
+        await trainModel(similarityMethod, paths, schema, debug);
+        logger.info(`Training done.`);
+      },
     },
   };
 };
 
+export { TfIdfUtils, BM25Utils };
 export default relatinatorIntegration;
-export { findRelated, tfIdf } from "relatinator";

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "astro-relatinator",
-  "version": "1.2.0",
+  "version": "2.0.1",
   "description": "Find related content in your Astro website or app",
   "author": "Darko Bozhinovski <dbozhinovski@gmail.com>",
   "license": "MIT",
@@ -28,16 +28,17 @@
   },
   "type": "module",
   "devDependencies": {
-    "@types/node": "^20.10.4",
-    "astro": ">=3",
-    "husky": "^8.0.3",
-    "natural": "^6.10.3"
+    "@types/node": "^22.13.4",
+    "astro": ">=5",
+    "husky": "^9.1.7",
+    "vite": "^6.1.1"
   },
   "dependencies": {
-    "glob": "^10.3.10",
-    "relatinator": "^1.1.0"
+    "glob": "^11.0.1",
+    "gray-matter": "^4.0.3",
+    "relatinator": "2.0.1"
   },
   "peerDependencies": {
-    "astro": ">=3"
+    "astro": ">=5"
   }
-}
+}