@alpic80/rivet-core 1.19.1-aidon.2 → 1.24.0-aidon.1

package/dist/esm/plugins/anthropic/nodes/ChatAnthropicNode.js

@@ -1,5 +1,5 @@
-import { uint8ArrayToBase64, } from '../../../index.js';
-import { anthropicModelOptions, anthropicModels, streamChatCompletions, AnthropicError, streamMessageApi, callMessageApi, } from '../anthropic.js';
+import {} from '../../../index.js';
+import { anthropicModelOptions, anthropicModels, streamChatCompletions, AnthropicError, streamMessageApi, } from '../anthropic.js';
 import { nanoid } from 'nanoid/non-secure';
 import { dedent } from 'ts-dedent';
 import retry from 'p-retry';
@@ -11,6 +11,8 @@ import { pluginNodeDefinition } from '../../../model/NodeDefinition.js';
 import { getScalarTypeOf, isArrayDataValue } from '../../../model/DataValue.js';
 import { assertNever } from '../../../utils/assertNever.js';
 import { isNotNull } from '../../../utils/genericUtilFunctions.js';
+import { uint8ArrayToBase64 } from '../../../utils/base64.js';
+import { getInputOrData } from '../../../utils/inputs.js';
 // Temporary
 const cache = new Map();
 export const ChatAnthropicNodeImpl = {
@@ -25,7 +27,7 @@ export const ChatAnthropicNodeImpl = {
                 width: 275,
             },
             data: {
-                model: 'claude-3-5-sonnet-20240620',
+                model: 'claude-3-7-sonnet-latest',
                 useModelInput: false,
                 temperature: 0.5,
                 useTemperatureInput: false,
@@ -43,6 +45,11 @@ export const ChatAnthropicNodeImpl = {
                 cache: false,
                 useAsGraphPartialOutput: true,
                 enableToolUse: false,
+                endpoint: '',
+                useEndpointInput: false,
+                overrideModel: undefined,
+                useOverrideModelInput: false,
+                enableCitations: false,
             },
         };
         return chartNode;
@@ -122,12 +129,20 @@ export const ChatAnthropicNodeImpl = {
             id: 'response',
             title: 'Response',
         });
+        if (data.enableCitations) {
+            outputs.push({
+                dataType: 'object[]',
+                id: 'citations',
+                title: 'Citations',
+                description: 'Citations from the response, if any.',
+            });
+        }
         if (data.enableToolUse) {
             outputs.push({
                 dataType: 'object[]',
                 id: 'function-calls',
-                title: 'Function Calls',
-                description: 'The function calls that were made, if any.',
+                title: 'Tool Calls',
+                description: 'The tool calls that were made, if any.',
             });
         }
         outputs.push({
@@ -139,7 +154,9 @@ export const ChatAnthropicNodeImpl = {
         return outputs;
     },
     getBody(data) {
-        const modelName = anthropicModels[data.model]?.displayName ?? 'Unknown Model';
+        const modelName = data.overrideModel
+            ? data.overrideModel
+            : anthropicModels[data.model]?.displayName ?? 'Unknown Model';
         return dedent `
       ${modelName}
       ${data.useTopP
@@ -152,65 +169,105 @@ export const ChatAnthropicNodeImpl = {
     getEditors() {
         return [
             {
-                type: 'dropdown',
-                label: 'Model',
-                dataKey: 'model',
-                useInputToggleDataKey: 'useModelInput',
-                options: anthropicModelOptions,
-            },
-            {
-                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: 'string',
-                label: 'Stop',
-                dataKey: 'stop',
-                useInputToggleDataKey: 'useStopInput',
-            },
-            {
-                type: 'toggle',
-                label: 'Cache (same inputs, same outputs)',
-                dataKey: 'cache',
+                type: 'group',
+                label: 'Parameters',
+                defaultOpen: true,
+                editors: [
+                    {
+                        type: 'dropdown',
+                        label: 'Model',
+                        dataKey: 'model',
+                        useInputToggleDataKey: 'useModelInput',
+                        options: anthropicModelOptions,
+                        disableIf: (d) => !!d.overrideModel?.trim(),
+                        helperMessage: (d) => (!!d.overrideModel ? `Model is overridden to: ${d.overrideModel}` : ''),
+                    },
+                    {
+                        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: 'string',
+                        label: 'Stop',
+                        dataKey: 'stop',
+                        useInputToggleDataKey: 'useStopInput',
+                    },
+                ],
             },
             {
-                type: 'toggle',
-                label: 'Use for subgraph partial output',
-                dataKey: 'useAsGraphPartialOutput',
+                type: 'group',
+                label: 'Tools',
+                editors: [
+                    {
+                        type: 'toggle',
+                        label: 'Enable Tool Use (disables streaming)',
+                        dataKey: 'enableToolUse',
+                    },
+                    {
+                        type: 'toggle',
+                        label: 'Enable Citations',
+                        dataKey: 'enableCitations',
+                    },
+                ],
             },
             {
-                type: 'toggle',
-                label: 'Enable Tool Use (disables streaming)',
-                dataKey: 'enableToolUse',
+                type: 'group',
+                label: 'Advanced',
+                editors: [
+                    {
+                        type: 'toggle',
+                        label: 'Cache (same inputs, same outputs)',
+                        dataKey: 'cache',
+                    },
+                    {
+                        type: 'toggle',
+                        label: 'Use for subgraph partial output',
+                        dataKey: 'useAsGraphPartialOutput',
+                    },
+                    {
+                        type: 'string',
+                        label: 'Endpoint',
+                        dataKey: 'endpoint',
+                        useInputToggleDataKey: 'useEndpointInput',
+                        helperMessage: 'Overrides the Anthropic API endpoint. Leave blank to use the default configured endpoint in settings, or https://api.anthropic.com/v1 if none is configured.',
+                    },
+                    {
+                        type: 'string',
+                        label: 'Override Model',
+                        dataKey: 'overrideModel',
+                        useInputToggleDataKey: 'useOverrideModelInput',
+                        helperMessage: 'Overrides the AI model used for the chat node to this value.',
+                    },
+                ],
             },
         ];
     },
@@ -225,14 +282,10 @@ export const ChatAnthropicNodeImpl = {
         };
     },
     async process(data, inputs, context) {
-        if (context.executor === 'browser') {
-            throw new Error('This node requires using the Node executor');
-        }
         const output = {};
-        const rawModel = data.useModelInput
-            ? coerceTypeOptional(inputs['model'], 'string') ?? data.model
-            : data.model;
-        const model = rawModel;
+        const rawModel = getInputOrData(data, inputs, 'model');
+        const overrideModel = getInputOrData(data, inputs, 'overrideModel');
+        const model = (overrideModel || rawModel);
         const temperature = data.useTemperatureInput
             ? coerceTypeOptional(inputs['temperature'], 'number') ?? data.temperature
             : data.temperature;
@@ -268,6 +321,9 @@ export const ChatAnthropicNodeImpl = {
         prompt += '\n\nAssistant:';
         // Get the "System" prompt input for Claude 3 models
         const system = data.model.startsWith('claude-3') ? getSystemPrompt(inputs) : undefined;
+        const systemInput = inputs['system'];
+        const includesCacheBreakpoint = rivetChatMessages.some((m) => m.isCacheBreakpoint) ||
+            (systemInput?.type === 'chat-message' && systemInput.value.isCacheBreakpoint);
         let { maxTokens } = data;
         const tokenizerInfo = {
             node: context.node,
@@ -306,7 +362,7 @@ export const ChatAnthropicNodeImpl = {
                     top_p: useTopP ? topP : undefined,
                     max_tokens: maxTokens ?? modelInfo.maxTokens,
                     stop_sequences: stop ? [stop] : undefined,
-                    system: system,
+                    system,
                     messages,
                     tools: tools
                         ? tools.map((tool) => ({ name: tool.name, description: tool.description, input_schema: tool.parameters }))
@@ -322,77 +378,79 @@ export const ChatAnthropicNodeImpl = {
                 }
                 const startTime = Date.now();
                 const apiKey = context.getPluginConfig('anthropicApiKey');
-                if (useMessageApi && data.enableToolUse) {
-                    // Streaming is not supported with tool usage.
-                    const response = await callMessageApi({
-                        apiKey: apiKey ?? '',
-                        ...messageOptions,
-                    });
-                    const { input_tokens: requestTokens, output_tokens: responseTokens } = response.usage;
-                    const responseText = response.content
-                        .map((c) => c.text)
-                        .filter(isNotNull)
-                        .join('');
-                    output['response'] = {
-                        type: 'string',
-                        value: responseText,
-                    };
-                    const functionCalls = response.content
-                        .filter((content) => content.name && content.id)
-                        .map((functionCall) => ({
-                        name: functionCall.name,
-                        arguments: functionCall.input, // Matches OpenAI ChatNode
-                        id: functionCall.id,
-                    }));
-                    if (functionCalls.length > 0) {
-                        output['function-calls'] = {
-                            type: 'object[]',
-                            value: functionCalls,
-                        };
-                    }
-                    output['all-messages'] = {
-                        type: 'chat-message[]',
-                        value: [
-                            ...rivetChatMessages,
-                            {
-                                type: 'assistant',
-                                message: responseText,
-                                function_call: functionCalls.length > 0
-                                    ? functionCalls.map((toolCall) => ({
-                                        name: toolCall.name,
-                                        arguments: JSON.stringify(toolCall.arguments),
-                                        id: toolCall.id,
-                                    }))[0]
-                                    : undefined,
-                                function_calls: functionCalls.map((toolCall) => ({
-                                    name: toolCall.name,
-                                    arguments: JSON.stringify(toolCall.arguments),
-                                    id: toolCall.id,
-                                })),
-                            },
-                        ],
-                    };
-                    output['requestTokens'] = { type: 'number', value: requestTokens ?? tokenCountEstimate };
-                    const responseTokenCount = responseTokens ?? context.tokenizer.getTokenCountForString(responseText, tokenizerInfo);
-                    output['responseTokens'] = { type: 'number', value: responseTokenCount };
-                }
-                else if (useMessageApi) {
+                const defaultApiEndpoint = context.getPluginConfig('anthropicApiEndpoint') || 'https://api.anthropic.com/v1';
+                const configuredEndpoint = getInputOrData(data, inputs, 'endpoint');
+                const apiEndpoint = configuredEndpoint?.trim() ? configuredEndpoint : defaultApiEndpoint;
+                if (useMessageApi) {
                     // Use the messages API for Claude 3 models
                     const chunks = streamMessageApi({
+                        apiEndpoint,
                         apiKey: apiKey ?? '',
                         signal: context.signal,
+                        beta: 'prompt-caching-2024-07-31',
                         ...messageOptions,
                     });
                     // Process the response chunks and update the output
                     const responseParts = [];
-                    let requestTokens = undefined, responseTokens = undefined;
+                    let requestTokens = undefined;
+                    let responseTokens = undefined;
+                    const citations = [];
+                    // Track tool calls
+                    const toolCalls = [];
+                    let currentToolCall = null;
+                    let accumulatedJsonString = '';
                     for await (const chunk of chunks) {
                         let completion = '';
                         if (chunk.type === 'content_block_start') {
-                            completion = chunk.content_block.text;
+                            if (chunk.content_block.type === 'text') {
+                                completion = chunk.content_block.text || '';
+                            }
+                            else if (chunk.content_block.type === 'tool_use') {
+                                currentToolCall = {
+                                    id: chunk.content_block.id,
+                                    name: chunk.content_block.name,
+                                    input: chunk.content_block.input || {},
+                                };
+                                accumulatedJsonString = '';
+                            }
                         }
                         else if (chunk.type === 'content_block_delta') {
-                            completion = chunk.delta.text;
+                            if (chunk.delta.type === 'text_delta') {
+                                completion = chunk.delta.text;
+                            }
+                            else if (chunk.delta.type === 'input_json_delta') {
+                                if (currentToolCall) {
+                                    accumulatedJsonString += chunk.delta.partial_json || '';
+                                    try {
+                                        // Try to parse the accumulated JSON
+                                        const parsedJson = JSON.parse(accumulatedJsonString);
+                                        currentToolCall.input = parsedJson;
+                                        accumulatedJsonString = '';
+                                    }
+                                    catch (e) {
+                                        // Not valid JSON yet, keep accumulating
+                                    }
+                                }
+                            }
+                            else if (chunk.delta.type === 'citations_delta') {
+                                citations.push(chunk.delta.citation);
+                            }
+                        }
+                        else if (chunk.type === 'content_block_stop') {
+                            if (currentToolCall) {
+                                if (accumulatedJsonString) {
+                                    try {
+                                        const parsedJson = JSON.parse(accumulatedJsonString);
+                                        currentToolCall.input = parsedJson;
+                                    }
+                                    catch (e) {
+                                        console.warn('Failed to parse tool call JSON input:', accumulatedJsonString);
+                                    }
+                                }
+                                toolCalls.push({ ...currentToolCall });
+                                currentToolCall = null;
+                                accumulatedJsonString = '';
+                            }
                         }
                         else if (chunk.type === 'message_start' && chunk.message?.usage?.input_tokens) {
                             requestTokens = chunk.message.usage.input_tokens;
@@ -400,14 +458,39 @@ export const ChatAnthropicNodeImpl = {
                         else if (chunk.type === 'message_delta' && chunk.delta?.usage?.output_tokens) {
                             responseTokens = chunk.delta.usage.output_tokens;
                         }
-                        if (!completion) {
-                            continue;
+                        if (completion) {
+                            responseParts.push(completion);
                         }
-                        responseParts.push(completion);
                         output['response'] = {
                             type: 'string',
                             value: responseParts.join('').trim(),
                         };
+                        if (toolCalls.length > 0) {
+                            output['function-calls'] = {
+                                type: 'object[]',
+                                value: toolCalls.map((tool) => ({
+                                    id: tool.id,
+                                    name: tool.name,
+                                    arguments: tool.input,
+                                })),
+                            };
+                        }
+                        else {
+                            output['function-calls'] = {
+                                type: 'control-flow-excluded',
+                                value: undefined,
+                            };
+                        }
+                        output['citations'] = {
+                            type: 'object[]',
+                            value: citations,
+                        };
+                        // Format function calls for the ChatMessage interface
+                        const functionCalls = toolCalls.map((tool) => ({
+                            name: tool.name,
+                            arguments: typeof tool.input === 'object' ? JSON.stringify(tool.input) : tool.input,
+                            id: tool.id,
+                        }));
                         output['all-messages'] = {
                             type: 'chat-message[]',
                             value: [
@@ -415,15 +498,16 @@ export const ChatAnthropicNodeImpl = {
                                 {
                                     type: 'assistant',
                                     message: responseParts.join('').trim(),
-                                    function_call: undefined,
-                                    function_calls: undefined,
+                                    function_call: functionCalls.length === 1 ? functionCalls[0] : undefined,
+                                    function_calls: functionCalls.length > 0 ? functionCalls : undefined,
                                 },
                             ],
                         };
                         context.onPartialOutputs?.(output);
                     }
-                    if (responseParts.length === 0) {
-                        throw new Error('No response from Anthropic');
+                    // Final validation
+                    if (responseParts.length === 0 && toolCalls.length === 0) {
+                        throw new Error('No response or tool calls received from Anthropic');
                     }
                     output['requestTokens'] = { type: 'number', value: requestTokens ?? tokenCountEstimate };
                     const responseTokenCount = responseTokens ?? (await context.tokenizer.getTokenCountForString(responseParts.join(''), tokenizerInfo));
@@ -432,6 +516,7 @@ export const ChatAnthropicNodeImpl = {
                 else {
                     // Use the normal chat completion method for non-Claude 3 models
                     const chunks = streamChatCompletions({
+                        apiEndpoint,
                         apiKey: apiKey ?? '',
                         signal: context.signal,
                         ...completionOptions,
@@ -513,22 +598,36 @@ export const ChatAnthropicNodeImpl = {
 };
 export const chatAnthropicNode = pluginNodeDefinition(ChatAnthropicNodeImpl, 'Chat');
 export function getSystemPrompt(inputs) {
-    const system = coerceTypeOptional(inputs['system'], 'string');
+    const systemInput = inputs['system'];
+    const system = coerceTypeOptional(systemInput, 'string');
     if (system) {
-        return system;
+        return [
+            {
+                type: 'text',
+                text: system,
+                cache_control: systemInput?.type === 'chat-message'
+                    ? systemInput.value.isCacheBreakpoint
+                        ? { type: 'ephemeral' }
+                        : null
+                    : null,
+            },
+        ];
     }
     const prompt = inputs['prompt'];
     if (prompt && prompt.type === 'chat-message[]') {
-        const systemMessage = prompt.value.find((message) => message.type === 'system');
-        if (systemMessage) {
-            if (typeof systemMessage.message === 'string') {
-                return systemMessage.message;
-            }
-            else if (Array.isArray(systemMessage.message)) {
-                return systemMessage.message.filter((p) => typeof p === 'string').join('');
-            }
+        const systemMessages = prompt.value.filter((message) => message.type === 'system');
+        if (systemMessages.length) {
+            const converted = systemMessages.map((message) => {
+                return {
+                    type: 'text',
+                    text: coerceType({ type: 'chat-message', value: message }, 'string'),
+                    cache_control: message.isCacheBreakpoint ? { type: 'ephemeral' } : null,
+                };
+            });
+            return converted;
         }
     }
+    return undefined;
 }
 function getChatMessages(inputs) {
     const prompt = inputs['prompt'];
@@ -567,8 +666,7 @@ export async function chatMessagesToClaude3ChatMessages(chatMessages) {
             message.content[0].type === 'tool_result') {
             const last = acc.at(-1);
             if (last?.role === 'user' && Array.isArray(last.content) && last.content.every((c) => c.type === 'tool_result')) {
-                const mc = message.content;
-                const content = last.content.concat(mc);
+                const content = last.content.concat(message.content);
                 return [...acc.slice(0, -1), { ...last, content }];
             }
         }
@@ -592,12 +690,13 @@ async function chatMessageToClaude3ChatMessage(message) {
                     type: 'tool_result',
                     tool_use_id: message.name,
                     content: content.length === 1 ? content[0].text : content,
+                    cache_control: message.isCacheBreakpoint ? { type: 'ephemeral' } : null,
                 },
             ],
         };
     }
     const content = Array.isArray(message.message)
-        ? await Promise.all(message.message.map(chatMessageContentToClaude3ChatMessage))
+        ? await Promise.all(message.message.map((part) => chatMessageContentToClaude3ChatMessage(part)))
         : [await chatMessageContentToClaude3ChatMessage(message.message)];
     if (message.type === 'assistant' && message.function_calls) {
         content.push(...message.function_calls.map((fc) => ({
@@ -605,6 +704,7 @@ async function chatMessageToClaude3ChatMessage(message) {
             id: fc.id,
             name: fc.name,
             input: JSON.parse(fc.arguments),
+            cache_control: message.isCacheBreakpoint ? { type: 'ephemeral' } : null,
         })));
     }
     else if (message.type === 'assistant' && message.function_call) {
@@ -613,8 +713,13 @@ async function chatMessageToClaude3ChatMessage(message) {
             id: message.function_call.id,
             name: message.function_call.name,
             input: JSON.parse(message.function_call.arguments),
+            cache_control: message.isCacheBreakpoint ? { type: 'ephemeral' } : null,
         });
     }
+    // If the message is a cache breakpoint, cache using the last content item of the message
+    if (message.isCacheBreakpoint) {
+        content.at(-1).cache_control = { type: 'ephemeral' };
+    }
     return {
         role: message.type,
         content,
@@ -625,6 +730,7 @@ async function chatMessageContentToClaude3ChatMessage(content) {
         return {
             type: 'text',
             text: content,
+            cache_control: null, // set later
         };
     }
     switch (content.type) {
@@ -636,9 +742,23 @@ async function chatMessageContentToClaude3ChatMessage(content) {
                     media_type: content.mediaType,
                     data: (await uint8ArrayToBase64(content.data)) ?? '',
                 },
+                cache_control: null, // set later
             };
         case 'url':
             throw new Error('unable to convert urls for Claude');
+        case 'document':
+            return {
+                type: 'document',
+                source: {
+                    type: 'base64',
+                    data: (await uint8ArrayToBase64(content.data)) ?? '',
+                    media_type: content.mediaType,
+                },
+                title: content.title?.trim() ? content.title.trim() : undefined,
+                context: content.context?.trim() ? content.context.trim() : undefined,
+                citations: content.enableCitations ? { enabled: true } : undefined,
+                cache_control: null, // set later
+            };
         default:
             assertNever(content);
     }

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

@@ -8,11 +8,19 @@ export const anthropicPlugin = {
     },
     configSpec: {
         anthropicApiKey: {
-            type: 'string',
+            type: 'secret',
             label: 'Anthropic API Key',
             description: 'The API key for the Anthropic service.',
             pullEnvironmentVariable: 'ANTHROPIC_API_KEY',
             helperText: 'You may also set the ANTHROPIC_API_KEY environment variable.',
         },
+        anthropicApiEndpoint: {
+            type: 'string',
+            label: 'Anthropic API Endpoint',
+            description: 'The API endpoint for the Anthropic service.',
+            pullEnvironmentVariable: 'ANTHROPIC_API_ENDPOINT',
+            helperText: 'Defaults to https://api.anthropic.com/v1. You may also set the ANTHROPIC_API_ENDPOINT environment variable.',
+            default: 'https://api.anthropic.com/v1',
+        },
     },
 };

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

@@ -49,9 +49,9 @@ export const runGentraceTests = async (gentracePipelineSlug, settings, project,
         }, rivetFormattedInputs);
         const fullRecording = recorder.getRecording();
         const stepRuns = convertRecordingToStepRuns(fullRecording, project, graphId);
-        stepRuns.forEach((stepRun) => {
-            runner.addStepRunNode(stepRun);
-        });
+        for (const stepRun of stepRuns) {
+            await runner.addStepRunNode(stepRun);
+        }
         if (stepRuns.length === 0) {
             throw new Error('No Rivet steps found. You need operations which are not Graph Input or Graph Output nodes.');
         }
@@ -77,9 +77,9 @@ export const runRemoteGentraceTests = async (gentracePipelineSlug, settings, pro
         const rivetFormattedInputs = mapValues(testCase.inputs, inferType);
         const fullRecording = await runAndRecord(rivetFormattedInputs);
         const stepRuns = convertRecordingToStepRuns(fullRecording, project, graphId);
-        stepRuns.forEach((stepRun) => {
-            runner.addStepRunNode(stepRun);
-        });
+        for (const stepRun of stepRuns) {
+            await runner.addStepRunNode(stepRun);
+        }
         if (stepRuns.length === 0) {
             throw new Error('No Rivet steps found. You need operations which are not Graph Input or Graph Output nodes.');
         }