Skip to content
Snippets Groups Projects
Unverified Commit e9a111d9 authored by Marcus Schiesser's avatar Marcus Schiesser Committed by GitHub
Browse files

fix: VectorIndexRetrieverOptions typing (#1366)


Co-authored-by: default avatarAlex Yang <himself65@outlook.com>
parent bb7622e4
No related branches found
No related tags found
No related merge requests found
---
"llamaindex": patch
---
fix: VectorIndexRetrieverOptions typing
......@@ -13,5 +13,6 @@
},
"[json]": {
"editor.defaultFormatter": "esbenp.prettier-vscode"
}
},
"prettier.prettierPath": "./node_modules/prettier"
}
......@@ -274,7 +274,7 @@ export class VectorStoreIndex extends BaseIndex<IndexDict> {
}
asRetriever(
options?: Omit<VectorIndexRetrieverOptions, "index">,
options?: OmitIndex<VectorIndexRetrieverOptions>,
): VectorIndexRetriever {
return new VectorIndexRetriever({ index: this, ...options });
}
......@@ -382,12 +382,19 @@ export class VectorStoreIndex extends BaseIndex<IndexDict> {
type TopKMap = { [P in ModalityType]: number };
type OmitIndex<T> = T extends { index: any } ? Omit<T, "index"> : never;
export type VectorIndexRetrieverOptions = {
index: VectorStoreIndex;
similarityTopK?: number | undefined;
topK?: TopKMap | undefined;
filters?: MetadataFilters | undefined;
};
} & (
| {
topK?: TopKMap | undefined;
}
| {
similarityTopK?: number | undefined;
}
);
export class VectorIndexRetriever extends BaseRetriever {
index: VectorStoreIndex;
......@@ -396,20 +403,22 @@ export class VectorIndexRetriever extends BaseRetriever {
serviceContext?: ServiceContext | undefined;
filters?: MetadataFilters | undefined;
constructor({
index,
similarityTopK,
topK,
filters,
}: VectorIndexRetrieverOptions) {
constructor(options: VectorIndexRetrieverOptions) {
super();
this.index = index;
this.index = options.index;
this.serviceContext = this.index.serviceContext;
this.topK = topK ?? {
[ModalityType.TEXT]: similarityTopK ?? DEFAULT_SIMILARITY_TOP_K,
[ModalityType.IMAGE]: DEFAULT_SIMILARITY_TOP_K,
};
this.filters = filters;
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,
};
}
this.filters = options.filters;
}
/**
......
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