@superagent-ai/reag
Version:
Reasoning Augmented Generation
105 lines (104 loc) • 4.59 kB
JavaScript
;
Object.defineProperty(exports, "__esModule", { value: true });
exports.ReagClient = void 0;
const ai_1 = require("ai");
const prompt_1 = require("./prompt");
const schema_1 = require("./schema");
const DEFAULT_BATCH_SIZE = 20;
/**
* The Client class that wraps a language model and exposes query methods.
*/
class ReagClient {
/**
* Constructs a new Client instance.
* @param options Configuration options for the Client.
*/
constructor(options) {
this.model = options.model;
this.system = options.system || prompt_1.REAG_SYSTEM_PROMPT;
this.batchSize = options.batchSize || DEFAULT_BATCH_SIZE;
this.schema = options.schema || schema_1.RESPONSE_SCHEMA;
}
/**
* Filters documents based on metadata criteria
*/
filterDocumentsByMetadata(documents, filter) {
if (!(filter === null || filter === void 0 ? void 0 : filter.length))
return documents;
return documents.filter((doc) => {
return filter.every((filter) => {
var _a;
const metadataValue = (_a = doc.metadata) === null || _a === void 0 ? void 0 : _a[filter.key];
if (!metadataValue)
return false;
// Handle string operations only if both values are strings
if (typeof metadataValue === "string" &&
typeof filter.value === "string") {
switch (filter.operator) {
case "contains":
return metadataValue.includes(filter.value);
case "startsWith":
return metadataValue.startsWith(filter.value);
case "endsWith":
return metadataValue.endsWith(filter.value);
case "regex":
return new RegExp(filter.value).test(metadataValue);
}
}
// Handle numeric comparisons and equality checks
switch (filter.operator) {
case "equals":
return metadataValue === filter.value;
case "notEquals":
return metadataValue !== filter.value;
case "greaterThan":
return metadataValue > filter.value;
case "lessThan":
return metadataValue < filter.value;
case "greaterThanOrEqual":
return metadataValue >= filter.value;
case "lessThanOrEqual":
return metadataValue <= filter.value;
default:
return metadataValue === filter.value;
}
});
});
}
/**
* Executes a query on the assigned language model with document batching
*/
async query(prompt, documents, options) {
try {
const filteredDocuments = this.filterDocumentsByMetadata(documents, options === null || options === void 0 ? void 0 : options.filter);
const formatDoc = (doc) => `Name: ${doc.name}\nMetadata: ${JSON.stringify(doc.metadata)}\nContent: ${doc.content}`;
const batches = Array.from({ length: Math.ceil(filteredDocuments.length / this.batchSize) }, (_, i) => filteredDocuments.slice(i * this.batchSize, (i + 1) * this.batchSize));
const batchResults = await Promise.all(batches.map(async (batch) => {
// Process each document in the batch individually
const batchResponses = await Promise.all(batch.map(async (document) => {
const system = `${this.system}\n\n# Available source\n\n${formatDoc(document)}`;
const response = await (0, ai_1.generateObject)({
model: this.model,
system,
prompt,
schema: this.schema,
});
return {
response,
document,
};
}));
return batchResponses;
}));
const results = batchResults.flat().map(({ response, document }) => ({
...response.object,
document,
}));
return results;
}
catch (error) {
throw new Error(`Query failed: ${error}`);
}
}
}
exports.ReagClient = ReagClient;