MCP Terminal Server
by dillip285
/**
* Copyright 2024 Google LLC
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
import {
CreateIndexOptions,
Pinecone,
PineconeConfiguration,
RecordMetadata,
} from '@pinecone-database/pinecone';
import { Genkit, z } from 'genkit';
import { GenkitPlugin, genkitPlugin } from 'genkit/plugin';
import { EmbedderArgument, Embedding } from 'genkit/embedder';
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(1000),
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';
/**
* pineconeRetrieverRef function creates a retriever for Pinecone.
* @param params The params for the new Pinecone retriever
* @param params.indexId The indexId for the Pinecone retriever
* @param params.displayName A display name for the retriever.
If not specified, the default label will be `Pinecone - <indexId>`
* @returns A reference to a Pinecone retriever.
*/
export const pineconeRetrieverRef = (params: {
indexId: string;
displayName?: string;
}) => {
return retrieverRef({
name: `pinecone/${params.indexId}`,
info: {
label: params.displayName ?? `Pinecone - ${params.indexId}`,
},
configSchema: PineconeRetrieverOptionsSchema,
});
};
/**
* pineconeIndexerRef function creates an indexer for Pinecone.
* @param params The params for the new Pinecone indexer.
* @param params.indexId The indexId for the Pinecone indexer.
* @param params.displayName A display name for the indexer.
If not specified, the default label will be `Pinecone - <indexId>`
* @returns A reference to a Pinecone indexer.
*/
export const pineconeIndexerRef = (params: {
indexId: string;
displayName?: string;
}) => {
return indexerRef({
name: `pinecone/${params.indexId}`,
info: {
label: params.displayName ?? `Pinecone - ${params.indexId}`,
},
configSchema: PineconeIndexerOptionsSchema.optional(),
});
};
/**
* Pinecone plugin that provides a Pinecone retriever and indexer
* @param params An array of params to set up Pinecone retrievers and indexers
* @param params.clientParams PineconeConfiguration containing the
PINECONE_API_KEY. If not set, the PINECONE_API_KEY environment variable will
be used instead.
* @param params.indexId The name of the index
* @param params.embedder The embedder to use for the indexer and retriever
* @param params.embedderOptions Options to customize the embedder
* @returns The Pinecone Genkit plugin
*/
export function pinecone<EmbedderCustomOptions extends z.ZodTypeAny>(
params: {
clientParams?: PineconeConfiguration;
indexId: string;
contentKey?: string;
embedder: EmbedderArgument<EmbedderCustomOptions>;
embedderOptions?: z.infer<EmbedderCustomOptions>;
}[]
): GenkitPlugin {
return genkitPlugin('pinecone', async (ai: Genkit) => {
params.map((i) => configurePineconeRetriever(ai, i));
params.map((i) => configurePineconeIndexer(ai, i));
});
}
export default pinecone;
/**
* Configures a Pinecone retriever.
* @param ai A Genkit instance
* @param params The params for the retriever
* @param params.indexId The name of the retriever
* @param params.clientParams PineconeConfiguration containing the
PINECONE_API_KEY. If not set, the PINECONE_API_KEY environment variable will
be used instead.
* @param params.textKey Deprecated. Please use contentKey.
* @param params.contentKey The metadata key that contains the
content. If not specified, the value '_content' is used by default.
* @param params.embedder The embedder to use for the retriever
* @param params.embedderOptions Options to customize the embedder
* @returns A Pinecone retriever
*/
export function configurePineconeRetriever<
EmbedderCustomOptions extends z.ZodTypeAny,
>(
ai: Genkit,
params: {
indexId: string;
clientParams?: PineconeConfiguration;
/**
* @deprecated use contentKey instead.
*/
textKey?: string;
contentKey?: string;
embedder: EmbedderArgument<EmbedderCustomOptions>;
embedderOptions?: z.infer<EmbedderCustomOptions>;
}
) {
const { indexId, embedder, embedderOptions } = {
...params,
};
const pineconeConfig = params.clientParams ?? getDefaultConfig();
const contentKey = params.contentKey ?? params.textKey ?? CONTENT_KEY;
const pinecone = new Pinecone(pineconeConfig);
const index = pinecone.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 is RecordMetadata => !!m)
.map((m) => {
const metadata = m;
return Document.fromData(
metadata[contentKey] as string,
metadata[CONTENT_TYPE] as string,
JSON.parse(metadata.docMetadata as string) as Record<
string,
unknown
>
);
}),
};
}
);
}
/**
* Configures a Pinecone indexer.
* @param ai A Genkit instance
* @param params The params for the indexer
* @param params.indexId The name of the indexer
* @param params.clientParams PineconeConfiguration containing the
PINECONE_API_KEY. If not set, the PINECONE_API_KEY environment variable will
be used instead.
* @param params.textKey Deprecated. Please use contentKey.
* @param params.contentKey The metadata key that contains the
content. If not specified, the value '_content' is used by default.
* @param params.embedder The embedder to use for the retriever
* @param params.embedderOptions Options to customize the embedder
* @returns A Genkit indexer
*/
export function configurePineconeIndexer<
EmbedderCustomOptions extends z.ZodTypeAny,
>(
ai: Genkit,
params: {
indexId: string;
clientParams?: PineconeConfiguration;
/**
* @deprecated use contentKey instead.
*/
textKey?: string;
contentKey?: string;
embedder: EmbedderArgument<EmbedderCustomOptions>;
embedderOptions?: z.infer<EmbedderCustomOptions>;
}
) {
const { indexId, embedder, embedderOptions } = {
...params,
};
const pineconeConfig = params.clientParams ?? getDefaultConfig();
const contentKey = params.contentKey ?? params.textKey ?? CONTENT_KEY;
const pinecone = new Pinecone(pineconeConfig);
const index = pinecone.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];
// The array of embeddings for this document
const docEmbeddings: Embedding[] = value;
// Create one doc per docEmbedding so we can store them 1:1.
// They should be unique because the embedding metadata is
// added to the new docs.
const embeddingDocs = doc.getEmbeddingDocuments(docEmbeddings);
return docEmbeddings.map((docEmbedding, j) => {
const metadata: RecordMetadata = {
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);
}, [])
);
}
);
}
/**
* Helper function for creating a Pinecone index.
* @param params The params for creating a Pinecone index
* @param params.clientParams The params to initialize Pinecone.
* @param params.options The options for creating the index.
* @returns A Pinecone index.
*/
export async function createPineconeIndex(params: {
clientParams?: PineconeConfiguration;
options: CreateIndexOptions;
}) {
const pineconeConfig = params.clientParams ?? getDefaultConfig();
const pinecone = new Pinecone(pineconeConfig);
return await pinecone.createIndex(params.options);
}
/**
* Helper function to describe a Pinecone index. Use it to check if a newly created index is ready for use.
* @param params The params for describing a Pinecone index.
* @param params.clientParams The params to initialize Pinecone.
* @param params.name The name of the Pinecone index to describe.
* @return A description of the Pinecone index.
*/
export async function describePineconeIndex(params: {
clientParams?: PineconeConfiguration;
name: string;
}) {
const pineconeConfig = params.clientParams ?? getDefaultConfig();
const pinecone = new Pinecone(pineconeConfig);
return await pinecone.describeIndex(params.name);
}
/**
* Helper function for deleting pinecone indices.
* @param params The params for deleting a Pinecone index.
* @param params.clientParams The params to initialize Pinecone.
* @param params.name The name of the Pinecone index to delete.
* @returns a void Promise that is fulfilled when the index has been deleted.
*/
export async function deletePineconeIndex(params: {
clientParams?: PineconeConfiguration;
name: string;
}) {
const pineconeConfig = params.clientParams ?? getDefaultConfig();
const pinecone = new Pinecone(pineconeConfig);
return await pinecone.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.\n' +
'For more details see https://firebase.google.com/docs/genkit/plugins/pinecone'
);
return { apiKey: maybeApiKey } as PineconeConfiguration;
}