mirror of
https://github.com/RooVetGit/Roo-Code.git
synced 2026-09-09 22:31:08 +00:00
Retry on provider failures
This commit is contained in:
parent
130c5ff779
commit
c9e3e60102
17 changed files with 1112 additions and 476 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
4
package-lock.json
generated
4
package-lock.json
generated
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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<typeof OpenAI>).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<typeof OpenAI>).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<typeof OpenAI>).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<typeof OpenAI>).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<typeof OpenAI>).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<typeof OpenAI>).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
|
||||
}))
|
||||
})
|
||||
})
|
||||
|
|
|
|||
|
|
@ -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<Anthropic.Beta.PromptCaching.Messages.RawPromptCachingBetaMessageStreamEvent>
|
||||
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<Anthropic.Beta.PromptCaching.Messages.RawPromptCachingBetaMessageStreamEvent>
|
||||
|
||||
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;
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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<ApiStreamChunk> {
|
||||
// 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<OpenAI.Chat.Completions.ChatCompletionChunk>) {
|
||||
// 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 }
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
export type ApiStream = AsyncGenerator<ApiStreamChunk>
|
||||
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
|
||||
}
|
||||
|
|
|
|||
166
src/api/utils/__tests__/retry.test.ts
Normal file
166
src/api/utils/__tests__/retry.test.ts
Normal file
|
|
@ -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
|
||||
);
|
||||
});
|
||||
});
|
||||
51
src/api/utils/retry.ts
Normal file
51
src/api/utils/retry.ts
Normal file
|
|
@ -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<AsyncGenerator<ApiStreamChunk>>,
|
||||
options: RetryOptions = {}
|
||||
): AsyncGenerator<ApiStreamChunk> {
|
||||
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);
|
||||
}
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue