genkitx-pinecone/lib/index.mjs

Genkit AI framework plugin for Pinecone vector database.

import {
  Pinecone
} from "@pinecone-database/pinecone";
import { z } from "genkit";
import { genkitPlugin } from "genkit/plugin";
import {
  CommonRetrieverOptionsSchema,
  Document,
  indexerRef,
  retrieverRef
} from "genkit/retriever";
import { Md5 } from "ts-md5";
const SparseVectorSchema = z.object({
  indices: z.number().array(),
  values: z.number().array()
}).refine(
  (input) => {
    return input.indices.length === input.values.length;
  },
  {
    message: "Indices and values must be of the same length"
  }
);
const PineconeRetrieverOptionsSchema = CommonRetrieverOptionsSchema.extend({
  k: z.number().max(1e3),
  namespace: z.string().optional(),
  filter: z.record(z.string(), z.any()).optional(),
  // includeValues is always false
  // includeMetadata is always true
  sparseVector: SparseVectorSchema.optional()
});
const PineconeIndexerOptionsSchema = z.object({
  namespace: z.string().optional()
});
const CONTENT_KEY = "_content";
const CONTENT_TYPE = "_contentType";
const pineconeRetrieverRef = (params) => {
  return retrieverRef({
    name: `pinecone/${params.indexId}`,
    info: {
      label: params.displayName ?? `Pinecone - ${params.indexId}`
    },
    configSchema: PineconeRetrieverOptionsSchema
  });
};
const pineconeIndexerRef = (params) => {
  return indexerRef({
    name: `pinecone/${params.indexId}`,
    info: {
      label: params.displayName ?? `Pinecone - ${params.indexId}`
    },
    configSchema: PineconeIndexerOptionsSchema.optional()
  });
};
function pinecone(params) {
  return genkitPlugin("pinecone", async (ai) => {
    params.map((i) => configurePineconeRetriever(ai, i));
    params.map((i) => configurePineconeIndexer(ai, i));
  });
}
var src_default = pinecone;
function configurePineconeRetriever(ai, params) {
  const { indexId, embedder, embedderOptions } = {
    ...params
  };
  const pineconeConfig = params.clientParams ?? getDefaultConfig();
  const contentKey = params.contentKey ?? params.textKey ?? CONTENT_KEY;
  const pinecone2 = new Pinecone(pineconeConfig);
  const index = pinecone2.index(indexId);
  return ai.defineRetriever(
    {
      name: `pinecone/${params.indexId}`,
      configSchema: PineconeRetrieverOptionsSchema
    },
    async (content, options) => {
      const queryEmbeddings = await ai.embed({
        embedder,
        content,
        options: embedderOptions
      });
      const scopedIndex = !!options.namespace ? index.namespace(options.namespace) : index;
      const response = await scopedIndex.query({
        topK: options.k,
        vector: queryEmbeddings[0].embedding,
        includeValues: false,
        includeMetadata: true
      });
      return {
        documents: response.matches.map((m) => m.metadata).filter((m) => !!m).map((m) => {
          const metadata = m;
          return Document.fromData(
            metadata[contentKey],
            metadata[CONTENT_TYPE],
            JSON.parse(metadata.docMetadata)
          );
        })
      };
    }
  );
}
function configurePineconeIndexer(ai, params) {
  const { indexId, embedder, embedderOptions } = {
    ...params
  };
  const pineconeConfig = params.clientParams ?? getDefaultConfig();
  const contentKey = params.contentKey ?? params.textKey ?? CONTENT_KEY;
  const pinecone2 = new Pinecone(pineconeConfig);
  const index = pinecone2.index(indexId);
  return ai.defineIndexer(
    {
      name: `pinecone/${params.indexId}`,
      configSchema: PineconeIndexerOptionsSchema.optional()
    },
    async (docs, options) => {
      const scopedIndex = !!options?.namespace ? index.namespace(options.namespace) : index;
      const embeddings = await Promise.all(
        docs.map(
          (doc) => ai.embed({
            embedder,
            content: doc,
            options: embedderOptions
          })
        )
      );
      await scopedIndex.upsert(
        embeddings.map((value, i) => {
          const doc = docs[i];
          const docEmbeddings = value;
          const embeddingDocs = doc.getEmbeddingDocuments(docEmbeddings);
          return docEmbeddings.map((docEmbedding, j) => {
            const metadata = {
              docMetadata: JSON.stringify(embeddingDocs[j].metadata)
            };
            metadata[contentKey] = embeddingDocs[j].data;
            metadata[CONTENT_TYPE] = embeddingDocs[j].dataType || "";
            const id = Md5.hashStr(JSON.stringify(embeddingDocs[j]));
            return {
              id,
              values: docEmbedding.embedding,
              metadata
            };
          });
        }).reduce((acc, val) => {
          return acc.concat(val);
        }, [])
      );
    }
  );
}
async function createPineconeIndex(params) {
  const pineconeConfig = params.clientParams ?? getDefaultConfig();
  const pinecone2 = new Pinecone(pineconeConfig);
  return await pinecone2.createIndex(params.options);
}
async function describePineconeIndex(params) {
  const pineconeConfig = params.clientParams ?? getDefaultConfig();
  const pinecone2 = new Pinecone(pineconeConfig);
  return await pinecone2.describeIndex(params.name);
}
async function deletePineconeIndex(params) {
  const pineconeConfig = params.clientParams ?? getDefaultConfig();
  const pinecone2 = new Pinecone(pineconeConfig);
  return await pinecone2.deleteIndex(params.name);
}
function getDefaultConfig() {
  const maybeApiKey = process.env.PINECONE_API_KEY;
  if (!maybeApiKey)
    throw new Error(
      "Please pass in the API key or set PINECONE_API_KEY environment variable.\nFor more details see https://firebase.google.com/docs/genkit/plugins/pinecone"
    );
  return { apiKey: maybeApiKey };
}
export {
  configurePineconeIndexer,
  configurePineconeRetriever,
  createPineconeIndex,
  src_default as default,
  deletePineconeIndex,
  describePineconeIndex,
  pinecone,
  pineconeIndexerRef,
  pineconeRetrieverRef
};
//# sourceMappingURL=index.mjs.map