Add PrismaVectorStore filter IN operator (#3304)

* add Prisma filter IN operator * better naming * Update langchain/src/vectorstores/prisma.ts Co-authored-by: Brace Sproul <[email protected]> * use typeguard * add operators doc * add in operator in doc * typo * Update langchain/src/vectorstores/prisma.ts Co-authored-by: Brace Sproul <[email protected]> * Update examples/src/indexes/vector_stores/prisma_vectorstore/prisma.ts Co-authored-by: Brace Sproul <[email protected]> * prettier * Update prisma.mdx * Update prisma.mdx --------- Co-authored-by: Brace Sproul <[email protected]> Co-authored-by: Jacob Lee <[email protected]>
langchain-ai · Nov 17, 2023 · da38143 · da38143 · vercel · Nov 17, 2023
1 parent 6d91724
commit da38143
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 3 deletions.
diff --git a/docs/core_docs/docs/integrations/vectorstores/prisma.mdx b/docs/core_docs/docs/integrations/vectorstores/prisma.mdx
@@ -82,6 +82,8 @@ import Schema from "@examples/indexes/vector_stores/prisma_vectorstore/prisma/sc
 
 <CodeBlock language="typescript">{Example}</CodeBlock>
 
+The following SQL operators are available as filters: `equals`, `in`, `lt`, `lte`, `gt`, `gte`, `not`.
+
 The samples above uses the following schema:
 
 <CodeBlock language="prisma">{Schema}</CodeBlock>

diff --git a/examples/src/indexes/vector_stores/prisma_vectorstore/prisma.ts b/examples/src/indexes/vector_stores/prisma_vectorstore/prisma.ts
@@ -62,7 +62,10 @@ export const run = async () => {
   const resultThree = await vectorStore.similaritySearchWithScore(
     "Hello world",
     1,
-    { content: { equals: "different_content" } }
+    {
+      content: { equals: "different_content" },
+      customForeignKey: { in: ["<id1>", "<id2>", "<id3>"] },
+    }
   );
   console.log(resultThree);
 };
diff --git a/langchain/src/vectorstores/prisma.ts b/langchain/src/vectorstores/prisma.ts
@@ -61,6 +61,7 @@ type ModelColumns<TModel extends Record<string, unknown>> = {
 export type PrismaSqlFilter<TModel extends Record<string, unknown>> = {
   [K in keyof TModel]?: {
     equals?: TModel[K];
+    in?: TModel[K][];
     lt?: TModel[K];
     lte?: TModel[K];
     gt?: TModel[K];
@@ -71,6 +72,7 @@ export type PrismaSqlFilter<TModel extends Record<string, unknown>> = {
 
 const OpMap = {
   equals: "=",
+  in: "IN",
   lt: "<",
   lte: "<=",
   gt: ">",
@@ -410,9 +412,29 @@ export class PrismaVectorStore<
         Object.entries(ops).map(([opName, value]) => {
           // column name, operators cannot be parametrised
           // these fields are thus not escaped by Prisma and can be dangerous if user input is used
+          const opNameKey = opName as keyof typeof OpMap;
           const colRaw = this.Prisma.raw(`"${key}"`);
-          const opRaw = this.Prisma.raw(OpMap[opName as keyof typeof OpMap]);
-          return this.Prisma.sql`${colRaw} ${opRaw} ${value}`;
+          const opRaw = this.Prisma.raw(OpMap[opNameKey]);
+
+          switch (OpMap[opNameKey]) {
+            case OpMap.in: {
+              if (
+                !Array.isArray(value) ||
+                !value.every((v) => typeof v === "string")
+              ) {
+                throw new Error(
+                  `Invalid filter: IN operator requires an array of strings. Received: ${JSON.stringify(
+                    value,
+                    null,
+                    2
+                  )}`
+                );
+              }
+              return this.Prisma.sql`${colRaw} ${opRaw} (${value.join(",")})`;
+            }
+            default:
+              return this.Prisma.sql`${colRaw} ${opRaw} ${value}`;
+          }
         })
       ),
       " AND ",