Преглед на файлове

Create logic for submitting prompt to deepinfra

Lee Morgan преди 1 месец
родител
ревизия
be93fd3ac6
променени са 3 файла, в които са добавени 64 реда и са изтрити 0 реда
  1. 2 0
      src/logic/mod.rs
  2. 60 0
      src/logic/prompt.rs
  3. 2 0
      src/main.rs

+ 2 - 0
src/logic/mod.rs

@@ -3,6 +3,8 @@ pub mod hash_password;
 pub mod compare_password;
 pub mod create_cookie;
 
+pub mod prompt;
+
 pub use valid_password::valid_password;
 pub use hash_password::hash_password;
 pub use compare_password::compare_password;

+ 60 - 0
src/logic/prompt.rs

@@ -0,0 +1,60 @@
+use reqwest::Client;
+use serde::Serialize;
+use serde_json::{json, Value};
+use crate::{
+    http_error::HttpError,
+    DEEPINFRA_TOKEN
+};
+
+#[derive(Serialize)]
+pub struct PromptMessage {
+    role: PromptRole,
+    content: String
+}
+
+#[derive(Serialize)]
+#[serde(rename_all = "lowercase")]
+pub enum PromptRole {
+    User,
+    System,
+    Assistant
+}
+
+impl PromptRole {
+    fn to_string(&self) -> String {
+        match self {
+            PromptRole::System => "system".to_string(),
+            PromptRole::User => "user".to_string(),
+            PromptRole::Assistant => "assistant".to_string()
+        }
+    }
+}
+
+pub async fn send_message(prompt: Vec<PromptMessage>) -> Result<String, HttpError> {
+    let client = Client::new();
+    let key = DEEPINFRA_TOKEN.get().unwrap();
+
+    let body = json!({
+        "model": "Qwen/Qwen3-32B",
+        "messages": prompt
+    });
+
+    let response = client
+        .post("https://api.deepinfra.com/v1/openai/chat/completions")
+        .header("Authorization", format!("Bearer {}", key))
+        .header("Content-Type", "application/json")
+        .json(&body)
+        .send()
+        .await
+        .map_err(|e| HttpError::InternalError(e.to_string()))?;
+
+    let data: Value = response
+        .json()
+        .await
+        .map_err(|e| HttpError::InternalError(e.to_string()))?;
+
+    let content = &data["choices"][0]["message"]["content"];
+    println!("{}", content);
+
+    Ok(content.to_string())
+}

+ 2 - 0
src/main.rs

@@ -16,6 +16,7 @@ static STRIPE_KEY: OnceLock<String> = OnceLock::new();
 static TOKEN_PACK_PRICE: OnceLock<i32> = OnceLock::new();
 static TOKENS_PER_PACK: OnceLock<i32> = OnceLock::new();
 static STRIPE_WEBHOOK_SECRET: OnceLock<String> = OnceLock::new();
+static DEEPINFRA_TOKEN: OnceLock<String> = OnceLock::new();
 
 #[actix_web::main]
 async fn main() -> std::io::Result<()> {
@@ -30,6 +31,7 @@ async fn main() -> std::io::Result<()> {
     TOKEN_PACK_PRICE.set(env_var("TOKEN_PACK_PRICE")).expect("Failed to set token pack price");
     TOKENS_PER_PACK.set(env_var("TOKENS_PER_PACK")).expect("Failed to set tokens per pack");
     STRIPE_WEBHOOK_SECRET.set(env_var("STRIPE_WEBHOOK_SECRET")).expect("Failed to set stripe webhook secret");
+    DEEPINFRA_TOKEN.set(env_var("DEEPINFRA_TOKEN")).expect("Failed to set DeepInfra token");
 
     //Database
     let pool = PgPool::connect(&database_url)