Compare commits

...

5 Commits

Author SHA1 Message Date
Alex Yang d038938272 chore: manual export modules 2024-01-24 15:18:07 -06:00
Alex Yang 76bb6c5f7f fix: set default fs 2024-01-24 14:53:09 -06:00
Alex Yang 38a5a3a04f fix: test 2024-01-24 14:43:02 -06:00
Alex Yang 86d780d580 fix: test 2024-01-24 14:42:48 -06:00
Alex Yang cb5141c410 feat: abstract file system api 2024-01-24 14:41:10 -06:00
21 changed files with 135 additions and 109 deletions
+11 -6
View File
@@ -136,6 +136,11 @@
"import": "./dist/Tool.mjs",
"require": "./dist/Tool.js"
},
"./readers": {
"types": "./dist/readers.d.mts",
"import": "./dist/readers.mjs",
"require": "./dist/readers.js"
},
"./readers/AssemblyAI": {
"types": "./dist/readers/AssemblyAI.d.mts",
"import": "./dist/readers/AssemblyAI.mjs",
@@ -146,10 +151,10 @@
"import": "./dist/readers/base.mjs",
"require": "./dist/readers/base.js"
},
"./readers/CSVReader": {
"types": "./dist/readers/CSVReader.d.mts",
"import": "./dist/readers/CSVReader.mjs",
"require": "./dist/readers/CSVReader.js"
"./readers/PapaCSVReader": {
"types": "./dist/readers/PapaCSVReader.d.mts",
"import": "./dist/readers/PapaCSVReader.mjs",
"require": "./dist/readers/PapaCSVReader.js"
},
"./readers/DocxReader": {
"types": "./dist/readers/DocxReader.d.mts",
@@ -207,8 +212,8 @@
"scripts": {
"lint": "eslint .",
"test": "jest",
"build": "bunchee",
"dev": "bunchee -w",
"build": "NODE_OPTIONS=--max_old_space_size=8192 bunchee",
"dev": "NODE_OPTIONS=--max_old_space_size=8192 bunchee -w",
"circular-check": "madge --circular ./src/*.ts"
}
}
+3 -3
View File
@@ -1,7 +1,8 @@
import _ from "lodash";
import { ImageType } from "../Node";
import { DEFAULT_SIMILARITY_TOP_K } from "../constants";
import { DEFAULT_FS, VectorStoreQueryMode } from "../storage";
import { defaultFS } from "../env";
import { VectorStoreQueryMode } from "../storage";
import { SimilarityType } from "./types";
/**
@@ -244,8 +245,7 @@ export async function imageToDataUrl(input: ImageType): Promise<string> {
_.isString(input)
) {
// string or file URL
const fs = DEFAULT_FS;
const dataBuffer = await fs.readFile(
const dataBuffer = await defaultFS.readFile(
input instanceof URL ? input.pathname : input,
);
input = new Blob([dataBuffer]);
+3
View File
@@ -1,4 +1,5 @@
import { Sha256 } from "@aws-crypto/sha256-js";
import { CompleteFileSystem, InMemoryFileSystem } from "../storage";
export interface SHA256 {
update(data: string | Uint8Array): void;
@@ -8,6 +9,8 @@ export interface SHA256 {
export const EOL = "\n";
export const defaultFS: CompleteFileSystem = new InMemoryFileSystem();
export function ok(value: unknown, message?: string): asserts value {
if (!value) {
const error = Error(message);
+4
View File
@@ -1,6 +1,8 @@
import { ok } from "node:assert";
import { createHash, randomUUID } from "node:crypto";
import fs from "node:fs/promises";
import { EOL } from "node:os";
import type { CompleteFileSystem } from "../storage";
import type { SHA256 } from "./index.edge-light";
export function createSHA256(): SHA256 {
@@ -15,4 +17,6 @@ export function createSHA256(): SHA256 {
};
}
export const defaultFS: CompleteFileSystem = fs;
export { EOL, ok, randomUUID };
+1 -10
View File
@@ -21,15 +21,6 @@ export * from "./ingestion";
export * from "./llm";
export * from "./nodeParsers";
export * from "./postprocessors";
export * from "./readers/AssemblyAI";
export * from "./readers/CSVReader";
export * from "./readers/DocxReader";
export * from "./readers/HTMLReader";
export * from "./readers/MarkdownReader";
export * from "./readers/NotionReader";
export * from "./readers/PDFReader";
export * from "./readers/SimpleDirectoryReader";
export * from "./readers/SimpleMongoReader";
export * from "./readers/base";
export * from "./readers";
export * from "./storage";
export * from "./synthesizers";
+2 -2
View File
@@ -1,14 +1,14 @@
import mammoth from "mammoth";
import { Document } from "../Node";
import { defaultFS } from "../env";
import { GenericFileSystem } from "../storage/FileSystem";
import { DEFAULT_FS } from "../storage/constants";
import { BaseReader } from "./base";
export class DocxReader implements BaseReader {
/** DocxParser */
async loadData(
file: string,
fs: GenericFileSystem = DEFAULT_FS,
fs: GenericFileSystem = defaultFS,
): Promise<Document[]> {
const dataBuffer = (await fs.readFile(file)) as any;
const { value } = await mammoth.extractRawText({ buffer: dataBuffer });
+2 -2
View File
@@ -1,5 +1,5 @@
import { Document } from "../Node";
import { DEFAULT_FS } from "../storage/constants";
import { defaultFS } from "../env";
import { GenericFileSystem } from "../storage/FileSystem";
import { BaseReader } from "./base";
@@ -20,7 +20,7 @@ export class HTMLReader implements BaseReader {
*/
async loadData(
file: string,
fs: GenericFileSystem = DEFAULT_FS,
fs: GenericFileSystem = defaultFS,
): Promise<Document[]> {
const dataBuffer = await fs.readFile(file, "utf-8");
const htmlOptions = this.getOptions();
+2 -2
View File
@@ -1,5 +1,5 @@
import { Document, ImageDocument } from "../Node";
import { DEFAULT_FS } from "../storage/constants";
import { defaultFS } from "../env";
import { GenericFileSystem } from "../storage/FileSystem";
import { BaseReader } from "./base";
@@ -16,7 +16,7 @@ export class ImageReader implements BaseReader {
*/
async loadData(
file: string,
fs: GenericFileSystem = DEFAULT_FS,
fs: GenericFileSystem = defaultFS,
): Promise<Document[]> {
const dataBuffer = await fs.readFile(file);
const blob = new Blob([dataBuffer]);
+3 -2
View File
@@ -1,5 +1,6 @@
import { Document } from "../Node";
import { DEFAULT_FS, GenericFileSystem } from "../storage";
import { defaultFS } from "../env";
import { GenericFileSystem } from "../storage";
import { BaseReader } from "./base";
type MarkdownTuple = [string | null, string];
@@ -89,7 +90,7 @@ export class MarkdownReader implements BaseReader {
async loadData(
file: string,
fs: GenericFileSystem = DEFAULT_FS,
fs: GenericFileSystem = defaultFS,
): Promise<Document[]> {
const content = await fs.readFile(file, { encoding: "utf-8" });
const tups = this.parseTups(content);
+2 -2
View File
@@ -1,6 +1,6 @@
import { Document } from "../Node";
import { defaultFS } from "../env";
import { GenericFileSystem } from "../storage/FileSystem";
import { DEFAULT_FS } from "../storage/constants";
import { BaseReader } from "./base";
/**
@@ -9,7 +9,7 @@ import { BaseReader } from "./base";
export class PDFReader implements BaseReader {
async loadData(
file: string,
fs: GenericFileSystem = DEFAULT_FS,
fs: GenericFileSystem = defaultFS,
): Promise<Document[]> {
const content = (await fs.readFile(file)) as any;
if (!(content instanceof Buffer)) {
@@ -1,6 +1,7 @@
import Papa, { ParseConfig } from "papaparse";
import { Document } from "../Node";
import { DEFAULT_FS, GenericFileSystem } from "../storage/FileSystem";
import { defaultFS } from "../env";
import { GenericFileSystem } from "../storage/FileSystem";
import { BaseReader } from "./base";
/**
@@ -40,7 +41,7 @@ export class PapaCSVReader implements BaseReader {
*/
async loadData(
file: string,
fs: GenericFileSystem = DEFAULT_FS,
fs: GenericFileSystem = defaultFS,
): Promise<Document[]> {
const fileContent: string = await fs.readFile(file, "utf-8");
const result = Papa.parse(fileContent, this.papaConfig);
@@ -1,13 +1,13 @@
import _ from "lodash";
import { Document } from "../Node";
import { defaultFS } from "../env";
import { CompleteFileSystem, walk } from "../storage/FileSystem";
import { DEFAULT_FS } from "../storage/constants";
import { PapaCSVReader } from "./CSVReader";
import { DocxReader } from "./DocxReader";
import { HTMLReader } from "./HTMLReader";
import { ImageReader } from "./ImageReader";
import { MarkdownReader } from "./MarkdownReader";
import { PDFReader } from "./PDFReader";
import { PapaCSVReader } from "./PapaCSVReader";
import { BaseReader } from "./base";
type ReaderCallback = (
@@ -28,7 +28,7 @@ enum ReaderStatus {
export class TextFileReader implements BaseReader {
async loadData(
file: string,
fs: CompleteFileSystem = DEFAULT_FS as CompleteFileSystem,
fs: CompleteFileSystem = defaultFS,
): Promise<Document[]> {
const dataBuffer = await fs.readFile(file, "utf-8");
return [new Document({ text: dataBuffer, id_: file })];
@@ -66,7 +66,7 @@ export class SimpleDirectoryReader implements BaseReader {
async loadData({
directoryPath,
fs = DEFAULT_FS as CompleteFileSystem,
fs = defaultFS,
defaultReader = new TextFileReader(),
fileExtToReader = FILE_EXT_TO_READER,
}: SimpleDirectoryReaderLoadDataProps): Promise<Document[]> {
+25
View File
@@ -0,0 +1,25 @@
export {
AudioSubtitlesReader,
AudioTranscriptParagraphsReader,
AudioTranscriptReader,
AudioTranscriptSentencesReader,
} from "./AssemblyAI";
export type {
AssemblyAIOptions,
SubtitleFormat,
TranscribeParams,
} from "./AssemblyAI";
export { DocxReader } from "./DocxReader";
export { HTMLReader } from "./HTMLReader";
export { MarkdownReader } from "./MarkdownReader";
export { NotionReader } from "./NotionReader";
export { PDFReader } from "./PDFReader";
export { PapaCSVReader } from "./PapaCSVReader";
export {
FILE_EXT_TO_READER,
SimpleDirectoryReader,
TextFileReader,
type SimpleDirectoryReaderLoadDataProps,
} from "./SimpleDirectoryReader";
export { SimpleMongoReader } from "./SimpleMongoReader";
export type { BaseReader } from "./base";
+35 -32
View File
@@ -1,34 +1,51 @@
import _ from "lodash";
import type nodeFS from "node:fs/promises";
/**
* A filesystem interface that is meant to be compatible with
* the 'fs' module from Node.js.
* Allows for the use of similar inteface implementation on
* browsers.
*/
export interface GenericFileSystem {
writeFile(path: string, content: string, options?: any): Promise<void>;
readFile(path: string, options?: any): Promise<string>;
export type GenericFileSystem = {
writeFile(
path: string,
content: string,
options?: Parameters<typeof nodeFS.writeFile>[2],
): Promise<void>;
readFile(
path: string,
options?: Parameters<typeof nodeFS.readFile>[1],
): Promise<string>;
access(path: string): Promise<void>;
mkdir(path: string, options?: any): Promise<void>;
}
mkdir(
path: string,
options?: Parameters<typeof nodeFS.mkdir>[1],
): Promise<void>;
};
export interface WalkableFileSystem {
export type WalkableFileSystem = {
readdir(path: string): Promise<string[]>;
stat(path: string): Promise<any>;
}
};
export type CompleteFileSystem = GenericFileSystem & WalkableFileSystem;
/**
* A filesystem implementation that stores files in memory.
*/
export class InMemoryFileSystem implements GenericFileSystem {
export class InMemoryFileSystem implements CompleteFileSystem {
private files: Record<string, any> = {};
async writeFile(path: string, content: string, options?: any): Promise<void> {
async writeFile(
path: string,
content: string,
options?: unknown,
): Promise<void> {
this.files[path] = _.cloneDeep(content);
}
async readFile(path: string, options?: any): Promise<string> {
async readFile(path: string, options?: unknown): Promise<string> {
if (!(path in this.files)) {
throw new Error(`File ${path} does not exist`);
}
@@ -41,27 +58,19 @@ export class InMemoryFileSystem implements GenericFileSystem {
}
}
async mkdir(path: string, options?: any): Promise<void> {
async mkdir(path: string, options?: unknown): Promise<void> {
this.files[path] = _.get(this.files, path, null);
}
}
export type CompleteFileSystem = GenericFileSystem & WalkableFileSystem;
async readdir(path: string): Promise<string[]> {
throw new Error("Not implemented");
}
export function getNodeFS(): CompleteFileSystem {
const fs = require("fs/promises");
return fs;
async stat(path: string): Promise<any> {
throw new Error("Not implemented");
}
}
let fs = null;
try {
fs = getNodeFS();
} catch (e) {
fs = new InMemoryFileSystem();
}
export const DEFAULT_FS: GenericFileSystem | CompleteFileSystem =
fs as GenericFileSystem;
// FS utility functions
/**
@@ -92,12 +101,6 @@ export async function* walk(
fs: WalkableFileSystem,
dirPath: string,
): AsyncIterable<string> {
if (fs instanceof InMemoryFileSystem) {
throw new Error(
"The InMemoryFileSystem does not support directory traversal.",
);
}
const entries = await fs.readdir(dirPath);
for (const entry of entries) {
const fullPath = `${dirPath}/${entry}`;
+12 -15
View File
@@ -1,10 +1,7 @@
import path from "path";
import { defaultFS } from "../env";
import { GenericFileSystem } from "./FileSystem";
import {
DEFAULT_FS,
DEFAULT_IMAGE_VECTOR_NAMESPACE,
DEFAULT_NAMESPACE,
} from "./constants";
import { DEFAULT_IMAGE_VECTOR_NAMESPACE, DEFAULT_NAMESPACE } from "./constants";
import { SimpleDocumentStore } from "./docStore/SimpleDocumentStore";
import { BaseDocumentStore } from "./docStore/types";
import { SimpleIndexStore } from "./indexStore/SimpleIndexStore";
@@ -19,14 +16,14 @@ export interface StorageContext {
imageVectorStore?: VectorStore;
}
type BuilderParams = {
docStore?: BaseDocumentStore;
indexStore?: BaseIndexStore;
vectorStore?: VectorStore;
imageVectorStore?: VectorStore;
storeImages?: boolean;
persistDir?: string;
fs?: GenericFileSystem;
export type BuilderParams = {
docStore: BaseDocumentStore;
indexStore: BaseIndexStore;
vectorStore: VectorStore;
imageVectorStore: VectorStore;
storeImages: boolean;
persistDir: string;
fs: GenericFileSystem;
};
export async function storageContextFromDefaults({
@@ -37,14 +34,14 @@ export async function storageContextFromDefaults({
storeImages,
persistDir,
fs,
}: BuilderParams): Promise<StorageContext> {
}: Partial<BuilderParams>): Promise<StorageContext> {
if (!persistDir) {
docStore = docStore || new SimpleDocumentStore();
indexStore = indexStore || new SimpleIndexStore();
vectorStore = vectorStore || new SimpleVectorStore();
imageVectorStore = storeImages ? new SimpleVectorStore() : imageVectorStore;
} else {
fs = fs || DEFAULT_FS;
fs = fs || defaultFS;
docStore =
docStore ||
(await SimpleDocumentStore.fromPersistDir(
-1
View File
@@ -6,4 +6,3 @@ export const DEFAULT_VECTOR_STORE_PERSIST_FILENAME = "vector_store.json";
export const DEFAULT_GRAPH_STORE_PERSIST_FILENAME = "graph_store.json";
export const DEFAULT_NAMESPACE = "docstore";
export const DEFAULT_IMAGE_VECTOR_NAMESPACE = "images";
export { DEFAULT_FS } from "./FileSystem";
@@ -1,9 +1,9 @@
import _ from "lodash";
import path from "path";
import { defaultFS } from "../../env";
import { GenericFileSystem } from "../FileSystem";
import {
DEFAULT_DOC_STORE_PERSIST_FILENAME,
DEFAULT_FS,
DEFAULT_NAMESPACE,
DEFAULT_PERSIST_DIR,
} from "../constants";
@@ -44,7 +44,7 @@ export class SimpleDocumentStore extends KVDocumentStore {
namespace?: string,
fs?: GenericFileSystem,
): Promise<SimpleDocumentStore> {
fs = fs || DEFAULT_FS;
fs = fs || defaultFS;
const simpleKVStore = await SimpleKVStore.fromPersistPath(persistPath, fs);
return new SimpleDocumentStore(simpleKVStore, namespace);
}
@@ -56,7 +56,7 @@ export class SimpleDocumentStore extends KVDocumentStore {
),
fs?: GenericFileSystem,
): Promise<void> {
fs = fs || DEFAULT_FS;
fs = fs || defaultFS;
if (
_.isObject(this.kvStore) &&
this.kvStore instanceof BaseInMemoryKVStore
@@ -1,7 +1,7 @@
import path from "path";
import { defaultFS } from "../../env";
import { GenericFileSystem } from "../FileSystem";
import {
DEFAULT_FS,
DEFAULT_INDEX_STORE_PERSIST_FILENAME,
DEFAULT_PERSIST_DIR,
} from "../constants";
@@ -20,7 +20,7 @@ export class SimpleIndexStore extends KVIndexStore {
static async fromPersistDir(
persistDir: string = DEFAULT_PERSIST_DIR,
fs: GenericFileSystem = DEFAULT_FS,
fs: GenericFileSystem = defaultFS,
): Promise<SimpleIndexStore> {
const persistPath = path.join(
persistDir,
@@ -31,7 +31,7 @@ export class SimpleIndexStore extends KVIndexStore {
static async fromPersistPath(
persistPath: string,
fs: GenericFileSystem = DEFAULT_FS,
fs: GenericFileSystem = defaultFS,
): Promise<SimpleIndexStore> {
let simpleKVStore = await SimpleKVStore.fromPersistPath(persistPath, fs);
return new SimpleIndexStore(simpleKVStore);
@@ -39,7 +39,7 @@ export class SimpleIndexStore extends KVIndexStore {
async persist(
persistPath: string = DEFAULT_PERSIST_DIR,
fs: GenericFileSystem = DEFAULT_FS,
fs: GenericFileSystem = defaultFS,
): Promise<void> {
await this.kvStore.persist(persistPath, fs);
}
@@ -1,19 +1,18 @@
import _ from "lodash";
import path from "path";
import { defaultFS } from "../../env";
import { GenericFileSystem, exists } from "../FileSystem";
import { DEFAULT_COLLECTION, DEFAULT_FS } from "../constants";
import { DEFAULT_COLLECTION } from "../constants";
import { BaseKVStore } from "./types";
export type DataType = Record<string, Record<string, any>>;
export class SimpleKVStore extends BaseKVStore {
private data: DataType;
private persistPath: string | undefined;
private fs: GenericFileSystem | undefined;
constructor(data?: DataType) {
constructor(private data: DataType = {}) {
super();
this.data = data || {};
}
async put(
@@ -61,7 +60,7 @@ export class SimpleKVStore extends BaseKVStore {
}
async persist(persistPath: string, fs?: GenericFileSystem): Promise<void> {
fs = fs || DEFAULT_FS;
fs = fs || defaultFS;
// TODO: decide on a way to polyfill path
let dirPath = path.dirname(persistPath);
if (!(await exists(fs, dirPath))) {
@@ -74,7 +73,7 @@ export class SimpleKVStore extends BaseKVStore {
persistPath: string,
fs?: GenericFileSystem,
): Promise<SimpleKVStore> {
fs = fs || DEFAULT_FS;
fs = fs || defaultFS;
let dirPath = path.dirname(persistPath);
if (!(await exists(fs, dirPath))) {
await fs.mkdir(dirPath);
@@ -6,8 +6,9 @@ import {
getTopKEmbeddingsLearner,
getTopKMMREmbeddings,
} from "../../embeddings";
import { defaultFS } from "../../env";
import { GenericFileSystem, exists } from "../FileSystem";
import { DEFAULT_FS, DEFAULT_PERSIST_DIR } from "../constants";
import { DEFAULT_PERSIST_DIR } from "../constants";
import {
VectorStore,
VectorStoreQuery,
@@ -31,17 +32,17 @@ class SimpleVectorStoreData {
export class SimpleVectorStore implements VectorStore {
storesText: boolean = false;
private data: SimpleVectorStoreData = new SimpleVectorStoreData();
private fs: GenericFileSystem = DEFAULT_FS;
private fs: GenericFileSystem = defaultFS;
private persistPath: string | undefined;
constructor(data?: SimpleVectorStoreData, fs?: GenericFileSystem) {
this.data = data || new SimpleVectorStoreData();
this.fs = fs || DEFAULT_FS;
this.fs = fs || defaultFS;
}
static async fromPersistDir(
persistDir: string = DEFAULT_PERSIST_DIR,
fs: GenericFileSystem = DEFAULT_FS,
fs: GenericFileSystem = defaultFS,
): Promise<SimpleVectorStore> {
let persistPath = `${persistDir}/vector_store.json`;
return await SimpleVectorStore.fromPersistPath(persistPath, fs);
@@ -160,7 +161,7 @@ export class SimpleVectorStore implements VectorStore {
persistPath: string,
fs?: GenericFileSystem,
): Promise<SimpleVectorStore> {
fs = fs || DEFAULT_FS;
fs = fs || defaultFS;
let dirPath = path.dirname(persistPath);
if (!(await exists(fs, dirPath))) {
@@ -1,10 +1,10 @@
import nodeFS from "node:fs/promises";
import os from "os";
import path from "path";
import {
GenericFileSystem,
InMemoryFileSystem,
exists,
getNodeFS,
walk,
} from "../storage/FileSystem";
@@ -16,8 +16,6 @@ type FileSystemUnderTest = {
tempDir: string;
};
const nodeFS = getNodeFS() as GenericFileSystem & any;
describe.each<FileSystemUnderTest>([
{
name: "InMemoryFileSystem",
@@ -102,14 +100,13 @@ describe.each<FileSystemUnderTest>([
});
describe("Test walk for Node.js fs", () => {
const fs = getNodeFS();
let tempDir: string;
beforeAll(async () => {
tempDir = await nodeFS.mkdtemp(path.join(os.tmpdir(), "jest-"));
await fs.writeFile(`${tempDir}/test.txt`, "Hello, world!");
await fs.mkdir(`${tempDir}/subDir`);
await fs.writeFile(`${tempDir}/subDir/test2.txt`, "Hello, again!");
await nodeFS.writeFile(`${tempDir}/test.txt`, "Hello, world!");
await nodeFS.mkdir(`${tempDir}/subDir`);
await nodeFS.writeFile(`${tempDir}/subDir/test2.txt`, "Hello, again!");
});
it("walks directory", async () => {
@@ -119,7 +116,7 @@ describe("Test walk for Node.js fs", () => {
]);
const actualFiles = new Set<string>();
for await (let file of walk(fs, tempDir)) {
for await (let file of walk(nodeFS, tempDir)) {
expect(file).toBeTruthy();
actualFiles.add(file);
}