42 lines
1.3 KiB
Rust
42 lines
1.3 KiB
Rust
use anyhow::Result;
|
|
use chat_gpt_lib_rs::{ChatGPTClient, ChatInput, Message, Model, Role};
|
|
|
|
pub struct AI {
|
|
chat_gpt_client: ChatGPTClient,
|
|
}
|
|
|
|
impl AI {
|
|
pub fn new(api_key: &str, base_url: &str) -> Self {
|
|
Self {
|
|
chat_gpt_client: ChatGPTClient::new(api_key, base_url),
|
|
}
|
|
}
|
|
|
|
pub async fn infer_country_of_living(
|
|
&self,
|
|
display_name: &str,
|
|
description: &str,
|
|
) -> Result<String> {
|
|
let chat_input = ChatInput {
|
|
model: Model::Gpt3_5Turbo,
|
|
messages: vec![
|
|
Message {
|
|
role: Role::System,
|
|
// TODO: Lol, prompt injection much?
|
|
content: "You are a tool that attempts to guess where a person is likely to be from based on their name and short bio. Please respond with two-letter country code only. If unable to determine, say xx.".to_string(),
|
|
},
|
|
Message {
|
|
role: Role::User,
|
|
content: format!("Name: {display_name}\nBio:\n{description}"),
|
|
},
|
|
],
|
|
..Default::default()
|
|
};
|
|
|
|
let response = self.chat_gpt_client.chat(chat_input).await?;
|
|
|
|
// TODO: Error handling?
|
|
return Ok(response.choices[0].message.content.to_lowercase());
|
|
}
|
|
}
|