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