diff --git a/Cargo.lock b/Cargo.lock index 4cbc4da..849e952 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -253,6 +253,7 @@ name = "doubleo7" version = "0.1.0" dependencies = [ "anyhow", + "reqwest 0.12.28", "rig-core", "serde", "tokio", @@ -325,6 +326,21 @@ version = "1.0.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" +[[package]] +name = "foreign-types" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f6f339eb8adc052cd2ca78910fda869aefa38d22d5cb648e6485e4d3fc06f3b1" +dependencies = [ + "foreign-types-shared", +] + +[[package]] +name = "foreign-types-shared" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "00b0228411908ca8685dba7fc2cdd70ec9990a6e753e89b6ac91a84c40fbaf4b" + [[package]] name = "form_urlencoded" version = "1.2.2" @@ -605,6 +621,22 @@ dependencies = [ "tower-service", ] +[[package]] +name = "hyper-tls" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "70206fc6890eaca9fde8a0bf71caa2ddfc9fe045ac9e5c70df101a7dbde866e0" +dependencies = [ + "bytes", + "http-body-util", + "hyper", + "hyper-util", + "native-tls", + "tokio", + "tokio-native-tls", + "tower-service", +] + [[package]] name = "hyper-util" version = "0.1.20" @@ -831,6 +863,12 @@ version = "0.2.189" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3eaf3ede3fee6db1a4c2ee091bf8a8b4dccdc6d17f656fb07896ee72867612f2" +[[package]] +name = "linux-raw-sys" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53" + [[package]] name = "litemap" version = "0.8.2" @@ -897,6 +935,23 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "native-tls" +version = "0.2.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "465500e14ea162429d264d44189adc38b199b62b1c21eea9f69e4b73cb03bbf2" +dependencies = [ + "libc", + "log", + "openssl", + "openssl-probe", + "openssl-sys", + "schannel", + "security-framework", + "security-framework-sys", + "tempfile", +] + [[package]] name = "nom" version = "7.1.3" @@ -922,12 +977,49 @@ version = "1.21.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" +[[package]] +name = "openssl" +version = "0.10.81" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77823a27f0babb03091cb9ed9ef80af3b39dbc82f97e8fa530374b7dafd87a45" +dependencies = [ + "bitflags", + "cfg-if", + "foreign-types", + "libc", + "openssl-macros", + "openssl-sys", +] + +[[package]] +name = "openssl-macros" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a948666b637a0f465e8564c73e89d4dde00d72d4d473cc972f390fc3dcee7d9c" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + [[package]] name = "openssl-probe" version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7c87def4c32ab89d880effc9e097653c8da5d6ef28e6b539d313baaacfbafcbe" +[[package]] +name = "openssl-sys" +version = "0.9.117" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b47e7e6bb2c38cd930d25a23b40fa52e068c10e85f3e03a7f5ba5aaca5713695" +dependencies = [ + "cc", + "libc", + "pkg-config", + "vcpkg", +] + [[package]] name = "ordered-float" version = "5.3.0" @@ -1196,6 +1288,46 @@ dependencies = [ "syn 3.0.3", ] +[[package]] +name = "reqwest" +version = "0.12.28" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147" +dependencies = [ + "base64", + "bytes", + "encoding_rs", + "futures-core", + "h2", + "http", + "http-body", + "http-body-util", + "hyper", + "hyper-rustls", + "hyper-tls", + "hyper-util", + "js-sys", + "log", + "mime", + "native-tls", + "percent-encoding", + "pin-project-lite", + "rustls-pki-types", + "serde", + "serde_json", + "serde_urlencoded", + "sync_wrapper", + "tokio", + "tokio-native-tls", + "tower", + "tower-http", + "tower-service", + "url", + "wasm-bindgen", + "wasm-bindgen-futures", + "web-sys", +] + [[package]] name = "reqwest" version = "0.13.4" @@ -1260,7 +1392,7 @@ dependencies = [ "mime_guess", "ordered-float", "pin-project-lite", - "reqwest", + "reqwest 0.13.4", "rig-derive", "schemars", "serde", @@ -1314,6 +1446,19 @@ dependencies = [ "semver", ] +[[package]] +name = "rustix" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6fe4565b9518b83ef4f91bb47ce29620ca828bd32cb7e408f0062e9930ba190" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys 0.61.2", +] + [[package]] name = "rustls" version = "0.23.43" @@ -1395,6 +1540,12 @@ version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cf54715a573b99ac80df0bc206da022bcd442c974952c7b9720069370852e21f" +[[package]] +name = "ryu" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" + [[package]] name = "same-file" version = "1.0.6" @@ -1542,6 +1693,18 @@ dependencies = [ "serde_core", ] +[[package]] +name = "serde_urlencoded" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3491c14715ca2294c4d6a88f15e84739788c1d030eed8c110436aafdaa2f3fd" +dependencies = [ + "form_urlencoded", + "itoa", + "ryu", + "serde", +] + [[package]] name = "sha1" version = "0.10.7" @@ -1682,6 +1845,19 @@ dependencies = [ "libc", ] +[[package]] +name = "tempfile" +version = "3.27.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" +dependencies = [ + "fastrand", + "getrandom 0.4.3", + "once_cell", + "rustix", + "windows-sys 0.61.2", +] + [[package]] name = "thiserror" version = "2.0.19" @@ -1755,6 +1931,16 @@ dependencies = [ "syn 3.0.3", ] +[[package]] +name = "tokio-native-tls" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbae76ab933c85776efabc971569dd6119c580d8f5d448769dec1764bf796ef2" +dependencies = [ + "native-tls", + "tokio", +] + [[package]] name = "tokio-rustls" version = "0.26.4" @@ -2013,6 +2199,12 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" +[[package]] +name = "vcpkg" +version = "0.2.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "accd4ea62f7bb7a82fe23066fb0957d48ef677f6eeb8215f372f52e48bb32426" + [[package]] name = "version_check" version = "0.9.5" diff --git a/Cargo.toml b/Cargo.toml index a00633b..44059af 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -5,6 +5,7 @@ edition = "2024" [dependencies] anyhow = "1.0.104" +reqwest = "0.12" rig-core = { git = "https://github.com/0xPlaygrounds/rig", branch = "main" } serde = "1.0.229" tokio = { version = "1.53.1", features = ["full"] } diff --git a/src/main.rs b/src/main.rs index 36e863e..fe5af3e 100644 --- a/src/main.rs +++ b/src/main.rs @@ -13,8 +13,7 @@ use rig_core::serde_json::json; use crate::models::Prompts; mod models; - -const LLAMA_SERVER_URL: &str = "http://127.0.0.1:8000"; +mod server; static PROMPTS: LazyLock = LazyLock::new(|| { toml::from_str(include_str!("prompts.toml")).expect("Could not parse prompts.toml") @@ -22,7 +21,9 @@ static PROMPTS: LazyLock = LazyLock::new(|| { #[tokio::main] async fn main() -> anyhow::Result<()> { - let gemma: ollama::CompletionModel = wire_gemma_client().await?; + // Independent setup steps (talk to unrelated backends, no data dependency) — run concurrently. + let (gemma, ()) = tokio::try_join!(wire_gemma_client(), server::ensure_running())?; + let shieldstral = wire_shieldstral().await?; let prompt_to_test = gemma.completion_request(&PROMPTS.gemma.prompt) @@ -53,7 +54,7 @@ async fn wire_gemma_client() -> anyhow::Result { } async fn wire_shieldstral() -> anyhow::Result> { - let client = llamafile::Client::from_url(LLAMA_SERVER_URL)?; + let client = llamafile::Client::from_url(&server::url())?; // Name doesn't matter here, server just uses whatever is running on it. let shieldstral = client.completion_model("shieldstral"); diff --git a/src/server.rs b/src/server.rs new file mode 100644 index 0000000..f3f197d --- /dev/null +++ b/src/server.rs @@ -0,0 +1,68 @@ +use std::sync::LazyLock; +use std::time::Duration; +use serde::Deserialize; +use tokio::process::Command; +use tokio::time::sleep; + +#[derive(Deserialize)] +struct ServerConfig { + llama_server: LlamaServerConfig, +} + +#[derive(Deserialize)] +struct LlamaServerConfig { + binary: String, + model_path: String, + host: String, + port: u16, + context_size: u32, +} + +static SERVER_CONFIG: LazyLock = LazyLock::new(|| { + toml::from_str(include_str!("server.toml")).expect("Could not parse server.toml") +}); + +static HTTP_CLIENT: LazyLock = LazyLock::new(reqwest::Client::new); + +pub(crate) fn url() -> String { + format!("http://{}:{}", SERVER_CONFIG.llama_server.host, SERVER_CONFIG.llama_server.port) +} + +async fn is_healthy(health_url: &str) -> bool { + HTTP_CLIENT.get(health_url).send().await.is_ok_and(|r| r.status().is_success()) +} + +/// Checks whether llama-server is already serving on the configured host/port, +/// and if not, spawns it from the configured binary/model path and waits for +/// it to report healthy before returning. +pub(crate) async fn ensure_running() -> anyhow::Result<()> { + let base_url = url(); + let health_url = format!("{base_url}/health"); + + if is_healthy(&health_url).await { + return Ok(()); + } + + println!("llama-server not running at {base_url}, starting it..."); + + Command::new(&SERVER_CONFIG.llama_server.binary) + .args([ + "-m", &SERVER_CONFIG.llama_server.model_path, + "--jinja", + "-c", &SERVER_CONFIG.llama_server.context_size.to_string(), + "--host", &SERVER_CONFIG.llama_server.host, + "--port", &SERVER_CONFIG.llama_server.port.to_string(), + ]) + .spawn() + .map_err(|e| anyhow::anyhow!("failed to spawn llama-server at {}: {e}", SERVER_CONFIG.llama_server.binary))?; + + for _ in 0..60 { + if is_healthy(&health_url).await { + println!("llama-server is up."); + return Ok(()); + } + sleep(Duration::from_secs(1)).await; + } + + anyhow::bail!("llama-server did not become healthy within 60s") +} diff --git a/src/server.toml b/src/server.toml new file mode 100644 index 0000000..97bf814 --- /dev/null +++ b/src/server.toml @@ -0,0 +1,6 @@ +[llama_server] +binary = "/home/austin/.local/share/llama.cpp/build/bin/llama-server" +model_path = "/home/austin/ai/Shieldstral-1.0-3B-BF16.gguf" +host = "127.0.0.1" +port = 8000 +context_size = 32768