lightrag 1.0.0

package/README.md

@@ -0,0 +1,79 @@
+# LightRAG
+
+Lightweight Retrieval-Augmented Generation library with Knowledge Graph.  
+Runs in Node.js (CommonJS). Browser usage requires a bundler.
+
+## Install
+
+```bash
+npm install lightrag
+```
+
+## Usage
+
+```js
+const { pipeline } = require('@huggingface/transformers');
+const { LightRAG, Embedder } = require('lightrag');
+
+(async () => {
+    const _pipe = await pipeline('feature-extraction', 'Xenova/multilingual-e5-small', { dtype: 'fp32' });
+    const _embedder = new Embedder(_pipe);
+    const _rag = new LightRAG({
+        embedder: _embedder,
+        tokenizer: _pipe.tokenizer,
+        llmFunc: async (prompt, opts = {}) => {
+            // call your LLM API here, return string
+        },
+    });
+
+    await _rag.insert('Your document text here...');
+
+    const _answer = await _rag.query('What is this about?', { mode: 'hybrid' });
+    console.log(_answer);
+})();
+```
+
+## API
+
+### `new LightRAG(opts)`
+
+| Option | Type | Default | Description |
+|--------|------|---------|-------------|
+| `opts.llmFunc` | `async (prompt, options) => string` | **required** | LLM call function |
+| `opts.embedder` | `Embedder` | **required** | Transformers.js embedder instance |
+| `opts.tokenizer` | tokenizer | — | Tokenizer for token-counting chunker (falls back to character-based) |
+| `opts.embeddingDim` | `number` | `384` | Embedding vector dimension |
+| `opts.chunkSize` | `number` | `480` | Max tokens per chunk (with tokenizer) or chars (without) |
+| `opts.chunkOverlap` | `number` | `50` | Overlap tokens (or chars) between chunks |
+| `opts.maxEntityTokens` | `number` | `96000` | Max characters sent to LLM for entity extraction per chunk |
+
+### `rag.insert(text)`
+
+Inserts document text: tokenizes → chunks → embeds each chunk into vector DB → extracts entities & relations into knowledge graph.
+
+### `rag.query(question, options)`
+
+Query with RAG.
+
+| Option | Type | Default | Description |
+|--------|------|---------|-------------|
+| `options.mode` | `string` | `'hybrid'` | `'naive'` \| `'local'` \| `'global'` \| `'hybrid'` \| `'mix'` |
+| `options.systemPrompt` | `string` | `''` | System prompt prepended to LLM context |
+| `options.history` | `array` | `[]` | Conversation history `[{role, content}]` |
+| `options.stream` | `boolean` | `false` | Whether the LLM should stream (passed to `llmFunc`) |
+
+### `rag.getGraphData()`
+
+Returns `{ nodes, edges }` — nodes with `id, entity_type, description, degree`; edges with `source, target, keywords, weight`.
+
+### `rag.getProgress()`
+
+Returns `{ total, ready, isInserting, isReady }` — insertion progress.
+
+### `rag.toJSON()` / `LightRAG.fromJSON(data, opts)`
+
+Serialize/deserialize the entire RAG state (vector DB + knowledge graph).
+
+## License
+
+MIT

package/package.json

@@ -0,0 +1,31 @@
+{
+  "name": "lightrag",
+  "version": "1.0.0",
+  "description": "Lightweight RAG library with knowledge graph — browser and Node.js, MIT licensed",
+  "main": "src/index.js",
+  "files": [
+    "src/"
+  ],
+  "scripts": {
+    "test": "echo \"Error: no test specified\" && exit 1"
+  },
+  "keywords": [
+    "rag",
+    "knowledge-graph",
+    "vector-db",
+    "llm",
+    "transformers.js",
+    "lightrag",
+    "entity-extraction"
+  ],
+  "author": "BlackboardLM",
+  "license": "MIT",
+  "repository": {
+    "type": "git",
+    "url": "https://github.com/winterist/BlackboardLM"
+  },
+  "type": "commonjs",
+  "dependencies": {
+    "@huggingface/transformers": "^4.2.0"
+  }
+}

package/src/chunker.js

@@ -0,0 +1,33 @@
+class TokenChunker {
+    constructor(tokenizer, tokenLimit = 480, overlapTokens = 50) {
+        this._tokenizer = tokenizer;
+        this._tokenLimit = tokenLimit;
+        this._overlapTokens = overlapTokens;
+    }
+    chunk(text) {
+        if (!this._tokenizer) {
+            const _chunks = [];
+            let _start = 0;
+            while (_start < text.length) {
+                const _end = Math.min(_start + 3000, text.length);
+                _chunks.push(text.slice(_start, _end));
+                if (_end >= text.length) break;
+                _start = _end - 150;
+            }
+            return _chunks;
+        }
+        const _tokens = this._tokenizer.encode(text);
+        const _chunks = [];
+        let _start = 0;
+        while (_start < _tokens.length) {
+            const _end = Math.min(_start + this._tokenLimit, _tokens.length);
+            const _chunkTokens = _tokens.slice(_start, _end);
+            const _chunkText = this._tokenizer.decode(_chunkTokens, { skip_special_tokens: true });
+            _chunks.push(_chunkText);
+            if (_end >= _tokens.length) break;
+            _start = _end - this._overlapTokens;
+        }
+        return _chunks;
+    }
+}
+module.exports = { TokenChunker };

package/src/embedder.js

@@ -0,0 +1,18 @@
+class Embedder {
+    constructor(pipeline) {
+        this._pipe = pipeline;
+    }
+    async embed(texts, context = 'document') {
+        const _prefix = context === 'query' ? 'query: ' : 'passage: ';
+        const _prefixed = texts.map(_t => _prefix + _t);
+        const _output = await this._pipe(_prefixed, { pooling: 'mean', normalize: true });
+        return _output.tolist();
+    }
+    async embedQuery(query) {
+        return (await this.embed([query], 'query'))[0];
+    }
+    async embedDocuments(docs) {
+        return await this.embed(docs, 'document');
+    }
+}
+module.exports = { Embedder };

package/src/graph.js

@@ -0,0 +1,67 @@
+class KnowledgeGraph {
+    constructor() {
+        this._nodes = new Map();
+        this._edges = [];
+        this._adj = new Map();
+    }
+    upsertNode(id, data = {}) {
+        if (!this._nodes.has(id)) {
+            this._nodes.set(id, { id, entity_type: data.entity_type || '', description: data.description || '' });
+            this._adj.set(id, new Map());
+        } else {
+            const _n = this._nodes.get(id);
+            if (data.description && data.description.length > _n.description.length) _n.description = data.description;
+            if (data.entity_type) _n.entity_type = data.entity_type;
+        }
+    }
+    upsertEdge(source, target, data = {}) {
+        if (!this._nodes.has(source)) this.upsertNode(source);
+        if (!this._nodes.has(target)) this.upsertNode(target);
+        const _e = { source, target, keywords: data.keywords || '', weight: data.weight || 1 };
+        this._edges.push(_e);
+        const _adjSource = this._adj.get(source);
+        if (_adjSource) _adjSource.set(target, _e);
+        const _adjTarget = this._adj.get(target);
+        if (_adjTarget) _adjTarget.set(source, _e);
+    }
+    getNeighbors(id, maxDegree = 10) {
+        const _adjNode = this._adj.get(id);
+        if (!_adjNode) return [];
+        return Array.from(_adjNode.entries())
+            .sort((_a, _b) => _b[1].weight - _a[1].weight)
+            .slice(0, maxDegree)
+            .map(([_nodeId, _edge]) => ({ nodeId: _nodeId, ..._edge }));
+    }
+    getAllNodes() {
+        return Array.from(this._nodes.values()).map(_n => ({
+            id: _n.id, entity_type: _n.entity_type, description: _n.description,
+            degree: (this._adj.get(_n.id)?.size || 0),
+        }));
+    }
+    getAllEdges() {
+        return this._edges
+            .filter(_e => this._nodes.has(_e.source) && this._nodes.has(_e.target))
+            .map(_e => ({ source: _e.source, target: _e.target, keywords: _e.keywords, weight: _e.weight }));
+    }
+    getNode(id) { return this._nodes.get(id) || null; }
+    get size() { return this._nodes.size; }
+    toJSON() {
+        return {
+            nodes: Array.from(this._nodes.entries()),
+            edges: this._edges,
+            adj: Array.from(this._adj.entries()).map(([k, v]) => [k, Array.from(v.entries())]),
+        };
+    }
+    static fromJSON(data) {
+        const _g = new KnowledgeGraph();
+        if (data.nodes) for (const [k, v] of data.nodes) _g._nodes.set(k, v);
+        if (data.edges) _g._edges = data.edges;
+        if (data.adj) for (const [k, entries] of data.adj) {
+            const _m = new Map();
+            for (const [k2, v] of entries) _m.set(k2, v);
+            _g._adj.set(k, _m);
+        }
+        return _g;
+    }
+}
+module.exports = { KnowledgeGraph };

package/src/index.js

@@ -0,0 +1,8 @@
+const { LightRAG } = require('./lightrag');
+const { Embedder } = require('./embedder');
+const { KnowledgeGraph } = require('./graph');
+const { VectorDB } = require('./vector_db');
+const { TokenChunker } = require('./chunker');
+const { ENTITY_EXTRACTION_PROMPT } = require('./prompts');
+
+module.exports = { LightRAG, Embedder, KnowledgeGraph, VectorDB, TokenChunker, ENTITY_EXTRACTION_PROMPT };

package/src/lightrag.js

@@ -0,0 +1,130 @@
+const { TokenChunker } = require('./chunker');
+const { VectorDB } = require('./vector_db');
+const { KnowledgeGraph } = require('./graph');
+const { ENTITY_EXTRACTION_PROMPT } = require('./prompts');
+
+class LightRAG {
+    constructor(opts = {}) {
+        this._llmFunc = opts.llmFunc;
+        this._embedder = opts.embedder;
+        this._embeddingDim = opts.embeddingDim || 384;
+        this._chunkSize = opts.chunkSize || 480;
+        this._chunkOverlap = opts.chunkOverlap || 50;
+        this._maxEntityTokens = opts.maxEntityTokens || 96000;
+        this._chunker = new TokenChunker(opts.tokenizer, this._chunkSize, this._chunkOverlap);
+        this._vdb = new VectorDB(this._embeddingDim);
+        this._graph = new KnowledgeGraph();
+        this._progress = { total: 0, ready: 0, isInserting: false, isReady: true };
+    }
+
+    async insert(text) {
+        const _chunks = this._chunker.chunk(text);
+        console.error(`[LightRAG] Chunked text (token_limit=${this._chunkSize}, overlap=${this._chunkOverlap}) into ${_chunks.length} chunks`);
+        this._progress = { total: _chunks.length, ready: 0, isInserting: true, isReady: false };
+        for (let _i = 0; _i < _chunks.length; _i++) {
+            const _chunk = _chunks[_i];
+            const _id = `chunk_${Date.now()}_${_i}`;
+            console.error(`[LightRAG] Embedding chunk ${_i + 1}/${_chunks.length} (${_chunk.length} chars)`);
+            const _vectors = await this._embedder.embedDocuments([_chunk]);
+            const _vec = new Float32Array(_vectors[0]);
+            await this._vdb.upsert([{ id: _id, vector: Array.from(_vec), text: _chunk }]);
+            console.error(`[LightRAG] Extracting entities from chunk ${_i + 1}/${_chunks.length}`);
+            await this._extractEntities(_chunk, _id);
+            this._progress.ready = _i + 1;
+        }
+        console.error(`[LightRAG] Insert complete: ${this._vdb.size} vectors, ${this._graph.size} entities`);
+        this._progress.isInserting = false;
+        this._progress.isReady = true;
+    }
+
+    async _extractEntities(text, sourceId) {
+        console.error(`[LightRAG] Calling LLM for entity extraction (${text.length} chars)`);
+        const _prompt = ENTITY_EXTRACTION_PROMPT.replace('{text}', text.slice(0, this._maxEntityTokens));
+        try {
+            const _raw = await this._llmFunc(_prompt);
+            const _json = this._parseJson(_raw);
+            if (_json && _json.entities) {
+                for (const _e of _json.entities) {
+                    const _name = (_e.name || '').trim().replace(/(?:^|\s+)([a-z])/g, (_m, _c) => _m[0] === ' ' ? ' ' + _c.toUpperCase() : _c.toUpperCase());
+                    if (!_name) continue;
+                    this._graph.upsertNode(_name, { entity_type: _e.type || '', description: _e.description || '', source_id: sourceId });
+                }
+            }
+            if (_json && _json.relations) {
+                for (const _r of _json.relations) {
+                    const _s = (_r.source || '').trim().replace(/(?:^|\s+)([a-z])/g, (_m, _c) => _m[0] === ' ' ? ' ' + _c.toUpperCase() : _c.toUpperCase());
+                    const _t = (_r.target || '').trim().replace(/(?:^|\s+)([a-z])/g, (_m, _c) => _m[0] === ' ' ? ' ' + _c.toUpperCase() : _c.toUpperCase());
+                    if (!_s || !_t) continue;
+                    this._graph.upsertEdge(_s, _t, { keywords: _r.keywords || '', weight: _r.weight || 1, source_id: sourceId });
+                }
+            }
+            const _ec = _json?.entities?.length || 0;
+            const _rc = _json?.relations?.length || 0;
+            console.error(`[LightRAG] Entity extraction done: ${_ec} entities, ${_rc} relations`);
+        } catch (_e) {
+            console.error(`[LightRAG] Entity extraction failed: ${_e.message}`);
+        }
+    }
+
+    _parseJson(raw) {
+        let _s = raw.trim();
+        if (_s.startsWith('```')) _s = _s.replace(/```\w*\n?/g, '').trim();
+        try { return JSON.parse(_s); } catch (_e) {}
+        const _m = _s.match(/\{[\s\S]*\}/);
+        if (_m) try { return JSON.parse(_m[0]); } catch (_e) {}
+        return null;
+    }
+
+    async query(question, options = {}) {
+        const _mode = options.mode || 'hybrid';
+        console.error(`[LightRAG] Query in "${_mode}" mode: "${question.slice(0, 80)}..."`);
+        if (_mode === 'naive' || this._vdb.size === 0) {
+            console.error(`[LightRAG] No retrieval needed (mode=${_mode}, vdb_size=${this._vdb.size})`);
+            return await this._llmFunc(question, { system_prompt: options.systemPrompt || '', history: options.history || [], stream: options.stream });
+        }
+        const _qVec = await this._embedder.embedQuery(question);
+        const _context = await this._buildContext(new Float32Array(_qVec), _mode);
+        console.error(`[LightRAG] Context built: ${_context.length} chars`);
+        const _sysPrompt = (options.systemPrompt || '') + (_context ? `\n\n---Context from your documents---\n${_context}` : '');
+        return await this._llmFunc(question, { system_prompt: _sysPrompt, history: options.history || [], stream: options.stream });
+    }
+
+    async _buildContext(_qVec, mode) {
+        let _parts = [];
+        if (mode === 'local' || mode === 'hybrid' || mode === 'mix') {
+            const _chunks = await this._vdb.query(Array.from(_qVec), 5);
+            if (_chunks.length) _parts.push('=== Relevant Document Chunks ===\n' + _chunks.map(_c => `[${_c.id}] ${_c.text}`).join('\n\n'));
+        }
+        if (mode === 'global' || mode === 'hybrid' || mode === 'mix') {
+            const _nodes = this._graph.getAllNodes();
+            const _sorted = _nodes.sort((_a, _b) => _b.degree - _a.degree).slice(0, 20);
+            if (_sorted.length) {
+                _parts.push('=== Knowledge Graph Entities ===\n' + _sorted.map(_n => `- ${_n.id} [${_n.entity_type}]: ${_n.description}`).join('\n'));
+            }
+        }
+        return _parts.join('\n\n');
+    }
+
+    getGraphData() {
+        return { nodes: this._graph.getAllNodes(), edges: this._graph.getAllEdges() };
+    }
+
+    getProgress() { return { ...this._progress }; }
+
+    toJSON() {
+        return {
+            embeddingDim: this._embeddingDim,
+            vdb: this._vdb.toJSON(),
+            graph: this._graph.toJSON(),
+        };
+    }
+
+    static fromJSON(data, opts = {}) {
+        const _rag = new LightRAG({ ...opts, embeddingDim: data.embeddingDim || 384 });
+        _rag._vdb = VectorDB.fromJSON(data.vdb || {});
+        _rag._graph = KnowledgeGraph.fromJSON(data.graph || {});
+        return _rag;
+    }
+}
+
+module.exports = { LightRAG };

package/src/prompts.js

@@ -0,0 +1,27 @@
+const ENTITY_EXTRACTION_PROMPT = `---Role---
+You are an expert knowledge graph builder. Extract entities and relationships from the text.
+
+---Goal---
+Given a text chunk, identify ALL named entities (people, places, organizations, concepts, events, dates, technologies) and the relationships between them.
+
+---Output Format---
+Return a JSON object with "entities" and "relations":
+{
+  "entities": [
+    {"name": "Entity", "type": "PERSON|ORGANIZATION|LOCATION|CONCEPT|EVENT|DATE|TECH", "description": "Brief description"}
+  ],
+  "relations": [
+    {"source": "Entity A", "target": "Entity B", "keywords": "relationship", "weight": 1.0}
+  ]
+}
+
+---Rules---
+- Entity name under 60 chars, description under 120 chars
+- Extract at least 3 entities per chunk
+- Weight between 0.5 (weak) and 1.5 (strong)
+- Return ONLY valid JSON, no other text.
+
+---Text---
+{text}`;
+
+module.exports = { ENTITY_EXTRACTION_PROMPT };

package/src/vector_db.js

@@ -0,0 +1,47 @@
+class VectorDB {
+    constructor(dim) {
+        this._dim = dim;
+        this._vectors = [];
+        this._meta = [];
+    }
+    async upsert(entries) {
+        for (const _e of entries) {
+            const _idx = this._meta.findIndex(_m => _m.id === _e.id);
+            if (_idx >= 0) {
+                this._vectors[_idx] = _e.vector;
+                this._meta[_idx] = { id: _e.id, text: _e.text || '' };
+            } else {
+                this._vectors.push(_e.vector);
+                this._meta.push({ id: _e.id, text: _e.text || '' });
+            }
+        }
+    }
+    async query(vector, topK = 10) {
+        if (this._vectors.length === 0) return [];
+        const _scores = this._vectors.map((_v, _i) => ({
+            index: _i,
+            score: this._cosine(vector, _v),
+            id: this._meta[_i].id,
+            text: this._meta[_i].text,
+        }));
+        _scores.sort((_a, _b) => _b.score - _a.score);
+        return _scores.slice(0, topK);
+    }
+    _cosine(a, b) {
+        let _dot = 0, _na = 0, _nb = 0;
+        for (let _i = 0; _i < a.length; _i++) { _dot += a[_i] * b[_i]; _na += a[_i] * a[_i]; _nb += b[_i] * b[_i]; }
+        _na = Math.sqrt(_na); _nb = Math.sqrt(_nb);
+        return (_na && _nb) ? _dot / (_na * _nb) : 0;
+    }
+    get size() { return this._vectors.length; }
+    toJSON() {
+        return { dim: this._dim, vectors: this._vectors, meta: this._meta };
+    }
+    static fromJSON(data) {
+        const _db = new VectorDB(data.dim);
+        _db._vectors = data.vectors || [];
+        _db._meta = data.meta || [];
+        return _db;
+    }
+}
+module.exports = { VectorDB };