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:
akumatus
2025-05-27 07:14:17 +00:00
parent 83caf98618
commit f4cba7d6ee
5 changed files with 355 additions and 224 deletions
@@ -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');
}