{"version":3,"sources":["../src/index.ts"],"sourcesContent":["/**\n * Copyright 2024 Google LLC\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n * http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\nimport * as neo4j_driver from \"neo4j-driver\";\nimport { Genkit, z } from \"genkit\";\nimport { GenkitPlugin, genkitPlugin } from \"genkit/plugin\";\n\nimport { EmbedderArgument, Embedding } from \"genkit/embedder\";\nimport {\n CommonRetrieverOptionsSchema,\n Document,\n indexerRef,\n retrieverRef,\n} from \"genkit/retriever\";\n\nconst Neo4jRetrieverOptionsSchema = CommonRetrieverOptionsSchema.extend({\n k: z.number().max(1000),\n // filter: z.record(z.string(), z.any()).optional(), later for metadata filtering\n});\n\nconst Neo4jIndexerOptionsSchema = z.object({\n namespace: z.string().optional(),\n});\n\nexport interface Neo4jGraphConfig {\n url: string;\n username: string;\n password: string;\n database?: string;\n}\n\n/**\n * neo4jRetrieverRef function creates a retriever for Neo4j.\n * @param params The params for the new Neo4j retriever\n * @param params.indexId The indexId for the Neo4j retriever\n * @param params.displayName A display name for the retriever.\nIf not specified, the default label will be `Neo4j - `\n * @returns A reference to a Neo4j retriever.\n */\nexport const neo4jRetrieverRef = (params: {\n indexId: string;\n displayName?: string;\n}) => {\n return retrieverRef({\n name: `neo4j/${params.indexId}`,\n info: {\n label: params.displayName ?? `Neo4j - ${params.indexId}`,\n },\n configSchema: Neo4jRetrieverOptionsSchema,\n });\n};\n\n/**\n * neo4jIndexerRef function creates an indexer for Neo4j.\n * @param params The params for the new Neo4j indexer.\n * @param params.indexId The indexId for the Neo4j indexer.\n * @param params.displayName A display name for the indexer.\nIf not specified, the default label will be `Neo4j - `\n * @returns A reference to a Neo4j indexer.\n */\nexport const neo4jIndexerRef = (params: {\n indexId: string;\n displayName?: string;\n}) => {\n return indexerRef({\n name: `neo4j/${params.indexId}`,\n info: {\n label: params.displayName ?? `Neo4j - ${params.indexId}`,\n },\n //configSchema: Neo4jIndexerOptionsSchema.optional(),\n });\n};\n\n/**\n * Neo4j plugin that provides a Neo4j retriever and indexer\n * @param params An array of params to set up Neo4j retrievers and indexers\n * @param params.clientParams Neo4jConfiguration containing the\nusername, password, and url. If not set, the NEO4J_URI, NEO4J_USERNAME,\nand NEO4J_PASSWORD environment variable will be used instead.\n * @param params.indexId The name of the index\n * @param params.embedder The embedder to use for the indexer and retriever\n * @param params.embedderOptions Options to customize the embedder\n * @returns The Neo4j Genkit plugin\n */\nexport function neo4j(\n params: {\n clientParams?: Neo4jGraphConfig;\n indexId: string;\n embedder: EmbedderArgument;\n embedderOptions?: z.infer;\n }[],\n): GenkitPlugin {\n return genkitPlugin(\"neo4j\", async (ai: Genkit) => {\n params.map((i) => configureNeo4jRetriever(ai, i));\n params.map((i) => configureNeo4jIndexer(ai, i));\n });\n}\n\nexport default neo4j;\n\n/**\n * Configures a Neo4j retriever.\n * @param ai A Genkit instance\n * @param params The params for the retriever\n * @param params.indexId The name of the retriever\n * @param params.clientParams PNeo4jConfiguration containing the\nusername, password, and url. If not set, the NEO4J_URI, NEO4J_USERNAME,\nand NEO4J_PASSWORD environment variable will be used instead.\n * @param params.embedder The embedder to use for the retriever\n * @param params.embedderOptions Options to customize the embedder\n * @returns A Pinecone retriever\n */\nexport function configureNeo4jRetriever<\n EmbedderCustomOptions extends z.ZodTypeAny,\n>(\n ai: Genkit,\n params: {\n indexId: string;\n clientParams?: Neo4jGraphConfig;\n embedder: EmbedderArgument;\n embedderOptions?: z.infer;\n },\n) {\n const { indexId, embedder, embedderOptions } = {\n ...params,\n };\n const neo4jConfig = params.clientParams ?? getDefaultConfig();\n const neo4j_instance = neo4j_driver.driver(\n neo4jConfig.url, // URL (protocol://host:port)\n neo4j_driver.auth.basic(neo4jConfig.username, neo4jConfig.password), // Authentication\n );\n return ai.defineRetriever(\n {\n name: `neo4j/${params.indexId}`,\n configSchema: Neo4jRetrieverOptionsSchema,\n },\n async (content, options) => {\n const queryEmbeddings = await ai.embed({\n embedder,\n content,\n options: embedderOptions,\n });\n\n const retriever_query = `\n CALL db.index.vector.queryNodes($index, $k, $embedding) YIELD node, score\n RETURN node.text AS text, node {.*, text: Null,\n embedding: Null, id: Null } AS metadata\n `;\n const response = await neo4j_instance.executeQuery(\n retriever_query,\n {\n k: options.k,\n embedding: queryEmbeddings[0].embedding,\n index: indexId,\n },\n {\n database: neo4jConfig.database,\n },\n );\n // Create documents properly by returning the result from map\n const documents = response.records.map((el) => {\n return Document.fromText(\n el.get(\"text\"),\n Object.fromEntries(\n Object.entries(el.get(\"metadata\")).filter(\n ([_, value]) => value !== null,\n ),\n ),\n );\n });\n neo4j_instance.close();\n return { documents: documents };\n },\n );\n}\n\n/**\n * Configures a Neo4j indexer.\n * @param ai A Genkit instance\n * @param params The params for the indexer\n * @param params.indexId The name of the indexer\n * @param params.clientParams Neo4jConfiguration containing the\nusername, password, and url. If not set, the NEO4J_URI, NEO4J_USERNAME,\nand NEO4J_PASSWORD environment variable will be used instead.\n * @param params.embedder The embedder to use for the retriever\n * @param params.embedderOptions Options to customize the embedder\n * @returns A Genkit indexer\n */\nexport function configureNeo4jIndexer<\n EmbedderCustomOptions extends z.ZodTypeAny,\n>(\n ai: Genkit,\n params: {\n indexId: string;\n clientParams?: Neo4jGraphConfig;\n embedder: EmbedderArgument;\n embedderOptions?: z.infer;\n },\n) {\n const { indexId, embedder, embedderOptions } = {\n ...params,\n };\n const neo4jConfig = params.clientParams ?? getDefaultConfig();\n const neo4j_instance = neo4j_driver.driver(\n neo4jConfig.url, // URL (protocol://host:port)\n neo4j_driver.auth.basic(neo4jConfig.username, neo4jConfig.password), // Authentication\n );\n\n return ai.defineIndexer(\n {\n name: `neo4j/${params.indexId}`,\n //configSchema: neo4jIndexerOptionsSchema.optional(),\n },\n async (docs, options) => {\n const embeddings = await Promise.all(\n docs.map((doc) =>\n ai.embed({\n embedder,\n content: doc,\n options: embedderOptions,\n }),\n ),\n );\n\n const BATCH_SIZE = 1000;\n\n for (let i = 0; i < docs.length; i += BATCH_SIZE) {\n const batchDocs = docs.slice(i, i + BATCH_SIZE);\n const batchEmbeddings = embeddings.slice(i, i + BATCH_SIZE);\n\n const batchParams = batchDocs.map((el, j) => ({\n text: el.content[0][\"text\"],\n metadata: el.metadata ?? {},\n embedding: batchEmbeddings[j][0][\"embedding\"],\n }));\n\n await neo4j_instance.executeQuery(\n `\n UNWIND $data AS row\n CREATE (t:\\`${indexId}\\`)\n SET t.text = row.text,\n t += row.metadata\n WITH t, row.embedding AS embedding\n CALL db.create.setNodeVectorProperty(t, 'embedding', embedding)\n `,\n { data: batchParams },\n { database: neo4jConfig.database },\n );\n }\n\n await neo4j_instance.executeQuery(\n `\n CREATE VECTOR INDEX $indexName IF NOT EXISTS\n FOR (n:\\`${indexId}\\`) ON n.embedding\n `,\n { indexName: indexId },\n { database: neo4jConfig.database },\n );\n neo4j_instance.close();\n },\n );\n}\n\nfunction getDefaultConfig() {\n const {\n NEO4J_URI: url,\n NEO4J_USERNAME: username,\n NEO4J_PASSWORD: password,\n NEO4J_DATABASE: database,\n } = process.env;\n\n if (!url || !username || !password) {\n throw new Error(\n \"Please provide Neo4j connection details through environment variables: NEO4J_URI, NEO4J_USERNAME, and NEO4J_PASSWORD are required.\\n\" +\n \"For more details see https://neo4j.com/docs/api/javascript-driver/current/\",\n );\n }\n\n return {\n url,\n username,\n password,\n ...(database && { database }),\n };\n}\n"],"mappings":";;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;AAgBA,YAAY,kBAAkB;AAC9B,SAAiB,SAAS;AAC1B,SAAuB,oBAAoB;AAG3C;AAAA,EACE;AAAA,EACA;AAAA,EACA;AAAA,EACA;AAAA,OACK;AAEP,MAAM,8BAA8B,6BAA6B,OAAO;AAAA,EACtE,GAAG,EAAE,OAAO,EAAE,IAAI,GAAI;AAAA;AAExB,CAAC;AAED,MAAM,4BAA4B,EAAE,OAAO;AAAA,EACzC,WAAW,EAAE,OAAO,EAAE,SAAS;AACjC,CAAC;AAiBM,MAAM,oBAAoB,CAAC,WAG5B;AAvDN;AAwDE,SAAO,aAAa;AAAA,IAClB,MAAM,SAAS,OAAO,OAAO;AAAA,IAC7B,MAAM;AAAA,MACJ,QAAO,YAAO,gBAAP,YAAsB,WAAW,OAAO,OAAO;AAAA,IACxD;AAAA,IACA,cAAc;AAAA,EAChB,CAAC;AACH;AAUO,MAAM,kBAAkB,CAAC,WAG1B;AA5EN;AA6EE,SAAO,WAAW;AAAA,IAChB,MAAM,SAAS,OAAO,OAAO;AAAA,IAC7B,MAAM;AAAA,MACJ,QAAO,YAAO,gBAAP,YAAsB,WAAW,OAAO,OAAO;AAAA,IACxD;AAAA;AAAA,EAEF,CAAC;AACH;AAaO,SAAS,MACd,QAMc;AACd,SAAO,aAAa,SAAS,CAAO,OAAe;AACjD,WAAO,IAAI,CAAC,MAAM,wBAAwB,IAAI,CAAC,CAAC;AAChD,WAAO,IAAI,CAAC,MAAM,sBAAsB,IAAI,CAAC,CAAC;AAAA,EAChD,EAAC;AACH;AAEA,IAAO,gBAAQ;AAcR,SAAS,wBAGd,IACA,QAMA;AAvIF;AAwIE,QAAM,EAAE,SAAS,UAAU,gBAAgB,IAAI,mBAC1C;AAEL,QAAM,eAAc,YAAO,iBAAP,YAAuB,iBAAiB;AAC5D,QAAM,iBAAiB,aAAa;AAAA,IAClC,YAAY;AAAA;AAAA,IACZ,aAAa,KAAK,MAAM,YAAY,UAAU,YAAY,QAAQ;AAAA;AAAA,EACpE;AACA,SAAO,GAAG;AAAA,IACR;AAAA,MACE,MAAM,SAAS,OAAO,OAAO;AAAA,MAC7B,cAAc;AAAA,IAChB;AAAA,IACA,CAAO,SAAS,YAAY;AAC1B,YAAM,kBAAkB,MAAM,GAAG,MAAM;AAAA,QACrC;AAAA,QACA;AAAA,QACA,SAAS;AAAA,MACX,CAAC;AAED,YAAM,kBAAkB;AAAA;AAAA;AAAA;AAAA;AAKxB,YAAM,WAAW,MAAM,eAAe;AAAA,QACpC;AAAA,QACA;AAAA,UACE,GAAG,QAAQ;AAAA,UACX,WAAW,gBAAgB,CAAC,EAAE;AAAA,UAC9B,OAAO;AAAA,QACT;AAAA,QACA;AAAA,UACE,UAAU,YAAY;AAAA,QACxB;AAAA,MACF;AAEA,YAAM,YAAY,SAAS,QAAQ,IAAI,CAAC,OAAO;AAC7C,eAAO,SAAS;AAAA,UACd,GAAG,IAAI,MAAM;AAAA,UACb,OAAO;AAAA,YACL,OAAO,QAAQ,GAAG,IAAI,UAAU,CAAC,EAAE;AAAA,cACjC,CAAC,CAAC,GAAG,KAAK,MAAM,UAAU;AAAA,YAC5B;AAAA,UACF;AAAA,QACF;AAAA,MACF,CAAC;AACD,qBAAe,MAAM;AACrB,aAAO,EAAE,UAAqB;AAAA,IAChC;AAAA,EACF;AACF;AAcO,SAAS,sBAGd,IACA,QAMA;AAnNF;AAoNE,QAAM,EAAE,SAAS,UAAU,gBAAgB,IAAI,mBAC1C;AAEL,QAAM,eAAc,YAAO,iBAAP,YAAuB,iBAAiB;AAC5D,QAAM,iBAAiB,aAAa;AAAA,IAClC,YAAY;AAAA;AAAA,IACZ,aAAa,KAAK,MAAM,YAAY,UAAU,YAAY,QAAQ;AAAA;AAAA,EACpE;AAEA,SAAO,GAAG;AAAA,IACR;AAAA,MACE,MAAM,SAAS,OAAO,OAAO;AAAA;AAAA,IAE/B;AAAA,IACA,CAAO,MAAM,YAAY;AACvB,YAAM,aAAa,MAAM,QAAQ;AAAA,QAC/B,KAAK;AAAA,UAAI,CAAC,QACR,GAAG,MAAM;AAAA,YACP;AAAA,YACA,SAAS;AAAA,YACT,SAAS;AAAA,UACX,CAAC;AAAA,QACH;AAAA,MACF;AAEA,YAAM,aAAa;AAEnB,eAAS,IAAI,GAAG,IAAI,KAAK,QAAQ,KAAK,YAAY;AAChD,cAAM,YAAY,KAAK,MAAM,GAAG,IAAI,UAAU;AAC9C,cAAM,kBAAkB,WAAW,MAAM,GAAG,IAAI,UAAU;AAE1D,cAAM,cAAc,UAAU,IAAI,CAAC,IAAI,MAAG;AAnPlD,cAAAA;AAmPsD;AAAA,YAC5C,MAAM,GAAG,QAAQ,CAAC,EAAE,MAAM;AAAA,YAC1B,WAAUA,MAAA,GAAG,aAAH,OAAAA,MAAe,CAAC;AAAA,YAC1B,WAAW,gBAAgB,CAAC,EAAE,CAAC,EAAE,WAAW;AAAA,UAC9C;AAAA,SAAE;AAEF,cAAM,eAAe;AAAA,UACnB;AAAA;AAAA,wBAEc,OAAO;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA,UAMrB,EAAE,MAAM,YAAY;AAAA,UACpB,EAAE,UAAU,YAAY,SAAS;AAAA,QACnC;AAAA,MACF;AAEA,YAAM,eAAe;AAAA,QACnB;AAAA;AAAA,mBAEW,OAAO;AAAA;AAAA,QAElB,EAAE,WAAW,QAAQ;AAAA,QACrB,EAAE,UAAU,YAAY,SAAS;AAAA,MACnC;AACA,qBAAe,MAAM;AAAA,IACvB;AAAA,EACF;AACF;AAEA,SAAS,mBAAmB;AAC1B,QAAM;AAAA,IACJ,WAAW;AAAA,IACX,gBAAgB;AAAA,IAChB,gBAAgB;AAAA,IAChB,gBAAgB;AAAA,EAClB,IAAI,QAAQ;AAEZ,MAAI,CAAC,OAAO,CAAC,YAAY,CAAC,UAAU;AAClC,UAAM,IAAI;AAAA,MACR;AAAA,IAEF;AAAA,EACF;AAEA,SAAO;AAAA,IACL;AAAA,IACA;AAAA,IACA;AAAA,KACI,YAAY,EAAE,SAAS;AAE/B;","names":["_a"]}