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
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
|
import { AssistantMessage, ASSISTANT_ROLE, TEXT_TYPE, Citation, CHUNK_TYPE } from './types';
import { v4 as uuidv4 } from 'uuid';
export class StreamParser {
private currentMessage: AssistantMessage;
private currentTag: string | null = null;
private buffer: string = '';
private citationIndex: number = 1;
constructor() {
this.currentMessage = {
role: ASSISTANT_ROLE.ASSISTANT,
content: [],
thoughts: [],
actions: [],
citations: [],
};
}
parse(chunk: string): AssistantMessage {
this.buffer += chunk;
while (this.buffer.length > 0) {
if (this.currentTag === null) {
const openTagMatch = this.buffer.match(/<(\w+)>/);
if (openTagMatch) {
this.currentTag = openTagMatch[1];
this.buffer = this.buffer.slice(openTagMatch.index! + openTagMatch[0].length);
} else {
break;
}
} else {
const closeTagIndex = this.buffer.indexOf(`</${this.currentTag}>`);
if (closeTagIndex !== -1) {
const content = this.buffer.slice(0, closeTagIndex);
this.processTag(this.currentTag, content);
this.buffer = this.buffer.slice(closeTagIndex + this.currentTag.length + 3);
this.currentTag = null;
} else {
break;
}
}
}
return this.currentMessage;
}
private processTag(tag: string, content: string) {
switch (tag) {
case 'thought':
this.currentMessage.thoughts!.push(content);
break;
case 'action':
this.currentMessage.actions!.push({ index: this.currentMessage.actions!.length, action: content, action_input: '' });
break;
case 'action_input':
if (this.currentMessage.actions!.length > 0) {
this.currentMessage.actions![this.currentMessage.actions!.length - 1].action_input = content;
}
break;
case 'answer':
this.processAnswer(content);
break;
}
}
private processAnswer(content: string) {
const groundedTextRegex = /<grounded_text citation_index="([^"]+)">([\s\S]*?)<\/grounded_text>/g;
let lastIndex = 0;
let match;
while ((match = groundedTextRegex.exec(content)) !== null) {
const [fullMatch, citationIndex, groundedText] = match;
// Add normal text before the grounded text
if (match.index > lastIndex) {
const normalText = content.slice(lastIndex, match.index).trim();
if (normalText) {
this.currentMessage.content.push({
index: this.currentMessage.content.length,
type: TEXT_TYPE.NORMAL,
text: normalText,
citation_ids: null,
});
}
}
// Add grounded text
const citation_id = uuidv4();
this.currentMessage.content.push({
index: this.currentMessage.content.length,
type: TEXT_TYPE.GROUNDED,
text: groundedText.trim(),
citation_ids: [citation_id],
});
// Add citation
this.currentMessage.citations!.push({
citation_id,
chunk_id: '',
type: CHUNK_TYPE.TEXT,
direct_text: '',
});
lastIndex = match.index + fullMatch.length;
}
// Add any remaining normal text after the last grounded text
if (lastIndex < content.length) {
const remainingText = content.slice(lastIndex).trim();
if (remainingText) {
this.currentMessage.content.push({
index: this.currentMessage.content.length,
type: TEXT_TYPE.NORMAL,
text: remainingText,
citation_ids: null,
});
}
}
}
getResult(): AssistantMessage {
return this.currentMessage;
}
}
|