1use anyhow::{anyhow, Context, Result};
2use futures::{io::BufReader, stream::BoxStream, AsyncBufReadExt, AsyncReadExt, Stream, StreamExt};
3use http_client::{AsyncBody, HttpClient, Method, Request as HttpRequest};
4use isahc::config::Configurable;
5use serde::{Deserialize, Serialize};
6use serde_json::{Map, Value};
7use std::{convert::TryFrom, future::Future, time::Duration};
8use strum::EnumIter;
9
10pub const OPEN_AI_API_URL: &str = "https://api.openai.com/v1";
11
12fn is_none_or_empty<T: AsRef<[U]>, U>(opt: &Option<T>) -> bool {
13 opt.as_ref().map_or(true, |v| v.as_ref().is_empty())
14}
15
16#[derive(Clone, Copy, Serialize, Deserialize, Debug, Eq, PartialEq)]
17#[serde(rename_all = "lowercase")]
18pub enum Role {
19 User,
20 Assistant,
21 System,
22 Tool,
23}
24
25impl TryFrom<String> for Role {
26 type Error = anyhow::Error;
27
28 fn try_from(value: String) -> Result<Self> {
29 match value.as_str() {
30 "user" => Ok(Self::User),
31 "assistant" => Ok(Self::Assistant),
32 "system" => Ok(Self::System),
33 "tool" => Ok(Self::Tool),
34 _ => Err(anyhow!("invalid role '{value}'")),
35 }
36 }
37}
38
39impl From<Role> for String {
40 fn from(val: Role) -> Self {
41 match val {
42 Role::User => "user".to_owned(),
43 Role::Assistant => "assistant".to_owned(),
44 Role::System => "system".to_owned(),
45 Role::Tool => "tool".to_owned(),
46 }
47 }
48}
49
50#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
51#[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq, EnumIter)]
52pub enum Model {
53 #[serde(rename = "gpt-3.5-turbo", alias = "gpt-3.5-turbo-0613")]
54 ThreePointFiveTurbo,
55 #[serde(rename = "gpt-4", alias = "gpt-4-0613")]
56 Four,
57 #[serde(rename = "gpt-4-turbo-preview", alias = "gpt-4-1106-preview")]
58 FourTurbo,
59 #[serde(rename = "gpt-4o", alias = "gpt-4o-2024-05-13")]
60 #[default]
61 FourOmni,
62 #[serde(rename = "gpt-4o-mini", alias = "gpt-4o-mini-2024-07-18")]
63 FourOmniMini,
64 #[serde(rename = "custom")]
65 Custom { name: String, max_tokens: usize },
66}
67
68impl Model {
69 pub fn from_id(id: &str) -> Result<Self> {
70 match id {
71 "gpt-3.5-turbo" => Ok(Self::ThreePointFiveTurbo),
72 "gpt-4" => Ok(Self::Four),
73 "gpt-4-turbo-preview" => Ok(Self::FourTurbo),
74 "gpt-4o" => Ok(Self::FourOmni),
75 "gpt-4o-mini" => Ok(Self::FourOmniMini),
76 _ => Err(anyhow!("invalid model id")),
77 }
78 }
79
80 pub fn id(&self) -> &str {
81 match self {
82 Self::ThreePointFiveTurbo => "gpt-3.5-turbo",
83 Self::Four => "gpt-4",
84 Self::FourTurbo => "gpt-4-turbo-preview",
85 Self::FourOmni => "gpt-4o",
86 Self::FourOmniMini => "gpt-4o-mini",
87 Self::Custom { name, .. } => name,
88 }
89 }
90
91 pub fn display_name(&self) -> &str {
92 match self {
93 Self::ThreePointFiveTurbo => "gpt-3.5-turbo",
94 Self::Four => "gpt-4",
95 Self::FourTurbo => "gpt-4-turbo",
96 Self::FourOmni => "gpt-4o",
97 Self::FourOmniMini => "gpt-4o-mini",
98 Self::Custom { name, .. } => name,
99 }
100 }
101
102 pub fn max_token_count(&self) -> usize {
103 match self {
104 Self::ThreePointFiveTurbo => 4096,
105 Self::Four => 8192,
106 Self::FourTurbo => 128000,
107 Self::FourOmni => 128000,
108 Self::FourOmniMini => 128000,
109 Self::Custom { max_tokens, .. } => *max_tokens,
110 }
111 }
112}
113
114#[derive(Debug, Serialize, Deserialize)]
115pub struct Request {
116 pub model: String,
117 pub messages: Vec<RequestMessage>,
118 pub stream: bool,
119 pub stop: Vec<String>,
120 pub temperature: f32,
121 #[serde(default, skip_serializing_if = "Option::is_none")]
122 pub tool_choice: Option<String>,
123 #[serde(default, skip_serializing_if = "Vec::is_empty")]
124 pub tools: Vec<ToolDefinition>,
125}
126
127#[derive(Debug, Deserialize, Serialize)]
128pub struct FunctionDefinition {
129 pub name: String,
130 pub description: Option<String>,
131 pub parameters: Option<Map<String, Value>>,
132}
133
134#[derive(Deserialize, Serialize, Debug)]
135#[serde(tag = "type", rename_all = "snake_case")]
136pub enum ToolDefinition {
137 #[allow(dead_code)]
138 Function { function: FunctionDefinition },
139}
140
141#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]
142#[serde(tag = "role", rename_all = "lowercase")]
143pub enum RequestMessage {
144 Assistant {
145 content: Option<String>,
146 #[serde(default, skip_serializing_if = "Vec::is_empty")]
147 tool_calls: Vec<ToolCall>,
148 },
149 User {
150 content: String,
151 },
152 System {
153 content: String,
154 },
155 Tool {
156 content: String,
157 tool_call_id: String,
158 },
159}
160
161#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]
162pub struct ToolCall {
163 pub id: String,
164 #[serde(flatten)]
165 pub content: ToolCallContent,
166}
167
168#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]
169#[serde(tag = "type", rename_all = "lowercase")]
170pub enum ToolCallContent {
171 Function { function: FunctionContent },
172}
173
174#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]
175pub struct FunctionContent {
176 pub name: String,
177 pub arguments: String,
178}
179
180#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]
181pub struct ResponseMessageDelta {
182 pub role: Option<Role>,
183 pub content: Option<String>,
184 #[serde(default, skip_serializing_if = "is_none_or_empty")]
185 pub tool_calls: Option<Vec<ToolCallChunk>>,
186}
187
188#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]
189pub struct ToolCallChunk {
190 pub index: usize,
191 pub id: Option<String>,
192
193 // There is also an optional `type` field that would determine if a
194 // function is there. Sometimes this streams in with the `function` before
195 // it streams in the `type`
196 pub function: Option<FunctionChunk>,
197}
198
199#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]
200pub struct FunctionChunk {
201 pub name: Option<String>,
202 pub arguments: Option<String>,
203}
204
205#[derive(Serialize, Deserialize, Debug)]
206pub struct Usage {
207 pub prompt_tokens: u32,
208 pub completion_tokens: u32,
209 pub total_tokens: u32,
210}
211
212#[derive(Serialize, Deserialize, Debug)]
213pub struct ChoiceDelta {
214 pub index: u32,
215 pub delta: ResponseMessageDelta,
216 pub finish_reason: Option<String>,
217}
218
219#[derive(Serialize, Deserialize, Debug)]
220pub struct ResponseStreamEvent {
221 pub created: u32,
222 pub model: String,
223 pub choices: Vec<ChoiceDelta>,
224 pub usage: Option<Usage>,
225}
226
227pub async fn stream_completion(
228 client: &dyn HttpClient,
229 api_url: &str,
230 api_key: &str,
231 request: Request,
232 low_speed_timeout: Option<Duration>,
233) -> Result<BoxStream<'static, Result<ResponseStreamEvent>>> {
234 let uri = format!("{api_url}/chat/completions");
235 let mut request_builder = HttpRequest::builder()
236 .method(Method::POST)
237 .uri(uri)
238 .header("Content-Type", "application/json")
239 .header("Authorization", format!("Bearer {}", api_key));
240
241 if let Some(low_speed_timeout) = low_speed_timeout {
242 request_builder = request_builder.low_speed_timeout(100, low_speed_timeout);
243 };
244
245 let request = request_builder.body(AsyncBody::from(serde_json::to_string(&request)?))?;
246 let mut response = client.send(request).await?;
247 if response.status().is_success() {
248 let reader = BufReader::new(response.into_body());
249 Ok(reader
250 .lines()
251 .filter_map(|line| async move {
252 match line {
253 Ok(line) => {
254 let line = line.strip_prefix("data: ")?;
255 if line == "[DONE]" {
256 None
257 } else {
258 match serde_json::from_str(line) {
259 Ok(response) => Some(Ok(response)),
260 Err(error) => Some(Err(anyhow!(error))),
261 }
262 }
263 }
264 Err(error) => Some(Err(anyhow!(error))),
265 }
266 })
267 .boxed())
268 } else {
269 let mut body = String::new();
270 response.body_mut().read_to_string(&mut body).await?;
271
272 #[derive(Deserialize)]
273 struct OpenAiResponse {
274 error: OpenAiError,
275 }
276
277 #[derive(Deserialize)]
278 struct OpenAiError {
279 message: String,
280 }
281
282 match serde_json::from_str::<OpenAiResponse>(&body) {
283 Ok(response) if !response.error.message.is_empty() => Err(anyhow!(
284 "Failed to connect to OpenAI API: {}",
285 response.error.message,
286 )),
287
288 _ => Err(anyhow!(
289 "Failed to connect to OpenAI API: {} {}",
290 response.status(),
291 body,
292 )),
293 }
294 }
295}
296
297#[derive(Copy, Clone, Serialize, Deserialize)]
298pub enum OpenAiEmbeddingModel {
299 #[serde(rename = "text-embedding-3-small")]
300 TextEmbedding3Small,
301 #[serde(rename = "text-embedding-3-large")]
302 TextEmbedding3Large,
303}
304
305#[derive(Serialize)]
306struct OpenAiEmbeddingRequest<'a> {
307 model: OpenAiEmbeddingModel,
308 input: Vec<&'a str>,
309}
310
311#[derive(Deserialize)]
312pub struct OpenAiEmbeddingResponse {
313 pub data: Vec<OpenAiEmbedding>,
314}
315
316#[derive(Deserialize)]
317pub struct OpenAiEmbedding {
318 pub embedding: Vec<f32>,
319}
320
321pub fn embed<'a>(
322 client: &dyn HttpClient,
323 api_url: &str,
324 api_key: &str,
325 model: OpenAiEmbeddingModel,
326 texts: impl IntoIterator<Item = &'a str>,
327) -> impl 'static + Future<Output = Result<OpenAiEmbeddingResponse>> {
328 let uri = format!("{api_url}/embeddings");
329
330 let request = OpenAiEmbeddingRequest {
331 model,
332 input: texts.into_iter().collect(),
333 };
334 let body = AsyncBody::from(serde_json::to_string(&request).unwrap());
335 let request = HttpRequest::builder()
336 .method(Method::POST)
337 .uri(uri)
338 .header("Content-Type", "application/json")
339 .header("Authorization", format!("Bearer {}", api_key))
340 .body(body)
341 .map(|request| client.send(request));
342
343 async move {
344 let mut response = request?.await?;
345 let mut body = String::new();
346 response.body_mut().read_to_string(&mut body).await?;
347
348 if response.status().is_success() {
349 let response: OpenAiEmbeddingResponse =
350 serde_json::from_str(&body).context("failed to parse OpenAI embedding response")?;
351 Ok(response)
352 } else {
353 Err(anyhow!(
354 "error during embedding, status: {:?}, body: {:?}",
355 response.status(),
356 body
357 ))
358 }
359 }
360}
361
362pub fn extract_text_from_events(
363 response: impl Stream<Item = Result<ResponseStreamEvent>>,
364) -> impl Stream<Item = Result<String>> {
365 response.filter_map(|response| async move {
366 match response {
367 Ok(mut response) => Some(Ok(response.choices.pop()?.delta.content?)),
368 Err(error) => Some(Err(error)),
369 }
370 })
371}