文章目录
前言
你好,我是醉墨居士,最近准备花一些时间来使用Rust语言实现一个智能助手,希望能够帮助到你。
环境准备
依赖
toml
ollama-rs = "0.2.2"
tokio = { version = "1.43.0", features = ["full"] }
tokio-stream = { version = "0.1.17" }
axum = "0.8.1"
代码
rust
use std::{collections::HashMap, sync::OnceLock};
use axum::{
extract::Query, routing::get, Json, Router
};
use ollama_rs::{
generation::completion::request::GenerationRequest, Ollama
};
#[derive(Debug)]
struct Context {
ollama: Ollama,
}
static CONTEXT: OnceLock<Context> = OnceLock::new();
impl Context {
fn global() -> &'static Self {
CONTEXT.get_or_init(|| {
println!("Initializing Ollama...");
Context {
ollama: Ollama::default(),
}
})
}
}
#[tokio::main]
async fn main() {
let app = Router::new().route("/", get(chat));
let listener = tokio::net::TcpListener::bind("0.0.0.0:3000").await.unwrap();
axum::serve(listener, app).await.unwrap();
}
async fn chat(Query(mut query): Query<HashMap<String, String>>) -> Json<String> {
let prompt = if let Some(prompt) = query.remove("prompt") {
prompt
} else {
return Json("Error: prompt not provided".to_string());
};
let req = GenerationRequest::new("qwen2.5:0.5b".into(), prompt);
let res = Context::global().ollama.generate(req).await;
Json(
if let Ok(res) = res {
res.response
} else {
"Error".to_string()
}
)
}
运行
bash
cargo run
使用
打开浏览器,访问http://localhost:3000?prompt=你好,即可得到AI回复
最后
我是醉墨居士,这篇博客就先到这里,后续我会继续更新,欢迎你提出宝贵意见,互相交流,共同进步