agentic_core/executor/
inference.rs1use std::sync::Arc;
7use std::time::Duration;
8
9use async_stream::stream;
10use futures::{Stream, StreamExt};
11
12use crate::executor::error::{ExecutorError, ExecutorResult};
13use crate::proxy::processed_response_headers;
14
15pub type BoxStream = std::pin::Pin<Box<dyn Stream<Item = String> + Send>>;
17
18pub(super) const DONE_MARKER: &str = "data: [DONE]\n\n";
20
21pub(super) async fn next_chunk<S>(stream: &mut S, timeout: Duration) -> ExecutorResult<Option<bytes::Bytes>>
26where
27 S: futures::Stream<Item = Result<bytes::Bytes, reqwest::Error>> + Unpin,
28{
29 let item = if timeout.is_zero() {
30 stream.next().await
31 } else {
32 tokio::time::timeout(timeout, stream.next()).await.map_err(|_| {
33 ExecutorError::StreamError("chunk timeout: no data received within the configured window".into())
34 })?
35 };
36 item.transpose().map_err(ExecutorError::NetworkError)
37}
38
39fn drain_complete_utf8_lines(buffer: &mut Vec<u8>) -> Vec<String> {
40 let mut lines = Vec::new();
41 while let Some(pos) = buffer.iter().position(|byte| *byte == b'\n') {
42 let line = buffer.drain(..=pos).collect::<Vec<_>>();
43 let line_end = if pos > 0 && line.get(pos - 1) == Some(&b'\r') {
44 pos - 1
45 } else {
46 pos
47 };
48 if let Ok(line) = std::str::from_utf8(&line[..line_end]) {
49 lines.push(line.to_string());
50 }
51 }
52 lines
53}
54
55pub(super) async fn send_request(
62 client: &reqwest::Client,
63 url: &str,
64 body: String,
65 auth: Option<&str>,
66 forwarded_headers: Option<&reqwest::header::HeaderMap>,
67) -> ExecutorResult<reqwest::Response> {
68 let mut headers = forwarded_headers.cloned().unwrap_or_default();
69 headers
70 .entry(reqwest::header::CONTENT_TYPE)
71 .or_insert(reqwest::header::HeaderValue::from_static("application/json"));
72 let mut req = client.post(url).headers(headers).body(body);
73 if let Some(key) = auth {
74 req = req.bearer_auth(key);
75 }
76
77 let resp = req.send().await.map_err(|e| ExecutorError::LLMTransport {
78 status: if e.is_timeout() {
79 http::StatusCode::GATEWAY_TIMEOUT
80 } else {
81 http::StatusCode::BAD_GATEWAY
82 },
83 message: if e.is_timeout() {
84 "LLM timeout"
85 } else {
86 "LLM unavailable"
87 },
88 })?;
89
90 if !resp.status().is_success() {
91 let status = resp.status().as_u16();
92 let headers = processed_response_headers(resp.headers());
93 let body = resp
96 .text()
97 .await
98 .inspect_err(|e| tracing::debug!("failed to read error response body: {e}"))
99 .unwrap_or_default();
100 return Err(ExecutorError::LLMRequest {
101 status: http::StatusCode::from_u16(status).unwrap_or(http::StatusCode::INTERNAL_SERVER_ERROR),
102 body,
103 headers,
104 });
105 }
106
107 Ok(resp)
108}
109
110pub(super) async fn fetch_response_json(
114 upstream_json: String,
115 url: &str,
116 client: &reqwest::Client,
117 auth: Option<&str>,
118) -> ExecutorResult<String> {
119 let resp = send_request(client, url, upstream_json, auth, None).await?;
120 resp.text().await.map_err(ExecutorError::NetworkError)
122}
123
124pub(super) async fn fetch_response_json_with_headers(
126 upstream_json: String,
127 url: &str,
128 client: &reqwest::Client,
129 headers: &reqwest::header::HeaderMap,
130) -> ExecutorResult<(String, http::HeaderMap)> {
131 let resp = send_request(client, url, upstream_json, None, Some(headers)).await?;
132 let response_headers = processed_response_headers(resp.headers());
133 let body = resp.text().await.map_err(ExecutorError::NetworkError)?;
134 Ok((body, response_headers))
135}
136
137pub fn call_inference(
147 upstream_json: String,
148 url: String,
149 client: Arc<reqwest::Client>,
150 auth: Option<String>,
151 chunk_timeout: Duration,
152) -> impl Stream<Item = Result<String, ExecutorError>> + Send + 'static {
153 stream! {
154 let resp = match send_request(&client, &url, upstream_json, auth.as_deref(), None).await {
155 Ok(r) => r,
156 Err(e) => { yield Err(e); return; }
157 };
158
159 let mut lines = Box::pin(response_lines(resp, chunk_timeout));
160 while let Some(line) = lines.next().await {
161 yield line;
162 }
163 }
164}
165
166pub(super) fn response_lines(
168 resp: reqwest::Response,
169 chunk_timeout: Duration,
170) -> impl Stream<Item = Result<String, ExecutorError>> + Send + 'static {
171 stream! {
172 let mut bytes = resp.bytes_stream();
173 let mut buf = Vec::with_capacity(8192);
174
175 loop {
176 let chunk = match next_chunk(&mut bytes, chunk_timeout).await {
177 Ok(Some(c)) => c,
178 Ok(None) => break,
179 Err(e) => { yield Err(e); return; }
180 };
181
182 buf.extend_from_slice(&chunk);
183
184 for line in drain_complete_utf8_lines(&mut buf) {
185 match line.as_str() {
186 "data: [DONE]" => return,
187 l if l.starts_with("data: ") => yield Ok(line),
188 _ => {}
189 }
190 }
191 }
192 }
193}
194
195#[cfg(test)]
196mod tests {
197 use super::*;
198
199 #[test]
200 fn utf8_line_reader_preserves_split_multibyte_characters() {
201 let snowman = "\u{2603}";
202 let line = format!(r#"data: {{"delta":"snow {snowman}"}}"#);
203 let bytes = format!("{line}\n").into_bytes();
204 let split_at = bytes
205 .windows(snowman.len())
206 .position(|window| window == snowman.as_bytes())
207 .expect("snowman bytes present")
208 + 1;
209 let mut buffer = bytes[..split_at].to_vec();
210
211 assert!(drain_complete_utf8_lines(&mut buffer).is_empty());
212
213 buffer.extend_from_slice(&bytes[split_at..]);
214 let lines = drain_complete_utf8_lines(&mut buffer);
215
216 assert!(buffer.is_empty());
217 assert_eq!(lines, vec![line]);
218 assert!(!lines[0].contains('\u{FFFD}'));
219 }
220}