diff --git a/.changeset/clever-hats-change.md b/.changeset/clever-hats-change.md new file mode 100644 index 0000000000000000000000000000000000000000..aa9714915e007295b060e70945aaf9144aa820ea --- /dev/null +++ b/.changeset/clever-hats-change.md @@ -0,0 +1,5 @@ +--- +"llamaindex": patch +--- + +feat: support metadata filters for AstraDB diff --git a/examples/astradb/example.ts b/examples/astradb/example.ts index 982dbf9854a67c05a65797e09c20f4f4b15fa40f..e576dcacc20368857b6bf789473b60c921041c8a 100644 --- a/examples/astradb/example.ts +++ b/examples/astradb/example.ts @@ -1,6 +1,7 @@ import { AstraDBVectorStore, Document, + MetadataFilters, storageContextFromDefaults, VectorStoreIndex, } from "llamaindex"; @@ -42,8 +43,10 @@ async function main() { const index = await VectorStoreIndex.fromDocuments(docs, { storageContext: ctx, }); - - const queryEngine = index.asQueryEngine(); + const preFilters: MetadataFilters = { + filters: [{ key: "id", operator: "in", value: [123, 789] }], + }; // try changing the filters to see the different results + const queryEngine = index.asQueryEngine({ preFilters }); const response = await queryEngine.query({ query: "Describe AstraDB.", }); diff --git a/packages/llamaindex/src/vector-store/AstraDBVectorStore.ts b/packages/llamaindex/src/vector-store/AstraDBVectorStore.ts index 6eda34a3f7ace850463325fad4345d176a470055..6fd3fa124da88791ba95dd61d4b38ba9f2179d09 100644 --- a/packages/llamaindex/src/vector-store/AstraDBVectorStore.ts +++ b/packages/llamaindex/src/vector-store/AstraDBVectorStore.ts @@ -2,19 +2,29 @@ import { Collection, DataAPIClient, Db, + type Filter, type FindOptions, + type SomeDoc, } from "@datastax/astra-db-ts"; import type { BaseNode } from "@llamaindex/core/schema"; import { MetadataMode } from "@llamaindex/core/schema"; import { getEnv } from "@llamaindex/env"; import { + FilterCondition, + FilterOperator, VectorStoreBase, type IEmbedModel, + type MetadataFilter, + type MetadataFilters, type VectorStoreNoEmbedModel, type VectorStoreQuery, type VectorStoreQueryResult, } from "./types.js"; -import { metadataDictToNode, nodeToMetadata } from "./utils.js"; +import { + metadataDictToNode, + nodeToMetadata, + parseArrayValue, +} from "./utils.js"; export class AstraDBVectorStore extends VectorStoreBase @@ -183,12 +193,8 @@ export class AstraDBVectorStore } const collection = this.collection; - const filters: Record<string, any> = {}; - query.filters?.filters?.forEach((f) => { - filters[f.key] = f.value; - }); - - const cursor = await collection.find(filters, <FindOptions>{ + const astraFilter = this.toAstraFilter(query.filters); + const cursor = await collection.find(astraFilter, <FindOptions>{ ...options, sort: query.queryEmbedding ? { $vector: query.queryEmbedding } @@ -230,4 +236,39 @@ export class AstraDBVectorStore nodes, }; } + + private toAstraFilter(filters?: MetadataFilters): Filter<SomeDoc> { + if (!filters || filters.filters?.length === 0) return {}; + const condition = filters.condition ?? FilterCondition.AND; + const listFilter = filters.filters.map((f) => this.buildFilterItem(f)); + if (condition === FilterCondition.OR) return { $or: listFilter }; + if (condition === FilterCondition.AND) return { $and: listFilter }; + throw new Error(`Not supported filter condition: ${condition}`); + } + + private buildFilterItem(filter: MetadataFilter): Filter<SomeDoc> { + const { key, operator, value } = filter; + switch (operator) { + case FilterOperator.EQ: + return { [key]: value }; + case FilterOperator.NE: + return { [key]: { $ne: value } }; + case FilterOperator.GT: + return { [key]: { $gt: value } }; + case FilterOperator.LT: + return { [key]: { $lt: value } }; + case FilterOperator.GTE: + return { [key]: { $gte: value } }; + case FilterOperator.LTE: + return { [key]: { $lte: value } }; + case FilterOperator.IN: + return { [key]: { $in: parseArrayValue(value) } }; + case FilterOperator.NIN: + return { [key]: { $nin: parseArrayValue(value) } }; + case FilterOperator.IS_EMPTY: + return { [key]: { $size: 0 } }; + default: + throw new Error(`Not supported filter operator: ${operator}`); + } + } }