feat(models): add streaming support to AnthropicClient
This commit is contained in:
@@ -10,6 +10,16 @@ vi.mock('@anthropic-ai/sdk', () => ({
|
|||||||
stop_reason: 'end_turn',
|
stop_reason: 'end_turn',
|
||||||
usage: { input_tokens: 10, output_tokens: 5 },
|
usage: { input_tokens: 10, output_tokens: 5 },
|
||||||
}),
|
}),
|
||||||
|
stream: vi.fn().mockReturnValue({
|
||||||
|
[Symbol.asyncIterator]: async function* () {
|
||||||
|
yield { type: 'content_block_delta', delta: { type: 'text_delta', text: 'Hello ' } };
|
||||||
|
yield { type: 'content_block_delta', delta: { type: 'text_delta', text: 'from ' } };
|
||||||
|
yield { type: 'content_block_delta', delta: { type: 'text_delta', text: 'Claude!' } };
|
||||||
|
},
|
||||||
|
finalMessage: vi.fn().mockResolvedValue({
|
||||||
|
usage: { input_tokens: 10, output_tokens: 5 },
|
||||||
|
}),
|
||||||
|
}),
|
||||||
},
|
},
|
||||||
})),
|
})),
|
||||||
}));
|
}));
|
||||||
@@ -31,3 +41,30 @@ describe('AnthropicClient', () => {
|
|||||||
expect(response.usage.outputTokens).toBe(5);
|
expect(response.usage.outputTokens).toBe(5);
|
||||||
});
|
});
|
||||||
});
|
});
|
||||||
|
|
||||||
|
describe('AnthropicClient streaming', () => {
|
||||||
|
it('streams messages chunk by chunk', async () => {
|
||||||
|
const client = new AnthropicClient({
|
||||||
|
apiKey: 'test-key',
|
||||||
|
model: 'claude-sonnet-4-20250514',
|
||||||
|
});
|
||||||
|
|
||||||
|
const chunks: string[] = [];
|
||||||
|
let finalUsage: { inputTokens: number; outputTokens: number } | undefined;
|
||||||
|
|
||||||
|
for await (const event of client.chatStream({
|
||||||
|
messages: [{ role: 'user', content: 'Hello' }],
|
||||||
|
})) {
|
||||||
|
if (event.type === 'content' && event.content) {
|
||||||
|
chunks.push(event.content);
|
||||||
|
}
|
||||||
|
if (event.type === 'done' && event.usage) {
|
||||||
|
finalUsage = event.usage;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
expect(chunks.length).toBeGreaterThan(0);
|
||||||
|
expect(chunks.join('')).toBe('Hello from Claude!');
|
||||||
|
expect(finalUsage).toEqual({ inputTokens: 10, outputTokens: 5 });
|
||||||
|
});
|
||||||
|
});
|
||||||
|
|||||||
+38
-2
@@ -1,8 +1,9 @@
|
|||||||
import Anthropic from '@anthropic-ai/sdk';
|
import Anthropic from '@anthropic-ai/sdk';
|
||||||
import type { ChatRequest, ChatResponse, ModelClient } from './types.js';
|
import type { ChatRequest, ChatResponse, ChatStreamEvent, ModelClient } from './types.js';
|
||||||
|
|
||||||
export interface AnthropicClientConfig {
|
export interface AnthropicClientConfig {
|
||||||
apiKey?: string; // Falls back to ANTHROPIC_API_KEY env var
|
apiKey?: string; // Falls back to ANTHROPIC_API_KEY env var
|
||||||
|
authToken?: string; // Alternative: use auth token instead of API key
|
||||||
model: string;
|
model: string;
|
||||||
maxTokens?: number;
|
maxTokens?: number;
|
||||||
}
|
}
|
||||||
@@ -15,6 +16,7 @@ export class AnthropicClient implements ModelClient {
|
|||||||
constructor(config: AnthropicClientConfig) {
|
constructor(config: AnthropicClientConfig) {
|
||||||
this.client = new Anthropic({
|
this.client = new Anthropic({
|
||||||
apiKey: config.apiKey,
|
apiKey: config.apiKey,
|
||||||
|
authToken: config.authToken,
|
||||||
});
|
});
|
||||||
this.model = config.model;
|
this.model = config.model;
|
||||||
this.defaultMaxTokens = config.maxTokens ?? 4096;
|
this.defaultMaxTokens = config.maxTokens ?? 4096;
|
||||||
@@ -43,4 +45,38 @@ export class AnthropicClient implements ModelClient {
|
|||||||
},
|
},
|
||||||
};
|
};
|
||||||
}
|
}
|
||||||
|
|
||||||
|
async *chatStream(request: ChatRequest): AsyncIterable<ChatStreamEvent> {
|
||||||
|
const stream = this.client.messages.stream({
|
||||||
|
model: this.model,
|
||||||
|
max_tokens: request.maxTokens ?? this.defaultMaxTokens,
|
||||||
|
system: request.system,
|
||||||
|
messages: request.messages.map((m) => ({
|
||||||
|
role: m.role,
|
||||||
|
content: m.content,
|
||||||
|
})),
|
||||||
|
});
|
||||||
|
|
||||||
|
try {
|
||||||
|
for await (const event of stream) {
|
||||||
|
if (event.type === 'content_block_delta' && event.delta.type === 'text_delta') {
|
||||||
|
yield { type: 'content', content: event.delta.text };
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
const finalMessage = await stream.finalMessage();
|
||||||
|
yield {
|
||||||
|
type: 'done',
|
||||||
|
usage: {
|
||||||
|
inputTokens: finalMessage.usage.input_tokens,
|
||||||
|
outputTokens: finalMessage.usage.output_tokens,
|
||||||
|
},
|
||||||
|
};
|
||||||
|
} catch (error) {
|
||||||
|
yield {
|
||||||
|
type: 'error',
|
||||||
|
error: error instanceof Error ? error : new Error(String(error)),
|
||||||
|
};
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user