diff --git a/packages/core/src/ChatEngine.ts b/packages/core/src/ChatEngine.ts index 5311ffa8a639288443b9097a92d781f4ed34f2c4..1afbb13d2908cd7e5d6b1820e3d67622f7e02d1c 100644 --- a/packages/core/src/ChatEngine.ts +++ b/packages/core/src/ChatEngine.ts @@ -13,8 +13,8 @@ import { Response } from "./Response"; import { BaseRetriever } from "./Retriever"; import { ServiceContext, serviceContextFromDefaults } from "./ServiceContext"; import { Event } from "./callbacks/CallbackManager"; -import { BaseNodePostprocessor } from "./indices/BaseNodePostprocessor"; -import { ChatMessage, LLM, OpenAI } from "./llm/LLM"; +import { ChatMessage, LLM, OpenAI } from "./llm"; +import { BaseNodePostprocessor } from "./postprocessors"; /** * A ChatEngine is used to handle back and forth chats between the application and the LLM. diff --git a/packages/core/src/QueryEngine.ts b/packages/core/src/QueryEngine.ts index 8e032fbf849f859f998348bb987055d83e266c1b..aedbb9be79e581ef6930ac8a4fa80aa3b1e9b18f 100644 --- a/packages/core/src/QueryEngine.ts +++ b/packages/core/src/QueryEngine.ts @@ -10,9 +10,12 @@ import { BaseRetriever } from "./Retriever"; import { ServiceContext, serviceContextFromDefaults } from "./ServiceContext"; import { QueryEngineTool, ToolMetadata } from "./Tool"; import { Event } from "./callbacks/CallbackManager"; -import { BaseNodePostprocessor } from "./indices/BaseNodePostprocessor"; -import { CompactAndRefine, ResponseSynthesizer } from "./synthesizers"; -import { BaseSynthesizer } from "./synthesizers/types"; +import { BaseNodePostprocessor } from "./postprocessors"; +import { + BaseSynthesizer, + CompactAndRefine, + ResponseSynthesizer, +} from "./synthesizers"; /** * A query engine is a question answerer that can use one or more steps. diff --git a/packages/core/src/indices/index.ts b/packages/core/src/indices/index.ts index ddfe185dc114c56aabf420abdd1431cb5f9f9ff5..8bda05b2d954cefd6a58893edc0efb095971a102 100644 --- a/packages/core/src/indices/index.ts +++ b/packages/core/src/indices/index.ts @@ -1,5 +1,4 @@ export * from "./BaseIndex"; -export * from "./BaseNodePostprocessor"; export * from "./keyword"; export * from "./summary"; export * from "./vectorStore"; diff --git a/packages/core/src/indices/keyword/KeywordTableIndex.ts b/packages/core/src/indices/keyword/KeywordTableIndex.ts index 406809d88f96c1f8a88b6755c574a2488567d603..5b67cec674933f9e08457ff9dbf3f009e7b3a623 100644 --- a/packages/core/src/indices/keyword/KeywordTableIndex.ts +++ b/packages/core/src/indices/keyword/KeywordTableIndex.ts @@ -6,8 +6,12 @@ import { ServiceContext, serviceContextFromDefaults, } from "../../ServiceContext"; -import { StorageContext, storageContextFromDefaults } from "../../storage"; -import { BaseDocumentStore } from "../../storage/docStore/types"; +import { BaseNodePostprocessor } from "../../postprocessors"; +import { + BaseDocumentStore, + StorageContext, + storageContextFromDefaults, +} from "../../storage"; import { BaseSynthesizer } from "../../synthesizers"; import { BaseIndex, @@ -15,7 +19,6 @@ import { IndexStructType, KeywordTable, } from "../BaseIndex"; -import { BaseNodePostprocessor } from "../BaseNodePostprocessor"; import { KeywordTableLLMRetriever, KeywordTableRAKERetriever, diff --git a/packages/core/src/indices/summary/SummaryIndex.ts b/packages/core/src/indices/summary/SummaryIndex.ts index eb6f753330e9bef150d0004058a6ae3bd53fae6a..391685dd895c494f362849a762fc50f77bd08216 100644 --- a/packages/core/src/indices/summary/SummaryIndex.ts +++ b/packages/core/src/indices/summary/SummaryIndex.ts @@ -6,11 +6,13 @@ import { ServiceContext, serviceContextFromDefaults, } from "../../ServiceContext"; +import { BaseNodePostprocessor } from "../../postprocessors"; import { + BaseDocumentStore, + RefDocInfo, StorageContext, storageContextFromDefaults, -} from "../../storage/StorageContext"; -import { BaseDocumentStore, RefDocInfo } from "../../storage/docStore/types"; +} from "../../storage"; import { BaseSynthesizer, CompactAndRefine, @@ -22,7 +24,6 @@ import { IndexList, IndexStructType, } from "../BaseIndex"; -import { BaseNodePostprocessor } from "../BaseNodePostprocessor"; import { SummaryIndexLLMRetriever, SummaryIndexRetriever, diff --git a/packages/core/src/indices/vectorStore/VectorStoreIndex.ts b/packages/core/src/indices/vectorStore/VectorStoreIndex.ts index 278209f2cf16f9e1c9d5a2c75721e6a67fb5ca97..59d177b34f1ed4b077ef3299e1d7599a43345af2 100644 --- a/packages/core/src/indices/vectorStore/VectorStoreIndex.ts +++ b/packages/core/src/indices/vectorStore/VectorStoreIndex.ts @@ -17,12 +17,13 @@ import { ClipEmbedding, MultiModalEmbedding, } from "../../embeddings"; +import { BaseNodePostprocessor } from "../../postprocessors"; import { + BaseIndexStore, StorageContext, + VectorStore, storageContextFromDefaults, -} from "../../storage/StorageContext"; -import { BaseIndexStore } from "../../storage/indexStore/types"; -import { VectorStore } from "../../storage/vectorStore/types"; +} from "../../storage"; import { BaseSynthesizer } from "../../synthesizers"; import { BaseIndex, @@ -30,7 +31,6 @@ import { IndexDict, IndexStructType, } from "../BaseIndex"; -import { BaseNodePostprocessor } from "../BaseNodePostprocessor"; import { VectorIndexRetriever } from "./VectorIndexRetriever"; interface IndexStructOptions { diff --git a/packages/core/src/postprocessors/MetadataReplacementPostProcessor.ts b/packages/core/src/postprocessors/MetadataReplacementPostProcessor.ts new file mode 100644 index 0000000000000000000000000000000000000000..d05d94a108f66f5a4ab388a45614e99e9ed5f523 --- /dev/null +++ b/packages/core/src/postprocessors/MetadataReplacementPostProcessor.ts @@ -0,0 +1,21 @@ +import { MetadataMode, NodeWithScore } from "../Node"; +import { BaseNodePostprocessor } from "./types"; + +export class MetadataReplacementPostProcessor implements BaseNodePostprocessor { + targetMetadataKey: string; + + constructor(targetMetadataKey: string) { + this.targetMetadataKey = targetMetadataKey; + } + + postprocessNodes(nodes: NodeWithScore[]): NodeWithScore[] { + for (let n of nodes) { + n.node.setContent( + n.node.metadata[this.targetMetadataKey] ?? + n.node.getContent(MetadataMode.NONE), + ); + } + + return nodes; + } +} diff --git a/packages/core/src/indices/BaseNodePostprocessor.ts b/packages/core/src/postprocessors/SimilarityPostprocessor.ts similarity index 81% rename from packages/core/src/indices/BaseNodePostprocessor.ts rename to packages/core/src/postprocessors/SimilarityPostprocessor.ts index b4707289425ba7abfa53ad33cf8e29cf8baa0aea..91674515ed7cbe67eeb23ba3c579af51a57f3fa9 100644 --- a/packages/core/src/indices/BaseNodePostprocessor.ts +++ b/packages/core/src/postprocessors/SimilarityPostprocessor.ts @@ -1,8 +1,5 @@ import { NodeWithScore } from "../Node"; - -export interface BaseNodePostprocessor { - postprocessNodes: (nodes: NodeWithScore[]) => NodeWithScore[]; -} +import { BaseNodePostprocessor } from "./types"; export class SimilarityPostprocessor implements BaseNodePostprocessor { similarityCutoff?: number; diff --git a/packages/core/src/postprocessors/index.ts b/packages/core/src/postprocessors/index.ts new file mode 100644 index 0000000000000000000000000000000000000000..f79e4ced0526bd211536df82353e40f00eaa10a9 --- /dev/null +++ b/packages/core/src/postprocessors/index.ts @@ -0,0 +1,3 @@ +export * from "./MetadataReplacementPostProcessor"; +export * from "./SimilarityPostprocessor"; +export * from "./types"; diff --git a/packages/core/src/postprocessors/types.ts b/packages/core/src/postprocessors/types.ts new file mode 100644 index 0000000000000000000000000000000000000000..2d0c73e78431a939646f57d30038c8a835b915f9 --- /dev/null +++ b/packages/core/src/postprocessors/types.ts @@ -0,0 +1,5 @@ +import { NodeWithScore } from "../Node"; + +export interface BaseNodePostprocessor { + postprocessNodes: (nodes: NodeWithScore[]) => NodeWithScore[]; +} diff --git a/packages/core/src/tests/postprocessors/MetadataReplacementPostProcessor.test.ts b/packages/core/src/tests/postprocessors/MetadataReplacementPostProcessor.test.ts new file mode 100644 index 0000000000000000000000000000000000000000..f9a76845db7e0bbb1e8314df11c5f85b4bc1f16d --- /dev/null +++ b/packages/core/src/tests/postprocessors/MetadataReplacementPostProcessor.test.ts @@ -0,0 +1,33 @@ +import { MetadataMode, NodeWithScore, TextNode } from "../../Node"; +import { MetadataReplacementPostProcessor } from "../../postprocessors"; + +describe("MetadataReplacementPostProcessor", () => { + let postProcessor: MetadataReplacementPostProcessor; + let nodes: NodeWithScore[]; + + beforeEach(() => { + postProcessor = new MetadataReplacementPostProcessor("targetKey"); + + nodes = [ + { + node: new TextNode({ + text: "OldContent", + }), + score: 5, + }, + ]; + }); + + test("Replaces the content of each node with specified metadata key if it exists", () => { + nodes[0].node.metadata = { targetKey: "NewContent" }; + const newNodes = postProcessor.postprocessNodes(nodes); + // Check if node content was replaced correctly + expect(newNodes[0].node.getContent(MetadataMode.NONE)).toBe("NewContent"); + }); + + test("Retains the original content of each node if no metadata key is found", () => { + const newNodes = postProcessor.postprocessNodes(nodes); + // Check if node content remained unchanged + expect(newNodes[0].node.getContent(MetadataMode.NONE)).toBe("OldContent"); + }); +});