Spaces:
Paused
Paused
File size: 2,020 Bytes
fb4d8fe | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 | import {
composeThinkingAndContent,
extractContentFromMessage,
extractThinkingFromMessage,
resolveFinalAssistantText,
} from "./tui-formatters.js";
type RunStreamState = {
thinkingText: string;
contentText: string;
displayText: string;
};
export class TuiStreamAssembler {
private runs = new Map<string, RunStreamState>();
private getOrCreateRun(runId: string): RunStreamState {
let state = this.runs.get(runId);
if (!state) {
state = {
thinkingText: "",
contentText: "",
displayText: "",
};
this.runs.set(runId, state);
}
return state;
}
private updateRunState(state: RunStreamState, message: unknown, showThinking: boolean) {
const thinkingText = extractThinkingFromMessage(message);
const contentText = extractContentFromMessage(message);
if (thinkingText) {
state.thinkingText = thinkingText;
}
if (contentText) {
state.contentText = contentText;
}
const displayText = composeThinkingAndContent({
thinkingText: state.thinkingText,
contentText: state.contentText,
showThinking,
});
state.displayText = displayText;
}
ingestDelta(runId: string, message: unknown, showThinking: boolean): string | null {
const state = this.getOrCreateRun(runId);
const previousDisplayText = state.displayText;
this.updateRunState(state, message, showThinking);
if (!state.displayText || state.displayText === previousDisplayText) {
return null;
}
return state.displayText;
}
finalize(runId: string, message: unknown, showThinking: boolean): string {
const state = this.getOrCreateRun(runId);
this.updateRunState(state, message, showThinking);
const finalComposed = state.displayText;
const finalText = resolveFinalAssistantText({
finalText: finalComposed,
streamedText: state.displayText,
});
this.runs.delete(runId);
return finalText;
}
drop(runId: string) {
this.runs.delete(runId);
}
}
|