diff --git a/.changeset/three-hounds-count.md b/.changeset/three-hounds-count.md
new file mode 100644
index 0000000000000000000000000000000000000000..a9ec11e53e3985296b02604fe7d9bd4f491f2a1b
--- /dev/null
+++ b/.changeset/three-hounds-count.md
@@ -0,0 +1,5 @@
+---
+"llamaindex": patch
+---
+
+feat(extractors): add keyword extractor and base extractor
diff --git a/examples/extractors/keywordExtractor.ts b/examples/extractors/keywordExtractor.ts
new file mode 100644
index 0000000000000000000000000000000000000000..0a329c898336267c2c2a33ce7162d101b9bab99d
--- /dev/null
+++ b/examples/extractors/keywordExtractor.ts
@@ -0,0 +1,24 @@
+import {
+ Document,
+ KeywordExtractor,
+ OpenAI,
+ SimpleNodeParser,
+} from "llamaindex";
+
+(async () => {
+ const openaiLLM = new OpenAI({ model: "gpt-3.5-turbo", temperature: 0 });
+
+ const nodeParser = new SimpleNodeParser();
+
+ const nodes = nodeParser.getNodesFromDocuments([
+ new Document({ text: "banana apple orange pear peach watermelon" }),
+ ]);
+
+ console.log(nodes);
+
+ const keywordExtractor = new KeywordExtractor(openaiLLM, 5);
+
+ const nodesWithKeywordMetadata = await keywordExtractor.processNodes(nodes);
+
+ process.stdout.write(JSON.stringify(nodesWithKeywordMetadata, null, 2));
+})();
diff --git a/examples/extractors/questionsAnsweredExtractor.ts b/examples/extractors/questionsAnsweredExtractor.ts
new file mode 100644
index 0000000000000000000000000000000000000000..3eb6161442602ddaa9f0b85958416bdf1c6d62b9
--- /dev/null
+++ b/examples/extractors/questionsAnsweredExtractor.ts
@@ -0,0 +1,31 @@
+import {
+ Document,
+ OpenAI,
+ QuestionsAnsweredExtractor,
+ SimpleNodeParser,
+} from "llamaindex";
+
+(async () => {
+ const openaiLLM = new OpenAI({ model: "gpt-3.5-turbo", temperature: 0 });
+
+ const nodeParser = new SimpleNodeParser();
+
+ const nodes = nodeParser.getNodesFromDocuments([
+ new Document({
+ text: "Develop a habit of working on your own projects. Don't let work mean something other people tell you to do. If you do manage to do great work one day, it will probably be on a project of your own. It may be within some bigger project, but you'll be driving your part of it.",
+ }),
+ new Document({
+ text: "The best way to get a good idea is to get a lot of ideas. The best way to get a lot of ideas is to get a lot of bad ideas. The best way to get a lot of bad ideas is to get a lot of ideas.",
+ }),
+ ]);
+
+ const questionsAnsweredExtractor = new QuestionsAnsweredExtractor(
+ openaiLLM,
+ 5,
+ );
+
+ const nodesWithQuestionsMetadata =
+ await questionsAnsweredExtractor.processNodes(nodes);
+
+ process.stdout.write(JSON.stringify(nodesWithQuestionsMetadata, null, 2));
+})();
diff --git a/examples/extractors/summaryExtractor.ts b/examples/extractors/summaryExtractor.ts
new file mode 100644
index 0000000000000000000000000000000000000000..b6c5bd217cba409af31a526646f6d5a7becc3f2c
--- /dev/null
+++ b/examples/extractors/summaryExtractor.ts
@@ -0,0 +1,24 @@
+import {
+ Document,
+ OpenAI,
+ SimpleNodeParser,
+ SummaryExtractor,
+} from "llamaindex";
+
+(async () => {
+ const openaiLLM = new OpenAI({ model: "gpt-3.5-turbo", temperature: 0 });
+
+ const nodeParser = new SimpleNodeParser();
+
+ const nodes = nodeParser.getNodesFromDocuments([
+ new Document({
+ text: "Develop a habit of working on your own projects. Don't let work mean something other people tell you to do. If you do manage to do great work one day, it will probably be on a project of your own. It may be within some bigger project, but you'll be driving your part of it.",
+ }),
+ ]);
+
+ const summaryExtractor = new SummaryExtractor(openaiLLM);
+
+ const nodesWithSummaryMetadata = await summaryExtractor.processNodes(nodes);
+
+ process.stdout.write(JSON.stringify(nodesWithSummaryMetadata, null, 2));
+})();
diff --git a/examples/extractors/titleExtractor.ts b/examples/extractors/titleExtractor.ts
new file mode 100644
index 0000000000000000000000000000000000000000..95fab03274092383fca748367e925672a4ff423c
--- /dev/null
+++ b/examples/extractors/titleExtractor.ts
@@ -0,0 +1,19 @@
+import { Document, OpenAI, SimpleNodeParser, TitleExtractor } from "llamaindex";
+
+(async () => {
+ const openaiLLM = new OpenAI({ model: "gpt-3.5-turbo", temperature: 0 });
+
+ const nodeParser = new SimpleNodeParser();
+
+ const nodes = nodeParser.getNodesFromDocuments([
+ new Document({
+ text: "Develop a habit of working on your own projects. Don't let work mean something other people tell you to do. If you do manage to do great work one day, it will probably be on a project of your own. It may be within some bigger project, but you'll be driving your part of it.",
+ }),
+ ]);
+
+ const titleExtractor = new TitleExtractor(openaiLLM, 1);
+
+ const nodesWithTitledMetadata = await titleExtractor.processNodes(nodes);
+
+ process.stdout.write(JSON.stringify(nodesWithTitledMetadata, null, 2));
+})();
diff --git a/packages/core/src/Node.ts b/packages/core/src/Node.ts
index a4a2dcfadc5035931c3f99dd7851c621047bff9f..806bba5d4340de6c2e31c773700ccb1befb45d61 100644
--- a/packages/core/src/Node.ts
+++ b/packages/core/src/Node.ts
@@ -168,6 +168,8 @@ export abstract class BaseNode<T extends Metadata = Metadata> {
*/
export class TextNode<T extends Metadata = Metadata> extends BaseNode<T> {
text: string = "";
+ textTemplate: string = "";
+
startCharIdx?: number;
endCharIdx?: number;
// textTemplate: NOTE write your own formatter if needed
diff --git a/packages/core/src/extractors/MetadataExtractors.ts b/packages/core/src/extractors/MetadataExtractors.ts
new file mode 100644
index 0000000000000000000000000000000000000000..8fc884c12b7f57bb25e42639055989efbbc87d75
--- /dev/null
+++ b/packages/core/src/extractors/MetadataExtractors.ts
@@ -0,0 +1,410 @@
+import { BaseNode, MetadataMode, TextNode } from "../Node";
+import { LLM } from "../llm";
+import {
+ defaultKeywordExtractorPromptTemplate,
+ defaultQuestionAnswerPromptTemplate,
+ defaultSummaryExtractorPromptTemplate,
+ defaultTitleCombinePromptTemplate,
+ defaultTitleExtractorPromptTemplate,
+} from "./prompts";
+import { BaseExtractor } from "./types";
+
+const STRIP_REGEX = /(\r\n|\n|\r)/gm;
+
+type ExtractKeyword = {
+ excerptKeywords: string;
+};
+
+/**
+ * Extract keywords from a list of nodes.
+ */
+export class KeywordExtractor extends BaseExtractor {
+ /**
+ * LLM instance.
+ * @type {LLM}
+ */
+ llm: LLM;
+
+ /**
+ * Number of keywords to extract.
+ * @type {number}
+ * @default 5
+ */
+ keywords: number = 5;
+
+ /**
+ * Constructor for the KeywordExtractor class.
+ * @param {LLM} llm LLM instance.
+ * @param {number} keywords Number of keywords to extract.
+ * @throws {Error} If keywords is less than 1.
+ */
+ constructor(llm: LLM, keywords: number = 5) {
+ if (keywords < 1) throw new Error("Keywords must be greater than 0");
+
+ super();
+ this.llm = llm;
+ this.keywords = keywords;
+ }
+
+ /**
+ *
+ * @param node Node to extract keywords from.
+ * @returns Keywords extracted from the node.
+ */
+ async extractKeywordsFromNodes(node: BaseNode): Promise<ExtractKeyword | {}> {
+ if (this.isTextNodeOnly && !(node instanceof TextNode)) {
+ return {};
+ }
+
+ const completion = await this.llm.complete({
+ prompt: defaultKeywordExtractorPromptTemplate({
+ contextStr: node.getContent(MetadataMode.ALL),
+ keywords: this.keywords,
+ }),
+ });
+
+ return {
+ excerptKeywords: completion.text,
+ };
+ }
+
+ /**
+ *
+ * @param nodes Nodes to extract keywords from.
+ * @returns Keywords extracted from the nodes.
+ */
+ async extract(nodes: BaseNode[]): Promise<Array<ExtractKeyword> | Array<{}>> {
+ const results = await Promise.all(
+ nodes.map((node) => this.extractKeywordsFromNodes(node)),
+ );
+ return results;
+ }
+}
+
+type ExtractTitle = {
+ documentTitle: string;
+};
+
+/**
+ * Extract title from a list of nodes.
+ */
+export class TitleExtractor extends BaseExtractor {
+ /**
+ * LLM instance.
+ * @type {LLM}
+ */
+ llm: LLM;
+
+ /**
+ * Can work for mixture of text and non-text nodes
+ * @type {boolean}
+ * @default false
+ */
+ isTextNodeOnly: boolean = false;
+
+ /**
+ * Number of nodes to extrct titles from.
+ * @type {number}
+ * @default 5
+ */
+ nodes: number = 5;
+
+ /**
+ * The prompt template to use for the title extractor.
+ * @type {string}
+ */
+ nodeTemplate: string;
+
+ /**
+ * The prompt template to merge title with..
+ * @type {string}
+ */
+ combineTemplate: string;
+
+ /**
+ * Constructor for the TitleExtractor class.
+ * @param {LLM} llm LLM instance.
+ * @param {number} nodes Number of nodes to extract titles from.
+ * @param {string} node_template The prompt template to use for the title extractor.
+ * @param {string} combine_template The prompt template to merge title with..
+ */
+ constructor(
+ llm: LLM,
+ nodes: number = 5,
+ node_template?: string,
+ combine_template?: string,
+ ) {
+ super();
+
+ this.llm = llm;
+ this.nodes = nodes;
+
+ this.nodeTemplate = node_template ?? defaultTitleExtractorPromptTemplate();
+ this.combineTemplate =
+ combine_template ?? defaultTitleCombinePromptTemplate();
+ }
+
+ /**
+ * Extract titles from a list of nodes.
+ * @param {BaseNode[]} nodes Nodes to extract titles from.
+ * @returns {Promise<BaseNode<ExtractTitle>[]>} Titles extracted from the nodes.
+ */
+ async extract(nodes: BaseNode[]): Promise<Array<ExtractTitle>> {
+ const nodesToExtractTitle: BaseNode[] = [];
+
+ for (let i = 0; i < this.nodes; i++) {
+ if (nodesToExtractTitle.length >= nodes.length) break;
+
+ if (this.isTextNodeOnly && !(nodes[i] instanceof TextNode)) continue;
+
+ nodesToExtractTitle.push(nodes[i]);
+ }
+
+ if (nodesToExtractTitle.length === 0) return [];
+
+ let titlesCandidates: string[] = [];
+ let title: string = "";
+
+ for (let i = 0; i < nodesToExtractTitle.length; i++) {
+ const completion = await this.llm.complete({
+ prompt: defaultTitleExtractorPromptTemplate({
+ contextStr: nodesToExtractTitle[i].getContent(MetadataMode.ALL),
+ }),
+ });
+
+ titlesCandidates.push(completion.text);
+ }
+
+ if (nodesToExtractTitle.length > 1) {
+ const combinedTitles = titlesCandidates.join(",");
+
+ const completion = await this.llm.complete({
+ prompt: defaultTitleCombinePromptTemplate({
+ contextStr: combinedTitles,
+ }),
+ });
+
+ title = completion.text;
+ }
+
+ if (nodesToExtractTitle.length === 1) {
+ title = titlesCandidates[0];
+ }
+
+ return nodes.map((_) => ({
+ documentTitle: title.trim().replace(STRIP_REGEX, ""),
+ }));
+ }
+}
+
+type ExtractQuestion = {
+ questionsThisExcerptCanAnswer: string;
+};
+
+/**
+ * Extract questions from a list of nodes.
+ */
+export class QuestionsAnsweredExtractor extends BaseExtractor {
+ /**
+ * LLM instance.
+ * @type {LLM}
+ */
+ llm: LLM;
+
+ /**
+ * Number of questions to generate.
+ * @type {number}
+ * @default 5
+ */
+ questions: number = 5;
+
+ /**
+ * The prompt template to use for the question extractor.
+ * @type {string}
+ */
+ promptTemplate: string;
+
+ /**
+ * Wheter to use metadata for embeddings only
+ * @type {boolean}
+ * @default false
+ */
+ embeddingOnly: boolean = false;
+
+ /**
+ * Constructor for the QuestionsAnsweredExtractor class.
+ * @param {LLM} llm LLM instance.
+ * @param {number} questions Number of questions to generate.
+ * @param {string} promptTemplate The prompt template to use for the question extractor.
+ * @param {boolean} embeddingOnly Wheter to use metadata for embeddings only.
+ */
+ constructor(
+ llm: LLM,
+ questions: number = 5,
+ promptTemplate?: string,
+ embeddingOnly: boolean = false,
+ ) {
+ if (questions < 1) throw new Error("Questions must be greater than 0");
+
+ super();
+
+ this.llm = llm;
+ this.questions = questions;
+ this.promptTemplate =
+ promptTemplate ??
+ defaultQuestionAnswerPromptTemplate({
+ numQuestions: questions,
+ contextStr: "",
+ });
+ this.embeddingOnly = embeddingOnly;
+ }
+
+ /**
+ * Extract answered questions from a node.
+ * @param {BaseNode} node Node to extract questions from.
+ * @returns {Promise<Array<ExtractQuestion> | Array<{}>>} Questions extracted from the node.
+ */
+ async extractQuestionsFromNode(
+ node: BaseNode,
+ ): Promise<ExtractQuestion | {}> {
+ if (this.isTextNodeOnly && !(node instanceof TextNode)) {
+ return {};
+ }
+
+ const contextStr = node.getContent(this.metadataMode);
+
+ const prompt = defaultQuestionAnswerPromptTemplate({
+ contextStr,
+ numQuestions: this.questions,
+ });
+
+ const questions = await this.llm.complete({
+ prompt,
+ });
+
+ return {
+ questionsThisExcerptCanAnswer: questions.text.replace(STRIP_REGEX, ""),
+ };
+ }
+
+ /**
+ * Extract answered questions from a list of nodes.
+ * @param {BaseNode[]} nodes Nodes to extract questions from.
+ * @returns {Promise<Array<ExtractQuestion> | Array<{}>>} Questions extracted from the nodes.
+ */
+ async extract(
+ nodes: BaseNode[],
+ ): Promise<Array<ExtractQuestion> | Array<{}>> {
+ const results = await Promise.all(
+ nodes.map((node) => this.extractQuestionsFromNode(node)),
+ );
+
+ return results;
+ }
+}
+
+type ExtractSummary = {
+ sectionSummary: string;
+ prevSectionSummary: string;
+ nextSectionSummary: string;
+};
+
+/**
+ * Extract summary from a list of nodes.
+ */
+export class SummaryExtractor extends BaseExtractor {
+ /**
+ * LLM instance.
+ * @type {LLM}
+ */
+ llm: LLM;
+
+ /**
+ * List of summaries to extract: 'self', 'prev', 'next'
+ * @type {string[]}
+ */
+ summaries: string[];
+
+ /**
+ * The prompt template to use for the summary extractor.
+ * @type {string}
+ */
+ promptTemplate: string;
+
+ private _selfSummary: boolean;
+ private _prevSummary: boolean;
+ private _nextSummary: boolean;
+
+ constructor(
+ llm: LLM,
+ summaries: string[] = ["self"],
+ promptTemplate?: string,
+ ) {
+ if (!summaries.some((s) => ["self", "prev", "next"].includes(s)))
+ throw new Error("Summaries must be one of 'self', 'prev', 'next'");
+
+ super();
+
+ this.llm = llm;
+ this.summaries = summaries;
+ this.promptTemplate =
+ promptTemplate ?? defaultSummaryExtractorPromptTemplate();
+
+ this._selfSummary = summaries.includes("self");
+ this._prevSummary = summaries.includes("prev");
+ this._nextSummary = summaries.includes("next");
+ }
+
+ /**
+ * Extract summary from a node.
+ * @param {BaseNode} node Node to extract summary from.
+ * @returns {Promise<string>} Summary extracted from the node.
+ */
+ async generateNodeSummary(node: BaseNode): Promise<string> {
+ if (this.isTextNodeOnly && !(node instanceof TextNode)) {
+ return "";
+ }
+
+ const contextStr = node.getContent(this.metadataMode);
+
+ const prompt = defaultSummaryExtractorPromptTemplate({
+ contextStr,
+ });
+
+ const summary = await this.llm.complete({
+ prompt,
+ });
+
+ return summary.text.replace(STRIP_REGEX, "");
+ }
+
+ /**
+ * Extract summaries from a list of nodes.
+ * @param {BaseNode[]} nodes Nodes to extract summaries from.
+ * @returns {Promise<Array<ExtractSummary> | Arry<{}>>} Summaries extracted from the nodes.
+ */
+ async extract(nodes: BaseNode[]): Promise<Array<ExtractSummary> | Array<{}>> {
+ if (!nodes.every((n) => n instanceof TextNode))
+ throw new Error("Only `TextNode` is allowed for `Summary` extractor");
+
+ const nodeSummaries = await Promise.all(
+ nodes.map((node) => this.generateNodeSummary(node)),
+ );
+
+ let metadataList: any[] = nodes.map(() => ({}));
+
+ for (let i = 0; i < nodes.length; i++) {
+ if (i > 0 && this._prevSummary && nodeSummaries[i - 1]) {
+ metadataList[i]["prevSectionSummary"] = nodeSummaries[i - 1];
+ }
+ if (i < nodes.length - 1 && this._nextSummary && nodeSummaries[i + 1]) {
+ metadataList[i]["nextSectionSummary"] = nodeSummaries[i + 1];
+ }
+ if (this._selfSummary && nodeSummaries[i]) {
+ metadataList[i]["sectionSummary"] = nodeSummaries[i];
+ }
+ }
+
+ return metadataList;
+ }
+}
diff --git a/packages/core/src/extractors/index.ts b/packages/core/src/extractors/index.ts
new file mode 100644
index 0000000000000000000000000000000000000000..41a324d26b90c6a1a00ccc6e935f55a5015d0567
--- /dev/null
+++ b/packages/core/src/extractors/index.ts
@@ -0,0 +1,6 @@
+export {
+ KeywordExtractor,
+ QuestionsAnsweredExtractor,
+ SummaryExtractor,
+ TitleExtractor,
+} from "./MetadataExtractors";
diff --git a/packages/core/src/extractors/prompts.ts b/packages/core/src/extractors/prompts.ts
new file mode 100644
index 0000000000000000000000000000000000000000..47a798c48f7a6b39b7b6f2545dd9972a77629579
--- /dev/null
+++ b/packages/core/src/extractors/prompts.ts
@@ -0,0 +1,81 @@
+export interface DefaultPromptTemplate {
+ contextStr: string;
+}
+
+export interface DefaultKeywordExtractorPromptTemplate
+ extends DefaultPromptTemplate {
+ keywords: number;
+}
+
+export interface DefaultQuestionAnswerPromptTemplate
+ extends DefaultPromptTemplate {
+ numQuestions: number;
+}
+
+export interface DefaultNodeTextTemplate {
+ metadataStr: string;
+ content: string;
+}
+
+export const defaultKeywordExtractorPromptTemplate = ({
+ contextStr = "",
+ keywords = 5,
+}: DefaultKeywordExtractorPromptTemplate) => `
+ ${contextStr}. Give ${keywords} unique keywords for thiss
+ document. Format as comma separated. Keywords:
+`;
+
+export const defaultTitleExtractorPromptTemplate = (
+ { contextStr = "" }: DefaultPromptTemplate = {
+ contextStr: "",
+ },
+) => `
+ ${contextStr}. Give a title that summarizes all of the unique entities, titles or themes found in the context. Title:
+`;
+
+export const defaultTitleCombinePromptTemplate = (
+ { contextStr = "" }: DefaultPromptTemplate = {
+ contextStr: "",
+ },
+) => `
+ ${contextStr}. Based on the above candidate titles and content,s
+ what is the comprehensive title for this document? Title:
+`;
+
+export const defaultQuestionAnswerPromptTemplate = (
+ { contextStr = "", numQuestions = 5 }: DefaultQuestionAnswerPromptTemplate = {
+ contextStr: "",
+ numQuestions: 5,
+ },
+) => `
+ ${contextStr}. Given the contextual information,s
+ generate ${numQuestions} questions this context can provides
+ specific answers to which are unlikely to be found elsewhere.
+
+ Higher-level summaries of surrounding context may be provideds
+ as well. Try using these summaries to generate better questionss
+ that this context can answer.
+`;
+
+export const defaultSummaryExtractorPromptTemplate = (
+ { contextStr = "" }: DefaultPromptTemplate = {
+ contextStr: "",
+ },
+) => `
+ ${contextStr}. Summarize the key topics and entities of the section.s
+ Summary:
+`;
+
+export const defaultNodeTextTemplate = ({
+ metadataStr = "",
+ content = "",
+}: {
+ metadataStr?: string;
+ content?: string;
+} = {}) => `[Excerpt from document]
+${metadataStr}
+Excerpt:
+-----
+${content}
+-----
+`;
diff --git a/packages/core/src/extractors/types.ts b/packages/core/src/extractors/types.ts
new file mode 100644
index 0000000000000000000000000000000000000000..28ec7f6a5b87a6862ed3b67bd053e6411d9a4a3c
--- /dev/null
+++ b/packages/core/src/extractors/types.ts
@@ -0,0 +1,64 @@
+import { BaseNode, MetadataMode, TextNode } from "../Node";
+import { defaultNodeTextTemplate } from "./prompts";
+
+/*
+ * Abstract class for all extractors.
+ */
+export abstract class BaseExtractor {
+ isTextNodeOnly: boolean = true;
+ showProgress: boolean = true;
+ metadataMode: MetadataMode = MetadataMode.ALL;
+ disableTemplateRewrite: boolean = false;
+ inPlace: boolean = true;
+ numWorkers: number = 4;
+
+ abstract extract(nodes: BaseNode[]): Promise<Record<string, any>[]>;
+
+ /**
+ *
+ * @param nodes Nodes to extract metadata from.
+ * @param excludedEmbedMetadataKeys Metadata keys to exclude from the embedding.
+ * @param excludedLlmMetadataKeys Metadata keys to exclude from the LLM.
+ * @returns Metadata extracted from the nodes.
+ */
+ async processNodes(
+ nodes: BaseNode[],
+ excludedEmbedMetadataKeys: string[] | undefined = undefined,
+ excludedLlmMetadataKeys: string[] | undefined = undefined,
+ ): Promise<BaseNode[]> {
+ let newNodes: BaseNode[];
+
+ if (this.inPlace) {
+ newNodes = nodes;
+ } else {
+ newNodes = nodes.slice();
+ }
+
+ let curMetadataList = await this.extract(newNodes);
+
+ for (let idx in newNodes) {
+ newNodes[idx].metadata = curMetadataList[idx];
+ }
+
+ for (let idx in newNodes) {
+ if (excludedEmbedMetadataKeys) {
+ newNodes[idx].excludedEmbedMetadataKeys.concat(
+ excludedEmbedMetadataKeys,
+ );
+ }
+ if (excludedLlmMetadataKeys) {
+ newNodes[idx].excludedLlmMetadataKeys.concat(excludedLlmMetadataKeys);
+ }
+ if (!this.disableTemplateRewrite) {
+ if (newNodes[idx] instanceof TextNode) {
+ newNodes[idx] = new TextNode({
+ ...newNodes[idx],
+ textTemplate: defaultNodeTextTemplate(),
+ });
+ }
+ }
+ }
+
+ return newNodes;
+ }
+}
diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts
index 3b0dfef57fc33cce1cb6ae330005895ac17b9971..bce66debf32d591c9a2162d3b500e191ed4fc1fc 100644
--- a/packages/core/src/index.ts
+++ b/packages/core/src/index.ts
@@ -15,6 +15,7 @@ export * from "./callbacks/CallbackManager";
export * from "./constants";
export * from "./embeddings";
export * from "./engines/chat";
+export * from "./extractors";
export * from "./indices";
export * from "./llm";
export * from "./nodeParsers";
diff --git a/packages/core/src/tests/MetadataExtractors.test.ts b/packages/core/src/tests/MetadataExtractors.test.ts
new file mode 100644
index 0000000000000000000000000000000000000000..5d8ef8a35dd88fd0394362bfdd68df7a87eaf7c9
--- /dev/null
+++ b/packages/core/src/tests/MetadataExtractors.test.ts
@@ -0,0 +1,138 @@
+import { Document } from "../Node";
+import { ServiceContext, serviceContextFromDefaults } from "../ServiceContext";
+import {
+ CallbackManager,
+ RetrievalCallbackResponse,
+ StreamCallbackResponse,
+} from "../callbacks/CallbackManager";
+import { OpenAIEmbedding } from "../embeddings";
+import {
+ KeywordExtractor,
+ QuestionsAnsweredExtractor,
+ SummaryExtractor,
+ TitleExtractor,
+} from "../extractors";
+import { OpenAI } from "../llm/LLM";
+import { SimpleNodeParser } from "../nodeParsers";
+import {
+ DEFAULT_LLM_TEXT_OUTPUT,
+ mockEmbeddingModel,
+ mockLlmGeneration,
+} from "./utility/mockOpenAI";
+
+// Mock the OpenAI getOpenAISession function during testing
+jest.mock("../llm/openai", () => {
+ return {
+ getOpenAISession: jest.fn().mockImplementation(() => null),
+ };
+});
+
+describe("[MetadataExtractor]: Extractors should populate the metadata", () => {
+ let serviceContext: ServiceContext;
+ let streamCallbackData: StreamCallbackResponse[] = [];
+ let retrieveCallbackData: RetrievalCallbackResponse[] = [];
+
+ beforeAll(async () => {
+ const callbackManager = new CallbackManager({
+ onLLMStream: (data) => {
+ streamCallbackData.push(data);
+ },
+ onRetrieve: (data) => {
+ retrieveCallbackData.push(data);
+ },
+ });
+
+ const languageModel = new OpenAI({
+ model: "gpt-3.5-turbo",
+ callbackManager,
+ });
+
+ mockLlmGeneration({ languageModel, callbackManager });
+
+ const embedModel = new OpenAIEmbedding();
+ mockEmbeddingModel(embedModel);
+
+ serviceContext = serviceContextFromDefaults({
+ callbackManager,
+ llm: languageModel,
+ embedModel,
+ });
+ });
+
+ beforeEach(() => {
+ streamCallbackData = [];
+ retrieveCallbackData = [];
+ });
+
+ afterAll(() => {
+ jest.clearAllMocks();
+ });
+
+ test("[MetadataExtractor] KeywordExtractor returns excerptKeywords metadata", async () => {
+ const nodeParser = new SimpleNodeParser();
+
+ const nodes = nodeParser.getNodesFromDocuments([
+ new Document({ text: DEFAULT_LLM_TEXT_OUTPUT }),
+ ]);
+
+ const keywordExtractor = new KeywordExtractor(serviceContext.llm, 5);
+
+ const nodesWithKeywordMetadata = await keywordExtractor.processNodes(nodes);
+
+ expect(nodesWithKeywordMetadata[0].metadata).toMatchObject({
+ excerptKeywords: DEFAULT_LLM_TEXT_OUTPUT,
+ });
+ });
+
+ test("[MetadataExtractor] TitleExtractor returns documentTitle metadata", async () => {
+ const nodeParser = new SimpleNodeParser();
+
+ const nodes = nodeParser.getNodesFromDocuments([
+ new Document({ text: DEFAULT_LLM_TEXT_OUTPUT }),
+ ]);
+
+ const titleExtractor = new TitleExtractor(serviceContext.llm, 5);
+
+ const nodesWithKeywordMetadata = await titleExtractor.processNodes(nodes);
+
+ expect(nodesWithKeywordMetadata[0].metadata).toMatchObject({
+ documentTitle: DEFAULT_LLM_TEXT_OUTPUT,
+ });
+ });
+
+ test("[MetadataExtractor] QuestionsAnsweredExtractor returns questionsThisExcerptCanAnswer metadata", async () => {
+ const nodeParser = new SimpleNodeParser();
+
+ const nodes = nodeParser.getNodesFromDocuments([
+ new Document({ text: DEFAULT_LLM_TEXT_OUTPUT }),
+ ]);
+
+ const questionsAnsweredExtractor = new QuestionsAnsweredExtractor(
+ serviceContext.llm,
+ 5,
+ );
+
+ const nodesWithKeywordMetadata =
+ await questionsAnsweredExtractor.processNodes(nodes);
+
+ expect(nodesWithKeywordMetadata[0].metadata).toMatchObject({
+ questionsThisExcerptCanAnswer: DEFAULT_LLM_TEXT_OUTPUT,
+ });
+ });
+
+ test("[MetadataExtractor] SumamryExtractor returns sectionSummary metadata", async () => {
+ const nodeParser = new SimpleNodeParser();
+
+ const nodes = nodeParser.getNodesFromDocuments([
+ new Document({ text: DEFAULT_LLM_TEXT_OUTPUT }),
+ ]);
+
+ const summaryExtractor = new SummaryExtractor(serviceContext.llm);
+
+ const nodesWithKeywordMetadata = await summaryExtractor.processNodes(nodes);
+
+ expect(nodesWithKeywordMetadata[0].metadata).toMatchObject({
+ sectionSummary: DEFAULT_LLM_TEXT_OUTPUT,
+ });
+ });
+});
diff --git a/packages/core/src/tests/utility/mockOpenAI.ts b/packages/core/src/tests/utility/mockOpenAI.ts
index 27e20e3bf2d5927c31eaa10b51b4fd9b468f447d..84f6925d746f6256c9300f203ca64297f6853abe 100644
--- a/packages/core/src/tests/utility/mockOpenAI.ts
+++ b/packages/core/src/tests/utility/mockOpenAI.ts
@@ -4,6 +4,8 @@ import { globalsHelper } from "../../GlobalsHelper";
import { OpenAI } from "../../llm/LLM";
import { LLMChatParamsBase } from "../../llm/types";
+export const DEFAULT_LLM_TEXT_OUTPUT = "MOCK_TOKEN_1-MOCK_TOKEN_2";
+
export function mockLlmGeneration({
languageModel,
callbackManager,
@@ -15,7 +17,7 @@ export function mockLlmGeneration({
.spyOn(languageModel, "chat")
.mockImplementation(
async ({ messages, parentEvent }: LLMChatParamsBase) => {
- const text = "MOCK_TOKEN_1-MOCK_TOKEN_2";
+ const text = DEFAULT_LLM_TEXT_OUTPUT;
const event = globalsHelper.createEvent({
parentEvent,
type: "llmPredict",
diff --git a/packages/core/tsconfig.json b/packages/core/tsconfig.json
index ce09c533fa989142d0b1ac6c330f9cea35f6c36c..074f60e7b98f450b66cf89d6caaaed70ac170474 100644
--- a/packages/core/tsconfig.json
+++ b/packages/core/tsconfig.json
@@ -14,6 +14,12 @@
"target": "ES2015",
"resolveJsonModule": true,
},
+ "ts-node": {
+ "files": true,
+ "compilerOptions": {
+ "module": "commonjs",
+ },
+ },
"include": ["./src"],
"exclude": ["node_modules"],
}