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
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
//! Streaming chat completion.
use serde_json::Value;
use tokio_util::sync::CancellationToken;
use uuid::Uuid;
use crate::clients::parsing::parser_for_transport;
use crate::core::errors::ConduitError;
use crate::core::results::AsyncTextStream;
use super::{LLM, build_messages, prepend_tape_history};
impl LLM {
/// Stream chat completion as an async `TextStream`.
pub async fn stream(
&mut self,
req: super::ChatRequest<'_>,
) -> Result<AsyncTextStream, ConduitError> {
let super::ChatRequest {
prompt,
user_content,
system_prompt,
model,
provider,
messages,
max_tokens,
tape,
cancellation,
..
} = req;
let tape_messages = match tape {
Some(tape_name) => self.build_tape_messages(tape_name, None).await,
None => Vec::new(),
};
let mut msgs = build_messages(
prompt,
user_content.as_deref(),
system_prompt,
messages.as_deref(),
);
prepend_tape_history(&mut msgs, tape_messages);
if let Some(tape_name) = tape {
let new_messages: Vec<Value> = msgs
.iter()
.filter(|m| m.get("role").and_then(|r| r.as_str()) == Some("user"))
.cloned()
.collect();
let run_id = Uuid::new_v4().to_string();
if let Err(e) = self
.async_tape
.record_chat(
tape_name,
&run_id,
system_prompt,
None,
&new_messages,
None,
None,
None,
None,
None,
Some(self.core.provider()),
Some(self.core.model()),
)
.await
{
tracing::error!(error = %e, tape = %tape_name, "failed to record streaming chat context");
}
}
let (response, transport, _prov, _model) = self
.core
.run_chat_stream(
msgs,
None,
model,
provider,
max_tokens,
None,
Default::default(),
)
.await?;
let parser = parser_for_transport(transport);
let (tx, rx) = tokio::sync::mpsc::channel::<String>(64);
tokio::spawn(Self::stream_sse_loop(response, parser, tx, cancellation));
let stream = tokio_stream::wrappers::ReceiverStream::new(rx);
Ok(AsyncTextStream::new(stream, None))
}
/// Consume SSE bytes from the response, parse text chunks, and forward
/// them through `tx`. Respects an optional `CancellationToken` — when
/// cancelled the loop stops and the channel closes, delivering whatever
/// partial content was already sent.
async fn stream_sse_loop(
response: reqwest::Response,
parser: &'static dyn crate::clients::parsing::types::BaseTransportParser,
tx: tokio::sync::mpsc::Sender<String>,
cancellation: Option<CancellationToken>,
) {
use futures::StreamExt;
let mut byte_stream = response.bytes_stream();
let mut buffer: Vec<u8> = Vec::new();
loop {
// Obtain the next chunk, racing against cancellation when a
// token was provided.
let chunk_result = match cancellation {
Some(ref token) => {
tokio::select! {
biased;
_ = token.cancelled() => {
tracing::info!("SSE stream cancelled");
break;
}
chunk = byte_stream.next() => chunk,
}
}
None => byte_stream.next().await,
};
let Some(chunk_result) = chunk_result else {
break; // stream finished
};
let bytes = match chunk_result {
Ok(b) => b,
Err(_) => break,
};
buffer.extend_from_slice(&bytes);
// Parse complete SSE lines from the byte buffer, leaving
// partial lines (which may contain incomplete multibyte
// UTF-8 sequences) for the next chunk.
let mut cursor = 0;
while let Some(rel) = buffer[cursor..].iter().position(|&b| b == b'\n') {
let line_end = cursor + rel;
let mut end = line_end;
if end > cursor && buffer[end - 1] == b'\r' {
end -= 1;
}
let line = String::from_utf8_lossy(&buffer[cursor..end]);
cursor = line_end + 1;
if let Some(data) = line.strip_prefix("data: ") {
if data == "[DONE]" {
break;
}
if let Ok(val) = serde_json::from_str::<Value>(data) {
let content = parser.extract_chunk_text(&val);
if !content.is_empty() && tx.send(content).await.is_err() {
return;
}
}
}
}
// Remove consumed bytes in one operation instead of per-line.
if cursor > 0 {
buffer.drain(..cursor);
}
}
}
}