82 lines
2.9 KiB
JavaScript
82 lines
2.9 KiB
JavaScript
import { Storage } from "@google-cloud/storage";
|
|
import { Document } from "@langchain/core/documents";
|
|
import { Docstore } from "langchain/stores/doc/base";
|
|
/**
|
|
* Class that provides an interface for interacting with Google Cloud
|
|
* Storage (GCS) as a document store. It extends the Docstore class and
|
|
* implements methods to search, add, and add a document to the GCS
|
|
* bucket.
|
|
*/
|
|
export class GoogleCloudStorageDocstore extends Docstore {
|
|
constructor(config) {
|
|
super();
|
|
Object.defineProperty(this, "bucket", {
|
|
enumerable: true,
|
|
configurable: true,
|
|
writable: true,
|
|
value: void 0
|
|
});
|
|
Object.defineProperty(this, "prefix", {
|
|
enumerable: true,
|
|
configurable: true,
|
|
writable: true,
|
|
value: ""
|
|
});
|
|
Object.defineProperty(this, "storage", {
|
|
enumerable: true,
|
|
configurable: true,
|
|
writable: true,
|
|
value: void 0
|
|
});
|
|
this.bucket = config.bucket;
|
|
this.prefix = config.prefix ?? this.prefix;
|
|
this.storage = new Storage();
|
|
}
|
|
/**
|
|
* Searches for a document in the GCS bucket and returns it as a Document
|
|
* instance.
|
|
* @param search The name of the document to search for in the GCS bucket
|
|
* @returns A Promise that resolves to a Document instance representing the found document
|
|
*/
|
|
async search(search) {
|
|
const file = this.getFile(search);
|
|
const [fileMetadata] = await file.getMetadata();
|
|
const metadata = fileMetadata?.metadata;
|
|
const [dataBuffer] = await file.download();
|
|
const pageContent = dataBuffer.toString();
|
|
const ret = new Document({
|
|
pageContent,
|
|
metadata,
|
|
});
|
|
return ret;
|
|
}
|
|
/**
|
|
* Adds multiple documents to the GCS bucket.
|
|
* @param texts An object where each key is the name of a document and the value is the Document instance to be added
|
|
* @returns A Promise that resolves when all documents have been added
|
|
*/
|
|
async add(texts) {
|
|
await Promise.all(Object.keys(texts).map((key) => this.addDocument(key, texts[key])));
|
|
}
|
|
/**
|
|
* Adds a single document to the GCS bucket.
|
|
* @param name The name of the document to be added
|
|
* @param document The Document instance to be added
|
|
* @returns A Promise that resolves when the document has been added
|
|
*/
|
|
async addDocument(name, document) {
|
|
const file = this.getFile(name);
|
|
await file.save(document.pageContent);
|
|
await file.setMetadata({ metadata: document.metadata });
|
|
}
|
|
/**
|
|
* Gets a file from the GCS bucket.
|
|
* @param name The name of the file to get from the GCS bucket
|
|
* @returns A File instance representing the fetched file
|
|
*/
|
|
getFile(name) {
|
|
const filename = this.prefix + name;
|
|
const file = this.storage.bucket(this.bucket).file(filename);
|
|
return file;
|
|
}
|
|
}
|