openai_client_cli/program/
entry.rs

1use crate::*;
2use clap::{arg, command};
3use eventsource_stream::Eventsource;
4use futures_util::StreamExt;
5use http::header::CONTENT_TYPE;
6use mime::Mime;
7use std::path::PathBuf;
8use std::io::stderr;
9use tracing::{info, Level};
10
11#[doc(hidden)]
12pub use clap::Parser;
13
14/// The main entry-point for the program.
15#[derive(Parser)]
16#[command(
17  author,
18  about,
19  bin_name = "openai-client",
20  help_template = "\
21{before-help}\
22{name} {version} by {author}
23{about}
24
25{usage-heading} {usage}
26
27{all-args}\
28{after-help}",
29  version,
30  next_line_help = true,
31)]
32pub struct Entry {
33  /// The file path where the API key is stored.
34  #[arg(
35    help = "\
36The file path where the API key is stored.
37The program will attempt the following steps to obtain a valid API key:
38 1. Read the file from the provided path <KEY_FILE_PATH>.
39 2. Read the environment variable `OPENAI_API_KEY`.
40 3. Read the file from the default paths in the following order:
41    `openai.env`, `.openai_profile`, `.env`,
42    `~/openai.env`, `~/.openai_profile` or `~/.env`.
43 4. Exit the program with a non-zero return code.
44",
45    long,
46    short = 'k',
47    value_name = "KEY_FILE_PATH",
48  )]
49  pub key_file: Option<PathBuf>,
50
51  /// The HTTP method used for the API request.
52  #[arg(
53    help = "\
54The HTTP method used for the API request.
55The program will attempt the following steps to determine a valid HTTP method:
56 1. Read the value of argument <METHOD>.
57 2. If the `parameter` object is successfully fetched from either
58    <PARAM_FILE_PATH> or one of the default paths, set <METHOD> to `POST`.
59 3. Otherwise, set <METHOD> to `GET`.
60",
61    long,
62    short = 'm',
63    value_name = "METHOD",
64  )]
65  pub method: Option<String>,
66
67  /// The file path where the organization ID is stored.
68  #[arg(
69    help = "\
70The file path where the organization ID is stored.
71The program will attempt the following steps to obtain a valid organization ID:
72 1. Read the file from the provided path <ORG_FILE_PATH>.
73 2. Read the file from provided path of key file <KEY_FILE_PATH>.
74 3. Read the environment variable `OPENAI_ORG_KEY`.
75 4. Read the file from the default paths in the following order:
76    `openai.env`, `.openai_profile`, `.env`,
77    `~/openai.env`, `~/.openai_profile` or `~/.env`.
78 5. Ignore the field and leave it empty.
79",
80    short = 'g',
81    long = "org-file",
82    value_name = "ORG_FILE_PATH",
83  )]
84  pub organization_file: Option<PathBuf>,
85
86  /// The file path where the API response will be stored.
87  #[arg(
88    help = "\
89The file path where the API response will be stored.
90The program will attempt the following steps to successfully store the response:
91 1. Export the output to the provided file path <OUTPUT_FILE_PATH>.
92 2. Export the output to the standard output.
93 3. Exit the program with a non-zero return code.
94",
95    long,
96    short = 'o',
97    value_name = "OUTPUT_FILE_PATH",
98  )]
99  pub output_file: Option<PathBuf>,
100
101  /// The file path where the API request parameters (body) are stored in JSON format.
102  #[arg(
103    help = "\
104The file path where the API request parameters (body) are stored in JSON format.
105The program will attempt the following steps to obtain a valid parameter object:
106 1. Read the file from the provided path <PARAM_FILE_PATH>.
107 2. Read the file from the default paths in the following order:
108    `openai.json`, `openai-parameters.json`, `openai_parameters.json`,
109    `openai-parameters`, `openai_parameters`, or `openai.config.json`.
110 3. Ignore the field and leave it empty
111",
112    long,
113    short = 'p',
114    value_name = "PARAM_FILE_PATH",
115  )]
116  pub parameter_file: Option<PathBuf>,
117
118  /// Hidden.
119  #[arg(hide = true, long, exclusive = true)]
120  pub _parameter: Option<Parameter>,
121
122  /// The API request path. (part of the URL)
123  #[arg(
124    help = "\
125The API request path. (part of the URL)
126The program will use regex to extract the matched segment in <PATH>.
127For example, the extracted strings will be the same when <PATH> is either
128`chat/completions`, `/chat/completions` or `https://api.openai.com/v1/chat/completions`.",
129    value_name = "PATH",
130  )]
131  pub path: String,
132
133  /// Switch for verbose logging mode.
134  #[arg(
135    default_value = "false",
136    help = "\
137Switch for verbose logging mode. This mode is useful for debugging purposes.
138It is disabled by default.
139",
140    long,
141    short = 'v',
142  )]
143  pub verbose: bool,
144}
145
146impl Entry {
147  /// Run the program.
148  pub async fn run(mut self) -> Result<()> {
149    let logger = tracing_subscriber::fmt()
150      .with_target(false)
151      .with_writer(stderr)
152      .without_time();
153    if self.verbose {
154      logger
155        .with_max_level(Level::DEBUG)
156        .with_file(true)
157        .with_line_number(true)
158        .init();
159    } else {
160      logger
161        .with_max_level(Level::WARN)
162        .init();
163    }
164
165    let key = Key::fetch(&self)?;
166    let organization = Organization::fetch(&self).ok();
167    if organization.is_none() {
168      info!("Ignored the field `organization` for not being fetched successfully");
169    }
170    let output = Output::fetch(&self)?;
171    // `parameter` should be fetched before `method`
172    let parameter = Parameter::fetch(&self).ok();
173    if parameter.is_none() {
174      info!("Ignored the field `parameter` for not being fetched successfully");
175    }
176    self._parameter = parameter;
177    let path = Path::fetch(&self)?;
178    let method = Method::fetch(&self)?;
179
180    let client = OpenAIClient::new(key, organization);
181    let request = OpenAIRequest::new(method, path, self._parameter)?;
182    let response = client.send(request).await?;
183    // debug!("\n{:#?}", response);
184
185    let status_error = response.error_for_status_ref().map(|_| ());
186    let content_type: Mime = response
187      .headers()
188      .get(CONTENT_TYPE)
189      .ok_or(Error::msg("The API response does not contain the header `Content-Type`"))?
190      .to_str()?
191      .parse()?;
192    info!("Resolving the API response in the content type: {content_type:?}");
193
194    let exporting_message = format!(
195      "Exporting the output to the {}",
196      if output.is_file() { "file" } else { "standard output" },
197    );
198
199    match content_type.subtype() {
200      mime::JSON => {
201        let response_json = response
202          .json::<serde_json::Value>()
203          .await
204          .map_err(Error::from)
205          .and_then(|object| {
206            serde_json::to_string_pretty(&object)
207              .map_err(Error::from)
208          });
209        if let Ok(response_json) = &response_json {
210          info!(
211            "Resolved the API response: <JSON Object ({} bytes)>",
212            response_json.len(),
213          );
214        }
215
216        if response_json.is_err() || status_error.is_err() {
217          Err(
218            Error::msg("\u{1b}[F")
219              .context(response_json.map_or_else(
220                |e| e.to_string(),
221                |json| format!("The API response in JSON format:\n{}", json)),
222              )
223              .context(status_error.map_or_else(
224                |e| e.to_string(),
225                |_| String::new(),
226              ))
227              .context("Failed to resolve the API response")
228          )
229        } else {
230          let response_json = response_json.unwrap();
231          let mut output = output.value();
232          info!("{}", exporting_message);
233          output.write_all(response_json.as_bytes())?;
234          Ok(())
235        }
236      },
237      mime::EVENT_STREAM => {
238        status_error?; // should not be an error
239
240        info!("{}", exporting_message);
241        let mut stream = response.bytes_stream().eventsource();
242        let mut output = output.value();
243        while let Some(chunk) = stream.next().await {
244          let chunk = chunk?;
245          let data = chunk.data;
246          info!(
247            "Resolved the API response: <Event Stream Data: ({} bytes)>",
248            data.len(),
249          );
250          if data == "[DONE]" {
251            info!("Reached the end of the API response");
252            break;
253          }
254          if chunk.retry.is_some() {
255            return Err(Error::msg("Failed to resolve API response: Retry occurred"));
256          }
257          output.write_all(&[data.as_bytes(), b"\n"].concat())?;
258        }
259        Ok(())
260      },
261      _ => Err(Error::msg(format!(
262        "Failed to resolve API response: {content_type:?} is an invalid format"
263      ))),
264    }
265  }
266}