client.rs

crates/cortex-core/src/provider/client.rs

1. //! One client for every OpenAI-compatible server: DeepInfra and local servers
2. //! differ only in base URL, key and pricing.
4. use std::time::{Duration, Instant};
6. use futures_util::StreamExt;
7. use tokio_util::sync::CancellationToken;
9. use super::error::ProviderError;
10. use super::openai::{parse_chunk, ChatRequest};
11. use super::sse::SseParser;
12. use crate::types::Usage;
14. pub const DEEPINFRA_BASE: &str = "https://api.deepinfra.com/v1/openai";
16. /// Something that arrived on the stream, in order.
17. #[derive(Debug, Clone, PartialEq)]
18. pub enum StreamItem {
19. Content(String),
20. Reasoning(String),
21. Usage(Usage),
22. }
24. /// How a stream ended. `content` and `reasoning` hold whatever arrived, even
25. /// when it was cancelled or failed part way, so a partial reply can be kept.
26. #[derive(Debug, Clone, Default, PartialEq)]
27. pub struct StreamOutcome {
28. pub content: String,
29. pub reasoning: String,
30. /// The last usage block seen, if the server sent one.
31. pub usage: Option<Usage>,
32. pub provider_cost_usd: Option<f64>,
33. pub finish_reason: Option<String>,
34. pub cancelled: bool,
35. pub error: Option<ProviderError>,
36. pub ttft_ms: Option<u64>,
37. pub duration_ms: u64,
38. }
40. #[derive(Clone)]
41. pub struct Client {
42. http: reqwest::Client,
43. base_url: String,
44. api_key: Option<String>,
45. }
47. fn is_loopback(url: &str) -> bool {
48. let rest = url.split("://").nth(1).unwrap_or(url);
49. let host = rest.split(['/', ':']).next().unwrap_or("");
50. matches!(host, "localhost" | "127.0.0.1" | "[::1]" | "::1") || rest.starts_with("[::1]")
51. }
53. impl Client {
54. pub fn new(base_url: &str, api_key: Option<String>) -> Self {
55. let mut b = reqwest::Client::builder()
56. .connect_timeout(Duration::from_secs(10))
57. .user_agent(concat!("Cortex/", env!("CARGO_PKG_VERSION")));
58. if is_loopback(base_url) {
59. b = b.no_proxy(); // a local model server is never behind the system proxy
60. }
61. Self {
62. http: b.build().expect("http client"),
63. base_url: base_url.trim_end_matches('/').to_string(),
64. api_key: api_key.filter(|k| !k.is_empty()),
65. }
66. }
68. pub fn base_url(&self) -> &str {
69. &self.base_url
70. }
72. fn auth(&self, rb: reqwest::RequestBuilder) -> reqwest::RequestBuilder {
73. match &self.api_key {
74. Some(k) => rb.bearer_auth(k),
75. None => rb,
76. }
77. }
79. fn transport_error(&self, e: reqwest::Error) -> ProviderError {
80. if e.is_connect() || e.is_timeout() {
81. ProviderError::Unreachable { url: self.base_url.clone() }
82. } else {
83. ProviderError::Other(e.to_string())
84. }
85. }
87. async fn check(&self, res: reqwest::Response) -> Result<reqwest::Response, ProviderError> {
88. if res.status().is_success() {
89. return Ok(res);
90. }
91. let status = res.status().as_u16();
92. let retry_after =
93. res.headers().get("retry-after").and_then(|v| v.to_str().ok()).and_then(|v| v.trim().parse::<u64>().ok());
94. let body = res.text().await.unwrap_or_default();
95. Err(ProviderError::from_status(status, &body, retry_after))
96. }
98. /// `GET {base}/models`: the raw JSON, for `pricing::parse_priced_models` or
99. /// `pricing::parse_local_models`.
100. pub async fn list_models(&self) -> Result<serde_json::Value, ProviderError> {
101. let res = self
102. .auth(self.http.get(format!("{}/models", self.base_url)))
103. .timeout(Duration::from_secs(20))
104. .send()
105. .await
106. .map_err(|e| self.transport_error(e))?;
107. let res = self.check(res).await?;
108. res.json().await.map_err(|e| ProviderError::Other(format!("unreadable model list: {e}")))
109. }
111. /// Stream a chat completion. `on` sees every item as it arrives. Cancelling
112. /// `cancel` drops the connection at once, which stops the provider billing.
113. pub async fn stream_chat(
114. &self,
115. req: &ChatRequest,
116. cancel: &CancellationToken,
117. mut on: impl FnMut(StreamItem),
118. ) -> StreamOutcome {
119. let started = Instant::now();
120. let mut out = StreamOutcome::default();
121. let send = self.auth(self.http.post(format!("{}/chat/completions", self.base_url))).json(req).send();
122. let res = tokio::select! {
123. _ = cancel.cancelled() => { out.cancelled = true; out.duration_ms = ms(started); return out; }
124. r = send => r,
125. };
126. let res = match res {
127. Ok(r) => match self.check(r).await {
128. Ok(r) => r,
129. Err(e) => {
130. out.error = Some(e);
131. out.duration_ms = ms(started);
132. return out;
133. }
134. },
135. Err(e) => {
136. out.error = Some(self.transport_error(e));
137. out.duration_ms = ms(started);
138. return out;
139. }
140. };
142. let mut body = res.bytes_stream();
143. let mut parser = SseParser::new();
144. 'read: loop {
145. let next = tokio::select! {
146. _ = cancel.cancelled() => { out.cancelled = true; break 'read; }
147. n = body.next() => n,
148. };
149. let (events, ended) = match next {
150. Some(Ok(bytes)) => (parser.push(&bytes), false),
151. Some(Err(e)) => {
152. out.error = Some(ProviderError::Stream(e.to_string()));
153. break;
154. }
155. None => (parser.finish(), true),
156. };
157. for ev in events {
158. match parse_chunk(&ev.data) {
159. Ok(c) => {
160. if c.done {
161. break 'read;
162. }
163. if let Some(r) = c.reasoning {
164. out.ttft_ms.get_or_insert(ms(started));
165. out.reasoning.push_str(&r);
166. on(StreamItem::Reasoning(r));
167. }
168. if let Some(t) = c.content {
169. out.ttft_ms.get_or_insert(ms(started));
170. out.content.push_str(&t);
171. on(StreamItem::Content(t));
172. }
173. if let Some(u) = c.usage {
174. out.usage = Some(u);
175. on(StreamItem::Usage(u));
176. }
177. if c.provider_cost_usd.is_some() {
178. out.provider_cost_usd = c.provider_cost_usd;
179. }
180. if c.finish_reason.is_some() {
181. out.finish_reason = c.finish_reason;
182. }
183. }
184. Err(e) => {
185. out.error = Some(e);
186. break 'read;
187. }
188. }
189. }
190. if ended {
191. break;
192. }
193. }
194. drop(body);
195. out.duration_ms = ms(started);
196. out
197. }
198. }
200. fn ms(since: Instant) -> u64 {
201. since.elapsed().as_millis() as u64
202. }
204. #[cfg(test)]
205. mod tests {
206. use super::*;
208. #[test]
209. fn loopback_detection() {
210. assert!(is_loopback("http://localhost:11434/v1"));
211. assert!(is_loopback("http://127.0.0.1:1234/v1"));
212. assert!(is_loopback("http://[::1]:8080/v1"));
213. assert!(!is_loopback(DEEPINFRA_BASE));
214. assert!(!is_loopback("http://localhost.example.com/v1"));
215. }
216. }