chat_gpt_rs/
api.rs

1//! OpenAI API client.
2//!
3//! Communicates with the OpenAI API to send chat requests and receive chat completions.
4
5use reqwest::{
6    header::{HeaderMap, HeaderValue, AUTHORIZATION, CONTENT_TYPE},
7    Client,
8};
9
10use crate::{
11    error::{Error, OpenAIErrorResponse, Result},
12    request::{Request, Response},
13    token::Token,
14};
15
16/// OpenAI API client.
17pub struct Api {
18    organization_id: Option<String>,
19    client: Client,
20}
21
22impl Api {
23    /// Base URL for the OpenAI API.
24    const BASE_URL: &'static str = "https://api.openai.com/v1";
25
26    /// URL for the chat endpoint.
27    const CHAT_URL: &'static str = "/chat/completions";
28
29    /// Create a new API client.
30    pub fn new(key: Token) -> Api {
31        let mut headers = HeaderMap::new();
32        headers.insert(AUTHORIZATION, key.into());
33        headers.insert(CONTENT_TYPE, HeaderValue::from_static("application/json"));
34        Api {
35            organization_id: None,
36            client: Client::builder().default_headers(headers).build().unwrap(),
37        }
38    }
39
40    /// Set the organization ID for the API client.
41    pub fn with_organization_id(mut self, organization_id: String) -> Self {
42        self.organization_id = Some(organization_id);
43        self
44    }
45
46    /// Send a chat request to the OpenAI API.
47    pub async fn chat(&self, r: Request) -> Result<Response> {
48        let url = format!("{}{}", Self::BASE_URL, Self::CHAT_URL);
49        let mut req = self.client.post(&url);
50        if let Some(organization_id) = &self.organization_id {
51            req = req.header("OpenAI-Organization", organization_id);
52        }
53        let res = req.json(&r).send().await?;
54
55        //if status is ok, return as Response json, otherwise return as ApiError
56        match res.status().is_success() {
57            true => Ok(res.json::<Response>().await?),
58            false => Err(Error::ApiError(res.json::<OpenAIErrorResponse>().await?)),
59        }
60    }
61}