Skip to content
Snippets Groups Projects
Commit e938a4d1 authored by yisding's avatar yisding
Browse files

minor changes

parent 64101926
No related branches found
No related tags found
No related merge requests found
import { SimpleDirectoryReader } from "llamaindex";
function callback(category: string, name: string, status: any, message?: string): boolean {
function callback(
category: string,
name: string,
status: any,
message?: string,
): boolean {
console.log(category, name, status, message);
if (name.endsWith('.pdf')) {
if (name.endsWith(".pdf")) {
console.log("I DON'T WANT PDF FILES!");
return false;
}
......@@ -12,7 +17,7 @@ function callback(category: string, name: string, status: any, message?: string)
async function main() {
// Load page
const reader = new SimpleDirectoryReader(callback);
const params = { directoryPath: "./data"};
const params = { directoryPath: "./data" };
await reader.loadData(params);
}
......
import _ from "lodash";
import { Document } from "../Node";
import { DEFAULT_FS } from "../storage/constants";
import { CompleteFileSystem, walk } from "../storage/FileSystem";
import { BaseReader } from "./base";
import { DEFAULT_FS } from "../storage/constants";
import { PapaCSVReader } from "./CSVReader";
import { DocxReader } from "./DocxReader";
import { HTMLReader } from "./HTMLReader";
import { MarkdownReader } from "./MarkdownReader";
import { PDFReader } from "./PDFReader";
import { BaseReader } from "./base";
type ReaderCallback = (category: string, name: string, status: ReaderStatus, message?: string) => boolean;
type ReaderCallback = (
category: "file" | "directory",
name: string,
status: ReaderStatus,
message?: string,
) => boolean;
enum ReaderStatus {
Started = 0,
Completed,
Error
STARTED = 0,
COMPLETE,
ERROR,
}
/**
......@@ -47,7 +52,7 @@ export type SimpleDirectoryReaderLoadDataProps = {
};
/**
* Read all of the documents in a directory.
* Read all of the documents in a directory.
* By default, supports the list of file types
* in the FILE_EXIT_TO_READER map.
*/
......@@ -60,12 +65,11 @@ export class SimpleDirectoryReader implements BaseReader {
defaultReader = new TextFileReader(),
fileExtToReader = FILE_EXT_TO_READER,
}: SimpleDirectoryReaderLoadDataProps): Promise<Document[]> {
// Observer can decide to skip the directory
if (this.doObserverCheck(
'Directory', directoryPath, ReaderStatus.Started
) == false) {
return Promise.reject('Cancelled');
if (
!this.doObserverCheck("directory", directoryPath, ReaderStatus.STARTED)
) {
return [];
}
let docs: Document[] = [];
......@@ -74,13 +78,11 @@ export class SimpleDirectoryReader implements BaseReader {
const fileExt = _.last(filePath.split(".")) || "";
// Observer can decide to skip each file
if (this.doObserverCheck(
'File', filePath, ReaderStatus.Started
) == false) {
if (!this.doObserverCheck("file", filePath, ReaderStatus.STARTED)) {
// Skip this file
continue;
}
}
let reader = null;
if (fileExt in fileExtToReader) {
......@@ -92,50 +94,45 @@ export class SimpleDirectoryReader implements BaseReader {
console.warn(msg);
// In an error condition, observer's false cancels the whole process.
if (this.doObserverCheck(
'File', filePath, ReaderStatus.Error, msg
) == false) {
return this.getCancelled();
if (
!this.doObserverCheck("file", filePath, ReaderStatus.ERROR, msg)
) {
return [];
}
continue;
}
const fileDocs = await reader.loadData(filePath, fs);
// Observer can still cancel addition of the resulting docs from this file
if (this.doObserverCheck(
'File', filePath, ReaderStatus.Completed
)) {
if (this.doObserverCheck("file", filePath, ReaderStatus.COMPLETE)) {
docs.push(...fileDocs);
}
}
} catch (e) {
const msg = `Error reading file ${filePath}: ${e}`;
console.error(msg);
// In an error condition, observer's false cancels the whole process.
if (this.doObserverCheck(
'File', filePath, ReaderStatus.Error, msg
) == false) {
return this.getCancelled();
if (!this.doObserverCheck("file", filePath, ReaderStatus.ERROR, msg)) {
return [];
}
}
}
// After successful import of all files, directory completion
// is only a notification for observer, cannot be cancelled.
this.doObserverCheck(
'Directory', directoryPath, ReaderStatus.Completed
);
this.doObserverCheck("directory", directoryPath, ReaderStatus.COMPLETE);
return docs;
}
private getCancelled() {
return Promise.reject('Cancelled');
}
private doObserverCheck(category: string, name: string, status: ReaderStatus, message?: string): boolean {
private doObserverCheck(
category: "file" | "directory",
name: string,
status: ReaderStatus,
message?: string,
): boolean {
if (this.observer) {
return this.observer(category, name, status, message);
}
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment