File size: 2,086 Bytes
f8b5d42 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 |
const fs = require("fs");
const path = require("path");
const documentsPath =
process.env.NODE_ENV === "development"
? path.resolve(__dirname, `../../storage/documents`)
: path.resolve(process.env.STORAGE_DIR, `documents`);
class DocumentManager {
constructor({ workspace = null, maxTokens = null }) {
this.workspace = workspace;
this.maxTokens = maxTokens || Number.POSITIVE_INFINITY;
this.documentStoragePath = documentsPath;
}
log(text, ...args) {
console.log(`\x1b[36m[DocumentManager]\x1b[0m ${text}`, ...args);
}
async pinnedDocuments() {
if (!this.workspace) return [];
const { Document } = require("../../models/documents");
return await Document.where({
workspaceId: Number(this.workspace.id),
pinned: true,
});
}
async pinnedDocs() {
if (!this.workspace) return [];
const docPaths = (await this.pinnedDocuments()).map((doc) => doc.docpath);
if (docPaths.length === 0) return [];
let tokens = 0;
const pinnedDocs = [];
for await (const docPath of docPaths) {
try {
const filePath = path.resolve(this.documentStoragePath, docPath);
const data = JSON.parse(
fs.readFileSync(filePath, { encoding: "utf-8" })
);
if (
!data.hasOwnProperty("pageContent") ||
!data.hasOwnProperty("token_count_estimate")
) {
this.log(
`Skipping document - Could not find page content or token_count_estimate in pinned source.`
);
continue;
}
if (tokens >= this.maxTokens) {
this.log(
`Skipping document - Token limit of ${this.maxTokens} has already been exceeded by pinned documents.`
);
continue;
}
pinnedDocs.push(data);
tokens += data.token_count_estimate || 0;
} catch {}
}
this.log(
`Found ${pinnedDocs.length} pinned sources - prepending to content with ~${tokens} tokens of content.`
);
return pinnedDocs;
}
}
module.exports.DocumentManager = DocumentManager;
|