rust_completions/
lib.rs

1use serde::{Deserialize, Serialize};
2use std::{env, error::Error};
3
4#[derive(Debug)]
5pub struct Config {
6    pub organization: String,
7    pub api_key: String,
8}
9
10impl Config {
11    pub fn build() -> Result<Config, &'static str> {
12        let env_vars = env::vars();
13
14        if env_vars.count() < 2 {
15            return Err("Not enough arguments have been supplied");
16        }
17
18        let organization = match env::var("ORGANIZATION_ID") {
19            Ok(value) => value,
20            Err(_) => return Err("ORGANIZATION_ID must exists"),
21        };
22
23        let api_key = match env::var("API_KEY") {
24            Ok(value) => value,
25            Err(_) => return Err("API_KEY must exists"),
26        };
27
28        Ok(Config {
29            organization,
30            api_key,
31        })
32    }
33}
34
35// TODO: Handle Request Options such as max_tokens
36
37#[derive(Serialize)]
38pub struct CompletionRequest {
39    pub model: String,
40    pub messages: Vec<Message>,
41    pub temperature: f64,
42}
43
44#[derive(Deserialize, Debug)]
45pub struct Usage {
46    pub prompt_tokens: u64,
47    pub completion_tokens: u64,
48    pub total_tokens: u64,
49}
50
51#[derive(Deserialize, Debug)]
52pub struct Choice {
53    pub message: Message,
54    pub finish_reason: String,
55    pub index: u64,
56}
57
58#[derive(Serialize, Deserialize, Debug)]
59pub struct Message {
60    pub role: String,
61    pub content: String,
62}
63
64#[derive(Deserialize, Debug)]
65pub struct CompletionResponse {
66    pub id: String,
67    pub object: String,
68    pub created: u64,
69    pub model: String,
70    pub usage: Usage,
71    pub choices: Vec<Choice>,
72}
73/// Get a chat completion from Open AI, you would need to borrow an instance a `Config` object
74/// and a Request body, defined by the struct `CompletionRequest`
75pub async fn new_chat_completion(
76    config: &Config,
77    body: CompletionRequest,
78) -> Result<CompletionResponse, Box<dyn Error>> {
79    let client = reqwest::Client::new();
80
81    let request = client
82        .post("https://api.openai.com/v1/chat/completions")
83        .bearer_auth(&config.api_key)
84        .header("OpenAI-Organization", &config.organization)
85        .json(&body)
86        .send()
87        .await?;
88
89    let response = request.json::<CompletionResponse>().await?;
90
91    Ok(response)
92}
93
94pub async fn run(config: Config) -> Result<(), Box<dyn Error>> {
95    let env_args: Vec<String> = env::args().collect();
96
97    let mut messages: Vec<Message> = Vec::new();
98    messages.push(Message {
99        role: String::from("user"),
100        content: String::from(&env_args[1]),
101    });
102
103    let body = CompletionRequest {
104        model: String::from("gpt-3.5-turbo"),
105        messages,
106        temperature: 0.0,
107    };
108
109    println!("Asking: \"{}\" with model: \"gpt-3.5-turbo\"", &env_args[1]);
110
111    let completion = new_chat_completion(&config, body).await?;
112
113    println!("Response: {}", completion.choices[0].message.content);
114
115    Ok(())
116}