use std::sync::LazyLock;
use anyhow::Result;
use ollama_rs::{
Ollama,
generation::completion::{GenerationResponseStream, request::GenerationRequest},
};
use tokio::io::{AsyncWriteExt, BufWriter};
use tokio_stream::StreamExt;
const DEF_MODEL: &str = "gemma3n:e4b";
fn ollama_client() -> &'static Ollama {
static OLLAMA_CLI: LazyLock<Ollama> = LazyLock::new(|| Ollama::default());
&OLLAMA_CLI
}
async fn generate_content_once(question: &str) -> Result<String> {
let cli = ollama_client();
let mut buf = vec![];
let mut bw = BufWriter::new(&mut buf);
let request = GenerationRequest::new(DEF_MODEL.into(), question.to_string());
let mut stream: GenerationResponseStream = cli.generate_stream(request).await?;
while let Some(Ok(res)) = stream.next().await {
for ele in res {
bw.write_all(ele.response.as_bytes()).await?;
bw.flush().await?;
}
}
Ok(String::from_utf8_lossy(&buf).to_string())
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_generate_content_once() {
let q = "我是谁?";
let res = generate_content_once(q).await.unwrap();
println!("res is\n{}", res)
}
}