@alpic80/rivet-core 1.19.1-aidon.1

package/dist/esm/plugins/google/nodes/ChatGoogleNode.js

@@ -0,0 +1,364 @@
+import { uint8ArrayToBase64, } from '../../../index.js';
+import { googleModelOptions, googleModels, streamChatCompletions, } from '../google.js';
+import { nanoid } from 'nanoid/non-secure';
+import { dedent } from 'ts-dedent';
+import retry from 'p-retry';
+import { match } from 'ts-pattern';
+import { coerceType, coerceTypeOptional } from '../../../utils/coerceType.js';
+import { addWarning } from '../../../utils/outputs.js';
+import { getError } from '../../../utils/errors.js';
+import { pluginNodeDefinition } from '../../../model/NodeDefinition.js';
+import { getScalarTypeOf, isArrayDataValue } from '../../../model/DataValue.js';
+// Temporary
+const cache = new Map();
+export const ChatGoogleNodeImpl = {
+    create() {
+        const chartNode = {
+            type: 'chatGoogle',
+            title: 'Chat (Google)',
+            id: nanoid(),
+            visualData: {
+                x: 0,
+                y: 0,
+                width: 275,
+            },
+            data: {
+                model: 'gemini-pro',
+                useModelInput: false,
+                temperature: 0.5,
+                useTemperatureInput: false,
+                top_p: 1,
+                useTopPInput: false,
+                top_k: undefined,
+                useTopKInput: false,
+                useTopP: false,
+                useUseTopPInput: false,
+                maxTokens: 1024,
+                useMaxTokensInput: false,
+                cache: false,
+                useAsGraphPartialOutput: true,
+            },
+        };
+        return chartNode;
+    },
+    getInputDefinitions(data) {
+        const inputs = [];
+        if (data.useModelInput) {
+            inputs.push({
+                id: 'model',
+                title: 'Model',
+                dataType: 'string',
+                required: false,
+            });
+        }
+        if (data.useTemperatureInput) {
+            inputs.push({
+                dataType: 'number',
+                id: 'temperature',
+                title: 'Temperature',
+            });
+        }
+        if (data.useTopPInput) {
+            inputs.push({
+                dataType: 'number',
+                id: 'top_p',
+                title: 'Top P',
+            });
+        }
+        if (data.useUseTopPInput) {
+            inputs.push({
+                dataType: 'boolean',
+                id: 'useTopP',
+                title: 'Use Top P',
+            });
+        }
+        if (data.useMaxTokensInput) {
+            inputs.push({
+                dataType: 'number',
+                id: 'maxTokens',
+                title: 'Max Tokens',
+            });
+        }
+        inputs.push({
+            dataType: ['chat-message', 'chat-message[]'],
+            id: 'prompt',
+            title: 'Prompt',
+        });
+        return inputs;
+    },
+    getOutputDefinitions(data) {
+        const outputs = [];
+        outputs.push({
+            dataType: 'string',
+            id: 'response',
+            title: 'Response',
+        });
+        outputs.push({
+            dataType: 'chat-message[]',
+            id: 'in-messages',
+            title: 'Messages Sent',
+            description: 'All messages sent to the model.',
+        });
+        outputs.push({
+            dataType: 'chat-message[]',
+            id: 'all-messages',
+            title: 'All Messages',
+            description: 'All messages, with the response appended.',
+        });
+        return outputs;
+    },
+    getBody(data) {
+        return dedent `
+      ${googleModels[data.model]?.displayName ?? `Google (${data.model})`}
+      ${data.useTopP
+            ? `Top P: ${data.useTopPInput ? '(Using Input)' : data.top_p}`
+            : `Temperature: ${data.useTemperatureInput ? '(Using Input)' : data.temperature}`}
+      Max Tokens: ${data.maxTokens}
+    `;
+    },
+    getEditors() {
+        return [
+            {
+                type: 'dropdown',
+                label: 'Model',
+                dataKey: 'model',
+                useInputToggleDataKey: 'useModelInput',
+                options: googleModelOptions,
+            },
+            {
+                type: 'number',
+                label: 'Temperature',
+                dataKey: 'temperature',
+                useInputToggleDataKey: 'useTemperatureInput',
+                min: 0,
+                max: 2,
+                step: 0.1,
+            },
+            {
+                type: 'number',
+                label: 'Top P',
+                dataKey: 'top_p',
+                useInputToggleDataKey: 'useTopPInput',
+                min: 0,
+                max: 1,
+                step: 0.1,
+            },
+            {
+                type: 'toggle',
+                label: 'Use Top P',
+                dataKey: 'useTopP',
+                useInputToggleDataKey: 'useUseTopPInput',
+            },
+            {
+                type: 'number',
+                label: 'Max Tokens',
+                dataKey: 'maxTokens',
+                useInputToggleDataKey: 'useMaxTokensInput',
+                min: 0,
+                max: Number.MAX_SAFE_INTEGER,
+                step: 1,
+            },
+            {
+                type: 'toggle',
+                label: 'Cache (same inputs, same outputs)',
+                dataKey: 'cache',
+            },
+            {
+                type: 'toggle',
+                label: 'Use for subgraph partial output',
+                dataKey: 'useAsGraphPartialOutput',
+            },
+        ];
+    },
+    getUIData() {
+        return {
+            infoBoxBody: dedent `
+        Makes a call to an Google chat model. The settings contains many options for tweaking the model's behavior.
+      `,
+            infoBoxTitle: 'Chat (Google) Node',
+            contextMenuTitle: 'Chat (Google)',
+            group: ['AI'],
+        };
+    },
+    async process(data, inputs, context) {
+        const output = {};
+        const rawModel = data.useModelInput
+            ? coerceTypeOptional(inputs['model'], 'string') ?? data.model
+            : data.model;
+        const model = rawModel;
+        const temperature = data.useTemperatureInput
+            ? coerceTypeOptional(inputs['temperature'], 'number') ?? data.temperature
+            : data.temperature;
+        const topP = data.useTopPInput ? coerceTypeOptional(inputs['top_p'], 'number') ?? data.top_p : data.top_p;
+        const useTopP = data.useUseTopPInput
+            ? coerceTypeOptional(inputs['useTopP'], 'boolean') ?? data.useTopP
+            : data.useTopP;
+        const { messages } = getChatGoogleNodeMessages(inputs);
+        const prompt = await Promise.all(messages.map(async (message) => {
+            return {
+                role: message.type === 'user' ? 'user' : 'assistant',
+                parts: await Promise.all([message.message].flat().map(async (part) => {
+                    if (typeof part === 'string') {
+                        return { text: part };
+                    }
+                    else if (part.type === 'image') {
+                        return {
+                            inline_data: {
+                                mime_type: part.mediaType,
+                                data: (await uint8ArrayToBase64(part.data)),
+                            },
+                        };
+                    }
+                    else {
+                        throw new Error(`Google Vertex AI does not support message parts of type ${part.type}`);
+                    }
+                })),
+            };
+        }));
+        let { maxTokens } = data;
+        const tokenizerInfo = {
+            node: context.node,
+            model,
+            endpoint: undefined,
+        };
+        // TODO Better token counting for Google models.
+        const tokenCount = await context.tokenizer.getTokenCountForMessages(messages, undefined, tokenizerInfo);
+        if (googleModels[model] && tokenCount >= googleModels[model].maxTokens) {
+            throw new Error(`The model ${model} can only handle ${googleModels[model].maxTokens} tokens, but ${tokenCount} were provided in the prompts alone.`);
+        }
+        if (googleModels[model] && tokenCount + maxTokens > googleModels[model].maxTokens) {
+            const message = `The model can only handle a maximum of ${googleModels[model].maxTokens} tokens, but the prompts and max tokens together exceed this limit. The max tokens has been reduced to ${googleModels[model].maxTokens - tokenCount}.`;
+            addWarning(output, message);
+            maxTokens = Math.floor((googleModels[model].maxTokens - tokenCount) * 0.95); // reduce max tokens by 5% to be safe, calculation is a little wrong.
+        }
+        const project = context.getPluginConfig('googleProjectId');
+        const location = context.getPluginConfig('googleRegion');
+        const applicationCredentials = context.getPluginConfig('googleApplicationCredentials');
+        if (project == null) {
+            throw new Error('Google Project ID is not defined.');
+        }
+        if (location == null) {
+            throw new Error('Google Region is not defined.');
+        }
+        if (applicationCredentials == null) {
+            throw new Error('Google Application Credentials is not defined.');
+        }
+        try {
+            return await retry(async () => {
+                const options = {
+                    prompt,
+                    model,
+                    temperature: useTopP ? undefined : temperature,
+                    top_p: useTopP ? topP : undefined,
+                    max_output_tokens: maxTokens,
+                };
+                const cacheKey = JSON.stringify(options);
+                if (data.cache) {
+                    const cached = cache.get(cacheKey);
+                    if (cached) {
+                        return cached;
+                    }
+                }
+                const startTime = Date.now();
+                const chunks = streamChatCompletions({
+                    signal: context.signal,
+                    project,
+                    location,
+                    applicationCredentials,
+                    ...options,
+                });
+                const responseParts = [];
+                for await (const chunk of chunks) {
+                    if (!chunk.completion) {
+                        // Could be error for some reason 🤷‍♂️ but ignoring has worked for me so far.
+                        continue;
+                    }
+                    responseParts.push(chunk.completion);
+                    output['response'] = {
+                        type: 'string',
+                        value: responseParts.join('').trim(),
+                    };
+                    context.onPartialOutputs?.(output);
+                }
+                const endTime = Date.now();
+                output['all-messages'] = {
+                    type: 'chat-message[]',
+                    value: [
+                        ...messages,
+                        {
+                            type: 'assistant',
+                            message: responseParts.join('').trim() ?? '',
+                            function_call: undefined,
+                            function_calls: undefined,
+                        },
+                    ],
+                };
+                output['in-messages'] = {
+                    type: 'chat-message[]',
+                    value: messages,
+                };
+                if (responseParts.length === 0) {
+                    throw new Error('No response from Google');
+                }
+                output['requestTokens'] = { type: 'number', value: tokenCount };
+                const responseTokenCount = await context.tokenizer.getTokenCountForString(responseParts.join(''), tokenizerInfo);
+                output['responseTokens'] = { type: 'number', value: responseTokenCount };
+                // TODO
+                // const cost =
+                //   getCostForPrompt(completionMessages, model) + getCostForTokens(responseTokenCount, 'completion', model);
+                // output['cost' as PortId] = { type: 'number', value: cost };
+                const duration = endTime - startTime;
+                output['duration'] = { type: 'number', value: duration };
+                Object.freeze(output);
+                cache.set(cacheKey, output);
+                return output;
+            }, {
+                retries: 10,
+                maxRetryTime: 1000 * 60 * 5,
+                factor: 2.5,
+                minTimeout: 500,
+                maxTimeout: 5000,
+                randomize: true,
+                signal: context.signal,
+                onFailedAttempt(err) {
+                    context.trace(`ChatGoogleNode failed, retrying: ${err.toString()}`);
+                    if (context.signal.aborted) {
+                        throw new Error('Aborted');
+                    }
+                },
+            });
+        }
+        catch (error) {
+            context.trace(getError(error).stack ?? 'Missing stack');
+            throw new Error(`Error processing ChatGoogleNode: ${error.message}`);
+        }
+    },
+};
+export const chatGoogleNode = pluginNodeDefinition(ChatGoogleNodeImpl, 'Chat');
+export function getChatGoogleNodeMessages(inputs) {
+    const prompt = inputs['prompt'];
+    if (!prompt) {
+        throw new Error('Prompt is required');
+    }
+    const messages = match(prompt)
+        .with({ type: 'chat-message' }, (p) => [p.value])
+        .with({ type: 'chat-message[]' }, (p) => p.value)
+        .with({ type: 'string' }, (p) => [{ type: 'user', message: p.value }])
+        .with({ type: 'string[]' }, (p) => p.value.map((v) => ({ type: 'user', message: v })))
+        .otherwise((p) => {
+        if (isArrayDataValue(p)) {
+            const stringValues = p.value.map((v) => coerceType({
+                type: getScalarTypeOf(p.type),
+                value: v,
+            }, 'string'));
+            return stringValues.filter((v) => v != null).map((v) => ({ type: 'user', message: v }));
+        }
+        const coercedMessage = coerceType(p, 'chat-message');
+        if (coercedMessage != null) {
+            return [coercedMessage];
+        }
+        const coercedString = coerceType(p, 'string');
+        return coercedString != null ? [{ type: 'user', message: coerceType(p, 'string') }] : [];
+    });
+    return { messages };
+}

package/dist/esm/plugins/google/plugin.js

@@ -0,0 +1,32 @@
+import {} from '../../index.js';
+import { chatGoogleNode } from './nodes/ChatGoogleNode.js';
+export const googlePlugin = {
+    id: 'google',
+    name: 'Google',
+    register: (register) => {
+        register(chatGoogleNode);
+    },
+    configSpec: {
+        googleProjectId: {
+            type: 'string',
+            label: 'Google Project ID',
+            description: 'The Google project ID.',
+            pullEnvironmentVariable: 'GCP_PROJECT',
+            helperText: 'You may also set the GCP_PROJECT environment variable.',
+        },
+        googleRegion: {
+            type: 'string',
+            label: 'Google Region',
+            description: 'The Google region.',
+            pullEnvironmentVariable: 'GCP_REGION',
+            helperText: 'You may also set the GCP_REGION environment variable.',
+        },
+        googleApplicationCredentials: {
+            type: 'string',
+            label: 'Google Application Credentials',
+            description: 'The path with the JSON file that contains your credentials.',
+            pullEnvironmentVariable: 'GOOGLE_APPLICATION_CREDENTIALS',
+            helperText: 'You may also set the GOOGLE_APPLICATION_CREDENTIALS environment variable. See https://cloud.google.com/vertex-ai/docs/start/client-libraries for more info.',
+        },
+    },
+};

package/dist/esm/plugins/huggingface/index.js

@@ -0,0 +1,2 @@
+import { huggingFacePlugin } from './plugin.js';
+export default huggingFacePlugin;

package/dist/esm/plugins/huggingface/nodes/ChatHuggingFace.js

@@ -0,0 +1,243 @@
+import { nanoid } from 'nanoid/non-secure';
+import {} from '../../../index.js';
+import { HfInference, HfInferenceEndpoint } from '@huggingface/inference';
+import { getInputOrData } from '../../../utils/inputs.js';
+import { coerceType } from '../../../utils/coerceType.js';
+import { dedent } from '../../../utils/misc.js';
+import { pluginNodeDefinition } from '../../../model/NodeDefinition.js';
+export const ChatHuggingFaceNodeImpl = {
+    create() {
+        return {
+            id: nanoid(),
+            type: 'chatHuggingFace',
+            data: {
+                model: '',
+                temperature: 0.5,
+                maxNewTokens: 1024,
+                doSample: false,
+            },
+            title: 'Chat (Hugging Face)',
+            visualData: {
+                x: 0,
+                y: 0,
+                width: 300,
+            },
+        };
+    },
+    getUIData() {
+        return {
+            group: ['AI', 'Hugging Face'],
+            contextMenuTitle: 'Chat (Hugging Face)',
+            infoBoxTitle: 'Chat (Hugging Face) Node',
+            infoBoxBody: 'Chat, using the hugging face inference API',
+        };
+    },
+    getInputDefinitions(data) {
+        const inputs = [];
+        inputs.push({
+            id: 'prompt',
+            dataType: 'string',
+            title: 'Prompt',
+            required: true,
+        });
+        if (data.useModelInput) {
+            inputs.push({
+                id: 'model',
+                dataType: 'string',
+                title: 'Model',
+            });
+        }
+        if (data.useEndpointInput) {
+            inputs.push({
+                id: 'endpoint',
+                dataType: 'string',
+                title: 'Endpoint',
+            });
+        }
+        if (data.useTemperatureInput) {
+            inputs.push({
+                id: 'temperature',
+                dataType: 'number',
+                title: 'Temperature',
+            });
+        }
+        if (data.useMaxNewTokensInput) {
+            inputs.push({
+                id: 'maxNewTokens',
+                dataType: 'number',
+                title: 'Max New Tokens',
+            });
+        }
+        if (data.useDoSampleInput) {
+            inputs.push({
+                id: 'doSample',
+                dataType: 'boolean',
+                title: 'Do Sample',
+            });
+        }
+        if (data.useMaxTimeInput) {
+            inputs.push({
+                id: 'maxTime',
+                dataType: 'number',
+                title: 'Max Time (s)',
+            });
+        }
+        if (data.useRepetitionPenaltyInput) {
+            inputs.push({
+                id: 'repetitionPenalty',
+                dataType: 'number',
+                title: 'Repetition Penalty',
+            });
+        }
+        if (data.useTopPInput) {
+            inputs.push({
+                id: 'topP',
+                dataType: 'number',
+                title: 'Top P',
+            });
+        }
+        if (data.useTopKInput) {
+            inputs.push({
+                id: 'topK',
+                dataType: 'number',
+                title: 'Top K',
+            });
+        }
+        return inputs;
+    },
+    getOutputDefinitions() {
+        return [
+            {
+                id: 'output',
+                dataType: 'string',
+                title: 'Output',
+            },
+        ];
+    },
+    getEditors() {
+        return [
+            {
+                type: 'string',
+                label: 'Model',
+                dataKey: 'model',
+                useInputToggleDataKey: 'useModelInput',
+            },
+            {
+                type: 'string',
+                label: 'Endpoint',
+                dataKey: 'endpoint',
+                useInputToggleDataKey: 'useEndpointInput',
+            },
+            {
+                type: 'number',
+                label: 'Temperature (0-100)',
+                dataKey: 'temperature',
+                useInputToggleDataKey: 'useTemperatureInput',
+                min: 0,
+                step: 50,
+                allowEmpty: true,
+            },
+            {
+                type: 'number',
+                label: 'Max New Tokens',
+                dataKey: 'maxNewTokens',
+                useInputToggleDataKey: 'useMaxNewTokensInput',
+                min: 0,
+                step: 1,
+            },
+            {
+                type: 'toggle',
+                label: 'Do Sample',
+                dataKey: 'doSample',
+                useInputToggleDataKey: 'useDoSampleInput',
+            },
+            {
+                type: 'number',
+                label: 'Max Time (s)',
+                dataKey: 'maxTime',
+                useInputToggleDataKey: 'useMaxTimeInput',
+                allowEmpty: true,
+            },
+            {
+                type: 'number',
+                label: 'Repetition Penalty (0-100)',
+                dataKey: 'repetitionPenalty',
+                useInputToggleDataKey: 'useRepetitionPenaltyInput',
+                allowEmpty: true,
+            },
+            {
+                type: 'number',
+                label: 'Top P (0-100)',
+                dataKey: 'topP',
+                useInputToggleDataKey: 'useTopPInput',
+                allowEmpty: true,
+            },
+            {
+                type: 'number',
+                label: 'Top K (0-100)',
+                dataKey: 'topK',
+                useInputToggleDataKey: 'useTopKInput',
+                allowEmpty: true,
+            },
+        ];
+    },
+    getBody(data) {
+        return dedent `
+      ${data.endpoint || data.useEndpointInput
+            ? `Endpoint: ${data.useEndpointInput ? '(Using Input)' : 'Yes'}`
+            : `Model: ${data.useModelInput ? '(Using Input)' : data.model}`}
+      ${data.useTemperatureInput
+            ? 'Temperature: (Using Input)'
+            : data.temperature != null
+                ? `Temperature: ${data.temperature}`
+                : ''}
+      Max New Tokens: ${data.useMaxNewTokensInput ? '(Using Input)' : data.maxNewTokens}
+    `;
+    },
+    async process(data, inputData, context) {
+        const accessToken = context.getPluginConfig('huggingFaceAccessToken');
+        const prompt = coerceType(inputData['prompt'], 'string');
+        const endpoint = getInputOrData(data, inputData, 'endpoint');
+        const model = getInputOrData(data, inputData, 'model');
+        const temperature = getInputOrData(data, inputData, 'temperature', 'number');
+        const maxNewTokens = getInputOrData(data, inputData, 'maxNewTokens', 'number');
+        const doSample = getInputOrData(data, inputData, 'doSample', 'boolean');
+        const maxTime = getInputOrData(data, inputData, 'maxTime', 'number');
+        const repetitionPenalty = getInputOrData(data, inputData, 'repetitionPenalty', 'number');
+        const topP = getInputOrData(data, inputData, 'topP', 'number');
+        const topK = getInputOrData(data, inputData, 'topK', 'number');
+        const hf = endpoint ? new HfInferenceEndpoint(endpoint, accessToken) : new HfInference(accessToken);
+        const generationStream = hf.textGenerationStream({
+            inputs: prompt,
+            model,
+            parameters: {
+                temperature,
+                max_new_tokens: maxNewTokens,
+                do_sample: doSample,
+                max_time: maxTime,
+                repetition_penalty: repetitionPenalty,
+                top_p: topP,
+                top_k: topK,
+            },
+        });
+        const parts = [];
+        for await (const { token } of generationStream) {
+            if (!token.special) {
+                parts.push(token.text);
+            }
+            context.onPartialOutputs?.({
+                ['output']: {
+                    type: 'string',
+                    value: parts.join(''),
+                },
+            });
+        }
+        return {
+            ['output']: {
+                type: 'string',
+                value: parts.join(''),
+            },
+        };
+    },
+};
+export const chatHuggingFaceNode = pluginNodeDefinition(ChatHuggingFaceNodeImpl, 'Chat (Hugging Face)');