Compare commits
4
Commits
a22560c337
...
fc391b5d0e
| Author | SHA1 | Date | |
|---|---|---|---|
 LucasDLTG
|
fc391b5d0e | ||
|
LucasDLTG
|
e31cdf131f | ||
|
LucasDLTG
|
d5856557b4 | ||
|
LucasDLTG
|
562d154480 |
Generated
+303
-2
@@ -11,6 +11,21 @@ dependencies = [
|
||||
"memchr",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "android_system_properties"
|
||||
version = "0.1.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "819e7219dbd41043ac279b19830f2efc897156490d7fd6ea916720117ee66311"
|
||||
dependencies = [
|
||||
"libc",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "anyhow"
|
||||
version = "1.0.102"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7f202df86484c868dbad7eaa557ef785d5c66295e41b460ef922eca0723b842c"
|
||||
|
||||
[[package]]
|
||||
name = "assert-json-diff"
|
||||
version = "2.0.2"
|
||||
@@ -167,6 +182,7 @@ name = "chat"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"axum",
|
||||
"chrono",
|
||||
"dotenvy",
|
||||
"futures",
|
||||
"jsonwebtoken",
|
||||
@@ -178,9 +194,24 @@ dependencies = [
|
||||
"tokio",
|
||||
"tokio-stream",
|
||||
"utoipa",
|
||||
"uuid",
|
||||
"wiremock",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "chrono"
|
||||
version = "0.4.44"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c673075a2e0e5f4a1dde27ce9dee1ea4558c7ffe648f576438a20ca1d2acc4b0"
|
||||
dependencies = [
|
||||
"iana-time-zone",
|
||||
"js-sys",
|
||||
"num-traits",
|
||||
"serde",
|
||||
"wasm-bindgen",
|
||||
"windows-link",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "cmake"
|
||||
version = "0.1.58"
|
||||
@@ -313,6 +344,12 @@ version = "1.0.7"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1"
|
||||
|
||||
[[package]]
|
||||
name = "foldhash"
|
||||
version = "0.1.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2"
|
||||
|
||||
[[package]]
|
||||
name = "form_urlencoded"
|
||||
version = "1.2.2"
|
||||
@@ -438,11 +475,24 @@ dependencies = [
|
||||
"cfg-if",
|
||||
"js-sys",
|
||||
"libc",
|
||||
"r-efi",
|
||||
"r-efi 5.3.0",
|
||||
"wasip2",
|
||||
"wasm-bindgen",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "getrandom"
|
||||
version = "0.4.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0de51e6874e94e7bf76d726fc5d13ba782deca734ff60d5bb2fb2607c7406555"
|
||||
dependencies = [
|
||||
"cfg-if",
|
||||
"libc",
|
||||
"r-efi 6.0.0",
|
||||
"wasip2",
|
||||
"wasip3",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "h2"
|
||||
version = "0.4.13"
|
||||
@@ -462,12 +512,27 @@ dependencies = [
|
||||
"tracing",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "hashbrown"
|
||||
version = "0.15.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1"
|
||||
dependencies = [
|
||||
"foldhash",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "hashbrown"
|
||||
version = "0.16.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "841d1cc9bed7f9236f321df977030373f4a4163ae1a7dbfe1a51a2c1a51d9100"
|
||||
|
||||
[[package]]
|
||||
name = "heck"
|
||||
version = "0.5.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea"
|
||||
|
||||
[[package]]
|
||||
name = "hermit-abi"
|
||||
version = "0.5.2"
|
||||
@@ -582,6 +647,30 @@ dependencies = [
|
||||
"windows-registry",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "iana-time-zone"
|
||||
version = "0.1.65"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e31bc9ad994ba00e440a8aa5c9ef0ec67d5cb5e5cb0cc7f8b744a35b389cc470"
|
||||
dependencies = [
|
||||
"android_system_properties",
|
||||
"core-foundation-sys",
|
||||
"iana-time-zone-haiku",
|
||||
"js-sys",
|
||||
"log",
|
||||
"wasm-bindgen",
|
||||
"windows-core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "iana-time-zone-haiku"
|
||||
version = "0.1.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f31827a206f56af32e590ba56d5d2d085f558508192593743f16b2306495269f"
|
||||
dependencies = [
|
||||
"cc",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "icu_collections"
|
||||
version = "2.2.0"
|
||||
@@ -664,6 +753,12 @@ dependencies = [
|
||||
"zerovec",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "id-arena"
|
||||
version = "2.3.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "3d3067d79b975e8844ca9eb072e16b31c3c1c36928edf9c6789548c524d0d954"
|
||||
|
||||
[[package]]
|
||||
name = "idna"
|
||||
version = "1.1.0"
|
||||
@@ -692,7 +787,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "45a8a2b9cb3e0b0c1803dbb0758ffac5de2f425b23c28f518faabd9d805342ff"
|
||||
dependencies = [
|
||||
"equivalent",
|
||||
"hashbrown",
|
||||
"hashbrown 0.16.1",
|
||||
"serde",
|
||||
"serde_core",
|
||||
]
|
||||
@@ -808,6 +903,12 @@ version = "1.5.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe"
|
||||
|
||||
[[package]]
|
||||
name = "leb128fmt"
|
||||
version = "0.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "09edd9e8b54e49e587e4f6295a7d29c3ea94d469cb40ab8ca70b288248a81db2"
|
||||
|
||||
[[package]]
|
||||
name = "libc"
|
||||
version = "0.2.184"
|
||||
@@ -995,6 +1096,16 @@ dependencies = [
|
||||
"zerocopy",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "prettyplease"
|
||||
version = "0.2.37"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"syn",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "proc-macro2"
|
||||
version = "1.0.106"
|
||||
@@ -1075,6 +1186,12 @@ version = "5.3.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f"
|
||||
|
||||
[[package]]
|
||||
name = "r-efi"
|
||||
version = "6.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf"
|
||||
|
||||
[[package]]
|
||||
name = "rand"
|
||||
version = "0.9.2"
|
||||
@@ -1348,6 +1465,12 @@ dependencies = [
|
||||
"libc",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "semver"
|
||||
version = "1.0.28"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8a7852d02fc848982e0c167ef163aaff9cd91dc640ba85e263cb1ce46fae51cd"
|
||||
|
||||
[[package]]
|
||||
name = "serde"
|
||||
version = "1.0.228"
|
||||
@@ -1774,6 +1897,12 @@ version = "1.0.24"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75"
|
||||
|
||||
[[package]]
|
||||
name = "unicode-xid"
|
||||
version = "0.2.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ebc1c04c71510c7f702b52b7c350734c9ff1295c464a03335b00bb84fc54f853"
|
||||
|
||||
[[package]]
|
||||
name = "untrusted"
|
||||
version = "0.7.1"
|
||||
@@ -1828,6 +1957,18 @@ dependencies = [
|
||||
"syn",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "uuid"
|
||||
version = "1.23.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "5ac8b6f42ead25368cf5b098aeb3dc8a1a2c05a3eee8a9a1a68c640edbfc79d9"
|
||||
dependencies = [
|
||||
"getrandom 0.4.2",
|
||||
"js-sys",
|
||||
"serde_core",
|
||||
"wasm-bindgen",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "walkdir"
|
||||
version = "2.5.0"
|
||||
@@ -1862,6 +2003,15 @@ dependencies = [
|
||||
"wit-bindgen",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "wasip3"
|
||||
version = "0.4.0+wasi-0.3.0-rc-2026-01-06"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "5428f8bf88ea5ddc08faddef2ac4a67e390b88186c703ce6dbd955e1c145aca5"
|
||||
dependencies = [
|
||||
"wit-bindgen",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "wasm-bindgen"
|
||||
version = "0.2.117"
|
||||
@@ -1917,6 +2067,28 @@ dependencies = [
|
||||
"unicode-ident",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "wasm-encoder"
|
||||
version = "0.244.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "990065f2fe63003fe337b932cfb5e3b80e0b4d0f5ff650e6985b1048f62c8319"
|
||||
dependencies = [
|
||||
"leb128fmt",
|
||||
"wasmparser",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "wasm-metadata"
|
||||
version = "0.244.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "bb0e353e6a2fbdc176932bbaab493762eb1255a7900fe0fea1a2f96c296cc909"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"indexmap",
|
||||
"wasm-encoder",
|
||||
"wasmparser",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "wasm-streams"
|
||||
version = "0.5.0"
|
||||
@@ -1930,6 +2102,18 @@ dependencies = [
|
||||
"web-sys",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "wasmparser"
|
||||
version = "0.244.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "47b807c72e1bac69382b3a6fb3dbe8ea4c0ed87ff5629b8685ae6b9a611028fe"
|
||||
dependencies = [
|
||||
"bitflags",
|
||||
"hashbrown 0.15.5",
|
||||
"indexmap",
|
||||
"semver",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "web-sys"
|
||||
version = "0.3.94"
|
||||
@@ -1968,6 +2152,41 @@ dependencies = [
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "windows-core"
|
||||
version = "0.62.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b8e83a14d34d0623b51dce9581199302a221863196a1dde71a7663a4c2be9deb"
|
||||
dependencies = [
|
||||
"windows-implement",
|
||||
"windows-interface",
|
||||
"windows-link",
|
||||
"windows-result",
|
||||
"windows-strings",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "windows-implement"
|
||||
version = "0.60.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "053e2e040ab57b9dc951b72c264860db7eb3b0200ba345b4e4c3b14f67855ddf"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "windows-interface"
|
||||
version = "0.59.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "3f316c4a2570ba26bbec722032c4099d8c8bc095efccdc15688708623367e358"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "windows-link"
|
||||
version = "0.2.1"
|
||||
@@ -2253,6 +2472,88 @@ name = "wit-bindgen"
|
||||
version = "0.51.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d7249219f66ced02969388cf2bb044a09756a083d0fab1e566056b04d9fbcaa5"
|
||||
dependencies = [
|
||||
"wit-bindgen-rust-macro",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "wit-bindgen-core"
|
||||
version = "0.51.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ea61de684c3ea68cb082b7a88508a8b27fcc8b797d738bfc99a82facf1d752dc"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"heck",
|
||||
"wit-parser",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "wit-bindgen-rust"
|
||||
version = "0.51.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b7c566e0f4b284dd6561c786d9cb0142da491f46a9fbed79ea69cdad5db17f21"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"heck",
|
||||
"indexmap",
|
||||
"prettyplease",
|
||||
"syn",
|
||||
"wasm-metadata",
|
||||
"wit-bindgen-core",
|
||||
"wit-component",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "wit-bindgen-rust-macro"
|
||||
version = "0.51.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0c0f9bfd77e6a48eccf51359e3ae77140a7f50b1e2ebfe62422d8afdaffab17a"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"prettyplease",
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
"wit-bindgen-core",
|
||||
"wit-bindgen-rust",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "wit-component"
|
||||
version = "0.244.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9d66ea20e9553b30172b5e831994e35fbde2d165325bec84fc43dbf6f4eb9cb2"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"bitflags",
|
||||
"indexmap",
|
||||
"log",
|
||||
"serde",
|
||||
"serde_derive",
|
||||
"serde_json",
|
||||
"wasm-encoder",
|
||||
"wasm-metadata",
|
||||
"wasmparser",
|
||||
"wit-parser",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "wit-parser"
|
||||
version = "0.244.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ecc8ac4bc1dc3381b7f59c34f00b67e18f910c2c0f50015669dde7def656a736"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"id-arena",
|
||||
"indexmap",
|
||||
"log",
|
||||
"semver",
|
||||
"serde",
|
||||
"serde_derive",
|
||||
"serde_json",
|
||||
"unicode-xid",
|
||||
"wasmparser",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "writeable"
|
||||
|
||||
@@ -20,3 +20,5 @@ dotenvy = "0.15"
|
||||
thiserror = "2.0.18"
|
||||
tokio-stream = "0.1"
|
||||
futures = "0.3"
|
||||
chrono = { version = "0.4.44", features = ["serde"] }
|
||||
uuid = { version = "1", features = ["v4", "serde"] }
|
||||
+169
@@ -0,0 +1,169 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
use utoipa::ToSchema;
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize, ToSchema)]
|
||||
pub struct ModelsResponse {
|
||||
pub models: Vec<ModelInfo>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize, ToSchema)]
|
||||
pub struct ModelInfo {
|
||||
pub name: String,
|
||||
pub family: Option<String>,
|
||||
pub parameter_size: Option<String>,
|
||||
pub quantization: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize, ToSchema)]
|
||||
pub struct LoadModelResponse {
|
||||
pub model: String,
|
||||
pub status: String,
|
||||
pub keep_alive: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize, ToSchema)]
|
||||
pub struct LoadModelBody {
|
||||
pub keep_alive: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize, ToSchema)]
|
||||
pub struct UnloadModelResponse {
|
||||
pub model: String,
|
||||
pub status: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize, ToSchema, Default)]
|
||||
pub struct BaseLLMRequest {
|
||||
pub model: String,
|
||||
|
||||
#[serde(default)]
|
||||
pub stream: bool,
|
||||
|
||||
pub temperature: Option<f32>,
|
||||
pub top_p: Option<f32>,
|
||||
|
||||
// Ollama-native
|
||||
pub top_k: Option<u32>,
|
||||
pub repeat_penalty: Option<f32>,
|
||||
pub seed: Option<i64>,
|
||||
|
||||
pub num_ctx: Option<u32>,
|
||||
pub num_predict: Option<u32>,
|
||||
|
||||
pub stop: Option<Vec<String>>,
|
||||
|
||||
pub keep_alive: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize, ToSchema)]
|
||||
pub struct CompletionRequest {
|
||||
#[serde(flatten)]
|
||||
pub base: BaseLLMRequest,
|
||||
|
||||
pub prompt: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize, ToSchema, PartialEq, Eq)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum CompletionObject {
|
||||
TextCompletion,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize, ToSchema, PartialEq, Eq)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum FinishReason {
|
||||
Stop,
|
||||
Length,
|
||||
ContentFilter,
|
||||
ToolCalls,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize, ToSchema)]
|
||||
pub struct CompletionResponse {
|
||||
pub id: String,
|
||||
pub object: CompletionObject,
|
||||
pub created: u64,
|
||||
pub model: String,
|
||||
pub choices: Vec<Choice>,
|
||||
pub usage: Usage,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize, ToSchema)]
|
||||
pub struct Choice {
|
||||
pub text: String,
|
||||
pub index: u32,
|
||||
pub finish_reason: FinishReason,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize, ToSchema)]
|
||||
pub struct Usage {
|
||||
pub prompt_tokens: u32,
|
||||
pub completion_tokens: u32,
|
||||
pub total_tokens: u32,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize, ToSchema)]
|
||||
pub struct CompletionChunk {
|
||||
pub id: String,
|
||||
pub object: String,
|
||||
pub choices: Vec<Choice>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize, ToSchema)]
|
||||
pub struct ChatRequest {
|
||||
#[serde(flatten)]
|
||||
pub base: BaseLLMRequest,
|
||||
|
||||
pub messages: Vec<Message>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize, ToSchema)]
|
||||
pub struct Message {
|
||||
pub role: Role,
|
||||
pub content: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize, ToSchema, PartialEq, Eq)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum Role {
|
||||
System,
|
||||
User,
|
||||
Assistant,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize, ToSchema)]
|
||||
pub struct ChatCompletionResponse {
|
||||
pub id: String,
|
||||
pub object: String,
|
||||
pub created: u64,
|
||||
pub model: String,
|
||||
pub choices: Vec<ChatChoice>,
|
||||
pub usage: Option<Usage>, // optional (Ollama may not always provide)
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize, ToSchema)]
|
||||
pub struct ChatChoice {
|
||||
pub index: u32,
|
||||
pub message: Message,
|
||||
pub finish_reason: FinishReason,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct ChatCompletionChunk {
|
||||
pub id: String,
|
||||
pub object: String,
|
||||
pub choices: Vec<ChatChunkChoice>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct ChatChunkChoice {
|
||||
pub index: u32,
|
||||
pub delta: ChatDelta,
|
||||
pub finish_reason: Option<FinishReason>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct ChatDelta {
|
||||
pub role: Option<Role>,
|
||||
pub content: Option<String>,
|
||||
}
|
||||
@@ -0,0 +1,2 @@
|
||||
pub mod api;
|
||||
pub mod ollama;
|
||||
@@ -0,0 +1,80 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::dto::api;
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
pub struct OllamaModels {
|
||||
pub models: Vec<OllamaModel>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
pub struct OllamaModel {
|
||||
pub name: String,
|
||||
|
||||
pub details: Option<OllamaModelDetails>,
|
||||
|
||||
pub size: Option<u64>,
|
||||
pub digest: Option<String>,
|
||||
pub modified_at: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
pub struct OllamaModelDetails {
|
||||
pub family: Option<String>,
|
||||
pub parameter_size: Option<String>,
|
||||
pub quantization_level: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
pub struct OllamaOptions {
|
||||
pub temperature: Option<f32>,
|
||||
pub top_p: Option<f32>,
|
||||
pub top_k: Option<u32>,
|
||||
pub repeat_penalty: Option<f32>,
|
||||
pub seed: Option<i64>,
|
||||
|
||||
pub num_ctx: Option<u32>,
|
||||
pub num_predict: Option<u32>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct OllamaGenerateRequest<'a> {
|
||||
pub model: &'a str,
|
||||
pub prompt: &'a str,
|
||||
pub stream: bool,
|
||||
pub options: OllamaOptions,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
pub struct OllamaGenerateResponse {
|
||||
pub model: String,
|
||||
pub created_at: Option<String>,
|
||||
pub response: String,
|
||||
pub done: bool,
|
||||
|
||||
#[serde(default)]
|
||||
pub context: Option<Vec<u64>>,
|
||||
|
||||
pub total_duration: Option<u64>,
|
||||
pub load_duration: Option<u64>,
|
||||
pub prompt_eval_count: Option<u32>,
|
||||
pub eval_count: Option<u32>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct OllamaChatRequest<'a> {
|
||||
pub model: &'a str,
|
||||
pub messages: &'a [api::Message],
|
||||
pub stream: bool,
|
||||
pub options: OllamaOptions,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct OllamaChatResponse {
|
||||
pub model: String,
|
||||
pub message: api::Message,
|
||||
pub done: bool,
|
||||
|
||||
pub prompt_eval_count: Option<u32>,
|
||||
pub eval_count: Option<u32>,
|
||||
}
|
||||
@@ -15,6 +15,9 @@ pub enum OllamaError {
|
||||
#[error("model '{0}' is not available — run `ollama pull {0}` first")]
|
||||
ModelNotFound(String),
|
||||
|
||||
#[error("keep_alive is required and cannot be empty")]
|
||||
MissingKeepAlive,
|
||||
|
||||
#[error(
|
||||
"invalid keep_alive format '{0}' — expected <number><unit> (e.g. 30s, 10m, 2h), a plain integer (seconds), or -1"
|
||||
)]
|
||||
|
||||
@@ -1,2 +1,3 @@
|
||||
pub mod dto;
|
||||
pub mod errors;
|
||||
pub mod providers;
|
||||
|
||||
+2
-1
@@ -1,11 +1,12 @@
|
||||
mod auth;
|
||||
mod dto;
|
||||
mod errors;
|
||||
mod openapi;
|
||||
mod providers;
|
||||
mod routes;
|
||||
mod state;
|
||||
|
||||
use crate::providers::ollama::OllamaProvider;
|
||||
use crate::providers::ollama::client::OllamaProvider;
|
||||
use crate::state::app_state::AppState;
|
||||
|
||||
use axum::Router;
|
||||
|
||||
+17
-17
@@ -1,18 +1,18 @@
|
||||
use utoipa::OpenApi;
|
||||
// use utoipa::OpenApi;
|
||||
|
||||
#[derive(OpenApi)]
|
||||
#[openapi(
|
||||
paths(
|
||||
crate::routes::v1::chat::completions
|
||||
),
|
||||
components(
|
||||
schemas(
|
||||
// add your request/response structs here later
|
||||
)
|
||||
),
|
||||
tags(
|
||||
(name = "chat", description = "Chat endpoints"),
|
||||
(name = "models", description = "Model management")
|
||||
)
|
||||
)]
|
||||
pub struct V1ApiDoc;
|
||||
// #[derive(OpenApi)]
|
||||
// #[openapi(
|
||||
// paths(
|
||||
// crate::routes::v1::chat::completions
|
||||
// ),
|
||||
// components(
|
||||
// schemas(
|
||||
// // add your request/response structs here later
|
||||
// )
|
||||
// ),
|
||||
// tags(
|
||||
// (name = "chat", description = "Chat endpoints"),
|
||||
// (name = "models", description = "Model management")
|
||||
// )
|
||||
// )]
|
||||
// pub struct V1ApiDoc;
|
||||
|
||||
@@ -1,400 +0,0 @@
|
||||
use crate::errors::OllamaError;
|
||||
use axum::response::sse::Event;
|
||||
use futures::StreamExt;
|
||||
use reqwest::Client;
|
||||
use serde_json::{Value, json};
|
||||
use tokio_stream::wrappers::ReceiverStream;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct OllamaProvider {
|
||||
pub client: Client,
|
||||
pub base_url: String,
|
||||
}
|
||||
|
||||
impl OllamaProvider {
|
||||
pub fn new(base_url: impl Into<String>) -> Self {
|
||||
Self {
|
||||
client: Client::new(),
|
||||
base_url: base_url.into(),
|
||||
}
|
||||
}
|
||||
|
||||
// ── private helpers ──────────────────────────────────────────────────────
|
||||
|
||||
fn build_options(body: &Value) -> Value {
|
||||
json!({
|
||||
"temperature": body.get("temperature"),
|
||||
"top_p": body.get("top_p"),
|
||||
"num_predict": body.get("max_tokens"),
|
||||
})
|
||||
}
|
||||
|
||||
async fn validate_model(&self, model: &str) -> Result<(), OllamaError> {
|
||||
let available = self.list_models().await?;
|
||||
let exists = available
|
||||
.get("models")
|
||||
.and_then(|m| m.as_array())
|
||||
.map(|arr| {
|
||||
arr.iter()
|
||||
.any(|m| m.get("name").and_then(|n| n.as_str()) == Some(model))
|
||||
})
|
||||
.unwrap_or(false);
|
||||
|
||||
if !exists {
|
||||
return Err(OllamaError::ModelNotFound(model.to_string()));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn parse_keep_alive(s: &str) -> Result<(), OllamaError> {
|
||||
let s = s.trim();
|
||||
|
||||
// Ollama also accepts plain integers (seconds) or "-1" (load forever)
|
||||
if s == "-1" || s.parse::<u64>().is_ok() {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
// Otherwise expect: <number><unit> e.g. "10m", "2h", "30s"
|
||||
let (num, unit) = s
|
||||
.find(|c: char| c.is_alphabetic())
|
||||
.map(|i| s.split_at(i))
|
||||
.ok_or_else(|| OllamaError::InvalidKeepAlive(s.to_string()))?;
|
||||
|
||||
num.parse::<u64>()
|
||||
.map_err(|_| OllamaError::InvalidKeepAlive(s.to_string()))?;
|
||||
|
||||
match unit {
|
||||
"s" | "m" | "h" => Ok(()),
|
||||
_ => Err(OllamaError::InvalidKeepAlive(s.to_string())),
|
||||
}
|
||||
}
|
||||
|
||||
fn extract_completion_params<'a>(
|
||||
&self,
|
||||
body: &'a Value,
|
||||
) -> Result<(&'a str, &'a str), OllamaError> {
|
||||
let prompt = body
|
||||
.get("prompt")
|
||||
.and_then(|v| v.as_str())
|
||||
.filter(|s| !s.trim().is_empty())
|
||||
.ok_or(OllamaError::MissingPrompt)?;
|
||||
|
||||
let model = body
|
||||
.get("model")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or(OllamaError::MissingModel)?;
|
||||
|
||||
Ok((prompt, model))
|
||||
}
|
||||
|
||||
fn format_completion_response(&self, res: &Value) -> Value {
|
||||
json!({
|
||||
"id": "cmpl-ollama",
|
||||
"object": "text_completion",
|
||||
"model": res.get("model"),
|
||||
"choices": [{
|
||||
"text": res.get("response"),
|
||||
"index": 0,
|
||||
"finish_reason": if res.get("done").and_then(|v| v.as_bool()).unwrap_or(false) {
|
||||
"stop"
|
||||
} else {
|
||||
"length"
|
||||
},
|
||||
}],
|
||||
"usage": {
|
||||
"prompt_tokens": res.get("prompt_eval_count"),
|
||||
"completion_tokens": res.get("eval_count"),
|
||||
"total_tokens": null,
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
fn extract_chat_params<'a>(
|
||||
&self,
|
||||
body: &'a Value,
|
||||
) -> Result<(&'a str, &'a Vec<Value>), OllamaError> {
|
||||
let model = body
|
||||
.get("model")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("llama3");
|
||||
|
||||
let messages = body
|
||||
.get("messages")
|
||||
.and_then(|v| v.as_array())
|
||||
.filter(|arr| !arr.is_empty())
|
||||
.ok_or(OllamaError::MissingMessages)?;
|
||||
|
||||
let has_user_msg = messages
|
||||
.iter()
|
||||
.any(|m| m.get("role").and_then(|r| r.as_str()) == Some("user"));
|
||||
|
||||
if !has_user_msg {
|
||||
return Err(OllamaError::MissingMessages);
|
||||
}
|
||||
|
||||
Ok((model, messages))
|
||||
}
|
||||
|
||||
fn format_chat_response(&self, res: &Value) -> Value {
|
||||
json!({
|
||||
"id": "chatcmpl-ollama",
|
||||
"object": "chat.completion",
|
||||
"model": res.get("model"),
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"message": {
|
||||
"role": res.get("message").and_then(|m| m.get("role")),
|
||||
"content": res.get("message").and_then(|m| m.get("content")),
|
||||
},
|
||||
"finish_reason": res
|
||||
.get("done_reason")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("stop"),
|
||||
}],
|
||||
"usage": {
|
||||
"prompt_tokens": res.get("prompt_eval_count"),
|
||||
"completion_tokens": res.get("eval_count"),
|
||||
"total_tokens": res.get("prompt_eval_count")
|
||||
.and_then(|p| p.as_u64())
|
||||
.zip(res.get("eval_count").and_then(|e| e.as_u64()))
|
||||
.map(|(p, e)| p + e),
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// ── public endpoints ─────────────────────────────────────────────────────
|
||||
|
||||
pub async fn list_models(&self) -> Result<Value, OllamaError> {
|
||||
let url = format!("{}/api/tags", self.base_url);
|
||||
let res = self.client.get(url).send().await?.json::<Value>().await?;
|
||||
Ok(res)
|
||||
}
|
||||
|
||||
pub async fn load_model(
|
||||
&self,
|
||||
model: &str,
|
||||
keep_alive: Option<&str>,
|
||||
) -> Result<Value, OllamaError> {
|
||||
self.validate_model(model).await?;
|
||||
|
||||
let keep_alive = keep_alive.unwrap_or("5m");
|
||||
Self::parse_keep_alive(keep_alive)?; // ← validated before any network call
|
||||
|
||||
let payload = json!({
|
||||
"model": model,
|
||||
"prompt": "",
|
||||
"keep_alive": keep_alive,
|
||||
"stream": false,
|
||||
});
|
||||
|
||||
let res = self
|
||||
.client
|
||||
.post(format!("{}/api/generate", self.base_url))
|
||||
.json(&payload)
|
||||
.send()
|
||||
.await?
|
||||
.json::<Value>()
|
||||
.await?;
|
||||
|
||||
Ok(json!({
|
||||
"model": res.get("model"),
|
||||
"status": "loaded",
|
||||
"keep_alive": keep_alive,
|
||||
}))
|
||||
}
|
||||
|
||||
pub async fn unload_model(&self, model: &str) -> Result<Value, OllamaError> {
|
||||
self.validate_model(model).await?;
|
||||
|
||||
let payload = json!({
|
||||
"model": model,
|
||||
"prompt": "",
|
||||
"keep_alive": "0",
|
||||
"stream": false,
|
||||
});
|
||||
|
||||
let res = self
|
||||
.client
|
||||
.post(format!("{}/api/generate", self.base_url))
|
||||
.json(&payload)
|
||||
.send()
|
||||
.await?
|
||||
.json::<Value>()
|
||||
.await?;
|
||||
|
||||
Ok(json!({
|
||||
"model": res.get("model"),
|
||||
"status": "unloaded",
|
||||
}))
|
||||
}
|
||||
|
||||
pub async fn completions(&self, body: Value) -> Result<Value, OllamaError> {
|
||||
let (prompt, model) = self.extract_completion_params(&body)?;
|
||||
self.validate_model(model).await?;
|
||||
|
||||
let payload = json!({
|
||||
"model": model,
|
||||
"prompt": prompt,
|
||||
"stream": false,
|
||||
"options": Self::build_options(&body),
|
||||
});
|
||||
|
||||
let res = self
|
||||
.client
|
||||
.post(format!("{}/api/generate", self.base_url))
|
||||
.json(&payload)
|
||||
.send()
|
||||
.await?
|
||||
.json::<Value>()
|
||||
.await?;
|
||||
|
||||
Ok(self.format_completion_response(&res))
|
||||
}
|
||||
|
||||
pub async fn completions_stream(
|
||||
&self,
|
||||
body: Value,
|
||||
) -> Result<ReceiverStream<Result<Event, OllamaError>>, OllamaError> {
|
||||
let (prompt, model) = self.extract_completion_params(&body)?;
|
||||
self.validate_model(model).await?;
|
||||
|
||||
let payload = json!({
|
||||
"model": model,
|
||||
"prompt": prompt,
|
||||
"stream": true,
|
||||
"options": Self::build_options(&body),
|
||||
});
|
||||
|
||||
let mut byte_stream = self
|
||||
.client
|
||||
.post(format!("{}/api/generate", self.base_url))
|
||||
.json(&payload)
|
||||
.send()
|
||||
.await?
|
||||
.bytes_stream();
|
||||
|
||||
let (tx, rx) = tokio::sync::mpsc::channel(32);
|
||||
|
||||
tokio::spawn(async move {
|
||||
while let Some(chunk) = byte_stream.next().await {
|
||||
let chunk = match chunk {
|
||||
Ok(b) => b,
|
||||
Err(e) => {
|
||||
let _ = tx.send(Err(OllamaError::Http(e))).await;
|
||||
break;
|
||||
}
|
||||
};
|
||||
|
||||
if let Ok(json) = serde_json::from_slice::<Value>(&chunk) {
|
||||
let token = json.get("response").and_then(|v| v.as_str()).unwrap_or("");
|
||||
let done = json.get("done").and_then(|v| v.as_bool()).unwrap_or(false);
|
||||
|
||||
// OpenAI-compatible SSE chunk
|
||||
let event_data = serde_json::to_string(&json!({
|
||||
"id": "cmpl-ollama",
|
||||
"object": "text_completion",
|
||||
"choices": [{ "text": token, "index": 0, "finish_reason": null }],
|
||||
}))
|
||||
.unwrap_or_default();
|
||||
|
||||
let _ = tx.send(Ok(Event::default().data(event_data))).await;
|
||||
|
||||
if done {
|
||||
// Final [DONE] sentinel — matches OpenAI streaming protocol
|
||||
let _ = tx.send(Ok(Event::default().data("[DONE]"))).await;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
Ok(ReceiverStream::new(rx))
|
||||
}
|
||||
|
||||
pub async fn chat_completions(&self, body: Value) -> Result<Value, OllamaError> {
|
||||
let (model, messages) = self.extract_chat_params(&body)?;
|
||||
self.validate_model(model).await?;
|
||||
|
||||
let payload = json!({
|
||||
"model": model,
|
||||
"messages": messages,
|
||||
"stream": false,
|
||||
"options": Self::build_options(&body),
|
||||
});
|
||||
|
||||
let res = self
|
||||
.client
|
||||
.post(format!("{}/api/chat", self.base_url))
|
||||
.json(&payload)
|
||||
.send()
|
||||
.await?
|
||||
.json::<Value>()
|
||||
.await?;
|
||||
|
||||
Ok(self.format_chat_response(&res))
|
||||
}
|
||||
|
||||
pub async fn chat_completions_stream(
|
||||
&self,
|
||||
body: Value,
|
||||
) -> Result<ReceiverStream<Result<Event, OllamaError>>, OllamaError> {
|
||||
let (model, messages) = self.extract_chat_params(&body)?;
|
||||
self.validate_model(model).await?;
|
||||
|
||||
let payload = json!({
|
||||
"model": model,
|
||||
"messages": messages,
|
||||
"stream": true,
|
||||
"options": Self::build_options(&body),
|
||||
});
|
||||
|
||||
let mut byte_stream = self
|
||||
.client
|
||||
.post(format!("{}/api/chat", self.base_url))
|
||||
.json(&payload)
|
||||
.send()
|
||||
.await?
|
||||
.bytes_stream();
|
||||
|
||||
let (tx, rx) = tokio::sync::mpsc::channel(32);
|
||||
|
||||
tokio::spawn(async move {
|
||||
while let Some(chunk) = byte_stream.next().await {
|
||||
let chunk = match chunk {
|
||||
Ok(b) => b,
|
||||
Err(e) => {
|
||||
let _ = tx.send(Err(OllamaError::Http(e))).await;
|
||||
break;
|
||||
}
|
||||
};
|
||||
|
||||
if let Ok(json) = serde_json::from_slice::<Value>(&chunk) {
|
||||
let done = json.get("done").and_then(|v| v.as_bool()).unwrap_or(false);
|
||||
|
||||
let event_data = serde_json::to_string(&json!({
|
||||
"id": "chatcmpl-ollama",
|
||||
"object": "chat.completion.chunk",
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"delta": {
|
||||
"role": json.get("message").and_then(|m| m.get("role")),
|
||||
"content": json.get("message").and_then(|m| m.get("content")),
|
||||
},
|
||||
"finish_reason": if done { json!("stop") } else { json!(null) },
|
||||
}],
|
||||
}))
|
||||
.unwrap_or_default();
|
||||
|
||||
let _ = tx.send(Ok(Event::default().data(event_data))).await;
|
||||
|
||||
if done {
|
||||
let _ = tx.send(Ok(Event::default().data("[DONE]"))).await;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
Ok(ReceiverStream::new(rx))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,432 @@
|
||||
use crate::dto::{api, ollama};
|
||||
use crate::errors::OllamaError;
|
||||
use axum::response::sse::Event;
|
||||
use futures::StreamExt;
|
||||
use reqwest::Client;
|
||||
use serde_json::json;
|
||||
use tokio_stream::wrappers::ReceiverStream;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct OllamaProvider {
|
||||
pub client: Client,
|
||||
pub base_url: String,
|
||||
}
|
||||
|
||||
impl OllamaProvider {
|
||||
pub fn new(base_url: impl Into<String>) -> Self {
|
||||
Self {
|
||||
client: Client::new(),
|
||||
base_url: base_url.into(),
|
||||
}
|
||||
}
|
||||
|
||||
// ── private helpers ──────────────────────────────────────────────────────
|
||||
|
||||
async fn model_exists(&self, model: &str) -> Result<bool, OllamaError> {
|
||||
let url = format!("{}/api/tags", self.base_url);
|
||||
|
||||
let res = self
|
||||
.client
|
||||
.get(url)
|
||||
.send()
|
||||
.await?
|
||||
.json::<ollama::OllamaModels>()
|
||||
.await?;
|
||||
|
||||
Ok(res.models.iter().any(|m| m.name == model))
|
||||
}
|
||||
|
||||
fn has_user_message(&self, messages: &[api::Message]) -> bool {
|
||||
messages.iter().any(|m| matches!(m.role, api::Role::User))
|
||||
}
|
||||
|
||||
fn extract_completion_params<'a>(
|
||||
&self,
|
||||
body: &'a api::CompletionRequest,
|
||||
) -> Result<(&'a str, &'a str), OllamaError> {
|
||||
let prompt = body.prompt.trim();
|
||||
|
||||
let model = body.base.model.as_str();
|
||||
|
||||
Ok((prompt, model))
|
||||
}
|
||||
|
||||
fn extract_chat_params<'a>(
|
||||
&self,
|
||||
body: &'a api::ChatRequest,
|
||||
) -> Result<(&'a [api::Message], &'a str), OllamaError> {
|
||||
let model = body.base.model.as_str();
|
||||
|
||||
Ok((&body.messages, model))
|
||||
}
|
||||
|
||||
pub fn parse_keep_alive(&self, s: &str) -> Result<(), OllamaError> {
|
||||
let s = s.trim();
|
||||
|
||||
if s == "-1" || s.parse::<u64>().is_ok() {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let split = s
|
||||
.find(|c: char| c.is_alphabetic())
|
||||
.ok_or_else(|| OllamaError::InvalidKeepAlive(s.to_string()))?;
|
||||
|
||||
let (num, unit) = s.split_at(split);
|
||||
|
||||
num.parse::<u64>()
|
||||
.map_err(|_| OllamaError::InvalidKeepAlive(s.to_string()))?;
|
||||
|
||||
match unit {
|
||||
"s" | "m" | "h" => Ok(()),
|
||||
_ => Err(OllamaError::InvalidKeepAlive(s.to_string())),
|
||||
}
|
||||
}
|
||||
|
||||
// // ── public endpoints ─────────────────────────────────────────────────────
|
||||
|
||||
pub async fn list_models(&self) -> Result<api::ModelsResponse, OllamaError> {
|
||||
let url = format!("{}/api/tags", self.base_url);
|
||||
|
||||
let res = self
|
||||
.client
|
||||
.get(url)
|
||||
.send()
|
||||
.await?
|
||||
.json::<ollama::OllamaModels>()
|
||||
.await?;
|
||||
|
||||
let models = res.models.into_iter().map(api::ModelInfo::from).collect();
|
||||
|
||||
Ok(api::ModelsResponse { models })
|
||||
}
|
||||
|
||||
pub async fn load_model(
|
||||
&self,
|
||||
model: &str,
|
||||
keep_alive: Option<&str>,
|
||||
) -> Result<api::LoadModelResponse, OllamaError> {
|
||||
let url = format!("{}/api/generate", self.base_url);
|
||||
|
||||
let keep_alive = keep_alive.ok_or(OllamaError::MissingKeepAlive)?;
|
||||
|
||||
self.parse_keep_alive(keep_alive)?;
|
||||
|
||||
let exists = self.model_exists(model).await?;
|
||||
if !exists {
|
||||
return Err(OllamaError::ModelNotFound(model.to_string()));
|
||||
}
|
||||
|
||||
let payload = json!({
|
||||
"model": model,
|
||||
"prompt": "",
|
||||
"keep_alive": keep_alive,
|
||||
"stream": false,
|
||||
});
|
||||
|
||||
let _res = self
|
||||
.client
|
||||
.post(url)
|
||||
.json(&payload)
|
||||
.send()
|
||||
.await?
|
||||
.text()
|
||||
.await?;
|
||||
|
||||
Ok(api::LoadModelResponse {
|
||||
model: model.to_string(),
|
||||
status: "loaded".to_string(),
|
||||
keep_alive: keep_alive.to_string(),
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn unload_model(&self, model: &str) -> Result<api::UnloadModelResponse, OllamaError> {
|
||||
let url = format!("{}/api/generate", self.base_url);
|
||||
|
||||
let exists = self.model_exists(model).await?;
|
||||
if !exists {
|
||||
return Err(OllamaError::ModelNotFound(model.to_string()));
|
||||
}
|
||||
|
||||
let payload = json!({
|
||||
"model": model,
|
||||
"prompt": "",
|
||||
"keep_alive": "0",
|
||||
"stream": false,
|
||||
});
|
||||
|
||||
let _res = self
|
||||
.client
|
||||
.post(url)
|
||||
.json(&payload)
|
||||
.send()
|
||||
.await?
|
||||
.text()
|
||||
.await?;
|
||||
|
||||
Ok(api::UnloadModelResponse {
|
||||
model: model.to_string(),
|
||||
status: "unloaded".to_string(),
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn completions(
|
||||
&self,
|
||||
body: &api::CompletionRequest,
|
||||
) -> Result<api::CompletionResponse, OllamaError> {
|
||||
let url = format!("{}/api/generate", self.base_url);
|
||||
|
||||
let (prompt, model) = self.extract_completion_params(body)?;
|
||||
|
||||
if prompt.is_empty() {
|
||||
return Err(OllamaError::MissingPrompt);
|
||||
}
|
||||
|
||||
if model.is_empty() {
|
||||
return Err(OllamaError::MissingModel);
|
||||
}
|
||||
|
||||
let exists = self.model_exists(model).await?;
|
||||
if !exists {
|
||||
return Err(OllamaError::ModelNotFound(model.to_string()));
|
||||
}
|
||||
|
||||
let options = ollama::OllamaOptions::from(&body.base);
|
||||
|
||||
let payload = ollama::OllamaGenerateRequest {
|
||||
model,
|
||||
prompt,
|
||||
stream: false,
|
||||
options,
|
||||
};
|
||||
|
||||
let res = self
|
||||
.client
|
||||
.post(url)
|
||||
.json(&payload)
|
||||
.send()
|
||||
.await?
|
||||
.json::<ollama::OllamaGenerateResponse>()
|
||||
.await?;
|
||||
|
||||
Ok(api::CompletionResponse::from(res))
|
||||
}
|
||||
|
||||
pub async fn completions_stream(
|
||||
&self,
|
||||
body: &api::CompletionRequest,
|
||||
) -> Result<ReceiverStream<Result<Event, OllamaError>>, OllamaError> {
|
||||
let url = format!("{}/api/generate", self.base_url);
|
||||
|
||||
let (prompt, model) = self.extract_completion_params(body)?;
|
||||
|
||||
if prompt.is_empty() {
|
||||
return Err(OllamaError::MissingPrompt);
|
||||
}
|
||||
|
||||
if model.is_empty() {
|
||||
return Err(OllamaError::MissingModel);
|
||||
}
|
||||
|
||||
let exists = self.model_exists(model).await?;
|
||||
if !exists {
|
||||
return Err(OllamaError::ModelNotFound(model.to_string()));
|
||||
}
|
||||
|
||||
let options = ollama::OllamaOptions::from(&body.base);
|
||||
|
||||
let payload = ollama::OllamaGenerateRequest {
|
||||
model,
|
||||
prompt,
|
||||
stream: true,
|
||||
options,
|
||||
};
|
||||
|
||||
let mut byte_stream = self
|
||||
.client
|
||||
.post(url)
|
||||
.json(&payload)
|
||||
.send()
|
||||
.await?
|
||||
.bytes_stream();
|
||||
|
||||
let (tx, rx) = tokio::sync::mpsc::channel(32);
|
||||
|
||||
tokio::spawn(async move {
|
||||
while let Some(chunk) = byte_stream.next().await {
|
||||
let chunk = match chunk {
|
||||
Ok(b) => b,
|
||||
Err(e) => {
|
||||
let _ = tx.send(Err(OllamaError::Http(e))).await;
|
||||
break;
|
||||
}
|
||||
};
|
||||
|
||||
// 🔥 IMPORTANT: typed deserialization
|
||||
let parsed: ollama::OllamaGenerateResponse = match serde_json::from_slice(&chunk) {
|
||||
Ok(v) => v,
|
||||
Err(_) => continue,
|
||||
};
|
||||
|
||||
// map → OpenAI chunk
|
||||
let event_data = serde_json::to_string(&api::CompletionChunk {
|
||||
id: "cmpl-ollama".to_string(),
|
||||
object: "text_completion".to_string(),
|
||||
choices: vec![api::Choice {
|
||||
text: parsed.response,
|
||||
index: 0,
|
||||
finish_reason: if parsed.done {
|
||||
api::FinishReason::Stop
|
||||
} else {
|
||||
api::FinishReason::Length
|
||||
},
|
||||
}],
|
||||
})
|
||||
.unwrap_or_default();
|
||||
|
||||
let _ = tx.send(Ok(Event::default().data(event_data))).await;
|
||||
|
||||
if parsed.done {
|
||||
let _ = tx.send(Ok(Event::default().data("[DONE]"))).await;
|
||||
break;
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
Ok(ReceiverStream::new(rx))
|
||||
}
|
||||
|
||||
pub async fn chat_completions(
|
||||
&self,
|
||||
body: &api::ChatRequest,
|
||||
) -> Result<api::ChatCompletionResponse, OllamaError> {
|
||||
let url = format!("{}/api/chat", self.base_url);
|
||||
|
||||
let (messages, model) = self.extract_chat_params(body)?;
|
||||
|
||||
if body.messages.is_empty() {
|
||||
return Err(OllamaError::MissingMessages);
|
||||
}
|
||||
|
||||
if !self.has_user_message(&body.messages) {
|
||||
return Err(OllamaError::MissingMessages);
|
||||
}
|
||||
|
||||
if model.is_empty() {
|
||||
return Err(OllamaError::MissingModel);
|
||||
}
|
||||
|
||||
let exists = self.model_exists(model).await?;
|
||||
if !exists {
|
||||
return Err(OllamaError::ModelNotFound(model.to_string()));
|
||||
}
|
||||
|
||||
// let options = ollama::OllamaOptions::from(body);
|
||||
let options = ollama::OllamaOptions::from(&body.base);
|
||||
|
||||
let payload = ollama::OllamaChatRequest {
|
||||
model,
|
||||
messages,
|
||||
stream: false,
|
||||
options,
|
||||
};
|
||||
|
||||
let res = self
|
||||
.client
|
||||
.post(url)
|
||||
.json(&payload)
|
||||
.send()
|
||||
.await?
|
||||
.json::<ollama::OllamaChatResponse>()
|
||||
.await?;
|
||||
|
||||
Ok(api::ChatCompletionResponse::from(res))
|
||||
}
|
||||
|
||||
pub async fn chat_completions_stream(
|
||||
&self,
|
||||
body: &api::ChatRequest,
|
||||
) -> Result<ReceiverStream<Result<Event, OllamaError>>, OllamaError> {
|
||||
let url = format!("{}/api/chat", self.base_url);
|
||||
|
||||
let (messages, model) = self.extract_chat_params(body)?;
|
||||
|
||||
if messages.is_empty() {
|
||||
return Err(OllamaError::MissingMessages);
|
||||
}
|
||||
|
||||
if model.is_empty() {
|
||||
return Err(OllamaError::MissingModel);
|
||||
}
|
||||
|
||||
let exists = self.model_exists(model).await?;
|
||||
if !exists {
|
||||
return Err(OllamaError::ModelNotFound(model.to_string()));
|
||||
}
|
||||
|
||||
let options = ollama::OllamaOptions::from(&body.base);
|
||||
|
||||
let payload = ollama::OllamaChatRequest {
|
||||
model,
|
||||
messages,
|
||||
stream: true,
|
||||
options,
|
||||
};
|
||||
|
||||
let mut byte_stream = self
|
||||
.client
|
||||
.post(url)
|
||||
.json(&payload)
|
||||
.send()
|
||||
.await?
|
||||
.bytes_stream();
|
||||
|
||||
let (tx, rx) = tokio::sync::mpsc::channel(32);
|
||||
|
||||
tokio::spawn(async move {
|
||||
let stream_id = format!("chatcmpl-{}", uuid::Uuid::new_v4());
|
||||
|
||||
while let Some(chunk) = byte_stream.next().await {
|
||||
let chunk = match chunk {
|
||||
Ok(b) => b,
|
||||
Err(e) => {
|
||||
let _ = tx.send(Err(OllamaError::Http(e))).await;
|
||||
break;
|
||||
}
|
||||
};
|
||||
|
||||
let parsed: ollama::OllamaChatResponse = match serde_json::from_slice(&chunk) {
|
||||
Ok(v) => v,
|
||||
Err(_) => continue,
|
||||
};
|
||||
|
||||
let event = api::ChatCompletionChunk {
|
||||
id: stream_id.clone(),
|
||||
object: "chat.completion.chunk".to_string(),
|
||||
choices: vec![api::ChatChunkChoice {
|
||||
index: 0,
|
||||
delta: api::ChatDelta {
|
||||
role: Some(parsed.message.role),
|
||||
content: Some(parsed.message.content),
|
||||
},
|
||||
finish_reason: if parsed.done {
|
||||
Some(api::FinishReason::Stop)
|
||||
} else {
|
||||
None
|
||||
},
|
||||
}],
|
||||
};
|
||||
|
||||
let event_data = serde_json::to_string(&event).unwrap_or_default();
|
||||
|
||||
let _ = tx.send(Ok(Event::default().data(event_data))).await;
|
||||
|
||||
if parsed.done {
|
||||
let _ = tx.send(Ok(Event::default().data("[DONE]"))).await;
|
||||
break;
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
Ok(ReceiverStream::new(rx))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,86 @@
|
||||
use crate::dto::{api, ollama};
|
||||
|
||||
use chrono::Utc;
|
||||
use uuid::Uuid;
|
||||
|
||||
impl From<ollama::OllamaModel> for api::ModelInfo {
|
||||
fn from(m: ollama::OllamaModel) -> Self {
|
||||
Self {
|
||||
name: m.name,
|
||||
|
||||
family: m.details.as_ref().and_then(|d| d.family.clone()),
|
||||
parameter_size: m.details.as_ref().and_then(|d| d.parameter_size.clone()),
|
||||
quantization: m
|
||||
.details
|
||||
.as_ref()
|
||||
.and_then(|d| d.quantization_level.clone()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<ollama::OllamaGenerateResponse> for api::CompletionResponse {
|
||||
fn from(res: ollama::OllamaGenerateResponse) -> Self {
|
||||
Self {
|
||||
id: Uuid::new_v4().to_string(),
|
||||
object: api::CompletionObject::TextCompletion,
|
||||
model: res.model,
|
||||
created: Utc::now().timestamp() as u64,
|
||||
|
||||
choices: vec![api::Choice {
|
||||
text: res.response,
|
||||
index: 0,
|
||||
finish_reason: api::FinishReason::Stop,
|
||||
}],
|
||||
|
||||
usage: api::Usage {
|
||||
prompt_tokens: res.prompt_eval_count.unwrap_or(0),
|
||||
completion_tokens: res.eval_count.unwrap_or(0),
|
||||
total_tokens: res.prompt_eval_count.unwrap_or(0) + res.eval_count.unwrap_or(0),
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<&api::BaseLLMRequest> for ollama::OllamaOptions {
|
||||
fn from(base: &api::BaseLLMRequest) -> Self {
|
||||
Self {
|
||||
temperature: base.temperature,
|
||||
top_p: base.top_p,
|
||||
top_k: base.top_k,
|
||||
repeat_penalty: base.repeat_penalty,
|
||||
seed: base.seed,
|
||||
num_ctx: base.num_ctx,
|
||||
num_predict: base.num_predict,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<ollama::OllamaChatResponse> for api::ChatCompletionResponse {
|
||||
fn from(res: ollama::OllamaChatResponse) -> Self {
|
||||
let prompt_tokens = res.prompt_eval_count.unwrap_or(0);
|
||||
let completion_tokens = res.eval_count.unwrap_or(0);
|
||||
|
||||
Self {
|
||||
id: Uuid::new_v4().to_string(),
|
||||
object: "chat.completion".to_string(),
|
||||
created: Utc::now().timestamp() as u64,
|
||||
model: res.model,
|
||||
|
||||
choices: vec![api::ChatChoice {
|
||||
index: 0,
|
||||
message: res.message,
|
||||
finish_reason: if res.done {
|
||||
api::FinishReason::Stop
|
||||
} else {
|
||||
api::FinishReason::Length
|
||||
},
|
||||
}],
|
||||
|
||||
usage: Some(api::Usage {
|
||||
prompt_tokens,
|
||||
completion_tokens,
|
||||
total_tokens: prompt_tokens + completion_tokens,
|
||||
}),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,2 @@
|
||||
pub mod client;
|
||||
pub mod mapper;
|
||||
+13
-19
@@ -6,8 +6,8 @@ use axum::{
|
||||
sse::{KeepAlive, Sse},
|
||||
},
|
||||
};
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::dto::api;
|
||||
use crate::errors::OllamaError;
|
||||
use crate::state::app_state::AppState;
|
||||
|
||||
@@ -20,17 +20,12 @@ use crate::state::app_state::AppState;
|
||||
)]
|
||||
pub async fn completions(
|
||||
State(state): State<AppState>,
|
||||
Json(body): Json<Value>,
|
||||
Json(body): Json<api::CompletionRequest>,
|
||||
) -> Result<Response, (axum::http::StatusCode, String)> {
|
||||
let wants_stream = body
|
||||
.get("stream")
|
||||
.and_then(|v| v.as_bool())
|
||||
.unwrap_or(false);
|
||||
|
||||
if wants_stream {
|
||||
if body.base.stream {
|
||||
let stream = state
|
||||
.ollama
|
||||
.completions_stream(body)
|
||||
.completions_stream(&body)
|
||||
.await
|
||||
.map_err(ollama_err)?;
|
||||
|
||||
@@ -38,7 +33,7 @@ pub async fn completions(
|
||||
.keep_alive(KeepAlive::default())
|
||||
.into_response())
|
||||
} else {
|
||||
let response = state.ollama.completions(body).await.map_err(ollama_err)?;
|
||||
let response = state.ollama.completions(&body).await.map_err(ollama_err)?;
|
||||
|
||||
Ok(Json(response).into_response())
|
||||
}
|
||||
@@ -46,17 +41,12 @@ pub async fn completions(
|
||||
|
||||
pub async fn chat_completions(
|
||||
State(state): State<AppState>,
|
||||
Json(body): Json<Value>,
|
||||
Json(body): Json<api::ChatRequest>,
|
||||
) -> Result<Response, (axum::http::StatusCode, String)> {
|
||||
let wants_stream = body
|
||||
.get("stream")
|
||||
.and_then(|v| v.as_bool())
|
||||
.unwrap_or(false);
|
||||
|
||||
if wants_stream {
|
||||
if body.base.stream {
|
||||
let stream = state
|
||||
.ollama
|
||||
.chat_completions_stream(body)
|
||||
.chat_completions_stream(&body)
|
||||
.await
|
||||
.map_err(ollama_err)?;
|
||||
|
||||
@@ -66,7 +56,7 @@ pub async fn chat_completions(
|
||||
} else {
|
||||
let response = state
|
||||
.ollama
|
||||
.chat_completions(body)
|
||||
.chat_completions(&body)
|
||||
.await
|
||||
.map_err(ollama_err)?;
|
||||
|
||||
@@ -88,6 +78,10 @@ fn ollama_err(e: OllamaError) -> (axum::http::StatusCode, String) {
|
||||
axum::http::StatusCode::UNPROCESSABLE_ENTITY,
|
||||
format!("model '{m}' is not available — run `ollama pull {m}` first"),
|
||||
),
|
||||
OllamaError::MissingKeepAlive => (
|
||||
axum::http::StatusCode::BAD_REQUEST,
|
||||
"keep alive is required and cannot be empty".to_string(),
|
||||
),
|
||||
OllamaError::InvalidKeepAlive(v) => (
|
||||
axum::http::StatusCode::BAD_REQUEST,
|
||||
format!("invalid keep_alive '{v}'"),
|
||||
|
||||
@@ -6,9 +6,9 @@ use crate::auth::middleware::auth_middleware;
|
||||
use crate::state::app_state::AppState;
|
||||
use axum::{Router, middleware, routing::get, routing::post};
|
||||
|
||||
fn public_router() -> Router<AppState> {
|
||||
Router::new().route("/openapi.json", get(openapi::openapi_json))
|
||||
}
|
||||
// fn public_router() -> Router<AppState> {
|
||||
// Router::new().route("/openapi.json", get(openapi::openapi_json))
|
||||
// }
|
||||
|
||||
pub fn protected_router() -> Router<AppState> {
|
||||
Router::new()
|
||||
@@ -21,6 +21,6 @@ pub fn protected_router() -> Router<AppState> {
|
||||
|
||||
pub fn router() -> Router<AppState> {
|
||||
Router::new()
|
||||
.merge(public_router())
|
||||
// .merge(public_router())
|
||||
.merge(protected_router().layer(middleware::from_fn(auth_middleware)))
|
||||
}
|
||||
|
||||
+20
-21
@@ -2,19 +2,13 @@ use axum::{
|
||||
Json,
|
||||
extract::{Path, State},
|
||||
};
|
||||
use serde::Deserialize;
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::dto::api;
|
||||
use crate::{errors::OllamaError, state::app_state::AppState};
|
||||
|
||||
#[derive(Deserialize)]
|
||||
pub struct LoadModelBody {
|
||||
pub keep_alive: Option<String>,
|
||||
}
|
||||
|
||||
pub async fn list_models(
|
||||
State(state): State<AppState>,
|
||||
) -> Result<Json<Value>, (axum::http::StatusCode, String)> {
|
||||
) -> Result<Json<api::ModelsResponse>, (axum::http::StatusCode, String)> {
|
||||
match state.ollama.list_models().await {
|
||||
Ok(models) => Ok(Json(models)),
|
||||
Err(e) => Err(ollama_err(e)),
|
||||
@@ -24,27 +18,28 @@ pub async fn list_models(
|
||||
pub async fn load_model(
|
||||
State(state): State<AppState>,
|
||||
Path(model): Path<String>,
|
||||
Json(body): Json<LoadModelBody>,
|
||||
) -> Result<Json<Value>, (axum::http::StatusCode, String)> {
|
||||
match state
|
||||
Json(body): Json<api::LoadModelBody>,
|
||||
) -> Result<Json<api::LoadModelResponse>, (axum::http::StatusCode, String)> {
|
||||
let response = state
|
||||
.ollama
|
||||
.load_model(&model, body.keep_alive.as_deref())
|
||||
.await
|
||||
{
|
||||
Ok(response) => Ok(Json(response)),
|
||||
Err(e) => Err(ollama_err(e)),
|
||||
}
|
||||
.map_err(ollama_err)?;
|
||||
|
||||
Ok(Json(response))
|
||||
}
|
||||
|
||||
pub async fn unload_model(
|
||||
State(state): State<AppState>,
|
||||
Path(model): Path<String>,
|
||||
) -> Result<Json<Value>, (axum::http::StatusCode, String)> {
|
||||
match state.ollama.unload_model(&model).await {
|
||||
// ← correct method
|
||||
Ok(response) => Ok(Json(response)),
|
||||
Err(e) => Err(ollama_err(e)),
|
||||
}
|
||||
) -> Result<Json<api::UnloadModelResponse>, (axum::http::StatusCode, String)> {
|
||||
let response = state
|
||||
.ollama
|
||||
.unload_model(&model)
|
||||
.await
|
||||
.map_err(ollama_err)?;
|
||||
|
||||
Ok(Json(response))
|
||||
}
|
||||
|
||||
fn ollama_err(e: OllamaError) -> (axum::http::StatusCode, String) {
|
||||
@@ -65,6 +60,10 @@ fn ollama_err(e: OllamaError) -> (axum::http::StatusCode, String) {
|
||||
axum::http::StatusCode::BAD_REQUEST,
|
||||
"messages array with at least one user message is required".to_string(),
|
||||
),
|
||||
OllamaError::MissingKeepAlive => (
|
||||
axum::http::StatusCode::BAD_REQUEST,
|
||||
"keep alive is required and cannot be empty".to_string(),
|
||||
),
|
||||
OllamaError::InvalidKeepAlive(v) => (
|
||||
axum::http::StatusCode::BAD_REQUEST,
|
||||
format!("invalid keep_alive '{v}' — use 30s / 10m / 2h, a plain integer, or -1"),
|
||||
|
||||
@@ -1,8 +1,8 @@
|
||||
use axum::Json;
|
||||
use utoipa::OpenApi;
|
||||
// use axum::Json;
|
||||
// use utoipa::OpenApi;
|
||||
|
||||
use crate::openapi::V1ApiDoc;
|
||||
// use crate::openapi::V1ApiDoc;
|
||||
|
||||
pub async fn openapi_json() -> Json<utoipa::openapi::OpenApi> {
|
||||
Json(V1ApiDoc::openapi())
|
||||
}
|
||||
// pub async fn openapi_json() -> Json<utoipa::openapi::OpenApi> {
|
||||
// Json(V1ApiDoc::openapi())
|
||||
// }
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
use crate::providers::ollama::OllamaProvider;
|
||||
use crate::providers::ollama::client::OllamaProvider;
|
||||
use std::sync::Arc;
|
||||
|
||||
#[derive(Clone)]
|
||||
|
||||
+134
-107
@@ -2,7 +2,9 @@ use serde_json::json;
|
||||
use wiremock::matchers::{method, path};
|
||||
use wiremock::{Mock, MockServer, ResponseTemplate};
|
||||
|
||||
use chat::providers::ollama::OllamaProvider;
|
||||
use chat::dto::api;
|
||||
use chat::errors::OllamaError;
|
||||
use chat::providers::ollama::client::OllamaProvider;
|
||||
|
||||
// ── helpers ──────────────────────────────────────────────────────────────────
|
||||
|
||||
@@ -31,7 +33,7 @@ async fn test_list_models_ok() {
|
||||
.await;
|
||||
|
||||
let res = provider.list_models().await.unwrap();
|
||||
assert_eq!(res["models"][0]["name"], "llama3");
|
||||
assert_eq!(res.models[0].name, "llama3");
|
||||
}
|
||||
|
||||
// ── completions ───────────────────────────────────────────────────────────────
|
||||
@@ -58,19 +60,25 @@ async fn test_completions_ok() {
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let res = provider
|
||||
.completions(json!({
|
||||
"model": "llama3",
|
||||
"prompt": "Who are you?"
|
||||
}))
|
||||
.await
|
||||
.unwrap();
|
||||
let req = api::CompletionRequest {
|
||||
base: api::BaseLLMRequest {
|
||||
model: "llama3".to_string(),
|
||||
..Default::default()
|
||||
},
|
||||
prompt: "Hello".to_string(),
|
||||
};
|
||||
|
||||
assert_eq!(res["object"], "text_completion");
|
||||
assert_eq!(res["choices"][0]["text"], "I am a helpful assistant.");
|
||||
assert_eq!(res["choices"][0]["finish_reason"], "stop");
|
||||
assert_eq!(res["usage"]["prompt_tokens"], 10);
|
||||
assert_eq!(res["usage"]["completion_tokens"], 8);
|
||||
let res = provider.completions(&req).await.unwrap();
|
||||
|
||||
assert_eq!(res.object, api::CompletionObject::TextCompletion);
|
||||
|
||||
assert_eq!(res.choices.len(), 1);
|
||||
assert_eq!(res.choices[0].text, "I am a helpful assistant.");
|
||||
assert_eq!(res.choices[0].finish_reason, api::FinishReason::Stop);
|
||||
|
||||
assert_eq!(res.usage.prompt_tokens, 10);
|
||||
assert_eq!(res.usage.completion_tokens, 8);
|
||||
assert_eq!(res.usage.total_tokens, 18);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -83,11 +91,17 @@ async fn test_completions_missing_prompt() {
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let err = provider
|
||||
.completions(json!({ "model": "llama3" }))
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(matches!(err, chat::errors::OllamaError::MissingPrompt));
|
||||
let req = api::CompletionRequest {
|
||||
base: api::BaseLLMRequest {
|
||||
model: "llama3".to_string(),
|
||||
..Default::default()
|
||||
},
|
||||
prompt: "".to_string(),
|
||||
};
|
||||
|
||||
let err = provider.completions(&req).await.unwrap_err();
|
||||
|
||||
assert!(matches!(err, OllamaError::MissingPrompt));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -100,13 +114,17 @@ async fn test_completions_empty_prompt() {
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let err = provider
|
||||
.completions(json!({
|
||||
"model": "llama3", "prompt": " "
|
||||
}))
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(matches!(err, chat::errors::OllamaError::MissingPrompt));
|
||||
let req = api::CompletionRequest {
|
||||
base: api::BaseLLMRequest {
|
||||
model: "llama3".to_string(),
|
||||
..Default::default()
|
||||
},
|
||||
prompt: " ".to_string(),
|
||||
};
|
||||
|
||||
let err = provider.completions(&req).await.unwrap_err();
|
||||
|
||||
assert!(matches!(err, OllamaError::MissingPrompt));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -119,12 +137,16 @@ async fn test_completions_model_not_found() {
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let err = provider
|
||||
.completions(json!({
|
||||
"model": "gpt-4", "prompt": "hello"
|
||||
}))
|
||||
.await
|
||||
.unwrap_err();
|
||||
let req = api::CompletionRequest {
|
||||
base: api::BaseLLMRequest {
|
||||
model: "gpt-4".to_string(),
|
||||
..Default::default()
|
||||
},
|
||||
prompt: "hello".to_string(),
|
||||
};
|
||||
|
||||
let err = provider.completions(&req).await.unwrap_err();
|
||||
|
||||
assert!(matches!(err, chat::errors::OllamaError::ModelNotFound(_)));
|
||||
}
|
||||
|
||||
@@ -146,28 +168,37 @@ async fn test_chat_completions_ok() {
|
||||
"model": "llama3",
|
||||
"message": { "role": "assistant", "content": "4." },
|
||||
"done": true,
|
||||
"done_reason": "stop",
|
||||
"prompt_eval_count": 5,
|
||||
"eval_count": 2,
|
||||
"eval_count": 2
|
||||
})))
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let res = provider
|
||||
.chat_completions(json!({
|
||||
"model": "llama3",
|
||||
"messages": [
|
||||
{ "role": "user", "content": "What is 2+2?" }
|
||||
]
|
||||
}))
|
||||
.await
|
||||
.unwrap();
|
||||
let req = api::ChatRequest {
|
||||
base: api::BaseLLMRequest {
|
||||
model: "llama3".to_string(),
|
||||
..Default::default()
|
||||
},
|
||||
messages: vec![api::Message {
|
||||
role: api::Role::User,
|
||||
content: "What is 2+2?".to_string(),
|
||||
}],
|
||||
};
|
||||
|
||||
assert_eq!(res["object"], "chat.completion");
|
||||
assert_eq!(res["choices"][0]["message"]["role"], "assistant");
|
||||
assert_eq!(res["choices"][0]["message"]["content"], "4.");
|
||||
assert_eq!(res["choices"][0]["finish_reason"], "stop");
|
||||
assert_eq!(res["usage"]["total_tokens"], 7);
|
||||
let res = provider.chat_completions(&req).await.unwrap();
|
||||
|
||||
assert_eq!(res.object, "chat.completion");
|
||||
assert_eq!(res.choices.len(), 1);
|
||||
|
||||
assert_eq!(res.choices[0].message.role, api::Role::Assistant);
|
||||
assert_eq!(res.choices[0].message.content, "4.");
|
||||
|
||||
assert_eq!(res.choices[0].finish_reason, api::FinishReason::Stop);
|
||||
|
||||
let usage = res.usage.unwrap();
|
||||
assert_eq!(usage.prompt_tokens, 5);
|
||||
assert_eq!(usage.completion_tokens, 2);
|
||||
assert_eq!(usage.total_tokens, 7);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -180,12 +211,16 @@ async fn test_chat_completions_missing_messages() {
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let err = provider
|
||||
.chat_completions(json!({
|
||||
"model": "llama3", "messages": []
|
||||
}))
|
||||
.await
|
||||
.unwrap_err();
|
||||
let req = api::ChatRequest {
|
||||
base: api::BaseLLMRequest {
|
||||
model: "llama3".to_string(),
|
||||
..Default::default()
|
||||
},
|
||||
messages: vec![],
|
||||
};
|
||||
|
||||
let err = provider.chat_completions(&req).await.unwrap_err();
|
||||
|
||||
assert!(matches!(err, chat::errors::OllamaError::MissingMessages));
|
||||
}
|
||||
|
||||
@@ -199,13 +234,19 @@ async fn test_chat_completions_no_user_message() {
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let err = provider
|
||||
.chat_completions(json!({
|
||||
"model": "llama3",
|
||||
"messages": [{ "role": "system", "content": "be helpful" }]
|
||||
}))
|
||||
.await
|
||||
.unwrap_err();
|
||||
let req = api::ChatRequest {
|
||||
base: api::BaseLLMRequest {
|
||||
model: "llama3".to_string(),
|
||||
..Default::default()
|
||||
},
|
||||
messages: vec![api::Message {
|
||||
role: api::Role::System,
|
||||
content: "be helpful".to_string(),
|
||||
}],
|
||||
};
|
||||
|
||||
let err = provider.chat_completions(&req).await.unwrap_err();
|
||||
|
||||
assert!(matches!(err, chat::errors::OllamaError::MissingMessages));
|
||||
}
|
||||
|
||||
@@ -219,17 +260,23 @@ async fn test_chat_completions_model_not_found() {
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let err = provider
|
||||
.chat_completions(json!({
|
||||
"model": "gpt-4",
|
||||
"messages": [{ "role": "user", "content": "hi" }]
|
||||
}))
|
||||
.await
|
||||
.unwrap_err();
|
||||
let req = api::ChatRequest {
|
||||
base: api::BaseLLMRequest {
|
||||
model: "gpt-4".to_string(),
|
||||
..Default::default()
|
||||
},
|
||||
messages: vec![api::Message {
|
||||
role: api::Role::User,
|
||||
content: "hi".to_string(),
|
||||
}],
|
||||
};
|
||||
|
||||
let err = provider.chat_completions(&req).await.unwrap_err();
|
||||
|
||||
assert!(matches!(err, chat::errors::OllamaError::ModelNotFound(_)));
|
||||
}
|
||||
|
||||
// ── load_model ────────────────────────────────────────────────────────────────
|
||||
// // ── load_model ────────────────────────────────────────────────────────────────
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_load_model_ok() {
|
||||
@@ -245,41 +292,17 @@ async fn test_load_model_ok() {
|
||||
.and(path("/api/generate"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
|
||||
"model": "llama3",
|
||||
"done": true,
|
||||
"response": "ok",
|
||||
"done": true,
|
||||
})))
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let res = provider.load_model("llama3", Some("10m")).await.unwrap();
|
||||
|
||||
assert_eq!(res["model"], "llama3");
|
||||
assert_eq!(res["status"], "loaded");
|
||||
assert_eq!(res["keep_alive"], "10m");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_load_model_default_keep_alive() {
|
||||
let (server, provider) = setup().await;
|
||||
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/api/tags"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(models_response(&["llama3"])))
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/api/generate"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
|
||||
"model": "llama3",
|
||||
"done": true,
|
||||
})))
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let res = provider.load_model("llama3", None).await.unwrap();
|
||||
|
||||
assert_eq!(res["status"], "loaded");
|
||||
assert_eq!(res["keep_alive"], "5m"); // default
|
||||
assert_eq!(res.model, "llama3");
|
||||
assert_eq!(res.status, "loaded");
|
||||
assert_eq!(res.keep_alive, "10m");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -293,6 +316,7 @@ async fn test_load_model_not_found() {
|
||||
.await;
|
||||
|
||||
let err = provider.load_model("gpt-4", Some("10m")).await.unwrap_err();
|
||||
|
||||
assert!(matches!(err, chat::errors::OllamaError::ModelNotFound(_)));
|
||||
}
|
||||
|
||||
@@ -310,6 +334,7 @@ async fn test_load_model_invalid_keep_alive() {
|
||||
.load_model("llama3", Some("10x"))
|
||||
.await
|
||||
.unwrap_err();
|
||||
|
||||
assert!(matches!(
|
||||
err,
|
||||
chat::errors::OllamaError::InvalidKeepAlive(_)
|
||||
@@ -330,20 +355,23 @@ async fn test_load_model_keep_alive_plain_integer() {
|
||||
.and(path("/api/generate"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
|
||||
"model": "llama3",
|
||||
"done": true,
|
||||
"done": true,
|
||||
})))
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
// plain integers (seconds) and "-1" are valid
|
||||
let res = provider.load_model("llama3", Some("3600")).await.unwrap();
|
||||
assert_eq!(res["status"], "loaded");
|
||||
|
||||
assert_eq!(res.status, "loaded");
|
||||
assert_eq!(res.keep_alive, "3600");
|
||||
|
||||
let res = provider.load_model("llama3", Some("-1")).await.unwrap();
|
||||
assert_eq!(res["status"], "loaded");
|
||||
|
||||
assert_eq!(res.status, "loaded");
|
||||
assert_eq!(res.keep_alive, "-1");
|
||||
}
|
||||
|
||||
// ── unload_model ──────────────────────────────────────────────────────────────
|
||||
// // ── unload_model ──────────────────────────────────────────────────────────────
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_unload_model_ok() {
|
||||
@@ -359,6 +387,7 @@ async fn test_unload_model_ok() {
|
||||
.and(path("/api/generate"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
|
||||
"model": "llama3",
|
||||
"response": "ok",
|
||||
"done": true,
|
||||
})))
|
||||
.mount(&server)
|
||||
@@ -366,10 +395,8 @@ async fn test_unload_model_ok() {
|
||||
|
||||
let res = provider.unload_model("llama3").await.unwrap();
|
||||
|
||||
assert_eq!(res["model"], "llama3");
|
||||
assert_eq!(res["status"], "unloaded");
|
||||
// no keep_alive field on unload response
|
||||
assert!(res.get("keep_alive").is_none());
|
||||
assert_eq!(res.model, "llama3");
|
||||
assert_eq!(res.status, "unloaded");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
|
||||
Reference in New Issue
Block a user