llamaindex
Version:
<p align="center"> <img height="100" width="100" alt="LlamaIndex logo" src="https://ts.llamaindex.ai/square.svg" /> </p> <h1 align="center">LlamaIndex.TS</h1> <h3 align="center"> Data framework for your LLM application. </h3>
2,642 lines • 96.5 kB
JavaScript
import { PromptHelper } from '@llamaindex/core/indices';
export * from '@llamaindex/core/indices';
import { MetadataMode, TransformComponent, splitNodesByType, ModalityType, ObjectType, ImageNode } from '@llamaindex/core/schema';
import { KVDocumentStore } from '@llamaindex/core/storage/doc-store';
import { SimpleKVStore, BaseInMemoryKVStore } from '@llamaindex/core/storage/kv-store';
import { createSHA256, AsyncLocalStorage, fs, path, consoleLogger } from '@llamaindex/env';
import { Settings as Settings$1, DEFAULT_PERSIST_DIR, DEFAULT_DOC_STORE_PERSIST_FILENAME, DEFAULT_NAMESPACE } from '@llamaindex/core/global';
import { SentenceSplitter } from '@llamaindex/core/node-parser';
import { RetrieverQueryEngine } from '@llamaindex/core/query-engine';
import { getResponseSynthesizer } from '@llamaindex/core/response-synthesizers';
import { messagesToHistory, extractText, streamReducer } from '@llamaindex/core/utils';
import '../../selectors/dist/index.js';
import { defaultCondenseQuestionPrompt, defaultKeywordExtractPrompt, defaultQueryKeywordExtractPrompt, defaultChoiceSelectPrompt } from '@llamaindex/core/prompts';
import { SimpleIndexStore } from '@llamaindex/core/storage/index-store';
import { getTopKMMREmbeddings, getTopKEmbeddings, DEFAULT_SIMILARITY_TOP_K } from '@llamaindex/core/embeddings';
import { VectorStoreQueryMode, FilterOperator, parsePrimitiveValue, parseArrayValue, BaseVectorStore, nodeToMetadata } from '@llamaindex/core/vector-store';
import _ from 'lodash';
import { IndexStructType, KeywordTable, IndexList, IndexDict } from '@llamaindex/core/data-structs';
import { BaseRetriever } from '@llamaindex/core/retriever';
import { BaseChatEngine, ContextChatEngine } from '@llamaindex/core/chat-engine';
import { wrapEventCaller } from '@llamaindex/core/decorator';
import { createMemory, Memory } from '@llamaindex/core/memory';
const transformToJSON = (obj)=>{
// eslint-disable-next-line @typescript-eslint/no-explicit-any
const seen = [];
// eslint-disable-next-line @typescript-eslint/no-explicit-any
const replacer = (key, value)=>{
if (value != null && typeof value == "object") {
if (seen.indexOf(value) >= 0) {
return;
}
seen.push(value);
}
return value;
};
// this is a custom replacer function that will allow us to handle circular references
const jsonStr = JSON.stringify(obj, replacer);
return jsonStr;
};
function getTransformationHash(nodes, transform) {
const nodesStr = nodes.map((node)=>node.getContent(MetadataMode.ALL)).join("");
const transformString = transformToJSON(transform);
const hash = createSHA256();
hash.update(nodesStr + transformString + transform.id);
return hash.digest();
}
async function classify(docStore, nodes) {
const existingDocIds = Object.values(await docStore.getAllDocumentHashes());
const docIdsFromNodes = new Set();
const dedupedNodes = [];
const unusedDocs = [];
for (const node of nodes){
const refDocId = node.sourceNode?.nodeId || node.id_;
docIdsFromNodes.add(refDocId);
const existingHash = await docStore.getDocumentHash(refDocId);
if (!existingHash) {
// document doesn't exist, so add it
dedupedNodes.push(node);
} else if (existingHash && existingHash !== node.hash) {
// document exists but hash is different, so mark doc as unused and add node as deduped
unusedDocs.push(refDocId);
dedupedNodes.push(node);
}
// otherwise, document exists and hash is the same, so do nothing
}
const missingDocs = existingDocIds.filter((id)=>!docIdsFromNodes.has(id));
return {
dedupedNodes,
missingDocs,
unusedDocs
};
}
class RollbackableTransformComponent extends TransformComponent {
// Remove unused docs from the doc store. It is useful in case
// generating embeddings fails and we want to remove the unused docs
// TODO: override this in UpsertsStrategy if we want to revert removed docs also
async rollback(docStore, nodes) {
const { unusedDocs } = await classify(docStore, nodes);
for (const docId of unusedDocs){
await docStore.deleteDocument(docId, false);
}
docStore.persist();
}
}
/**
* Handle doc store duplicates by checking all hashes.
*/ class DuplicatesStrategy extends RollbackableTransformComponent {
constructor(docStore){
super(async (nodes)=>{
const hashes = await this.docStore.getAllDocumentHashes();
const currentHashes = new Set();
const nodesToRun = [];
for (const node of nodes){
if (!(node.hash in hashes) && !currentHashes.has(node.hash)) {
await this.docStore.setDocumentHash(node.id_, node.hash);
nodesToRun.push(node);
currentHashes.add(node.hash);
}
}
await this.docStore.addDocuments(nodesToRun, true);
return nodesToRun;
});
this.docStore = docStore;
}
}
/**
* Handle docstore upserts by checking hashes and ids.
* Identify missing docs and delete them from docstore and vector store
*/ class UpsertsAndDeleteStrategy extends RollbackableTransformComponent {
constructor(docStore, vectorStores){
super(async (nodes)=>{
const { dedupedNodes, missingDocs, unusedDocs } = await classify(this.docStore, nodes);
// remove unused docs
for (const refDocId of unusedDocs){
await this.docStore.deleteRefDoc(refDocId, false);
if (this.vectorStores) {
for (const vectorStore of this.vectorStores){
await vectorStore.delete(refDocId);
}
}
}
// remove missing docs
for (const docId of missingDocs){
await this.docStore.deleteDocument(docId, true);
if (this.vectorStores) {
for (const vectorStore of this.vectorStores){
await vectorStore.delete(docId);
}
}
}
await this.docStore.addDocuments(dedupedNodes, true);
return dedupedNodes;
});
this.docStore = docStore;
this.vectorStores = vectorStores;
}
}
/**
* Handles doc store upserts by checking hashes and ids.
*/ class UpsertsStrategy extends RollbackableTransformComponent {
constructor(docStore, vectorStores){
super(async (nodes)=>{
const { dedupedNodes, unusedDocs } = await classify(this.docStore, nodes);
// remove unused docs
for (const refDocId of unusedDocs){
await this.docStore.deleteRefDoc(refDocId, false);
if (this.vectorStores) {
for (const vectorStore of this.vectorStores){
await vectorStore.delete(refDocId);
}
}
}
// add non-duplicate docs
await this.docStore.addDocuments(dedupedNodes, true);
return dedupedNodes;
});
this.docStore = docStore;
this.vectorStores = vectorStores;
}
}
/**
* Document de-deduplication strategies work by comparing the hashes or ids stored in the document store.
* They require a document store to be set which must be persisted across pipeline runs.
*/ var DocStoreStrategy = /*#__PURE__*/ function(DocStoreStrategy) {
// Use upserts to handle duplicates. Checks if the a document is already in the doc store based on its id. If it is not, or if the hash of the document is updated, it will update the document in the doc store and run the transformations.
DocStoreStrategy["UPSERTS"] = "upserts";
// Only handle duplicates. Checks if the hash of a document is already in the doc store. Only then it will add the document to the doc store and run the transformations
DocStoreStrategy["DUPLICATES_ONLY"] = "duplicates_only";
// Use upserts and delete to handle duplicates. Like the upsert strategy but it will also delete non-existing documents from the doc store
DocStoreStrategy["UPSERTS_AND_DELETE"] = "upserts_and_delete";
DocStoreStrategy["NONE"] = "none";
return DocStoreStrategy;
}({});
class NoOpStrategy extends RollbackableTransformComponent {
constructor(){
super(async (nodes)=>nodes);
}
}
function createDocStoreStrategy(docStoreStrategy, docStore, vectorStores = []) {
if (docStoreStrategy === "none") {
return new NoOpStrategy();
}
if (!docStore) {
throw new Error("docStore is required to create a doc store strategy.");
}
if (vectorStores.length > 0) {
if (docStoreStrategy === "upserts") {
return new UpsertsStrategy(docStore, vectorStores);
} else if (docStoreStrategy === "upserts_and_delete") {
return new UpsertsAndDeleteStrategy(docStore, vectorStores);
} else if (docStoreStrategy === "duplicates_only") {
return new DuplicatesStrategy(docStore);
} else {
throw new Error(`Invalid docstore strategy: ${docStoreStrategy}`);
}
} else {
if (docStoreStrategy === "upserts") {
console.warn("Docstore strategy set to upserts, but no vector store. Switching to duplicates_only strategy.");
} else if (docStoreStrategy === "upserts_and_delete") {
console.warn("Docstore strategy set to upserts and delete, but no vector store. Switching to duplicates_only strategy.");
}
return new DuplicatesStrategy(docStore);
}
}
async function runTransformations(nodesToRun, transformations, // eslint-disable-next-line @typescript-eslint/no-explicit-any
transformOptions = {}, { inPlace = true, cache, docStoreStrategy } = {}) {
let nodes = nodesToRun;
if (!inPlace) {
nodes = [
...nodesToRun
];
}
if (docStoreStrategy) {
nodes = await docStoreStrategy(nodes);
}
for (const transform of transformations){
if (cache) {
const hash = getTransformationHash(nodes, transform);
const cachedNodes = await cache.get(hash);
if (cachedNodes) {
nodes = cachedNodes;
} else {
nodes = await transform(nodes, transformOptions);
await cache.put(hash, nodes);
}
} else {
nodes = await transform(nodes, transformOptions);
}
}
return nodes;
}
async function addNodesToVectorStores(nodes, vectorStores, nodesAdded) {
const nodeMap = splitNodesByType(nodes);
for(const type in nodeMap){
const nodes = nodeMap[type];
if (nodes) {
const vectorStore = vectorStores[type];
if (!vectorStore) {
throw new Error(`Cannot insert nodes of type ${type} without assigned vector store`);
}
const newIds = await vectorStore.add(nodes);
if (nodesAdded) {
await nodesAdded(newIds, nodes, vectorStore);
}
}
}
}
/**
* @internal
*/ class GlobalSettings {
#prompt;
#promptHelper;
#nodeParser;
#chunkOverlap;
#promptHelperAsyncLocalStorage;
#nodeParserAsyncLocalStorage;
#chunkOverlapAsyncLocalStorage;
#promptAsyncLocalStorage;
get debug() {
return Settings$1.debug;
}
get llm() {
return Settings$1.llm;
}
set llm(llm) {
Settings$1.llm = llm;
}
withLLM(llm, fn) {
return Settings$1.withLLM(llm, fn);
}
get promptHelper() {
if (this.#promptHelper === null) {
this.#promptHelper = new PromptHelper();
}
return this.#promptHelperAsyncLocalStorage.getStore() ?? this.#promptHelper;
}
set promptHelper(promptHelper) {
this.#promptHelper = promptHelper;
}
withPromptHelper(promptHelper, fn) {
return this.#promptHelperAsyncLocalStorage.run(promptHelper, fn);
}
get embedModel() {
return Settings$1.embedModel;
}
set embedModel(embedModel) {
Settings$1.embedModel = embedModel;
}
withEmbedModel(embedModel, fn) {
return Settings$1.withEmbedModel(embedModel, fn);
}
get nodeParser() {
if (this.#nodeParser === null) {
this.#nodeParser = new SentenceSplitter({
chunkSize: this.chunkSize,
chunkOverlap: this.chunkOverlap
});
}
return this.#nodeParserAsyncLocalStorage.getStore() ?? this.#nodeParser;
}
set nodeParser(nodeParser) {
this.#nodeParser = nodeParser;
}
withNodeParser(nodeParser, fn) {
return this.#nodeParserAsyncLocalStorage.run(nodeParser, fn);
}
get callbackManager() {
return Settings$1.callbackManager;
}
set callbackManager(callbackManager) {
Settings$1.callbackManager = callbackManager;
}
withCallbackManager(callbackManager, fn) {
return Settings$1.withCallbackManager(callbackManager, fn);
}
set chunkSize(chunkSize) {
Settings$1.chunkSize = chunkSize;
}
get chunkSize() {
return Settings$1.chunkSize;
}
withChunkSize(chunkSize, fn) {
return Settings$1.withChunkSize(chunkSize, fn);
}
get chunkOverlap() {
return this.#chunkOverlapAsyncLocalStorage.getStore() ?? this.#chunkOverlap;
}
set chunkOverlap(chunkOverlap) {
if (typeof chunkOverlap === "number") {
this.#chunkOverlap = chunkOverlap;
}
}
withChunkOverlap(chunkOverlap, fn) {
return this.#chunkOverlapAsyncLocalStorage.run(chunkOverlap, fn);
}
get prompt() {
return this.#promptAsyncLocalStorage.getStore() ?? this.#prompt;
}
set prompt(prompt) {
this.#prompt = prompt;
}
withPrompt(prompt, fn) {
return this.#promptAsyncLocalStorage.run(prompt, fn);
}
constructor(){
this.#prompt = {};
this.#promptHelper = null;
this.#nodeParser = null;
this.#promptHelperAsyncLocalStorage = new AsyncLocalStorage();
this.#nodeParserAsyncLocalStorage = new AsyncLocalStorage();
this.#chunkOverlapAsyncLocalStorage = new AsyncLocalStorage();
this.#promptAsyncLocalStorage = new AsyncLocalStorage();
}
}
const Settings = new GlobalSettings();
const DEFAULT_NAME = "query_engine_tool";
const DEFAULT_DESCRIPTION = "Useful for running a natural language query against a knowledge base and get back a natural language response.";
const DEFAULT_PARAMETERS = {
type: "object",
properties: {
query: {
type: "string",
description: "The query to search for"
}
},
required: [
"query"
]
};
class QueryEngineTool {
constructor({ queryEngine, metadata, includeSourceNodes }){
this.queryEngine = queryEngine;
this.metadata = {
name: metadata?.name ?? DEFAULT_NAME,
description: metadata?.description ?? DEFAULT_DESCRIPTION,
parameters: metadata?.parameters ?? DEFAULT_PARAMETERS
};
this.includeSourceNodes = includeSourceNodes ?? false;
}
async call({ query }) {
const response = await this.queryEngine.query({
query
});
if (!this.includeSourceNodes) {
return {
content: response.message.content
};
}
// Use JSON.parse(JSON.stringify()) to remove undefined values from sourceNodes
// since undefined is not a valid JSONValue
return {
content: response.message.content,
sourceNodes: JSON.parse(JSON.stringify(response.sourceNodes ?? []))
};
}
}
/**
* Indexes are the data structure that we store our nodes and embeddings in so
* they can be retrieved for our queries.
*/ class BaseIndex {
constructor(init){
this.storageContext = init.storageContext;
this.docStore = init.docStore;
this.indexStore = init.indexStore;
this.indexStruct = init.indexStruct;
}
/**
* Returns a query tool by calling asQueryEngine.
* Either options or retriever can be passed, but not both.
* If options are provided, they are passed to generate a retriever.
*/ asQueryTool(params) {
if (params.options) {
params.retriever = this.asRetriever(params.options);
}
return new QueryEngineTool({
queryEngine: this.asQueryEngine(params),
metadata: params?.metadata,
includeSourceNodes: params?.includeSourceNodes ?? false
});
}
/**
* Insert a document into the index.
* @param document
*/ async insert(document) {
const nodes = await runTransformations([
document
], [
Settings.nodeParser
]);
await this.insertNodes(nodes);
await this.docStore.setDocumentHash(document.id_, document.hash);
}
/**
* Alias for asRetriever
* @param options
*/ // eslint-disable-next-line @typescript-eslint/no-explicit-any
retriever(options) {
return this.asRetriever(options);
}
/**
* Alias for asQueryEngine
* @param options you can supply your own custom Retriever and ResponseSynthesizer
*/ queryEngine(options) {
return this.asQueryEngine(options);
}
/**
* Alias for asQueryTool
* Either options or retriever can be passed, but not both.
* If options are provided, they are passed to generate a retriever.
*/ queryTool(params) {
return this.asQueryTool(params);
}
}
// FS utility helpers
/**
* Checks if a file exists.
* Analogous to the os.path.exists function from Python.
* @param path The path to the file to check.
* @returns A promise that resolves to true if the file exists, false otherwise.
*/ async function exists(path) {
try {
await fs.access(path);
return true;
} catch {
return false;
}
}
const LEARNER_MODES = new Set([
VectorStoreQueryMode.SVM,
VectorStoreQueryMode.LINEAR_REGRESSION,
VectorStoreQueryMode.LOGISTIC_REGRESSION
]);
const MMR_MODE = VectorStoreQueryMode.MMR;
// Mapping of filter operators to metadata filter functions
const OPERATOR_TO_FILTER = {
[FilterOperator.EQ]: ({ key, value }, metadata)=>{
return metadata[key] === parsePrimitiveValue(value);
},
[FilterOperator.NE]: ({ key, value }, metadata)=>{
return metadata[key] !== parsePrimitiveValue(value);
},
[FilterOperator.IN]: ({ key, value }, metadata)=>{
return !!parseArrayValue(value).find((v)=>metadata[key] === v);
},
[FilterOperator.NIN]: ({ key, value }, metadata)=>{
return !parseArrayValue(value).find((v)=>metadata[key] === v);
},
[FilterOperator.ANY]: ({ key, value }, metadata)=>{
if (!Array.isArray(metadata[key])) return false;
return parseArrayValue(value).some((v)=>metadata[key].includes(v));
},
[FilterOperator.ALL]: ({ key, value }, metadata)=>{
if (!Array.isArray(metadata[key])) return false;
return parseArrayValue(value).every((v)=>metadata[key].includes(v));
},
[FilterOperator.TEXT_MATCH]: ({ key, value }, metadata)=>{
return metadata[key].includes(parsePrimitiveValue(value));
},
[FilterOperator.CONTAINS]: ({ key, value }, metadata)=>{
if (!Array.isArray(metadata[key])) return false;
return !!parseArrayValue(metadata[key]).find((v)=>v === value);
},
[FilterOperator.GT]: ({ key, value }, metadata)=>{
return metadata[key] > parsePrimitiveValue(value);
},
[FilterOperator.LT]: ({ key, value }, metadata)=>{
return metadata[key] < parsePrimitiveValue(value);
},
[FilterOperator.GTE]: ({ key, value }, metadata)=>{
return metadata[key] >= parsePrimitiveValue(value);
},
[FilterOperator.LTE]: ({ key, value }, metadata)=>{
return metadata[key] <= parsePrimitiveValue(value);
}
};
// Build a filter function based on the metadata and the preFilters
const buildFilterFn = (metadata, preFilters)=>{
if (!preFilters) return true;
if (!metadata) return false;
const { filters, condition } = preFilters;
const queryCondition = condition || "and"; // default to and
const itemFilterFn = (filter)=>{
if (filter.operator === FilterOperator.IS_EMPTY) {
// for `is_empty` operator, return true if the metadata key is not present or the value is empty
const value = metadata[filter.key];
return value === undefined || value === null || value === "" || Array.isArray(value) && value.length === 0;
}
if (metadata[filter.key] === undefined) {
// for other operators, always return false if the metadata key is not present
return false;
}
const metadataLookupFn = OPERATOR_TO_FILTER[filter.operator];
if (!metadataLookupFn) throw new Error(`Unsupported operator: ${filter.operator}`);
return metadataLookupFn(filter, metadata);
};
if (queryCondition === "and") return filters.every(itemFilterFn);
return filters.some(itemFilterFn);
};
class SimpleVectorStoreData {
constructor(){
this.embeddingDict = {};
this.textIdToRefDocId = {};
this.metadataDict = {};
}
}
class SimpleVectorStore extends BaseVectorStore {
constructor(init){
super(init), this.storesText = false;
this.data = init?.data || new SimpleVectorStoreData();
}
static async fromPersistDir(persistDir = DEFAULT_PERSIST_DIR, embedModel, options) {
const persistPath = path.join(persistDir, "vector_store.json");
return await SimpleVectorStore.fromPersistPath(persistPath, embedModel, options);
}
client() {
return null;
}
async get(textId) {
return this.data.embeddingDict[textId];
}
async add(embeddingResults) {
for (const node of embeddingResults){
this.data.embeddingDict[node.id_] = node.getEmbedding();
if (!node.sourceNode) {
continue;
}
this.data.textIdToRefDocId[node.id_] = node.sourceNode?.nodeId;
// Add metadata to the metadataDict
const metadata = nodeToMetadata(node, true, undefined, false);
delete metadata["_node_content"];
this.data.metadataDict[node.id_] = metadata;
}
if (this.persistPath) {
await this.persist(this.persistPath);
}
return embeddingResults.map((result)=>result.id_);
}
async delete(refDocId) {
const textIdsToDelete = Object.keys(this.data.textIdToRefDocId).filter((textId)=>this.data.textIdToRefDocId[textId] === refDocId);
for (const textId of textIdsToDelete){
delete this.data.embeddingDict[textId];
delete this.data.textIdToRefDocId[textId];
if (this.data.metadataDict) delete this.data.metadataDict[textId];
}
if (this.persistPath) {
await this.persist(this.persistPath);
}
return Promise.resolve();
}
async filterNodes(query) {
const items = Object.entries(this.data.embeddingDict);
const queryFilterFn = (nodeId)=>{
const metadata = this.data.metadataDict[nodeId];
return buildFilterFn(metadata, query.filters);
};
const nodeFilterFn = (nodeId)=>{
if (!query.docIds) return true;
const availableIds = new Set(query.docIds);
return availableIds.has(nodeId);
};
const queriedItems = items.filter((item)=>nodeFilterFn(item[0]) && queryFilterFn(item[0]));
const nodeIds = queriedItems.map((item)=>item[0]);
const embeddings = queriedItems.map((item)=>item[1]);
return {
nodeIds,
embeddings
};
}
async query(query) {
const { nodeIds, embeddings } = await this.filterNodes(query);
const queryEmbedding = query.queryEmbedding;
let topSimilarities, topIds;
if (LEARNER_MODES.has(query.mode)) {
// fixme: unfinished
throw new Error("Learner modes not implemented for SimpleVectorStore yet.");
} else if (query.mode === MMR_MODE) {
const mmrThreshold = query.mmrThreshold;
[topSimilarities, topIds] = getTopKMMREmbeddings(queryEmbedding, embeddings, null, query.similarityTopK, nodeIds, mmrThreshold);
} else if (query.mode === VectorStoreQueryMode.DEFAULT) {
[topSimilarities, topIds] = getTopKEmbeddings(queryEmbedding, embeddings, query.similarityTopK, nodeIds);
} else {
throw new Error(`Invalid query mode: ${query.mode}`);
}
return Promise.resolve({
similarities: topSimilarities,
ids: topIds
});
}
async persist(persistPath = path.join(DEFAULT_PERSIST_DIR, "vector_store.json")) {
await SimpleVectorStore.persistData(persistPath, this.data);
}
static async persistData(persistPath, data) {
const dirPath = path.dirname(persistPath);
if (!await exists(dirPath)) {
await fs.mkdir(dirPath);
}
await fs.writeFile(persistPath, JSON.stringify(data));
}
static async fromPersistPath(persistPath, embedModel, options) {
const logger = options?.logger ?? consoleLogger;
const dirPath = path.dirname(persistPath);
if (!await exists(dirPath)) {
await fs.mkdir(dirPath, {
recursive: true
});
}
let dataDict = {};
if (!await exists(persistPath)) {
logger.log(`Starting new store from path: ${persistPath}`);
} else {
try {
const fileData = await fs.readFile(persistPath);
dataDict = JSON.parse(fileData.toString());
} catch (e) {
throw new Error(`Failed to load data from path: ${persistPath}`, {
cause: e
});
}
}
const data = new SimpleVectorStoreData();
// @ts-expect-error TS2322
data.embeddingDict = dataDict.embeddingDict ?? {};
// @ts-expect-error TS2322
data.textIdToRefDocId = dataDict.textIdToRefDocId ?? {};
// @ts-expect-error TS2322
data.metadataDict = dataDict.metadataDict ?? {};
const store = new SimpleVectorStore({
data,
embedModel
});
store.persistPath = persistPath;
return store;
}
static fromDict(saveDict, embedModel) {
const data = new SimpleVectorStoreData();
data.embeddingDict = saveDict.embeddingDict;
data.textIdToRefDocId = saveDict.textIdToRefDocId;
data.metadataDict = saveDict.metadataDict;
return new SimpleVectorStore({
data,
embedModel
});
}
toDict() {
return {
embeddingDict: this.data.embeddingDict,
textIdToRefDocId: this.data.textIdToRefDocId,
metadataDict: this.data.metadataDict
};
}
}
class SimpleDocumentStore extends KVDocumentStore {
constructor(kvStore, namespace){
kvStore = kvStore || new SimpleKVStore();
namespace = namespace || DEFAULT_NAMESPACE;
super(kvStore, namespace);
this.kvStore = kvStore;
}
static async fromPersistDir(persistDir = DEFAULT_PERSIST_DIR, namespace, options) {
const persistPath = path.join(persistDir, DEFAULT_DOC_STORE_PERSIST_FILENAME);
return await SimpleDocumentStore.fromPersistPath(persistPath, namespace, options);
}
static async fromPersistPath(persistPath, namespace, options) {
const simpleKVStore = await SimpleKVStore.fromPersistPath(persistPath, options);
return new SimpleDocumentStore(simpleKVStore, namespace);
}
async persist(persistPath = path.join(DEFAULT_PERSIST_DIR, DEFAULT_DOC_STORE_PERSIST_FILENAME)) {
if (_.isObject(this.kvStore) && this.kvStore instanceof BaseInMemoryKVStore) {
await this.kvStore.persist(persistPath);
}
}
static fromDict(saveDict, namespace) {
const simpleKVStore = SimpleKVStore.fromDict(saveDict);
return new SimpleDocumentStore(simpleKVStore, namespace);
}
toDict() {
if (_.isObject(this.kvStore) && this.kvStore instanceof SimpleKVStore) {
return this.kvStore.toDict();
}
// If the kvstore is not a SimpleKVStore, you might want to throw an error or return a default value.
throw new Error("KVStore is not a SimpleKVStore");
}
}
async function storageContextFromDefaults({ docStore, indexStore, vectorStore, vectorStores, persistDir }) {
vectorStores = vectorStores ?? {};
if (!persistDir) {
docStore = docStore ?? new SimpleDocumentStore();
indexStore = indexStore ?? new SimpleIndexStore();
if (!(ModalityType.TEXT in vectorStores)) {
vectorStores[ModalityType.TEXT] = vectorStore ?? new SimpleVectorStore();
}
} else {
const embedModel = Settings.embedModel;
docStore = docStore || await SimpleDocumentStore.fromPersistDir(persistDir, DEFAULT_NAMESPACE);
indexStore = indexStore || await SimpleIndexStore.fromPersistDir(persistDir);
if (!(ObjectType.TEXT in vectorStores)) {
vectorStores[ModalityType.TEXT] = vectorStore ?? await SimpleVectorStore.fromPersistDir(persistDir, embedModel);
}
}
return {
docStore,
indexStore,
vectorStores
};
}
// generate from "tsup ./src/index.js --format esm"
var __getOwnPropNames = Object.getOwnPropertyNames;
var __commonJS = (cb, mod)=>function __require() {
return mod || (0, cb[__getOwnPropNames(cb)[0]])((mod = {
exports: {}
}).exports, mod), mod.exports;
};
// src/stopwords.js
var require_stopwords = __commonJS({
"src/stopwords.js" (exports, module) {
module.exports = {
stopwords: [
"a",
"about",
"above",
"across",
"after",
"again",
"against",
"all",
"almost",
"alone",
"along",
"already",
"also",
"although",
"always",
"among",
"an",
"and",
"another",
"any",
"anybody",
"anyone",
"anything",
"anywhere",
"are",
"area",
"areas",
"around",
"as",
"ask",
"asked",
"asking",
"asks",
"at",
"away",
"b",
"back",
"backed",
"backing",
"backs",
"be",
"because",
"become",
"becomes",
"became",
"been",
"before",
"began",
"behind",
"being",
"beings",
"best",
"better",
"between",
"big",
"both",
"but",
"by",
"c",
"came",
"can",
"cannot",
"case",
"cases",
"certain",
"certainly",
"clear",
"clearly",
"come",
"contains",
"could",
"d",
"did",
"differ",
"different",
"differently",
"do",
"does",
"done",
"down",
"downed",
"downing",
"downs",
"during",
"e",
"each",
"early",
"either",
"end",
"ended",
"ending",
"ends",
"enough",
"even",
"evenly",
"ever",
"every",
"everybody",
"everyone",
"everything",
"everywhere",
"f",
"face",
"faces",
"fact",
"facts",
"far",
"felt",
"few",
"find",
"finds",
"first",
"for",
"four",
"from",
"full",
"fully",
"further",
"furthered",
"furthering",
"furthers",
"g",
"gave",
"general",
"generally",
"get",
"gets",
"give",
"given",
"gives",
"go",
"going",
"good",
"goods",
"got",
"great",
"greater",
"greatest",
"group",
"grouped",
"grouping",
"groups",
"h",
"had",
"has",
"have",
"having",
"he",
"her",
"herself",
"here",
"high",
"higher",
"highest",
"him",
"himself",
"his",
"how",
"however",
"i",
"if",
"important",
"in",
"interest",
"interested",
"interesting",
"interests",
"into",
"is",
"it",
"its",
"itself",
"j",
"just",
"k",
"keep",
"keeps",
"kind",
"knew",
"know",
"known",
"knows",
"l",
"large",
"largely",
"last",
"later",
"latest",
"least",
"less",
"let",
"lets",
"like",
"likely",
"long",
"longer",
"longest",
"m",
"made",
"make",
"making",
"man",
"many",
"may",
"me",
"member",
"members",
"men",
"might",
"more",
"most",
"mostly",
"mr",
"mrs",
"much",
"must",
"my",
"myself",
"n",
"necessary",
"need",
"needed",
"needing",
"needs",
"never",
"new",
"newer",
"newest",
"next",
"no",
"non",
"not",
"nobody",
"noone",
"nothing",
"now",
"nowhere",
"number",
"numbers",
"o",
"of",
"off",
"often",
"old",
"older",
"oldest",
"on",
"once",
"one",
"only",
"open",
"opened",
"opening",
"opens",
"or",
"order",
"ordered",
"ordering",
"orders",
"other",
"others",
"our",
"out",
"over",
"p",
"part",
"parted",
"parting",
"parts",
"per",
"perhaps",
"place",
"places",
"point",
"pointed",
"pointing",
"points",
"possible",
"present",
"presented",
"presenting",
"presents",
"problem",
"problems",
"put",
"puts",
"q",
"quite",
"r",
"rather",
"really",
"right",
"room",
"rooms",
"s",
"said",
"same",
"saw",
"say",
"says",
"second",
"seconds",
"see",
"sees",
"seem",
"seemed",
"seeming",
"seems",
"several",
"shall",
"she",
"should",
"show",
"showed",
"showing",
"shows",
"side",
"sides",
"since",
"small",
"smaller",
"smallest",
"so",
"some",
"somebody",
"someone",
"something",
"somewhere",
"state",
"states",
"still",
"such",
"sure",
"t",
"take",
"taken",
"than",
"that",
"the",
"their",
"them",
"then",
"there",
"therefore",
"these",
"they",
"thing",
"things",
"think",
"thinks",
"this",
"those",
"though",
"thought",
"thoughts",
"three",
"through",
"thus",
"to",
"today",
"together",
"too",
"took",
"toward",
"turn",
"turned",
"turning",
"turns",
"two",
"u",
"under",
"until",
"up",
"upon",
"us",
"use",
"uses",
"used",
"v",
"very",
"w",
"want",
"wanted",
"wanting",
"wants",
"was",
"way",
"ways",
"we",
"well",
"wells",
"went",
"were",
"what",
"when",
"where",
"whether",
"which",
"while",
"who",
"whole",
"whose",
"why",
"will",
"with",
"within",
"without",
"work",
"worked",
"working",
"works",
"would",
"y",
"year",
"years",
"yet",
"you",
"young",
"younger",
"youngest",
"your",
"yours",
"eoc",
"mu",
"sigma",
"mu sigma",
"musigma",
"client",
"clients",
"capabilities",
"capability",
"firm",
"firms",
"biggest",
"-"
]
};
}
});
const { fromPairs, sortBy, toPairs } = _;
var stopwords = require_stopwords();
function isNumber(str) {
return /\d/.test(str);
}
function isAcceptable(phrase, minCharLength, maxWordsLength) {
if (phrase < minCharLength) {
return false;
}
let words = phrase.split(" ");
if (words.length > maxWordsLength) {
return false;
}
let digits = 0;
let alpha = 0;
for(let i = 0; i < phrase.length; i++){
if (/\d/.test(phrase[i])) digits += 1;
if (/[a-zA-Z]/.test(phrase[i])) alpha += 1;
}
if (alpha == 0) {
return false;
}
if (digits > alpha) {
return false;
}
return true;
}
function countOccurances(haystack, needle) {
return haystack.reduce((n, value)=>{
return n + (value === needle);
}, 0);
}
function generateCandidateKeywordScores(phraseList, wordScore, minKeywordFrequency = 1) {
let keywordCandidates = {};
phraseList.forEach((phrase)=>{
if (minKeywordFrequency > 1) {
if (countOccurances(phraseList, phrase) < minKeywordFrequency) {
return;
}
}
phrase in keywordCandidates || (keywordCandidates[phrase] = 0);
let wordList = separateWords(phrase, 0);
let candidateScore = 0;
wordList.forEach((word)=>{
candidateScore += wordScore[word];
keywordCandidates[phrase] = candidateScore;
});
});
return keywordCandidates;
}
function separateWords(text, minWordReturnSize) {
let wordDelimiters = /[^a-zA-Z0-9_\+\-/]/;
let words = [];
text.split(wordDelimiters).forEach((singleWord)=>{
let currentWord = singleWord.trim().toLowerCase();
if (currentWord.length > minWordReturnSize && currentWord != "" && !isNumber(currentWord)) {
words.push(currentWord);
}
});
return words;
}
function calculateWordScores(phraseList) {
let wordFrequency = {};
let wordDegree = {};
phraseList.forEach((phrase)=>{
let wordList = separateWords(phrase, 0);
let wordListLength = wordList.length;
let wordListDegree = wordListLength - 1;
wordList.forEach((word)=>{
word in wordFrequency || (wordFrequency[word] = 0);
wordFrequency[word] += 1;
word in wordDegree || (wordDegree[word] = 0);
wordDegree[word] += wordListDegree;
});
});
Object.keys(wordFrequency).forEach((item)=>{
wordDegree[item] = wordDegree[item] + wordFrequency[item];
});
let wordScore = {};
Object.keys(wordFrequency).forEach((item)=>{
item in wordScore || (wordScore[item] = 0);
wordScore[item] = wordDegree[item] / (wordFrequency[item] * 1);
});
return wordScore;
}
function generateCandidateKeywords(sentenceList, stopWordPattern, minCharLength = 1, maxWordsLength = 5) {
let phraseList = [];
sentenceList.forEach((sentence)=>{
let tmp = stopWordPattern[Symbol.replace](sentence, "|");
let phrases = tmp.split("|");
phrases.forEach((ph)=>{
let phrase = ph.trim().toLowerCase();
if (phrase != "" && isAcceptable(phrase, minCharLength, maxWordsLength)) {
phraseList.push(phrase);
}
});
});
return phraseList;
}
function buildStopWordRegex(path) {
let stopWordList = loadStopWords();
let stopWordRegexList = [];
stopWordList.forEach((word)=>{
if (/\w+/.test(word)) {
let wordRegex = `\\b${word}\\b`;
stopWordRegexList.push(wordRegex);
}
});
let stopWordPattern = new RegExp(stopWordRegexList.join("|"), "ig");
return stopWordPattern;
}
function splitSentences(text) {
let sentenceDelimiters = /[\[\]\n.!?,;:\t\\-\\"\\(\\)\\\'\u2019\u2013]/;
return text.split(sentenceDelimiters);
}
function loadStopWords(path) {
let contents = stopwords.stopwords;
return contents;
}
function rake(text, stopWordsPath, minCharLength = 3, maxWordsLength = 5, minKeywordFrequency = 1) {
let stopWordPattern = buildStopWordRegex();
let sentenceList = splitSentences(text);
let phraseList = generateCandidateKeywords(sentenceList, stopWordPattern, minCharLength, maxWordsLength);
let wordScores = calculateWordScores(phraseList);
let keywordCandidates = generateCandidateKeywordScores(phraseList, wordScores, minKeywordFrequency);
let sortedKeywords = fromPairs(sortBy(toPairs(keywordCandidates), (pair)=>pair[1]).reverse());
return sortedKeywords;
}
// Get subtokens from a list of tokens., filtering for stopwords.
function expandTokensWithSubtokens(tokens) {
const results = new Set();
const regex = /\w+/g;
for (const token of tokens){
results.add(token);
const subTokens = token.match(regex);
if (subTokens && subTokens.length > 1) {
for (const w of subTokens){
results.add(w);
}
}
}
return results;
}
function extractKeywordsGivenResponse(response, startToken = "", lowercase = true) {
const results = [];
response = response.trim();
if (response.startsWith(startToken)) {
response = response.substring(startToken.length);
}
const keywords = response.split(",");
for (const k of keywords){
let rk = k;
if (lowercase) {
rk = rk.toLowerCase();
}
results.push(rk.trim());
}
return expandTokensWithSubtokens(new Set(results));
}
function simpleExtractKeywords(textChunk, maxKeywords) {
const regex = /\w+/g;
const tokens = [
...textChunk.matchAll(regex)
].map((token)=>token[0].toLowerCase().trim());
// Creating a frequency map
const valueCounts = {};
for (const token of tokens){
valueCounts[token] = (valueCounts[token] || 0) + 1;
}
// Sorting tokens by frequency
const sortedTokens = Object.keys(valueCounts).sort((a, b)=>valueCounts[b] - valueCounts[a]);
const keywords = maxKeywords ? sortedTokens.slice(0, maxKeywords) : sortedTokens;
return new Set(keywords);
}
function rakeExtractKeywords(textChunk, maxKeywords) {
const keywords = Object.keys(rake(textChunk));
const limitedKeywords = maxKeywords ? keywords.slice(0, maxKeywords) : keywords;
return new Set(limitedKeywords);
}
function applyDecs2203RFactory() {
function createAddInitializerMethod(initializers, decoratorFinishedRef) {
return function addInitializer(initializer) {
assertNotFinished(decoratorFinishedRef, "addInitializer");
assertCallable(initializer, "An initializer");
initializers.push(initializer);
};
}
function memberDec(dec, name, desc, initializers, kind, isStatic, isPrivate, metadata, value) {
var kindStr;
switch(kind){
case 1:
kindStr = "accessor";
break;
case 2:
kindStr = "method";
break;
case 3:
kindStr = "getter";
break;
case 4:
kindStr = "setter";
break;
default:
kindStr = "field";
}
var ctx = {
kind: kindStr,
name: isPrivate ? "#" + name : name,
static: isStatic,
private: isPrivate,
metadata: metadata
};
var decoratorFinishedRef = {
v: false
};
ctx.addInitializer = createAddInitializerMethod(initializers, decoratorFinishedRef);
var get, set;
if (kind === 0) {
if (isPrivate) {
get = desc.get;
set = desc.set;
} else {
get = function() {
return this[name];
};
set = function(v) {
this[name] = v;
};
}
} else if (kind === 2) {
get = function() {
return desc.value;
};
} else {
if (kind === 1 || kind === 3) {
get = function() {
return desc.get.call(this);
};
}
if (kind === 1 || kind === 4) {
set = function(v) {
desc.set.call(this, v);
};
}
}
ctx.access = get && set ? {
get: get,
set: set
} : get ? {
get: get
} : {
set: set
};
try {
return dec(value, ctx);
} finally{
decoratorFinishedRef.v = true;
}
}
function assertNotFinished(decoratorFinishedRef, fnName) {
if (decoratorFinishedRef.v) {
throw new Error("attempted to call " + fnName + " after decoration was finished");
}
}
function assertCallable(fn, hint) {
if (typeof fn !== "function") {
throw new TypeError(hint + " must be a function");
}
}
function assertValidReturnValue(kind, value) {
var type = typeof value;
if (kind === 1) {
if (type !== "object" || value === null) {
throw new TypeError("accessor decorators must return an object with get, set, or init properties or void 0");
}
if (value.get !== undefined) {
assertCallable(value.get, "accessor.get");
}
if (value.set !== undefined) {
assertCallable(value.set, "accessor.set");
}
if (value.init !== undefined) {
assertCallable(value.init, "accessor.init");
}
} else if (type !== "function") {
var hint;
if (kind === 0) {
hint = "field";
} else if (kind === 10) {
hint = "class";
} else {
hint = "method";
}
throw new TypeError(hint + " decorators must return a function or void 0");
}
}
function applyMemberDec(ret, base, decInfo, name, kind, isStatic, isPrivate, initializers, metadata) {
var decs = decInfo[0];
var desc, init, value;
if (isPrivate) {
if (kind === 0 || kind === 1) {
desc = {
get: decInfo[3],
set: decInfo[4]
};
} else if (kind === 3) {
desc = {
get: decInfo[3]
};
} else if (kind === 4) {
desc = {
set: decInfo[3]
};
} else {
desc = {
value: decInfo[3]
};
}
} else if (kind !== 0) {
desc = Object.getOwnPropertyDescriptor(base, name);
}
if (kind === 1) {
value = {
get: desc.get,
set: desc.set
};
} else if (kind === 2) {
value = desc.value;
} else if (kind === 3) {
value = desc.get;
} else if (kind === 4) {
value = desc.set;
}
var newValue, get, set;
if (typeof decs === "function") {
newValue = memberDec(decs, name, desc, initializers, kind, isStatic, isPrivate, metadata, value);
if (newValue !== void 0) {
assertValidReturnValue(kind, newValue);
if (kind === 0) {
init = newValue;
} else if (kind === 1) {
init = newValue.init;
get = newValue.get || value.get;
set = newValue.set || value.set;
value = {
get: get,
set: set
};
} else {
value = newValue;
}
}
} else {
for(var i = decs.length - 1; i >= 0; i--){
var dec = decs[i];
newValue = memberDec(dec, name, desc, initializers, kind, isStatic, isPrivate, metadata, value);
if (newValue !== void 0) {
assertValidReturnValue(kind, newValue);
var newInit;
if (kind === 0) {
newInit = newValue;
} else if (kind === 1) {
newInit = newValue.init;
get = newValue.get || value.get;
set = newValue.set || value.set;
value = {
get: get,
set: set
};
} else {
value = newValue;
}
if (newInit !== void 0) {
if (init === void 0) {
init = newInit;
} else if (typeof init === "function") {
init = [
init,
newInit
];
} else {
init.push(newInit);
}
}
}
}
}
if (kind === 0 || kind === 1) {
if (init === void 0) {
init = function(instance, init) {
return init;
};
} else if (typeof init !== "function") {
var ownInitializers = init;
init = function(instance, init) {
var value = init;
for(var i = 0; i < ownInitializers.length; i++){
value = ownInitializers[i].call(instance, value);
}
return value;
};
} else {
var originalInitializer = init;
init = function(instance, init) {
return originalInitializer.call(instance, init);
};
}
ret.push(init);
}
if (kind !== 0) {
if (kind === 1) {
desc.get = value.get;
desc.set = value.set;
} else if (kind === 2) {
desc.value = value;
} else if (kind === 3) {
desc.get = value;
} else if (kind === 4) {
desc.set = value;
}
if (isPrivate) {
if (kind === 1) {
ret.push(function(instance, args) {
return value.get.call(instance, args);
});
ret.push(function(instance, args) {
return value.set.call(instance, args);
});
} else if (kind === 2) {
ret.push(value);
} else {
ret.push(function(instance, args) {
return value.call(instance, args);
});
}
} else {
Object.defineProperty(base, name, desc);
}
}
}
function applyMemberDecs(Class, decInfos, metadata) {
var ret = [];
var protoInitializers;
var staticInitializers;
var existingProtoNonFields = new Map();
var existingStaticNonFields = new Map();
for(var i = 0; i < decInfos.length; i++){
var decInfo = decInfos[i];
if (!Array.isArray(decInfo)) continue;
var kind = decInfo[1];
var name = decInfo[2];
var isPrivate = decInfo.length > 3;
var isStatic = kind >= 5;
var base;
var initializers;
if (isStatic) {
base = Class;
kind = kind - 5;
staticInitializers = staticInitializers || [];
initializers = staticInitializers;
} else {
base = Class.prototype;
protoInitializers = protoInitializers || [];
initializers = protoInitializers;
}
if (kind !== 0 && !isPrivate) {
var existingNonFields = isStatic ? existingStaticNonFields : existingProtoNonFields;
var existingKind = existingNonFields.get(name) || 0;
if (existingKind === true || existingKind === 3 && kind !== 4 || existingKind === 4 && kind !== 3) {
throw new Error("Attempted to decorate a public method/accessor that has the same name as a previously decorated public method/accessor. This is not currently supported by the decorators plugin. Property name was: " + name);
} else if (!existingKind && kind > 2) {
existingNonFields.set(name, kind);
} else {
existingNonFields.set(name, true);
}
}
applyMemberDec(ret, base, decInfo, name, kind, isStatic, isPrivate, initializers, metadata);
}
pushInitializers(ret, protoInitializers);
pushInitializers(ret, staticInitializers);
return ret;
}
function pushInitializers(ret, initializers) {
if (initializers) {
ret.push(function(instance) {
for(var i = 0; i < initializers.length; i++){
initializers[i].call(instance);
}
return instance;
});
}
}
function applyClassDecs(targetClass, classDecs, metadata) {
if (classDecs.length > 0) {
var initializers = [];
var newClass = targetClass;
var name = targetClass.name;
for(var i = classDecs.length - 1; i >= 0; i--){
var decoratorFinishedRef = {
v: false
};
try {
var nextNewClass = classDecs[i](newClass, {
kind: "class",
name: name,
addInitializer: createAddInitializerMethod(initializers, decoratorFinishedRef),
metadata
});
} finally{
decoratorFinishedRef.v = true;
}
if (nextNewClass !== undefined) {
assertValidReturnValue(10, nextNewClass);
newClass = nextNewClass;
}
}
return [
defineMetadata(newClass, metadata),
function() {
for(var i = 0; i < initializers.length; i++){
initializers[i].call(newClass);
}
}
];
}
}
function defineMetadata(Class, metadata) {
return Object.defineProperty(Class, Symbol.metadata || Symbol.for("Symbol.metadata"), {
configurable: true,
enumerable: true,
value: metadata
});
}
return function applyDecs2203R(targetClass, memberDecs, classDecs, parentClass) {
if (parentClass !== void 0) {
var parentMetadata = parentClass[Symbol.metadata || Symbol.for("Symbol.metadata")];
}
var metadata = Object.create(parentMetadata === void 0 ? null : parentMetadata);
var e = applyMemberDecs(targetClass, memberDecs, metadata);
if (!classDecs.length) defineMetadata(targetClass, metadata);
return {
e: e,
get c () {
return applyClassDecs(targetClass, classDecs, metadata);
}
};
};
}
function _apply_decs_2203_r(targetClass, memberDecs, classDecs, parentClass) {
return (_apply_decs_2203_r = applyDecs2203RFactory())(targetClass, memberDecs, classDecs, parentClass);
}
var _initProto;
/**
* CondenseQuestionChatEngine is used in conjunction with a Index (for example VectorStoreIndex).
* It does two steps on taking a user's chat message: first, it condenses the chat message
* with the previous chat history into a question with more context.
* Then, it queries the underlying Index using the new question with context and returns
* the response.
* CondenseQuestionChatEngine performs well when the input is primarily questions about the
* underlying data. It performs less well when the chat messages are not questions about the
* data, or are very referential to previous context.
*/ class CondenseQuestionChatEngine extends BaseChatEngine {
static{
({ e: [_initProto] } = _apply_decs_2203_r(this, [
[
wrapEventCaller,
2,
"chat"
]
], []));
}
get chatHistory() {
return this.memory.getLLM();
}
constructor(init){
super(), _initProto(this);
this.queryEngine = init.queryEngine;
this.memory = createMemory(init.chatHistory);
this.llm = Settings.llm;
this.condenseMessagePrompt = init?.condenseMessagePrompt ?? defaultCondenseQuestionPrompt;
}
_getPromptModules() {
return {};
}
_getPrompts() {
return {
condenseMessagePrompt: this.condenseMessagePrompt
};
}
_updatePrompts(promptsDict) {
if (promptsDict.condenseMessagePrompt) {
this.condenseMessagePrompt = promptsDict.condenseMessagePrompt;
}
}
async condenseQuestion(chatHistory, question) {
const chatHistoryStr = messagesToHistory(await chatHistory.getLLM());
return this.llm.complete({
prompt: this.condenseMessagePrompt.format({
question: question,
chatHistory: chatHistoryStr
})
});
}
async chat(params) {
const { message, stream } = params;
const chatHistory = params.chatHistory ? params.chatHistory instanceof Memory ? params.chatHistory : createMemory(params.chatHistory) : this.memory;
const condensedQuestion = (await this.condenseQuestion(chatHistory, extractText(message))).text;
await chatHistory.add({
content: message,
role: "user"
});
if (stream) {
const stream = await this.queryEngine.query({
query: condensedQuestion,
stream: true
});
return streamReducer({
stream,
initialValue: "",
reducer: (accumulator, part)=>accumulator += extractText(part.message.content),
finished: (accumulator)=>{
void chatHistory.add({
content: accumulator,
role: "assistant"
});
}
});
}
const response = await this.queryEngine.query({
query: condensedQuestion
});
await chatHistory.add({
content: response.message.content,
role: "assistant"
});
return response;
}
reset() {
void this.memory.clear();
}
}
var KeywordTableRetrieverMode = /*#__PURE__*/ function(KeywordTableRetrieverMode) {
KeywordTableRetrieverMode["DEFAULT"] = "DEFAULT";
KeywordTableRetrieverMode["SIMPLE"] = "SIMPLE";
KeywordTableRetrieverMode["RAKE"] = "RAKE";
return KeywordTableRetrieverMode;
}({});
// Base Keyword Table Retriever
class BaseKeywordTableRetriever extends BaseRetriever {
constructor({ index, keywordExtractTemplate, queryKeywordExtractTemplate, maxKeywordsPerQuery = 10, numChunksPerQuery = 10 }){
super();
this.index = index;
this.indexStruct = index.indexStruct;
this.docstore = index.docStore;
this.llm = Settings.llm;
this.maxKeywordsPerQuery = maxKeywordsPerQuery;
this.numChunksPerQuery = numChunksPerQuery;
this.keywordExtractTemplate = keywordExtractTemplate || defaultKeywordExtractPrompt;
this.queryKeywordExtractTemplate = queryKeywordExtractTemplate || defaultQueryKeywordExtractPrompt;
}
async _retrieve(query) {
const keywords = await this.getKeywords(extractText(query));
const chunkIndicesCount = {};
const filteredKeywords = keywords.filter((keyword)=>this.indexStruct.table.has(keyword));
for (const keyword of filteredKeywords){
for (const nodeId of this.indexStruct.table.get(keyword) || []){
chunkIndicesCount[nodeId] = (chunkIndicesCount[nodeId] ?? 0) + 1;
}
}
const sortedChunkIndices = Object.keys(chunkIndicesCount).sort((a, b)=>chunkIndicesCount[b] - chunkIndicesCount[a]).slice(0, this.numChunksPerQuery);
const sortedNodes = await this.docstore.getNodes(sortedChunkIndices);
return sortedNodes.map((node)=>({
node
}));
}
}
// Extracts keywords using LLMs.
class KeywordTableLLMRetriever extends BaseKeywordTableRetriever {
async getKeywords(query) {
const response = await this.llm.complete({
prompt: this.queryKeywordExtractTemplate.format({
question: query,
maxKeywords: `${this.maxKeywordsPerQuery}`
})
});
const keywords = extractKeywordsGivenResponse(response.text, "KEYWORDS:");
return [
...keywords
];
}
}
// Extracts keywords using simple regex-based keyword extractor.
class KeywordTableSimpleRetriever extends BaseKeywordTableRetriever {
getKeywords(query) {
return Promise.resolve([
...simpleExtractKeywords(query, this.maxKeywordsPerQuery)
]);
}
}
// Extracts keywords using RAKE keyword extractor
class KeywordTableRAKERetriever extends BaseKeywordTableRetriever {
getKeywords(query) {
return Promise.resolve([
...rakeExtractKeywords(query, this.maxKeywordsPerQuery)
]);
}
}
const KeywordTableRetrieverMap = {
["DEFAULT"]: KeywordTableLLMRetriever,
["SIMPLE"]: KeywordTableSimpleRetriever,
["RAKE"]: KeywordTableRAKERetriever
};
/**
* The KeywordTableIndex, an index that extracts keywords from each Node and builds a mapping from each keyword to the corresponding Nodes of that keyword.
*/ class KeywordTableIndex extends BaseIndex {
constructor(init){
super(init);
}
static async init(options) {
const storageContext = options.storageContext ?? await storageContextFromDefaults({});
const { docStore, indexStore } = storageContext;
// Setup IndexStruct from storage
const indexStructs = await indexStore.getIndexStructs();
let indexStruct;
if (options.indexStruct && indexStructs.length > 0) {
throw new Error("Cannot initialize index with both indexStruct and indexStore");
}
if (options.indexStruct) {
indexStruct = options.indexStruct;
} else if (indexStructs.length == 1) {
indexStruct = indexStructs[0];
} else if (indexStructs.length > 1 && options.indexId) {
indexStruct = await indexStore.getIndexStruct(options.indexId);
} else {
indexStruct = null;
}
// check indexStruct type
if (indexStruct && indexStruct.type !== IndexStructType.KEYWORD_TABLE) {
throw new Error("Attempting to initialize KeywordTableIndex with non-keyword table indexStruct");
}
if (indexStruct) {
if (options.nodes) {
throw new Error("Cannot initialize KeywordTableIndex with both nodes and indexStruct");
}
} else {
if (!options.nodes) {
throw new Error("Cannot initialize KeywordTableIndex without nodes or indexStruct");
}
indexStruct = await KeywordTableIndex.buildIndexFromNodes(options.nodes, storageContext.docStore);
await indexStore.addIndexStruct(indexStruct);
}
return new KeywordTableIndex({
storageContext,
docStore,
indexStore,
indexStruct
});
}
// eslint-disable-next-line @typescript-eslint/no-explicit-any
asRetriever(options) {
const { mode = "DEFAULT", ...otherOptions } = options ?? {};
const KeywordTableRetriever = KeywordTableRetrieverMap[mode];
if (KeywordTableRetriever) {
return new KeywordTableRetriever({
index: this,
...otherOptions
});
}
throw new Error(`Unknown retriever mode: ${mode}`);
}
asQueryEngine(options) {
const { retriever, responseSynthesizer } = options ?? {};
return new RetrieverQueryEngine(retriever ?? this.asRetriever(), responseSynthesizer, options?.nodePostprocessors);
}
asChatEngine(options) {
const { retriever, ...contextChatEngineOptions } = options ?? {};
return new ContextChatEngine({
retriever: retriever ?? this.asRetriever(),
...contextChatEngineOptions
});
}
static async extractKeywords(text) {
const llm = Settings.llm;
const response = await llm.complete({
prompt: defaultKeywordExtractPrompt.format({
context: text
})
});
return extractKeywordsGivenResponse(response.text, "KEYWORDS:");
}
/**
* High level API: split documents, get keywords, and build index.
* @param documents
* @param args
* @param args.storageContext
* @returns
*/ static async fromDocuments(documents, args = {}) {
let { storageContext } = args;
storageContext = storageContext ?? await storageContextFromDefaults({});
const docStore = storageContext.docStore;
await docStore.addDocuments(documents, true);
for (const doc of documents){
await docStore.setDocumentHash(doc.id_, doc.hash);
}
const nodes = await Settings.nodeParser.getNodesFromDocuments(documents);
const index = await KeywordTableIndex.init({
nodes,
storageContext
});
return index;
}
/**
* Get keywords for nodes and place them into the index.
* @param nodes
* @param docStore
* @returns
*/ static async buildIndexFromNodes(nodes, docStore) {
const indexStruct = new KeywordTable();
await docStore.addDocuments(nodes, true);
for (const node of nodes){
const keywords = await KeywordTableIndex.extractKeywords(node.getContent(MetadataMode.LLM));
indexStruct.addNode([
...keywords
], node.id_);
}
return indexStruct;
}
async insertNodes(nodes) {
for (const node of nodes){
const keywords = await KeywordTableIndex.extractKeywords(node.getContent(MetadataMode.LLM));
this.indexStruct.addNode([
...keywords
], node.id_);
}
}
deleteNode(nodeId) {
const keywordsToDelete = new Set();
for (const [keyword, existingNodeIds] of Object.entries(this.indexStruct.table)){
const index = existingNodeIds.indexOf(nodeId);
if (index !== -1) {
existingNodeIds.splice(index, 1);
// Delete keywords that have zero nodes
if (existingNodeIds.length === 0) {
keywordsToDelete.add(keyword);
}
}
}
this.indexStruct.deleteNode([
...keywordsToDelete
], nodeId);
}
async deleteNodes(nodeIds, deleteFromDocStore) {
nodeIds.forEach((nodeId)=>{
this.deleteNode(nodeId);
});
if (deleteFromDocStore) {
for (const nodeId of nodeIds){
await this.docStore.deleteDocument(nodeId, false);
}
}
await this.storageContext.indexStore.addIndexStruct(this.indexStruct);
}
async deleteRefDoc(refDocId, deleteFromDocStore) {
const refDocInfo = await this.docStore.getRefDocInfo(refDocId);
if (!refDocInfo) {
return;
}
await this.deleteNodes(refDocInfo.nodeIds, false);
if (deleteFromDocStore) {
await this.docStore.deleteRefDoc(refDocId, false);
}
return;
}
}
const defaultFormatNodeBatchFn = (summaryNodes)=>{
return summaryNodes.map((node, idx)=>{
return `
Document ${idx + 1}:
${node.getContent(MetadataMode.LLM)}
`.trim();
}).join("\n\n");
};
const defaultParseChoiceSelectAnswerFn = (answer, numChoices, raiseErr = false)=>{
// split the line into the answer number and relevance score portions
const lineTokens = answer.split("\n").map((line)=>{
const lineTokens = line.split(",");
if (lineTokens.length !== 2) {
if (raiseErr) {
throw new Error(`Invalid answer line: ${line}. Answer line must be of the form: answer_num: <int>, answer_relevance: <float>`);
} else {
return null;
}
}
return lineTokens;
}).filter((lineTokens)=>!_.isNil(lineTokens));
// parse the answer number and relevance score
return lineTokens.reduce((parseResult, lineToken)=>{
try {
const docNum = parseInt(lineToken[0].split(":")[1].trim());
const answerRelevance = parseFloat(lineToken[1].split(":")[1].trim());
if (docNum < 1 || docNum > numChoices) {
if (raiseErr) {
throw new Error(`Invalid answer number: ${docNum}. Answer number must be between 1 and ${numChoices}`);
}
} else {
parseResult[docNum] = answerRelevance;
}
} catch (e) {
if (raiseErr) {
throw e;
}
}
return parseResult;
}, {});
};
var SummaryRetrieverMode = /*#__PURE__*/ function(SummaryRetrieverMode) {
SummaryRetrieverMode["DEFAULT"] = "default";
// EMBEDDING = "embedding",
SummaryRetrieverMode["LLM"] = "llm";
return SummaryRetrieverMode;
}({});
/**
* A SummaryIndex keeps nodes in a sequential order for use with summarization.
*/ class SummaryIndex extends BaseIndex {
constructor(init){
super(init);
}
static async init(options) {
const storageContext = options.storageContext ?? await storageContextFromDefaults({});
const { docStore, indexStore } = storageContext;
// Setup IndexStruct from storage
const indexStructs = await indexStore.getIndexStructs();
let indexStruct;
if (options.indexStruct && indexStructs.length > 0) {
throw new Error("Cannot initialize index with both indexStruct and indexStore");
}
if (options.indexStruct) {
indexStruct = options.indexStruct;
} else if (indexStructs.length == 1) {
indexStruct = indexStructs[0].type === IndexStructType.LIST ? indexStructs[0] : null;
} else if (indexStructs.length > 1 && options.indexId) {
indexStruct = await indexStore.getIndexStruct(options.indexId);
} else {
indexStruct = null;
}
// check indexStruct type
if (indexStruct && indexStruct.type !== IndexStructType.LIST) {
throw new Error("Attempting to initialize SummaryIndex with non-list indexStruct");
}
if (indexStruct) {
if (options.nodes) {
throw new Error("Cannot initialize SummaryIndex with both nodes and indexStruct");
}
} else {
if (!options.nodes) {
throw new Error("Cannot initialize SummaryIndex without nodes or indexStruct");
}
indexStruct = await SummaryIndex.buildIndexFromNodes(options.nodes, storageContext.docStore);
await indexStore.addIndexStruct(indexStruct);
}
return new SummaryIndex({
storageContext,
docStore,
indexStore,
indexStruct
});
}
static async fromDocuments(documents, args = {}) {
let { storageContext } = args;
storageContext = storageContext ?? await storageContextFromDefaults({});
const docStore = storageContext.docStore;
await docStore.addDocuments(documents, true);
for (const doc of documents){
await docStore.setDocumentHash(doc.id_, doc.hash);
}
const nodes = await Settings.nodeParser.getNodesFromDocuments(documents);
const index = await SummaryIndex.init({
nodes,
storageContext
});
return index;
}
asRetriever(options) {
const { mode = "default" } = options ?? {};
switch(mode){
case "default":
return new SummaryIndexRetriever(this);
case "llm":
return new SummaryIndexLLMRetriever(this);
default:
throw new Error(`Unknown retriever mode: ${mode}`);
}
}
asQueryEngine(options) {
let { retriever, responseSynthesizer } = options ?? {};
if (!retriever) {
retriever = this.asRetriever();
}
if (!responseSynthesizer) {
responseSynthesizer = getResponseSynthesizer("compact");
}
return new RetrieverQueryEngine(retriever, responseSynthesizer, options?.nodePostprocessors);
}
asChatEngine(options) {
const { retriever, mode, ...contextChatEngineOptions } = options ?? {};
return new ContextChatEngine({
retriever: retriever ?? this.asRetriever({
mode: mode ?? "default"
}),
...contextChatEngineOptions
});
}
static async buildIndexFromNodes(nodes, docStore, indexStruct) {
indexStruct = indexStruct || new IndexList();
await docStore.addDocuments(nodes, true);
for (const node of nodes){
indexStruct.addNode(node);
}
return indexStruct;
}
async insertNodes(nodes) {
for (const node of nodes){
this.indexStruct.addNode(node);
}
}
async deleteRefDoc(refDocId, deleteFromDocStore) {
const refDocInfo = await this.docStore.getRefDocInfo(refDocId);
if (!refDocInfo) {
return;
}
await this.deleteNodes(refDocInfo.nodeIds, false);
if (deleteFromDocStore) {
await this.docStore.deleteRefDoc(refDocId, false);
}
return;
}
async deleteNodes(nodeIds, deleteFromDocStore) {
this.indexStruct.nodes = this.indexStruct.nodes.filter((existingNodeId)=>!nodeIds.includes(existingNodeId));
if (deleteFromDocStore) {
for (const nodeId of nodeIds){
await this.docStore.deleteDocument(nodeId, false);
}
}
await this.storageContext.indexStore.addIndexStruct(this.indexStruct);
}
async getRefDocInfo() {
const nodeDocIds = this.indexStruct.nodes;
const nodes = await this.docStore.getNodes(nodeDocIds);
const refDocInfoMap = {};
for (const node of nodes){
const refNode = node.sourceNode;
if (_.isNil(refNode)) {
continue;
}
const refDocInfo = await this.docStore.getRefDocInfo(refNode.nodeId);
if (_.isNil(refDocInfo)) {
continue;
}
refDocInfoMap[refNode.nodeId] = refDocInfo;
}
return refDocInfoMap;
}
}
/**
* Simple retriever for SummaryIndex that returns all nodes
*/ class SummaryIndexRetriever extends BaseRetriever {
constructor(index){
super();
this.index = index;
}
async _retrieve(queryBundle) {
const nodeIds = this.index.indexStruct.nodes;
const nodes = await this.index.docStore.getNodes(nodeIds);
return nodes.map((node)=>({
node: node,
score: 1
}));
}
}
/**
* LLM retriever for SummaryIndex which lets you select the most relevant chunks.
*/ class SummaryIndexLLMRetriever extends BaseRetriever {
constructor(index, choiceSelectPrompt, choiceBatchSize = 10, formatNodeBatchFn, parseChoiceSelectAnswerFn){
super();
this.index = index;
this.choiceSelectPrompt = choiceSelectPrompt || defaultChoiceSelectPrompt;
this.choiceBatchSize = choiceBatchSize;
this.formatNodeBatchFn = formatNodeBatchFn || defaultFormatNodeBatchFn;
this.parseChoiceSelectAnswerFn = parseChoiceSelectAnswerFn || defaultParseChoiceSelectAnswerFn;
}
async _retrieve(query) {
const nodeIds = this.index.indexStruct.nodes;
const results = [];
for(let idx = 0; idx < nodeIds.length; idx += this.choiceBatchSize){
const nodeIdsBatch = nodeIds.slice(idx, idx + this.choiceBatchSize);
const nodesBatch = await this.index.docStore.getNodes(nodeIdsBatch);
const fmtBatchStr = this.formatNodeBatchFn(nodesBatch);
const input = {
context: fmtBatchStr,
query: extractText(query)
};
const llm = Settings.llm;
const rawResponse = (await llm.complete({
prompt: this.choiceSelectPrompt.format(input)
})).text;
// parseResult is a map from doc number to relevance score
const parseResult = this.parseChoiceSelectAnswerFn(rawResponse, nodesBatch.length);
const choiceNodeIds = nodeIdsBatch.filter((nodeId, idx)=>{
return `${idx}` in parseResult;
});
const choiceNodes = await this.index.docStore.getNodes(choiceNodeIds);
const nodeWithScores = choiceNodes.map((node, i)=>({
node: node,
score: _.get(parseResult, `${i + 1}`, 1)
}));
results.push(...nodeWithScores);
}
return results;
}
}
/**
* The VectorStoreIndex, an index that stores the nodes only according to their vector embeddings.
*/ class VectorStoreIndex extends BaseIndex {
constructor(init){
super(init);
this.indexStore = init.indexStore;
this.vectorStores = init.vectorStores ?? init.storageContext.vectorStores;
this.embedModel = Settings.embedModel;
}
/**
* The async init function creates a new VectorStoreIndex.
* @param options
* @returns
*/ static async init(options) {
const storageContext = options.storageContext ?? await storageContextFromDefaults({});
const indexStore = storageContext.indexStore;
const docStore = storageContext.docStore;
let indexStruct = await VectorStoreIndex.setupIndexStructFromStorage(indexStore, options);
if (!options.nodes && !indexStruct) {
throw new Error("Cannot initialize VectorStoreIndex without nodes or indexStruct");
}
indexStruct = indexStruct ?? new IndexDict();
const index = new this({
storageContext,
docStore,
indexStruct,
indexStore,
vectorStores: options.vectorStores
});
if (options.nodes) {
// If nodes are passed in, then we need to update the index
await index.buildIndexFromNodes(options.nodes, {
logProgress: options.logProgress,
progressCallback: options.progressCallback
});
}
return index;
}
static async setupIndexStructFromStorage(indexStore, options) {
const indexStructs = await indexStore.getIndexStructs();
let indexStruct;
if (options.indexStruct && indexStructs.length > 0) {
throw new Error("Cannot initialize index with both indexStruct and indexStore");
}
if (options.indexStruct) {
indexStruct = options.indexStruct;
} else if (indexStructs.length == 1) {
indexStruct = indexStructs[0].type === IndexStructType.SIMPLE_DICT ? indexStructs[0] : undefined;
indexStruct = indexStructs[0];
} else if (indexStructs.length > 1 && options.indexId) {
indexStruct = await indexStore.getIndexStruct(options.indexId);
}
// Check indexStruct type
if (indexStruct && indexStruct.type !== IndexStructType.SIMPLE_DICT) {
throw new Error("Attempting to initialize VectorStoreIndex with non-vector indexStruct");
}
return indexStruct;
}
/**
* Calculates the embeddings for the given nodes.
*
* @param nodes - An array of BaseNode objects representing the nodes for which embeddings are to be calculated.
* @param {Object} [options] - An optional object containing additional parameters.
* @param {boolean} [options.logProgress] - A boolean indicating whether to log progress to the console (useful for debugging).
*/ async getNodeEmbeddingResults(nodes, options) {
const nodeMap = splitNodesByType(nodes);
for(const type in nodeMap){
const nodes = nodeMap[type];
const embedModel = this.vectorStores[type]?.embedModel ?? this.embedModel;
if (embedModel && nodes) {
await embedModel(nodes, {
logProgress: options?.logProgress,
progressCallback: options?.progressCallback
});
}
}
return nodes;
}
/**
* Get embeddings for nodes and place them into the index.
* @param nodes
* @returns
*/ async buildIndexFromNodes(nodes, options) {
await this.insertNodes(nodes, options);
}
/**
* High level API: split documents, get embeddings, and build index.
* @param documents
* @param args
* @returns
*/ static async fromDocuments(documents, args = {}) {
args.storageContext = args.storageContext ?? await storageContextFromDefaults({});
args.vectorStores = args.vectorStores ?? args.storageContext.vectorStores;
args.docStoreStrategy = args.docStoreStrategy ?? // set doc store strategy defaults to the same as for the IngestionPipeline
(args.vectorStores ? DocStoreStrategy.UPSERTS : DocStoreStrategy.DUPLICATES_ONLY);
const docStore = args.storageContext.docStore;
if (args.logProgress) {
console.log("Using node parser on documents...");
}
// use doc store strategy to avoid duplicates
const vectorStores = Object.values(args.vectorStores ?? {});
const docStoreStrategy = createDocStoreStrategy(args.docStoreStrategy, docStore, vectorStores);
args.nodes = await runTransformations(documents, [
Settings.nodeParser
], {}, {
docStoreStrategy
});
if (args.logProgress) {
console.log("Finished parsing documents.");
}
try {
return await this.init(args);
} catch (error) {
await docStoreStrategy.rollback(args.storageContext.docStore, args.nodes);
throw error;
}
}
static async fromVectorStores(vectorStores) {
if (!vectorStores[ModalityType.TEXT]?.storesText) {
throw new Error("Cannot initialize from a vector store that does not store text");
}
const storageContext = await storageContextFromDefaults({
vectorStores
});
const index = await this.init({
nodes: [],
storageContext
});
return index;
}
static async fromVectorStore(vectorStore) {
return this.fromVectorStores({
[ModalityType.TEXT]: vectorStore
});
}
asRetriever(options) {
return new VectorIndexRetriever({
index: this,
...options
});
}
/**
* Create a RetrieverQueryEngine.
* similarityTopK is only used if no existing retriever is provided.
*/ asQueryEngine(options) {
const { retriever, responseSynthesizer, preFilters, customParams, nodePostprocessors, similarityTopK } = options ?? {};
return new RetrieverQueryEngine(retriever ?? this.asRetriever({
similarityTopK,
filters: preFilters,
customParams: customParams ?? undefined
}), responseSynthesizer, nodePostprocessors);
}
/**
* Convert the index to a chat engine.
* @param options The options for creating the chat engine
* @returns A ContextChatEngine that uses the index's retriever to get context for each query
*/ asChatEngine(options = {}) {
const { retriever, similarityTopK, preFilters, customParams, ...contextChatEngineOptions } = options;
return new ContextChatEngine({
retriever: retriever ?? this.asRetriever({
similarityTopK,
filters: preFilters,
customParams: customParams ?? undefined
}),
...contextChatEngineOptions
});
}
async insertNodesToStore(newIds, nodes, vectorStore) {
// NOTE: if the vector store doesn't store text,
// we need to add the nodes to the index struct and document store
// NOTE: if the vector store keeps text,
// we only need to add image and index nodes
for(let i = 0; i < nodes.length; ++i){
const { type } = nodes[i];
if (!vectorStore.storesText || type === ObjectType.INDEX || type === ObjectType.IMAGE) {
const nodeWithoutEmbedding = nodes[i].clone();
nodeWithoutEmbedding.embedding = undefined;
this.indexStruct.addNode(nodeWithoutEmbedding, newIds[i]);
await this.docStore.addDocuments([
nodeWithoutEmbedding
], true);
}
}
}
async insertNodes(nodes, options) {
if (!nodes || nodes.length === 0) {
return;
}
nodes = await this.getNodeEmbeddingResults(nodes, options);
await addNodesToVectorStores(nodes, this.vectorStores, this.insertNodesToStore.bind(this));
await this.indexStore.addIndexStruct(this.indexStruct);
}
async deleteRefDoc(refDocId, deleteFromDocStore = true) {
for (const vectorStore of Object.values(this.vectorStores)){
await this.deleteRefDocFromStore(vectorStore, refDocId);
}
if (deleteFromDocStore) {
await this.docStore.deleteDocument(refDocId, false);
}
}
async deleteRefDocFromStore(vectorStore, refDocId) {
await vectorStore.delete(refDocId);
if (!vectorStore.storesText) {
const refDocInfo = await this.docStore.getRefDocInfo(refDocId);
if (refDocInfo) {
for (const nodeId of refDocInfo.nodeIds){
this.indexStruct.delete(nodeId);
await vectorStore.delete(nodeId);
}
}
await this.indexStore.addIndexStruct(this.indexStruct);
}
}
}
class VectorIndexRetriever extends BaseRetriever {
constructor(options){
super();
this.index = options.index;
this.queryMode = options.mode ?? VectorStoreQueryMode.DEFAULT;
if ("topK" in options && options.topK) {
this.topK = options.topK;
} else {
this.topK = {
[ModalityType.TEXT]: "similarityTopK" in options && options.similarityTopK ? options.similarityTopK : DEFAULT_SIMILARITY_TOP_K,
[ModalityType.IMAGE]: DEFAULT_SIMILARITY_TOP_K,
[ModalityType.AUDIO]: DEFAULT_SIMILARITY_TOP_K
};
}
this.filters = options.filters;
this.customParams = options.customParams;
}
/**
* @deprecated, pass similarityTopK or topK in constructor instead or directly modify topK
*/ set similarityTopK(similarityTopK) {
this.topK[ModalityType.TEXT] = similarityTopK;
}
async _retrieve(params) {
const { query } = params;
const vectorStores = this.index.vectorStores;
let nodesWithScores = [];
for(const type in vectorStores){
const vectorStore = vectorStores[type];
nodesWithScores = nodesWithScores.concat(await this.retrieveQuery(query, type, vectorStore));
}
return nodesWithScores;
}
async retrieveQuery(query, type, vectorStore, filters, customParams) {
// convert string message to multi-modal format
let queryStr = query;
if (typeof query === "string") {
queryStr = query;
query = [
{
type: "text",
text: queryStr
}
];
} else {
queryStr = extractText(query);
}
// overwrite embed model if specified, otherwise use the one from the vector store
const embedModel = this.index.embedModel ?? vectorStore.embedModel;
let nodes = [];
// query each content item (e.g. text or image) separately
for (const item of query){
const queryEmbedding = await embedModel.getQueryEmbedding(item);
if (queryEmbedding) {
const result = await vectorStore.query({
queryStr,
queryEmbedding,
mode: this.queryMode ?? VectorStoreQueryMode.DEFAULT,
similarityTopK: this.topK[type],
filters: this.filters ?? filters ?? undefined,
customParams: this.customParams ?? customParams ?? undefined
});
nodes = nodes.concat(this.buildNodeListFromQueryResult(result));
}
}
return nodes;
}
buildNodeListFromQueryResult(result) {
const nodesWithScores = [];
for(let i = 0; i < result.ids.length; i++){
const nodeFromResult = result.nodes?.[i];
if (!this.index.indexStruct.nodesDict[result.ids[i]] && nodeFromResult) {
this.index.indexStruct.nodesDict[result.ids[i]] = nodeFromResult;
}
const node = this.index.indexStruct.nodesDict[result.ids[i]];
// XXX: Hack, if it's an image node, we reconstruct the image from the URL
// Alternative: Store image in doc store and retrieve it here
if (node instanceof ImageNode) {
node.image = node.getUrl();
}
nodesWithScores.push({
node: node,
score: result.similarities[i]
});
}
return nodesWithScores;
}
}
export { BaseIndex, KeywordTableIndex, KeywordTableLLMRetriever, KeywordTableRAKERetriever, KeywordTableRetrieverMode, KeywordTableSimpleRetriever, SummaryIndex, SummaryIndexLLMRetriever, SummaryIndexRetriever, SummaryRetrieverMode, VectorIndexRetriever, VectorStoreIndex, defaultFormatNodeBatchFn, defaultParseChoiceSelectAnswerFn, expandTokensWithSubtokens, extractKeywordsGivenResponse, rakeExtractKeywords, simpleExtractKeywords };