feat(m2): add llm module — build_system_prompt and Ollama streaming client
This commit is contained in:
parent
9c45015052
commit
c149c598ae
1 changed files with 142 additions and 0 deletions
142
src-tauri/src/llm.rs
Normal file
142
src-tauri/src/llm.rs
Normal file
|
|
@ -0,0 +1,142 @@
|
||||||
|
use anyhow::Context;
|
||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
use tauri::{AppHandle, Emitter};
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize)]
|
||||||
|
pub struct ChatMessage {
|
||||||
|
pub role: String,
|
||||||
|
pub content: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Deserialize)]
|
||||||
|
struct OllamaChunk {
|
||||||
|
message: Option<OllamaChunkMessage>,
|
||||||
|
done: bool,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Deserialize)]
|
||||||
|
struct OllamaChunkMessage {
|
||||||
|
content: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn build_system_prompt(source_os: &str, distro: &str) -> String {
|
||||||
|
format!(
|
||||||
|
"You are Robin, a friendly Linux migration companion built into the user's desktop. \
|
||||||
|
The user migrated from {source_os} and is currently using {distro}. \
|
||||||
|
Help them with Linux questions. When concepts differ from their previous OS, \
|
||||||
|
explain the Linux equivalent clearly. Be concise, practical, and warm. \
|
||||||
|
Use plain language and avoid unexplained jargon. \
|
||||||
|
Never make the user feel bad for not knowing something."
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn chat_stream(
|
||||||
|
base_url: &str,
|
||||||
|
model: &str,
|
||||||
|
messages: Vec<ChatMessage>,
|
||||||
|
app: &AppHandle,
|
||||||
|
) -> anyhow::Result<()> {
|
||||||
|
let client = reqwest::Client::new();
|
||||||
|
let url = format!("{base_url}/api/chat");
|
||||||
|
let body = serde_json::json!({
|
||||||
|
"model": model,
|
||||||
|
"messages": messages,
|
||||||
|
"stream": true
|
||||||
|
});
|
||||||
|
|
||||||
|
let mut response = client
|
||||||
|
.post(&url)
|
||||||
|
.json(&body)
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.context("failed to reach Ollama — is it running at the configured URL?")?;
|
||||||
|
|
||||||
|
if !response.status().is_success() {
|
||||||
|
let status = response.status();
|
||||||
|
let text = response.text().await.unwrap_or_default();
|
||||||
|
anyhow::bail!("Ollama returned {status}: {text}");
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut buffer = String::new();
|
||||||
|
|
||||||
|
loop {
|
||||||
|
match response.chunk().await {
|
||||||
|
Ok(Some(bytes)) => {
|
||||||
|
buffer.push_str(&String::from_utf8_lossy(&bytes));
|
||||||
|
while let Some(nl) = buffer.find('\n') {
|
||||||
|
let line = buffer[..nl].trim().to_string();
|
||||||
|
buffer.drain(..=nl);
|
||||||
|
if line.is_empty() {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
match serde_json::from_str::<OllamaChunk>(&line) {
|
||||||
|
Ok(chunk) if chunk.done => {
|
||||||
|
let _ = app.emit("robin:chat-done", ());
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
Ok(chunk) => {
|
||||||
|
if let Some(msg) = chunk.message {
|
||||||
|
if !msg.content.is_empty() {
|
||||||
|
let _ = app.emit("robin:chat-token", msg.content);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
log::warn!("llm: failed to parse Ollama chunk: {e} — raw: {line}");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(None) => break,
|
||||||
|
Err(e) => {
|
||||||
|
return Err(anyhow::anyhow!("stream read error: {e}"));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Stream ended without a done:true line
|
||||||
|
let _ = app.emit("robin:chat-done", ());
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn prompt_includes_source_os_and_distro() {
|
||||||
|
let prompt = build_system_prompt("macOS", "cachyos");
|
||||||
|
assert!(prompt.contains("macOS"), "prompt must mention source OS");
|
||||||
|
assert!(prompt.contains("cachyos"), "prompt must mention distro");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn prompt_includes_robin_name() {
|
||||||
|
let prompt = build_system_prompt("Windows", "linuxmint");
|
||||||
|
assert!(prompt.contains("Robin"), "prompt must identify Robin as the assistant");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn parse_token_chunk() {
|
||||||
|
let json = r#"{"model":"llama3.2","created_at":"2024-01-01T00:00:00Z","message":{"role":"assistant","content":"Hello"},"done":false}"#;
|
||||||
|
let chunk: OllamaChunk = serde_json::from_str(json).unwrap();
|
||||||
|
assert!(!chunk.done);
|
||||||
|
assert_eq!(chunk.message.unwrap().content, "Hello");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn parse_done_chunk() {
|
||||||
|
let json = r#"{"model":"llama3.2","created_at":"2024-01-01T00:00:00Z","done":true,"done_reason":"stop"}"#;
|
||||||
|
let chunk: OllamaChunk = serde_json::from_str(json).unwrap();
|
||||||
|
assert!(chunk.done);
|
||||||
|
assert!(chunk.message.is_none());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn parse_empty_token_is_handled() {
|
||||||
|
let json = r#"{"model":"llama3.2","created_at":"2024-01-01T00:00:00Z","message":{"role":"assistant","content":""},"done":false}"#;
|
||||||
|
let chunk: OllamaChunk = serde_json::from_str(json).unwrap();
|
||||||
|
assert!(!chunk.done);
|
||||||
|
assert_eq!(chunk.message.unwrap().content, "");
|
||||||
|
}
|
||||||
|
}
|
||||||
Loading…
Reference in a new issue