1use anyhow::{anyhow, Context, Result};
2use futures::{io::BufReader, stream::BoxStream, AsyncBufReadExt, AsyncReadExt, 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
114fn serialize_model<S>(model: &Model, serializer: S) -> Result<S::Ok, S::Error>
115where
116 S: serde::Serializer,
117{
118 match model {
119 Model::Custom { name, .. } => serializer.serialize_str(name),
120 _ => serializer.serialize_str(model.id()),
121 }
122}
123
124#[derive(Debug, Serialize)]
125pub struct Request {
126 #[serde(serialize_with = "serialize_model")]
127 pub model: Model,
128 pub messages: Vec<RequestMessage>,
129 pub stream: bool,
130 pub stop: Vec<String>,
131 pub temperature: f32,
132 #[serde(skip_serializing_if = "Option::is_none")]
133 pub tool_choice: Option<String>,
134 #[serde(skip_serializing_if = "Vec::is_empty")]
135 pub tools: Vec<ToolDefinition>,
136}
137
138#[derive(Debug, Serialize)]
139pub struct FunctionDefinition {
140 pub name: String,
141 pub description: Option<String>,
142 pub parameters: Option<Map<String, Value>>,
143}
144
145#[derive(Serialize, Debug)]
146#[serde(tag = "type", rename_all = "snake_case")]
147pub enum ToolDefinition {
148 #[allow(dead_code)]
149 Function { function: FunctionDefinition },
150}
151
152#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]
153#[serde(tag = "role", rename_all = "lowercase")]
154pub enum RequestMessage {
155 Assistant {
156 content: Option<String>,
157 #[serde(default, skip_serializing_if = "Vec::is_empty")]
158 tool_calls: Vec<ToolCall>,
159 },
160 User {
161 content: String,
162 },
163 System {
164 content: String,
165 },
166 Tool {
167 content: String,
168 tool_call_id: String,
169 },
170}
171
172#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]
173pub struct ToolCall {
174 pub id: String,
175 #[serde(flatten)]
176 pub content: ToolCallContent,
177}
178
179#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]
180#[serde(tag = "type", rename_all = "lowercase")]
181pub enum ToolCallContent {
182 Function { function: FunctionContent },
183}
184
185#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]
186pub struct FunctionContent {
187 pub name: String,
188 pub arguments: String,
189}
190
191#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]
192pub struct ResponseMessageDelta {
193 pub role: Option<Role>,
194 pub content: Option<String>,
195 #[serde(default, skip_serializing_if = "is_none_or_empty")]
196 pub tool_calls: Option<Vec<ToolCallChunk>>,
197}
198
199#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]
200pub struct ToolCallChunk {
201 pub index: usize,
202 pub id: Option<String>,
203
204 // There is also an optional `type` field that would determine if a
205 // function is there. Sometimes this streams in with the `function` before
206 // it streams in the `type`
207 pub function: Option<FunctionChunk>,
208}
209
210#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]
211pub struct FunctionChunk {
212 pub name: Option<String>,
213 pub arguments: Option<String>,
214}
215
216#[derive(Deserialize, Debug)]
217pub struct Usage {
218 pub prompt_tokens: u32,
219 pub completion_tokens: u32,
220 pub total_tokens: u32,
221}
222
223#[derive(Deserialize, Debug)]
224pub struct ChoiceDelta {
225 pub index: u32,
226 pub delta: ResponseMessageDelta,
227 pub finish_reason: Option<String>,
228}
229
230#[derive(Deserialize, Debug)]
231pub struct ResponseStreamEvent {
232 pub created: u32,
233 pub model: String,
234 pub choices: Vec<ChoiceDelta>,
235 pub usage: Option<Usage>,
236}
237
238pub async fn stream_completion(
239 client: &dyn HttpClient,
240 api_url: &str,
241 api_key: &str,
242 request: Request,
243 low_speed_timeout: Option<Duration>,
244) -> Result<BoxStream<'static, Result<ResponseStreamEvent>>> {
245 let uri = format!("{api_url}/chat/completions");
246 let mut request_builder = HttpRequest::builder()
247 .method(Method::POST)
248 .uri(uri)
249 .header("Content-Type", "application/json")
250 .header("Authorization", format!("Bearer {}", api_key));
251
252 if let Some(low_speed_timeout) = low_speed_timeout {
253 request_builder = request_builder.low_speed_timeout(100, low_speed_timeout);
254 };
255
256 let request = request_builder.body(AsyncBody::from(serde_json::to_string(&request)?))?;
257 let mut response = client.send(request).await?;
258 if response.status().is_success() {
259 let reader = BufReader::new(response.into_body());
260 Ok(reader
261 .lines()
262 .filter_map(|line| async move {
263 match line {
264 Ok(line) => {
265 let line = line.strip_prefix("data: ")?;
266 if line == "[DONE]" {
267 None
268 } else {
269 match serde_json::from_str(line) {
270 Ok(response) => Some(Ok(response)),
271 Err(error) => Some(Err(anyhow!(error))),
272 }
273 }
274 }
275 Err(error) => Some(Err(anyhow!(error))),
276 }
277 })
278 .boxed())
279 } else {
280 let mut body = String::new();
281 response.body_mut().read_to_string(&mut body).await?;
282
283 #[derive(Deserialize)]
284 struct OpenAiResponse {
285 error: OpenAiError,
286 }
287
288 #[derive(Deserialize)]
289 struct OpenAiError {
290 message: String,
291 }
292
293 match serde_json::from_str::<OpenAiResponse>(&body) {
294 Ok(response) if !response.error.message.is_empty() => Err(anyhow!(
295 "Failed to connect to OpenAI API: {}",
296 response.error.message,
297 )),
298
299 _ => Err(anyhow!(
300 "Failed to connect to OpenAI API: {} {}",
301 response.status(),
302 body,
303 )),
304 }
305 }
306}
307
308#[derive(Copy, Clone, Serialize, Deserialize)]
309pub enum OpenAiEmbeddingModel {
310 #[serde(rename = "text-embedding-3-small")]
311 TextEmbedding3Small,
312 #[serde(rename = "text-embedding-3-large")]
313 TextEmbedding3Large,
314}
315
316#[derive(Serialize)]
317struct OpenAiEmbeddingRequest<'a> {
318 model: OpenAiEmbeddingModel,
319 input: Vec<&'a str>,
320}
321
322#[derive(Deserialize)]
323pub struct OpenAiEmbeddingResponse {
324 pub data: Vec<OpenAiEmbedding>,
325}
326
327#[derive(Deserialize)]
328pub struct OpenAiEmbedding {
329 pub embedding: Vec<f32>,
330}
331
332pub fn embed<'a>(
333 client: &dyn HttpClient,
334 api_url: &str,
335 api_key: &str,
336 model: OpenAiEmbeddingModel,
337 texts: impl IntoIterator<Item = &'a str>,
338) -> impl 'static + Future<Output = Result<OpenAiEmbeddingResponse>> {
339 let uri = format!("{api_url}/embeddings");
340
341 let request = OpenAiEmbeddingRequest {
342 model,
343 input: texts.into_iter().collect(),
344 };
345 let body = AsyncBody::from(serde_json::to_string(&request).unwrap());
346 let request = HttpRequest::builder()
347 .method(Method::POST)
348 .uri(uri)
349 .header("Content-Type", "application/json")
350 .header("Authorization", format!("Bearer {}", api_key))
351 .body(body)
352 .map(|request| client.send(request));
353
354 async move {
355 let mut response = request?.await?;
356 let mut body = String::new();
357 response.body_mut().read_to_string(&mut body).await?;
358
359 if response.status().is_success() {
360 let response: OpenAiEmbeddingResponse =
361 serde_json::from_str(&body).context("failed to parse OpenAI embedding response")?;
362 Ok(response)
363 } else {
364 Err(anyhow!(
365 "error during embedding, status: {:?}, body: {:?}",
366 response.status(),
367 body
368 ))
369 }
370 }
371}