55 lines
1.7 KiB
TypeScript
55 lines
1.7 KiB
TypeScript
import { ModelOption, ExpertResult } from '../../types';
|
|
import { getExpertSystemInstruction } from './prompts';
|
|
import { withRetry } from '../utils/retry';
|
|
|
|
export const streamExpertResponse = async (
|
|
ai: any,
|
|
model: ModelOption,
|
|
expert: ExpertResult,
|
|
context: string,
|
|
budget: number,
|
|
signal: AbortSignal,
|
|
onChunk: (text: string, thought: string) => void
|
|
): Promise<void> => {
|
|
// We wrap the stream initiation in retry.
|
|
// If the stream is successfully established but fails during iteration,
|
|
// we catch that separately.
|
|
const streamResult = await withRetry(() => ai.models.generateContentStream({
|
|
model: model,
|
|
contents: expert.prompt,
|
|
config: {
|
|
systemInstruction: getExpertSystemInstruction(expert.role, expert.description, context),
|
|
temperature: expert.temperature,
|
|
thinkingConfig: {
|
|
thinkingBudget: budget,
|
|
includeThoughts: true
|
|
}
|
|
}
|
|
}));
|
|
|
|
try {
|
|
for await (const chunk of streamResult) {
|
|
if (signal.aborted) break;
|
|
|
|
let chunkText = "";
|
|
let chunkThought = "";
|
|
|
|
if (chunk.candidates?.[0]?.content?.parts) {
|
|
for (const part of chunk.candidates[0].content.parts) {
|
|
if (part.thought) {
|
|
chunkThought += (part.text || "");
|
|
} else if (part.text) {
|
|
chunkText += part.text;
|
|
}
|
|
}
|
|
onChunk(chunkText, chunkThought);
|
|
}
|
|
}
|
|
} catch (streamError) {
|
|
console.error(`Stream interrupted for expert ${expert.role}:`, streamError);
|
|
// We don't retry mid-stream automatically here to avoid complex state management,
|
|
// but the initial connection is protected by withRetry.
|
|
throw streamError;
|
|
}
|
|
};
|