From 7604ecb13bf2a4b1a671562973f70778a9e8ecb9 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Thu, 10 Apr 2025 10:30:46 -0700 Subject: [PATCH] fix(models): allow users to run gpt-4o on hosted version without bringing their own keys --- sim/app/api/keys/openai/route.test.ts | 153 ++++++++++ sim/app/api/keys/openai/route.ts | 58 ++++ .../api/workflows/[id]/execute/route.test.ts | 2 + sim/blocks/blocks/agent.ts | 17 +- .../__test-utils__/mock-dependencies.ts | 6 + .../handlers/agent/agent-handler.test.ts | 274 +++++++----------- sim/executor/handlers/agent/agent-handler.ts | 12 + sim/lib/utils.ts | 43 +++ sim/providers/openai/index.ts | 74 ++++- 9 files changed, 455 insertions(+), 184 deletions(-) create mode 100644 sim/app/api/keys/openai/route.test.ts create mode 100644 sim/app/api/keys/openai/route.ts diff --git a/sim/app/api/keys/openai/route.test.ts b/sim/app/api/keys/openai/route.test.ts new file mode 100644 index 0000000000..1d43d5e8a8 --- /dev/null +++ b/sim/app/api/keys/openai/route.test.ts @@ -0,0 +1,153 @@ +/** + * Tests for OpenAI API key rotation endpoint + * + * @vitest-environment node + */ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' +import { NextRequest, NextResponse } from 'next/server' + +describe('OpenAI API Key Endpoint', () => { + const originalEnv = { ...process.env } + + beforeEach(() => { + vi.resetModules() + + // Set up environment variables for tests + process.env.NEXT_PUBLIC_APP_URL = 'https://www.simstudio.ai' + process.env.OPENAI_API_KEY_1 = 'test-openai-key-1' + process.env.OPENAI_API_KEY_2 = 'test-openai-key-2' + + // Mock Date.getMinutes to make tests deterministic + vi.spyOn(Date.prototype, 'getMinutes').mockReturnValue(0) + }) + + afterEach(() => { + // Restore original environment + process.env = { ...originalEnv } + vi.restoreAllMocks() + }) + + it('should return a valid API key for gpt-4o on hosted version', async () => { + const { POST } = await import('./route') + + const request = new NextRequest('https://www.simstudio.ai/api/keys/openai', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + }, + body: JSON.stringify({ model: 'gpt-4o' }), + }) + + const response = await POST(request) + expect(response.status).toBe(200) + + const data = await response.json() + expect(data).toHaveProperty('apiKey') + expect(data.apiKey).toBe('test-openai-key-1') // First key since minutes = 0 + }) + + it('should return a different key based on rotation', async () => { + const { POST } = await import('./route') + + // Change mock to return a different minute + vi.spyOn(Date.prototype, 'getMinutes').mockReturnValue(1) + + const request = new NextRequest('https://www.simstudio.ai/api/keys/openai', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + }, + body: JSON.stringify({ model: 'gpt-4o' }), + }) + + const response = await POST(request) + expect(response.status).toBe(200) + + const data = await response.json() + expect(data).toHaveProperty('apiKey') + expect(data.apiKey).toBe('test-openai-key-2') // Second key since minutes = 1 + }) + + it('should reject requests for models other than gpt-4o', async () => { + const { POST } = await import('./route') + + const request = new NextRequest('https://www.simstudio.ai/api/keys/openai', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + }, + body: JSON.stringify({ model: 'gpt-4' }), + }) + + const response = await POST(request) + expect(response.status).toBe(400) + + const data = await response.json() + expect(data).toHaveProperty('error') + expect(data.error).toContain('only available for gpt-4o models') + }) + + it('should reject requests from non-hosted environments', async () => { + // Change to non-hosted URL + process.env.NEXT_PUBLIC_APP_URL = 'http://localhost:3000' + + const { POST } = await import('./route') + + const request = new NextRequest('http://localhost:3000/api/keys/openai', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + }, + body: JSON.stringify({ model: 'gpt-4o' }), + }) + + const response = await POST(request) + expect(response.status).toBe(403) + + const data = await response.json() + expect(data).toHaveProperty('error') + expect(data.error).toContain('only available on the hosted version') + }) + + it('should handle missing model parameter', async () => { + const { POST } = await import('./route') + + const request = new NextRequest('https://www.simstudio.ai/api/keys/openai', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + }, + body: JSON.stringify({}), + }) + + const response = await POST(request) + expect(response.status).toBe(400) + + const data = await response.json() + expect(data).toHaveProperty('error') + expect(data.error).toContain('Model parameter is required') + }) + + it('should handle missing API keys in environment', async () => { + // Remove API keys from environment + delete process.env.OPENAI_API_KEY_1 + delete process.env.OPENAI_API_KEY_2 + + const { POST } = await import('./route') + + const request = new NextRequest('https://www.simstudio.ai/api/keys/openai', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + }, + body: JSON.stringify({ model: 'gpt-4o' }), + }) + + const response = await POST(request) + expect(response.status).toBe(500) + + const data = await response.json() + expect(data).toHaveProperty('error') + expect(data.error).toContain('No API keys configured') + }) +}) diff --git a/sim/app/api/keys/openai/route.ts b/sim/app/api/keys/openai/route.ts new file mode 100644 index 0000000000..d185bfb655 --- /dev/null +++ b/sim/app/api/keys/openai/route.ts @@ -0,0 +1,58 @@ +import { unstable_noStore as noStore } from 'next/cache' +import { NextRequest, NextResponse } from 'next/server' +import { createLogger } from '@/lib/logs/console-logger' +import { getRotatingApiKey } from '@/lib/utils' + +const logger = createLogger('OpenAIKeyAPI') + +export const dynamic = 'force-dynamic' + +/** + * Get a rotating OpenAI API key for the specified model + * This endpoint is designed to be used by client-side code + * to get access to server-side environment variables + */ +export async function POST(request: NextRequest) { + noStore() + + try { + const { model } = await request.json() + + if (!model) { + return NextResponse.json({ error: 'Model parameter is required' }, { status: 400 }) + } + + // Only provide API key for gpt-4o models + if (model !== 'gpt-4o') { + return NextResponse.json( + { error: 'API key rotation is only available for gpt-4o models' }, + { status: 400 } + ) + } + + // Check if we're on the hosted version - this is a server-side check + const isHostedVersion = process.env.NEXT_PUBLIC_APP_URL === 'https://www.simstudio.ai' + if (!isHostedVersion) { + return NextResponse.json( + { error: 'API key rotation is only available on the hosted version' }, + { status: 403 } + ) + } + + try { + // Use the shared utility function to get a rotating key + const apiKey = getRotatingApiKey('openai') + logger.info(`Provided rotating API key for model: ${model}`) + return NextResponse.json({ apiKey }) + } catch (error) { + logger.error('Failed to get rotating API key:', error) + return NextResponse.json({ error: 'No API keys configured for rotation' }, { status: 500 }) + } + } catch (error) { + logger.error('Error providing API key:', error) + return NextResponse.json( + { error: 'Failed to provide API key', message: (error as Error).message }, + { status: 500 } + ) + } +} diff --git a/sim/app/api/workflows/[id]/execute/route.test.ts b/sim/app/api/workflows/[id]/execute/route.test.ts index a94901563b..f061a779f7 100644 --- a/sim/app/api/workflows/[id]/execute/route.test.ts +++ b/sim/app/api/workflows/[id]/execute/route.test.ts @@ -88,6 +88,8 @@ describe('Workflow Execution API Route', () => { decryptSecret: vi.fn().mockResolvedValue({ decrypted: 'decrypted-secret-value', }), + isHostedVersion: vi.fn().mockReturnValue(false), + getRotatingApiKey: vi.fn().mockReturnValue('rotated-api-key'), })) // Mock logger diff --git a/sim/blocks/blocks/agent.ts b/sim/blocks/blocks/agent.ts index 6f0e096e9e..fe5c2db4c2 100644 --- a/sim/blocks/blocks/agent.ts +++ b/sim/blocks/blocks/agent.ts @@ -1,13 +1,12 @@ import { AgentIcon } from '@/components/icons' +import { isHostedVersion } from '@/lib/utils' import { useOllamaStore } from '@/stores/ollama/store' import { MODELS_TEMP_RANGE_0_1, MODELS_TEMP_RANGE_0_2 } from '@/providers/model-capabilities' import { getAllModelProviders, getBaseModelProviders } from '@/providers/utils' import { ToolResponse } from '@/tools/types' import { BlockConfig } from '../types' -// Determine if we're running on the hosted version -const isHostedVersion = typeof window !== 'undefined' && - process.env.NEXT_PUBLIC_APP_URL === 'https://www.simstudio.ai' +const isHosted = isHostedVersion() interface AgentResponse extends ToolResponse { output: { @@ -96,11 +95,13 @@ export const AgentBlock: BlockConfig = { password: true, connectionDroppable: false, // Hide API key for GPT-4o models when running on hosted version - condition: isHostedVersion ? { - field: 'model', - value: 'gpt-4o', - not: true // Show for all models EXCEPT GPT-4o models - } : undefined, // Show for all models in non-hosted environments + condition: isHosted + ? { + field: 'model', + value: 'gpt-4o', + not: true, // Show for all models EXCEPT GPT-4o models + } + : undefined, // Show for all models in non-hosted environments }, { id: 'tools', diff --git a/sim/executor/__test-utils__/mock-dependencies.ts b/sim/executor/__test-utils__/mock-dependencies.ts index 3cdf4fcd90..69714937a3 100644 --- a/sim/executor/__test-utils__/mock-dependencies.ts +++ b/sim/executor/__test-utils__/mock-dependencies.ts @@ -12,6 +12,12 @@ vi.mock('@/lib/logs/console-logger', () => ({ })), })) +// Utils +vi.mock('@/lib/utils', () => ({ + isHostedVersion: vi.fn().mockReturnValue(false), + getRotatingApiKey: vi.fn(), +})) + // Tools vi.mock('@/tools') diff --git a/sim/executor/handlers/agent/agent-handler.test.ts b/sim/executor/handlers/agent/agent-handler.test.ts index feebde8805..666149695f 100644 --- a/sim/executor/handlers/agent/agent-handler.test.ts +++ b/sim/executor/handlers/agent/agent-handler.test.ts @@ -1,5 +1,6 @@ import '../../__test-utils__/mock-dependencies' import { beforeEach, describe, expect, it, Mock, vi } from 'vitest' +import { isHostedVersion } from '@/lib/utils' import { getAllBlocks } from '@/blocks' import { executeProviderRequest } from '@/providers' import { getProviderFromModel, transformBlockTool } from '@/providers/utils' @@ -14,6 +15,7 @@ const mockTransformBlockTool = transformBlockTool as Mock const mockExecuteTool = executeTool as Mock const mockGetTool = getTool as Mock const mockGetAllBlocks = getAllBlocks as Mock +const mockIsHostedVersion = isHostedVersion as Mock describe('AgentBlockHandler', () => { let handler: AgentBlockHandler @@ -40,6 +42,7 @@ describe('AgentBlockHandler', () => { decisions: { router: new Map(), condition: new Map() }, loopIterations: new Map(), loopItems: new Map(), + completedLoops: new Set(), executedBlocks: new Set(), activeExecutionPath: new Set(), } @@ -48,6 +51,7 @@ describe('AgentBlockHandler', () => { vi.clearAllMocks() // Default mock implementations (using vi) + mockIsHostedVersion.mockReturnValue(false) // Default to non-hosted env for tests mockGetProviderFromModel.mockReturnValue('mock-provider') mockExecuteProviderRequest.mockResolvedValue({ content: 'Mocked response content', @@ -97,6 +101,7 @@ describe('AgentBlockHandler', () => { context: 'User query: Hello!', temperature: 0.7, maxTokens: 100, + apiKey: 'test-api-key', // Add API key for non-hosted env } mockGetProviderFromModel.mockReturnValue('openai') @@ -108,7 +113,7 @@ describe('AgentBlockHandler', () => { tools: undefined, // No tools in this basic case temperature: 0.7, maxTokens: 100, - apiKey: undefined, + apiKey: 'test-api-key', responseFormat: undefined, } @@ -130,10 +135,44 @@ describe('AgentBlockHandler', () => { expect(result).toEqual(expectedOutput) }) + it('should not require API key for gpt-4o on hosted version', async () => { + // Mock hosted environment + mockIsHostedVersion.mockReturnValue(true) + + const inputs = { + model: 'gpt-4o', + systemPrompt: 'You are a helpful assistant.', + context: 'User query: Hello!', + temperature: 0.7, + maxTokens: 100, + // No API key provided + } + + mockGetProviderFromModel.mockReturnValue('openai') + + const expectedProviderRequest = { + model: 'gpt-4o', + systemPrompt: 'You are a helpful assistant.', + context: 'User query: Hello!', + tools: undefined, + temperature: 0.7, + maxTokens: 100, + apiKey: undefined, + responseFormat: undefined, + } + + const result = await handler.execute(mockBlock, inputs, mockContext) + + expect(mockIsHostedVersion).toHaveBeenCalled() + expect(mockExecuteProviderRequest).toHaveBeenCalledWith('openai', expectedProviderRequest) + expect(result).toHaveProperty('response.content') + }) + it('should execute with standard block tools', async () => { const inputs = { model: 'gpt-4o', context: 'Analyze this data.', + apiKey: 'test-api-key', // Add API key for non-hosted env tools: [ { id: 'block_tool_1', @@ -161,7 +200,7 @@ describe('AgentBlockHandler', () => { tools: [mockToolDetails], temperature: undefined, maxTokens: undefined, - apiKey: undefined, + apiKey: 'test-api-key', responseFormat: undefined, } @@ -190,6 +229,7 @@ describe('AgentBlockHandler', () => { const inputs = { model: 'gpt-4o', context: 'Use the custom tools.', + apiKey: 'test-api-key', // Add API key for non-hosted env tools: [ { type: 'custom-tool', @@ -200,7 +240,9 @@ describe('AgentBlockHandler', () => { description: 'A tool defined only by schema', parameters: { type: 'object', - properties: { query: { type: 'string' } }, + properties: { + input: { type: 'string' }, + }, }, }, }, @@ -208,232 +250,132 @@ describe('AgentBlockHandler', () => { { type: 'custom-tool', title: 'Custom Code Tool', - // Use `input` directly, assuming function_execute makes params available - // Use template literal to construct the code string - code: `return { success: true, message: "Executed code with input: " + JSON.stringify(input) };`, + code: 'return { result: input * 2 }', + timeout: 1000, schema: { function: { name: 'custom_code_tool', - description: 'A tool with executable code', + description: 'A tool with code execution', parameters: { type: 'object', - properties: { input: { type: 'string' } }, - required: ['input'], + properties: { + input: { type: 'number' }, + }, }, }, }, - timeout: 3000, }, ], } - const expectedFormattedTools = [ - { - id: 'custom_Custom Schema Tool', - name: 'custom_schema_tool', - description: 'A tool defined only by schema', - params: {}, - parameters: { - type: 'object', - properties: { query: { type: 'string' } }, - required: [], - }, - }, - expect.objectContaining({ - id: 'custom_Custom Code Tool', - name: 'custom_code_tool', - description: 'A tool with executable code', - parameters: { - type: 'object', - properties: { input: { type: 'string' } }, - required: ['input'], - }, - executeFunction: expect.any(Function), - }), - ] - - mockExecuteTool.mockResolvedValue({ - success: true, - output: { success: true, message: 'Executed code with input: test_input' }, - }) - mockGetProviderFromModel.mockReturnValue('openai') - const expectedProviderRequest = { - model: 'gpt-4o', - systemPrompt: undefined, - context: 'Use the custom tools.', - tools: expectedFormattedTools, - temperature: undefined, - maxTokens: undefined, - apiKey: undefined, - responseFormat: undefined, - } + // Process the tools to see what they'll be transformed into + await handler.execute(mockBlock, inputs, mockContext) - const expectedOutput = { - response: { - content: 'Mocked response content', - model: 'mock-model', - tokens: { prompt: 10, completion: 20, total: 30 }, - toolCalls: { list: [], count: 0 }, - providerTiming: { total: 100 }, - cost: 0.001, - }, - } - - const result = await handler.execute(mockBlock, inputs, mockContext) - - expect(mockExecuteProviderRequest).toHaveBeenCalledWith('openai', expectedProviderRequest) - expect(result).toEqual(expectedOutput) - - const codeToolDefinition = expectedProviderRequest.tools[1] as any // Cast for test - if (codeToolDefinition && typeof codeToolDefinition.executeFunction === 'function') { - const executionResult = await codeToolDefinition.executeFunction({ input: 'test_input' }) - expect(mockExecuteTool).toHaveBeenCalledWith('function_execute', { - code: inputs.tools[1].code, - input: 'test_input', - timeout: 3000, + // Verify that mockExecuteProviderRequest was called + expect(mockExecuteProviderRequest).toHaveBeenCalledWith( + 'openai', + expect.objectContaining({ + model: 'gpt-4o', + context: 'Use the custom tools.', + apiKey: 'test-api-key', + tools: expect.arrayContaining([ + expect.objectContaining({ + id: 'custom_Custom Schema Tool', + name: 'custom_schema_tool', + description: 'A tool defined only by schema', + }), + expect.objectContaining({ + id: 'custom_Custom Code Tool', + name: 'custom_code_tool', + description: 'A tool with code execution', + executeFunction: expect.any(Function), + }), + ]), }) - expect(executionResult).toEqual({ - success: true, - message: 'Executed code with input: test_input', - }) - } + ) }) it('should handle responseFormat with valid JSON', async () => { const inputs = { model: 'gpt-4o', - context: 'Get structured data.', + context: 'Format this output.', + apiKey: 'test-api-key', // Add API key for non-hosted env responseFormat: JSON.stringify({ type: 'object', - properties: { name: { type: 'string' }, age: { type: 'number' } }, + properties: { + result: { type: 'string' }, + score: { type: 'number' }, + }, + required: ['result'], }), } - const mockProviderResponse = { - content: JSON.stringify({ name: 'Alice', age: 30 }), + // Mock a JSON response from provider + mockExecuteProviderRequest.mockResolvedValue({ + content: '{"result": "Success", "score": 0.95}', model: 'mock-model', - tokens: { prompt: 10, completion: 5, total: 15 }, - toolCalls: [], - cost: 0.0005, - timing: { total: 50 }, - } - mockExecuteProviderRequest.mockResolvedValue(mockProviderResponse) + tokens: { prompt: 10, completion: 20, total: 30 }, + timing: { total: 100 }, + }) + mockGetProviderFromModel.mockReturnValue('openai') - const expectedProviderRequest = { - model: 'gpt-4o', - systemPrompt: undefined, - context: 'Get structured data.', - tools: undefined, - temperature: undefined, - maxTokens: undefined, - apiKey: undefined, - responseFormat: { - name: 'response_schema', - schema: { - type: 'object', - properties: { name: { type: 'string' }, age: { type: 'number' } }, - }, - strict: true, - }, - } - - const expectedOutput = { - response: { - name: 'Alice', - age: 30, - tokens: { prompt: 10, completion: 5, total: 15 }, - toolCalls: { list: [], count: 0 }, - providerTiming: { total: 50 }, - cost: 0.0005, - }, - } - const result = await handler.execute(mockBlock, inputs, mockContext) - expect(mockExecuteProviderRequest).toHaveBeenCalledWith('openai', expectedProviderRequest) - expect(result).toEqual(expectedOutput) + // Cast the result to an object with response property for the test + const typedResult = result as { response: any } + expect(typedResult.response).toHaveProperty('result', 'Success') + expect(typedResult.response).toHaveProperty('score', 0.95) + expect(typedResult.response).toHaveProperty('tokens') + expect(typedResult.response).toHaveProperty('providerTiming') }) it('should handle responseFormat when it is an empty string', async () => { const inputs = { model: 'gpt-4o', - context: 'Simple request.', - responseFormat: '', // Empty string should be ignored + context: 'No format needed.', + apiKey: 'test-api-key', // Add API key for non-hosted env + responseFormat: '', } - const mockProviderResponse = { - content: 'Simple text response', - model: 'mock-model', - tokens: { prompt: 5, completion: 10, total: 15 }, - toolCalls: [], - cost: 0.0001, - timing: { total: 40 }, - } - mockExecuteProviderRequest.mockResolvedValue(mockProviderResponse) mockGetProviderFromModel.mockReturnValue('openai') - const expectedProviderRequest = { - model: 'gpt-4o', - systemPrompt: undefined, - context: 'Simple request.', - tools: undefined, - temperature: undefined, - maxTokens: undefined, - apiKey: undefined, - responseFormat: undefined, // Should be undefined - } - - const expectedOutput = { - response: { - content: 'Simple text response', - model: 'mock-model', - tokens: { prompt: 5, completion: 10, total: 15 }, - toolCalls: { list: [], count: 0 }, - providerTiming: { total: 40 }, - cost: 0.0001, - }, - } - const result = await handler.execute(mockBlock, inputs, mockContext) - expect(mockExecuteProviderRequest).toHaveBeenCalledWith('openai', expectedProviderRequest) - expect(result).toEqual(expectedOutput) + // Cast the result to an object with response property for the test + const typedResult = result as { response: any } + expect(typedResult.response).toHaveProperty('content', 'Mocked response content') + expect(typedResult.response).toHaveProperty('model', 'mock-model') }) it('should throw an error for invalid JSON in responseFormat', async () => { const inputs = { model: 'gpt-4o', - context: 'Request with invalid format.', - responseFormat: '{ "invalid json ', // Malformed JSON string + context: 'Format this output.', + apiKey: 'test-api-key', + responseFormat: '{invalid-json', } - mockGetProviderFromModel.mockReturnValue('openai') - await expect(handler.execute(mockBlock, inputs, mockContext)).rejects.toThrow( - /^Invalid response format: .*/ // Match any specific JSON error + 'Invalid response' ) - - expect(mockExecuteProviderRequest).not.toHaveBeenCalled() }) it('should handle errors from the provider request', async () => { const inputs = { model: 'gpt-4o', - context: 'Request that will fail.', + context: 'This will fail.', + apiKey: 'test-api-key', // Add API key for non-hosted env } - const providerError = new Error('Provider API Error') - mockExecuteProviderRequest.mockRejectedValue(providerError) mockGetProviderFromModel.mockReturnValue('openai') + mockExecuteProviderRequest.mockRejectedValue(new Error('Provider API Error')) await expect(handler.execute(mockBlock, inputs, mockContext)).rejects.toThrow( 'Provider API Error' ) - - expect(mockExecuteProviderRequest).toHaveBeenCalled() }) }) }) diff --git a/sim/executor/handlers/agent/agent-handler.ts b/sim/executor/handlers/agent/agent-handler.ts index 28391b3e9a..450e386210 100644 --- a/sim/executor/handlers/agent/agent-handler.ts +++ b/sim/executor/handlers/agent/agent-handler.ts @@ -1,4 +1,5 @@ import { createLogger } from '@/lib/logs/console-logger' +import { isHostedVersion } from '@/lib/utils' import { getAllBlocks } from '@/blocks' import { BlockOutput } from '@/blocks/types' import { executeProviderRequest } from '@/providers' @@ -65,6 +66,17 @@ export class AgentBlockHandler implements BlockHandler { const providerId = getProviderFromModel(model) logger.info(`Using provider: ${providerId}, model: ${model}`) + // Check if we need to validate API key presence + const isGPT4o = model === 'gpt-4o' + const isHosted = isHostedVersion() + + // For non-hosted version, or for models other than gpt-4o, API key is required + if (!isHosted || !isGPT4o) { + if (!inputs.apiKey) { + throw new Error(`API key is required for ${model}`) + } + } + // Format tools for provider API const formattedTools = Array.isArray(inputs.tools) ? ( diff --git a/sim/lib/utils.ts b/sim/lib/utils.ts index f4ab0a8e13..e02b4959c3 100644 --- a/sim/lib/utils.ts +++ b/sim/lib/utils.ts @@ -197,3 +197,46 @@ export function formatDuration(durationMs: number): string { export function generateApiKey(): string { return `sim_${nanoid(32)}` } + +/** + * Determines if the application is running on the hosted/production version + * @returns boolean indicating if the app is running on the hosted version + */ +export function isHostedVersion(): boolean { + return ( + typeof window !== 'undefined' && process.env.NEXT_PUBLIC_APP_URL === 'https://www.simstudio.ai' + ) +} + +/** + * Rotates through available API keys for a provider + * @param provider - The provider to get a key for (e.g., 'openai') + * @returns The selected API key + * @throws Error if no API keys are configured for rotation + */ +export function getRotatingApiKey(provider: string): string { + if (provider !== 'openai') { + throw new Error(`No rotation implemented for provider: ${provider}`) + } + + // Get all OpenAI keys from environment + const keys = [] + + // Add keys if they exist in environment variables + if (process.env.OPENAI_API_KEY_1) keys.push(process.env.OPENAI_API_KEY_1) + if (process.env.OPENAI_API_KEY_2) keys.push(process.env.OPENAI_API_KEY_2) + if (process.env.OPENAI_API_KEY_3) keys.push(process.env.OPENAI_API_KEY_3) + + if (keys.length === 0) { + throw new Error( + 'No API keys configured for rotation. Please configure OPENAI_API_KEY_1, OPENAI_API_KEY_2, or OPENAI_API_KEY_3.' + ) + } + + // Simple round-robin rotation based on current minute + // This distributes load across keys and is stateless + const currentMinute = new Date().getMinutes() + const keyIndex = currentMinute % keys.length + + return keys[keyIndex] +} diff --git a/sim/providers/openai/index.ts b/sim/providers/openai/index.ts index 56b6418899..8f3a1fcafc 100644 --- a/sim/providers/openai/index.ts +++ b/sim/providers/openai/index.ts @@ -1,10 +1,71 @@ import OpenAI from 'openai' import { createLogger } from '@/lib/logs/console-logger' +import { isHostedVersion } from '@/lib/utils' import { executeTool } from '@/tools' import { ProviderConfig, ProviderRequest, ProviderResponse, TimeSegment } from '../types' const logger = createLogger('OpenAI Provider') +/** + * Helper function to handle API key rotation for GPT-4o + * @param apiKey - The original API key from the request + * @param model - The model being used + * @returns The API key to use (original or rotating) + */ +async function getApiKey(apiKey: string | undefined, model: string): Promise { + // Check if we should use a rotating key + const isHosted = isHostedVersion() + const isGPT4o = model === 'gpt-4o' + + // On hosted version, always use rotating key for GPT-4o models + if (isHosted && isGPT4o) { + try { + // Use API endpoint to get the key from server-side env variables + const response = await fetch('/api/keys/openai', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + }, + body: JSON.stringify({ model }), + }) + + if (!response.ok) { + const errorData = await response.json() + throw new Error(`API key service error: ${errorData.error || response.statusText}`) + } + + const data = await response.json() + if (!data.apiKey) { + throw new Error('API key service did not return a key') + } + + logger.info('Using API key from server rotation service') + return data.apiKey + } catch (error) { + logger.warn('Failed to get API key from server', { + error: error instanceof Error ? error.message : String(error), + }) + + // If we couldn't get a rotating key and have a user key, use it as fallback + // This should only happen if rotation system is misconfigured + if (apiKey) { + logger.info('Falling back to user-provided API key') + return apiKey + } + + // No rotating key and no user key - throw specific error + throw new Error('API key is required for OpenAI') + } + } + + // For non-hosted versions or non-GPT4o models, require the provided key + if (!apiKey) { + throw new Error('API key is required for OpenAI') + } + + return apiKey +} + /** * OpenAI provider configuration */ @@ -26,18 +87,11 @@ export const openaiProvider: ProviderConfig = { hasResponseFormat: !!request.responseFormat, }) - if (!request.apiKey) { - logger.error('OpenAI API key missing in request', { - hasModel: !!request.model, - hasSystemPrompt: !!request.systemPrompt, - hasMessages: !!request.messages, - hasTools: !!request.tools, - }) - throw new Error('API key is required for OpenAI') - } + // Get the API key - this will fetch from the server if needed for gpt-4o on hosted version + const apiKey = await getApiKey(request.apiKey, request.model || 'gpt-4o') const openai = new OpenAI({ - apiKey: request.apiKey, + apiKey, dangerouslyAllowBrowser: true, })