feat: support metadata filters for Astra (#1330)

Co-authored-by: Marcus Schiesser <mail@marcusschiesser.de>

.changeset/clever-hats-change.md

@@ -0,0 +1,5 @@
+---
+"llamaindex": patch
+---
+
+feat: support metadata filters for AstraDB

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.",
     });

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}`);
+    }
+  }
 }