mirror of
https://github.com/toeverything/AFFiNE.git
synced 2026-07-19 02:56:23 +08:00
refactor(core): add text stream parser (#12459)
Support [AI-82](https://linear.app/affine-design/issue/AI-82). Added a `TextStreamParser` class to standardize formatting of different types of AI stream chunks across providers. ### What changed? - Created a new `TextStreamParser` class in `utils.ts` that handles formatting of various chunk types (text-delta, reasoning, tool-call, tool-result, error) - Refactored the Anthropic, Gemini, and OpenAI providers to use this shared parser instead of duplicating formatting logic - Added comprehensive tests for the new `TextStreamParser` class, including tests for individual chunk types and sequences of chunks - Defined a common `AITools` type to standardize tool interfaces across providers <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **New Features** - Enhanced formatting and structure for streamed AI responses, including improved handling of callouts, web search, and web crawl results. - **Refactor** - Streamlined and unified the processing of streamed AI response chunks across providers for more consistent output. - **Bug Fixes** - Improved error handling and display for streamed responses. - **Tests** - Added comprehensive tests to ensure correct formatting and handling of various streamed message types. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
This commit is contained in:
@@ -25,7 +25,7 @@ import type {
|
||||
PromptMessage,
|
||||
} from '../types';
|
||||
import { ModelOutputType } from '../types';
|
||||
import { chatToGPTMessage } from '../utils';
|
||||
import { chatToGPTMessage, TextStreamParser } from '../utils';
|
||||
|
||||
export const DEFAULT_DIMENSIONS = 256;
|
||||
|
||||
@@ -37,8 +37,6 @@ export type GeminiConfig = {
|
||||
export abstract class GeminiProvider<T> extends CopilotProvider<T> {
|
||||
private readonly MAX_STEPS = 20;
|
||||
|
||||
private readonly CALLOUT_PREFIX = '\n> [!]\n> ';
|
||||
|
||||
protected abstract instance:
|
||||
| GoogleGenerativeAIProvider
|
||||
| GoogleVertexProvider;
|
||||
@@ -167,42 +165,13 @@ export abstract class GeminiProvider<T> extends CopilotProvider<T> {
|
||||
},
|
||||
});
|
||||
|
||||
let lastType;
|
||||
// reasoning, tool-call, tool-result need to mark as callout
|
||||
let prefix: string | null = this.CALLOUT_PREFIX;
|
||||
const parser = new TextStreamParser();
|
||||
for await (const chunk of fullStream) {
|
||||
if (chunk) {
|
||||
switch (chunk.type) {
|
||||
case 'text-delta': {
|
||||
let result = chunk.textDelta;
|
||||
if (lastType !== chunk.type) {
|
||||
result = '\n\n' + result;
|
||||
}
|
||||
yield result;
|
||||
break;
|
||||
}
|
||||
case 'reasoning': {
|
||||
if (prefix) {
|
||||
yield prefix;
|
||||
prefix = null;
|
||||
}
|
||||
let result = chunk.textDelta;
|
||||
if (lastType !== chunk.type) {
|
||||
result = '\n\n' + result;
|
||||
}
|
||||
yield this.markAsCallout(result);
|
||||
break;
|
||||
}
|
||||
case 'error': {
|
||||
const error = chunk.error as { type: string; message: string };
|
||||
throw new Error(error.message);
|
||||
}
|
||||
}
|
||||
if (options.signal?.aborted) {
|
||||
await fullStream.cancel();
|
||||
break;
|
||||
}
|
||||
lastType = chunk.type;
|
||||
const result = parser.parse(chunk);
|
||||
yield result;
|
||||
if (options.signal?.aborted) {
|
||||
await fullStream.cancel();
|
||||
break;
|
||||
}
|
||||
}
|
||||
} catch (e: any) {
|
||||
@@ -222,10 +191,6 @@ export abstract class GeminiProvider<T> extends CopilotProvider<T> {
|
||||
return result;
|
||||
}
|
||||
|
||||
private markAsCallout(text: string) {
|
||||
return text.replaceAll('\n', '\n> ');
|
||||
}
|
||||
|
||||
private isReasoningModel(model: string) {
|
||||
return model.startsWith('gemini-2.5');
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user