Jayden/AFFiNE-Mirror
0
0
Fork 0
mirror of https://github.com/toeverything/AFFiNE.git synced 2026-07-18 02:26:21 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor(core): add text stream parser (#12459)

Browse Source 
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
Download Patch File Download Diff File Expand all files Collapse all files
packages/backend/server/src/plugins/copilot/providers/openai.ts
+25 -99
View File
@@ -11,6 +11,7 @@ import {
generateObject,
generateText,
streamText,
ToolSet,
} from 'ai';
import {
@@ -30,7 +31,7 @@ import type {
PromptMessage,
} from './types';
import { CopilotProviderType, ModelInputType, ModelOutputType } from './types';
import { chatToGPTMessage, CitationParser } from './utils';
import { chatToGPTMessage, CitationParser, TextStreamParser } from './utils';
export const DEFAULT_DIMENSIONS = 256;
@@ -39,12 +40,6 @@ export type OpenAIConfig = {
baseUrl?: string;
};
type OpenAITools = {
web_search_preview: ReturnType<typeof openai.tools.webSearchPreview>;
web_search_exa: ReturnType<typeof createExaSearchTool>;
web_crawl_exa: ReturnType<typeof createExaCrawlTool>;
};
export class OpenAIProvider extends CopilotProvider<OpenAIConfig> {
readonly type = CopilotProviderType.OpenAI;
@@ -187,8 +182,6 @@ export class OpenAIProvider extends CopilotProvider<OpenAIConfig> {
private readonly MAX_STEPS = 20;
private readonly CALLOUT_PREFIX = '\n> [!]\n> ';
#instance!: VercelOpenAIProvider;
override configured(): boolean {
@@ -231,11 +224,8 @@ export class OpenAIProvider extends CopilotProvider<OpenAIConfig> {
}
}
private getTools(
options: CopilotChatOptions,
model: string
): Partial<OpenAITools> {
const tools: Partial<OpenAITools> = {};
private getTools(options: CopilotChatOptions, model: string): ToolSet {
const tools: ToolSet = {};
if (options?.tools?.length) {
for (const tool of options.tools) {
switch (tool) {
@@ -324,82 +314,34 @@ export class OpenAIProvider extends CopilotProvider<OpenAIConfig> {
providerOptions: {
openai: this.getOpenAIOptions(options, model.id),
},
tools: this.getTools(options, model.id) as OpenAITools,
tools: this.getTools(options, model.id),
maxSteps: this.MAX_STEPS,
abortSignal: options.signal,
});
const parser = new CitationParser();
let lastType;
// reasoning, tool-call, tool-result need to mark as callout
let prefix: string | null = this.CALLOUT_PREFIX;
const citationParser = new CitationParser();
const textParser = new TextStreamParser();
for await (const chunk of fullStream) {
if (chunk) {
switch (chunk.type) {
case 'text-delta': {
let result = parser.parse(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 'tool-call': {
if (prefix) {
yield prefix;
prefix = null;
}
if (chunk.toolName === 'web_search_exa') {
yield this.markAsCallout(
`\nSearching the web "${chunk.args.query}"\n`
);
}
if (chunk.toolName === 'web_crawl_exa') {
yield this.markAsCallout(
`\nCrawling the web "${chunk.args.url}"\n`
);
}
break;
}
case 'tool-result': {
if (
chunk.toolName === 'web_search_exa' &&
Array.isArray(chunk.result)
) {
yield this.markAsCallout(
`\n${this.getWebSearchLinks(chunk.result)}\n`
);
}
break;
}
case 'finish': {
const result = parser.end();
yield 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();
switch (chunk.type) {
case 'text-delta': {
let result = textParser.parse(chunk);
result = citationParser.parse(result);
yield result;
break;
}
lastType = chunk.type;
case 'finish': {
const result = citationParser.end();
yield result;
break;
}
default: {
yield textParser.parse(chunk);
break;
}
}
if (options.signal?.aborted) {
await fullStream.cancel();
break;
}
}
} catch (e: any) {
@@ -539,22 +481,6 @@ export class OpenAIProvider extends CopilotProvider<OpenAIConfig> {
return result;
}
private getWebSearchLinks(
list: {
title: string | null;
url: string;
}[]
): string {
const links = list.reduce((acc, result) => {
return acc + `\n[${result.title ?? result.url}](${result.url})\n\n`;
}, '');
return links;
}
private markAsCallout(text: string) {
return text.replaceAll('\n', '\n> ');
}
private isReasoningModel(model: string) {
// o series reasoning models
return model.startsWith('o');