DX-2255: Allow Metadata in Filter

src/client/metadata.ts

@@ -46,20 +46,25 @@ type ValidOperations<T> = T extends number
         ? MutuallyExclusives<ArrayOperation, any>
         : never;
 
+// Merge TContent and TMetadata, prefixing metadata keys with @metadata.
+type MergedFields<TContent, TMetadata> = TContent & {
+  [K in keyof TMetadata as `@metadata.${string & K}`]: TMetadata[K];
+};
+
 // Type definitions for FilterTree with strict operations
 // A leaf must have exactly one field with exactly one operation
-type Leaf<TContent> = {
-  [Field in keyof TContent]: {
-    [K in Field]: ValidOperations<TContent[K]>;
+type Leaf<TFields> = {
+  [Field in keyof TFields]: {
+    [K in Field]: ValidOperations<TFields[K]>;
   } & {
-    [K in Exclude<keyof TContent, Field>]?: never;
+    [K in Exclude<keyof TFields, Field>]?: never;
   };
-}[keyof TContent];
+}[keyof TFields];
 
-export type TreeNode<TContent> =
-  | Leaf<TContent>
-  | { OR: TreeNode<TContent>[] }
-  | { AND: TreeNode<TContent>[] };
+export type TreeNode<TContent, TMetadata> =
+  | Leaf<MergedFields<TContent, TMetadata>>
+  | { OR: TreeNode<TContent, TMetadata>[] }
+  | { AND: TreeNode<TContent, TMetadata>[] };
 
 const valueFormatter = (value: string | boolean | number | any[]): string | number | boolean => {
   return Array.isArray(value)
@@ -70,16 +75,18 @@ const valueFormatter = (value: string | boolean | number | any[]): string | numb
 };
 
 // Recursive function to construct filter string from FilterTree
-export function constructFilterString<TContent>(filterTree: TreeNode<TContent>): string {
+export function constructFilterString<TContent, TMetadata>(
+  filterTree: TreeNode<TContent, TMetadata>
+): string {
   if ("OR" in filterTree) {
-    return `(${filterTree.OR.map((node: TreeNode<TContent>) => constructFilterString(node)).join(" OR ")})`;
+    return `(${filterTree.OR.map((node: TreeNode<TContent, TMetadata>) => constructFilterString(node)).join(" OR ")})`;
   }
   if ("AND" in filterTree) {
-    return `(${filterTree.AND.map((node: TreeNode<TContent>) => constructFilterString(node)).join(" AND ")})`;
+    return `(${filterTree.AND.map((node: TreeNode<TContent, TMetadata>) => constructFilterString(node)).join(" AND ")})`;
   }
 
-  const field = Object.keys(filterTree)[0] as keyof TContent;
-  const operationObj = (filterTree as Leaf<TContent>)[field];
+  const field = Object.keys(filterTree)[0];
+  const operationObj = (filterTree as Record<string, any>)[field];
   const operation = Object.keys(operationObj)[0];
   const value = operationObj[operation as keyof typeof operationObj];
 

src/search-index.test.ts

@@ -3,7 +3,7 @@ import { Search } from "./platforms/nodejs";
 
 const client = Search.fromEnv();
 const indexName = "test-index-name";
-const searchIndex = client.index<{ text: string }, { key: string }>(indexName);
+const searchIndex = client.index<{ text: string }, { key: string; count: number }>(indexName);
 
 describe("SearchIndex", () => {
   beforeEach(async () => {
@@ -12,12 +12,12 @@ describe("SearchIndex", () => {
 
     // Insert test data
     await searchIndex.upsert([
-      { id: "id1", content: { text: "test-data-1" }, metadata: { key: "value1" } },
-      { id: "id2", content: { text: "test-data-2" }, metadata: { key: "value2" } },
+      { id: "id1", content: { text: "test-data-1" }, metadata: { key: "value1", count: 1 } },
+      { id: "id2", content: { text: "test-data-2" }, metadata: { key: "value2", count: 2 } },
       {
         id: "different-id3",
         content: { text: "different-test-data-3" },
-        metadata: { key: "value3" },
+        metadata: { key: "value3", count: 3 },
       },
     ]);
 
@@ -42,8 +42,8 @@ describe("SearchIndex", () => {
     const results = await searchIndex.fetch({ ids: ["id1", "id2"] });
 
     expect(results).toEqual([
-      { id: "id1", content: { text: "test-data-1" }, metadata: { key: "value1" } },
-      { id: "id2", content: { text: "test-data-2" }, metadata: { key: "value2" } },
+      { id: "id1", content: { text: "test-data-1" }, metadata: { key: "value1", count: 1 } },
+      { id: "id2", content: { text: "test-data-2" }, metadata: { key: "value2", count: 2 } },
     ]);
   });
 
@@ -52,21 +52,22 @@ describe("SearchIndex", () => {
       query: "test-data-1",
       limit: 2,
       filter: "text GLOB 'test*'",
-      keepOriginalQueryAfterEnrichment: true
+      keepOriginalQueryAfterEnrichment: true,
     });
 
     expect(results).toEqual([
       {
         id: "id1",
         content: { text: "test-data-1" },
-        metadata: { key: "value1" },
+        metadata: { key: "value1", count: 1 },
         score: expect.any(Number),
       },
 
       {
         content: { text: "test-data-2" },
         metadata: {
           key: "value2",
+          count: 2,
         },
         id: "id2",
         score: expect.any(Number),
@@ -87,7 +88,28 @@ describe("SearchIndex", () => {
       {
         id: "id1",
         content: { text: "test-data-1" },
-        metadata: { key: "value1" },
+        metadata: { key: "value1", count: 1 },
+        score: expect.any(Number),
+      },
+    ]);
+  });
+
+  test("should search with a metadata filter", async () => {
+    const results = await searchIndex.search({
+      query: "test-data",
+      limit: 2,
+      filter: {
+        AND: [{ text: { glob: "*test-data*" } }, { "@metadata.count": { greaterThanOrEquals: 3 } }],
+      },
+      semanticWeight: 0.5,
+      inputEnrichment: false,
+    });
+
+    expect(results).toEqual([
+      {
+        id: "different-id3",
+        content: { text: "different-test-data-3" },
+        metadata: { key: "value3", count: 3 },
         score: expect.any(Number),
       },
     ]);
@@ -125,7 +147,7 @@ describe("SearchIndex", () => {
     });
 
     expect(documents).toEqual([
-      { id: "id1", content: { text: "test-data-1" }, metadata: { key: "value1" } },
+      { id: "id1", content: { text: "test-data-1" }, metadata: { key: "value1", count: 1 } },
     ]);
 
     const { documents: nextDocuments } = await searchIndex.range({
@@ -139,6 +161,7 @@ describe("SearchIndex", () => {
         content: { text: "test-data-2" },
         metadata: {
           key: "value2",
+          count: 2,
         },
         id: "id2",
       },

src/search-index.ts

@@ -76,13 +76,21 @@ export class SearchIndex<TContent extends Dict = Dict, TIndexMetadata extends Di
   search = async (params: {
     query: string;
     limit?: number;
-    filter?: string | TreeNode<TContent>;
+    filter?: string | TreeNode<TContent, TIndexMetadata>;
     reranking?: boolean;
     semanticWeight?: number;
     inputEnrichment?: boolean;
     keepOriginalQueryAfterEnrichment?: boolean;
   }): Promise<SearchResult<TContent, TIndexMetadata>> => {
-    const { query, limit = 5, filter, reranking, semanticWeight, inputEnrichment, keepOriginalQueryAfterEnrichment } = params;
+    const {
+      query,
+      limit = 5,
+      filter,
+      reranking,
+      semanticWeight,
+      inputEnrichment,
+      keepOriginalQueryAfterEnrichment,
+    } = params;
 
     if (semanticWeight && (semanticWeight < 0 || semanticWeight > 1)) {
       throw new UpstashError("semanticWeight must be between 0 and 1");