@jakmro/filtering feature (#10)

* feat: result filtering based on predicate function

* chore: bump version

Author: jakmro

README.md

@@ -21,6 +21,7 @@ Private, local RAGs. Supercharge LLMs with your own knowledge base.
       - [Interfaces (for Custom Components)](#interfaces-for-custom-components)
       - [Text Splitters](#text-splitters)
       - [Utilities](#utilities)
+      - [Types](#types)
   - [:jigsaw: Using Custom Components](#jigsaw-using-custom-components)
   - [:electric_plug: Plugins](#electric_plug-plugins)
   - [:handshake: Contributing](#handshake-contributing)
@@ -272,11 +273,12 @@ The core class for managing the RAG workflow.
 
   * `async load(): Promise<this>`: Initializes the vector store and loads the LLM.
   * `async unload(): Promise<void>`: Unloads the vector store and LLM.
-  * `async generate(input: Message[] | string, options?: { augmentedGeneration?: boolean; k?: number; questionGenerator?: Function; promptGenerator?: Function; callback?: (token: string) => void }): Promise<string>` Generates a response.
+  * `async generate(input: Message[] | string, options?: { augmentedGeneration?: boolean; k?: number; predicate?: (value: SearchResult) => boolean; questionGenerator?: Function; promptGenerator?: Function; callback?: (token: string) => void }): Promise<string>` Generates a response.
     * `input` (`Message[] | string`): A string or an array of `Message` objects.
     * `options` (object, optional): Generation options.
       * `augmentedGeneration` (`boolean`, optional): If `true` (default), retrieves context from the vector store to augment the prompt.
       * `k` (`number`, optional): Number of documents to retrieve (default: `3`).
+      * `predicate` (`function`, optional): Function to filter retrieved documents (default: includes all).
       * `questionGenerator` (`function`, optional): Custom question generator.
       * `promptGenerator` (`function`, optional): Custom prompt generator.
       * `callback` (`function`, optional): A function that receives tokens as they are generated.
@@ -322,7 +324,7 @@ These interfaces define the contracts for creating your own custom components.
   * `add(document: string, metadata?: Record<string, any>): Promise<string>`: Adds a document.
   * `update(id: string, document?: string, metadata?: Record<string, any>): Promise<void>`: Updates a document.
   * `delete(id: string): Promise<void>`: Deletes a document.
-  * `similaritySearch(query: string, k?: number): Promise<{ id: string; content: string; ... }[]>`: Searches for `k` similar documents.
+  * `similaritySearch(query: string, k?: number, predicate?: (value: SearchResult) => boolean): Promise<SearchResult[]>`: Searches for `k` similar documents. Which can be filtered with an optional `predicate` function.
 
 #### `TextSplitter`
 
@@ -345,6 +347,24 @@ The library provides wrappers around common `langchain` text splitters. All spli
   * `dotProduct(a: number[], b: number[]): number`: Calculates the dot product of two vectors.
   * `magnitude(a: number[]): number`: Calculates the Euclidean magnitude of a vector.
 
+### Types
+
+```typescript
+interface Message {
+  role: 'user' | 'assistant' | 'system';
+  content: string;
+}
+
+type ResourceSource = string | number | object;
+
+interface SearchResult {
+  id: string;
+  content: string;
+  metadata?: Record<string, any>;
+  similarity: number;
+}
+```
+
 ## :jigsaw: Using Custom Components
 
 Bring your own components by creating classes that implement the `LLM`, `Embeddings`, `VectorStore` and `TextSplitter` interfaces. This allows you to use any model or service that fits your needs.
@@ -382,15 +402,9 @@ interface VectorStore {
   delete(id: string): Promise<void>;
   similaritySearch(
     query: string,
-    k?: number
-  ): Promise<
-    {
-      id: string;
-      content: string;
-      metadata?: Record<string, any>;
-      similarity: number;
-    }[]
-  >;
+    k?: number,
+    predicate?: (value: SearchResult) => boolean
+  ): Promise<SearchResult[]>;
 }
 ```
 

package.json

@@ -1,6 +1,6 @@
 {
   "name": "react-native-rag",
-  "version": "0.1.1",
+  "version": "0.1.2",
   "description": "Private, local RAGs. Supercharge LLMs with your own knowledge base.",
   "main": "./lib/module/index.js",
   "types": "./lib/typescript/src/index.d.ts",

packages/op-sqlite/README.md

@@ -94,12 +94,13 @@ The `OPSQLiteVectorStore` class implements the `VectorStore` interface from `rea
 
       * **Throws**: An error if the document with the specified `id` does not exist.
 
-  * #### `similaritySearch(query: string, k?: number): Promise<object[]>`
+  * #### `similaritySearch(query: string, k?: number, predicate?: (value: SearchResult) => boolean): Promise<SearchResult[]>`
 
     Finds the most relevant documents in the store based on a query string. It generates an embedding for the query and returns the top `k` documents that are most similar.
 
       * `query`: The text to search for.
       * `k`: The number of similar documents to return. Defaults to `3`.
+      * `predicate`: An optional function to filter the search results. It receives each result object and should return `true` to include the result or `false` to exclude it. If not provided, all results are included.
       * **Returns**: A `Promise` that resolves with an array of result objects, each containing `{ id, content, metadata, similarity }`. The `similarity` score ranges from 0 to 1, where 1 means a perfect match.
 
   * #### `deleteVectorStore(): Promise<void>`

packages/op-sqlite/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@react-native-rag/op-sqlite",
-  "version": "0.1.1",
+  "version": "0.1.2",
   "main": "src/index.ts",
   "scripts": {
     "test": "echo \"Error: no test specified\" && exit 1"

packages/op-sqlite/src/wrappers/op-sqlite.ts

@@ -1,4 +1,9 @@
-import { type VectorStore, type Embeddings, uuidv4 } from 'react-native-rag';
+import {
+  type VectorStore,
+  type Embeddings,
+  uuidv4,
+  type SearchResult,
+} from 'react-native-rag';
 import { open, type DB } from '@op-engineering/op-sqlite';
 
 export interface OPSQLiteVectorStoreParams {
@@ -107,27 +112,35 @@ export class OPSQLiteVectorStore implements VectorStore {
     }
   }
 
-  async similaritySearch(query: string, k: number = 3) {
+  async similaritySearch(
+    query: string,
+    k: number = 3,
+    predicate: (value: SearchResult) => boolean = () => true
+  ): Promise<SearchResult[]> {
     const queryEmbedding = await this.embeddings.embed(query);
     const vectorStr = `[${queryEmbedding.join(',')}]`;
-    const results = await this.db.execute(
+    const queryResult = await this.db.execute(
       `
         SELECT id, content, vector_distance_cos(embedding, vector(?)) as cosine_distance, metadata
         FROM vectors
-        ORDER BY cosine_distance ASC
-        LIMIT ?;
+        ORDER BY cosine_distance ASC;
       `,
-      [vectorStr, k]
+      [vectorStr]
     );
-    return results.rows.map((row) => {
-      const id = row.id as string;
-      const content = row.content as string;
-      const metadata = row.metadata
-        ? (JSON.parse(row.metadata as string) as Record<string, any>)
-        : undefined;
-      const similarity = 1 - (row.cosine_distance as number); // Convert cosine distance to similarity
-      return { id, content, metadata, similarity };
-    });
+    const results = Array.from<SearchResult>(
+      queryResult.rows.map((row: any) => {
+        const id = row.id as string;
+        const content = row.content as string;
+        const metadata = row.metadata
+          ? (JSON.parse(row.metadata as string) as Record<string, any>)
+          : undefined;
+        const similarity = 1 - (row.cosine_distance as number); // Convert cosine distance to similarity
+        return { id, content, metadata, similarity };
+      })
+    )
+      .filter(predicate)
+      .slice(0, k);
+    return results;
   }
 
   async deleteVectorStore() {

src/hooks/rag.ts

@@ -1,6 +1,6 @@
 import { useEffect, useState, useCallback, useMemo } from 'react';
 import { RAG, type RAGParams } from '../rag/rag';
-import type { Message } from '../types/common';
+import type { Message, SearchResult } from '../types/common';
 import type { TextSplitter } from '../interfaces/textSplitter';
 
 /**
@@ -75,7 +75,7 @@ export function useRAG({ vectorStore, llm, preventLoad }: UseRAGParams) {
   /**
    * Generates a text response.
    * @param {Message[] | string} input - Input messages or query.
-   * @param {object} [options={}] - Generation options (augmentedGeneration, k, questionGenerator, promptGenerator).
+   * @param {object} [options={}] - Generation options (augmentedGeneration, k, predicate, questionGenerator, promptGenerator).
    * @returns {Promise<string>} Complete generated string.
    * @throws {Error} If not ready or busy.
    */
@@ -85,6 +85,7 @@ export function useRAG({ vectorStore, llm, preventLoad }: UseRAGParams) {
       options: {
         augmentedGeneration?: boolean;
         k?: number;
+        predicate?: (value: SearchResult) => boolean;
         questionGenerator?: (messages: Message[]) => string;
         promptGenerator?: (
           messages: Message[],

src/interfaces/vectorStore.ts

@@ -1,3 +1,5 @@
+import type { SearchResult } from '../types/common';
+
 /**
  * Defines the essential operations for a vector store.
  * A vector store efficiently stores and retrieves high-dimensional vectors,
@@ -49,18 +51,13 @@ export interface VectorStore {
   /**
    * Performs a similarity search against the stored vectors.
    * @param query The query string to search for.
-   * @param k The number of top similar results to return. Defaults to 3.
+   * @param k The number of top similar results to return.
+   * @param predicate An optional function to filter results based on custom criteria.
    * @returns An array of objects containing the ID, content, metadata, and similarity score for each result.
    */
   similaritySearch(
     query: string,
-    k?: number
-  ): Promise<
-    {
-      id: string;
-      content: string;
-      metadata?: Record<string, any>;
-      similarity: number;
-    }[]
-  >;
+    k?: number,
+    predicate?: (value: SearchResult) => boolean
+  ): Promise<SearchResult[]>;
 }

src/rag/rag.ts

@@ -1,5 +1,5 @@
 import { RecursiveCharacterTextSplitter } from '../text-splitters/langchain';
-import type { Message } from '../types/common';
+import type { Message, SearchResult } from '../types/common';
 import type { LLM } from '../interfaces/llm';
 import type { TextSplitter } from '../interfaces/textSplitter';
 import type { VectorStore } from '../interfaces/vectorStore';
@@ -127,6 +127,7 @@ Context: ${retrievedDocs.map((doc) => doc.content).join('\n')}`;
    * @param options Optional parameters for generation:
    * - `augmentedGeneration`: Whether to use augmented generation with retrieved documents (default: true).
    * - `k`: Number of documents to retrieve (default: 3).
+   * - `predicate`: Function to filter retrieved documents (default: includes all).
    * - `questionGenerator`: Function to generate the question from messages for the document retrieval step. (default: uses last user message).
    * - `promptGenerator`: Function to generate the prompt from messages and retrieved documents (default: uses last user message and retrieved docs).
    * - `callback`: Optional callback function to handle token generation.
@@ -137,6 +138,7 @@ Context: ${retrievedDocs.map((doc) => doc.content).join('\n')}`;
     options: {
       augmentedGeneration?: boolean;
       k?: number;
+      predicate?: (value: SearchResult) => boolean;
       questionGenerator?: (messages: Message[]) => string;
       promptGenerator?: (
         messages: Message[],
@@ -155,6 +157,7 @@ Context: ${retrievedDocs.map((doc) => doc.content).join('\n')}`;
     const {
       augmentedGeneration = true,
       k = 3,
+      predicate = () => true,
       questionGenerator = this.questionGenerator,
       promptGenerator = this.promptGenerator,
       callback = () => {},
@@ -170,7 +173,8 @@ Context: ${retrievedDocs.map((doc) => doc.content).join('\n')}`;
     }
     const retrievedDocs = await this.vectorStore.similaritySearch(
       questionGenerator(input),
-      k
+      k,
+      predicate
     );
     const prompt = promptGenerator(input, retrievedDocs);
     const augmentedInput: Message[] = [

src/types/common.ts

@@ -21,3 +21,13 @@ export interface Message {
  * This type is flexible to cover different model loading strategies.
  */
 export type ResourceSource = string | number | object;
+
+/**
+ * Defines the structure of a search result from a vector store.
+ */
+export interface SearchResult {
+  id: string;
+  content: string;
+  metadata?: Record<string, any>;
+  similarity: number;
+}

src/vector_stores/memoryVectorStore.ts

@@ -1,5 +1,6 @@
 import type { Embeddings } from '../interfaces/embeddings';
 import type { VectorStore } from '../interfaces/vectorStore';
+import type { SearchResult } from '../types/common';
 import { uuidv4 } from '../utils/uuidv4';
 import { cosine } from '../utils/vectorMath';
 
@@ -86,24 +87,20 @@ export class MemoryVectorStore implements VectorStore {
 
   async similaritySearch(
     query: string,
-    k: number = 3
-  ): Promise<
-    {
-      id: string;
-      content: string;
-      metadata?: Record<string, any>;
-      similarity: number;
-    }[]
-  > {
+    k: number = 3,
+    predicate: (value: SearchResult) => boolean = () => true
+  ): Promise<SearchResult[]> {
     const queryEmbedding = await this.embeddings.embed(query);
-    const results = Array.from(this.memoryVectors.values()).map(
-      (memoryVector) => ({
+    const results = Array.from(this.memoryVectors.values())
+      .map((memoryVector) => ({
         id: memoryVector.id,
         content: memoryVector.content,
         metadata: memoryVector.metadata,
         similarity: cosine(queryEmbedding, memoryVector.embedding),
-      })
-    );
-    return results.sort((a, b) => b.similarity - a.similarity).slice(0, k);
+      }))
+      .filter(predicate)
+      .sort((a, b) => b.similarity - a.similarity)
+      .slice(0, k);
+    return results;
   }
 }