Skip to content

Commit 48e756e

Browse files
JeanMertzJeanMertz
authored
fix(llm, ollama): Ensure extractor is drained before mapping (#138)
The `ReasoningExtractor.handle()` call has been moved from the `map_content` function to the main stream processing loop and `map_response` function. This ensures reasoning extraction and finalization happens immediately when content is received, whereas previously finalization was accidentally skipped. This ensures all content is processed/flushed fully, instead of partially. Signed-off-by: Jean Mertz <git@jeanmertz.com>
1 parent cf61494 commit 48e756e

File tree

5 files changed

+33
-35
lines changed

5 files changed

+33
-35
lines changed

‎crates/jp_llm/src/provider/ollama.rs‎

Lines changed: 23 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -68,7 +68,10 @@ impl Provider for Ollama {
6868
tokio::pin!(stream);
6969
while let Some(event) = stream.next().await {
7070
let events = event
71-
.map(|event| map_event(event, &mut current_state, &mut extractor))
71+
.map(|event| {
72+
extractor.handle(&event.message.content);
73+
map_event(event, &mut current_state, &mut extractor)
74+
})
7275
.unwrap_or_default();
7376

7477
for event in events {
@@ -77,7 +80,7 @@ impl Provider for Ollama {
7780
}
7881

7982
extractor.finalize();
80-
for event in map_content("", &mut current_state, &mut extractor) {
83+
for event in map_content(&mut current_state, &mut extractor) {
8184
yield event;
8285
}
8386
});
@@ -98,23 +101,23 @@ fn map_model(model: LocalModel) -> ModelDetails {
98101
}
99102

100103
fn map_response(response: ChatMessageResponse) -> Result<Vec<Event>> {
101-
map_event(
102-
response,
103-
&mut AccumulationState::default(),
104-
&mut ReasoningExtractor::default(),
105-
)
106-
.into_iter()
107-
.map(|v| {
108-
v.map(|e| match e {
109-
StreamEvent::ChatChunk(content) => match content {
110-
CompletionChunk::Content(s) => Event::Content(s),
111-
CompletionChunk::Reasoning(s) => Event::Reasoning(s),
112-
},
113-
StreamEvent::ToolCall(request) => Event::ToolCall(request),
114-
StreamEvent::Metadata(key, metadata) => Event::Metadata(key, metadata),
104+
let mut extractor = ReasoningExtractor::default();
105+
extractor.handle(&response.message.content);
106+
extractor.finalize();
107+
108+
map_event(response, &mut AccumulationState::default(), &mut extractor)
109+
.into_iter()
110+
.map(|v| {
111+
v.map(|e| match e {
112+
StreamEvent::ChatChunk(content) => match content {
113+
CompletionChunk::Content(s) => Event::Content(s),
114+
CompletionChunk::Reasoning(s) => Event::Reasoning(s),
115+
},
116+
StreamEvent::ToolCall(request) => Event::ToolCall(request),
117+
StreamEvent::Metadata(key, metadata) => Event::Metadata(key, metadata),
118+
})
115119
})
116-
})
117-
.collect::<Result<Vec<_>>>()
120+
.collect::<Result<Vec<_>>>()
118121
}
119122

120123
fn map_event(
@@ -135,18 +138,15 @@ fn map_event(
135138
events.extend(handle_delta(delta, state).transpose());
136139
}
137140

138-
events.extend(map_content(&event.message.content, state, extractor));
141+
events.extend(map_content(state, extractor));
139142
events
140143
}
141144

142145
fn map_content(
143-
content: &str,
144146
state: &mut AccumulationState,
145147
extractor: &mut ReasoningExtractor,
146148
) -> Vec<Result<StreamEvent>> {
147149
let mut events = Vec::new();
148-
extractor.handle(content);
149-
150150
if !extractor.reasoning.is_empty() {
151151
let reasoning = mem::take(&mut extractor.reasoning);
152152
events.extend(handle_delta(Delta::reasoning(reasoning), state).transpose());
@@ -504,7 +504,7 @@ impl From<ToolCall> for Delta {
504504

505505
#[derive(Default, Debug)]
506506
/// A parser that segments a stream of text into 'reasoning' and 'other' buckets.
507-
/// It handles streams with or without a <think> block.
507+
/// It handles streams with or without a `<think>` block.
508508
pub struct ReasoningExtractor {
509509
pub other: String,
510510
pub reasoning: String,
@@ -681,10 +681,6 @@ mod tests {
681681
.unwrap()
682682
.chat_completion(&model.into(), query)
683683
.await
684-
.map(|mut v| {
685-
v.truncate(10);
686-
v
687-
})
688684
},
689685
)
690686
.await

‎crates/jp_llm/tests/fixtures/test_ollama_chat_completion.snap‎

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -3,9 +3,11 @@ source: crates/jp_test/src/mock.rs
33
expression: expr
44
---
55
Ok(
6-
[
7-
Content(
8-
"It looks like you're just sending a test message! That's perfectly fine. If you're ready to start our conversation, I'm here to chat and help with any questions or topics you'd like to discuss. What's on your mind?",
9-
),
10-
],
6+
Reply(
7+
[
8+
Content(
9+
"This is an automated response from the AI. You sent a test message, and I'm here to acknowledge it! If you have any questions or need assistance with anything, feel free to ask.",
10+
),
11+
],
12+
),
1113
)

‎crates/jp_llm/tests/fixtures/test_ollama_chat_completion.yml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,4 +7,4 @@ then:
77
header:
88
- name: content-type
99
value: application/json; charset=utf-8
10-
body: "{\"model\":\"llama3:latest\",\"created_at\":\"2025-05-22T07:57:38.792897Z\",\"message\":{\"role\":\"assistant\",\"content\":\"It looks like you're just sending a test message! That's perfectly fine. If you're ready to start our conversation, I'm here to chat and help with any questions or topics you'd like to discuss. What's on your mind?\"},\"done_reason\":\"stop\",\"done\":true,\"total_duration\":6774768167,\"load_duration\":5236091667,\"prompt_eval_count\":12,\"prompt_eval_duration\":363966666,\"eval_count\":50,\"eval_duration\":1173898417}"
10+
body: "{\"model\":\"llama3:latest\",\"created_at\":\"2025-06-06T13:23:20.747198Z\",\"message\":{\"role\":\"assistant\",\"content\":\"This is an automated response from the AI. You sent a test message, and I'm here to acknowledge it! If you have any questions or need assistance with anything, feel free to ask.\"},\"done_reason\":\"stop\",\"done\":true,\"total_duration\":3708365708,\"load_duration\":2493309291,\"prompt_eval_count\":12,\"prompt_eval_duration\":386527500,\"eval_count\":40,\"eval_duration\":827594542}"

‎crates/jp_llm/tests/fixtures/test_ollama_chat_completion_stream.snap‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,4 +2,4 @@
22
source: crates/jp_test/src/mock.rs
33
expression: expr
44
---
5-
"This appears to be a test message. I'm here and ready to assist you with any questions or topics you'd like to discuss!"
5+
"This is a test message. I'm receiving it just fine!"

‎crates/jp_llm/tests/fixtures/test_ollama_chat_completion_stream.yml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,4 +7,4 @@ then:
77
header:
88
- name: content-type
99
value: application/x-ndjson
10-
body: "{\"model\":\"llama3:latest\",\"created_at\":\"2025-05-22T07:57:37.619275Z\",\"message\":{\"role\":\"assistant\",\"content\":\"This\"},\"done\":false}\n{\"model\":\"llama3:latest\",\"created_at\":\"2025-05-22T07:57:37.649229Z\",\"message\":{\"role\":\"assistant\",\"content\":\" appears\"},\"done\":false}\n{\"model\":\"llama3:latest\",\"created_at\":\"2025-05-22T07:57:37.677752Z\",\"message\":{\"role\":\"assistant\",\"content\":\" to\"},\"done\":false}\n{\"model\":\"llama3:latest\",\"created_at\":\"2025-05-22T07:57:37.706284Z\",\"message\":{\"role\":\"assistant\",\"content\":\" be\"},\"done\":false}\n{\"model\":\"llama3:latest\",\"created_at\":\"2025-05-22T07:57:37.73435Z\",\"message\":{\"role\":\"assistant\",\"content\":\" a\"},\"done\":false}\n{\"model\":\"llama3:latest\",\"created_at\":\"2025-05-22T07:57:37.763966Z\",\"message\":{\"role\":\"assistant\",\"content\":\" test\"},\"done\":false}\n{\"model\":\"llama3:latest\",\"created_at\":\"2025-05-22T07:57:37.792762Z\",\"message\":{\"role\":\"assistant\",\"content\":\" message\"},\"done\":false}\n{\"model\":\"llama3:latest\",\"created_at\":\"2025-05-22T07:57:37.822461Z\",\"message\":{\"role\":\"assistant\",\"content\":\".\"},\"done\":false}\n{\"model\":\"llama3:latest\",\"created_at\":\"2025-05-22T07:57:37.851783Z\",\"message\":{\"role\":\"assistant\",\"content\":\" I\"},\"done\":false}\n{\"model\":\"llama3:latest\",\"created_at\":\"2025-05-22T07:57:37.880345Z\",\"message\":{\"role\":\"assistant\",\"content\":\"'m\"},\"done\":false}\n{\"model\":\"llama3:latest\",\"created_at\":\"2025-05-22T07:57:37.908243Z\",\"message\":{\"role\":\"assistant\",\"content\":\" here\"},\"done\":false}\n{\"model\":\"llama3:latest\",\"created_at\":\"2025-05-22T07:57:37.937261Z\",\"message\":{\"role\":\"assistant\",\"content\":\" and\"},\"done\":false}\n{\"model\":\"llama3:latest\",\"created_at\":\"2025-05-22T07:57:37.967325Z\",\"message\":{\"role\":\"assistant\",\"content\":\" ready\"},\"done\":false}\n{\"model\":\"llama3:latest\",\"created_at\":\"2025-05-22T07:57:37.996367Z\",\"message\":{\"role\":\"assistant\",\"content\":\" to\"},\"done\":false}\n{\"model\":\"llama3:latest\",\"created_at\":\"2025-05-22T07:57:38.02522Z\",\"message\":{\"role\":\"assistant\",\"content\":\" assist\"},\"done\":false}\n{\"model\":\"llama3:latest\",\"created_at\":\"2025-05-22T07:57:38.05475Z\",\"message\":{\"role\":\"assistant\",\"content\":\" you\"},\"done\":false}\n{\"model\":\"llama3:latest\",\"created_at\":\"2025-05-22T07:57:38.084333Z\",\"message\":{\"role\":\"assistant\",\"content\":\" with\"},\"done\":false}\n{\"model\":\"llama3:latest\",\"created_at\":\"2025-05-22T07:57:38.113514Z\",\"message\":{\"role\":\"assistant\",\"content\":\" any\"},\"done\":false}\n{\"model\":\"llama3:latest\",\"created_at\":\"2025-05-22T07:57:38.142456Z\",\"message\":{\"role\":\"assistant\",\"content\":\" questions\"},\"done\":false}\n{\"model\":\"llama3:latest\",\"created_at\":\"2025-05-22T07:57:38.171743Z\",\"message\":{\"role\":\"assistant\",\"content\":\" or\"},\"done\":false}\n{\"model\":\"llama3:latest\",\"created_at\":\"2025-05-22T07:57:38.200941Z\",\"message\":{\"role\":\"assistant\",\"content\":\" topics\"},\"done\":false}\n{\"model\":\"llama3:latest\",\"created_at\":\"2025-05-22T07:57:38.230011Z\",\"message\":{\"role\":\"assistant\",\"content\":\" you\"},\"done\":false}\n{\"model\":\"llama3:latest\",\"created_at\":\"2025-05-22T07:57:38.25913Z\",\"message\":{\"role\":\"assistant\",\"content\":\"'d\"},\"done\":false}\n{\"model\":\"llama3:latest\",\"created_at\":\"2025-05-22T07:57:38.289317Z\",\"message\":{\"role\":\"assistant\",\"content\":\" like\"},\"done\":false}\n{\"model\":\"llama3:latest\",\"created_at\":\"2025-05-22T07:57:38.318546Z\",\"message\":{\"role\":\"assistant\",\"content\":\" to\"},\"done\":false}\n{\"model\":\"llama3:latest\",\"created_at\":\"2025-05-22T07:57:38.347245Z\",\"message\":{\"role\":\"assistant\",\"content\":\" discuss\"},\"done\":false}\n{\"model\":\"llama3:latest\",\"created_at\":\"2025-05-22T07:57:38.37708Z\",\"message\":{\"role\":\"assistant\",\"content\":\"!\"},\"done\":false}\n{\"model\":\"llama3:latest\",\"created_at\":\"2025-05-22T07:57:38.405906Z\",\"message\":{\"role\":\"assistant\",\"content\":\"\"},\"done_reason\":\"stop\",\"done\":true,\"total_duration\":6387758208,\"load_duration\":5236307708,\"prompt_eval_count\":12,\"prompt_eval_duration\":363200583,\"eval_count\":28,\"eval_duration\":787724917}\n"
10+
body: "{\"model\":\"llama3:latest\",\"created_at\":\"2025-06-06T13:23:19.810545Z\",\"message\":{\"role\":\"assistant\",\"content\":\"This\"},\"done\":false}\n{\"model\":\"llama3:latest\",\"created_at\":\"2025-06-06T13:23:19.828065Z\",\"message\":{\"role\":\"assistant\",\"content\":\" is\"},\"done\":false}\n{\"model\":\"llama3:latest\",\"created_at\":\"2025-06-06T13:23:19.919572Z\",\"message\":{\"role\":\"assistant\",\"content\":\" a\"},\"done\":false}\n{\"model\":\"llama3:latest\",\"created_at\":\"2025-06-06T13:23:19.949477Z\",\"message\":{\"role\":\"assistant\",\"content\":\" test\"},\"done\":false}\n{\"model\":\"llama3:latest\",\"created_at\":\"2025-06-06T13:23:19.978451Z\",\"message\":{\"role\":\"assistant\",\"content\":\" message\"},\"done\":false}\n{\"model\":\"llama3:latest\",\"created_at\":\"2025-06-06T13:23:20.008737Z\",\"message\":{\"role\":\"assistant\",\"content\":\".\"},\"done\":false}\n{\"model\":\"llama3:latest\",\"created_at\":\"2025-06-06T13:23:20.037548Z\",\"message\":{\"role\":\"assistant\",\"content\":\" I\"},\"done\":false}\n{\"model\":\"llama3:latest\",\"created_at\":\"2025-06-06T13:23:20.066217Z\",\"message\":{\"role\":\"assistant\",\"content\":\"'m\"},\"done\":false}\n{\"model\":\"llama3:latest\",\"created_at\":\"2025-06-06T13:23:20.096003Z\",\"message\":{\"role\":\"assistant\",\"content\":\" receiving\"},\"done\":false}\n{\"model\":\"llama3:latest\",\"created_at\":\"2025-06-06T13:23:20.124694Z\",\"message\":{\"role\":\"assistant\",\"content\":\" it\"},\"done\":false}\n{\"model\":\"llama3:latest\",\"created_at\":\"2025-06-06T13:23:20.153746Z\",\"message\":{\"role\":\"assistant\",\"content\":\" just\"},\"done\":false}\n{\"model\":\"llama3:latest\",\"created_at\":\"2025-06-06T13:23:20.182879Z\",\"message\":{\"role\":\"assistant\",\"content\":\" fine\"},\"done\":false}\n{\"model\":\"llama3:latest\",\"created_at\":\"2025-06-06T13:23:20.213119Z\",\"message\":{\"role\":\"assistant\",\"content\":\"!\"},\"done\":false}\n{\"model\":\"llama3:latest\",\"created_at\":\"2025-06-06T13:23:20.242749Z\",\"message\":{\"role\":\"assistant\",\"content\":\"\"},\"done_reason\":\"stop\",\"done\":true,\"total_duration\":3203965042,\"load_duration\":2493534250,\"prompt_eval_count\":12,\"prompt_eval_duration\":276055500,\"eval_count\":14,\"eval_duration\":433676750}\n"

0 commit comments

Comments
 (0)