diff --git a/apps/simple/listIndex.ts b/apps/simple/listIndex.ts index db8e30fc8aee45fdcdffa4776c29beeaf5e34321..0bb1c65c71697b029f4de667f0aea6f07bd015c3 100644 --- a/apps/simple/listIndex.ts +++ b/apps/simple/listIndex.ts @@ -1,9 +1,23 @@ -import { Document, ListIndex, ListRetrieverMode } from "llamaindex"; +import { + Document, + ListIndex, + ListRetrieverMode, + serviceContextFromDefaults, + SimpleNodeParser, +} from "llamaindex"; import essay from "./essay"; async function main() { + const serviceContext = serviceContextFromDefaults({ + nodeParser: new SimpleNodeParser({ + chunkSize: 40, + }), + }); const document = new Document({ text: essay }); - const index = await ListIndex.fromDocuments([document]); + const index = await ListIndex.fromDocuments({ + documents: [document], + serviceContext, + }); const queryEngine = index.asQueryEngine(ListRetrieverMode.LLM); const response = await queryEngine.aquery( "What did the author do growing up?" diff --git a/packages/core/src/QueryEngine.ts b/packages/core/src/QueryEngine.ts index 8ea7378b35d9c6be6c18df4c88a9becc2bbc0c5d..e574c6cf5e110b8e8ac3228d890da05cee050e0b 100644 --- a/packages/core/src/QueryEngine.ts +++ b/packages/core/src/QueryEngine.ts @@ -26,11 +26,15 @@ export class RetrieverQueryEngine implements BaseQueryEngine { retriever: BaseRetriever; responseSynthesizer: ResponseSynthesizer; - constructor(retriever: BaseRetriever) { + constructor( + retriever: BaseRetriever, + responseSynthesizer?: ResponseSynthesizer + ) { this.retriever = retriever; const serviceContext: ServiceContext | undefined = this.retriever.getServiceContext(); - this.responseSynthesizer = new ResponseSynthesizer({ serviceContext }); + this.responseSynthesizer = + responseSynthesizer || new ResponseSynthesizer({ serviceContext }); } async aquery(query: string, parentEvent?: Event) { diff --git a/packages/core/src/index/list/ListIndex.ts b/packages/core/src/index/list/ListIndex.ts index 7b2aff279d0bd8a1b76c2e0c275f09a905534412..54bd62e62c2c51072fec069f527c1bc993dffc24 100644 --- a/packages/core/src/index/list/ListIndex.ts +++ b/packages/core/src/index/list/ListIndex.ts @@ -16,6 +16,10 @@ import { } from "../../ServiceContext"; import { BaseDocumentStore, RefDocInfo } from "../../storage/docStore/types"; import _ from "lodash"; +import { + ResponseSynthesizer, + CompactAndRefine, +} from "../../ResponseSynthesizer"; export enum ListRetrieverMode { DEFAULT = "default", @@ -74,11 +78,12 @@ export class ListIndex extends BaseIndex<IndexList> { }); } - static async fromDocuments( - documents: Document[], - storageContext?: StorageContext, - serviceContext?: ServiceContext - ): Promise<ListIndex> { + static async fromDocuments(args: { + documents: Document[]; + storageContext?: StorageContext; + serviceContext?: ServiceContext; + }): Promise<ListIndex> { + let { documents, storageContext, serviceContext } = args; storageContext = storageContext ?? (await storageContextFromDefaults({})); serviceContext = serviceContext ?? serviceContextFromDefaults({}); const docStore = storageContext.docStore; @@ -111,9 +116,20 @@ export class ListIndex extends BaseIndex<IndexList> { } asQueryEngine( - mode: ListRetrieverMode = ListRetrieverMode.DEFAULT + mode: ListRetrieverMode = ListRetrieverMode.DEFAULT, + responseSynthesizer?: ResponseSynthesizer ): BaseQueryEngine { - return new RetrieverQueryEngine(this.asRetriever(mode)); + if (_.isNil(responseSynthesizer)) { + let responseBuilder = new CompactAndRefine(this.serviceContext); + responseSynthesizer = new ResponseSynthesizer({ + serviceContext: this.serviceContext, + responseBuilder, + }); + } + return new RetrieverQueryEngine( + this.asRetriever(mode), + responseSynthesizer + ); } static async _buildIndexFromNodes( diff --git a/packages/core/src/tests/CallbackManager.test.ts b/packages/core/src/tests/CallbackManager.test.ts index e66d3e6501c386899a586cff2bca1289747fe169..ea790ea1180d93cdcb51c25ff30771290000bf24 100644 --- a/packages/core/src/tests/CallbackManager.test.ts +++ b/packages/core/src/tests/CallbackManager.test.ts @@ -8,7 +8,11 @@ import { RetrievalCallbackResponse, StreamCallbackResponse, } from "../callbacks/CallbackManager"; -import { ListIndex } from "../index/list"; +import { ListIndex, ListRetrieverMode } from "../index/list"; +import { + ResponseSynthesizer, + SimpleResponseBuilder, +} from "../ResponseSynthesizer"; import { mockEmbeddingModel, mockLlmGeneration } from "./utility/mockOpenAI"; // Mock the OpenAI getOpenAISession function during testing @@ -134,12 +138,19 @@ describe("CallbackManager: onLLMStream and onRetrieve", () => { }); test("For ListIndex w/ a ListIndexRetriever", async () => { - const listIndex = await ListIndex.fromDocuments( - [document], - undefined, - serviceContext + const listIndex = await ListIndex.fromDocuments({ + documents: [document], + serviceContext, + }); + const responseBuilder = new SimpleResponseBuilder(serviceContext); + const responseSynthesizer = new ResponseSynthesizer({ + serviceContext: serviceContext, + responseBuilder, + }); + const queryEngine = listIndex.asQueryEngine( + ListRetrieverMode.DEFAULT, + responseSynthesizer ); - const queryEngine = listIndex.asQueryEngine(); const query = "What is the author's name?"; const response = await queryEngine.aquery(query); expect(response.toString()).toBe("MOCK_TOKEN_1-MOCK_TOKEN_2");