Skip to content

Commit 9b7cd5c

Browse files
pavel-esirpavel-esir
committed
Fix pytests
1 parent 2be242f commit 9b7cd5c

File tree

3 files changed

+118
-50
lines changed

3 files changed

+118
-50
lines changed

‎src/cpp/src/parsers.cpp‎

Lines changed: 49 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -25,12 +25,12 @@ class ReasoningIncrementalParser::ReasoningParserImpl {
2525
* @brief Ensure required fields exist in the message container.
2626
*/
2727
void ensure_message_fields(JsonContainer& message) {
28-
if (!message.contains("reasoning_content")) {
28+
// if (!message.contains("reasoning_content")) {
2929
message["reasoning_content"] = "";
30-
}
31-
if (!message.contains("content")) {
30+
// }
31+
// if (!message.contains("content")) {
3232
message["content"] = "";
33-
}
33+
// }
3434
}
3535

3636
/**
@@ -59,11 +59,15 @@ class ReasoningIncrementalParser::ReasoningParserImpl {
5959
* @brief Handle the case where both open and close tags are found in the same chunk.
6060
*/
6161
void handle_complete_reasoning(JsonContainer& message, std::string_view txt_chunk,
62-
size_t open_idx, size_t close_idx) {
62+
size_t open_idx, size_t close_idx, std::string& delta_text) {
6363
// Extract reasoning content between tags
6464
message["reasoning_content"] = std::string(txt_chunk.substr(open_idx + m_open_tag.size(), close_idx - (open_idx + m_open_tag.size())));
6565
message["content"] = std::string(txt_chunk.substr(close_idx + m_close_tag.size()));
6666

67+
if (!m_keep_original_content) {
68+
delta_text = std::string(txt_chunk.substr(close_idx + m_close_tag.size()));
69+
}
70+
6771
m_think_tag_opened = false;
6872
m_deactivated = true;
6973
m_text_cache.clear();
@@ -73,10 +77,16 @@ class ReasoningIncrementalParser::ReasoningParserImpl {
7377
* @brief Handle the case where only the open tag is found.
7478
*/
7579
void handle_open_tag(JsonContainer& message, std::string& reason_str,
76-
std::string_view txt_chunk, size_t open_idx) {
80+
std::string_view txt_chunk, size_t open_idx, std::string& delta_text) {
7781
// Start accumulating reasoning content
78-
reason_str.append(txt_chunk.substr(open_idx + m_open_tag.size()));
79-
message["reasoning_content"] = std::move(reason_str);
82+
// reason_str.append(txt_chunk.substr(open_idx + m_open_tag.size()));
83+
message["reasoning_content"] = std::string(txt_chunk.substr(open_idx + m_open_tag.size()));
84+
85+
if (!m_keep_original_content) {
86+
delta_text.clear();
87+
} else {
88+
delta_text = txt_chunk;
89+
}
8090

8191
m_think_tag_opened = true;
8292
m_text_cache.clear();
@@ -85,10 +95,21 @@ class ReasoningIncrementalParser::ReasoningParserImpl {
8595
/**
8696
* @brief Handle the case where the close tag is found.
8797
*/
88-
void handle_close_tag(JsonContainer& message, std::string_view txt_chunk, size_t close_idx) {
98+
void handle_close_tag(JsonContainer& message, std::string_view txt_chunk, size_t close_idx, std::string& delta_text) {
8999
// Append text before close tag to reasoning content
90100
message["reasoning_content"] = std::move(std::string(txt_chunk.substr(0, close_idx)));
91-
message["content"] = std::string(txt_chunk.substr(close_idx + m_close_tag.size()));;
101+
auto content = std::string(txt_chunk.substr(close_idx + m_close_tag.size()));
102+
message["content"] = content;
103+
104+
if (!m_keep_original_content) {
105+
// Despite the fact that we put txt_chung to delta_text it's correct.
106+
// Since txt_chunk contains some cached parts from the previous calls that were not yet processed yet
107+
// and we kept them in cache until we decide what to do with them. Here we definitely know that that cached parts
108+
// belonged to reasoning_content so we can discard them.
109+
delta_text = content;
110+
} else {
111+
delta_text = txt_chunk;
112+
}
92113

93114
m_text_cache.clear();
94115
m_think_tag_opened = false;
@@ -98,21 +119,28 @@ class ReasoningIncrementalParser::ReasoningParserImpl {
98119
/**
99120
* @brief Handle accumulating text while inside reasoning tags.
100121
*/
101-
void handle_inside_reasoning(JsonContainer& message, std::string& reason_str, std::string_view txt_chunk) {
122+
void handle_inside_reasoning(JsonContainer& message, std::string& reason_str, std::string_view txt_chunk, std::string& delta_text) {
102123
// Find if the end of txt_chunk might be the start of a close tag
103124
const size_t num_chars_to_keep = find_close_tag_prefix_length(txt_chunk);
104125

105126
if (num_chars_to_keep > 0) {
106127
// Keep potential partial close tag in cache
107128
m_text_cache = std::string(txt_chunk.substr(txt_chunk.size() - num_chars_to_keep));
108-
reason_str.append(txt_chunk.substr(0, txt_chunk.size() - num_chars_to_keep));
129+
// reason_str.append(txt_chunk.substr(0, txt_chunk.size() - num_chars_to_keep));
130+
reason_str = txt_chunk.substr(0, txt_chunk.size() - num_chars_to_keep);
131+
132+
// TODO:
133+
delta_text = std::string(txt_chunk.substr(0, txt_chunk.size() - num_chars_to_keep));
109134
} else {
110135
// No partial close tag, accumulate all text
111-
reason_str.append(txt_chunk);
136+
// reason_str.append(txt_chunk);
137+
reason_str = txt_chunk;
112138
m_text_cache.clear();
113139
}
114-
115140
message["reasoning_content"] = std::move(reason_str);
141+
if (!m_keep_original_content) {
142+
delta_text.clear();
143+
}
116144
}
117145

118146
public:
@@ -132,6 +160,7 @@ class ReasoningIncrementalParser::ReasoningParserImpl {
132160
std::string& delta_text,
133161
const std::optional<std::vector<int64_t>>& delta_tokens
134162
) {
163+
ensure_message_fields(message);
135164
if (m_deactivated) {
136165
message["content"] = delta_text;
137166
return delta_text;
@@ -141,9 +170,8 @@ class ReasoningIncrementalParser::ReasoningParserImpl {
141170
}
142171
m_first_run = false;
143172

144-
ensure_message_fields(message);
145173

146-
const std::string txt_chunk = m_text_cache + delta_text;
174+
std::string txt_chunk = m_text_cache + delta_text;
147175
std::string reason_str = std::move(message["reasoning_content"].get_string());
148176

149177
// Cache find() results to avoid redundant searches
@@ -156,18 +184,19 @@ class ReasoningIncrementalParser::ReasoningParserImpl {
156184
? close_idx : std::string::npos;
157185

158186
if (close_idx_after_open != std::string::npos) {
159-
handle_complete_reasoning(message, txt_chunk, open_idx, close_idx_after_open);
187+
handle_complete_reasoning(message, txt_chunk, open_idx, close_idx_after_open, delta_text);
160188
} else {
161-
handle_open_tag(message, reason_str, txt_chunk, open_idx);
189+
handle_open_tag(message, reason_str, txt_chunk, open_idx, delta_text);
162190
}
163191
} else if (m_think_tag_opened && close_idx != std::string::npos) {
164-
handle_close_tag(message, txt_chunk, close_idx);
192+
handle_close_tag(message, txt_chunk, close_idx, delta_text);
165193
} else if (m_think_tag_opened) {
166-
handle_inside_reasoning(message, reason_str, txt_chunk);
194+
handle_inside_reasoning(message, reason_str, txt_chunk, delta_text);
167195
} else {
168196
// Think tag was not opened yet and not found in the current delta_text.
169197
// Accumulate text in the cache to detect if <think> is split between several delta_text pieces.
170198
m_text_cache += delta_text;
199+
delta_text.clear();
171200
}
172201

173202
return delta_text;

‎src/cpp/src/text_streamer.cpp‎

Lines changed: 15 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -144,7 +144,7 @@ TextParserStreamerImpl(std::vector<std::shared_ptr<IncrementalParser>> parsers)
144144

145145
};
146146

147-
void concatenate_json_containers(JsonContainer& from, const JsonContainer& to, std::vector<std::string> keys_to_concatenate) {
147+
void concatenate_json_containers(const JsonContainer& from, JsonContainer& to, std::vector<std::string> keys_to_concatenate) {
148148
for (const auto& key : keys_to_concatenate) {
149149
if (to.contains(key) && from.contains(key)) {
150150
// If both are strings, concatenate
@@ -161,7 +161,9 @@ void concatenate_json_containers(JsonContainer& from, const JsonContainer& to, s
161161
TextParserStreamer::TextParserStreamer(const Tokenizer& tokenizer, std::vector<std::shared_ptr<IncrementalParser>> parsers)
162162
: TextStreamer(tokenizer, [this](std::string s) -> CallbackTypeVariant {
163163
return this->write(s);
164-
}), m_pimpl{std::make_unique<TextParserStreamerImpl>(parsers)} {}
164+
}), m_pimpl{std::make_unique<TextParserStreamerImpl>(parsers)} {
165+
m_pimpl->m_parsed_message["content"] = "";
166+
}
165167

166168
CallbackTypeVariant TextParserStreamer::write(std::string message) {
167169
// When 'write' is called with string, it means new chunk of tokens is decoded into text
@@ -199,10 +201,18 @@ CallbackTypeVariant TextParserStreamer::write(std::string message) {
199201
// Message can be modified inside parser, if parser for example extracted tool calling from message content
200202
}
201203

204+
// std::cout << msg["content"].get_string() << std::endl;
205+
// std::cout << msg["reasoning_content"].get_string() << std::endl;
206+
202207
// concatenate msg with m_parsed_message
203-
concatenate_json_containers(msg, m_pimpl->m_parsed_message, {"content", "reasoning_content"});
208+
concatenate_json_containers(msg, m_pimpl->m_parsed_message, {"reasoning_content"});
209+
210+
// We have to put into DeltaMessage's "content" fields only chunks that belong to content.
211+
// But into m_parsed_message["content"] we need to accumulate full content if m_keep_original_content == True
212+
// and if m_keep_original_content == False only part that is outside reasoning tags and outside tool calls.
213+
214+
m_pimpl->m_parsed_message["content"] = m_pimpl->m_parsed_message["content"].get_string() + message;
204215

205-
// return write(m_pimpl->m_parsed_message);
206216
return write(msg);
207217
}
208218

@@ -212,6 +222,7 @@ JsonContainer TextParserStreamer::get_parsed_message() const {
212222

213223
void TextParserStreamer::reset() {
214224
m_pimpl->m_parsed_message = JsonContainer();
225+
m_pimpl->m_parsed_message["content"] = "";
215226
for (auto& parser : m_pimpl->m_parsers) {
216227
parser->reset();
217228
}

0 commit comments

Comments
 (0)