feat: migration to local NLP pipeline

.gitignore

@@ -57,4 +57,7 @@ build/
 .eslintcache
 
 # VSCode
-.vscode/
+.vscode/
+
+# MacOS
+.DS_Store

apis/websocket/.env.example

@@ -1,2 +1,2 @@
-# Your Wit.ai token
-WIT_AI_TOKEN="YOUR_TOKEN_HERE"
+# LLM API configuration (optional)
+LLM_API_KEY="YOUR_LLM_API_KEY_HERE"

apis/websocket/.gitignore

@@ -0,0 +1,3 @@
+# temporary measures
+models/*/*
+indexes/*/*

apis/websocket/config.json

@@ -4,5 +4,60 @@
   "address": "127.0.0.1",
   "port": 3000,
   "ocrConcurrentQueues": 1,
-  "logLevel": "debug"
+  "logLevel": "debug",
+  "ai": {
+    "indexes": {
+      "documentation": "./indexes/documentation",
+      "qa": "./indexes/qa",
+      "productRelevance": "./indexes/product-relevance"
+    },
+    "embeddings": {
+      "modelPath": "Xenova/all-MiniLM-L6-v2",
+      "quant": "q8"
+    },
+    "intentClassification": {
+      "modelPath": "./models/intent-classification",
+      "quant": "q4",
+      "thresholds": {
+        "question": 0.7,
+        "problem": 0.7,
+        "complete": 0.7
+      }
+    },
+    "productRelevance": {
+      "threshold": 0.85
+    },
+    "answerValidation": {
+      "modelPath": "./models/answer-validation",
+      "quant": "q4",
+      "thresholds": {
+        "answer": 0.7,
+        "partial": 0.7,
+        "counter": 0.7,
+        "backchannel": 0.7
+      }
+    },
+    "llm": {
+      "enabled": false,
+      "baseURL": "",
+      "model": "",
+      "temperature": 0.7,
+      "systemPrompt": "You are a helpful assistant for ReVanced, an open-source project that provides patches for Android apps."
+    },
+    "rag": {
+      "confidence": {
+        "high": 0.9,
+        "medium": 0.3
+      },
+      "thresholds": {
+        "qa": 0.9,
+        "docs": 0.8
+      }
+    },
+    "documentParser": {
+      "maxChunkSize": 1000,
+      "chunkOverlap": 100,
+      "timeout": 30000
+    }
+  }
 }

apis/websocket/config.schema.json

@@ -22,6 +22,110 @@
       "type": "string",
       "enum": ["error", "warn", "info", "log", "debug", "trace", "none"],
       "default": "info"
+    },
+    "ai": {
+      "description": "AI configuration",
+      "type": "object",
+      "properties": {
+        "indexes": {
+          "description": "Vector store index paths",
+          "type": "object",
+          "properties": {
+            "documentation": { "type": "string" },
+            "qa": { "type": "string" },
+            "productRelevance": { "type": "string" }
+          }
+        },
+        "embeddings": {
+          "description": "Embeddings model configuration",
+          "type": "object",
+          "properties": {
+            "modelPath": { "type": "string" },
+            "quant": { "type": "string" }
+          }
+        },
+        "intentClassification": {
+          "description": "Intent classification model configuration",
+          "type": "object",
+          "properties": {
+            "modelPath": { "type": "string" },
+            "quant": { "type": "string" },
+            "thresholds": {
+              "type": "object",
+              "properties": {
+                "question": { "type": "number" },
+                "problem": { "type": "number" },
+                "complete": { "type": "number" }
+              }
+            }
+          }
+        },
+        "productRelevance": {
+          "description": "Product relevance configuration",
+          "type": "object",
+          "properties": {
+            "threshold": { "type": "number" }
+          }
+        },
+        "answerValidation": {
+          "description": "Answer validation model configuration",
+          "type": "object",
+          "properties": {
+            "modelPath": { "type": "string" },
+            "quant": { "type": "string" },
+            "thresholds": {
+              "type": "object",
+              "properties": {
+                "answer": { "type": "number" },
+                "partial": { "type": "number" },
+                "counter": { "type": "number" },
+                "backchannel": { "type": "number" }
+              }
+            }
+          }
+        },
+        "llm": {
+          "description": "LLM provider configuration",
+          "type": "object",
+          "properties": {
+            "enabled": { "type": "boolean" },
+            "baseURL": { "type": "string" },
+            "model": { "type": "string" },
+            "temperature": { "type": "number" },
+            "systemPrompt": { "type": "string" }
+          }
+        },
+        "rag": {
+          "description": "RAG configuration",
+          "type": "object",
+          "properties": {
+            "confidence": {
+              "type": "object",
+              "properties": {
+                "high": { "type": "number" },
+                "medium": { "type": "number" }
+              }
+            },
+            "thresholds": {
+              "type": "object",
+              "properties": {
+                "qa": { "type": "number" },
+                "docs": { "type": "number" },
+                "returnEmpty": { "type": "boolean" }
+              }
+            }
+          }
+        },
+        "documentParser": {
+          "description": "Document parsing and chunking configuration",
+          "type": "object",
+          "properties": {
+            "maxChunkSize": { "type": "integer" },
+            "chunkOverlap": { "type": "integer" },
+            "timeout": { "type": "integer" }
+          }
+        }
+      }
     }
   }
 }

apis/websocket/package.json

@@ -27,6 +27,7 @@
   },
   "homepage": "https://github.com/revanced/revanced-bots#readme",
   "dependencies": {
+    "@revanced/bot-ai": "workspace:*",
     "@revanced/bot-shared": "workspace:*",
     "@sapphire/async-queue": "^1.5.5",
     "chalk": "^5.6.2",

apis/websocket/src/context.ts

@@ -1,69 +1,23 @@
-import { createLogger } from '@revanced/bot-shared'
-import { exists as pathExists } from 'fs/promises'
-import { join as joinPath } from 'path'
-import { createWorker as createTesseractWorker, OEM } from 'tesseract.js'
-import { getConfig } from './utils/config'
-
-export const config = getConfig()
-
-export const logger = createLogger({
-    level: config.logLevel === 'none' ? Number.MAX_SAFE_INTEGER : config.logLevel,
-})
-
-export const wit = {
-    token: process.env['WIT_AI_TOKEN']!,
-    async fetch(route: string, options?: RequestInit) {
-        const res = await fetch(`https://api.wit.ai${route}`, {
-            headers: {
-                Authorization: `Bearer ${this.token}`,
-                'Content-Type': 'application/json',
-            },
-            ...options,
-        })
-
-        if (!res.ok) throw new Error(`Failed to fetch from Wit.ai: ${res.statusText} (${res.status})`)
-
-        return (await res.json()) as WitMessageResponse
-    },
-    message(text: string) {
-        return this.fetch(`/message?q=${encodeURIComponent(text)}&n=8`) as Promise<WitMessageResponse>
-    },
-    async train(text: string, label?: string) {
-        await this.fetch('/utterances', {
-            body: JSON.stringify([
-                {
-                    text,
-                    intent: label,
-                    entities: [],
-                    traits: [],
-                },
-            ]),
-            method: 'POST',
-        })
-    },
-} satisfies Wit
-
-export interface Wit {
-    token: string
-    fetch(route: string, options?: RequestInit): Promise<WitMessageResponse>
-    message(text: string): Promise<WitMessageResponse>
-    train(text: string, label?: string): Promise<void>
-}
-
-export interface WitMessageResponse {
-    text: string
-    intents: Array<{
-        id: string
-        name: string
-        confidence: number
-    }>
-}
-
-const TesseractWorkerDirPath = joinPath(import.meta.dir, 'worker')
-const TesseractWorkerPath = joinPath(TesseractWorkerDirPath, 'index.js')
-
-export const tesseract = await createTesseractWorker(
-    'eng',
-    OEM.DEFAULT,
-    (await pathExists(TesseractWorkerDirPath)) ? { workerPath: TesseractWorkerPath } : undefined,
-)
+import { AI } from '@revanced/bot-ai'
+import { createLogger } from '@revanced/bot-shared'
+import { exists as pathExists } from 'fs/promises'
+import { join as joinPath } from 'path'
+import { createWorker as createTesseractWorker, OEM } from 'tesseract.js'
+import { getConfig } from './utils/config'
+
+export const config = getConfig()
+
+export const logger = createLogger({
+    level: config.logLevel === 'none' ? Number.MAX_SAFE_INTEGER : config.logLevel,
+})
+
+export const ai = new AI(config.ai)
+
+const TesseractWorkerDirPath = joinPath(import.meta.dir, 'worker')
+const TesseractWorkerPath = joinPath(TesseractWorkerDirPath, 'index.js')
+
+export const tesseract = await createTesseractWorker(
+    'eng',
+    OEM.DEFAULT,
+    (await pathExists(TesseractWorkerDirPath)) ? { workerPath: TesseractWorkerPath } : undefined,
+)

apis/websocket/src/events/addDocumentation.ts

@@ -0,0 +1,38 @@
+import { type ClientOperation, ServerOperation } from '@revanced/bot-shared'
+import type { EventHandler } from '.'
+
+const addDocumentationEventHandler: EventHandler<ClientOperation.AddDocumentation> = async (packet, { ai, logger }) => {
+    const {
+        client,
+        d: { text, url, title },
+    } = packet
+
+    const nextSeq = client.currentSequence++
+
+    logger.debug(`Client ${client.id} requested to add documentation:`, text ? text.slice(0, 50) : url)
+
+    try {
+        await ai.addDocumentation(text, url, title ?? '')
+
+        client.send(
+            {
+                op: ServerOperation.AddedDocumentation,
+                d: true,
+            },
+            nextSeq,
+        )
+    } catch (e) {
+        if (!client.disconnected)
+            client.send(
+                {
+                    op: ServerOperation.AddDocumentationFailed,
+                    d: null,
+                },
+                nextSeq,
+            )
+        else logger.warn(`Client disconnected before the failed packet could be sent (${nextSeq})`)
+        logger.error(`Failed to add documentation (${nextSeq}):`, e)
+    }
+}
+
+export default addDocumentationEventHandler

apis/websocket/src/events/addDocumentationFromUrl.ts

@@ -0,0 +1,47 @@
+import { type ClientOperation, ServerOperation } from '@revanced/bot-shared'
+import type { EventHandler } from '.'
+
+const addDocumentationFromUrlEventHandler: EventHandler<ClientOperation.AddDocumentationFromUrl> = async (
+    packet,
+    { ai, logger },
+) => {
+    const {
+        client,
+        d: { url },
+    } = packet
+
+    const nextSeq = client.currentSequence++
+
+    logger.debug(`Client ${client.id} requested to add documentation from URL:`, url)
+
+    try {
+        const result = await ai.addDocumentationFromUrl(url)
+
+        client.send(
+            {
+                op: ServerOperation.AddedDocumentationFromUrl,
+                d: result,
+            },
+            nextSeq,
+        )
+
+        if (result.success) {
+            logger.info(`Successfully added ${result.chunksAdded} chunks from URL: ${url} (title: ${result.title})`)
+        } else {
+            logger.warn(`Failed to add documentation from URL: ${url} - ${result.error}`)
+        }
+    } catch (e) {
+        if (!client.disconnected)
+            client.send(
+                {
+                    op: ServerOperation.AddDocumentationFromUrlFailed,
+                    d: null,
+                },
+                nextSeq,
+            )
+        else logger.warn(`Client disconnected before the failed packet could be sent (${nextSeq})`)
+        logger.error(`Failed to add documentation from URL (${nextSeq}):`, e)
+    }
+}
+
+export default addDocumentationFromUrlEventHandler