Skip to content
Snippets Groups Projects
Unverified Commit 33b56293 authored by Alex Yang's avatar Alex Yang Committed by GitHub
Browse files

refactor: move `data-structs` module (#1343)

parent 723b41c2
No related branches found
No related tags found
No related merge requests found
{
"type": "module",
"main": "./dist/index.cjs",
"module": "./dist/index.js",
"types": "./dist/index.d.ts",
"exports": "./dist/index.js",
"private": true
}
......@@ -283,6 +283,20 @@
"types": "./tools/dist/index.d.ts",
"default": "./tools/dist/index.js"
}
},
"./data-structs": {
"require": {
"types": "./data-structs/dist/index.d.cts",
"default": "./data-structs/dist/index.cjs"
},
"import": {
"types": "./data-structs/dist/index.d.ts",
"default": "./data-structs/dist/index.js"
},
"default": {
"types": "./data-structs/dist/index.d.ts",
"default": "./data-structs/dist/index.js"
}
}
},
"files": [
......@@ -305,7 +319,8 @@
"./chat-engine",
"./retriever",
"./vector-store",
"./tools"
"./tools",
"./data-structs"
],
"scripts": {
"dev": "bunchee --watch",
......
import { randomUUID } from "@llamaindex/env";
import type { UUID } from "../global";
import { IndexStructType } from "./struct-type";
export abstract class IndexStruct {
indexId: string;
summary: string | undefined;
constructor(
indexId: UUID = randomUUID(),
summary: string | undefined = undefined,
) {
this.indexId = indexId;
this.summary = summary;
}
toJson(): Record<string, unknown> {
return {
indexId: this.indexId,
summary: this.summary,
};
}
getSummary(): string {
if (this.summary === undefined) {
throw new Error("summary field of the index struct is not set");
}
return this.summary;
}
}
// A table of keywords mapping keywords to text chunks.
export class KeywordTable extends IndexStruct {
table: Map<string, Set<string>> = new Map();
type: IndexStructType = IndexStructType.KEYWORD_TABLE;
addNode(keywords: string[], nodeId: string): void {
keywords.forEach((keyword) => {
if (!this.table.has(keyword)) {
this.table.set(keyword, new Set());
}
this.table.get(keyword)!.add(nodeId);
});
}
deleteNode(keywords: string[], nodeId: string) {
keywords.forEach((keyword) => {
if (this.table.has(keyword)) {
this.table.get(keyword)!.delete(nodeId);
}
});
}
toJson(): Record<string, unknown> {
return {
...super.toJson(),
table: Array.from(this.table.entries()).reduce(
(acc, [keyword, nodeIds]) => {
acc[keyword] = Array.from(nodeIds);
return acc;
},
{} as Record<string, string[]>,
),
type: this.type,
};
}
}
export { IndexStruct, KeywordTable } from "./data-structs";
export { IndexStructType } from "./struct-type";
export const IndexStructType = {
NODE: "node",
TREE: "tree",
LIST: "list",
KEYWORD_TABLE: "keyword_table",
DICT: "dict",
SIMPLE_DICT: "simple_dict",
WEAVIATE: "weaviate",
PINECONE: "pinecone",
QDRANT: "qdrant",
LANCEDB: "lancedb",
MILVUS: "milvus",
CHROMA: "chroma",
MYSCALE: "myscale",
CLICKHOUSE: "clickhouse",
VECTOR_STORE: "vector_store",
OPENSEARCH: "opensearch",
DASHVECTOR: "dashvector",
CHATGPT_RETRIEVAL_PLUGIN: "chatgpt_retrieval_plugin",
DEEPLAKE: "deeplake",
EPSILLA: "epsilla",
MULTIMODAL_VECTOR_STORE: "multimodal",
SQL: "sql",
KG: "kg",
SIMPLE_KG: "simple_kg",
SIMPLE_LPG: "simple_lpg",
NEBULAGRAPH: "nebulagraph",
FALKORDB: "falkordb",
EMPTY: "empty",
COMPOSITE: "composite",
PANDAS: "pandas",
DOCUMENT_SUMMARY: "document_summary",
VECTARA: "vectara",
ZILLIZ_CLOUD_PIPELINE: "zilliz_cloud_pipeline",
POSTGRESML: "postgresml",
} as const;
export type IndexStructType =
(typeof IndexStructType)[keyof typeof IndexStructType];
......@@ -8,39 +8,6 @@ import { runTransformations } from "../ingestion/IngestionPipeline.js";
import type { StorageContext } from "../storage/StorageContext.js";
import type { BaseDocumentStore } from "../storage/docStore/types.js";
import type { BaseIndexStore } from "../storage/indexStore/types.js";
import { IndexStruct } from "./IndexStruct.js";
import { IndexStructType } from "./json-to-index-struct.js";
// A table of keywords mapping keywords to text chunks.
export class KeywordTable extends IndexStruct {
table: Map<string, Set<string>> = new Map();
type: IndexStructType = IndexStructType.KEYWORD_TABLE;
addNode(keywords: string[], nodeId: string): void {
keywords.forEach((keyword) => {
if (!this.table.has(keyword)) {
this.table.set(keyword, new Set());
}
this.table.get(keyword)!.add(nodeId);
});
}
deleteNode(keywords: string[], nodeId: string) {
keywords.forEach((keyword) => {
if (this.table.has(keyword)) {
this.table.get(keyword)!.delete(nodeId);
}
});
}
toJson(): Record<string, unknown> {
return {
...super.toJson(),
table: this.table,
type: this.type,
};
}
}
export interface BaseIndexInit<T> {
serviceContext?: ServiceContext | undefined;
......
......@@ -13,7 +13,7 @@ import type { StorageContext } from "../../storage/StorageContext.js";
import { storageContextFromDefaults } from "../../storage/StorageContext.js";
import type { BaseDocumentStore } from "../../storage/docStore/types.js";
import type { BaseIndexInit } from "../BaseIndex.js";
import { BaseIndex, KeywordTable } from "../BaseIndex.js";
import { BaseIndex } from "../BaseIndex.js";
import { IndexStructType } from "../json-to-index-struct.js";
import {
extractKeywordsGivenResponse,
......@@ -21,6 +21,7 @@ import {
simpleExtractKeywords,
} from "./utils.js";
import { KeywordTable } from "@llamaindex/core/data-structs";
import type { LLM } from "@llamaindex/core/llms";
import {
defaultKeywordExtractPrompt,
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment