From 37cf852c57627234b25b1e7315f6e222f4b78ff7 Mon Sep 17 00:00:00 2001 From: HassanBahati Date: Mon, 8 Sep 2025 17:38:12 +0300 Subject: [PATCH 1/2] feat(js/plugins/chroma): migrate chroma plugin to v2 plugin api --- js/plugins/chroma/src/index.ts | 248 ++++++++++++++++++++++++++++++++- 1 file changed, 243 insertions(+), 5 deletions(-) diff --git a/js/plugins/chroma/src/index.ts b/js/plugins/chroma/src/index.ts index 696bcf5586..e490929f7d 100644 --- a/js/plugins/chroma/src/index.ts +++ b/js/plugins/chroma/src/index.ts @@ -32,11 +32,19 @@ import { retrieverRef, z, type EmbedderArgument, + type EmbedderAction, type Embedding, type Genkit, } from 'genkit'; -import { genkitPlugin, type GenkitPlugin } from 'genkit/plugin'; +import { + genkitPluginV2, + type GenkitPluginV2, + type ResolvableAction, + retriever, + indexer, +} from 'genkit/plugin'; import { CommonRetrieverOptionsSchema } from 'genkit/retriever'; +import type { ActionType } from 'genkit/registry'; import { Md5 } from 'ts-md5'; export { IncludeEnum }; @@ -74,12 +82,53 @@ type ChromaPluginParams< /** * Chroma plugin that provides the Chroma retriever and indexer */ + export function chroma( params: ChromaPluginParams -): GenkitPlugin { - return genkitPlugin('chroma', async (ai: Genkit) => { - params.map((i) => chromaRetriever(ai, i)); - params.map((i) => chromaIndexer(ai, i)); +): GenkitPluginV2 { + return genkitPluginV2({ + name: 'chroma', + async init() { + const actions: ResolvableAction[] = []; + for (const param of params) { + actions.push(createChromaRetriever(param)); + actions.push(createChromaIndexer(param)); + } + return actions; + }, + async resolve(actionType: ActionType, name: string) { + // Find the matching param by collection name + const collectionName = name.replace('chroma/', ''); + const param = params.find(p => p.collectionName === collectionName); + if (!param) return undefined; + + switch (actionType) { + case 'retriever': + return createChromaRetriever(param); + case 'indexer': + return createChromaIndexer(param); + default: + return undefined; + } + }, + async list() { + return params.flatMap(param => [ + { + name: `chroma/${param.collectionName}`, + type: 'retriever' as const, + info: { + label: `Chroma DB - ${param.collectionName}`, + }, + }, + { + name: `chroma/${param.collectionName}`, + type: 'indexer' as const, + info: { + label: `Chroma DB - ${param.collectionName}`, + }, + }, + ]); + }, }); } @@ -369,6 +418,195 @@ export async function deleteChromaCollection(params: { }); } +/** + * Standalone Chroma retriever action (v2 API) + */ +function createChromaRetriever( + params: { + clientParams?: ChromaClientParams; + collectionName: string; + createCollectionIfMissing?: boolean; + embedder: EmbedderArgument; + embedderOptions?: z.infer; + } +) { + const { embedder, collectionName, embedderOptions } = params; + + return retriever( + { + name: `chroma/${collectionName}`, + configSchema: ChromaRetrieverOptionsSchema.optional(), + }, + async (content, options) => { + const clientParams = await resolve(params.clientParams); + const client = new ChromaClient(clientParams); + let collection: Collection; + if (params.createCollectionIfMissing) { + collection = await client.getOrCreateCollection({ + name: collectionName, + }); + } else { + collection = await client.getCollection({ + name: collectionName, + }); + } + + // For v2 API, we need to handle embedding differently + // The embedder will be resolved at runtime + const queryEmbeddings = await resolveEmbedder(embedder, { + content, + options: embedderOptions, + }); + + const results = await collection.query({ + nResults: options?.k, + include: getIncludes(options?.include), + where: options?.where, + whereDocument: options?.whereDocument, + queryEmbeddings: queryEmbeddings[0].embedding, + }); + + const documents = results.documents[0]; + const metadatas = results.metadatas; + const embeddings = results.embeddings; + const distances = results.distances; + + const combined = documents + .map((d, i) => { + if (d !== null) { + return { + document: d, + metadata: constructMetadata(i, metadatas, embeddings, distances), + }; + } + return undefined; + }) + .filter( + (r): r is { document: string; metadata: Record } => !!r + ); + + return { + documents: combined.map((result) => { + const data = result.document; + const metadata = result.metadata.metadata[0]; + const dataType = metadata.dataType; + const docMetadata = metadata.docMetadata + ? JSON.parse(metadata.docMetadata) + : undefined; + return Document.fromData(data, dataType, docMetadata).toJSON(); + }), + }; + } + ); +} + +/** + * Standalone Chroma indexer action (v2 API) + */ +function createChromaIndexer( + params: { + clientParams?: ChromaClientParams; + collectionName: string; + createCollectionIfMissing?: boolean; + embedder: EmbedderArgument; + embedderOptions?: z.infer; + } +) { + const { collectionName, embedder, embedderOptions } = { + ...params, + }; + + return indexer( + { + name: `chroma/${params.collectionName}`, + configSchema: ChromaIndexerOptionsSchema, + }, + async (docs) => { + const clientParams = await resolve(params.clientParams); + const client = new ChromaClient(clientParams); + + let collection: Collection; + if (params.createCollectionIfMissing) { + collection = await client.getOrCreateCollection({ + name: collectionName, + }); + } else { + collection = await client.getCollection({ + name: collectionName, + }); + } + + const embeddings = await Promise.all( + docs.map((doc) => + resolveEmbedder(embedder, { + content: doc, + options: embedderOptions, + }) + ) + ); + + const entries = embeddings + .map((value, i) => { + const doc = docs[i]; + // The array of embeddings for this document + const docEmbeddings: Embedding[] = value; + const embeddingDocs = doc.getEmbeddingDocuments(docEmbeddings); + return docEmbeddings.map((docEmbedding, j) => { + const metadata: Metadata = { + docMetadata: JSON.stringify(embeddingDocs[j].metadata), + dataType: embeddingDocs[j].dataType || '', + }; + + const data = embeddingDocs[j].data; + const id = Md5.hashStr(JSON.stringify(embeddingDocs[j])); + return { + id, + value: docEmbedding.embedding, + document: data, + metadata, + }; + }); + }) + .reduce((acc, val) => { + return acc.concat(val); + }, []); + + await collection.add({ + ids: entries.map((e) => e.id), + embeddings: entries.map((e) => e.value), + metadatas: entries.map((e) => e.metadata), + documents: entries.map((e) => e.document), + }); + } + ); +} + +/** + * Helper function to resolve embedder and get embeddings + * Call embedder actions directly + */ +async function resolveEmbedder( + embedder: EmbedderArgument, + params: { + content: Document; + options?: z.infer; + } +): Promise { + // If embedder is an action (function with __action property), call it directly + if (typeof embedder === 'function' && '__action' in embedder) { + const embedderAction = embedder as EmbedderAction; + const response = await embedderAction({ + input: [params.content], + options: params.options, + }); + return response.embeddings; + } + + // If embedder is a string reference, we need to resolve it + // throw an error as this requires registry access + throw new Error(`Embedder resolution for string references not supported in v2 API: ${embedder}`); +} + async function resolve( params?: ChromaClientParams ): Promise { From efc5cf3a8f84fe8b0fc842bc4a797422fce64763 Mon Sep 17 00:00:00 2001 From: HassanBahati Date: Tue, 9 Sep 2025 17:27:51 +0300 Subject: [PATCH 2/2] chore(js/plugins/chroma): format --- js/plugins/chroma/src/index.ts | 60 ++++++++++++++++++---------------- 1 file changed, 31 insertions(+), 29 deletions(-) diff --git a/js/plugins/chroma/src/index.ts b/js/plugins/chroma/src/index.ts index e490929f7d..94e272441c 100644 --- a/js/plugins/chroma/src/index.ts +++ b/js/plugins/chroma/src/index.ts @@ -31,20 +31,20 @@ import { indexerRef, retrieverRef, z, - type EmbedderArgument, type EmbedderAction, + type EmbedderArgument, type Embedding, type Genkit, } from 'genkit'; -import { +import { genkitPluginV2, + indexer, + retriever, type GenkitPluginV2, type ResolvableAction, - retriever, - indexer, } from 'genkit/plugin'; -import { CommonRetrieverOptionsSchema } from 'genkit/retriever'; import type { ActionType } from 'genkit/registry'; +import { CommonRetrieverOptionsSchema } from 'genkit/retriever'; import { Md5 } from 'ts-md5'; export { IncludeEnum }; @@ -99,7 +99,7 @@ export function chroma( async resolve(actionType: ActionType, name: string) { // Find the matching param by collection name const collectionName = name.replace('chroma/', ''); - const param = params.find(p => p.collectionName === collectionName); + const param = params.find((p) => p.collectionName === collectionName); if (!param) return undefined; switch (actionType) { @@ -112,7 +112,7 @@ export function chroma( } }, async list() { - return params.flatMap(param => [ + return params.flatMap((param) => [ { name: `chroma/${param.collectionName}`, type: 'retriever' as const, @@ -421,17 +421,17 @@ export async function deleteChromaCollection(params: { /** * Standalone Chroma retriever action (v2 API) */ -function createChromaRetriever( - params: { - clientParams?: ChromaClientParams; - collectionName: string; - createCollectionIfMissing?: boolean; - embedder: EmbedderArgument; - embedderOptions?: z.infer; - } -) { +function createChromaRetriever< + EmbedderCustomOptions extends z.ZodTypeAny, +>(params: { + clientParams?: ChromaClientParams; + collectionName: string; + createCollectionIfMissing?: boolean; + embedder: EmbedderArgument; + embedderOptions?: z.infer; +}) { const { embedder, collectionName, embedderOptions } = params; - + return retriever( { name: `chroma/${collectionName}`, @@ -457,7 +457,7 @@ function createChromaRetriever( content, options: embedderOptions, }); - + const results = await collection.query({ nResults: options?.k, include: getIncludes(options?.include), @@ -503,15 +503,15 @@ function createChromaRetriever( /** * Standalone Chroma indexer action (v2 API) */ -function createChromaIndexer( - params: { - clientParams?: ChromaClientParams; - collectionName: string; - createCollectionIfMissing?: boolean; - embedder: EmbedderArgument; - embedderOptions?: z.infer; - } -) { +function createChromaIndexer< + EmbedderCustomOptions extends z.ZodTypeAny, +>(params: { + clientParams?: ChromaClientParams; + collectionName: string; + createCollectionIfMissing?: boolean; + embedder: EmbedderArgument; + embedderOptions?: z.infer; +}) { const { collectionName, embedder, embedderOptions } = { ...params, }; @@ -601,10 +601,12 @@ async function resolveEmbedder( }); return response.embeddings; } - + // If embedder is a string reference, we need to resolve it // throw an error as this requires registry access - throw new Error(`Embedder resolution for string references not supported in v2 API: ${embedder}`); + throw new Error( + `Embedder resolution for string references not supported in v2 API: ${embedder}` + ); } async function resolve(