diff --git a/CHANGELOG.md b/CHANGELOG.md index 76fdaf74f7..6302f4473b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,9 @@ # Roo Cline Changelog +## [2.2.8] + +- Exponential backoff and retry on errors from providers + ## [2.2.7] - More fixes to search/replace diffs diff --git a/README.md b/README.md index e7e0b78b31..ce554e9c49 100644 --- a/README.md +++ b/README.md @@ -13,6 +13,7 @@ A fork of Cline, an autonomous coding agent, optimized for speed and flexibility - Support for dragging and dropping images into chats - Support for auto-approving MCP tools - Support for enabling/disabling MCP servers +- Exponential backoff and retry on errors from providers ## Disclaimer diff --git a/package-lock.json b/package-lock.json index 945e8e7103..733c11455e 100644 --- a/package-lock.json +++ b/package-lock.json @@ -1,12 +1,12 @@ { "name": "roo-cline", - "version": "2.2.7", + "version": "2.2.8", "lockfileVersion": 3, "requires": true, "packages": { "": { "name": "roo-cline", - "version": "2.2.7", + "version": "2.2.8", "dependencies": { "@anthropic-ai/bedrock-sdk": "^0.10.2", "@anthropic-ai/sdk": "^0.26.0", diff --git a/package.json b/package.json index 9fda9aa89a..62c46fa41b 100644 --- a/package.json +++ b/package.json @@ -3,7 +3,7 @@ "displayName": "Roo Cline", "description": "A fork of Cline, an autonomous coding agent, with some added experimental configuration and automation features.", "publisher": "RooVeterinaryInc", - "version": "2.2.7", + "version": "2.2.8", "icon": "assets/icons/rocket.png", "galleryBanner": { "color": "#617A91", diff --git a/src/api/providers/__tests__/openrouter.test.ts b/src/api/providers/__tests__/openrouter.test.ts index d2df3eecb1..9c87916931 100644 --- a/src/api/providers/__tests__/openrouter.test.ts +++ b/src/api/providers/__tests__/openrouter.test.ts @@ -8,6 +8,14 @@ import { Anthropic } from '@anthropic-ai/sdk' jest.mock('openai') jest.mock('axios') jest.mock('delay', () => jest.fn(() => Promise.resolve())) +jest.mock('../../utils/retry', () => ({ + withRetry: jest.fn().mockImplementation(async function*(fn) { + const generator = await fn(); + for await (const chunk of generator) { + yield chunk; + } + }) +})) describe('OpenRouterHandler', () => { const mockOptions: ApiHandlerOptions = { @@ -41,81 +49,300 @@ describe('OpenRouterHandler', () => { }) }) - test('getModel returns correct model info when options are provided', () => { - const handler = new OpenRouterHandler(mockOptions) - const result = handler.getModel() - - expect(result).toEqual({ - id: mockOptions.openRouterModelId, - info: mockOptions.openRouterModelInfo + describe('getModel', () => { + test('returns correct model info when options are provided', () => { + const handler = new OpenRouterHandler(mockOptions) + const result = handler.getModel() + + expect(result).toEqual({ + id: mockOptions.openRouterModelId, + info: mockOptions.openRouterModelInfo + }) + }) + + test('returns default model info when options are not provided', () => { + const handlerWithoutModelOptions = new OpenRouterHandler({ + openRouterApiKey: 'test-key' + }) + const result = handlerWithoutModelOptions.getModel() + + expect(result).toEqual({ + id: 'anthropic/claude-3.5-sonnet:beta', + info: expect.any(Object) + }) }) }) - test('createMessage generates correct stream chunks', async () => { - const handler = new OpenRouterHandler(mockOptions) - const mockStream = { - async *[Symbol.asyncIterator]() { - yield { - id: 'test-id', - choices: [{ - delta: { - content: 'test response' - } - }] + describe('createMessage', () => { + test('applies correct formatting for Claude models', async () => { + const claudeOptions = { + ...mockOptions, + openRouterModelId: 'anthropic/claude-3-haiku' + }; + const handler = new OpenRouterHandler(claudeOptions); + const mockStream = { + async *[Symbol.asyncIterator]() { + yield { + id: 'test-id', + choices: [{ + delta: { + content: 'test response' + } + }] + }; } - } - } - - // Mock OpenAI chat.completions.create - const mockCreate = jest.fn().mockResolvedValue(mockStream) - ;(OpenAI as jest.MockedClass).prototype.chat = { - completions: { create: mockCreate } - } as any - - // Mock axios.get for generation details - ;(axios.get as jest.Mock).mockResolvedValue({ - data: { - data: { - native_tokens_prompt: 10, - native_tokens_completion: 20, - total_cost: 0.001 - } - } - }) - - const systemPrompt = 'test system prompt' - const messages: Anthropic.Messages.MessageParam[] = [{ role: 'user' as const, content: 'test message' }] - - const generator = handler.createMessage(systemPrompt, messages) - const chunks = [] + }; - for await (const chunk of generator) { - chunks.push(chunk) - } - - // Verify stream chunks - expect(chunks).toHaveLength(2) // One text chunk and one usage chunk - expect(chunks[0]).toEqual({ - type: 'text', - text: 'test response' + // Mock OpenAI chat.completions.create + const mockCreate = jest.fn().mockResolvedValue(mockStream); + (OpenAI as jest.MockedClass).prototype.chat = { + completions: { create: mockCreate } + } as any; + + // Mock axios.get for generation details + (axios.get as jest.Mock).mockResolvedValue({ + data: { + data: { + native_tokens_prompt: 10, + native_tokens_completion: 20, + total_cost: 0.001 + } + } + }); + + const systemPrompt = 'test system prompt'; + const messages: Anthropic.Messages.MessageParam[] = [ + { role: 'user' as const, content: 'message 1' }, + { role: 'assistant' as const, content: 'response 1' }, + { role: 'user' as const, content: 'message 2' } + ]; + + const generator = handler.createMessage(systemPrompt, messages); + for await (const _ of generator) { /* consume generator */ } + + // Verify OpenAI client was called with correct Claude-specific parameters + const callArgs = mockCreate.mock.calls[0][0]; + expect(callArgs.model).toBe('anthropic/claude-3-haiku'); + expect(callArgs.max_tokens).toBeUndefined(); + expect(callArgs.temperature).toBe(0); + expect(callArgs.stream).toBe(true); + + // Verify system message has ephemeral cache control + const systemMessage = callArgs.messages[0]; + expect(systemMessage.role).toBe('system'); + expect(systemMessage.content[0]).toEqual({ + type: 'text', + text: systemPrompt, + cache_control: { type: 'ephemeral' } + }); + + // Verify user messages have ephemeral cache control + const userMessages = callArgs.messages.filter((m: OpenAI.Chat.ChatCompletionMessageParam) => m.role === 'user'); + const lastTwoUserMessages = userMessages.slice(-2); + + lastTwoUserMessages.forEach((msg: OpenAI.Chat.ChatCompletionMessageParam) => { + expect(Array.isArray(msg.content)).toBe(true); + const content = msg.content as Array<{type: string; text: string; cache_control?: {type: string}}>; + const textParts = content.filter(part => part.type === 'text'); + expect(textParts.length).toBeGreaterThan(0); + const lastTextPart = textParts[textParts.length - 1]; + expect(lastTextPart.cache_control).toEqual({ type: 'ephemeral' }); + }); }) - expect(chunks[1]).toEqual({ - type: 'usage', - inputTokens: 10, - outputTokens: 20, - totalCost: 0.001, - fullResponseText: 'test response' + + test('generates correct stream chunks with default options', async () => { + const handler = new OpenRouterHandler(mockOptions) + const mockStream = { + async *[Symbol.asyncIterator]() { + yield { + id: 'test-id', + choices: [{ + delta: { + content: 'test response' + } + }] + } + } + } + + // Mock OpenAI chat.completions.create + const mockCreate = jest.fn().mockResolvedValue(mockStream) + ;(OpenAI as jest.MockedClass).prototype.chat = { + completions: { create: mockCreate } + } as any + + // Mock axios.get for generation details + ;(axios.get as jest.Mock).mockResolvedValue({ + data: { + data: { + native_tokens_prompt: 10, + native_tokens_completion: 20, + total_cost: 0.001 + } + } + }) + + const systemPrompt = 'test system prompt' + const messages: Anthropic.Messages.MessageParam[] = [{ role: 'user' as const, content: 'test message' }] + + const generator = handler.createMessage(systemPrompt, messages) + const chunks = [] + + for await (const chunk of generator) { + chunks.push(chunk) + } + + // Verify stream chunks + expect(chunks).toHaveLength(3) // text chunk, usage chunk, and final newlines + expect(chunks[0]).toEqual({ + type: 'text', + text: 'test response' + }) + expect(chunks[1]).toEqual({ + type: 'usage', + inputTokens: 10, + outputTokens: 20, + totalCost: 0.001, + fullResponseText: 'test response' + }) + expect(chunks[2]).toEqual({ + type: 'text', + text: '\n\n' + }) + + // Verify OpenAI client was called with correct parameters + expect(mockCreate).toHaveBeenCalledWith(expect.objectContaining({ + model: mockOptions.openRouterModelId, + temperature: 0, + messages: expect.arrayContaining([ + { role: 'system', content: systemPrompt }, + { role: 'user', content: 'test message' } + ]), + stream: true + })) + }) + + test('includes middle-out transform when enabled', async () => { + const optionsWithMiddleOut = { + ...mockOptions, + openRouterUseMiddleOutTransform: true + } + const handler = new OpenRouterHandler(optionsWithMiddleOut) + const mockStream = { + async *[Symbol.asyncIterator]() { + yield { + id: 'test-id', + choices: [{ + delta: { + content: 'test response' + } + }] + } + } + } + + // Mock OpenAI chat.completions.create + const mockCreate = jest.fn().mockResolvedValue(mockStream) + ;(OpenAI as jest.MockedClass).prototype.chat = { + completions: { create: mockCreate } + } as any + + // Mock axios.get for generation details + ;(axios.get as jest.Mock).mockResolvedValue({ + data: { + data: { + native_tokens_prompt: 10, + native_tokens_completion: 20, + total_cost: 0.001 + } + } + }) + + const systemPrompt = 'test system prompt' + const messages: Anthropic.Messages.MessageParam[] = [{ role: 'user' as const, content: 'test message' }] + + const generator = handler.createMessage(systemPrompt, messages) + for await (const _ of generator) { /* consume generator */ } + + // Verify OpenAI client was called with middle-out transform + expect(mockCreate).toHaveBeenCalledWith(expect.objectContaining({ + transforms: ['middle-out'] + })) + }) + + test('handles generation details fetch failure gracefully', async () => { + const handler = new OpenRouterHandler(mockOptions) + const mockStream = { + async *[Symbol.asyncIterator]() { + yield { + id: 'test-id', + choices: [{ + delta: { + content: 'test response' + } + }] + } + } + } + + // Mock OpenAI chat.completions.create + const mockCreate = jest.fn().mockResolvedValue(mockStream) + ;(OpenAI as jest.MockedClass).prototype.chat = { + completions: { create: mockCreate } + } as any + + // Mock axios.get to fail + ;(axios.get as jest.Mock).mockRejectedValue(new Error('Network error')) + + const systemPrompt = 'test system prompt' + const messages: Anthropic.Messages.MessageParam[] = [{ role: 'user' as const, content: 'test message' }] + + const generator = handler.createMessage(systemPrompt, messages) + const chunks = [] + + for await (const chunk of generator) { + chunks.push(chunk) + } + + // Should still get text chunks even if usage info fails + expect(chunks).toHaveLength(2) + expect(chunks[0]).toEqual({ + type: 'text', + text: 'test response' + }) + expect(chunks[1]).toEqual({ + type: 'text', + text: '\n\n' + }) + }) + + test('handles API errors gracefully', async () => { + const handler = new OpenRouterHandler(mockOptions) + const mockStream = { + async *[Symbol.asyncIterator]() { + yield { + error: { + message: 'Test error', + code: 500 + } + } + } + } + + // Mock OpenAI chat.completions.create + const mockCreate = jest.fn().mockResolvedValue(mockStream) + ;(OpenAI as jest.MockedClass).prototype.chat = { + completions: { create: mockCreate } + } as any + + const systemPrompt = 'test system prompt' + const messages: Anthropic.Messages.MessageParam[] = [{ role: 'user' as const, content: 'test message' }] + + const generator = handler.createMessage(systemPrompt, messages) + + await expect(async () => { + for await (const _ of generator) { /* consume generator */ } + }).rejects.toThrow('OpenRouter API Error 500: Test error') }) - - // Verify OpenAI client was called with correct parameters - expect(mockCreate).toHaveBeenCalledWith(expect.objectContaining({ - model: mockOptions.openRouterModelId, - temperature: 0, - messages: expect.arrayContaining([ - { role: 'system', content: systemPrompt }, - { role: 'user', content: 'test message' } - ]), - stream: true - })) }) }) diff --git a/src/api/providers/anthropic.ts b/src/api/providers/anthropic.ts index c090f17c63..60102f8e7c 100644 --- a/src/api/providers/anthropic.ts +++ b/src/api/providers/anthropic.ts @@ -8,7 +8,9 @@ import { ModelInfo, } from "../../shared/api" import { ApiHandler } from "../index" -import { ApiStream } from "../transform/stream" +import { ApiStream, ApiStreamChunk } from "../transform/stream" +import delay from "delay" +import { withRetry } from "../utils/retry" export class AnthropicHandler implements ApiHandler { private options: ApiHandlerOptions @@ -23,145 +25,164 @@ export class AnthropicHandler implements ApiHandler { } async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream { - let stream: AnthropicStream - const modelId = this.getModel().id - switch (modelId) { - // 'latest' alias does not support cache_control - case "claude-3-5-sonnet-20241022": - case "claude-3-5-haiku-20241022": - case "claude-3-opus-20240229": - case "claude-3-haiku-20240307": { - /* - The latest message will be the new user message, one before will be the assistant message from a previous request, and the user message before that will be a previously cached user message. So we need to mark the latest user message as ephemeral to cache it for the next request, and mark the second to last user message as ephemeral to let the server know the last message to retrieve from the cache for the current request.. - */ - const userMsgIndices = messages.reduce( - (acc, msg, index) => (msg.role === "user" ? [...acc, index] : acc), - [] as number[], - ) - const lastUserMsgIndex = userMsgIndices[userMsgIndices.length - 1] ?? -1 - const secondLastMsgUserIndex = userMsgIndices[userMsgIndices.length - 2] ?? -1 - stream = await this.client.beta.promptCaching.messages.create( - { + const self = this; + const modelId = this.getModel().id; + + const gen = withRetry(async () => { + let stream: AnthropicStream + + switch (modelId) { + // 'latest' alias does not support cache_control + case "claude-3-5-sonnet-20241022": + case "claude-3-5-haiku-20241022": + case "claude-3-opus-20240229": + case "claude-3-haiku-20240307": { + /* + The latest message will be the new user message, one before will be the assistant message from a previous request, and the user message before that will be a previously cached user message. So we need to mark the latest user message as ephemeral to cache it for the next request, and mark the second to last user message as ephemeral to let the server know the last message to retrieve from the cache for the current request.. + */ + const userMsgIndices = messages.reduce( + (acc, msg, index) => (msg.role === "user" ? [...acc, index] : acc), + [] as number[], + ) + const lastUserMsgIndex = userMsgIndices[userMsgIndices.length - 1] ?? -1 + const secondLastMsgUserIndex = userMsgIndices[userMsgIndices.length - 2] ?? -1 + stream = await self.client.beta.promptCaching.messages.create( + { + model: modelId, + max_tokens: self.getModel().info.maxTokens || 8192, + temperature: 0, + system: [{ text: systemPrompt, type: "text", cache_control: { type: "ephemeral" } }], // setting cache breakpoint for system prompt so new tasks can reuse it + messages: messages.map((message, index) => { + if (index === lastUserMsgIndex || index === secondLastMsgUserIndex) { + return { + ...message, + content: + typeof message.content === "string" + ? [ + { + type: "text", + text: message.content, + cache_control: { type: "ephemeral" }, + }, + ] + : message.content.map((content, contentIndex) => + contentIndex === message.content.length - 1 + ? { ...content, cache_control: { type: "ephemeral" } } + : content, + ), + } + } + return message + }), + // tools, // cache breakpoints go from tools > system > messages, and since tools dont change, we can just set the breakpoint at the end of system (this avoids having to set a breakpoint at the end of tools which by itself does not meet min requirements for haiku caching) + // tool_choice: { type: "auto" }, + // tools: tools, + stream: true, + }, + (() => { + // prompt caching: https://x.com/alexalbert__/status/1823751995901272068 + // https://github.com/anthropics/anthropic-sdk-typescript?tab=readme-ov-file#default-headers + // https://github.com/anthropics/anthropic-sdk-typescript/commit/c920b77fc67bd839bfeb6716ceab9d7c9bbe7393 + switch (modelId) { + case "claude-3-5-sonnet-20241022": + case "claude-3-5-haiku-20241022": + case "claude-3-opus-20240229": + case "claude-3-haiku-20240307": + return { + headers: { "anthropic-beta": "prompt-caching-2024-07-31" }, + } + default: + return undefined + } + })(), + ) + break + } + default: { + stream = (await self.client.messages.create({ model: modelId, - max_tokens: this.getModel().info.maxTokens || 8192, + max_tokens: self.getModel().info.maxTokens || 8192, temperature: 0, - system: [{ text: systemPrompt, type: "text", cache_control: { type: "ephemeral" } }], // setting cache breakpoint for system prompt so new tasks can reuse it - messages: messages.map((message, index) => { - if (index === lastUserMsgIndex || index === secondLastMsgUserIndex) { - return { - ...message, - content: - typeof message.content === "string" - ? [ - { - type: "text", - text: message.content, - cache_control: { type: "ephemeral" }, - }, - ] - : message.content.map((content, contentIndex) => - contentIndex === message.content.length - 1 - ? { ...content, cache_control: { type: "ephemeral" } } - : content, - ), - } - } - return message - }), - // tools, // cache breakpoints go from tools > system > messages, and since tools dont change, we can just set the breakpoint at the end of system (this avoids having to set a breakpoint at the end of tools which by itself does not meet min requirements for haiku caching) + system: [{ text: systemPrompt, type: "text" }], + messages, + // tools, // tool_choice: { type: "auto" }, - // tools: tools, stream: true, - }, - (() => { - // prompt caching: https://x.com/alexalbert__/status/1823751995901272068 - // https://github.com/anthropics/anthropic-sdk-typescript?tab=readme-ov-file#default-headers - // https://github.com/anthropics/anthropic-sdk-typescript/commit/c920b77fc67bd839bfeb6716ceab9d7c9bbe7393 - switch (modelId) { - case "claude-3-5-sonnet-20241022": - case "claude-3-5-haiku-20241022": - case "claude-3-opus-20240229": - case "claude-3-haiku-20240307": - return { - headers: { "anthropic-beta": "prompt-caching-2024-07-31" }, - } - default: - return undefined - } - })(), - ) - break + })) as any + break + } } - default: { - stream = (await this.client.messages.create({ - model: modelId, - max_tokens: this.getModel().info.maxTokens || 8192, - temperature: 0, - system: [{ text: systemPrompt, type: "text" }], - messages, - // tools, - // tool_choice: { type: "auto" }, - stream: true, - })) as any - break - } - } - for await (const chunk of stream) { - switch (chunk.type) { - case "message_start": - // tells us cache reads/writes/input/output - const usage = chunk.message.usage - yield { - type: "usage", - inputTokens: usage.input_tokens || 0, - outputTokens: usage.output_tokens || 0, - cacheWriteTokens: usage.cache_creation_input_tokens || undefined, - cacheReadTokens: usage.cache_read_input_tokens || undefined, - } - break - case "message_delta": - // tells us stop_reason, stop_sequence, and output tokens along the way and at the end of the message - - yield { - type: "usage", - inputTokens: 0, - outputTokens: chunk.usage.output_tokens || 0, - } - break - case "message_stop": - // no usage data, just an indicator that the message is done - break - case "content_block_start": - switch (chunk.content_block.type) { - case "text": - // we may receive multiple text blocks, in which case just insert a line break between them - if (chunk.index > 0) { - yield { - type: "text", - text: "\n", - } - } + return (async function*() { + for await (const chunk of stream) { + switch (chunk.type) { + case "message_start": + // tells us cache reads/writes/input/output + const usage = chunk.message.usage yield { - type: "text", - text: chunk.content_block.text, + type: "usage", + inputTokens: usage.input_tokens || 0, + outputTokens: usage.output_tokens || 0, + cacheWriteTokens: usage.cache_creation_input_tokens || undefined, + cacheReadTokens: usage.cache_read_input_tokens || undefined, + } as ApiStreamChunk; + break + case "message_delta": + // tells us stop_reason, stop_sequence, and output tokens along the way and at the end of the message + + yield { + type: "usage", + inputTokens: 0, + outputTokens: chunk.usage.output_tokens || 0, + } as ApiStreamChunk; + break + case "message_stop": + // no usage data, just an indicator that the message is done + break + case "content_block_start": + switch (chunk.content_block.type) { + case "text": + // we may receive multiple text blocks, in which case just insert a line break between them + if (chunk.index > 0) { + yield { + type: "text", + text: "\n", + } as ApiStreamChunk; + } + yield { + type: "text", + text: chunk.content_block.text, + } as ApiStreamChunk; + break } break - } - break - case "content_block_delta": - switch (chunk.delta.type) { - case "text_delta": - yield { - type: "text", - text: chunk.delta.text, + case "content_block_delta": + switch (chunk.delta.type) { + case "text_delta": + yield { + type: "text", + text: chunk.delta.text, + } as ApiStreamChunk; + break } break + case "content_block_stop": + break } - break - case "content_block_stop": - break + } + })(); + }, { + maxRetries: 5, + initialDelayMs: 2000, + onRetry: (error, attempt, delayMs) => { + console.log(`Anthropic request failed (attempt ${attempt})`); + console.log(`Error:`, error); + console.log(`Retrying in ${delayMs}ms...`); } + }); + + for await (const chunk of gen) { + yield chunk; } } diff --git a/src/api/providers/bedrock.ts b/src/api/providers/bedrock.ts index 58f75ad4ac..e765fda571 100644 --- a/src/api/providers/bedrock.ts +++ b/src/api/providers/bedrock.ts @@ -2,7 +2,8 @@ import AnthropicBedrock from "@anthropic-ai/bedrock-sdk" import { Anthropic } from "@anthropic-ai/sdk" import { ApiHandler } from "../" import { ApiHandlerOptions, bedrockDefaultModelId, BedrockModelId, bedrockModels, ModelInfo } from "../../shared/api" -import { ApiStream } from "../transform/stream" +import { ApiStream, ApiStreamChunk } from "../transform/stream" +import { withRetry } from "../utils/retry" // https://docs.anthropic.com/en/api/claude-on-amazon-bedrock export class AwsBedrockHandler implements ApiHandler { @@ -25,6 +26,8 @@ export class AwsBedrockHandler implements ApiHandler { } async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream { + const self = this; + // cross region inference requires prefixing the model id with the region let modelId: string if (this.options.awsUseCrossRegionInference) { @@ -45,59 +48,77 @@ export class AwsBedrockHandler implements ApiHandler { modelId = this.getModel().id } - const stream = await this.client.messages.create({ - model: modelId, - max_tokens: this.getModel().info.maxTokens || 8192, - temperature: 0, - system: systemPrompt, - messages, - stream: true, - }) - for await (const chunk of stream) { - switch (chunk.type) { - case "message_start": - const usage = chunk.message.usage - yield { - type: "usage", - inputTokens: usage.input_tokens || 0, - outputTokens: usage.output_tokens || 0, - } - break - case "message_delta": - yield { - type: "usage", - inputTokens: 0, - outputTokens: chunk.usage.output_tokens || 0, - } - break + const gen = withRetry(async () => { + const stream = await self.client.messages.create({ + model: modelId, + max_tokens: self.getModel().info.maxTokens || 8192, + temperature: 0, + system: systemPrompt, + messages, + stream: true, + }) - case "content_block_start": - switch (chunk.content_block.type) { - case "text": - if (chunk.index > 0) { - yield { - type: "text", - text: "\n", - } - } + return (async function*() { + for await (const chunk of stream) { + switch (chunk.type) { + case "message_start": + // tells us cache reads/writes/input/output + const usage = chunk.message.usage yield { - type: "text", - text: chunk.content_block.text, + type: "usage", + inputTokens: usage.input_tokens || 0, + outputTokens: usage.output_tokens || 0, + } as ApiStreamChunk; + break + case "message_delta": + yield { + type: "usage", + inputTokens: 0, + outputTokens: chunk.usage.output_tokens || 0, + } as ApiStreamChunk; + break + + case "content_block_start": + switch (chunk.content_block.type) { + case "text": + if (chunk.index > 0) { + yield { + type: "text", + text: "\n", + } as ApiStreamChunk; + } + yield { + type: "text", + text: chunk.content_block.text, + } as ApiStreamChunk; + break + } + break + case "content_block_delta": + switch (chunk.delta.type) { + case "text_delta": + yield { + type: "text", + text: chunk.delta.text, + } as ApiStreamChunk; + break } break } - break - case "content_block_delta": - switch (chunk.delta.type) { - case "text_delta": - yield { - type: "text", - text: chunk.delta.text, - } - break - } - break + } + })(); + }, { + maxRetries: 5, + initialDelayMs: 2000, + onRetry: (error, attempt, delayMs) => { + console.log(`Bedrock request failed (attempt ${attempt})`); + console.log(`Error:`, error); + console.log(`Retrying in ${delayMs}ms...`); } + }); + + for await (const chunk of gen) { + yield chunk; } } diff --git a/src/api/providers/gemini.ts b/src/api/providers/gemini.ts index d7ac5ec67d..2c64a3387a 100644 --- a/src/api/providers/gemini.ts +++ b/src/api/providers/gemini.ts @@ -3,7 +3,8 @@ import { GoogleGenerativeAI } from "@google/generative-ai" import { ApiHandler } from "../" import { ApiHandlerOptions, geminiDefaultModelId, GeminiModelId, geminiModels, ModelInfo } from "../../shared/api" import { convertAnthropicMessageToGemini } from "../transform/gemini-format" -import { ApiStream } from "../transform/stream" +import { ApiStream, ApiStreamChunk } from "../transform/stream" +import { withRetry } from "../utils/retry" export class GeminiHandler implements ApiHandler { private options: ApiHandlerOptions @@ -18,30 +19,48 @@ export class GeminiHandler implements ApiHandler { } async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream { + const self = this; const model = this.client.getGenerativeModel({ model: this.getModel().id, systemInstruction: systemPrompt, }) - const result = await model.generateContentStream({ - contents: messages.map(convertAnthropicMessageToGemini), - generationConfig: { - // maxOutputTokens: this.getModel().info.maxTokens, - temperature: 0, - }, - }) - for await (const chunk of result.stream) { - yield { - type: "text", - text: chunk.text(), + const gen = withRetry(async () => { + const result = await model.generateContentStream({ + contents: messages.map(convertAnthropicMessageToGemini), + generationConfig: { + // maxOutputTokens: this.getModel().info.maxTokens, + temperature: 0, + }, + }) + + return (async function*() { + for await (const chunk of result.stream) { + yield { + type: "text", + text: chunk.text(), + } as ApiStreamChunk; + } + + const response = await result.response + yield { + type: "usage", + inputTokens: response.usageMetadata?.promptTokenCount ?? 0, + outputTokens: response.usageMetadata?.candidatesTokenCount ?? 0, + } as ApiStreamChunk; + })(); + }, { + maxRetries: 5, + initialDelayMs: 2000, + onRetry: (error, attempt, delayMs) => { + console.log(`Gemini request failed (attempt ${attempt})`); + console.log(`Error:`, error); + console.log(`Retrying in ${delayMs}ms...`); } - } + }); - const response = await result.response - yield { - type: "usage", - inputTokens: response.usageMetadata?.promptTokenCount ?? 0, - outputTokens: response.usageMetadata?.candidatesTokenCount ?? 0, + for await (const chunk of gen) { + yield chunk; } } diff --git a/src/api/providers/lmstudio.ts b/src/api/providers/lmstudio.ts index 868ef7da13..08c434ac82 100644 --- a/src/api/providers/lmstudio.ts +++ b/src/api/providers/lmstudio.ts @@ -3,7 +3,8 @@ import OpenAI from "openai" import { ApiHandler } from "../" import { ApiHandlerOptions, ModelInfo, openAiModelInfoSaneDefaults } from "../../shared/api" import { convertToOpenAiMessages } from "../transform/openai-format" -import { ApiStream } from "../transform/stream" +import { ApiStream, ApiStreamChunk } from "../transform/stream" +import { withRetry } from "../utils/retry" export class LmStudioHandler implements ApiHandler { private options: ApiHandlerOptions @@ -18,32 +19,48 @@ export class LmStudioHandler implements ApiHandler { } async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream { + const self = this; const openAiMessages: OpenAI.Chat.ChatCompletionMessageParam[] = [ { role: "system", content: systemPrompt }, ...convertToOpenAiMessages(messages), ] - try { - const stream = await this.client.chat.completions.create({ - model: this.getModel().id, + const gen = withRetry(async () => { + const stream = await self.client.chat.completions.create({ + model: self.getModel().id, messages: openAiMessages, temperature: 0, stream: true, - }) - for await (const chunk of stream) { - const delta = chunk.choices[0]?.delta - if (delta?.content) { - yield { - type: "text", - text: delta.content, + }).catch(error => { + // LM Studio doesn't return an error code/body for now + throw new Error( + "Please check the LM Studio developer logs to debug what went wrong. You may need to load the model with a larger context length to work with Cline's prompts.", + ) + }); + + return (async function*() { + for await (const chunk of stream) { + const delta = chunk.choices[0]?.delta + if (delta?.content) { + yield { + type: "text", + text: delta.content, + } as ApiStreamChunk; } } + })(); + }, { + maxRetries: 5, + initialDelayMs: 2000, + onRetry: (error, attempt, delayMs) => { + console.log(`LM Studio request failed (attempt ${attempt})`); + console.log(`Error:`, error); + console.log(`Retrying in ${delayMs}ms...`); } - } catch (error) { - // LM Studio doesn't return an error code/body for now - throw new Error( - "Please check the LM Studio developer logs to debug what went wrong. You may need to load the model with a larger context length to work with Cline's prompts.", - ) + }); + + for await (const chunk of gen) { + yield chunk; } } diff --git a/src/api/providers/ollama.ts b/src/api/providers/ollama.ts index 7668bd395f..45151d5b35 100644 --- a/src/api/providers/ollama.ts +++ b/src/api/providers/ollama.ts @@ -3,7 +3,8 @@ import OpenAI from "openai" import { ApiHandler } from "../" import { ApiHandlerOptions, ModelInfo, openAiModelInfoSaneDefaults } from "../../shared/api" import { convertToOpenAiMessages } from "../transform/openai-format" -import { ApiStream } from "../transform/stream" +import { ApiStream, ApiStreamChunk } from "../transform/stream" +import { withRetry } from "../utils/retry" export class OllamaHandler implements ApiHandler { private options: ApiHandlerOptions @@ -18,25 +19,49 @@ export class OllamaHandler implements ApiHandler { } async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream { + const self = this; const openAiMessages: OpenAI.Chat.ChatCompletionMessageParam[] = [ { role: "system", content: systemPrompt }, ...convertToOpenAiMessages(messages), ] - const stream = await this.client.chat.completions.create({ - model: this.getModel().id, - messages: openAiMessages, - temperature: 0, - stream: true, - }) - for await (const chunk of stream) { - const delta = chunk.choices[0]?.delta - if (delta?.content) { - yield { - type: "text", - text: delta.content, + const gen = withRetry(async () => { + const stream = await self.client.chat.completions.create({ + model: self.getModel().id, + messages: openAiMessages, + temperature: 0, + stream: true, + }).catch(error => { + // Check if it's a connection error, which likely means Ollama isn't running + if (error instanceof Error && error.message.includes('ECONNREFUSED')) { + throw new Error('Could not connect to Ollama. Please make sure Ollama is running on your system.'); } + throw error; + }); + + return (async function*() { + for await (const chunk of stream) { + const delta = chunk.choices[0]?.delta + if (delta?.content) { + yield { + type: "text", + text: delta.content, + } as ApiStreamChunk; + } + } + })(); + }, { + maxRetries: 5, + initialDelayMs: 2000, + onRetry: (error, attempt, delayMs) => { + console.log(`Ollama request failed (attempt ${attempt})`); + console.log(`Error:`, error); + console.log(`Retrying in ${delayMs}ms...`); } + }); + + for await (const chunk of gen) { + yield chunk; } } diff --git a/src/api/providers/openai-native.ts b/src/api/providers/openai-native.ts index 70d55b7abe..6b9d893bcf 100644 --- a/src/api/providers/openai-native.ts +++ b/src/api/providers/openai-native.ts @@ -9,7 +9,8 @@ import { openAiNativeModels, } from "../../shared/api" import { convertToOpenAiMessages } from "../transform/openai-format" -import { ApiStream } from "../transform/stream" +import { ApiStream, ApiStreamChunk } from "../transform/stream" +import { withRetry } from "../utils/retry" export class OpenAiNativeHandler implements ApiHandler { private options: ApiHandlerOptions @@ -23,52 +24,87 @@ export class OpenAiNativeHandler implements ApiHandler { } async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream { + const self = this; + switch (this.getModel().id) { case "o1-preview": case "o1-mini": { - // o1 doesnt support streaming, non-1 temp, or system prompt - const response = await this.client.chat.completions.create({ - model: this.getModel().id, - messages: [{ role: "user", content: systemPrompt }, ...convertToOpenAiMessages(messages)], - }) - yield { - type: "text", - text: response.choices[0]?.message.content || "", - } - yield { - type: "usage", - inputTokens: response.usage?.prompt_tokens || 0, - outputTokens: response.usage?.completion_tokens || 0, - } - break - } - default: { - const stream = await this.client.chat.completions.create({ - model: this.getModel().id, - // max_completion_tokens: this.getModel().info.maxTokens, - temperature: 0, - messages: [{ role: "system", content: systemPrompt }, ...convertToOpenAiMessages(messages)], - stream: true, - stream_options: { include_usage: true }, - }) + const gen = withRetry(async () => { + // o1 doesnt support streaming, non-1 temp, or system prompt + const response = await self.client.chat.completions.create({ + model: self.getModel().id, + messages: [{ role: "user", content: systemPrompt }, ...convertToOpenAiMessages(messages)], + }); - for await (const chunk of stream) { - const delta = chunk.choices[0]?.delta - if (delta?.content) { + return (async function*() { yield { type: "text", - text: delta.content, - } - } - - // contains a null value except for the last chunk which contains the token usage statistics for the entire request - if (chunk.usage) { + text: response.choices[0]?.message.content || "", + } as ApiStreamChunk; yield { type: "usage", - inputTokens: chunk.usage.prompt_tokens || 0, - outputTokens: chunk.usage.completion_tokens || 0, - } + inputTokens: response.usage?.prompt_tokens || 0, + outputTokens: response.usage?.completion_tokens || 0, + } as ApiStreamChunk; + })(); + }, { + maxRetries: 5, + initialDelayMs: 2000, + onRetry: (error, attempt, delayMs) => { + console.log(`OpenAI Native request failed (attempt ${attempt})`); + console.log(`Error:`, error); + console.log(`Retrying in ${delayMs}ms...`); } + }); + + for await (const chunk of gen) { + yield chunk; + } + break; + } + default: { + const gen = withRetry(async () => { + const stream = await self.client.chat.completions.create({ + model: self.getModel().id, + // max_completion_tokens: this.getModel().info.maxTokens, + temperature: 0, + messages: [{ role: "system", content: systemPrompt }, ...convertToOpenAiMessages(messages)], + stream: true, + stream_options: { include_usage: true }, + }); + + return (async function*() { + for await (const chunk of stream) { + const delta = chunk.choices[0]?.delta + if (delta?.content) { + yield { + type: "text", + text: delta.content, + } as ApiStreamChunk; + } + + // contains a null value except for the last chunk which contains the token usage statistics for the entire request + if (chunk.usage) { + yield { + type: "usage", + inputTokens: chunk.usage.prompt_tokens || 0, + outputTokens: chunk.usage.completion_tokens || 0, + } as ApiStreamChunk; + } + } + })(); + }, { + maxRetries: 5, + initialDelayMs: 2000, + onRetry: (error, attempt, delayMs) => { + console.log(`OpenAI Native request failed (attempt ${attempt})`); + console.log(`Error:`, error); + console.log(`Retrying in ${delayMs}ms...`); + } + }); + + for await (const chunk of gen) { + yield chunk; } } } diff --git a/src/api/providers/openai.ts b/src/api/providers/openai.ts index 71308ed180..00ca147515 100644 --- a/src/api/providers/openai.ts +++ b/src/api/providers/openai.ts @@ -8,7 +8,8 @@ import { } from "../../shared/api" import { ApiHandler } from "../index" import { convertToOpenAiMessages } from "../transform/openai-format" -import { ApiStream } from "../transform/stream" +import { ApiStream, ApiStreamChunk } from "../transform/stream" +import { withRetry } from "../utils/retry" export class OpenAiHandler implements ApiHandler { private options: ApiHandlerOptions @@ -49,22 +50,41 @@ export class OpenAiHandler implements ApiHandler { requestOptions.stream_options = { include_usage: true } } - const stream = await this.client.chat.completions.create(requestOptions) - for await (const chunk of stream) { - const delta = chunk.choices[0]?.delta - if (delta?.content) { - yield { - type: "text", - text: delta.content, - } - } - if (chunk.usage) { - yield { - type: "usage", - inputTokens: chunk.usage.prompt_tokens || 0, - outputTokens: chunk.usage.completion_tokens || 0, + const self = this; + + const gen = withRetry(async () => { + const stream = await self.client.chat.completions.create(requestOptions); + + return (async function*() { + for await (const chunk of stream) { + const delta = chunk.choices[0]?.delta + if (delta?.content) { + yield { + type: "text", + text: delta.content, + } as ApiStreamChunk; + } + if (chunk.usage) { + yield { + type: "usage", + inputTokens: chunk.usage.prompt_tokens || 0, + outputTokens: chunk.usage.completion_tokens || 0, + } as ApiStreamChunk; + } } + })(); + }, { + maxRetries: 5, + initialDelayMs: 2000, + onRetry: (error, attempt, delayMs) => { + console.log(`OpenAI request failed (attempt ${attempt})`); + console.log(`Error:`, error); + console.log(`Retrying in ${delayMs}ms...`); } + }); + + for await (const chunk of gen) { + yield chunk; } } diff --git a/src/api/providers/openrouter.ts b/src/api/providers/openrouter.ts index c2c34d80f0..cfb9c9a0d5 100644 --- a/src/api/providers/openrouter.ts +++ b/src/api/providers/openrouter.ts @@ -6,6 +6,7 @@ import { ApiHandlerOptions, ModelInfo, openRouterDefaultModelId, openRouterDefau import { convertToOpenAiMessages } from "../transform/openai-format" import { ApiStreamChunk, ApiStreamUsageChunk } from "../transform/stream" import delay from "delay" +import { withRetry } from "../utils/retry" // Add custom interface for OpenRouter params interface OpenRouterChatCompletionParams extends OpenAI.Chat.ChatCompletionCreateParamsStreaming { @@ -18,6 +19,7 @@ interface OpenRouterApiStreamUsageChunk extends ApiStreamUsageChunk { } export class OpenRouterHandler implements ApiHandler { + private static requestCount = 0 private options: ApiHandlerOptions private client: OpenAI @@ -27,21 +29,28 @@ export class OpenRouterHandler implements ApiHandler { baseURL: "https://openrouter.ai/api/v1", apiKey: this.options.openRouterApiKey, defaultHeaders: { - "HTTP-Referer": "https://github.com/RooVetGit/Roo-Cline", // Optional, for including your app on openrouter.ai rankings. - "X-Title": "Roo-Cline", // Optional. Shows in rankings on openrouter.ai. + "HTTP-Referer": "https://github.com/RooVetGit/Roo-Cline", + "X-Title": "Roo-Cline", }, }) } + getModel(): { id: string; info: ModelInfo } { + const modelId = this.options.openRouterModelId + const modelInfo = this.options.openRouterModelInfo + if (modelId && modelInfo) { + return { id: modelId, info: modelInfo } + } + return { id: openRouterDefaultModelId, info: openRouterDefaultModelInfo } + } + async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): AsyncGenerator { - // Convert Anthropic messages to OpenAI format const openAiMessages: OpenAI.Chat.ChatCompletionMessageParam[] = [ { role: "system", content: systemPrompt }, ...convertToOpenAiMessages(messages), ] // prompt caching: https://openrouter.ai/docs/prompt-caching - // this is specifically for claude models (some models may 'support prompt caching' automatically without this) switch (this.getModel().id) { case "anthropic/claude-3.5-sonnet": case "anthropic/claude-3.5-sonnet:beta": @@ -66,17 +75,13 @@ export class OpenRouterHandler implements ApiHandler { }, ], } - // Add cache_control to the last two user messages - // (note: this works because we only ever add one user message at a time, but if we added multiple we'd need to mark the user message before the last assistant message) const lastTwoUserMessages = openAiMessages.filter((msg) => msg.role === "user").slice(-2) lastTwoUserMessages.forEach((msg) => { if (typeof msg.content === "string") { msg.content = [{ type: "text", text: msg.content }] } if (Array.isArray(msg.content)) { - // NOTE: this is fine since env details will always be added at the end. but if it weren't there, and the user added a image_url type message, it would pop a text part before it and then move it after to the end. let lastTextPart = msg.content.filter((part) => part.type === "text").pop() - if (!lastTextPart) { lastTextPart = { type: "text", text: "..." } msg.content.push(lastTextPart) @@ -86,12 +91,8 @@ export class OpenRouterHandler implements ApiHandler { } }) break - default: - break } - // Not sure how openrouter defaults max tokens when no value is provided, but the anthropic api requires this value and since they offer both 4096 and 8192 variants, we should ensure 8192. - // (models usually default to max tokens allowed) let maxTokens: number | undefined switch (this.getModel().id) { case "anthropic/claude-3.5-sonnet": @@ -105,83 +106,85 @@ export class OpenRouterHandler implements ApiHandler { maxTokens = 8_192 break } - // https://openrouter.ai/docs/transforms + let fullResponseText = ""; - const stream = await this.client.chat.completions.create({ - model: this.getModel().id, - max_tokens: maxTokens, - temperature: 0, - messages: openAiMessages, - stream: true, - // This way, the transforms field will only be included in the parameters when openRouterUseMiddleOutTransform is true. - ...(this.options.openRouterUseMiddleOutTransform && { transforms: ["middle-out"] }) - } as OpenRouterChatCompletionParams); + const self = this; - let genId: string | undefined + const gen = withRetry(async () => { + const stream = await self.client.chat.completions.create({ + model: self.getModel().id, + max_tokens: maxTokens, + temperature: 0, + messages: openAiMessages, + stream: true, + ...(self.options.openRouterUseMiddleOutTransform && { transforms: ["middle-out"] }) + } as OpenRouterChatCompletionParams); - for await (const chunk of stream as unknown as AsyncIterable) { - // openrouter returns an error object instead of the openai sdk throwing an error - if ("error" in chunk) { - const error = chunk.error as { message?: string; code?: number } - console.error(`OpenRouter API Error: ${error?.code} - ${error?.message}`) - throw new Error(`OpenRouter API Error ${error?.code}: ${error?.message}`) - } + let genId: string | undefined; - if (!genId && chunk.id) { - genId = chunk.id - } + return (async function*() { + for await (const chunk of stream) { + if ("error" in chunk) { + const error = chunk.error as { message?: string; code?: number } + console.error(`OpenRouter API Error: ${error?.code} - ${error?.message}`) + throw new Error(`OpenRouter API Error ${error?.code}: ${error?.message}`) + } - const delta = chunk.choices[0]?.delta - if (delta?.content) { - fullResponseText += delta.content; + if (!genId && chunk.id) { + genId = chunk.id + } + + const delta = chunk.choices[0]?.delta + if (delta?.content) { + fullResponseText += delta.content; + yield { + type: "text", + text: delta.content, + } as ApiStreamChunk; + } + } + + await delay(500) + + try { + const response = await axios.get(`https://openrouter.ai/api/v1/generation?id=${genId}`, { + headers: { + Authorization: `Bearer ${self.options.openRouterApiKey}`, + }, + timeout: 5_000, + }) + + const generation = response.data?.data + console.log("OpenRouter generation details:", response.data) + yield { + type: "usage", + inputTokens: generation?.native_tokens_prompt || 0, + outputTokens: generation?.native_tokens_completion || 0, + totalCost: generation?.total_cost || 0, + fullResponseText + } as ApiStreamChunk; + } catch (error) { + console.error("Error fetching OpenRouter generation details:", error) + } + + // Add newlines before starting content yield { type: "text", - text: delta.content, + text: "\n\n" } as ApiStreamChunk; + })(); + }, { + maxRetries: 5, + initialDelayMs: 2000, + onRetry: (error, attempt, delayMs) => { + console.log(`OpenRouter request failed (attempt ${attempt})`); + console.log(`Error:`, error); + console.log(`Retrying in ${delayMs}ms...`); } - // if (chunk.usage) { - // yield { - // type: "usage", - // inputTokens: chunk.usage.prompt_tokens || 0, - // outputTokens: chunk.usage.completion_tokens || 0, - // } - // } + }); + + for await (const chunk of gen) { + yield chunk; } - - await delay(500) // FIXME: necessary delay to ensure generation endpoint is ready - - try { - const response = await axios.get(`https://openrouter.ai/api/v1/generation?id=${genId}`, { - headers: { - Authorization: `Bearer ${this.options.openRouterApiKey}`, - }, - timeout: 5_000, // this request hangs sometimes - }) - - const generation = response.data?.data - console.log("OpenRouter generation details:", response.data) - yield { - type: "usage", - // cacheWriteTokens: 0, - // cacheReadTokens: 0, - // openrouter generation endpoint fails often - inputTokens: generation?.native_tokens_prompt || 0, - outputTokens: generation?.native_tokens_completion || 0, - totalCost: generation?.total_cost || 0, - fullResponseText - } as OpenRouterApiStreamUsageChunk; - } catch (error) { - // ignore if fails - console.error("Error fetching OpenRouter generation details:", error) - } - - } - getModel(): { id: string; info: ModelInfo } { - const modelId = this.options.openRouterModelId - const modelInfo = this.options.openRouterModelInfo - if (modelId && modelInfo) { - return { id: modelId, info: modelInfo } - } - return { id: openRouterDefaultModelId, info: openRouterDefaultModelInfo } } } diff --git a/src/api/providers/vertex.ts b/src/api/providers/vertex.ts index 60e6967dd6..9fe1f2d50b 100644 --- a/src/api/providers/vertex.ts +++ b/src/api/providers/vertex.ts @@ -2,7 +2,8 @@ import { Anthropic } from "@anthropic-ai/sdk" import { AnthropicVertex } from "@anthropic-ai/vertex-sdk" import { ApiHandler } from "../" import { ApiHandlerOptions, ModelInfo, vertexDefaultModelId, VertexModelId, vertexModels } from "../../shared/api" -import { ApiStream } from "../transform/stream" +import { ApiStream, ApiStreamChunk } from "../transform/stream" +import { withRetry } from "../utils/retry" // https://docs.anthropic.com/en/api/claude-on-vertex-ai export class VertexHandler implements ApiHandler { @@ -19,59 +20,78 @@ export class VertexHandler implements ApiHandler { } async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream { - const stream = await this.client.messages.create({ - model: this.getModel().id, - max_tokens: this.getModel().info.maxTokens || 8192, - temperature: 0, - system: systemPrompt, - messages, - stream: true, - }) - for await (const chunk of stream) { - switch (chunk.type) { - case "message_start": - const usage = chunk.message.usage - yield { - type: "usage", - inputTokens: usage.input_tokens || 0, - outputTokens: usage.output_tokens || 0, - } - break - case "message_delta": - yield { - type: "usage", - inputTokens: 0, - outputTokens: chunk.usage.output_tokens || 0, - } - break + const self = this; - case "content_block_start": - switch (chunk.content_block.type) { - case "text": - if (chunk.index > 0) { - yield { - type: "text", - text: "\n", - } - } + const gen = withRetry(async () => { + const stream = await self.client.messages.create({ + model: self.getModel().id, + max_tokens: self.getModel().info.maxTokens || 8192, + temperature: 0, + system: systemPrompt, + messages, + stream: true, + }) + + return (async function*() { + for await (const chunk of stream) { + switch (chunk.type) { + case "message_start": + const usage = chunk.message.usage yield { - type: "text", - text: chunk.content_block.text, + type: "usage", + inputTokens: usage.input_tokens || 0, + outputTokens: usage.output_tokens || 0, + } as ApiStreamChunk; + break + case "message_delta": + yield { + type: "usage", + inputTokens: 0, + outputTokens: chunk.usage.output_tokens || 0, + } as ApiStreamChunk; + break + + case "content_block_start": + switch (chunk.content_block.type) { + case "text": + if (chunk.index > 0) { + yield { + type: "text", + text: "\n", + } as ApiStreamChunk; + } + yield { + type: "text", + text: chunk.content_block.text, + } as ApiStreamChunk; + break + } + break + case "content_block_delta": + switch (chunk.delta.type) { + case "text_delta": + yield { + type: "text", + text: chunk.delta.text, + } as ApiStreamChunk; + break } break } - break - case "content_block_delta": - switch (chunk.delta.type) { - case "text_delta": - yield { - type: "text", - text: chunk.delta.text, - } - break - } - break + } + })(); + }, { + maxRetries: 5, + initialDelayMs: 2000, + onRetry: (error, attempt, delayMs) => { + console.log(`Vertex request failed (attempt ${attempt})`); + console.log(`Error:`, error); + console.log(`Retrying in ${delayMs}ms...`); } + }); + + for await (const chunk of gen) { + yield chunk; } } diff --git a/src/api/transform/stream.ts b/src/api/transform/stream.ts index 0290201dad..fa534e221b 100644 --- a/src/api/transform/stream.ts +++ b/src/api/transform/stream.ts @@ -1,5 +1,5 @@ export type ApiStream = AsyncGenerator -export type ApiStreamChunk = ApiStreamTextChunk | ApiStreamUsageChunk +export type ApiStreamChunk = ApiStreamTextChunk | ApiStreamUsageChunk | ApiStreamStatusChunk export interface ApiStreamTextChunk { type: "text" @@ -14,3 +14,8 @@ export interface ApiStreamUsageChunk { cacheReadTokens?: number totalCost?: number // openrouter } + +export interface ApiStreamStatusChunk { + type: "status" + text: string +} diff --git a/src/api/utils/__tests__/retry.test.ts b/src/api/utils/__tests__/retry.test.ts new file mode 100644 index 0000000000..9778b52779 --- /dev/null +++ b/src/api/utils/__tests__/retry.test.ts @@ -0,0 +1,166 @@ +import { withRetry } from '../retry'; +import { ApiStreamChunk, ApiStreamTextChunk } from '../../transform/stream'; +import delay from 'delay'; + +jest.mock('delay'); + +describe('withRetry', () => { + beforeEach(() => { + jest.clearAllMocks(); + (delay as jest.Mock).mockResolvedValue(undefined); + }); + + it('should complete successfully with no retries', async () => { + const mockChunks: ApiStreamTextChunk[] = [ + { type: 'text', text: 'chunk1' }, + { type: 'text', text: 'chunk2' } + ]; + + const operation = async () => { + const generator = async function* () { + for (const chunk of mockChunks) { + yield chunk; + } + }; + return generator(); + }; + + const result: ApiStreamChunk[] = []; + for await (const chunk of withRetry(operation)) { + result.push(chunk); + } + + expect(result).toEqual(mockChunks); + expect(delay).not.toHaveBeenCalled(); + }); + + it('should retry on failure and eventually succeed', async () => { + let attempts = 0; + const mockChunks: ApiStreamTextChunk[] = [ + { type: 'text', text: 'success' } + ]; + + const operation = async () => { + const generator = async function* () { + attempts++; + if (attempts < 3) { + throw new Error('Temporary failure'); + } + yield mockChunks[0]; + }; + return generator(); + }; + + const onRetry = jest.fn(); + const result: ApiStreamChunk[] = []; + + for await (const chunk of withRetry(operation, { + maxRetries: 5, + initialDelayMs: 1000, + onRetry + })) { + result.push(chunk); + } + + expect(attempts).toBe(3); + expect(onRetry).toHaveBeenCalledTimes(2); + expect(delay).toHaveBeenCalledTimes(2); + expect(delay).toHaveBeenNthCalledWith(1, 1000); // First retry + expect(delay).toHaveBeenNthCalledWith(2, 2000); // Second retry with exponential backoff + + // Filter out retry status messages + const actualResults = result.filter( + (chunk): chunk is ApiStreamTextChunk => + chunk.type === 'text' && chunk.text === 'success' + ); + expect(actualResults).toEqual(mockChunks); + }); + + it('should throw after max retries exceeded', async () => { + const operation = async () => { + const generator = async function* () { + throw new Error('Persistent failure'); + }; + return generator(); + }; + + const onRetry = jest.fn(); + const generator = withRetry(operation, { + maxRetries: 3, + initialDelayMs: 1000, + onRetry + }); + + const result: ApiStreamChunk[] = []; + await expect(async () => { + for await (const chunk of generator) { + result.push(chunk); + } + }).rejects.toThrow('Persistent failure'); + + expect(onRetry).toHaveBeenCalledTimes(3); + expect(delay).toHaveBeenCalledTimes(3); + expect(delay).toHaveBeenNthCalledWith(1, 1000); + expect(delay).toHaveBeenNthCalledWith(2, 2000); + expect(delay).toHaveBeenNthCalledWith(3, 4000); + + // Should have yielded retry status messages + expect(result.length).toBe(3); + result.forEach(chunk => { + expect(chunk.type).toBe('text'); + if (chunk.type === 'text') { + expect(chunk.text).toMatch(/Request failed\. Retrying in \d+ seconds\.\.\./); + } + }); + }); + + it('should use default options when not provided', async () => { + let attempts = 0; + const operation = async () => { + const generator = async function* () { + attempts++; + if (attempts === 1) { + throw new Error('First attempt failure'); + } + yield { type: 'text' as const, text: 'success' }; + }; + return generator(); + }; + + const result: ApiStreamChunk[] = []; + for await (const chunk of withRetry(operation)) { + result.push(chunk); + } + + expect(attempts).toBe(2); + expect(delay).toHaveBeenCalledTimes(1); + expect(delay).toHaveBeenCalledWith(2000); // Default initialDelayMs + }); + + it('should handle custom onRetry callback', async () => { + const customOnRetry = jest.fn(); + let attempts = 0; + + const operation = async () => { + const generator = async function* () { + attempts++; + if (attempts === 1) { + throw new Error('Test error'); + } + yield { type: 'text' as const, text: 'success' }; + }; + return generator(); + }; + + for await (const chunk of withRetry(operation, { onRetry: customOnRetry })) { + // Consume chunks + } + + expect(customOnRetry).toHaveBeenCalledTimes(1); + expect(customOnRetry).toHaveBeenCalledWith( + expect.any(Error), + 1, + 2000 + ); + }); +}); \ No newline at end of file diff --git a/src/api/utils/retry.ts b/src/api/utils/retry.ts new file mode 100644 index 0000000000..3718ea1fbe --- /dev/null +++ b/src/api/utils/retry.ts @@ -0,0 +1,51 @@ +import delay from "delay" +import { ApiStreamChunk } from "../transform/stream" + +export interface RetryOptions { + maxRetries?: number; + initialDelayMs?: number; + onRetry?: (error: unknown, attempt: number, delayMs: number) => void; +} + +export async function* withRetry( + operation: () => Promise>, + options: RetryOptions = {} +): AsyncGenerator { + const { + maxRetries = 5, + initialDelayMs = 2000, + onRetry = (error, attempt, delayMs) => { + console.log(`Operation failed, attempt ${attempt}/${maxRetries}`); + console.log(`Error:`, error); + console.log(`Retrying in ${delayMs}ms...`); + }, + } = options; + + let attempt = 0; + + while (true) { + try { + const stream = await operation(); + for await (const chunk of stream) { + yield chunk; + } + return; + } catch (error) { + attempt++; + if (attempt > maxRetries) { + console.log(`Max retries (${maxRetries}) exceeded, giving up`); + throw error; + } + + const delayMs = initialDelayMs * Math.pow(2, attempt - 1); + onRetry(error, attempt, delayMs); + + yield { + type: "text", + text: `Request failed. Retrying in ${delayMs/1000} seconds... (attempt ${attempt}/${maxRetries})\n\n` + }; + + await delay(delayMs); + } + } +} \ No newline at end of file