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',
|
||||
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);
|
||||
});
|
||||
});
|
||||
|
||||
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 type { ChatRequest, ChatResponse, ModelClient } from './types.js';
|
||||
import type { ChatRequest, ChatResponse, ChatStreamEvent, ModelClient } from './types.js';
|
||||
|
||||
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;
|
||||
maxTokens?: number;
|
||||
}
|
||||
@@ -15,6 +16,7 @@ export class AnthropicClient implements ModelClient {
|
||||
constructor(config: AnthropicClientConfig) {
|
||||
this.client = new Anthropic({
|
||||
apiKey: config.apiKey,
|
||||
authToken: config.authToken,
|
||||
});
|
||||
this.model = config.model;
|
||||
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