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::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 #[serde(default, skip_serializing_if = "Option::is_none")]
120 pub max_tokens: Option<usize>,
121 pub stop: Vec<String>,
122 pub temperature: f32,
123 #[serde(default, skip_serializing_if = "Option::is_none")]
124 pub tool_choice: Option<ToolChoice>,
125 #[serde(default, skip_serializing_if = "Vec::is_empty")]
126 pub tools: Vec<ToolDefinition>,
127}
128
129#[derive(Debug, Serialize, Deserialize)]
130#[serde(untagged)]
131pub enum ToolChoice {
132 Auto,
133 Required,
134 None,
135 Other(ToolDefinition),
136}
137
138#[derive(Clone, Deserialize, Serialize, Debug)]
139#[serde(tag = "type", rename_all = "snake_case")]
140pub enum ToolDefinition {
141 #[allow(dead_code)]
142 Function { function: FunctionDefinition },
143}
144
145#[derive(Clone, Debug, Serialize, Deserialize)]
146pub struct FunctionDefinition {
147 pub name: String,
148 pub description: Option<String>,
149 pub parameters: Option<Value>,
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(Serialize, Deserialize, Debug)]
217pub struct Usage {
218 pub prompt_tokens: u32,
219 pub completion_tokens: u32,
220 pub total_tokens: u32,
221}
222
223#[derive(Serialize, Deserialize, Debug)]
224pub struct ChoiceDelta {
225 pub index: u32,
226 pub delta: ResponseMessageDelta,
227 pub finish_reason: Option<String>,
228}
229
230#[derive(Serialize, Deserialize, Debug)]
231#[serde(untagged)]
232pub enum ResponseStreamResult {
233 Ok(ResponseStreamEvent),
234 Err { error: String },
235}
236
237#[derive(Serialize, Deserialize, Debug)]
238pub struct ResponseStreamEvent {
239 pub created: u32,
240 pub model: String,
241 pub choices: Vec<ChoiceDelta>,
242 pub usage: Option<Usage>,
243}
244
245pub async fn stream_completion(
246 client: &dyn HttpClient,
247 api_url: &str,
248 api_key: &str,
249 request: Request,
250 low_speed_timeout: Option<Duration>,
251) -> Result<BoxStream<'static, Result<ResponseStreamEvent>>> {
252 let uri = format!("{api_url}/chat/completions");
253 let mut request_builder = HttpRequest::builder()
254 .method(Method::POST)
255 .uri(uri)
256 .header("Content-Type", "application/json")
257 .header("Authorization", format!("Bearer {}", api_key));
258
259 if let Some(low_speed_timeout) = low_speed_timeout {
260 request_builder = request_builder.low_speed_timeout(100, low_speed_timeout);
261 };
262
263 let request = request_builder.body(AsyncBody::from(serde_json::to_string(&request)?))?;
264 let mut response = client.send(request).await?;
265 if response.status().is_success() {
266 let reader = BufReader::new(response.into_body());
267 Ok(reader
268 .lines()
269 .filter_map(|line| async move {
270 match line {
271 Ok(line) => {
272 let line = line.strip_prefix("data: ")?;
273 if line == "[DONE]" {
274 None
275 } else {
276 match serde_json::from_str(line) {
277 Ok(ResponseStreamResult::Ok(response)) => Some(Ok(response)),
278 Ok(ResponseStreamResult::Err { error }) => {
279 Some(Err(anyhow!(error)))
280 }
281 Err(error) => Some(Err(anyhow!(error))),
282 }
283 }
284 }
285 Err(error) => Some(Err(anyhow!(error))),
286 }
287 })
288 .boxed())
289 } else {
290 let mut body = String::new();
291 response.body_mut().read_to_string(&mut body).await?;
292
293 #[derive(Deserialize)]
294 struct OpenAiResponse {
295 error: OpenAiError,
296 }
297
298 #[derive(Deserialize)]
299 struct OpenAiError {
300 message: String,
301 }
302
303 match serde_json::from_str::<OpenAiResponse>(&body) {
304 Ok(response) if !response.error.message.is_empty() => Err(anyhow!(
305 "Failed to connect to OpenAI API: {}",
306 response.error.message,
307 )),
308
309 _ => Err(anyhow!(
310 "Failed to connect to OpenAI API: {} {}",
311 response.status(),
312 body,
313 )),
314 }
315 }
316}
317
318#[derive(Copy, Clone, Serialize, Deserialize)]
319pub enum OpenAiEmbeddingModel {
320 #[serde(rename = "text-embedding-3-small")]
321 TextEmbedding3Small,
322 #[serde(rename = "text-embedding-3-large")]
323 TextEmbedding3Large,
324}
325
326#[derive(Serialize)]
327struct OpenAiEmbeddingRequest<'a> {
328 model: OpenAiEmbeddingModel,
329 input: Vec<&'a str>,
330}
331
332#[derive(Deserialize)]
333pub struct OpenAiEmbeddingResponse {
334 pub data: Vec<OpenAiEmbedding>,
335}
336
337#[derive(Deserialize)]
338pub struct OpenAiEmbedding {
339 pub embedding: Vec<f32>,
340}
341
342pub fn embed<'a>(
343 client: &dyn HttpClient,
344 api_url: &str,
345 api_key: &str,
346 model: OpenAiEmbeddingModel,
347 texts: impl IntoIterator<Item = &'a str>,
348) -> impl 'static + Future<Output = Result<OpenAiEmbeddingResponse>> {
349 let uri = format!("{api_url}/embeddings");
350
351 let request = OpenAiEmbeddingRequest {
352 model,
353 input: texts.into_iter().collect(),
354 };
355 let body = AsyncBody::from(serde_json::to_string(&request).unwrap());
356 let request = HttpRequest::builder()
357 .method(Method::POST)
358 .uri(uri)
359 .header("Content-Type", "application/json")
360 .header("Authorization", format!("Bearer {}", api_key))
361 .body(body)
362 .map(|request| client.send(request));
363
364 async move {
365 let mut response = request?.await?;
366 let mut body = String::new();
367 response.body_mut().read_to_string(&mut body).await?;
368
369 if response.status().is_success() {
370 let response: OpenAiEmbeddingResponse =
371 serde_json::from_str(&body).context("failed to parse OpenAI embedding response")?;
372 Ok(response)
373 } else {
374 Err(anyhow!(
375 "error during embedding, status: {:?}, body: {:?}",
376 response.status(),
377 body
378 ))
379 }
380 }
381}
382
383pub fn extract_text_from_events(
384 response: impl Stream<Item = Result<ResponseStreamEvent>>,
385) -> impl Stream<Item = Result<String>> {
386 response.filter_map(|response| async move {
387 match response {
388 Ok(mut response) => Some(Ok(response.choices.pop()?.delta.content?)),
389 Err(error) => Some(Err(error)),
390 }
391 })
392}