Add context filter to VectorPayload (#94)

* feat(database): add context filter to VectorPayload

* fix: add context length check to context filtering test

* fix: pass contextFilter to retrieve

---------

Co-authored-by: Ronaldo Lima <[email protected]>

src/context-service/index.ts

@@ -88,6 +88,7 @@ export class ContextService {
           similarityThreshold: optionsWithDefault.similarityThreshold,
           topK: optionsWithDefault.topK,
           namespace: optionsWithDefault.namespace,
+          contextFilter: optionsWithDefault.contextFilter,
         });
 
         // Log the result, which will be captured by the outer traceable

src/database.ts

@@ -73,6 +73,7 @@ export type VectorPayload = {
   similarityThreshold?: number;
   topK?: number;
   namespace?: string;
+  contextFilter?: string;
 };
 
 export type ResetOptions = {
@@ -106,6 +107,7 @@ export class Database {
     similarityThreshold = DEFAULT_SIMILARITY_THRESHOLD,
     topK = DEFAULT_TOP_K,
     namespace,
+    contextFilter,
   }: VectorPayload): Promise<{ data: string; id: string; metadata: TMetadata }[]> {
     const index = this.index;
     const result = await index.query<Record<string, string>>(
@@ -114,6 +116,7 @@ export class Database {
         topK,
         includeData: true,
         includeMetadata: true,
+        ...(typeof contextFilter === "string" && { filter: contextFilter }),
       },
       { namespace }
     );

src/rag-chat.test.ts

@@ -1014,3 +1014,57 @@ describe("RAG Chat with non-embedding db", () => {
     expect(called).toBeTrue();
   });
 });
+
+describe("RAGChat - context filtering", () => {
+  const namespace = "context-filtering";
+  const vector = new Index({
+    token: process.env.UPSTASH_VECTOR_REST_TOKEN!,
+    url: process.env.UPSTASH_VECTOR_REST_URL!,
+  });
+
+  const ragChat = new RAGChat({
+    vector,
+    namespace,
+    streaming: true,
+    model: upstash("meta-llama/Meta-Llama-3-8B-Instruct"),
+  });
+
+  afterAll(async () => {
+    await vector.reset({ namespace });
+    await vector.deleteNamespace(namespace);
+  });
+
+  test(
+    "should return metadata",
+    async () => {
+      await ragChat.context.add({
+        type: "text",
+        data: "Tokyo is the Capital of Japan.",
+        options: { namespace, metadata: { unit: "Samurai" } },
+      });
+      await ragChat.context.add({
+        type: "text",
+        data: "Shakuhachi is a traditional wind instrument",
+        options: { namespace, metadata: { unit: "Shakuhachi" } },
+      });
+      await awaitUntilIndexed(vector);
+
+      const result = await ragChat.chat<{ unit: string }>("Where is the capital of Japan?", {
+        namespace,
+        topK: 5,
+        contextFilter: "unit = 'Samurai'",
+        onContextFetched(context) {
+          expect(context.length).toBe(1);
+          return context;
+        },
+      });
+
+      expect(result.metadata).toEqual([
+        {
+          unit: "Samurai",
+        },
+      ]);
+    },
+    { timeout: 30_000 }
+  );
+});

src/rag-chat.ts

@@ -291,6 +291,7 @@ export class RAGChat {
       promptFn: isRagDisabledAndPromptFunctionMissing
         ? DEFAULT_PROMPT_WITHOUT_RAG
         : (options?.promptFn ?? this.config.prompt),
+      contextFilter: options?.contextFilter ?? undefined,
     };
   }
 }

src/types.ts

@@ -85,6 +85,13 @@ export type ChatOptions = {
    * Must be provided if the Vector Database doesn't have default embeddings.
    */
   embedding?: number[];
+
+  /**
+   * Allows filtering metadata from the vector database.
+   * @example "population >= 1000000 AND geography.continent = 'Asia'"
+   * https://upstash.com/docs/vector/features/filtering#metadata-filtering
+   */
+  contextFilter?: string;
 } & CommonChatAndRAGOptions;
 
 export type PrepareChatResult = { data: string; id: string; metadata: unknown }[];