UNPKG

@just-every/ensemble

Version:

LLM provider abstraction layer with unified streaming interface

838 lines 40 kB
import Anthropic from '@anthropic-ai/sdk'; import { v4 as uuidv4 } from 'uuid'; import { createCitationTracker, formatCitation, generateFootnotes } from '../utils/citation_tracker.js'; function formatWebSearchResults(results) { if (!Array.isArray(results)) return ''; return results .filter(r => r.type === 'web_search_result') .map((r, i) => `${i + 1}. ${r.title || 'Untitled'}${r.url}`) .join('\n'); } import { BaseModelProvider } from './base_provider.js'; import { costTracker } from '../index.js'; import { log_llm_error, log_llm_request, log_llm_response } from '../utils/llm_logger.js'; import { isPaused } from '../utils/pause_controller.js'; import { findModel } from '../data/model_data.js'; import { appendMessageWithImage, resizeAndTruncateForClaude } from '../utils/image_utils.js'; import { bufferDelta, flushBufferedDeltas } from '../utils/delta_buffer.js'; import { hasEventHandler } from '../utils/event_controller.js'; const THINKING_BUDGET_CONFIGS = { '-low': 0, '-medium': 8000, '-high': 15000, '-max': 30000, }; function contentToString(content) { if (content) { if (Array.isArray(content)) { let results = ''; for (const eachContent of content) { const convertedContent = contentToString(eachContent); if (convertedContent.length > 0) { if (results.length > 0) { results += '\n\n'; } results += convertedContent; } } return results.trim(); } else if (typeof content === 'string') { return content.trim(); } else if (typeof content.text === 'string') { return content.text.trim(); } return JSON.stringify(content); } return ''; } async function resolveAsyncEnums(params) { if (!params || typeof params !== 'object') { return params; } const resolved = { ...params }; if (resolved.properties) { const resolvedProps = {}; for (const [key, value] of Object.entries(resolved.properties)) { if (value && typeof value === 'object') { const propCopy = { ...value }; if (typeof propCopy.enum === 'function') { try { const enumValue = await propCopy.enum(); if (Array.isArray(enumValue) && enumValue.length > 0) { propCopy.enum = enumValue; } else { delete propCopy.enum; } } catch { delete propCopy.enum; } } resolvedProps[key] = await resolveAsyncEnums(propCopy); } else { resolvedProps[key] = value; } } resolved.properties = resolvedProps; } return resolved; } async function convertToClaudeTools(tools) { return await Promise.all(tools.map(async (tool) => { if (tool.definition.function.name === 'claude_web_search') { return { type: 'web_search_20250305', name: 'web_search', }; } return { name: tool.definition.function.name, description: tool.definition.function.description, input_schema: await resolveAsyncEnums(tool.definition.function.parameters), }; })); } function getImageMediaType(imageData) { if (imageData.includes('data:image/jpeg')) return 'image/jpeg'; if (imageData.includes('data:image/png')) return 'image/png'; if (imageData.includes('data:image/gif')) return 'image/gif'; if (imageData.includes('data:image/webp')) return 'image/webp'; return 'image/jpeg'; } function cleanBase64Data(imageData) { return imageData.replace(/^data:image\/[a-z]+;base64,/, ''); } async function addImagesToInput(input, images) { for (const [, imageData] of Object.entries(images)) { const processedImageData = await resizeAndTruncateForClaude(imageData); const mediaType = getImageMediaType(processedImageData); const cleanedImageData = cleanBase64Data(processedImageData); input.push({ type: 'image', source: { type: 'base64', media_type: mediaType, data: cleanedImageData, }, }); } return input; } async function convertToClaudeMessage(model, role, content, msg, result) { if (!msg) return null; if (msg.type === 'function_call') { let inputArgs = {}; try { const argsString = msg.arguments || '{}'; if (argsString.includes('}{')) { console.warn(`Malformed concatenated JSON arguments for ${msg.name}: ${argsString}`); const firstBraceIndex = argsString.indexOf('{'); const firstCloseBraceIndex = argsString.indexOf('}') + 1; if (firstBraceIndex !== -1 && firstCloseBraceIndex > firstBraceIndex) { const firstJsonStr = argsString.substring(firstBraceIndex, firstCloseBraceIndex); try { inputArgs = JSON.parse(firstJsonStr); console.log(`Successfully extracted first JSON object: ${firstJsonStr}`); } catch (innerE) { console.error(`Failed to parse extracted JSON: ${firstJsonStr}`, innerE); inputArgs = {}; } } else { inputArgs = {}; } } else { inputArgs = JSON.parse(argsString); } } catch (e) { console.error(`Error parsing function call arguments for ${msg.name}: ${msg.arguments}`, e); inputArgs = {}; } const toolUseBlock = { type: 'tool_use', id: msg.call_id, name: msg.name, input: inputArgs, }; return { role: 'assistant', content: [toolUseBlock] }; } else if (msg.type === 'function_call_output') { const toolResultBlock = { type: 'tool_result', tool_use_id: msg.call_id, content: msg.output || '', ...(msg.status === 'incomplete' ? { is_error: true } : {}), }; let contentBlocks = []; contentBlocks = await appendMessageWithImage(model, contentBlocks, toolResultBlock, 'content', addImagesToInput); return { role: 'user', content: contentBlocks }; } else if (msg.type === 'thinking') { if (!content) { return null; } if ('signature' in msg && msg.signature) { return { role: 'assistant', content: [ { type: 'thinking', thinking: content.trim(), signature: msg.signature, }, ], }; } return { role: 'assistant', content: 'Thinking: ' + content.trim() }; } else { if (!content) { return null; } let messageRole = role; if (messageRole === 'developer') { if (!result?.length) { messageRole = 'system'; } else { messageRole = 'user'; } } if (!['user', 'assistant', 'system'].includes(messageRole)) { messageRole = 'user'; } let contentBlocks = []; contentBlocks = await appendMessageWithImage(model, contentBlocks, { type: 'text', text: content, }, 'text', addImagesToInput); return { role: messageRole, content: contentBlocks, }; } } export class ClaudeProvider extends BaseModelProvider { _client; apiKey; constructor(apiKey) { super('anthropic'); this.apiKey = apiKey; } get client() { if (!this._client) { const apiKey = this.apiKey || process.env.ANTHROPIC_API_KEY; if (!apiKey) { throw new Error('Failed to initialize Claude client. Make sure ANTHROPIC_API_KEY is set.'); } this._client = new Anthropic({ apiKey: apiKey, }); } return this._client; } async prepareClaudeMessages(messages, modelId, thinkingEnabled = false) { const result = []; const seenToolUseIds = new Set(); for (const msg of messages) { const role = 'role' in msg && msg.role !== 'developer' ? msg.role : 'system'; let content = ''; if ('content' in msg) { content = contentToString(msg.content); } const structuredMsg = await convertToClaudeMessage(modelId, role, content, msg, result); if (structuredMsg) { if (structuredMsg.role === 'assistant' && Array.isArray(structuredMsg.content)) { let hasDuplicateToolUse = false; for (const contentBlock of structuredMsg.content) { if (contentBlock.type === 'tool_use') { if (seenToolUseIds.has(contentBlock.id)) { console.warn(`Skipping duplicate tool_use ID: ${contentBlock.id}`); hasDuplicateToolUse = true; break; } else { seenToolUseIds.add(contentBlock.id); } } } if (!hasDuplicateToolUse) { result.push(structuredMsg); } } else { result.push(structuredMsg); } } } if (thinkingEnabled && result.length > 0) { const convertedToolIds = new Set(); for (let i = 0; i < result.length; i++) { const msg = result[i]; if (msg.role === 'assistant' && Array.isArray(msg.content)) { const hasToolUse = msg.content.some(block => block.type === 'tool_use'); if (hasToolUse) { const hasThinkingBlock = msg.content.length > 0 && (msg.content[0].type === 'thinking' || msg.content[0].type === 'redacted_thinking'); if (!hasThinkingBlock) { const toolUseBlocks = msg.content.filter(block => block.type === 'tool_use'); toolUseBlocks.forEach(block => { if (block.id) { convertedToolIds.add(block.id); } }); const toolCalls = toolUseBlocks .map(block => { const args = typeof block.input === 'string' ? block.input : JSON.stringify(block.input); return `Called tool '${block.name}' with arguments: ${args}`; }) .join('\n'); msg.role = 'user'; msg.content = [ { type: 'text', text: `[Previous assistant action]\n${toolCalls}`, }, ]; } } } } for (let i = 0; i < result.length; i++) { const msg = result[i]; if (msg.role === 'user' && Array.isArray(msg.content)) { for (const block of msg.content) { if (block.type === 'text' && typeof block.text === 'string') { const toolUseMatches = block.text.matchAll(/"id"\s*:\s*"(call_[^"]+)"/g); for (const match of toolUseMatches) { if (match[1]) { convertedToolIds.add(match[1]); } } } } } } for (let i = 0; i < result.length; i++) { const msg = result[i]; if (msg.role === 'user' && Array.isArray(msg.content)) { const convertedBlocks = []; let hasConvertedToolResult = false; for (const block of msg.content) { if (block.type === 'tool_result' && convertedToolIds.has(block.tool_use_id)) { hasConvertedToolResult = true; convertedBlocks.push({ type: 'text', text: `[Tool Result for ${block.tool_use_id}]\n${block.content || '(empty result)'}`, }); } else { convertedBlocks.push(block); } } if (hasConvertedToolResult) { msg.content = convertedBlocks; } } } for (let i = 1; i < result.length; i++) { const prevMsg = result[i - 1]; const currentMsg = result[i]; if (prevMsg.role === 'assistant' && currentMsg.role === 'assistant') { let hasThinkingBlock = false; if (Array.isArray(currentMsg.content)) { hasThinkingBlock = currentMsg.content.length > 0 && (currentMsg.content[0].type === 'thinking' || currentMsg.content[0].type === 'redacted_thinking'); } if (!hasThinkingBlock) { const contentStr = contentToString(currentMsg.content); currentMsg.role = 'user'; currentMsg.content = [ { type: 'text', text: `Previous thoughts:\n\n${contentStr}`, }, ]; } } } for (let i = 0; i < result.length; i++) { const msg = result[i]; if (msg.role === 'assistant' && Array.isArray(msg.content)) { const hasThinkingBlock = msg.content.length > 0 && (msg.content[0].type === 'thinking' || msg.content[0].type === 'redacted_thinking'); if (!hasThinkingBlock && msg.content.every(block => block.type === 'text')) { const contentStr = contentToString(msg.content); msg.role = 'user'; msg.content = [ { type: 'text', text: `[Previous assistant response]\n${contentStr}`, }, ]; } } } } return result; } async *createResponseStream(messages, model, agent) { let totalInputTokens = 0; let totalOutputTokens = 0; let totalCacheCreationInputTokens = 0; let totalCacheReadInputTokens = 0; let streamCompletedSuccessfully = false; let messageCompleteYielded = false; let requestId; try { const { getToolsFromAgent } = await import('../utils/agent.js'); const tools = agent ? await getToolsFromAgent(agent) : []; const settings = agent?.modelSettings; let headers = undefined; if (model.startsWith('claude-sonnet-4') || model.startsWith('claude-opus-4')) { headers = { 'anthropic-beta': 'interleaved-thinking-2025-05-14', }; } let thinking = undefined; let thinkingSet = false; for (const [suffix, budget] of Object.entries(THINKING_BUDGET_CONFIGS)) { if (model.endsWith(suffix)) { thinkingSet = true; if (budget > 0) { thinking = { type: 'enabled', budget_tokens: budget, }; } model = model.slice(0, -suffix.length); break; } } const modelData = findModel(model); let max_tokens = settings?.max_tokens || modelData?.features?.max_output_tokens || 8192; if (modelData?.features?.max_output_tokens) { max_tokens = Math.min(max_tokens, modelData.features.max_output_tokens); } if (!thinkingSet && (model.startsWith('claude-sonnet-4') || model.startsWith('claude-opus-4') || model.startsWith('claude-3-7-sonnet'))) { thinking = { type: 'enabled', budget_tokens: 8000, }; } if (settings?.json_schema) { messages.push({ type: 'message', role: 'system', content: `Your response MUST be a valid JSON object that conforms to this schema:\n${JSON.stringify(settings.json_schema, null, 2)}`, }); } const thinkingEnabled = thinking !== undefined && thinking.type === 'enabled'; const claudeMessages = await this.prepareClaudeMessages(messages, model, thinkingEnabled); const systemPrompt = claudeMessages.reduce((acc, msg) => { if (msg.role === 'system' && msg.content) { if (acc.length > 0) { acc += '\n\n'; } acc += contentToString(msg.content); } return acc; }, ''); const requestParams = { model: model, messages: claudeMessages.filter(m => m.role === 'user' || m.role === 'assistant'), ...(systemPrompt ? { system: systemPrompt.trim() } : {}), stream: true, max_tokens, ...(thinking ? { thinking } : {}), ...(settings?.temperature !== undefined ? { temperature: settings.temperature } : {}), }; if (tools && tools.length > 0) { requestParams.tools = await convertToClaudeTools(tools); } if (!requestParams.messages || requestParams.messages.length === 0) { console.warn('Claude API Warning: No user or assistant messages provided after filtering. Adding default message.'); requestParams.messages = [ { role: 'user', content: "Let's think this through step by step.", }, ]; } requestId = log_llm_request(agent.agent_id, 'anthropic', model, requestParams, new Date()); let currentToolCall = null; let toolCallStarted = false; let accumulatedSignature = ''; let accumulatedThinking = ''; let accumulatedContent = ''; const messageId = uuidv4(); let deltaPosition = 0; const deltaBuffers = new Map(); const citationTracker = createCitationTracker(); const { waitWhilePaused } = await import('../utils/pause_controller.js'); await waitWhilePaused(100, agent.abortSignal); const stream = await this.client.messages.create(requestParams, { ...(headers ? { headers } : {}), }); const events = []; try { for await (const event of stream) { events.push(event); if (isPaused()) { console.log(`[Claude] System paused during stream for model ${model}. Waiting...`); await waitWhilePaused(100, agent.abortSignal); console.log(`[Claude] System resumed, continuing stream for model ${model}`); } if (event.type === 'message_start' && event.message?.usage) { const usage = event.message.usage; totalInputTokens += usage.input_tokens || 0; totalOutputTokens += usage.output_tokens || 0; totalCacheCreationInputTokens += usage.cache_creation_input_tokens || 0; totalCacheReadInputTokens += usage.cache_read_input_tokens || 0; } else if (event.type === 'message_delta' && event.usage) { const usage = event.usage; totalInputTokens += usage.input_tokens || 0; totalOutputTokens += usage.output_tokens || 0; totalCacheCreationInputTokens += usage.cache_creation_input_tokens || 0; totalCacheReadInputTokens += usage.cache_read_input_tokens || 0; } if (event.type === 'content_block_delta') { if (event.delta.type === 'signature_delta' && event.delta.signature) { accumulatedSignature += event.delta.signature; } else if (event.delta.type === 'thinking_delta' && event.delta.thinking) { yield { type: 'message_delta', content: '', thinking_content: event.delta.thinking, message_id: messageId, order: deltaPosition++, }; accumulatedThinking += event.delta.thinking; } else if (event.delta.type === 'text_delta' && event.delta.text) { for (const ev of bufferDelta(deltaBuffers, messageId, event.delta.text, content => ({ type: 'message_delta', content, message_id: messageId, order: deltaPosition++, }))) { yield ev; } accumulatedContent += event.delta.text; } else if (event.delta.type === 'input_json_delta' && currentToolCall && event.delta.partial_json) { try { if (!currentToolCall.function._partialArguments) { currentToolCall.function._partialArguments = ''; } currentToolCall.function._partialArguments += event.delta.partial_json; yield { type: 'tool_delta', tool_call: { ...currentToolCall, function: { ...currentToolCall.function, arguments: '{}', }, }, }; } catch (err) { console.error('Error processing tool_use delta (input_json_delta):', err, event); } } else if (event.delta.type === 'citations_delta' && event.delta.citation) { const citationMarker = formatCitation(citationTracker, { title: event.delta.citation.title, url: event.delta.citation.url, citedText: event.delta.citation.cited_text, }); yield { type: 'message_delta', content: citationMarker, message_id: messageId, order: deltaPosition++, }; accumulatedContent += citationMarker; } } else if (event.type === 'content_block_start' && event.content_block?.type === 'text') { if (event.content_block.text) { for (const ev of bufferDelta(deltaBuffers, messageId, event.content_block.text, content => ({ type: 'message_delta', content, message_id: messageId, order: deltaPosition++, }))) { yield ev; } accumulatedContent += event.content_block.text; } } else if (event.type === 'content_block_stop' && event.content_block?.type === 'text') { } else if (event.type === 'content_block_start' && event.content_block?.type === 'web_search_tool_result') { if (event.content_block.content) { const formatted = formatWebSearchResults(event.content_block.content); if (formatted) { yield { type: 'message_delta', content: '\n\nSearch Results:\n' + formatted + '\n', message_id: messageId, order: deltaPosition++, }; accumulatedContent += '\n\nSearch Results:\n' + formatted + '\n'; } } } else if (event.type === 'content_block_start' && event.content_block?.type === 'tool_use') { const toolUse = event.content_block; const toolId = toolUse.id || `call_${Date.now()}_${Math.random().toString(36).substr(2, 9)}`; const toolName = toolUse.name; const toolInput = toolUse.input !== undefined ? toolUse.input : {}; currentToolCall = { id: toolId, type: 'function', function: { name: toolName, arguments: typeof toolInput === 'string' ? toolInput : JSON.stringify(toolInput), }, }; toolCallStarted = false; } else if (event.type === 'content_block_stop' && event.content_block?.type === 'tool_use' && currentToolCall) { try { if (currentToolCall.function._partialArguments) { const partialArgs = currentToolCall.function._partialArguments; try { JSON.parse(partialArgs); currentToolCall.function.arguments = partialArgs; } catch (jsonError) { console.warn(`Invalid JSON in partial arguments for ${currentToolCall.function.name}: ${partialArgs}`, jsonError); if (partialArgs.includes('}{')) { const firstBraceIndex = partialArgs.indexOf('{'); const firstCloseBraceIndex = partialArgs.indexOf('}') + 1; if (firstBraceIndex !== -1 && firstCloseBraceIndex > firstBraceIndex) { const firstJsonStr = partialArgs.substring(firstBraceIndex, firstCloseBraceIndex); try { JSON.parse(firstJsonStr); currentToolCall.function.arguments = firstJsonStr; console.log(`Extracted valid JSON from partial arguments: ${firstJsonStr}`); } catch (extractError) { console.error(`Failed to extract valid JSON: ${firstJsonStr}`, extractError); currentToolCall.function.arguments = '{}'; } } else { currentToolCall.function.arguments = '{}'; } } else { currentToolCall.function.arguments = '{}'; } } delete currentToolCall.function._partialArguments; } yield { type: 'tool_start', tool_call: currentToolCall, }; toolCallStarted = true; } catch (err) { console.error('Error finalizing tool call:', err, event); } finally { currentToolCall = null; } } else if (event.type === 'message_stop') { if (event['amazon-bedrock-invocationMetrics']) { const metrics = event['amazon-bedrock-invocationMetrics']; totalInputTokens += metrics.inputTokenCount || 0; totalOutputTokens += metrics.outputTokenCount || 0; } else if (event.usage) { const usage = event.usage; totalInputTokens += usage.input_tokens || 0; totalOutputTokens += usage.output_tokens || 0; totalCacheCreationInputTokens += usage.cache_creation_input_tokens || 0; totalCacheReadInputTokens += usage.cache_read_input_tokens || 0; } if (currentToolCall && !toolCallStarted) { if (currentToolCall.function._partialArguments) { const partialArgs = currentToolCall.function._partialArguments; try { JSON.parse(partialArgs); currentToolCall.function.arguments = partialArgs; } catch (jsonError) { console.warn(`Invalid JSON in partial arguments at message_stop for ${currentToolCall.function.name}: ${partialArgs}`, jsonError); if (partialArgs.includes('}{')) { const firstBraceIndex = partialArgs.indexOf('{'); const firstCloseBraceIndex = partialArgs.indexOf('}') + 1; if (firstBraceIndex !== -1 && firstCloseBraceIndex > firstBraceIndex) { const firstJsonStr = partialArgs.substring(firstBraceIndex, firstCloseBraceIndex); try { JSON.parse(firstJsonStr); currentToolCall.function.arguments = firstJsonStr; console.log(`Extracted valid JSON at message_stop: ${firstJsonStr}`); } catch (extractError) { console.error(`Failed to extract valid JSON at message_stop: ${firstJsonStr}`, extractError); currentToolCall.function.arguments = '{}'; } } else { currentToolCall.function.arguments = '{}'; } } else { currentToolCall.function.arguments = '{}'; } } delete currentToolCall.function._partialArguments; } yield { type: 'tool_start', tool_call: currentToolCall, }; } for (const ev of flushBufferedDeltas(deltaBuffers, (_id, content) => ({ type: 'message_delta', content, message_id: messageId, order: deltaPosition++, }))) { yield ev; } if (accumulatedContent || accumulatedThinking) { if (citationTracker.citations.size > 0) { const footnotes = generateFootnotes(citationTracker); accumulatedContent += footnotes; } yield { type: 'message_complete', message_id: messageId, content: accumulatedContent, thinking_content: accumulatedThinking, thinking_signature: accumulatedSignature, }; messageCompleteYielded = true; } streamCompletedSuccessfully = true; } else if (event.type === 'error') { log_llm_error(requestId, event); console.error('Claude API error event:', event.error); yield { type: 'error', error: 'Claude API error: ' + (event.error ? event.error.message || JSON.stringify(event.error) : 'Unknown error'), }; streamCompletedSuccessfully = false; break; } } if (streamCompletedSuccessfully && (accumulatedContent || accumulatedThinking) && !messageCompleteYielded) { console.warn('Stream finished successfully but message_stop might not have triggered message_complete emission. Emitting now.'); for (const ev of flushBufferedDeltas(deltaBuffers, (_id, content) => ({ type: 'message_delta', content, message_id: messageId, order: deltaPosition++, }))) { yield ev; } if (citationTracker.citations.size > 0) { const footnotes = generateFootnotes(citationTracker); accumulatedContent += footnotes; } yield { type: 'message_complete', message_id: messageId, content: accumulatedContent, thinking_content: accumulatedThinking, thinking_signature: accumulatedSignature, }; messageCompleteYielded = true; } } catch (streamError) { log_llm_error(requestId, streamError); console.error('Error processing Claude stream:', streamError); yield { type: 'error', error: `Claude stream error (${model}): ${streamError}`, }; } finally { log_llm_response(requestId, events); } } catch (error) { log_llm_error(requestId, error); console.error('Error in Claude streaming completion setup:', error); yield { type: 'error', error: `Claude request error (${model}): ${error}`, }; } finally { if (totalInputTokens > 0 || totalOutputTokens > 0) { const cachedTokens = totalCacheCreationInputTokens + totalCacheReadInputTokens; const calculatedUsage = costTracker.addUsage({ model, input_tokens: totalInputTokens, output_tokens: totalOutputTokens, cached_tokens: cachedTokens, metadata: { cache_creation_input_tokens: totalCacheCreationInputTokens, cache_read_input_tokens: totalCacheReadInputTokens, total_tokens: totalInputTokens + totalOutputTokens, }, }); if (!hasEventHandler()) { yield { type: 'cost_update', usage: { ...calculatedUsage, total_tokens: totalInputTokens + totalOutputTokens, }, }; } } } } } export const claudeProvider = new ClaudeProvider(); //# sourceMappingURL=claude.js.map