diff --git a/.gitignore b/.gitignore index 6b28a0c0829ddda2b0ca5515b4bd6e2081688756..2641a07fa5489a4a28b76c5e83474c459749f8c9 100644 --- a/.gitignore +++ b/.gitignore @@ -35,4 +35,4 @@ yarn-error.log* # vercel .vercel -storage/ \ No newline at end of file +storage/ diff --git a/apps/simple/listIndex.ts b/apps/simple/listIndex.ts new file mode 100644 index 0000000000000000000000000000000000000000..5e03cfa019c06c1083f241ff0a6df1180bfa3c6c --- /dev/null +++ b/apps/simple/listIndex.ts @@ -0,0 +1,15 @@ +import { Document } from "@llamaindex/core/src/Node"; +import { ListIndex } from "@llamaindex/core/src/index/list"; +import essay from "./essay"; + +async function main() { + const document = new Document({ text: essay }); + const index = await ListIndex.fromDocuments([document]); + const queryEngine = index.asQueryEngine(); + const response = await queryEngine.aquery( + "What did the author do growing up?" + ); + console.log(response.toString()); +} + +main().catch(console.error); diff --git a/apps/simple/index.ts b/apps/simple/vectorIndex.ts similarity index 88% rename from apps/simple/index.ts rename to apps/simple/vectorIndex.ts index 733bb7f07fedcf7c8b74309fe384b214f7cec6cc..d05b5874984e3846ee2de6a55ceb1c29875a827a 100644 --- a/apps/simple/index.ts +++ b/apps/simple/vectorIndex.ts @@ -2,7 +2,7 @@ import { Document } from "@llamaindex/core/src/Node"; import { VectorStoreIndex } from "@llamaindex/core/src/BaseIndex"; import essay from "./essay"; -(async () => { +async function main() { const document = new Document({ text: essay }); const index = await VectorStoreIndex.fromDocuments([document]); const queryEngine = index.asQueryEngine(); @@ -10,4 +10,6 @@ import essay from "./essay"; "What did the author do growing up?" ); console.log(response.toString()); -})(); +} + +main().catch(console.error); diff --git a/packages/core/src/BaseIndex.ts b/packages/core/src/BaseIndex.ts index 65452fbd94152564aa8e35bb79d2c10081b29ad0..d8512f1b63a0a26a6b5bcab3d8260d59ed3f9a8e 100644 --- a/packages/core/src/BaseIndex.ts +++ b/packages/core/src/BaseIndex.ts @@ -9,6 +9,7 @@ import { } from "./storage/StorageContext"; import { BaseDocumentStore } from "./storage/docStore/types"; import { VectorStore } from "./storage/vectorStore/types"; +import { BaseIndexStore } from "./storage/indexStore/types"; export abstract class IndexStruct { indexId: string; @@ -56,14 +57,16 @@ export interface BaseIndexInit<T> { serviceContext: ServiceContext; storageContext: StorageContext; docStore: BaseDocumentStore; - vectorStore: VectorStore; + vectorStore?: VectorStore; + indexStore?: BaseIndexStore; indexStruct: T; } export abstract class BaseIndex<T> { serviceContext: ServiceContext; storageContext: StorageContext; docStore: BaseDocumentStore; - vectorStore: VectorStore; + vectorStore?: VectorStore; + indexStore?: BaseIndexStore; indexStruct: T; constructor(init: BaseIndexInit<T>) { @@ -71,6 +74,7 @@ export abstract class BaseIndex<T> { this.storageContext = init.storageContext; this.docStore = init.docStore; this.vectorStore = init.vectorStore; + this.indexStore = init.indexStore; this.indexStruct = init.indexStruct; } diff --git a/packages/core/src/index/list/ListIndex.ts b/packages/core/src/index/list/ListIndex.ts index 44496db2bca2ca2f73e5c543d1f3c73a19b0dd94..78c1971b9b750a01e101bc27f50980c72d8f1d16 100644 --- a/packages/core/src/index/list/ListIndex.ts +++ b/packages/core/src/index/list/ListIndex.ts @@ -1,8 +1,16 @@ -import { BaseNode } from "../../Node"; +import { BaseNode, Document } from "../../Node"; import { BaseIndex, BaseIndexInit, IndexList } from "../../BaseIndex"; +import { BaseQueryEngine, RetrieverQueryEngine } from "../../QueryEngine"; +import { + StorageContext, + storageContextFromDefaults, +} from "../../storage/StorageContext"; import { BaseRetriever } from "../../Retriever"; import { ListIndexRetriever } from "./ListIndexRetriever"; -import { ServiceContext } from "../../ServiceContext"; +import { + ServiceContext, + serviceContextFromDefaults, +} from "../../ServiceContext"; import { RefDocInfo } from "../../storage/docStore/types"; import _ from "lodash"; @@ -12,17 +20,73 @@ export enum ListRetrieverMode { LLM = "llm", } -export interface ListIndexInit extends BaseIndexInit<IndexList> { +export interface ListIndexOptions { nodes?: BaseNode[]; - indexStruct: IndexList; - serviceContext: ServiceContext; + indexStruct?: IndexList; + serviceContext?: ServiceContext; + storageContext?: StorageContext; } export class ListIndex extends BaseIndex<IndexList> { - constructor(init: ListIndexInit) { + constructor(init: BaseIndexInit<IndexList>) { super(init); } + static async init(options: ListIndexOptions): Promise<ListIndex> { + const storageContext = + options.storageContext ?? (await storageContextFromDefaults({})); + const serviceContext = + options.serviceContext ?? serviceContextFromDefaults({}); + const { docStore, indexStore } = storageContext; + + let indexStruct: IndexList; + if (options.indexStruct) { + if (options.nodes) { + throw new Error( + "Cannot initialize VectorStoreIndex with both nodes and indexStruct" + ); + } + indexStruct = options.indexStruct; + } else { + if (!options.nodes) { + throw new Error( + "Cannot initialize VectorStoreIndex without nodes or indexStruct" + ); + } + indexStruct = ListIndex._buildIndexFromNodes(options.nodes); + } + + return new ListIndex({ + storageContext, + serviceContext, + docStore, + indexStore, + indexStruct, + }); + } + + static async fromDocuments( + documents: Document[], + storageContext?: StorageContext, + serviceContext?: ServiceContext + ): Promise<ListIndex> { + storageContext = storageContext ?? (await storageContextFromDefaults({})); + serviceContext = serviceContext ?? serviceContextFromDefaults({}); + const docStore = storageContext.docStore; + + for (const doc of documents) { + docStore.setDocumentHash(doc.id_, doc.hash); + } + + const nodes = serviceContext.nodeParser.getNodesFromDocuments(documents); + const index = await ListIndex.init({ + nodes, + storageContext, + serviceContext, + }); + return index; + } + asRetriever( mode: ListRetrieverMode = ListRetrieverMode.DEFAULT ): BaseRetriever { @@ -36,8 +100,17 @@ export class ListIndex extends BaseIndex<IndexList> { } } - protected _buildIndexFromNodes(nodes: BaseNode[]): IndexList { - const indexStruct = new IndexList(); + asQueryEngine( + mode: ListRetrieverMode = ListRetrieverMode.DEFAULT + ): BaseQueryEngine { + return new RetrieverQueryEngine(this.asRetriever()); + } + + static _buildIndexFromNodes( + nodes: BaseNode[], + indexStruct?: IndexList + ): IndexList { + indexStruct = indexStruct || new IndexList(); for (const node of nodes) { indexStruct.addNode(node); diff --git a/packages/core/src/index/list/index.ts b/packages/core/src/index/list/index.ts new file mode 100644 index 0000000000000000000000000000000000000000..f8d0b8d5eae44cf561fd8483fbabf6dc716260d0 --- /dev/null +++ b/packages/core/src/index/list/index.ts @@ -0,0 +1,5 @@ +export { ListIndex, ListRetrieverMode } from "./ListIndex"; +export { + ListIndexRetriever, + ListIndexLLMRetriever, +} from "./ListIndexRetriever";