70 lines
1.7 KiB
Rust
70 lines
1.7 KiB
Rust
// src/llm.rs
|
||
use serde::{Deserialize, Serialize};
|
||
|
||
use crate::messages::ChatMsg;
|
||
|
||
#[derive(Serialize)]
|
||
struct LlmRequest {
|
||
model: String,
|
||
messages: Vec<Message>,
|
||
}
|
||
|
||
#[derive(Serialize, Deserialize)]
|
||
struct Message {
|
||
role: String, // "user" or "assistant"
|
||
content: String,
|
||
}
|
||
|
||
pub struct LlmWorker {
|
||
uri: String,
|
||
}
|
||
|
||
impl LlmWorker {
|
||
pub fn new(uri: String) -> Self {
|
||
Self { uri }
|
||
}
|
||
|
||
pub async fn query(&self, message: &ChatMsg) -> Result<ChatMsg, String> {
|
||
let client = reqwest::Client::new();
|
||
|
||
// Build the request body
|
||
let payload = LlmRequest {
|
||
model: "gemma2-9b-it".into(), // whatever model you run locally
|
||
messages: vec![Message {
|
||
role: "user".into(),
|
||
content: message.text.clone().into(),
|
||
}],
|
||
};
|
||
|
||
// POST to lm‑studio (default 127.0.0.1:1234)
|
||
let resp = client
|
||
.post(self.uri.clone())
|
||
.json(&payload)
|
||
.send()
|
||
.await
|
||
.map_err(|_| String::from("Failed to make request to LLM server"))?;
|
||
|
||
// The API returns a JSON with `choices[].message.content`
|
||
#[derive(Deserialize)]
|
||
struct LlmResponse {
|
||
choices: Vec<Choice>,
|
||
}
|
||
#[derive(Deserialize)]
|
||
struct Choice {
|
||
message: Message,
|
||
}
|
||
|
||
let llm_resp: LlmResponse = resp
|
||
.json()
|
||
.await
|
||
.map_err(|_| String::from("Failed to make request to LLM server"))?;
|
||
|
||
Ok(ChatMsg {
|
||
display_name: message.display_name.clone(),
|
||
user_id: 0,
|
||
text: llm_resp.choices[0].message.content.clone(),
|
||
timestamp: chrono::Utc::now().timestamp() as usize,
|
||
})
|
||
}
|
||
}
|