104 lines
3.4 KiB
JavaScript
104 lines
3.4 KiB
JavaScript
import { BaseRetriever, } from "@langchain/core/retrievers";
|
|
import { Document } from "@langchain/core/documents";
|
|
import { getEnvironmentVariable } from "@langchain/core/utils/env";
|
|
import { Dria } from "dria";
|
|
/**
|
|
* Class for retrieving documents from knowledge uploaded to Dria.
|
|
*
|
|
* @example
|
|
* ```typescript
|
|
* // contract of TypeScript Handbook v4.9 uploaded to Dria
|
|
* const contractId = "-B64DjhUtCwBdXSpsRytlRQCu-bie-vSTvTIT8Ap3g0";
|
|
* const retriever = new DriaRetriever({ contractId });
|
|
*
|
|
* const docs = await retriever.getRelevantDocuments("What is a union type?");
|
|
* console.log(docs);
|
|
* ```
|
|
*/
|
|
export class DriaRetriever extends BaseRetriever {
|
|
static lc_name() {
|
|
return "DriaRetriever";
|
|
}
|
|
get lc_secrets() {
|
|
return { apiKey: "DRIA_API_KEY" };
|
|
}
|
|
get lc_aliases() {
|
|
return { apiKey: "api_key" };
|
|
}
|
|
constructor(fields) {
|
|
super(fields);
|
|
Object.defineProperty(this, "lc_namespace", {
|
|
enumerable: true,
|
|
configurable: true,
|
|
writable: true,
|
|
value: ["langchain", "retrievers", "dria"]
|
|
});
|
|
Object.defineProperty(this, "apiKey", {
|
|
enumerable: true,
|
|
configurable: true,
|
|
writable: true,
|
|
value: void 0
|
|
});
|
|
Object.defineProperty(this, "driaClient", {
|
|
enumerable: true,
|
|
configurable: true,
|
|
writable: true,
|
|
value: void 0
|
|
});
|
|
Object.defineProperty(this, "searchOptions", {
|
|
enumerable: true,
|
|
configurable: true,
|
|
writable: true,
|
|
value: void 0
|
|
});
|
|
const apiKey = fields.apiKey ?? getEnvironmentVariable("DRIA_API_KEY");
|
|
if (!apiKey)
|
|
throw new Error("Missing DRIA_API_KEY.");
|
|
this.apiKey = apiKey;
|
|
this.searchOptions = {
|
|
topK: fields.topK,
|
|
field: fields.field,
|
|
rerank: fields.rerank,
|
|
level: fields.level,
|
|
};
|
|
this.driaClient = new Dria({
|
|
contractId: fields.contractId,
|
|
apiKey: this.apiKey,
|
|
});
|
|
}
|
|
/**
|
|
* Currently connected knowledge on Dria.
|
|
*
|
|
* Retriever will use this contract ID while retrieving documents,
|
|
* and will throw an error if `undefined`.
|
|
*
|
|
* In the case that this is `undefined`, the user is expected to
|
|
* set contract ID manually, such as after creating a new knowledge & inserting
|
|
* data there with the Dria client.
|
|
*/
|
|
get contractId() {
|
|
return this.driaClient.contractId;
|
|
}
|
|
set contractId(value) {
|
|
this.driaClient.contractId = value;
|
|
}
|
|
/**
|
|
* Retrieves documents from Dria with respect to the configured contract ID, based on
|
|
* the given query string.
|
|
*
|
|
* @param query The query string
|
|
* @returns A promise that resolves to an array of documents, with page content as text,
|
|
* along with `id` and the relevance `score` within the metadata.
|
|
*/
|
|
async _getRelevantDocuments(query) {
|
|
const docs = await this.driaClient.search(query, this.searchOptions);
|
|
return docs.map((d) => new Document({
|
|
// dria.search returns a string within the metadata as the content
|
|
pageContent: d.metadata,
|
|
metadata: {
|
|
id: d.id,
|
|
score: d.score,
|
|
},
|
|
}));
|
|
}
|
|
}
|