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