chatgpt_private_test/
lib.rs

1pub mod utils;
2
3use dotenv::dotenv;
4use http_req::{request::Method, request::Request, uri::Uri};
5use serde::Deserialize;
6use serde_json::Value;
7use std::env;
8
9pub async fn custom_gpt(sys_prompt: &str, u_prompt: &str, m_token: u16) -> Option<String> {
10    let system_prompt = serde_json::json!(
11        {"role": "system", "content": sys_prompt}
12    );
13    let user_prompt = serde_json::json!(
14        {"role": "user", "content": u_prompt}
15    );
16
17    match chat(vec![system_prompt, user_prompt], m_token).await {
18        Ok((res, _count)) => Some(res),
19        Err(_) => None,
20    }
21}
22
23pub async fn chat(message_obj: Vec<Value>, m_token: u16) -> Result<(String, u32), anyhow::Error> {
24    dotenv().ok();
25    let api_token = env::var("OPENAI_API_TOKEN")?;
26
27    let params = serde_json::json!({
28      "model": "gpt-3.5-turbo-16k",
29      "messages": message_obj,
30      "temperature": 0.7,
31      "top_p": 1,
32      "n": 1,
33      "stream": false,
34      "max_tokens": m_token,
35      "presence_penalty": 0,
36      "frequency_penalty": 0,
37      "stop": "\n"
38    });
39
40    let uri = "https://api.openai.com/v1/chat/completions";
41
42    let uri = Uri::try_from(uri)?;
43    let mut writer = Vec::new();
44    let body = serde_json::to_vec(&params)?;
45
46    let bearer_token = format!("Bearer {}", api_token);
47    let _response = Request::new(&uri)
48        .method(Method::POST)
49        .header("Authorization", &bearer_token)
50        .header("Content-Type", "application/json")
51        .header("Content-Length", &body.len())
52        .body(&body)
53        .send(&mut writer)?;
54
55    let res = serde_json::from_slice::<ChatResponse>(&writer)?;
56    let token_count = res.usage.total_tokens;
57    Ok((res.choices[0].message.content.to_string(), token_count))
58}
59
60#[derive(Deserialize)]
61pub struct ChatResponse {
62    pub id: String,
63    pub choices: Vec<Choice>,
64    pub usage: Usage,
65}
66
67#[derive(Deserialize)]
68pub struct Choice {
69    pub index: u32,
70    pub message: Message,
71    pub finish_reason: String,
72}
73
74#[derive(Deserialize)]
75pub struct Message {
76    pub role: String,
77    pub content: String,
78}
79
80#[derive(Deserialize)]
81pub struct Usage {
82    pub prompt_tokens: u32,
83    pub completion_tokens: u32,
84    pub total_tokens: u32,
85}