92 lines
2.6 KiB
JavaScript
92 lines
2.6 KiB
JavaScript
const Anthropic = require('@anthropic-ai/sdk');
|
|
const OpenAI = require('openai');
|
|
|
|
/**
|
|
* Normalize Anthropic + OpenAI-compatible streams into one AsyncIterable.
|
|
* Yields: { token, done, fullText } events.
|
|
* Handles errors gracefully — emits error event, doesn't crash.
|
|
*
|
|
* Usage:
|
|
* for await (const chunk of streamCompletion(model, messages, systemPrompt)) {
|
|
* // chunk = { token: string, done: boolean, fullText: string }
|
|
* }
|
|
*/
|
|
async function* streamCompletion(model, messages, systemPrompt) {
|
|
if (!model || !model.provider) {
|
|
yield { token: '', done: true, fullText: '', error: 'Model or provider not specified' };
|
|
return;
|
|
}
|
|
|
|
try {
|
|
if (model.provider === 'anthropic') {
|
|
yield* streamAnthropic(model, messages, systemPrompt);
|
|
} else if (model.provider === 'openai') {
|
|
yield* streamOpenAI(model, messages, systemPrompt);
|
|
} else {
|
|
yield { token: '', done: true, fullText: '', error: `Unknown provider: ${model.provider}` };
|
|
}
|
|
} catch (err) {
|
|
console.error('[llm] streamCompletion error:', err.message);
|
|
yield { token: '', done: true, fullText: '', error: err.message };
|
|
}
|
|
}
|
|
|
|
async function* streamAnthropic(model, messages, systemPrompt) {
|
|
const client = new Anthropic({
|
|
apiKey: model.api_key || process.env.ANTHROPIC_API_KEY,
|
|
});
|
|
|
|
const stream = await client.messages.create({
|
|
model: model.name,
|
|
max_tokens: 4096,
|
|
system: systemPrompt || undefined,
|
|
messages: messages.map(m => ({ role: m.role, content: m.content })),
|
|
stream: true,
|
|
});
|
|
|
|
let fullText = '';
|
|
|
|
for await (const event of stream) {
|
|
if (event.type === 'content_block_delta' && event.delta?.type === 'text_delta') {
|
|
const token = event.delta.text || '';
|
|
fullText += token;
|
|
yield { token, done: false, fullText };
|
|
}
|
|
}
|
|
|
|
yield { token: '', done: true, fullText };
|
|
}
|
|
|
|
async function* streamOpenAI(model, messages, systemPrompt) {
|
|
const client = new OpenAI({
|
|
apiKey: model.api_key || process.env.OPENAI_API_KEY,
|
|
baseURL: model.api_base,
|
|
});
|
|
|
|
const openaiMessages = [];
|
|
if (systemPrompt) {
|
|
openaiMessages.push({ role: 'system', content: systemPrompt });
|
|
}
|
|
openaiMessages.push(...messages.map(m => ({ role: m.role, content: m.content })));
|
|
|
|
const stream = await client.chat.completions.create({
|
|
model: model.name,
|
|
messages: openaiMessages,
|
|
stream: true,
|
|
});
|
|
|
|
let fullText = '';
|
|
|
|
for await (const chunk of stream) {
|
|
const token = chunk.choices?.[0]?.delta?.content || '';
|
|
if (token) {
|
|
fullText += token;
|
|
yield { token, done: false, fullText };
|
|
}
|
|
}
|
|
|
|
yield { token: '', done: true, fullText };
|
|
}
|
|
|
|
module.exports = { streamCompletion };
|