使用 Rust 透過 OpenAI API 生成文字補全

Others
, , , ,

前置作業

在 OpenAI 註冊一個帳號，並且在 API keys 頁面產生一個 API 金鑰。

建立專案

建立專案。

1
2
cargo new gpt-cli-rust
cd gpt-cli-rust

安裝依賴套件。

1
2
3
4
5
[dependencies]
reqwest = { version = "0.11", features = ["json"] }
serde = { version = "1.0", features = ["derive"] }
tokio = { version = "1.23", features = ["full"] }
dotenvy = { version = "0.15" }

新增 .env 檔。

1
2
OPENAI_API_URL=https://api.openai.com/v1
OPENAI_API_KEY=

新增 .gitignore 檔。

1
2
/target
.env

實作

新增 src/lib.rs 檔。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
use reqwest::header::HeaderMap;
use serde::{Deserialize, Serialize};
use std::{env, error::Error};

pub async fn fetch() -> Result<Response, Box<dyn Error>> {
    let api_url = env::var("OPENAI_API_URL").unwrap();
    let api_key = env::var("OPENAI_API_KEY").unwrap();

    let mut headers = HeaderMap::new();
    headers.insert("Content-Type", "application/json".parse().unwrap());
    headers.insert(
        "Authorization",
        format!("Bearer {}", api_key).parse().unwrap(),
    );

    let body = Request{
        model: String::from("text-davinci-003"),
        prompt: String::from("\n\nHuman: Hello, who are you?\nAI: I am an AI created by OpenAI. How can I help you today?\nHuman: Hello, who are you?"),
        temperature: 0.9,
        max_tokens: 150,
        top_p: 1.0,
        frequency_penalty: 0.0,
        presence_penalty: 0.6,
        stop: vec![String::from(" Human:"), String::from(" AI:")],
    };

    let client = reqwest::Client::new();
    let res = client
        .post(format!("{}/completions", api_url))
        .headers(headers)
        .json(&body)
        .send()
        .await?
        .json::<Response>()
        .await?;

    return Ok(res);
}

#[derive(Deserialize, Debug)]
pub struct Response {
    pub id: Option<String>,
    pub object: Option<String>,
    pub created: Option<usize>,
    pub model: Option<String>,
    pub choices: Option<Vec<Choice>>,
    pub usage: Option<Usage>,
}

#[derive(Deserialize, Debug)]
pub struct Choice {
    pub text: Option<String>,
    pub index: Option<usize>,
    pub logprobs: Option<usize>,
    pub finish_reason: Option<String>,
}

#[derive(Deserialize, Debug)]
pub struct Usage {
    pub prompt_tokens: Option<usize>,
    pub completion_tokens: Option<usize>,
    pub total_tokens: Option<usize>,
}

#[derive(Serialize)]
struct Request {
    model: String,
    prompt: String,
    temperature: f32,
    max_tokens: usize,
    top_p: f32,
    frequency_penalty: f32,
    presence_penalty: f32,
    stop: Vec<String>,
}

新增 src/main.rs 檔。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
use dotenvy::dotenv;
use openai_cli_rust::fetch;
use std::error::Error;

#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    dotenv().ok();

    let resp = fetch().await?;
    if let Some(choices) = resp.choices {
        for choice in choices.iter() {
            println!("{:?}", choice.text);
        }
    }

    Ok(())
}

執行程式。

1
cargo run

程式碼

參考資料