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 similarity from 'compute-cosine-similarity';
import * as fs from 'fs';
import { Embedding, Genkit, z } from 'genkit';
import { EmbedderArgument } from 'genkit/embedder';
import { GenkitPlugin, genkitPlugin } from 'genkit/plugin';
import {
CommonRetrieverOptionsSchema,
Document,
DocumentData,
indexerRef,
retrieverRef,
} from 'genkit/retriever';
import { Md5 } from 'ts-md5';
const _LOCAL_FILESTORE = '__db_{INDEX_NAME}.json';
interface DbValue {
doc: DocumentData;
embedding: Embedding;
}
function loadFilestore(indexName: string) {
let existingData = {};
const indexFileName = _LOCAL_FILESTORE.replace('{INDEX_NAME}', indexName);
if (fs.existsSync(indexFileName)) {
existingData = JSON.parse(fs.readFileSync(indexFileName).toString());
}
return existingData;
}
function addDocument(
embedding: Embedding,
doc: Document,
contents: Record<string, DbValue>
) {
const id = Md5.hashStr(JSON.stringify(doc));
if (!(id in contents)) {
contents[id] = { doc, embedding };
} else {
console.debug(`Skipping ${id} since it is already present`);
}
}
interface Params<EmbedderCustomOptions extends z.ZodTypeAny> {
indexName: string;
embedder: EmbedderArgument<EmbedderCustomOptions>;
embedderOptions?: z.infer<EmbedderCustomOptions>;
}
/**
* Local file-based vectorstore plugin that provides retriever and indexer.
*
* NOT INTENDED FOR USE IN PRODUCTION
*/
export function devLocalVectorstore<EmbedderCustomOptions extends z.ZodTypeAny>(
params: Params<EmbedderCustomOptions>[]
): GenkitPlugin {
return genkitPlugin('devLocalVectorstore', async (ai) => {
params.map((p) => configureDevLocalRetriever(ai, p));
params.map((p) => configureDevLocalIndexer(ai, p));
});
}
export default devLocalVectorstore;
/**
* Local file-based vectorstore retriever reference
*/
export function devLocalRetrieverRef(indexName: string) {
return retrieverRef({
name: `devLocalVectorstore/${indexName}`,
info: {
label: `Local file-based Retriever - ${indexName}`,
},
configSchema: CommonRetrieverOptionsSchema.optional(),
});
}
/**
* Local file-based indexer reference
*/
export function devLocalIndexerRef(indexName: string) {
return indexerRef({
name: `devLocalVectorstore/${indexName}`,
info: {
label: `Local file-based Indexer - ${indexName}`,
},
configSchema: z.null().optional(),
});
}
async function importDocumentsToLocalVectorstore<
EmbedderCustomOptions extends z.ZodTypeAny,
>(
ai: Genkit,
params: {
indexName: string;
docs: Array<Document>;
embedder: EmbedderArgument<EmbedderCustomOptions>;
embedderOptions?: z.infer<EmbedderCustomOptions>;
}
) {
const { docs, embedder, embedderOptions } = { ...params };
const data = loadFilestore(params.indexName);
await Promise.all(
docs.map(async (doc) => {
const embeddings = await ai.embed({
embedder,
content: doc,
options: embedderOptions,
});
const embeddingDocs = doc.getEmbeddingDocuments(embeddings);
for (const i in embeddingDocs) {
addDocument(embeddings[i], embeddingDocs[i], data);
}
})
);
// Update the file
fs.writeFileSync(
_LOCAL_FILESTORE.replace('{INDEX_NAME}', params.indexName),
JSON.stringify(data, null, 2)
);
}
async function getClosestDocuments<
I extends z.ZodTypeAny,
EmbedderCustomOptions extends z.ZodTypeAny,
>(params: {
queryEmbeddings: Array<number>;
db: Record<string, DbValue>;
k: number;
}): Promise<Document[]> {
const scoredDocs: { score: number; doc: Document }[] = [];
// Very dumb way to check for similar docs.
for (const value of Object.values(params.db)) {
const thisEmbedding = value.embedding.embedding;
const score = similarity(params.queryEmbeddings, thisEmbedding) ?? 0;
scoredDocs.push({
score,
doc: new Document(value.doc),
});
}
scoredDocs.sort((a, b) => (a.score > b.score ? -1 : 1));
return scoredDocs.slice(0, params.k).map((o) => o.doc);
}
/**
* Configures a local vectorstore retriever
*/
function configureDevLocalRetriever<EmbedderCustomOptions extends z.ZodTypeAny>(
ai: Genkit,
params: {
indexName: string;
embedder: EmbedderArgument<EmbedderCustomOptions>;
embedderOptions?: z.infer<EmbedderCustomOptions>;
}
) {
const { embedder, embedderOptions } = params;
const vectorstore = ai.defineRetriever(
{
name: `devLocalVectorstore/${params.indexName}`,
configSchema: CommonRetrieverOptionsSchema,
},
async (content, options) => {
const db = loadFilestore(params.indexName);
const embeddings = await ai.embed({
embedder,
content,
options: embedderOptions,
});
return {
documents: await getClosestDocuments({
k: options?.k ?? 3,
queryEmbeddings: embeddings[0].embedding,
db,
}),
};
}
);
return vectorstore;
}
/**
* Configures a local vectorstore indexer.
*/
function configureDevLocalIndexer<EmbedderCustomOptions extends z.ZodTypeAny>(
ai: Genkit,
params: {
indexName: string;
embedder: EmbedderArgument<EmbedderCustomOptions>;
embedderOptions?: z.infer<EmbedderCustomOptions>;
}
) {
const { embedder, embedderOptions } = params;
const vectorstore = ai.defineIndexer(
{ name: `devLocalVectorstore/${params.indexName}` },
async (docs) => {
await importDocumentsToLocalVectorstore(ai, {
indexName: params.indexName,
docs,
embedder,
embedderOptions: embedderOptions,
});
}
);
return vectorstore;
}