diff --git a/.claude/rules/code-style.md b/.claude/rules/code-style.md new file mode 100644 index 00000000..34f774a8 --- /dev/null +++ b/.claude/rules/code-style.md @@ -0,0 +1,23 @@ +# Coding Standards + +- Keep at most a single public struct per module. +- Keep at most a single public function per module (multiple public struct methods are OK). +- Keep module names elegant and clearly readable. The name of the module, or any file, should be enough to determine its contents unambiguously. +- Keep modules structure as flat as possible, avoid logical grouping of modules, instead keep the naming consistent. +- Keep standalone, private functions and structs above the public struct or function that is exported. +- Group the modules by name prefix. For example, `client_foo`, `client_bar`, etc., wherever it makes sense to do so. +- Decide to group the modules based on software architecture, messaging hierarchy, or inheritance. Do not group modules just for the sake of it. +- Maintain a tree-like structure of modules, avoid circular dependencies at all costs. Extract common functions or structs into separate modules, or separate subprojects in the workspace. +- Name files the same way as the struct or function they contain. +- Be explicit, do not use general import statements that involve "*", prefer to import everything explicitly. +- Do not use copy-pasted or copied code in any capacity. If you have issues extracting something into a module, discuss the steps first. +- Keeping slightly different message types, or other kinds of structs that are only slightly different, because of the context they are used in, is fine. +- Each function or method should do just a single thing. The single responsibility principle is really important. +- Always use descriptive and explicit variable names, even in anonymous functions. Never use single-letter variable names. +- Instead of writing comments that explain what the code does, make the code self-documenting. +- Handle all the errors; never ignore them. Make sure the application does not panic. +- Use object-oriented style and composition. Avoid functions that take a struct as a parameter; move it to the struct implementation instead. +- Avoid unnecessary abstractions. +- Before using vendor crates or modules, make sure they are well-maintained, secure, and documented. +- Always make sure there is only one valid way to do a specific task in the codebase. Make sure everything has a single source of truth. +- Prefer using data/value objects instead of inline types diff --git a/.claude/rules/commits.md b/.claude/rules/commits.md new file mode 100644 index 00000000..e660b3a0 --- /dev/null +++ b/.claude/rules/commits.md @@ -0,0 +1,5 @@ +# Committing Changes + +- Always keep the commit messages short, human-readable, and descriptive. Keep commit messages as one-liners. +- Do not add any metadata to commits. +- Describe what the changes actually do instead of listing the changed files. diff --git a/.claude/rules/rust.md b/.claude/rules/rust.md new file mode 100644 index 00000000..be678983 --- /dev/null +++ b/.claude/rules/rust.md @@ -0,0 +1,19 @@ +--- +paths: + - "**/*.rs" + - "**/Cargo.toml" +--- + +# Rust Standards + +- Do not inline import paths unless necessary. Prefer to use `use` statements in Rust files instead of inline paths to imported modules. The exception would be `error.rs` type modules that handle lib-level error structs. +- Always use explicit lifetime variable names (do not use `'a` and such, use descriptive names like `'message` or similar) +- Always use explicit generic parameter names (never use single letter names like `T` for generics, prefix all of them with `T`, however). For example, use `TMessage` instead of `T`, etc. +- Do not use `pub(crate)` in Rust; in case of doubt, just make things public. +- In Rust, never ignore errors with `Err(_)`; always make sure you are matching an expected error variant instead. +- Never use `.expect`, or `.unwrap`. In Rust, if a function can fail, use a matching Result (can be from the anyhow crate) instead. In case of doubt on this, ask. Allow `.expect` in mutex lock poison checks, or when integrating CPP libraries into Rust. +- Always make sure mutex locks are held for the shortest possible time. +- Always specify Rust dependencies in root Cargo.toml, then use workspace versions of packages in workspace members. +- In Rust, when implementing a `new` method in a struct, prefer to use a struct with a parameter list instead of multiple function arguments. It should be easier to maintain. +- Always check the project with Clippy. +- Always format the code with `cargo fmt`. diff --git a/.claude/rules/teamwork.md b/.claude/rules/teamwork.md new file mode 100644 index 00000000..34752f56 --- /dev/null +++ b/.claude/rules/teamwork.md @@ -0,0 +1,7 @@ +# Teamwork and Project Organization + +Team members own one module each. The project needs to be organized around small self-contained modules. + +Each class, struct, function, interface, trait, and alike needs to be named after its functionality in self-descriptive English. The goal is to name things in a way that will allow anyone to understand the project organization, and goals by just listing the directory of files. + +Developers need to be able to own their own modules without stepping on another's work. diff --git a/.claude/rules/testing.md b/.claude/rules/testing.md new file mode 100644 index 00000000..c60bf4b2 --- /dev/null +++ b/.claude/rules/testing.md @@ -0,0 +1,14 @@ +# Unit Tests and Quality Control + +- Always check that the unit tests pass. +- Always test the code, make sure tests work after the changes. +- Always write tests that check the algorithms, or meaningful edge cases. Never write tests that check things that can be handled by types instead. +- If some piece of code can be handled by proper types, use types instead. Write tests as a last resort. +- In unit tests, make sure there is always just a single correct way to do a specific thing. Never accept fuzzy inputs from end users. +- When working on tests, if you notice that the tested code can be better, you can suggest changes. +- Maintain 100% test coverage across the codebase. No file, branch, or line may be excluded from coverage reports. +- Reach 100% coverage with the minimum number of tests. Each test must cover a unique code path, behavior, or edge case that no other test already covers. +- If two tests cover overlapping paths, remove the weaker one. Redundant tests waste maintenance effort without improving correctness signal. +- Tests must exercise actual functionality and observable behavior. Never write a test purely to hit lines for the sake of coverage. +- Design tests deliberately before writing them. Identify the feature or branch under test, then write the smallest test that verifies it. +- Coverage gaps signal missing tests, never permission to exclude files. Write the test instead of suppressing the gap. diff --git a/.github/workflows/unit-tests.yml b/.github/workflows/unit-tests.yml index a1e2a152..864d5133 100644 --- a/.github/workflows/unit-tests.yml +++ b/.github/workflows/unit-tests.yml @@ -17,7 +17,7 @@ jobs: with: submodules: recursive - - uses: dtolnay/rust-toolchain@stable + - uses: dtolnay/rust-toolchain@29eef336d9b2848a0b548edc03f92a220660cdb8 # stable - uses: Swatinem/rust-cache@v2 @@ -34,7 +34,7 @@ jobs: - name: install system dependencies run: sudo apt-get update && sudo apt-get install -y cmake libclang-dev - - uses: dtolnay/rust-toolchain@stable + - uses: dtolnay/rust-toolchain@29eef336d9b2848a0b548edc03f92a220660cdb8 # stable - uses: Swatinem/rust-cache@v2 diff --git a/CLAUDE.md b/CLAUDE.md index a1373d68..ccdbce0d 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -6,61 +6,6 @@ Keep it simple, be opinionated, follow best practices. Avoid using configurable Keep the code beautiful. Always optimize the code for a great developer experience. -Be proactive and fix preexisting issues if you encounter them. +Codebase needs to be architected in a way to make it easy for multiple team members to work in parallel on multiple modules, so the concerns always need clear separation. -Be uncompromising when it comes to the code quality and architecture. Any compromises, coverage gaps, or quality gaps are not acceptable. - -Never make assumptions or guesses about code behavior; always investigate. Always make sure everything works. - -## Coding Standards - -- Do not inline import paths unless necessary. Prefer to use `use` statements in Rust files instead of inline paths to imported modules. The exception would be `error.rs` type modules that handle lib-level error structs. -- Keep at most a single public struct per Rust module. -- Keep at most a single public function per Rust module (multiple public struct methods are OK). -- Keep module names elegant and clearly readable. The name of the module, or any file, should be enough to determine its contents unambiguously. -- Keep modules structure as flat as possible, avoid logical grouping of modules, instead keep the naming consistent. -- Keep standalone, private functions and structs above the public struct or function that is exported. -- Group the modules by name prefix. For example, `client_foo`, `client_bar`, etc., wherever it makes sense to do so. -- Decide to group the modules based on software architecture, messaging hierarchy, or inheritance. Do not group modules just for the sake of it. -- Maintain a tree-like structure of modules, avoid circular dependencies at all costs. Extract common functions or structs into separate modules, or separate subprojects in the workspace. -- Name files the same way as the struct or function they contain. -- Be explicit, do not use general import statements that involve "*", prefer to import everything explicitly. -- Do not use copy-pasted or copied code in any capacity. If you have issues extracting something into a module, discuss the steps first. -- Keeping slightly different message types, or other kinds of structs that are only slightly different, because of the context they are used in, is fine. -- Each function or method should do just a single thing. The single responsibility principle is really important. -- Always use explicit lifetime variable names (do not use `'a` and such, use descriptive names like `'message` or similar) -- Always use explicit generic parameter names (never use single letter names like `T` for generics, prefix all of them with `T`, however). For example, use `TMessage` instead of `T`, etc. -- Always use descriptive and explicit variable names, even in anonymous functions. Never use single-letter variable names. -- Instead of writing comments that explain what the code does, make the code self-documenting. -- Do not use `pub(crate)` in Rust; in case of doubt, just make things public. -- Add an empty line before return statements that end the function or a method. -- Add an empty line between loops and preceding statements from the same scope. -- Handle all the errors; never ignore them. Make sure the application does not panic. -- In Rust, never ignore errors with `Err(_)`; always make sure you are matching an expected error variant instead. -- Never use `.expect`, or `.unwrap`. In Rust, if a function can fail, use a matching Result (can be from the anyhow crate) instead. In case of doubt on this, ask. Allow `.expect` in mutex lock poison checks, unit tests, or when integrating CPP libraries into Rust, and there is no way to use Result instead. -- Use object-oriented style and composition. Avoid functions that take a struct as a parameter; move it to the struct implementation instead. -- Always make sure mutex locks are held for the shortest possible time. -- Always specify Rust dependencies in root Cargo.toml, then use workspace versions of packages in workspace members. -- Avoid unnecessary abstractions. -- Before using vendor crates or modules, make sure they are well-maintained, secure, and documented. -- Always make sure there is only one valid way to do a specific task in the codebase. Make sure everything has a single source of truth. -- In Rust, when implementing `new` method in a struct, prefer to use a struct with parameters list instead of multiple function arguments. It should be easier to maintain. -- Use only the most precise error variants to cover a Result error case. If nothing suitable is available, add a new error variant. - -## Unit Tests and Quality Control - -- Always check the project with Clippy. -- Always format the code with `cargo fmt`. -- Always check that the unit tests pass. -- Always test the code, make sure tests work after the changes. -- Always write tests that check the algorithms, or meaningful edge cases. Never write tests that check things that can be handled by types instead. -- If some piece of code can be handled by proper types, use types instead. Write tests as a last resort. -- In unit tests, make sure there is always just a single correct way to do a specific thing. Never accept fuzzy inputs from end users. -- When working on tests, if you notice that the tested code can be better, you can suggest changes. -- When running tests, always save output to a temporary file, so you won't need to re-run them to analyze it. - -## Committing Changes - -- Always keep the commit messages short, human readable, descriptive. Keep commit messages as one-liners. -- Do not add any metadata to commits. -- Describe what the changes actually do instead of listing the changed files. +Be proactive and fix any preexisting issues you encounter. diff --git a/Cargo.lock b/Cargo.lock index 0a67d242..1ca757cc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -140,7 +140,7 @@ version = "0.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6fac387a98bb7c37292057cffc56d62ecb629900026402633ae9160df93a8766" dependencies = [ - "nom", + "nom 7.1.3", ] [[package]] @@ -1045,7 +1045,10 @@ dependencies = [ "encoding_rs", "enumflags2", "llama-cpp-bindings-sys", + "llama-cpp-bindings-types", "llguidance", + "nom 8.0.0", + "serde_json", "serial_test", "thiserror", "toktrie", @@ -1063,6 +1066,7 @@ dependencies = [ "cmake", "find_cuda_helper", "glob", + "thiserror", "walkdir", ] @@ -1088,6 +1092,15 @@ dependencies = [ "tracing-subscriber", ] +[[package]] +name = "llama-cpp-bindings-types" +version = "0.5.0" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "llguidance" version = "1.7.0" @@ -1186,6 +1199,15 @@ dependencies = [ "minimal-lexical", ] +[[package]] +name = "nom" +version = "8.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df9761775871bdef83bee530e60050f7e54b1105350d6884eb0fb4f46c2f9405" +dependencies = [ + "memchr", +] + [[package]] name = "nu-ansi-term" version = "0.50.3" diff --git a/Cargo.toml b/Cargo.toml index fe9db1cb..fe4ad0c3 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -3,6 +3,7 @@ resolver = "2" members = [ "llama-cpp-bindings-build", "llama-cpp-bindings-sys", + "llama-cpp-bindings-types", "llama-cpp-bindings", "llama-cpp-bindings-tests", ] @@ -11,9 +12,27 @@ members = [ edition = "2024" [workspace.dependencies] -encoding_rs = "0.8.35" -llama-cpp-bindings = { path = "llama-cpp-bindings", version = "0.5.0" } -llama-cpp-bindings-build = { path = "llama-cpp-bindings-build", version = "0.5.0" } -llama-cpp-bindings-sys = { path = "llama-cpp-bindings-sys", version = "0.5.0" } -tracing = "0.1" - +anyhow = "=1.0.102" +bindgen = "=0.72.1" +cc = { version = "=1.2.58", features = ["parallel"] } +cmake = "=0.1.58" +encoding_rs = "=0.8.35" +enumflags2 = "=0.7.12" +find_cuda_helper = "=0.2.0" +glob = "=0.3.3" +hf-hub = "=0.5.0" +llama-cpp-bindings = { path = "llama-cpp-bindings", version = "=0.5.0" } +llama-cpp-bindings-build = { path = "llama-cpp-bindings-build", version = "=0.5.0" } +llama-cpp-bindings-sys = { path = "llama-cpp-bindings-sys", version = "=0.5.0" } +llama-cpp-bindings-types = { path = "llama-cpp-bindings-types", version = "=0.5.0" } +llguidance = "=1.7.0" +nom = "=8.0.0" +serde = { version = "=1.0.228", features = ["derive"] } +serde_json = "=1.0.149" +serial_test = "=3.4.0" +thiserror = "=2.0.18" +toktrie = "=1.7.0" +tracing = "=0.1.44" +tracing-core = "=0.1.36" +tracing-subscriber = { version = "=0.3.23", features = ["json"] } +walkdir = "=2.5.0" diff --git a/Makefile b/Makefile index 6f9134bd..bae3fbf6 100644 --- a/Makefile +++ b/Makefile @@ -1,7 +1,8 @@ -FEATURES = sampler,llguidance +FEATURES = sampler TEST_FEATURES = +QWEN_CAPABLE_FEATURES = multimodal_capable,mrope_model CARGO_TEST_LLM_FLAGS = --no-fail-fast -p llama-cpp-bindings-tests $(if $(TEST_FEATURES),--features $(TEST_FEATURES),) -- --test-threads=1 -CARGO_COV_LLM_FLAGS = -p llama-cpp-bindings-tests $(if $(TEST_FEATURES),--features $(TEST_FEATURES),) +CARGO_TEST_LLM_FLAGS_QWEN_CAPABLE = --no-fail-fast -p llama-cpp-bindings-tests $(if $(TEST_FEATURES),--features $(TEST_FEATURES),) --features $(QWEN_CAPABLE_FEATURES) -- --test-threads=1 QWEN3_5_0_8B_ENV = \ LLAMA_TEST_HF_REPO=unsloth/Qwen3.5-0.8B-GGUF \ @@ -21,37 +22,48 @@ QWEN3_6_35B_A3B_ENV = \ LLAMA_TEST_HF_ENCODER_REPO=Xiaojian9992024/t5-small-GGUF \ LLAMA_TEST_HF_ENCODER_MODEL=t5-small.bf16.gguf +GLM4_7_FLASH_ENV = \ + LLAMA_TEST_HF_REPO=unsloth/GLM-4.7-Flash-GGUF \ + LLAMA_TEST_HF_MODEL=GLM-4.7-Flash-Q4_K_M.gguf \ + LLAMA_TEST_HF_EMBED_REPO=Qwen/Qwen3-Embedding-0.6B-GGUF \ + LLAMA_TEST_HF_EMBED_MODEL=Qwen3-Embedding-0.6B-Q8_0.gguf \ + LLAMA_TEST_HF_ENCODER_REPO=Xiaojian9992024/t5-small-GGUF \ + LLAMA_TEST_HF_ENCODER_MODEL=t5-small.bf16.gguf + +DEEPSEEK_R1_DISTILL_LLAMA_8B_ENV = \ + LLAMA_TEST_HF_REPO=unsloth/DeepSeek-R1-Distill-Llama-8B-GGUF \ + LLAMA_TEST_HF_MODEL=DeepSeek-R1-Distill-Llama-8B-Q4_K_M.gguf \ + LLAMA_TEST_HF_EMBED_REPO=Qwen/Qwen3-Embedding-0.6B-GGUF \ + LLAMA_TEST_HF_EMBED_MODEL=Qwen3-Embedding-0.6B-Q8_0.gguf \ + LLAMA_TEST_HF_ENCODER_REPO=Xiaojian9992024/t5-small-GGUF \ + LLAMA_TEST_HF_ENCODER_MODEL=t5-small.bf16.gguf + .PHONY: test.unit test.unit: clippy cargo test -p llama-cpp-bindings --features $(FEATURES) +.PHONY: test.deepseek_r1_distill_llama_8b +test.deepseek_r1_distill_llama_8b: clippy + $(DEEPSEEK_R1_DISTILL_LLAMA_8B_ENV) cargo test $(CARGO_TEST_LLM_FLAGS) + +.PHONY: test.glm4_7_flash +test.glm4_7_flash: clippy + $(GLM4_7_FLASH_ENV) cargo test $(CARGO_TEST_LLM_FLAGS) + .PHONY: test.qwen3.5_0.8B test.qwen3.5_0.8B: clippy - $(QWEN3_5_0_8B_ENV) cargo test $(CARGO_TEST_LLM_FLAGS) + $(QWEN3_5_0_8B_ENV) cargo test $(CARGO_TEST_LLM_FLAGS_QWEN_CAPABLE) .PHONY: test.qwen3.6_35b_a3b test.qwen3.6_35b_a3b: clippy - $(QWEN3_6_35B_A3B_ENV) cargo test $(CARGO_TEST_LLM_FLAGS) - -.PHONY: test.qwen3.5_0.8B.coverage.run -test.qwen3.5_0.8B.coverage.run: clippy - $(QWEN3_5_0_8B_ENV) cargo llvm-cov $(CARGO_COV_LLM_FLAGS) -- --test-threads=1 - -.PHONY: test.qwen3.5_0.8B.coverage - -test.qwen3.5_0.8B.coverage: clippy - $(QWEN3_5_0_8B_ENV) cargo llvm-cov $(CARGO_COV_LLM_FLAGS) --fail-under-lines 99.5 -- --test-threads=1 - -.PHONY: test.qwen3.5_0.8B.coverage.json -test.qwen3.5_0.8B.coverage.json: test.qwen3.5_0.8B.coverage.run - cargo llvm-cov report -p llama-cpp-bindings --json --output-path target/coverage.json - -.PHONY: test.qwen3.5_0.8B.coverage.html -test.qwen3.5_0.8B.coverage.html: test.qwen3.5_0.8B.coverage.run - cargo llvm-cov report -p llama-cpp-bindings --html + $(QWEN3_6_35B_A3B_ENV) cargo test $(CARGO_TEST_LLM_FLAGS_QWEN_CAPABLE) .PHONY: test.llms -test.llms: test.qwen3.5_0.8B +test.llms: \ + test.deepseek_r1_distill_llama_8b \ + test.glm4_7_flash \ + test.qwen3.5_0.8B \ + test.qwen3.6_35b_a3b .PHONY: test test: test.unit test.llms diff --git a/llama-cpp-bindings-build/Cargo.toml b/llama-cpp-bindings-build/Cargo.toml index b137ba93..2ecfbad9 100644 --- a/llama-cpp-bindings-build/Cargo.toml +++ b/llama-cpp-bindings-build/Cargo.toml @@ -7,12 +7,13 @@ license = "Apache-2.0" repository = "https://github.com/intentee/llama-cpp-bindings" [dependencies] -bindgen = "0.72.1" -cc = { version = "1.2.58", features = ["parallel"] } -cmake = "0.1" -find_cuda_helper = "0.2.0" -glob = "0.3.3" -walkdir = "2" +bindgen = { workspace = true } +cc = { workspace = true } +cmake = { workspace = true } +find_cuda_helper = { workspace = true } +glob = { workspace = true } +thiserror = { workspace = true } +walkdir = { workspace = true } [features] cuda = [] diff --git a/llama-cpp-bindings-build/src/android_ndk.rs b/llama-cpp-bindings-build/src/android_ndk.rs index d8bc2ec8..5c6c193f 100644 --- a/llama-cpp-bindings-build/src/android_ndk.rs +++ b/llama-cpp-bindings-build/src/android_ndk.rs @@ -1,6 +1,32 @@ use std::env; use std::path::{Path, PathBuf}; +use thiserror::Error; + +const DEFAULT_ANDROID_API_LEVEL: &str = "28"; + +#[derive(Debug, Error)] +pub enum AndroidNdkDetectionError { + #[error( + "Android NDK not found for target {target_triple}. Set ANDROID_NDK, ANDROID_NDK_ROOT, NDK_ROOT, or CARGO_NDK_ANDROID_NDK." + )] + NdkRootNotConfigured { + target_triple: String, + #[source] + source: env::VarError, + }, + #[error("Android NDK path does not exist: {path}")] + NdkRootMissing { path: PathBuf }, + #[error("Android NDK toolchain file not found: {path}")] + NdkToolchainFileMissing { path: PathBuf }, + #[error("Android NDK toolchain not found at: {path}")] + NdkToolchainDirectoryMissing { path: PathBuf }, + #[error("Unsupported host platform for Android NDK")] + UnsupportedHostPlatform, + #[error("Unsupported Android target triple: {target_triple}")] + UnsupportedAndroidTarget { target_triple: String }, +} + /// Consolidated Android NDK configuration, computed once and shared between /// bindgen and `CMake` configuration steps. #[derive(Debug)] @@ -16,7 +42,12 @@ pub struct AndroidNdk { } impl AndroidNdk { - pub fn detect(target_triple: &str) -> Result { + /// # Errors + /// + /// Returns [`AndroidNdkDetectionError`] when the NDK installation cannot be + /// located, an environment variable is missing, the target triple is + /// unsupported, or the host platform is not supported by the NDK. + pub fn detect(target_triple: &str) -> Result { let ndk_path = detect_ndk_path(target_triple)?; validate_ndk_installation(&ndk_path)?; @@ -28,10 +59,9 @@ impl AndroidNdk { let toolchain_path = format!("{ndk_path}/toolchains/llvm/prebuilt/{host_tag}"); if !Path::new(&toolchain_path).exists() { - return Err(format!( - "Android NDK toolchain not found at: {toolchain_path}\n\ - Please ensure you have the correct Android NDK for your platform." - )); + return Err(AndroidNdkDetectionError::NdkToolchainDirectoryMissing { + path: PathBuf::from(toolchain_path), + }); } let sysroot = format!("{toolchain_path}/sysroot"); @@ -58,35 +88,35 @@ impl AndroidNdk { } } -fn detect_ndk_path(target_triple: &str) -> Result { +fn detect_ndk_path(target_triple: &str) -> Result { env::var("ANDROID_NDK") - .or_else(|_| env::var("ANDROID_NDK_ROOT")) - .or_else(|_| env::var("NDK_ROOT")) - .or_else(|_| env::var("CARGO_NDK_ANDROID_NDK")) - .or_else(|_| detect_ndk_from_sdk()) - .map_err(|_| { - format!( - "Android NDK not found. Please set one of: ANDROID_NDK, NDK_ROOT, ANDROID_NDK_ROOT\n\ - Current target: {target_triple}\n\ - Download from: https://developer.android.com/ndk/downloads" - ) + .or_else(|_android_ndk_unset| env::var("ANDROID_NDK_ROOT")) + .or_else(|_android_ndk_root_unset| env::var("NDK_ROOT")) + .or_else(|_ndk_root_unset| env::var("CARGO_NDK_ANDROID_NDK")) + .or_else(|_cargo_ndk_android_ndk_unset| detect_ndk_from_sdk()) + .map_err(|source| AndroidNdkDetectionError::NdkRootNotConfigured { + target_triple: target_triple.to_owned(), + source, }) } fn detect_ndk_from_sdk() -> Result { - #[allow(deprecated)] let home = env::home_dir().ok_or(env::VarError::NotPresent)?; - let android_home = env::var("ANDROID_HOME") - .or_else(|_| env::var("ANDROID_SDK_ROOT")) - .unwrap_or_else(|_| format!("{}/Android/Sdk", home.display())); + let android_home = match env::var("ANDROID_HOME") + .or_else(|_android_home_unset| env::var("ANDROID_SDK_ROOT")) + { + Ok(value) => value, + Err(_neither_env_var_set) => format!("{}/Android/Sdk", home.display()), + }; let ndk_dir = format!("{android_home}/ndk"); - let entries = std::fs::read_dir(&ndk_dir).map_err(|_| env::VarError::NotPresent)?; + let entries = + std::fs::read_dir(&ndk_dir).map_err(|_directory_unreadable| env::VarError::NotPresent)?; let mut versions: Vec = entries .filter_map(std::result::Result::ok) - .filter(|entry| entry.file_type().map(|ft| ft.is_dir()).unwrap_or(false)) + .filter(|entry| entry.file_type().is_ok_and(|file_type| file_type.is_dir())) .filter_map(|entry| { entry .file_name() @@ -103,24 +133,21 @@ fn detect_ndk_from_sdk() -> Result { .ok_or(env::VarError::NotPresent) } -fn validate_ndk_installation(ndk_path: &str) -> Result<(), String> { +fn validate_ndk_installation(ndk_path: &str) -> Result<(), AndroidNdkDetectionError> { let ndk_path = Path::new(ndk_path); if !ndk_path.exists() { - return Err(format!( - "Android NDK path does not exist: {}", - ndk_path.display() - )); + return Err(AndroidNdkDetectionError::NdkRootMissing { + path: ndk_path.to_path_buf(), + }); } let toolchain_file = ndk_path.join("build/cmake/android.toolchain.cmake"); if !toolchain_file.exists() { - return Err(format!( - "Android NDK toolchain file not found: {}\n\ - This indicates an incomplete NDK installation.", - toolchain_file.display() - )); + return Err(AndroidNdkDetectionError::NdkToolchainFileMissing { + path: toolchain_file, + }); } Ok(()) @@ -128,14 +155,16 @@ fn validate_ndk_installation(ndk_path: &str) -> Result<(), String> { fn detect_api_level() -> String { env::var("ANDROID_API_LEVEL") - .or_else(|_| env::var("ANDROID_PLATFORM").map(|platform| platform.replace("android-", ""))) - .or_else(|_| { + .or_else(|_android_api_level_unset| { + env::var("ANDROID_PLATFORM").map(|platform| platform.replace("android-", "")) + }) + .or_else(|_android_platform_unset| { env::var("CARGO_NDK_ANDROID_PLATFORM").map(|platform| platform.replace("android-", "")) }) - .unwrap_or_else(|_| "28".to_string()) + .unwrap_or_else(|_no_api_level_configured| DEFAULT_ANDROID_API_LEVEL.to_string()) } -fn detect_host_tag() -> Result<&'static str, String> { +fn detect_host_tag() -> Result<&'static str, AndroidNdkDetectionError> { if cfg!(target_os = "macos") { Ok("darwin-x86_64") } else if cfg!(target_os = "linux") { @@ -143,11 +172,11 @@ fn detect_host_tag() -> Result<&'static str, String> { } else if cfg!(target_os = "windows") { Ok("windows-x86_64") } else { - Err("Unsupported host platform for Android NDK".to_string()) + Err(AndroidNdkDetectionError::UnsupportedHostPlatform) } } -fn target_triple_to_abi(target_triple: &str) -> Result<&'static str, String> { +fn target_triple_to_abi(target_triple: &str) -> Result<&'static str, AndroidNdkDetectionError> { if target_triple.contains("aarch64") { Ok("arm64-v8a") } else if target_triple.contains("armv7") { @@ -157,14 +186,15 @@ fn target_triple_to_abi(target_triple: &str) -> Result<&'static str, String> { } else if target_triple.contains("i686") { Ok("x86") } else { - Err(format!( - "Unsupported Android target: {target_triple}\n\ - Supported targets: aarch64-linux-android, armv7-linux-androideabi, i686-linux-android, x86_64-linux-android" - )) + Err(AndroidNdkDetectionError::UnsupportedAndroidTarget { + target_triple: target_triple.to_owned(), + }) } } -fn target_triple_to_ndk_prefix(target_triple: &str) -> Result<&'static str, String> { +fn target_triple_to_ndk_prefix( + target_triple: &str, +) -> Result<&'static str, AndroidNdkDetectionError> { if target_triple.contains("aarch64") { Ok("aarch64-linux-android") } else if target_triple.contains("armv7") { @@ -174,7 +204,9 @@ fn target_triple_to_ndk_prefix(target_triple: &str) -> Result<&'static str, Stri } else if target_triple.contains("i686") { Ok("i686-linux-android") } else { - Err(format!("Unsupported Android target: {target_triple}")) + Err(AndroidNdkDetectionError::UnsupportedAndroidTarget { + target_triple: target_triple.to_owned(), + }) } } @@ -183,11 +215,14 @@ fn find_clang_builtin_includes(toolchain_path: &str) -> Option { let entries = std::fs::read_dir(&clang_lib_path).ok()?; let version_dir = entries.filter_map(std::result::Result::ok).find(|entry| { - entry.file_type().map(|ft| ft.is_dir()).unwrap_or(false) + entry + .file_type() + .map(|file_type| file_type.is_dir()) + .unwrap_or(false) && entry .file_name() .to_str() - .is_some_and(|name| name.starts_with(|ch: char| ch.is_ascii_digit())) + .is_some_and(|name| name.starts_with(|character: char| character.is_ascii_digit())) })?; let include_path = PathBuf::from(&clang_lib_path) diff --git a/llama-cpp-bindings-build/src/cmake_config.rs b/llama-cpp-bindings-build/src/cmake_config.rs index 6d5a20f6..a52521e3 100644 --- a/llama-cpp-bindings-build/src/cmake_config.rs +++ b/llama-cpp-bindings-build/src/cmake_config.rs @@ -200,6 +200,7 @@ fn configure_platform_specific( match target_os { TargetOs::Apple(_) => { config.define("GGML_BLAS", "OFF"); + override_archive_commands_for_apple_ar(config); } TargetOs::Windows(WindowsVariant::Msvc) => { config.cflag("/w"); @@ -267,6 +268,31 @@ fn configure_android_cmake(config: &mut Config, ndk: &AndroidNdk, _target_triple println!("cargo:rustc-link-lib=android"); } +/// macOS BSD ar (from cctools) does not accept GNU ar's `-D` (deterministic) +/// flag. cmake's default archive recipe is ` qcD …`, which produces +/// `illegal option -- D` warnings during every static-library link. +/// +/// We override the archive command for every language used by llama.cpp's +/// build — C, C++, Objective-C and Objective-C++ (the latter two appear once +/// `GGML_METAL=ON` enables the Metal backend). Plain `qc` keeps the +/// quick-create semantics; `` still runs as ARCHIVE_FINISH. +fn override_archive_commands_for_apple_ar(config: &mut Config) { + for language in ["C", "CXX", "OBJC", "OBJCXX"] { + config.define( + format!("CMAKE_{language}_ARCHIVE_CREATE"), + " qc ", + ); + config.define( + format!("CMAKE_{language}_ARCHIVE_APPEND"), + " q ", + ); + config.define( + format!("CMAKE_{language}_ARCHIVE_FINISH"), + " ", + ); + } +} + fn configure_android_arch_flags(config: &mut Config, abi: &str) { match abi { "arm64-v8a" => { diff --git a/llama-cpp-bindings-build/src/cpp_wrapper.rs b/llama-cpp-bindings-build/src/cpp_wrapper.rs index e29cf9be..e85e4fe2 100644 --- a/llama-cpp-bindings-build/src/cpp_wrapper.rs +++ b/llama-cpp-bindings-build/src/cpp_wrapper.rs @@ -8,9 +8,15 @@ pub fn compile_cpp_wrappers(llama_src: &Path, target_os: &TargetOs) { build .cpp(true) .warnings(false) + .file("wrapper_chat_parse.cpp") .file("wrapper_common.cpp") .file("wrapper_fit.cpp") .file("wrapper_reasoning.cpp") + .file("wrapper_token_text.cpp") + .file("wrapper_tool_calls.cpp") + .file("marker_probes/chunked_thinking.cpp") + .file("marker_probes/registry.cpp") + .include(".") .include(llama_src) .include(llama_src.join("common")) .include(llama_src.join("include")) diff --git a/llama-cpp-bindings-build/src/rebuild_tracking.rs b/llama-cpp-bindings-build/src/rebuild_tracking.rs index 4d6565d1..43f8295f 100644 --- a/llama-cpp-bindings-build/src/rebuild_tracking.rs +++ b/llama-cpp-bindings-build/src/rebuild_tracking.rs @@ -19,14 +19,24 @@ fn is_cmake_file(entry: &DirEntry) -> bool { pub fn register_rebuild_triggers(llama_src: &Path) { println!("cargo:rerun-if-changed=build.rs"); println!("cargo:rerun-if-changed=wrapper.h"); + println!("cargo:rerun-if-changed=wrapper_chat_parse.h"); + println!("cargo:rerun-if-changed=wrapper_chat_parse.cpp"); println!("cargo:rerun-if-changed=wrapper_common.h"); println!("cargo:rerun-if-changed=wrapper_common.cpp"); println!("cargo:rerun-if-changed=wrapper_fit.h"); println!("cargo:rerun-if-changed=wrapper_fit.cpp"); println!("cargo:rerun-if-changed=wrapper_reasoning.h"); println!("cargo:rerun-if-changed=wrapper_reasoning.cpp"); + println!("cargo:rerun-if-changed=wrapper_token_text.h"); + println!("cargo:rerun-if-changed=wrapper_token_text.cpp"); + println!("cargo:rerun-if-changed=wrapper_tool_calls.h"); + println!("cargo:rerun-if-changed=wrapper_tool_calls.cpp"); println!("cargo:rerun-if-changed=wrapper_utils.h"); println!("cargo:rerun-if-changed=wrapper_mtmd.h"); + println!("cargo:rerun-if-changed=marker_probes/marker_probe.h"); + println!("cargo:rerun-if-changed=marker_probes/registry.cpp"); + println!("cargo:rerun-if-changed=marker_probes/chunked_thinking.h"); + println!("cargo:rerun-if-changed=marker_probes/chunked_thinking.cpp"); println!("cargo:rerun-if-env-changed=LLAMA_LIB_PROFILE"); println!("cargo:rerun-if-env-changed=LLAMA_BUILD_SHARED_LIBS"); diff --git a/llama-cpp-bindings-sys/llama.cpp b/llama-cpp-bindings-sys/llama.cpp index 278521c3..846262d7 160000 --- a/llama-cpp-bindings-sys/llama.cpp +++ b/llama-cpp-bindings-sys/llama.cpp @@ -1 +1 @@ -Subproject commit 278521c33a11b89d9d7ed2afe5c20502840816b1 +Subproject commit 846262d7875dcabf502a150fa3d7b9c770dde7eb diff --git a/llama-cpp-bindings-sys/marker_probes/chunked_thinking.cpp b/llama-cpp-bindings-sys/marker_probes/chunked_thinking.cpp new file mode 100644 index 00000000..d29e49ae --- /dev/null +++ b/llama-cpp-bindings-sys/marker_probes/chunked_thinking.cpp @@ -0,0 +1,144 @@ +#include "chunked_thinking.h" + +#include "llama.cpp/common/chat-auto-parser.h" +#include "llama.cpp/common/chat.h" + +#include +#include +#include +#include +#include + +namespace marker_probes { + +namespace { + +constexpr std::string_view REASON_PROBE = "__PADDLER_REASON_PROBE_3F4A8C__"; +constexpr std::string_view RESPONSE_PROBE = "__PADDLER_RESPONSE_PROBE_3F4A8C__"; + +std::string trim_copy(std::string_view input) { + auto first = input.find_first_not_of(" \t\r\n"); + if (first == std::string_view::npos) { + return {}; + } + auto last = input.find_last_not_of(" \t\r\n"); + return std::string(input.substr(first, last - first + 1)); +} + +bool render_template(const common_chat_template & tmpl, + const autoparser::generation_params & params, + std::string & out) { + try { + out = common_chat_template_direct_apply(tmpl, params); + return true; + } catch (const std::exception &) { + return false; + } catch (...) { + return false; + } +} + +autoparser::generation_params plain_text_params() { + autoparser::generation_params params; + params.add_generation_prompt = false; + params.enable_thinking = true; + params.is_inference = false; + params.add_inference = false; + params.mark_input = false; + params.messages = nlohmann::ordered_json::array({ + nlohmann::ordered_json{ { "role", "user" }, { "content", "U" } }, + nlohmann::ordered_json{ { "role", "assistant" }, { "content", std::string(RESPONSE_PROBE) } }, + }); + return params; +} + +autoparser::generation_params chunked_thinking_params() { + autoparser::generation_params params; + params.add_generation_prompt = false; + params.enable_thinking = true; + params.is_inference = false; + params.add_inference = false; + params.mark_input = false; + params.messages = nlohmann::ordered_json::array({ + nlohmann::ordered_json{ { "role", "user" }, { "content", "U" } }, + nlohmann::ordered_json{ + { "role", "assistant" }, + { "content", nlohmann::ordered_json::array({ + nlohmann::ordered_json{ { "type", "thinking" }, { "thinking", std::string(REASON_PROBE) } }, + nlohmann::ordered_json{ { "type", "text" }, { "text", std::string(RESPONSE_PROBE) } }, + }) }, + }, + }); + return params; +} + +bool contains(std::string_view haystack, std::string_view needle) { + return haystack.find(needle) != std::string_view::npos; +} + +} // namespace + +probe_result chunked_thinking(const common_chat_template & tmpl) { + probe_result result; + + std::string render_plain; + if (!render_template(tmpl, plain_text_params(), render_plain)) { + return result; + } + + std::string render_chunked; + if (!render_template(tmpl, chunked_thinking_params(), render_chunked)) { + return result; + } + + if (!contains(render_chunked, REASON_PROBE) || !contains(render_chunked, RESPONSE_PROBE)) { + return result; + } + + const std::size_t plain_size = render_plain.size(); + const std::size_t chunked_size = render_chunked.size(); + const std::size_t min_size = std::min(plain_size, chunked_size); + + std::size_t common_prefix = 0; + while (common_prefix < min_size && render_plain[common_prefix] == render_chunked[common_prefix]) { + ++common_prefix; + } + + std::size_t common_suffix = 0; + while (common_suffix < min_size - common_prefix + && render_plain[plain_size - 1 - common_suffix] == render_chunked[chunked_size - 1 - common_suffix]) { + ++common_suffix; + } + + if (common_prefix + common_suffix > chunked_size) { + return result; + } + + std::string_view diff_slice(render_chunked); + diff_slice = diff_slice.substr(common_prefix, chunked_size - common_prefix - common_suffix); + + auto reason_pos = diff_slice.find(REASON_PROBE); + if (reason_pos == std::string_view::npos) { + return result; + } + + std::string start = trim_copy(diff_slice.substr(0, reason_pos)); + std::string end = trim_copy(diff_slice.substr(reason_pos + REASON_PROBE.size())); + + if (start.empty() || end.empty()) { + return result; + } + if (contains(start, REASON_PROBE) || contains(start, RESPONSE_PROBE)) { + return result; + } + if (contains(end, REASON_PROBE) || contains(end, RESPONSE_PROBE)) { + return result; + } + + result.start = std::move(start); + result.end = std::move(end); + result.found = true; + return result; +} + +} // namespace marker_probes diff --git a/llama-cpp-bindings-sys/marker_probes/chunked_thinking.h b/llama-cpp-bindings-sys/marker_probes/chunked_thinking.h new file mode 100644 index 00000000..9128f68b --- /dev/null +++ b/llama-cpp-bindings-sys/marker_probes/chunked_thinking.h @@ -0,0 +1,9 @@ +#pragma once + +#include "marker_probe.h" + +namespace marker_probes { + +probe_result chunked_thinking(const common_chat_template & tmpl); + +} // namespace marker_probes diff --git a/llama-cpp-bindings-sys/marker_probes/marker_probe.h b/llama-cpp-bindings-sys/marker_probes/marker_probe.h new file mode 100644 index 00000000..3df72c39 --- /dev/null +++ b/llama-cpp-bindings-sys/marker_probes/marker_probe.h @@ -0,0 +1,20 @@ +#pragma once + +#include "llama.cpp/common/chat.h" + +#include +#include + +namespace marker_probes { + +struct probe_result { + std::string start; + std::string end; + bool found = false; +}; + +using probe_fn = probe_result (*)(const common_chat_template &); + +const std::vector & registered(); + +} // namespace marker_probes diff --git a/llama-cpp-bindings-sys/marker_probes/registry.cpp b/llama-cpp-bindings-sys/marker_probes/registry.cpp new file mode 100644 index 00000000..315bc56c --- /dev/null +++ b/llama-cpp-bindings-sys/marker_probes/registry.cpp @@ -0,0 +1,16 @@ +#include "marker_probe.h" + +#include "chunked_thinking.h" + +#include + +namespace marker_probes { + +const std::vector & registered() { + static const std::vector probes = { + chunked_thinking, + }; + return probes; +} + +} // namespace marker_probes diff --git a/llama-cpp-bindings-sys/wrapper.h b/llama-cpp-bindings-sys/wrapper.h index f371d6e5..eb98bc49 100644 --- a/llama-cpp-bindings-sys/wrapper.h +++ b/llama-cpp-bindings-sys/wrapper.h @@ -1,5 +1,7 @@ #include "llama.cpp/include/llama.h" #include "llama.cpp/ggml/include/gguf.h" +#include "wrapper_chat_parse.h" #include "wrapper_common.h" #include "wrapper_fit.h" #include "wrapper_reasoning.h" +#include "wrapper_tool_calls.h" diff --git a/llama-cpp-bindings-sys/wrapper_chat_parse.cpp b/llama-cpp-bindings-sys/wrapper_chat_parse.cpp new file mode 100644 index 00000000..f60cada6 --- /dev/null +++ b/llama-cpp-bindings-sys/wrapper_chat_parse.cpp @@ -0,0 +1,153 @@ +#include "wrapper_chat_parse.h" +#include "wrapper_token_text.h" + +#include "llama.cpp/common/chat-auto-parser.h" +#include "llama.cpp/common/chat.h" +#include "llama.cpp/include/llama.h" +#include "marker_probes/marker_probe.h" + +#include +#include +#include + +using wrapper_helpers::token_text_or_empty; + +struct llama_rs_parsed_chat { + common_chat_msg message; +}; + +extern "C" llama_rs_status llama_rs_parse_chat_message( + const struct llama_model * model, + const char * tools_json, + const char * input, + int is_partial, + llama_rs_parsed_chat_handle * out_handle, + char ** out_error) { + if (out_handle) { + *out_handle = nullptr; + } + if (out_error) { + *out_error = nullptr; + } + + if (!model || !input || !out_handle || !out_error) { + return LLAMA_RS_STATUS_INVALID_ARGUMENT; + } + + try { + const char * tmpl_src = llama_model_chat_template(model, nullptr); + if (!tmpl_src) { + return LLAMA_RS_STATUS_INVALID_ARGUMENT; + } + + const llama_vocab * vocab = llama_model_get_vocab(model); + if (!vocab) { + return LLAMA_RS_STATUS_INVALID_ARGUMENT; + } + + std::string bos_token = token_text_or_empty(vocab, llama_vocab_bos(vocab)); + std::string eos_token = token_text_or_empty(vocab, llama_vocab_eos(vocab)); + + common_chat_template tmpl(tmpl_src, bos_token, eos_token); + + autoparser::autoparser parser; + parser.analyze_template(tmpl); + + if (parser.reasoning.mode == autoparser::reasoning_mode::NONE) { + for (auto probe : marker_probes::registered()) { + auto fallback = probe(tmpl); + if (fallback.found) { + parser.reasoning.mode = autoparser::reasoning_mode::TAG_BASED; + parser.reasoning.start = std::move(fallback.start); + parser.reasoning.end = std::move(fallback.end); + break; + } + } + } + + autoparser::generation_params inputs; + inputs.add_generation_prompt = true; + inputs.enable_thinking = true; + inputs.messages = nlohmann::ordered_json::array({ + { { "role", "user" }, { "content", "ping" } } + }); + + if (tools_json && tools_json[0] != '\0') { + inputs.tools = nlohmann::ordered_json::parse(tools_json); + } else { + inputs.tools = nlohmann::ordered_json::array(); + } + + common_chat_params chat_params = + autoparser::peg_generator::generate_parser(tmpl, inputs, parser); + + common_chat_parser_params parser_params(chat_params); + parser_params.parser.load(chat_params.parser); + + common_chat_msg parsed = common_chat_parse(input, is_partial != 0, parser_params); + + auto * handle = new llama_rs_parsed_chat{}; + handle->message = std::move(parsed); + + *out_handle = handle; + + return LLAMA_RS_STATUS_OK; + } catch (const std::exception & ex) { + *out_error = llama_rs_dup_string(std::string(ex.what())); + + return LLAMA_RS_STATUS_EXCEPTION; + } catch (...) { + *out_error = llama_rs_dup_string(std::string("unknown c++ exception")); + + return LLAMA_RS_STATUS_EXCEPTION; + } +} + +extern "C" void llama_rs_parsed_chat_free(llama_rs_parsed_chat_handle handle) { + delete handle; +} + +extern "C" size_t llama_rs_parsed_chat_tool_call_count(llama_rs_parsed_chat_handle handle) { + if (!handle) { + return 0; + } + return handle->message.tool_calls.size(); +} + +extern "C" char * llama_rs_parsed_chat_tool_call_id( + llama_rs_parsed_chat_handle handle, size_t index) { + if (!handle || index >= handle->message.tool_calls.size()) { + return nullptr; + } + return llama_rs_dup_string(handle->message.tool_calls[index].id); +} + +extern "C" char * llama_rs_parsed_chat_tool_call_name( + llama_rs_parsed_chat_handle handle, size_t index) { + if (!handle || index >= handle->message.tool_calls.size()) { + return nullptr; + } + return llama_rs_dup_string(handle->message.tool_calls[index].name); +} + +extern "C" char * llama_rs_parsed_chat_tool_call_arguments( + llama_rs_parsed_chat_handle handle, size_t index) { + if (!handle || index >= handle->message.tool_calls.size()) { + return nullptr; + } + return llama_rs_dup_string(handle->message.tool_calls[index].arguments); +} + +extern "C" char * llama_rs_parsed_chat_content(llama_rs_parsed_chat_handle handle) { + if (!handle) { + return nullptr; + } + return llama_rs_dup_string(handle->message.content); +} + +extern "C" char * llama_rs_parsed_chat_reasoning_content(llama_rs_parsed_chat_handle handle) { + if (!handle) { + return nullptr; + } + return llama_rs_dup_string(handle->message.reasoning_content); +} diff --git a/llama-cpp-bindings-sys/wrapper_chat_parse.h b/llama-cpp-bindings-sys/wrapper_chat_parse.h new file mode 100644 index 00000000..12fed5d9 --- /dev/null +++ b/llama-cpp-bindings-sys/wrapper_chat_parse.h @@ -0,0 +1,58 @@ +#pragma once + +#include "llama.cpp/include/llama.h" +#include "wrapper_utils.h" + +#include + +#ifdef __cplusplus +extern "C" { +#endif + +struct llama_rs_parsed_chat; +typedef struct llama_rs_parsed_chat * llama_rs_parsed_chat_handle; + +/** + * Parse a chat-completion turn from raw assistant output using llama.cpp's + * `common_chat_parse`, driven by the model's autoparser-built peg parser. + * + * `tools_json` is a serialized JSON array of OpenAI-style tool definitions + * (or empty / null when the request had no tools). `is_partial` switches + * between mid-stream parses (partial accepts incomplete payloads) and final + * parses (rejects malformed input). + * + * On success, `*out_handle` owns the parsed message; free via + * `llama_rs_parsed_chat_free`. On failure, `*out_error` carries an + * exception message; free via `llama_rs_string_free`. + */ +llama_rs_status llama_rs_parse_chat_message( + const struct llama_model * model, + const char * tools_json, + const char * input, + int is_partial, + llama_rs_parsed_chat_handle * out_handle, + char ** out_error); + +void llama_rs_parsed_chat_free(llama_rs_parsed_chat_handle handle); + +size_t llama_rs_parsed_chat_tool_call_count(llama_rs_parsed_chat_handle handle); + +/** + * Returns a heap-allocated UTF-8 string for the i-th tool call's `id`, + * `name`, or `arguments` field. Free with `llama_rs_string_free`. Returns + * nullptr if `handle` is null or `index` is out of bounds. + * + * `arguments` is the raw JSON string emitted by the parser — the caller is + * expected to feed it into a schema validator or hand it back to clients + * verbatim. + */ +char * llama_rs_parsed_chat_tool_call_id(llama_rs_parsed_chat_handle handle, size_t index); +char * llama_rs_parsed_chat_tool_call_name(llama_rs_parsed_chat_handle handle, size_t index); +char * llama_rs_parsed_chat_tool_call_arguments(llama_rs_parsed_chat_handle handle, size_t index); + +char * llama_rs_parsed_chat_content(llama_rs_parsed_chat_handle handle); +char * llama_rs_parsed_chat_reasoning_content(llama_rs_parsed_chat_handle handle); + +#ifdef __cplusplus +} +#endif diff --git a/llama-cpp-bindings-sys/wrapper_reasoning.cpp b/llama-cpp-bindings-sys/wrapper_reasoning.cpp index 6e7edd7c..36b0763e 100644 --- a/llama-cpp-bindings-sys/wrapper_reasoning.cpp +++ b/llama-cpp-bindings-sys/wrapper_reasoning.cpp @@ -3,8 +3,10 @@ #include "llama.cpp/common/chat-auto-parser.h" #include "llama.cpp/common/chat.h" #include "llama.cpp/include/llama.h" +#include "marker_probes/marker_probe.h" #include +#include #include namespace { @@ -59,17 +61,62 @@ extern "C" llama_rs_status llama_rs_detect_reasoning_markers( common_chat_template tmpl(tmpl_src, bos_token, eos_token); - autoparser::autoparser parser; - parser.analyze_template(tmpl); + std::string detected_start; + std::string detected_end; + bool detected = false; + + autoparser::generation_params probe_params; + probe_params.add_generation_prompt = true; + probe_params.enable_thinking = true; + probe_params.is_inference = false; + probe_params.add_inference = false; + probe_params.mark_input = false; + probe_params.messages = nlohmann::ordered_json::array({ + nlohmann::ordered_json{ { "role", "user" }, { "content", "ping" } }, + }); + + const std::string tmpl_src_str = tmpl_src; + if (auto specialized = common_chat_try_specialized_template(tmpl, tmpl_src_str, probe_params)) { + if (specialized->supports_thinking + && !specialized->thinking_start_tag.empty() + && !specialized->thinking_end_tag.empty()) { + detected_start = std::move(specialized->thinking_start_tag); + detected_end = std::move(specialized->thinking_end_tag); + detected = true; + } + } + + if (!detected) { + autoparser::autoparser parser; + parser.analyze_template(tmpl); + + if (parser.reasoning.mode != autoparser::reasoning_mode::NONE + && !parser.reasoning.start.empty() + && !parser.reasoning.end.empty()) { + detected_start = std::move(parser.reasoning.start); + detected_end = std::move(parser.reasoning.end); + detected = true; + } + } - if (parser.reasoning.mode == autoparser::reasoning_mode::NONE - || parser.reasoning.start.empty() - || parser.reasoning.end.empty()) { + if (!detected) { + for (auto probe : marker_probes::registered()) { + auto fallback = probe(tmpl); + if (fallback.found) { + detected_start = std::move(fallback.start); + detected_end = std::move(fallback.end); + detected = true; + break; + } + } + } + + if (!detected) { return LLAMA_RS_STATUS_OK; } - char * open_dup = llama_rs_dup_string(parser.reasoning.start); - char * close_dup = llama_rs_dup_string(parser.reasoning.end); + char * open_dup = llama_rs_dup_string(detected_start); + char * close_dup = llama_rs_dup_string(detected_end); if (!open_dup || !close_dup) { std::free(open_dup); @@ -92,3 +139,4 @@ extern "C" llama_rs_status llama_rs_detect_reasoning_markers( return LLAMA_RS_STATUS_EXCEPTION; } } + diff --git a/llama-cpp-bindings-sys/wrapper_token_text.cpp b/llama-cpp-bindings-sys/wrapper_token_text.cpp new file mode 100644 index 00000000..78fbcddf --- /dev/null +++ b/llama-cpp-bindings-sys/wrapper_token_text.cpp @@ -0,0 +1,18 @@ +#include "wrapper_token_text.h" + +namespace wrapper_helpers { + +std::string token_text_or_empty(const llama_vocab * vocab, llama_token token) { + if (token == LLAMA_TOKEN_NULL) { + return {}; + } + + const char * text = llama_vocab_get_text(vocab, token); + if (!text) { + return {}; + } + + return std::string(text); +} + +} diff --git a/llama-cpp-bindings-sys/wrapper_token_text.h b/llama-cpp-bindings-sys/wrapper_token_text.h new file mode 100644 index 00000000..231527e1 --- /dev/null +++ b/llama-cpp-bindings-sys/wrapper_token_text.h @@ -0,0 +1,11 @@ +#pragma once + +#include "llama.cpp/include/llama.h" + +#include + +namespace wrapper_helpers { + +std::string token_text_or_empty(const llama_vocab * vocab, llama_token token); + +} diff --git a/llama-cpp-bindings-sys/wrapper_tool_calls.cpp b/llama-cpp-bindings-sys/wrapper_tool_calls.cpp new file mode 100644 index 00000000..eb869201 --- /dev/null +++ b/llama-cpp-bindings-sys/wrapper_tool_calls.cpp @@ -0,0 +1,278 @@ +#include "wrapper_tool_calls.h" +#include "wrapper_token_text.h" + +#include "llama.cpp/common/chat-auto-parser.h" +#include "llama.cpp/common/chat-auto-parser-helpers.h" +#include "llama.cpp/common/chat.h" +#include "llama.cpp/include/llama.h" + +#include +#include +#include + +using wrapper_helpers::token_text_or_empty; + +namespace { + +// Render the chat template with a deterministic tool-call assistant turn and +// diff it against the no-tool-call variant. Returns the raw section between +// the model's tool-call open/close markers — i.e. the `<...>{...}` +// fragment the model is expected to emit, with any reasoning prelude removed. +// +// We deliberately reproduce the autoparser's diff-based approach (so the +// detected markers come from the model's actual template behavior, not from a +// hardcoded list), but use plain-ASCII synthetic names where the upstream +// autoparser uses sentinel strings that some Jinja templates choke on. +std::string detect_tool_call_haystack( + const common_chat_template & tmpl, + const autoparser::analyze_reasoning & reasoning) { + nlohmann::ordered_json user_msg = { + { "role", "user" }, + { "content", "Please use the tool" } + }; + nlohmann::ordered_json assistant_no_tools = { + { "role", "assistant" }, + { "content", "Sure, calling." } + }; + nlohmann::ordered_json first_tool_call = { + { "id", "call_001" }, + { "type", "function" }, + { "function", { + { "name", "tool_first" }, + { "arguments", { + { "arg_first", "XXXX" }, + { "arg_second", "YYYY" }, + }} + }} + }; + nlohmann::ordered_json assistant_with_tools = { + { "role", "assistant" }, + { "content", "" }, + { "tool_calls", nlohmann::ordered_json::array({ first_tool_call }) } + }; + nlohmann::ordered_json tool_definition = { + { "type", "function" }, + { "function", { + { "name", "tool_first" }, + { "description", "First test tool" }, + { "parameters", { + { "type", "object" }, + { "properties", { + { "arg_first", { { "type", "string" }, { "description", "first arg" } } }, + { "arg_second", { { "type", "string" }, { "description", "second arg" } } }, + }}, + { "required", nlohmann::ordered_json::array({ "arg_first", "arg_second" }) }, + }} + }} + }; + + template_params params_no_tools; + params_no_tools.messages = nlohmann::ordered_json::array({ user_msg, assistant_no_tools }); + params_no_tools.tools = nlohmann::ordered_json::array({ tool_definition }); + params_no_tools.add_generation_prompt = false; + params_no_tools.enable_thinking = true; + + template_params params_with_tools = params_no_tools; + params_with_tools.messages = + nlohmann::ordered_json::array({ user_msg, assistant_with_tools }); + + std::string output_no_tools = autoparser::apply_template(tmpl, params_no_tools); + std::string output_with_tools = autoparser::apply_template(tmpl, params_with_tools); + + if (output_no_tools.empty() || output_with_tools.empty()) { + return {}; + } + + diff_split diff = calculate_diff_split(output_no_tools, output_with_tools); + std::string haystack = diff.right; + + // Strip reasoning markers so the surrounding tool-call markers can be + // located reliably — the autoparser does the same for the JSON-native + // path. + auto remove_first = [&haystack](const std::string & needle) { + if (needle.empty()) { + return; + } + auto pos = haystack.find(needle); + if (pos != std::string::npos) { + haystack = haystack.substr(0, pos) + haystack.substr(pos + needle.length()); + } + }; + + remove_first(reasoning.start); + remove_first(reasoning.end); + + return haystack; +} + +} // namespace + +extern "C" llama_rs_status llama_rs_compute_tool_call_haystack( + const struct llama_model * model, + char ** out_haystack, + char ** out_error) { + if (out_haystack) { + *out_haystack = nullptr; + } + if (out_error) { + *out_error = nullptr; + } + + if (!model || !out_haystack || !out_error) { + return LLAMA_RS_STATUS_INVALID_ARGUMENT; + } + + try { + const char * tmpl_src = llama_model_chat_template(model, nullptr); + if (!tmpl_src) { + return LLAMA_RS_STATUS_OK; + } + + const llama_vocab * vocab = llama_model_get_vocab(model); + if (!vocab) { + return LLAMA_RS_STATUS_OK; + } + + std::string bos_token = token_text_or_empty(vocab, llama_vocab_bos(vocab)); + std::string eos_token = token_text_or_empty(vocab, llama_vocab_eos(vocab)); + + common_chat_template tmpl(tmpl_src, bos_token, eos_token); + auto jinja_caps = tmpl.original_caps(); + autoparser::analyze_reasoning reasoning(tmpl, jinja_caps.supports_tool_calls); + + std::string haystack = detect_tool_call_haystack(tmpl, reasoning); + if (haystack.empty()) { + return LLAMA_RS_STATUS_OK; + } + + char * haystack_dup = llama_rs_dup_string(haystack); + if (!haystack_dup) { + return LLAMA_RS_STATUS_ALLOCATION_FAILED; + } + + *out_haystack = haystack_dup; + + return LLAMA_RS_STATUS_OK; + } catch (const std::exception & ex) { + *out_error = llama_rs_dup_string(std::string(ex.what())); + + return LLAMA_RS_STATUS_EXCEPTION; + } catch (...) { + *out_error = llama_rs_dup_string(std::string("unknown c++ exception")); + + return LLAMA_RS_STATUS_EXCEPTION; + } +} + +extern "C" llama_rs_status llama_rs_diagnose_tool_call_synthetic_renders( + const struct llama_model * model, + char ** out_no_tools, + char ** out_with_tools, + char ** out_error) { + if (out_no_tools) { + *out_no_tools = nullptr; + } + if (out_with_tools) { + *out_with_tools = nullptr; + } + if (out_error) { + *out_error = nullptr; + } + + if (!model || !out_no_tools || !out_with_tools || !out_error) { + return LLAMA_RS_STATUS_INVALID_ARGUMENT; + } + + try { + const char * tmpl_src = llama_model_chat_template(model, nullptr); + if (!tmpl_src) { + return LLAMA_RS_STATUS_OK; + } + + const llama_vocab * vocab = llama_model_get_vocab(model); + if (!vocab) { + return LLAMA_RS_STATUS_OK; + } + + std::string bos_token = token_text_or_empty(vocab, llama_vocab_bos(vocab)); + std::string eos_token = token_text_or_empty(vocab, llama_vocab_eos(vocab)); + + common_chat_template tmpl(tmpl_src, bos_token, eos_token); + + nlohmann::ordered_json user_msg = { + { "role", "user" }, + { "content", "Please use the tool" } + }; + nlohmann::ordered_json assistant_no_tools = { + { "role", "assistant" }, + { "content", "Sure, calling." } + }; + nlohmann::ordered_json first_tool_call = { + { "id", "call_001" }, + { "type", "function" }, + { "function", { + { "name", "tool_first" }, + { "arguments", { + { "arg_first", "XXXX" }, + { "arg_second", "YYYY" }, + }} + }} + }; + nlohmann::ordered_json assistant_with_tools = { + { "role", "assistant" }, + { "content", "" }, + { "tool_calls", nlohmann::ordered_json::array({ first_tool_call }) } + }; + nlohmann::ordered_json tool_definition = { + { "type", "function" }, + { "function", { + { "name", "tool_first" }, + { "description", "First test tool" }, + { "parameters", { + { "type", "object" }, + { "properties", { + { "arg_first", { { "type", "string" }, { "description", "first arg" } } }, + { "arg_second", { { "type", "string" }, { "description", "second arg" } } }, + }}, + { "required", nlohmann::ordered_json::array({ "arg_first", "arg_second" }) }, + }} + }} + }; + + template_params params_no_tools; + params_no_tools.messages = nlohmann::ordered_json::array({ user_msg, assistant_no_tools }); + params_no_tools.tools = nlohmann::ordered_json::array({ tool_definition }); + params_no_tools.add_generation_prompt = false; + params_no_tools.enable_thinking = true; + + template_params params_with_tools = params_no_tools; + params_with_tools.messages = + nlohmann::ordered_json::array({ user_msg, assistant_with_tools }); + + std::string output_a = autoparser::apply_template(tmpl, params_no_tools); + std::string output_b = autoparser::apply_template(tmpl, params_with_tools); + + char * a_dup = llama_rs_dup_string(output_a); + char * b_dup = llama_rs_dup_string(output_b); + + if (!a_dup || !b_dup) { + std::free(a_dup); + std::free(b_dup); + + return LLAMA_RS_STATUS_ALLOCATION_FAILED; + } + + *out_no_tools = a_dup; + *out_with_tools = b_dup; + + return LLAMA_RS_STATUS_OK; + } catch (const std::exception & ex) { + *out_error = llama_rs_dup_string(std::string(ex.what())); + + return LLAMA_RS_STATUS_EXCEPTION; + } catch (...) { + *out_error = llama_rs_dup_string(std::string("unknown c++ exception")); + + return LLAMA_RS_STATUS_EXCEPTION; + } +} diff --git a/llama-cpp-bindings-sys/wrapper_tool_calls.h b/llama-cpp-bindings-sys/wrapper_tool_calls.h new file mode 100644 index 00000000..e6a59e20 --- /dev/null +++ b/llama-cpp-bindings-sys/wrapper_tool_calls.h @@ -0,0 +1,51 @@ +#pragma once + +#include "llama.cpp/include/llama.h" +#include "wrapper_utils.h" + +#ifdef __cplusplus +extern "C" { +#endif + +/** + * Render the model's chat template with the autoparser's standard tool-call + * vs. plain-assistant synthetic turns and return the diff slice that surrounds + * the tool-call payload. The returned haystack is the text that lives between + * the model's tool-call open/close markers (with any reasoning prelude + * stripped). Marker extraction from the haystack is performed in Rust. + * + * On success (LLAMA_RS_STATUS_OK): + * - If the model declares no tool-call markers (or an empty haystack), + * *out_haystack is left as nullptr. + * - Otherwise *out_haystack is a heap-allocated null-terminated string owned + * by the caller. Free via llama_rs_string_free. + * + * On LLAMA_RS_STATUS_EXCEPTION, *out_error is set to a heap-allocated message; + * free via llama_rs_string_free. + */ +llama_rs_status llama_rs_compute_tool_call_haystack( + const struct llama_model * model, + char ** out_haystack, + char ** out_error); + +/** + * Render the model's chat template with the autoparser's standard synthetic + * inputs (assistant_no_tools vs assistant_with_tools). Useful for diagnosing + * why marker detection fails. + * + * On success (LLAMA_RS_STATUS_OK): + * - *out_no_tools and *out_with_tools point to heap-allocated rendered + * outputs (free via llama_rs_string_free). Either can be empty when the + * template throws during rendering. + * + * On LLAMA_RS_STATUS_EXCEPTION, *out_error is set. + */ +llama_rs_status llama_rs_diagnose_tool_call_synthetic_renders( + const struct llama_model * model, + char ** out_no_tools, + char ** out_with_tools, + char ** out_error); + +#ifdef __cplusplus +} +#endif diff --git a/llama-cpp-bindings-tests/Cargo.toml b/llama-cpp-bindings-tests/Cargo.toml index 1f1b210e..81ce6f39 100644 --- a/llama-cpp-bindings-tests/Cargo.toml +++ b/llama-cpp-bindings-tests/Cargo.toml @@ -7,15 +7,15 @@ license = "Apache-2.0" publish = false [dependencies] -anyhow = "1.0.102" +anyhow = { workspace = true } encoding_rs = { workspace = true } -hf-hub = "0.5.0" -llama-cpp-bindings = { workspace = true, features = ["sampler", "llguidance"] } +hf-hub = { workspace = true } +llama-cpp-bindings = { workspace = true, features = ["sampler"] } llama-cpp-bindings-sys = { workspace = true } -serde_json = "1.0" -serial_test = "3" +serde_json = { workspace = true } +serial_test = { workspace = true } tracing = { workspace = true } -tracing-subscriber = { version = "0.3", features = ["json"] } +tracing-subscriber = { workspace = true } [features] cuda = ["llama-cpp-bindings/cuda"] @@ -23,6 +23,8 @@ cuda-no-vmm = ["llama-cpp-bindings/cuda-no-vmm"] metal = ["llama-cpp-bindings/metal"] vulkan = ["llama-cpp-bindings/vulkan"] rocm = ["llama-cpp-bindings/rocm"] +multimodal_capable = [] +mrope_model = [] [lints.rust] unsafe_op_in_unsafe_fn = "warn" diff --git a/llama-cpp-bindings-tests/src/classify_sample_loop.rs b/llama-cpp-bindings-tests/src/classify_sample_loop.rs new file mode 100644 index 00000000..03ad1551 --- /dev/null +++ b/llama-cpp-bindings-tests/src/classify_sample_loop.rs @@ -0,0 +1,117 @@ +use anyhow::Result; +use llama_cpp_bindings::context::LlamaContext; +use llama_cpp_bindings::llama_batch::LlamaBatch; +use llama_cpp_bindings::model::LlamaModel; +use llama_cpp_bindings::sampled_token::SampledToken; +use llama_cpp_bindings::sampled_token_classifier::IngestOutcome; +use llama_cpp_bindings::sampled_token_classifier::SampledTokenClassifier; +use llama_cpp_bindings::sampling::LlamaSampler; + +/// Drives a classifier through the full sample/decode/flush loop. +/// +/// Suppresses EOG outcomes (so `generated_raw` and the per-section streams +/// never contain end-of-generation marker text) and captures per-section +/// counts. Tests that need to exercise classifier behaviour during real +/// inference should construct one of these and call +/// [`ClassifySampleLoop::run`] instead of re-implementing the loop. The +/// strict per-test assertions then run on [`ClassifySampleLoopOutcome`]. +pub struct ClassifySampleLoop<'borrow, 'model, 'tokens> { + pub model: &'model LlamaModel, + pub classifier: &'borrow mut SampledTokenClassifier<'model>, + pub sampler: &'borrow mut LlamaSampler, + pub context: &'borrow mut LlamaContext<'model>, + pub batch: &'borrow mut LlamaBatch<'tokens>, + pub initial_position: i32, + pub max_generated_tokens: i32, +} + +#[derive(Debug, Default)] +pub struct ClassifySampleLoopOutcome { + pub generated_raw: String, + pub content_stream: String, + pub reasoning_stream: String, + pub observed_content: u64, + pub observed_reasoning: u64, + pub observed_tool_call: u64, + pub observed_undeterminable: u64, + pub eog_seen: bool, +} + +impl ClassifySampleLoop<'_, '_, '_> { + /// # Errors + /// Forwards [`SampledTokenClassifier::sample`] / [`LlamaContext::decode`] / + /// [`LlamaBatch::add`] errors verbatim. Stops on EOG, on + /// `max_generated_tokens` exhaustion, or on the first error. + pub fn run(self) -> Result { + let mut outcome = ClassifySampleLoopOutcome::default(); + let mut position = self.initial_position; + let max_position = position + self.max_generated_tokens; + + while position < max_position { + let (raw_token, ingest_outcomes) = + self.classifier + .sample(self.sampler, self.context, self.batch.n_tokens() - 1)?; + + for ingest_outcome in &ingest_outcomes { + let is_eog = self.model.is_eog_token(&ingest_outcome.sampled_token); + if is_eog { + outcome.eog_seen = true; + } else { + outcome.generated_raw.push_str(&ingest_outcome.raw_piece); + } + // Counters always include EOG so they match the classifier's + // internal usage counters (which include every sampled token). + // EOG text is suppressed from `generated_raw` and the per-section + // streams so callers can assert exact textual equality. + record_outcome(ingest_outcome, &mut outcome, is_eog); + } + + let raw_as_sampled = SampledToken::Content(raw_token); + if self.model.is_eog_token(&raw_as_sampled) { + outcome.eog_seen = true; + break; + } + + self.batch.clear(); + self.batch.add(&raw_as_sampled, position, &[0], true)?; + position += 1; + + self.context.decode(self.batch)?; + } + + for ingest_outcome in self.classifier.flush() { + let is_eog = self.model.is_eog_token(&ingest_outcome.sampled_token); + if is_eog { + outcome.eog_seen = true; + } else { + outcome.generated_raw.push_str(&ingest_outcome.raw_piece); + } + record_outcome(&ingest_outcome, &mut outcome, is_eog); + } + + Ok(outcome) + } +} + +fn record_outcome(ingest: &IngestOutcome, outcome: &mut ClassifySampleLoopOutcome, is_eog: bool) { + match ingest.sampled_token { + SampledToken::Content(_) => { + outcome.observed_content += 1; + if !is_eog { + outcome.content_stream.push_str(&ingest.visible_piece); + } + } + SampledToken::Reasoning(_) => { + outcome.observed_reasoning += 1; + if !is_eog { + outcome.reasoning_stream.push_str(&ingest.visible_piece); + } + } + SampledToken::ToolCall(_) => { + outcome.observed_tool_call += 1; + } + SampledToken::Undeterminable(_) => { + outcome.observed_undeterminable += 1; + } + } +} diff --git a/llama-cpp-bindings-tests/src/test_fixture.rs b/llama-cpp-bindings-tests/src/fixture_session.rs similarity index 56% rename from llama-cpp-bindings-tests/src/test_fixture.rs rename to llama-cpp-bindings-tests/src/fixture_session.rs index b747f02f..37993878 100644 --- a/llama-cpp-bindings-tests/src/test_fixture.rs +++ b/llama-cpp-bindings-tests/src/fixture_session.rs @@ -1,6 +1,7 @@ -//! Process-wide cached fixture for LLM-backed integration tests. - +use std::sync::Arc; +use std::sync::Mutex; use std::sync::OnceLock; +use std::sync::Weak; use anyhow::Result; use llama_cpp_bindings::llama_backend::LlamaBackend; @@ -9,135 +10,148 @@ use llama_cpp_bindings::model::params::LlamaModelParams; use llama_cpp_bindings::mtmd::MtmdContext; use llama_cpp_bindings::mtmd::MtmdContextParams; +use crate::gpu_backend::inference_model_params; +use crate::gpu_backend::require_compiled_backends_present; use crate::test_model; -/// Shared test resources reused across LLM-backed integration tests in a single process. -/// -/// The backend and the default model load eagerly on first access; the embedding model and -/// multimodal context load lazily, only when a test asks for them. The fixture lives for the -/// duration of the test process so the GGUF files are mapped into memory exactly once. -pub struct TestFixture { - backend: LlamaBackend, - default_model: LlamaModel, - embedding_model: OnceLock, +static SHARED: Mutex> = Mutex::new(Weak::new()); + +struct FixtureSessionInner { mtmd_context: OnceLock, + embedding_model: OnceLock, + default_model: LlamaModel, + backend: LlamaBackend, } -impl TestFixture { - /// Returns the process-wide fixture, loading on first call. - /// - /// # Panics - /// Panics if the backend or default model cannot be loaded — that is an - /// unrecoverable test-setup failure and there is no meaningful continuation. - #[must_use] - pub fn shared() -> &'static Self { - static FIXTURE: OnceLock = OnceLock::new(); - - if let Some(fixture) = FIXTURE.get() { - return fixture; - } - - let fixture = Self::load().expect("test fixture: load failed"); - let _ = FIXTURE.set(fixture); - - FIXTURE.get().expect("test fixture: just set above") - } - +impl FixtureSessionInner { fn load() -> Result { let backend = LlamaBackend::init()?; + require_compiled_backends_present()?; let default_model = Self::load_default_model(&backend)?; Ok(Self { - backend, - default_model, - embedding_model: OnceLock::new(), mtmd_context: OnceLock::new(), + embedding_model: OnceLock::new(), + default_model, + backend, }) } fn load_default_model(backend: &LlamaBackend) -> Result { let path = test_model::download_model()?; - let params = LlamaModelParams::default(); + let params = inference_model_params(); Ok(LlamaModel::load_from_file(backend, &path, ¶ms)?) } - /// Returns the backend shared by every cached resource on this fixture. + fn load_embedding_model(&self) -> Result { + let path = test_model::download_embedding_model()?; + let params = LlamaModelParams::default(); + + Ok(LlamaModel::load_from_file(&self.backend, &path, ¶ms)?) + } + + fn load_mtmd_context(&self) -> Result { + let mmproj_path = test_model::download_mmproj()?; + let mmproj_str = mmproj_path + .to_str() + .ok_or_else(|| anyhow::anyhow!("mmproj path is not valid UTF-8"))?; + let params = MtmdContextParams::default(); + + Ok(MtmdContext::init_from_file( + mmproj_str, + &self.default_model, + ¶ms, + )?) + } +} + +pub struct FixtureSession { + inner: Arc, +} + +impl FixtureSession { + /// Opens a session against the shared fixture, loading on first call or + /// after the previous session has been fully dropped. + /// + /// # Errors + /// Returns an error if the backend or default model cannot be loaded. + /// + /// # Panics + /// Panics if the shared mutex is poisoned by a prior load failure. + pub fn open() -> Result { + let inner = { + let mut shared = SHARED.lock().expect("fixture singleton mutex poisoned"); + if let Some(existing) = shared.upgrade() { + existing + } else { + let new_inner = Arc::new(FixtureSessionInner::load()?); + *shared = Arc::downgrade(&new_inner); + new_inner + } + }; + + Ok(Self { inner }) + } + #[must_use] - pub const fn backend(&self) -> &LlamaBackend { - &self.backend + pub fn backend(&self) -> &LlamaBackend { + &self.inner.backend } - /// Returns the default test model. #[must_use] - pub const fn default_model(&self) -> &LlamaModel { - &self.default_model + pub fn default_model(&self) -> &LlamaModel { + &self.inner.default_model } /// Returns the embedding model, loading it on first call. /// /// # Errors - /// Returns an error if the required environment variables are not set or the model - /// cannot be downloaded or loaded. + /// Returns an error if the required environment variables are not set or the + /// model cannot be downloaded or loaded. /// /// # Panics - /// Panics only if the just-stored value cannot be read back (impossible in practice). + /// Panics only if the just-stored value cannot be read back, which cannot + /// happen in practice. pub fn embedding_model(&self) -> Result<&LlamaModel> { - if let Some(model) = self.embedding_model.get() { + if let Some(model) = self.inner.embedding_model.get() { return Ok(model); } - let model = self.load_embedding_model()?; - let _ = self.embedding_model.set(model); + let model = self.inner.load_embedding_model()?; + let _ = self.inner.embedding_model.set(model); Ok(self + .inner .embedding_model .get() - .expect("test fixture: embedding model just set")) - } - - fn load_embedding_model(&self) -> Result { - let path = test_model::download_embedding_model()?; - let params = LlamaModelParams::default(); - - Ok(LlamaModel::load_from_file(&self.backend, &path, ¶ms)?) + .expect("embedding model just set")) } /// Returns the multimodal context, loading it on first call. /// /// # Errors - /// Returns an error if `LLAMA_TEST_HF_MMPROJ` is unset or the context cannot be initialized. + /// Returns an error if `LLAMA_TEST_HF_MMPROJ` is unset or the context cannot + /// be initialized. /// /// # Panics - /// Panics only if the just-stored value cannot be read back (impossible in practice). + /// Panics only if the just-stored value cannot be read back, which cannot + /// happen in practice. pub fn mtmd_context(&self) -> Result<&MtmdContext> { if !test_model::has_mmproj() { anyhow::bail!("mtmd tests require LLAMA_TEST_HF_MMPROJ to be set"); } - if let Some(ctx) = self.mtmd_context.get() { + if let Some(ctx) = self.inner.mtmd_context.get() { return Ok(ctx); } - let ctx = self.load_mtmd_context()?; - let _ = self.mtmd_context.set(ctx); + let ctx = self.inner.load_mtmd_context()?; + let _ = self.inner.mtmd_context.set(ctx); Ok(self + .inner .mtmd_context .get() - .expect("test fixture: mtmd context just set")) - } - - fn load_mtmd_context(&self) -> Result { - let mmproj_path = test_model::download_mmproj()?; - let mmproj_str = mmproj_path - .to_str() - .ok_or_else(|| anyhow::anyhow!("mmproj path is not valid UTF-8"))?; - let params = MtmdContextParams::default(); - - Ok(MtmdContext::init_from_file( - mmproj_str, - &self.default_model, - ¶ms, - )?) + .expect("mtmd context just set")) } } diff --git a/llama-cpp-bindings-tests/src/gpu_backend.rs b/llama-cpp-bindings-tests/src/gpu_backend.rs new file mode 100644 index 00000000..bd9b5f8e --- /dev/null +++ b/llama-cpp-bindings-tests/src/gpu_backend.rs @@ -0,0 +1,166 @@ +use anyhow::Result; +#[cfg(any( + test, + feature = "cuda", + feature = "cuda-no-vmm", + feature = "metal", + feature = "vulkan", + feature = "rocm", +))] +use llama_cpp_bindings::llama_backend_device::LlamaBackendDevice; +use llama_cpp_bindings::llama_backend_device::list_llama_ggml_backend_devices; +use llama_cpp_bindings::model::params::LlamaModelParams; + +#[must_use] +pub fn inference_model_params() -> LlamaModelParams { + let params = LlamaModelParams::default(); + + #[cfg(any( + feature = "cuda", + feature = "cuda-no-vmm", + feature = "metal", + feature = "vulkan", + feature = "rocm", + ))] + let params = params.with_n_gpu_layers(999); + + params +} + +/// Confirms every compile-time backend feature has a matching ggml backend registered at runtime. +/// +/// Always asserts at least the CPU backend is registered (any llama.cpp build registers it); +/// when a GPU backend feature is enabled, also asserts the corresponding GPU backend is present. +/// +/// # Errors +/// +/// Returns an error when no ggml backends are registered, or when a compiled-in GPU backend +/// feature has no matching device. The error message names the missing backend(s) and lists +/// the backends that *are* registered, so misconfiguration is easy to diagnose. +pub fn require_compiled_backends_present() -> Result<()> { + let devices = list_llama_ggml_backend_devices(); + + if devices.is_empty() { + anyhow::bail!("no ggml backends registered; even CPU-only builds register a CPU backend"); + } + + #[cfg(feature = "cuda")] + require_backend(&devices, "cuda", &["CUDA"])?; + #[cfg(feature = "cuda-no-vmm")] + require_backend(&devices, "cuda-no-vmm", &["CUDA"])?; + #[cfg(feature = "metal")] + require_backend(&devices, "metal", &["Metal", "MTL"])?; + #[cfg(feature = "vulkan")] + require_backend(&devices, "vulkan", &["Vulkan"])?; + #[cfg(feature = "rocm")] + require_backend(&devices, "rocm", &["HIP", "ROCm"])?; + + Ok(()) +} + +#[cfg(any( + test, + feature = "cuda", + feature = "cuda-no-vmm", + feature = "metal", + feature = "vulkan", + feature = "rocm", +))] +fn require_backend( + devices: &[LlamaBackendDevice], + feature: &str, + accepted_names: &[&str], +) -> Result<()> { + let found = devices.iter().any(|device| { + accepted_names + .iter() + .any(|wanted| device.backend.eq_ignore_ascii_case(wanted)) + }); + + if !found { + let summary: Vec = devices + .iter() + .map(|device| format!("{}/{:?}", device.backend, device.device_type)) + .collect(); + + anyhow::bail!( + "feature `{feature}` enabled but no matching backend ({}) is registered; available: [{}]", + accepted_names.join(" / "), + summary.join(", ") + ); + } + + Ok(()) +} + +#[cfg(test)] +mod tests { + use llama_cpp_bindings::llama_backend_device::LlamaBackendDevice; + use llama_cpp_bindings::llama_backend_device::LlamaBackendDeviceType; + + use super::require_backend; + + fn synthetic_device(backend: &str, device_type: LlamaBackendDeviceType) -> LlamaBackendDevice { + LlamaBackendDevice { + index: 0, + name: format!("{backend}0"), + description: "synthetic test device".to_owned(), + backend: backend.to_owned(), + memory_total: 0, + memory_free: 0, + device_type, + } + } + + use anyhow::Result; + use anyhow::anyhow; + + #[test] + fn require_backend_succeeds_when_backend_name_matches_case_insensitively() -> Result<()> { + let devices = vec![synthetic_device("cuda", LlamaBackendDeviceType::Gpu)]; + + require_backend(&devices, "cuda", &["CUDA"]) + } + + #[test] + fn require_backend_succeeds_with_any_of_multiple_accepted_names() -> Result<()> { + let devices = vec![synthetic_device("HIP", LlamaBackendDeviceType::Gpu)]; + + require_backend(&devices, "rocm", &["HIP", "ROCm"]) + } + + #[test] + fn require_backend_fails_with_message_naming_feature_and_accepted_names_when_missing() + -> Result<()> { + let devices = vec![synthetic_device("Vulkan", LlamaBackendDeviceType::Gpu)]; + + let error = require_backend(&devices, "cuda", &["CUDA"]) + .err() + .ok_or_else(|| anyhow!("expected error when CUDA missing"))?; + + let message = format!("{error:#}"); + + if !message.contains("`cuda`") { + return Err(anyhow!("missing feature name: {message}")); + } + if !message.contains("CUDA") { + return Err(anyhow!("missing accepted name: {message}")); + } + if !message.contains("Vulkan") { + return Err(anyhow!("missing actual-backend summary: {message}")); + } + + Ok(()) + } + + #[test] + fn require_backend_fails_when_devices_list_is_empty() -> Result<()> { + let devices: Vec = Vec::new(); + + if require_backend(&devices, "metal", &["Metal"]).is_ok() { + return Err(anyhow!("expected Err for empty device list")); + } + + Ok(()) + } +} diff --git a/llama-cpp-bindings-tests/src/lib.rs b/llama-cpp-bindings-tests/src/lib.rs index 50c951f8..bda23c56 100644 --- a/llama-cpp-bindings-tests/src/lib.rs +++ b/llama-cpp-bindings-tests/src/lib.rs @@ -4,7 +4,9 @@ //! exists so production code in `llama-cpp-bindings` stays free of test-only //! dependencies (`anyhow`, `hf-hub`, `serial_test`, …) and helpers. -pub mod test_fixture; +pub mod classify_sample_loop; +pub mod fixture_session; +pub mod gpu_backend; pub mod test_model; -pub use test_fixture::TestFixture; +pub use fixture_session::FixtureSession; diff --git a/llama-cpp-bindings-tests/src/test_model.rs b/llama-cpp-bindings-tests/src/test_model.rs index e4ceb7d8..934f1d9e 100644 --- a/llama-cpp-bindings-tests/src/test_model.rs +++ b/llama-cpp-bindings-tests/src/test_model.rs @@ -167,6 +167,12 @@ mod tests { #[test] #[serial_test::serial] fn download_model_returns_path_with_env_set() { + if std::env::var("LLAMA_TEST_HF_REPO").is_err() + || std::env::var("LLAMA_TEST_HF_MODEL").is_err() + { + return; + } + let result = super::download_model(); assert!(result.is_ok()); @@ -175,6 +181,12 @@ mod tests { #[test] #[serial_test::serial] fn download_embedding_model_returns_path_with_env_set() { + if std::env::var("LLAMA_TEST_HF_EMBED_REPO").is_err() + || std::env::var("LLAMA_TEST_HF_EMBED_MODEL").is_err() + { + return; + } + let result = super::download_embedding_model(); assert!(result.is_ok()); @@ -183,14 +195,25 @@ mod tests { #[test] #[serial_test::serial] fn download_encoder_model_returns_path_with_env_set() { + if std::env::var("LLAMA_TEST_HF_ENCODER_REPO").is_err() + || std::env::var("LLAMA_TEST_HF_ENCODER_MODEL").is_err() + { + return; + } + let result = super::download_encoder_model(); assert!(result.is_ok()); } + #[cfg(feature = "multimodal_capable")] #[test] #[serial_test::serial] fn download_mmproj_returns_path_when_env_set() { + if std::env::var("LLAMA_TEST_HF_REPO").is_err() { + return; + } + let _guard = EnvVarGuard::set("LLAMA_TEST_HF_MMPROJ", "mmproj-F16.gguf"); let result = super::download_mmproj(); diff --git a/llama-cpp-bindings-tests/tests/constrained_decoding.rs b/llama-cpp-bindings-tests/tests/constrained_decoding.rs index 79c855b2..6be1014f 100644 --- a/llama-cpp-bindings-tests/tests/constrained_decoding.rs +++ b/llama-cpp-bindings-tests/tests/constrained_decoding.rs @@ -1,23 +1,24 @@ use std::io::Write; use anyhow::Result; +use llama_cpp_bindings::context::LlamaContext; use llama_cpp_bindings::context::params::LlamaContextParams; use llama_cpp_bindings::llama_batch::LlamaBatch; use llama_cpp_bindings::model::AddBos; use llama_cpp_bindings::sampled_token::SampledToken; use llama_cpp_bindings::sampling::LlamaSampler; -use llama_cpp_bindings_tests::TestFixture; +use llama_cpp_bindings_tests::FixtureSession; #[test] fn json_schema_constrains_output() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let prompt = "The weather in Paris is sunny and 22 degrees. Extract as JSON:\n"; let ctx_params = LlamaContextParams::default(); - let mut ctx = model.new_context(backend, ctx_params)?; + let mut ctx = LlamaContext::from_model(model, backend, ctx_params)?; let tokens_list = model.str_to_token(prompt, AddBos::Always)?; diff --git a/llama-cpp-bindings-tests/tests/context.rs b/llama-cpp-bindings-tests/tests/context.rs index 934b55d2..fe7ba7c8 100644 --- a/llama-cpp-bindings-tests/tests/context.rs +++ b/llama-cpp-bindings-tests/tests/context.rs @@ -6,24 +6,25 @@ use std::sync::atomic::AtomicBool; use anyhow::Result; use llama_cpp_bindings::DecodeError; use llama_cpp_bindings::LogitsError; +use llama_cpp_bindings::context::LlamaContext; use llama_cpp_bindings::context::params::LlamaContextParams; use llama_cpp_bindings::llama_batch::LlamaBatch; use llama_cpp_bindings::model::AddBos; use llama_cpp_bindings::model::LlamaLoraAdapter; use llama_cpp_bindings::model::LlamaModel; -use llama_cpp_bindings::model::params::LlamaModelParams; -use llama_cpp_bindings_tests::TestFixture; +use llama_cpp_bindings_tests::FixtureSession; +use llama_cpp_bindings_tests::gpu_backend::inference_model_params; use llama_cpp_bindings_tests::test_model; use serial_test::serial; #[test] #[serial] fn context_creation_and_properties() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let context = model.new_context(backend, ctx_params)?; + let context = LlamaContext::from_model(model, backend, ctx_params)?; assert!(context.n_ctx() > 0); assert!(context.n_batch() > 0); @@ -35,11 +36,11 @@ fn context_creation_and_properties() -> Result<()> { #[test] #[serial] fn decode_and_get_logits() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let tokens = model.str_to_token("hello", AddBos::Always)?; let mut batch = LlamaBatch::new(512, 1)?; batch.add_sequence(&tokens, 0, false)?; @@ -56,11 +57,11 @@ fn decode_and_get_logits() -> Result<()> { #[test] #[serial] fn timings_work() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; context.reset_timings(); let timings = context.timings(); @@ -72,11 +73,11 @@ fn timings_work() -> Result<()> { #[test] #[serial] fn token_data_array_has_entries_after_decode() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let tokens = model.str_to_token("hello", AddBos::Always)?; let mut batch = LlamaBatch::new(512, 1)?; batch.add_sequence(&tokens, 0, false)?; @@ -92,11 +93,11 @@ fn token_data_array_has_entries_after_decode() -> Result<()> { #[test] #[serial] fn get_logits_ith_returns_valid_slice() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let tokens = model.str_to_token("hello", AddBos::Always)?; let last_index = i32::try_from(tokens.len() - 1)?; let mut batch = LlamaBatch::new(512, 1)?; @@ -113,11 +114,11 @@ fn get_logits_ith_returns_valid_slice() -> Result<()> { #[test] #[serial] fn token_data_array_ith_returns_valid_data() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let tokens = model.str_to_token("hello", AddBos::Always)?; let last_index = i32::try_from(tokens.len() - 1)?; let mut batch = LlamaBatch::new(512, 1)?; @@ -137,13 +138,13 @@ fn token_data_array_ith_returns_valid_data() -> Result<()> { #[test] #[serial] fn embeddings_ith_returns_error_when_embeddings_disabled() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default() .with_n_ctx(NonZeroU32::new(512)) .with_embeddings(false); - let context = model.new_context(backend, ctx_params)?; + let context = LlamaContext::from_model(model, backend, ctx_params)?; let result = context.embeddings_ith(0); @@ -155,13 +156,13 @@ fn embeddings_ith_returns_error_when_embeddings_disabled() -> Result<()> { #[test] #[serial] fn embeddings_seq_ith_returns_error_when_embeddings_disabled() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default() .with_n_ctx(NonZeroU32::new(512)) .with_embeddings(false); - let context = model.new_context(backend, ctx_params)?; + let context = LlamaContext::from_model(model, backend, ctx_params)?; let result = context.embeddings_seq_ith(0); @@ -173,11 +174,11 @@ fn embeddings_seq_ith_returns_error_when_embeddings_disabled() -> Result<()> { #[test] #[serial] fn candidates_returns_n_vocab_entries() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let tokens = model.str_to_token("hello", AddBos::Always)?; let mut batch = LlamaBatch::new(512, 1)?; batch.add_sequence(&tokens, 0, false)?; @@ -193,11 +194,11 @@ fn candidates_returns_n_vocab_entries() -> Result<()> { #[test] #[serial] fn debug_format_contains_struct_name() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let context = model.new_context(backend, ctx_params)?; + let context = LlamaContext::from_model(model, backend, ctx_params)?; let debug_output = format!("{context:?}"); assert!(debug_output.contains("LlamaContext")); @@ -208,13 +209,13 @@ fn debug_format_contains_struct_name() -> Result<()> { #[test] #[serial] fn decode_with_embeddings_enabled() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.embedding_model()?; let ctx_params = LlamaContextParams::default() .with_n_ctx(NonZeroU32::new(512)) .with_embeddings(true); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let tokens = model.str_to_token("hello", AddBos::Always)?; let mut batch = LlamaBatch::new(512, 1)?; batch.add_sequence(&tokens, 0, false)?; @@ -229,13 +230,13 @@ fn decode_with_embeddings_enabled() -> Result<()> { #[test] #[serial] fn embeddings_seq_ith_returns_valid_embeddings() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.embedding_model()?; let ctx_params = LlamaContextParams::default() .with_n_ctx(NonZeroU32::new(512)) .with_embeddings(true); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let tokens = model.str_to_token("hello", AddBos::Always)?; let mut batch = LlamaBatch::new(512, 1)?; batch.add_sequence(&tokens, 0, false)?; @@ -251,14 +252,14 @@ fn embeddings_seq_ith_returns_valid_embeddings() -> Result<()> { #[test] #[serial] fn multi_sequence_embeddings_returns_one_embedding_per_sequence() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.embedding_model()?; let ctx_params = LlamaContextParams::default() .with_n_ctx(NonZeroU32::new(512)) .with_n_seq_max(4) .with_embeddings(true); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let inputs = [ "alpha is here", @@ -316,14 +317,14 @@ fn multi_sequence_embeddings_returns_one_embedding_per_sequence() -> Result<()> #[test] #[serial] fn embeddings_returns_distinct_values_when_reused_batch_has_extra_capacity() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.embedding_model()?; let ctx_params = LlamaContextParams::default() .with_n_ctx(NonZeroU32::new(512)) .with_n_seq_max(4) .with_embeddings(true); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let iterations = [ [ @@ -388,13 +389,13 @@ fn embeddings_returns_distinct_values_when_reused_batch_has_extra_capacity() -> #[test] #[serial] fn embeddings_ith_returns_valid_embeddings() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.embedding_model()?; let ctx_params = LlamaContextParams::default() .with_n_ctx(NonZeroU32::new(512)) .with_embeddings(true); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let tokens = model.str_to_token("hello", AddBos::Always)?; let last_index = i32::try_from(tokens.len() - 1)?; let mut batch = LlamaBatch::new(512, 1)?; @@ -411,11 +412,11 @@ fn embeddings_ith_returns_valid_embeddings() -> Result<()> { #[test] #[serial] fn candidates_ith_returns_n_vocab_entries() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let tokens = model.str_to_token("hello", AddBos::Always)?; let last_index = i32::try_from(tokens.len() - 1)?; let mut batch = LlamaBatch::new(512, 1)?; @@ -432,11 +433,11 @@ fn candidates_ith_returns_n_vocab_entries() -> Result<()> { #[test] #[serial] fn lora_adapter_remove_succeeds_with_no_adapters() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let context = model.new_context(backend, ctx_params)?; + let context = LlamaContext::from_model(model, backend, ctx_params)?; let mut adapter = LlamaLoraAdapter { lora_adapter: NonNull::dangling(), }; @@ -451,11 +452,11 @@ fn lora_adapter_remove_succeeds_with_no_adapters() -> Result<()> { #[test] #[serial] fn encode_on_non_encoder_model_returns_error() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let tokens = model.str_to_token("hello", AddBos::Always)?; let mut batch = LlamaBatch::new(512, 1)?; batch.add_sequence(&tokens, 0, false)?; @@ -470,11 +471,11 @@ fn encode_on_non_encoder_model_returns_error() -> Result<()> { #[test] #[serial] fn lora_adapter_set_with_dangling_pointer_succeeds_or_errors() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let context = model.new_context(backend, ctx_params)?; + let context = LlamaContext::from_model(model, backend, ctx_params)?; let mut adapter = LlamaLoraAdapter { lora_adapter: NonNull::dangling(), }; @@ -489,13 +490,13 @@ fn lora_adapter_set_with_dangling_pointer_succeeds_or_errors() -> Result<()> { #[test] #[serial] fn embeddings_ith_returns_null_embedding_error_for_non_embedding_token() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.embedding_model()?; let ctx_params = LlamaContextParams::default() .with_n_ctx(NonZeroU32::new(512)) .with_embeddings(true); - let context = model.new_context(backend, ctx_params)?; + let context = LlamaContext::from_model(model, backend, ctx_params)?; let result = context.embeddings_ith(999); @@ -507,13 +508,13 @@ fn embeddings_ith_returns_null_embedding_error_for_non_embedding_token() -> Resu #[test] #[serial] fn embeddings_seq_ith_returns_null_embedding_error_for_invalid_seq() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default() .with_n_ctx(NonZeroU32::new(512)) .with_embeddings(true); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let tokens = model.str_to_token("hello", AddBos::Always)?; let mut batch = LlamaBatch::new(512, 1)?; batch.add_sequence(&tokens, 0, false)?; @@ -529,11 +530,11 @@ fn embeddings_seq_ith_returns_null_embedding_error_for_invalid_seq() -> Result<( #[test] #[serial] fn decode_empty_batch_returns_error() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let mut batch = LlamaBatch::new(512, 1)?; let result = context.decode(&mut batch); @@ -546,15 +547,15 @@ fn decode_empty_batch_returns_error() -> Result<()> { #[test] #[serial] fn encode_succeeds_with_encoder_model() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model_path = test_model::download_encoder_model()?; - let model_params = LlamaModelParams::default(); + let model_params = inference_model_params(); let model = LlamaModel::load_from_file(backend, &model_path, &model_params)?; let ctx_params = LlamaContextParams::default() .with_n_ctx(NonZeroU32::new(512)) .with_embeddings(true); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(&model, backend, ctx_params)?; let tokens = model.str_to_token("hello", AddBos::Never)?; let mut batch = LlamaBatch::new(512, 1)?; batch.add_sequence(&tokens, 0, false)?; @@ -569,11 +570,11 @@ fn encode_succeeds_with_encoder_model() -> Result<()> { #[test] #[serial] fn set_abort_flag_aborts_decode() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let abort_flag = Arc::new(AtomicBool::new(true)); context.set_abort_flag(abort_flag); @@ -591,11 +592,11 @@ fn set_abort_flag_aborts_decode() -> Result<()> { #[test] #[serial] fn set_abort_flag_false_allows_decode() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let abort_flag = Arc::new(AtomicBool::new(false)); context.set_abort_flag(abort_flag); @@ -613,11 +614,11 @@ fn set_abort_flag_false_allows_decode() -> Result<()> { #[test] #[serial] fn clear_abort_callback_allows_decode_with_flag_true() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let abort_flag = Arc::new(AtomicBool::new(true)); context.set_abort_flag(abort_flag); context.clear_abort_callback(); @@ -636,11 +637,11 @@ fn clear_abort_callback_allows_decode_with_flag_true() -> Result<()> { #[test] #[serial] fn synchronize_completes_without_panic() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let context = model.new_context(backend, ctx_params)?; + let context = LlamaContext::from_model(model, backend, ctx_params)?; context.synchronize(); @@ -650,11 +651,11 @@ fn synchronize_completes_without_panic() -> Result<()> { #[test] #[serial] fn detach_threadpool_completes_without_panic() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let context = model.new_context(backend, ctx_params)?; + let context = LlamaContext::from_model(model, backend, ctx_params)?; context.detach_threadpool(); @@ -664,11 +665,11 @@ fn detach_threadpool_completes_without_panic() -> Result<()> { #[test] #[serial] fn get_logits_ith_returns_token_not_initialized_for_unknown_index() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let context = model.new_context(backend, ctx_params)?; + let context = LlamaContext::from_model(model, backend, ctx_params)?; let result = context.get_logits_ith(7); @@ -680,11 +681,11 @@ fn get_logits_ith_returns_token_not_initialized_for_unknown_index() -> Result<() #[test] #[serial] fn get_logits_ith_returns_token_index_exceeds_context_for_huge_index() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(64)); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let huge_index = i32::try_from(context.n_ctx())?; context.mark_logits_initialized(huge_index); diff --git a/llama-cpp-bindings-tests/tests/context_kv_cache.rs b/llama-cpp-bindings-tests/tests/context_kv_cache.rs index 69cfa9ee..0095bff6 100644 --- a/llama-cpp-bindings-tests/tests/context_kv_cache.rs +++ b/llama-cpp-bindings-tests/tests/context_kv_cache.rs @@ -2,21 +2,22 @@ use std::num::NonZeroU8; use std::num::NonZeroU32; use anyhow::Result; +use llama_cpp_bindings::context::LlamaContext; use llama_cpp_bindings::context::kv_cache::KvCacheConversionError; use llama_cpp_bindings::context::params::LlamaContextParams; use llama_cpp_bindings::llama_batch::LlamaBatch; use llama_cpp_bindings::model::AddBos; -use llama_cpp_bindings_tests::TestFixture; +use llama_cpp_bindings_tests::FixtureSession; use serial_test::serial; #[test] #[serial] fn clear_kv_cache_resets_positions() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let tokens = model.str_to_token("Hello world", AddBos::Always)?; let mut batch = LlamaBatch::new(512, 1)?; @@ -32,11 +33,11 @@ fn clear_kv_cache_resets_positions() -> Result<()> { #[test] #[serial] fn kv_cache_seq_pos_max_is_non_negative_after_decode() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let tokens = model.str_to_token("Hello world", AddBos::Always)?; let mut batch = LlamaBatch::new(512, 1)?; @@ -51,11 +52,11 @@ fn kv_cache_seq_pos_max_is_non_negative_after_decode() -> Result<()> { #[test] #[serial] fn clear_kv_cache_seq_with_range() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let tokens = model.str_to_token("Hello world", AddBos::Always)?; let mut batch = LlamaBatch::new(512, 1)?; @@ -71,11 +72,11 @@ fn clear_kv_cache_seq_with_range() -> Result<()> { #[test] #[serial] fn copy_kv_cache_seq_succeeds() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let tokens = model.str_to_token("Hello world", AddBos::Always)?; let mut batch = LlamaBatch::new(512, 1)?; @@ -91,11 +92,11 @@ fn copy_kv_cache_seq_succeeds() -> Result<()> { #[test] #[serial] fn copy_cache_executes_without_crash() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let tokens = model.str_to_token("Hello world", AddBos::Always)?; let mut batch = LlamaBatch::new(512, 1)?; @@ -108,14 +109,15 @@ fn copy_cache_executes_without_crash() -> Result<()> { Ok(()) } +#[cfg(feature = "mrope_model")] #[test] #[serial] fn kv_cache_seq_add_returns_error_for_mrope_model() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let tokens = model.str_to_token("Hello world", AddBos::Always)?; let mut batch = LlamaBatch::new(512, 1)?; @@ -129,14 +131,15 @@ fn kv_cache_seq_add_returns_error_for_mrope_model() -> Result<()> { Ok(()) } +#[cfg(feature = "mrope_model")] #[test] #[serial] fn kv_cache_seq_div_returns_error_for_mrope_model() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let tokens = model.str_to_token("Hello world", AddBos::Always)?; let mut batch = LlamaBatch::new(512, 1)?; @@ -154,11 +157,11 @@ fn kv_cache_seq_div_returns_error_for_mrope_model() -> Result<()> { #[test] #[serial] fn kv_cache_seq_keep_retains_specified_sequence() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let tokens = model.str_to_token("Hello world", AddBos::Always)?; let mut batch = LlamaBatch::new(512, 1)?; @@ -175,11 +178,11 @@ fn kv_cache_seq_keep_retains_specified_sequence() -> Result<()> { #[test] #[serial] fn copy_kv_cache_seq_with_explicit_range() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let tokens = model.str_to_token("Hello world", AddBos::Always)?; let mut batch = LlamaBatch::new(512, 1)?; @@ -196,11 +199,11 @@ fn copy_kv_cache_seq_with_explicit_range() -> Result<()> { #[test] #[serial] fn kv_cache_seq_add_succeeds_on_embedding_model() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.embedding_model()?; let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let tokens = model.str_to_token("Hello world", AddBos::Always)?; let mut batch = LlamaBatch::new(512, 1)?; @@ -217,11 +220,11 @@ fn kv_cache_seq_add_succeeds_on_embedding_model() -> Result<()> { #[test] #[serial] fn kv_cache_seq_div_succeeds_on_embedding_model() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.embedding_model()?; let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let tokens = model.str_to_token("Hello world", AddBos::Always)?; let mut batch = LlamaBatch::new(512, 1)?; @@ -239,11 +242,11 @@ fn kv_cache_seq_div_succeeds_on_embedding_model() -> Result<()> { #[test] #[serial] fn kv_cache_seq_pos_max_returns_negative_one_for_unused_seq() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let context = model.new_context(backend, ctx_params)?; + let context = LlamaContext::from_model(model, backend, ctx_params)?; let result = context.kv_cache_seq_pos_max(999); @@ -255,11 +258,11 @@ fn kv_cache_seq_pos_max_returns_negative_one_for_unused_seq() -> Result<()> { #[test] #[serial] fn copy_kv_cache_seq_rejects_p0_exceeding_i32_max() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let result = context.copy_kv_cache_seq(0, 1, Some(u32::MAX), None); @@ -274,11 +277,11 @@ fn copy_kv_cache_seq_rejects_p0_exceeding_i32_max() -> Result<()> { #[test] #[serial] fn copy_kv_cache_seq_rejects_p1_exceeding_i32_max() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let result = context.copy_kv_cache_seq(0, 1, Some(0), Some(u32::MAX)); @@ -293,11 +296,11 @@ fn copy_kv_cache_seq_rejects_p1_exceeding_i32_max() -> Result<()> { #[test] #[serial] fn clear_kv_cache_seq_rejects_src_exceeding_i32_max() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let result = context.clear_kv_cache_seq(Some(u32::MAX), None, None); @@ -312,11 +315,11 @@ fn clear_kv_cache_seq_rejects_src_exceeding_i32_max() -> Result<()> { #[test] #[serial] fn clear_kv_cache_seq_rejects_p0_exceeding_i32_max() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let result = context.clear_kv_cache_seq(Some(0), Some(u32::MAX), None); @@ -331,11 +334,11 @@ fn clear_kv_cache_seq_rejects_p0_exceeding_i32_max() -> Result<()> { #[test] #[serial] fn clear_kv_cache_seq_rejects_p1_exceeding_i32_max() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let result = context.clear_kv_cache_seq(Some(0), Some(0), Some(u32::MAX)); @@ -350,11 +353,11 @@ fn clear_kv_cache_seq_rejects_p1_exceeding_i32_max() -> Result<()> { #[test] #[serial] fn kv_cache_seq_add_rejects_p0_exceeding_i32_max() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let result = context.kv_cache_seq_add(0, Some(u32::MAX), None, 1); @@ -369,11 +372,11 @@ fn kv_cache_seq_add_rejects_p0_exceeding_i32_max() -> Result<()> { #[test] #[serial] fn kv_cache_seq_add_rejects_p1_exceeding_i32_max() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let result = context.kv_cache_seq_add(0, Some(0), Some(u32::MAX), 1); @@ -388,11 +391,11 @@ fn kv_cache_seq_add_rejects_p1_exceeding_i32_max() -> Result<()> { #[test] #[serial] fn kv_cache_seq_div_rejects_p0_exceeding_i32_max() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let divisor = NonZeroU8::new(2).ok_or_else(|| anyhow::anyhow!("2 is non-zero"))?; let result = context.kv_cache_seq_div(0, Some(u32::MAX), None, divisor); @@ -408,11 +411,11 @@ fn kv_cache_seq_div_rejects_p0_exceeding_i32_max() -> Result<()> { #[test] #[serial] fn kv_cache_seq_div_rejects_p1_exceeding_i32_max() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let divisor = NonZeroU8::new(2).ok_or_else(|| anyhow::anyhow!("2 is non-zero"))?; let result = context.kv_cache_seq_div(0, Some(0), Some(u32::MAX), divisor); diff --git a/llama-cpp-bindings-tests/tests/context_session.rs b/llama-cpp-bindings-tests/tests/context_session.rs index 95ecfbc6..4c52260f 100644 --- a/llama-cpp-bindings-tests/tests/context_session.rs +++ b/llama-cpp-bindings-tests/tests/context_session.rs @@ -1,20 +1,21 @@ use std::num::NonZeroU32; use anyhow::Result; +use llama_cpp_bindings::context::LlamaContext; use llama_cpp_bindings::context::params::LlamaContextParams; use llama_cpp_bindings::llama_batch::LlamaBatch; use llama_cpp_bindings::model::AddBos; -use llama_cpp_bindings_tests::TestFixture; +use llama_cpp_bindings_tests::FixtureSession; use serial_test::serial; #[test] #[serial] fn save_and_load_session_file() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let tokens = model.str_to_token("Hello world", AddBos::Always)?; let mut batch = LlamaBatch::new(512, 1)?; @@ -35,11 +36,11 @@ fn save_and_load_session_file() -> Result<()> { #[test] #[serial] fn get_state_size_is_positive() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let context = model.new_context(backend, ctx_params)?; + let context = LlamaContext::from_model(model, backend, ctx_params)?; assert!(context.get_state_size() > 0); @@ -49,11 +50,11 @@ fn get_state_size_is_positive() -> Result<()> { #[test] #[serial] fn state_seq_save_and_load_file_roundtrip() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let tokens = model.str_to_token("Hello world", AddBos::Always)?; let mut batch = LlamaBatch::new(512, 1)?; @@ -76,11 +77,11 @@ fn state_seq_save_and_load_file_roundtrip() -> Result<()> { #[test] #[serial] fn copy_state_data_and_set_state_data_roundtrip() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let tokens = model.str_to_token("Hello world", AddBos::Always)?; let mut batch = LlamaBatch::new(512, 1)?; @@ -101,11 +102,11 @@ fn copy_state_data_and_set_state_data_roundtrip() -> Result<()> { #[test] #[serial] fn state_load_file_with_nonexistent_file_returns_error() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let result = context.state_load_file("/nonexistent/session.bin", 512); @@ -117,11 +118,11 @@ fn state_load_file_with_nonexistent_file_returns_error() -> Result<()> { #[test] #[serial] fn state_seq_load_file_with_nonexistent_file_returns_error() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let result = context.state_seq_load_file("/nonexistent/seq_state.bin", 0, 512); @@ -133,11 +134,11 @@ fn state_seq_load_file_with_nonexistent_file_returns_error() -> Result<()> { #[test] #[serial] fn state_save_file_to_invalid_directory_returns_failed_to_save() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let context = model.new_context(backend, ctx_params)?; + let context = LlamaContext::from_model(model, backend, ctx_params)?; let result = context.state_save_file("/nonexistent_dir/session.bin", &[]); @@ -149,11 +150,11 @@ fn state_save_file_to_invalid_directory_returns_failed_to_save() -> Result<()> { #[test] #[serial] fn state_seq_save_file_to_invalid_directory_returns_failed_to_save() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let context = model.new_context(backend, ctx_params)?; + let context = LlamaContext::from_model(model, backend, ctx_params)?; let result = context.state_seq_save_file("/nonexistent_dir/seq_state.bin", 0, &[]); @@ -165,11 +166,11 @@ fn state_seq_save_file_to_invalid_directory_returns_failed_to_save() -> Result<( #[test] #[serial] fn state_load_file_with_zero_max_tokens_returns_error() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let tokens = model.str_to_token("Hello world", AddBos::Always)?; let mut batch = LlamaBatch::new(512, 1)?; @@ -190,11 +191,11 @@ fn state_load_file_with_zero_max_tokens_returns_error() -> Result<()> { #[test] #[serial] fn state_seq_load_file_with_zero_max_tokens_returns_error() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let tokens = model.str_to_token("Hello world", AddBos::Always)?; let mut batch = LlamaBatch::new(512, 1)?; @@ -215,11 +216,11 @@ fn state_seq_load_file_with_zero_max_tokens_returns_error() -> Result<()> { #[test] #[serial] fn state_load_file_with_insufficient_max_tokens_returns_length_error() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let tokens = model.str_to_token( "Hello world this is a longer string for more tokens", @@ -243,11 +244,11 @@ fn state_load_file_with_insufficient_max_tokens_returns_length_error() -> Result #[test] #[serial] fn state_seq_load_file_with_insufficient_max_tokens_returns_length_error() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let tokens = model.str_to_token( "Hello world this is a longer string for more tokens", @@ -275,11 +276,11 @@ fn state_save_file_with_non_utf8_path_returns_error() -> Result<()> { use std::ffi::OsStr; use std::os::unix::ffi::OsStrExt; - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let context = model.new_context(backend, ctx_params)?; + let context = LlamaContext::from_model(model, backend, ctx_params)?; let non_utf8_path = std::path::Path::new(OsStr::from_bytes(b"/tmp/\xff\xfe.bin")); let result = context.state_save_file(non_utf8_path, &[]); @@ -296,11 +297,11 @@ fn state_load_file_with_non_utf8_path_returns_error() -> Result<()> { use std::ffi::OsStr; use std::os::unix::ffi::OsStrExt; - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let non_utf8_path = std::path::Path::new(OsStr::from_bytes(b"/tmp/\xff\xfe.bin")); let result = context.state_load_file(non_utf8_path, 512); @@ -317,11 +318,11 @@ fn state_seq_save_file_with_non_utf8_path_returns_error() -> Result<()> { use std::ffi::OsStr; use std::os::unix::ffi::OsStrExt; - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let context = model.new_context(backend, ctx_params)?; + let context = LlamaContext::from_model(model, backend, ctx_params)?; let non_utf8_path = std::path::Path::new(OsStr::from_bytes(b"/tmp/\xff\xfe.bin")); let result = context.state_seq_save_file(non_utf8_path, 0, &[]); @@ -338,11 +339,11 @@ fn state_seq_load_file_with_non_utf8_path_returns_error() -> Result<()> { use std::ffi::OsStr; use std::os::unix::ffi::OsStrExt; - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let non_utf8_path = std::path::Path::new(OsStr::from_bytes(b"/tmp/\xff\xfe.bin")); let result = context.state_seq_load_file(non_utf8_path, 0, 512); @@ -355,11 +356,11 @@ fn state_seq_load_file_with_non_utf8_path_returns_error() -> Result<()> { #[test] #[serial] fn state_save_file_with_null_byte_in_path_returns_error() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let context = model.new_context(backend, ctx_params)?; + let context = LlamaContext::from_model(model, backend, ctx_params)?; let path_with_null = std::path::Path::new("/tmp/foo\0bar.bin"); let result = context.state_save_file(path_with_null, &[]); @@ -372,11 +373,11 @@ fn state_save_file_with_null_byte_in_path_returns_error() -> Result<()> { #[test] #[serial] fn state_load_file_with_null_byte_in_path_returns_error() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let path_with_null = std::path::Path::new("/tmp/foo\0bar.bin"); let result = context.state_load_file(path_with_null, 512); @@ -389,11 +390,11 @@ fn state_load_file_with_null_byte_in_path_returns_error() -> Result<()> { #[test] #[serial] fn state_seq_save_file_with_null_byte_in_path_returns_error() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let context = model.new_context(backend, ctx_params)?; + let context = LlamaContext::from_model(model, backend, ctx_params)?; let path_with_null = std::path::Path::new("/tmp/foo\0bar.bin"); let result = context.state_seq_save_file(path_with_null, 0, &[]); @@ -406,11 +407,11 @@ fn state_seq_save_file_with_null_byte_in_path_returns_error() -> Result<()> { #[test] #[serial] fn state_seq_load_file_with_null_byte_in_path_returns_error() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let path_with_null = std::path::Path::new("/tmp/foo\0bar.bin"); let result = context.state_seq_load_file(path_with_null, 0, 512); @@ -425,11 +426,11 @@ fn state_seq_load_file_with_null_byte_in_path_returns_error() -> Result<()> { fn state_seq_get_size_ext_returns_size_for_decoded_sequence() -> Result<()> { use llama_cpp_bindings::context::llama_state_seq_flags::LlamaStateSeqFlags; - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let tokens = model.str_to_token("Hello world", AddBos::Always)?; let mut batch = LlamaBatch::new(512, 1)?; @@ -449,11 +450,11 @@ fn state_seq_get_size_ext_returns_size_for_decoded_sequence() -> Result<()> { fn state_seq_get_data_ext_and_set_data_ext_round_trip() -> Result<()> { use llama_cpp_bindings::context::llama_state_seq_flags::LlamaStateSeqFlags; - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let tokens = model.str_to_token("Hello world", AddBos::Always)?; let mut batch = LlamaBatch::new(512, 1)?; diff --git a/llama-cpp-bindings-tests/tests/deepseek_r1_8b_classifier_does_not_emit_reasoning_for_thinking_disabled_prompt.rs b/llama-cpp-bindings-tests/tests/deepseek_r1_8b_classifier_does_not_emit_reasoning_for_thinking_disabled_prompt.rs new file mode 100644 index 00000000..364717a7 --- /dev/null +++ b/llama-cpp-bindings-tests/tests/deepseek_r1_8b_classifier_does_not_emit_reasoning_for_thinking_disabled_prompt.rs @@ -0,0 +1,128 @@ +use std::num::NonZeroU32; + +use anyhow::Result; +use llama_cpp_bindings::context::LlamaContext; +use llama_cpp_bindings::context::params::LlamaContextParams; +use llama_cpp_bindings::llama_backend::LlamaBackend; +use llama_cpp_bindings::llama_batch::LlamaBatch; +use llama_cpp_bindings::model::AddBos; +use llama_cpp_bindings::model::LlamaModel; +use llama_cpp_bindings::sampling::LlamaSampler; +use llama_cpp_bindings_tests::classify_sample_loop::ClassifySampleLoop; +use llama_cpp_bindings_tests::gpu_backend::inference_model_params; +use llama_cpp_bindings_tests::gpu_backend::require_compiled_backends_present; +use llama_cpp_bindings_tests::test_model::download_file_from; + +const DEEPSEEK_R1_8B_REPO: &str = "unsloth/DeepSeek-R1-Distill-Llama-8B-GGUF"; +const DEEPSEEK_R1_8B_FILE: &str = "DeepSeek-R1-Distill-Llama-8B-Q4_K_M.gguf"; + +const MAX_GENERATED_TOKENS: i32 = 200; + +// DeepSeek-R1-Distill-Llama-8B has no native thinking-disabled mode in its +// chat template (R1 is a pure reasoner). This prompt manually closes the +// `` block before generation so the classifier starts in CONTENT — +// verifies the "spurious close in content section" path with this model's +// tokenizer and still produces zero Reasoning tokens. +const DEEPSEEK_R1_8B_THINKING_DISABLED_PROMPT: &str = "\ +<｜User｜>What is 2 + 2?<｜Assistant｜> + + + +"; + +const FORBIDDEN_MARKERS: &[&str] = &["", ""]; + +#[test] +fn deepseek_r1_8b_classifier_does_not_emit_reasoning_for_thinking_disabled_prompt() -> Result<()> { + let backend = LlamaBackend::init()?; + require_compiled_backends_present()?; + + let path = download_file_from(DEEPSEEK_R1_8B_REPO, DEEPSEEK_R1_8B_FILE)?; + let params = inference_model_params(); + let model = LlamaModel::load_from_file(&backend, &path, ¶ms)?; + + let mut classifier = model.sampled_token_classifier(); + let prompt_tokens = + model.str_to_token(DEEPSEEK_R1_8B_THINKING_DISABLED_PROMPT, AddBos::Never)?; + let prompt_token_count = u64::try_from(prompt_tokens.len())?; + + let mut batch = LlamaBatch::new(2048, 1)?; + classifier.feed_prompt_sequence_to_batch(&mut batch, &prompt_tokens, 0, false)?; + + let context_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(8192)); + let mut context = LlamaContext::from_model(&model, &backend, context_params)?; + + context.decode(&mut batch)?; + + let promoted = classifier.commit_prompt_tokens(); + assert_eq!(promoted, prompt_token_count); + + let mut sampler = LlamaSampler::chain_simple([ + LlamaSampler::penalties(64, 1.1, 0.0, 0.0), + LlamaSampler::top_k(40), + LlamaSampler::top_p(0.9, 1), + LlamaSampler::min_p(0.05, 1), + LlamaSampler::temp(0.7), + LlamaSampler::dist(0x00C0_FFEE), + ]); + let initial_position = batch.n_tokens(); + let outcome = ClassifySampleLoop { + model: &model, + classifier: &mut classifier, + sampler: &mut sampler, + context: &mut context, + batch: &mut batch, + initial_position, + max_generated_tokens: MAX_GENERATED_TOKENS, + } + .run()?; + + let usage = classifier.usage(); + + assert!( + !outcome.generated_raw.is_empty(), + "DeepSeek-R1-8B: must generate at least one token" + ); + assert_eq!( + outcome.observed_reasoning, 0, + "DeepSeek-R1-8B thinking-disabled: classifier must not emit any Reasoning token \ + when the prompt closes the think block before generation begins; \ + generated={:?}", + outcome.generated_raw + ); + assert_eq!( + outcome.observed_undeterminable, 0, + "DeepSeek-R1-8B thinking-disabled: prompt-token replay must move section to Content \ + before generation, so no Undeterminable tokens may be emitted; \ + generated={:?}", + outcome.generated_raw + ); + assert_eq!( + usage.reasoning_tokens, 0, + "DeepSeek-R1-8B thinking-disabled: usage.reasoning_tokens must be zero; usage={usage:?}" + ); + assert_eq!( + usage.undeterminable_tokens, 0, + "DeepSeek-R1-8B thinking-disabled: usage.undeterminable_tokens must be zero; usage={usage:?}" + ); + assert!( + outcome.observed_content > 0, + "DeepSeek-R1-8B thinking-disabled: classifier must emit at least one Content token" + ); + assert_eq!( + usage.completion_tokens(), + outcome.observed_content, + "DeepSeek-R1-8B thinking-disabled: completion tokens must equal observed Content tokens" + ); + + for forbidden in FORBIDDEN_MARKERS { + assert!( + !outcome.content_stream.contains(forbidden), + "DeepSeek-R1-8B thinking-disabled: content_stream leaked marker {forbidden:?}; \ + content_stream={:?}", + outcome.content_stream + ); + } + + Ok(()) +} diff --git a/llama-cpp-bindings-tests/tests/deepseek_r1_8b_classifier_emits_reasoning.rs b/llama-cpp-bindings-tests/tests/deepseek_r1_8b_classifier_emits_reasoning.rs new file mode 100644 index 00000000..6b8f34bc --- /dev/null +++ b/llama-cpp-bindings-tests/tests/deepseek_r1_8b_classifier_emits_reasoning.rs @@ -0,0 +1,148 @@ +use std::num::NonZeroU32; + +use anyhow::Result; +use anyhow::bail; +use llama_cpp_bindings::ChatMessageParseOutcome; +use llama_cpp_bindings::context::LlamaContext; +use llama_cpp_bindings::context::params::LlamaContextParams; +use llama_cpp_bindings::llama_backend::LlamaBackend; +use llama_cpp_bindings::llama_batch::LlamaBatch; +use llama_cpp_bindings::model::AddBos; +use llama_cpp_bindings::model::LlamaModel; +use llama_cpp_bindings::sampling::LlamaSampler; +use llama_cpp_bindings_tests::classify_sample_loop::ClassifySampleLoop; +use llama_cpp_bindings_tests::gpu_backend::inference_model_params; +use llama_cpp_bindings_tests::gpu_backend::require_compiled_backends_present; +use llama_cpp_bindings_tests::test_model::download_file_from; + +const DEEPSEEK_R1_8B_REPO: &str = "unsloth/DeepSeek-R1-Distill-Llama-8B-GGUF"; +const DEEPSEEK_R1_8B_FILE: &str = "DeepSeek-R1-Distill-Llama-8B-Q4_K_M.gguf"; + +const MAX_GENERATED_TOKENS: i32 = 1500; + +// DeepSeek-R1-Distill-Llama-8B uses `...` reasoning markers +// and full-width-bar role tokens `<｜User｜>` / `<｜Assistant｜>` (U+FF5C, +// not ASCII `|`). The chat template's `add_generation_prompt` ALWAYS appends +// `<｜Assistant｜>\n` — DeepSeek-R1 is a pure reasoner with no +// thinking-disabled mode — so the model resumes generation already inside +// the reasoning block. +const DEEPSEEK_R1_8B_THINKING_PROMPT: &str = "\ +<｜User｜>What is 2 + 2?<｜Assistant｜> +"; + +const FORBIDDEN_MARKERS: &[&str] = &["", ""]; + +#[test] +fn deepseek_r1_8b_classifier_emits_reasoning_for_thinking_enabled_prompt() -> Result<()> { + let backend = LlamaBackend::init()?; + require_compiled_backends_present()?; + + let path = download_file_from(DEEPSEEK_R1_8B_REPO, DEEPSEEK_R1_8B_FILE)?; + let params = inference_model_params(); + let model = LlamaModel::load_from_file(&backend, &path, ¶ms)?; + + let mut classifier = model.sampled_token_classifier(); + let prompt_tokens = model.str_to_token(DEEPSEEK_R1_8B_THINKING_PROMPT, AddBos::Never)?; + let prompt_token_count = u64::try_from(prompt_tokens.len())?; + + let mut batch = LlamaBatch::new(2048, 1)?; + classifier.feed_prompt_sequence_to_batch(&mut batch, &prompt_tokens, 0, false)?; + + let context_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(8192)); + let mut context = LlamaContext::from_model(&model, &backend, context_params)?; + + context.decode(&mut batch)?; + + let promoted = classifier.commit_prompt_tokens(); + assert_eq!(promoted, prompt_token_count); + + let mut sampler = LlamaSampler::chain_simple([ + LlamaSampler::penalties(64, 1.1, 0.0, 0.0), + LlamaSampler::top_k(40), + LlamaSampler::top_p(0.9, 1), + LlamaSampler::min_p(0.05, 1), + LlamaSampler::temp(0.7), + LlamaSampler::dist(0x00C0_FFEE), + ]); + let initial_position = batch.n_tokens(); + let outcome = ClassifySampleLoop { + model: &model, + classifier: &mut classifier, + sampler: &mut sampler, + context: &mut context, + batch: &mut batch, + initial_position, + max_generated_tokens: MAX_GENERATED_TOKENS, + } + .run()?; + + let usage = classifier.usage(); + let parse_outcome = model.parse_chat_message("[]", &outcome.generated_raw, false)?; + let ChatMessageParseOutcome::Recognized(parsed) = parse_outcome else { + bail!("DeepSeek-R1-8B chat template must be recognised by the parser; got Unrecognized"); + }; + + assert!( + !outcome.generated_raw.is_empty(), + "DeepSeek-R1-8B: must generate at least one token" + ); + assert!( + outcome.observed_reasoning > 0, + "DeepSeek-R1-8B: classifier must emit at least one Reasoning token when the prompt \ + opens a block; outcome={outcome:?}", + ); + assert!( + usage.reasoning_tokens > 0, + "DeepSeek-R1-8B: usage.reasoning_tokens must be non-zero when the prompt opens a \ + block; usage was {usage:?}" + ); + assert_eq!( + outcome.observed_undeterminable, 0, + "DeepSeek-R1-8B: prompt-token replay must move section to Reasoning before generation, \ + so no Undeterminable tokens may be emitted; outcome={outcome:?}" + ); + assert_eq!( + usage.undeterminable_tokens, 0, + "DeepSeek-R1-8B: usage.undeterminable_tokens must be zero; usage={usage:?}" + ); + assert_eq!( + usage.completion_tokens(), + outcome.observed_content + outcome.observed_reasoning, + "DeepSeek-R1-8B: completion tokens must equal observed Content + Reasoning" + ); + + if parsed.reasoning_content.is_empty() { + eprintln!( + "DeepSeek-R1-8B didn't close its reasoning block within {MAX_GENERATED_TOKENS} \ + tokens — skipping strict parser-equality assertions" + ); + } else { + assert_eq!( + outcome.reasoning_stream, parsed.reasoning_content, + "DeepSeek-R1-8B: per-token reasoning stream must equal parser-side reasoning_content \ + (any difference means a marker leaked into the user-visible stream)", + ); + assert_eq!( + outcome.content_stream, parsed.content, + "DeepSeek-R1-8B: per-token content stream must equal parser-side content \ + (any difference means a marker leaked into the user-visible stream)", + ); + } + + for forbidden in FORBIDDEN_MARKERS { + assert!( + !outcome.reasoning_stream.contains(forbidden), + "DeepSeek-R1-8B: reasoning_stream leaked marker {forbidden:?}; \ + reasoning_stream={:?}", + outcome.reasoning_stream + ); + assert!( + !outcome.content_stream.contains(forbidden), + "DeepSeek-R1-8B: content_stream leaked marker {forbidden:?}; \ + content_stream={:?}", + outcome.content_stream + ); + } + + Ok(()) +} diff --git a/llama-cpp-bindings-tests/tests/deepseek_r1_8b_duck_types_gemma_paired_quote.rs b/llama-cpp-bindings-tests/tests/deepseek_r1_8b_duck_types_gemma_paired_quote.rs new file mode 100644 index 00000000..329111a6 --- /dev/null +++ b/llama-cpp-bindings-tests/tests/deepseek_r1_8b_duck_types_gemma_paired_quote.rs @@ -0,0 +1,69 @@ +use anyhow::Result; +use anyhow::bail; +use llama_cpp_bindings::ChatMessageParseOutcome; +use llama_cpp_bindings::ToolCallArguments; +use llama_cpp_bindings::llama_backend::LlamaBackend; +use llama_cpp_bindings::model::LlamaModel; +use llama_cpp_bindings_tests::gpu_backend::inference_model_params; +use llama_cpp_bindings_tests::gpu_backend::require_compiled_backends_present; +use llama_cpp_bindings_tests::test_model::download_file_from; + +const DEEPSEEK_R1_8B_REPO: &str = "unsloth/DeepSeek-R1-Distill-Llama-8B-GGUF"; +const DEEPSEEK_R1_8B_FILE: &str = "DeepSeek-R1-Distill-Llama-8B-Q4_K_M.gguf"; + +const TOOLS_JSON: &str = r#"[ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather for a location", + "parameters": { + "type": "object", + "properties": { + "location": {"type": "string", "description": "The city name"} + }, + "required": ["location"] + } + } + } +]"#; + +const GEMMA_PAIRED_QUOTE_PAYLOAD: &str = "<|tool_call>call:get_weather{location:<|\"|>Paris<|\"|>}"; + +#[test] +fn deepseek_r1_8b_duck_types_gemma_paired_quote() -> Result<()> { + let backend = LlamaBackend::init()?; + require_compiled_backends_present()?; + + let path = download_file_from(DEEPSEEK_R1_8B_REPO, DEEPSEEK_R1_8B_FILE)?; + let params = inference_model_params(); + let model = LlamaModel::load_from_file(&backend, &path, ¶ms)?; + + let outcome = model.parse_chat_message(TOOLS_JSON, GEMMA_PAIRED_QUOTE_PAYLOAD, false)?; + + let ChatMessageParseOutcome::Recognized(parsed) = outcome else { + bail!( + "duck-type pass must recognise Gemma paired-quote on a model with no registered \ + template; got Unrecognized" + ); + }; + assert_eq!( + parsed.tool_calls.len(), + 1, + "expected one tool call; got {:?}", + parsed.tool_calls + ); + assert_eq!(parsed.tool_calls[0].name, "get_weather"); + let location = match &parsed.tool_calls[0].arguments { + ToolCallArguments::ValidJson(value) => value + .get("location") + .and_then(|v| v.as_str()) + .map(str::to_owned), + ToolCallArguments::InvalidJson(raw) => { + bail!("expected ValidJson, got InvalidJson: {raw}"); + } + }; + assert_eq!(location.as_deref(), Some("Paris")); + + Ok(()) +} diff --git a/llama-cpp-bindings-tests/tests/deepseek_r1_8b_duck_types_glm_key_value_tags.rs b/llama-cpp-bindings-tests/tests/deepseek_r1_8b_duck_types_glm_key_value_tags.rs new file mode 100644 index 00000000..c2aa85a6 --- /dev/null +++ b/llama-cpp-bindings-tests/tests/deepseek_r1_8b_duck_types_glm_key_value_tags.rs @@ -0,0 +1,72 @@ +use anyhow::Result; +use anyhow::bail; +use llama_cpp_bindings::ChatMessageParseOutcome; +use llama_cpp_bindings::ToolCallArguments; +use llama_cpp_bindings::llama_backend::LlamaBackend; +use llama_cpp_bindings::model::LlamaModel; +use llama_cpp_bindings_tests::gpu_backend::inference_model_params; +use llama_cpp_bindings_tests::gpu_backend::require_compiled_backends_present; +use llama_cpp_bindings_tests::test_model::download_file_from; + +const DEEPSEEK_R1_8B_REPO: &str = "unsloth/DeepSeek-R1-Distill-Llama-8B-GGUF"; +const DEEPSEEK_R1_8B_FILE: &str = "DeepSeek-R1-Distill-Llama-8B-Q4_K_M.gguf"; + +const TOOLS_JSON: &str = r#"[ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather for a location", + "parameters": { + "type": "object", + "properties": { + "location": {"type": "string", "description": "The city name"} + }, + "required": ["location"] + } + } + } +]"#; + +const GLM_KEY_VALUE_PAYLOAD: &str = "get_weather\ +location\ +Paris\ +"; + +#[test] +fn deepseek_r1_8b_duck_types_glm_key_value_tags() -> Result<()> { + let backend = LlamaBackend::init()?; + require_compiled_backends_present()?; + + let path = download_file_from(DEEPSEEK_R1_8B_REPO, DEEPSEEK_R1_8B_FILE)?; + let params = inference_model_params(); + let model = LlamaModel::load_from_file(&backend, &path, ¶ms)?; + + let outcome = model.parse_chat_message(TOOLS_JSON, GLM_KEY_VALUE_PAYLOAD, false)?; + + let ChatMessageParseOutcome::Recognized(parsed) = outcome else { + bail!( + "duck-type pass must recognise GLM key-value tags on a model with no registered \ + template; got Unrecognized" + ); + }; + assert_eq!( + parsed.tool_calls.len(), + 1, + "expected one tool call; got {:?}", + parsed.tool_calls + ); + assert_eq!(parsed.tool_calls[0].name, "get_weather"); + let location = match &parsed.tool_calls[0].arguments { + ToolCallArguments::ValidJson(value) => value + .get("location") + .and_then(|v| v.as_str()) + .map(str::to_owned), + ToolCallArguments::InvalidJson(raw) => { + bail!("expected ValidJson, got InvalidJson: {raw}"); + } + }; + assert_eq!(location.as_deref(), Some("Paris")); + + Ok(()) +} diff --git a/llama-cpp-bindings-tests/tests/deepseek_r1_8b_duck_types_mistral_bracketed_json.rs b/llama-cpp-bindings-tests/tests/deepseek_r1_8b_duck_types_mistral_bracketed_json.rs new file mode 100644 index 00000000..25a38992 --- /dev/null +++ b/llama-cpp-bindings-tests/tests/deepseek_r1_8b_duck_types_mistral_bracketed_json.rs @@ -0,0 +1,69 @@ +use anyhow::Result; +use anyhow::bail; +use llama_cpp_bindings::ChatMessageParseOutcome; +use llama_cpp_bindings::ToolCallArguments; +use llama_cpp_bindings::llama_backend::LlamaBackend; +use llama_cpp_bindings::model::LlamaModel; +use llama_cpp_bindings_tests::gpu_backend::inference_model_params; +use llama_cpp_bindings_tests::gpu_backend::require_compiled_backends_present; +use llama_cpp_bindings_tests::test_model::download_file_from; + +const DEEPSEEK_R1_8B_REPO: &str = "unsloth/DeepSeek-R1-Distill-Llama-8B-GGUF"; +const DEEPSEEK_R1_8B_FILE: &str = "DeepSeek-R1-Distill-Llama-8B-Q4_K_M.gguf"; + +const TOOLS_JSON: &str = r#"[ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather for a location", + "parameters": { + "type": "object", + "properties": { + "location": {"type": "string", "description": "The city name"} + }, + "required": ["location"] + } + } + } +]"#; + +const MISTRAL_BRACKETED_JSON_PAYLOAD: &str = r#"[TOOL_CALLS]get_weather[ARGS]{"location":"Paris"}"#; + +#[test] +fn deepseek_r1_8b_duck_types_mistral_bracketed_json() -> Result<()> { + let backend = LlamaBackend::init()?; + require_compiled_backends_present()?; + + let path = download_file_from(DEEPSEEK_R1_8B_REPO, DEEPSEEK_R1_8B_FILE)?; + let params = inference_model_params(); + let model = LlamaModel::load_from_file(&backend, &path, ¶ms)?; + + let outcome = model.parse_chat_message(TOOLS_JSON, MISTRAL_BRACKETED_JSON_PAYLOAD, false)?; + + let ChatMessageParseOutcome::Recognized(parsed) = outcome else { + bail!( + "duck-type pass must recognise Mistral bracketed-JSON on a model with no registered \ + template; got Unrecognized" + ); + }; + assert_eq!( + parsed.tool_calls.len(), + 1, + "expected one tool call; got {:?}", + parsed.tool_calls + ); + assert_eq!(parsed.tool_calls[0].name, "get_weather"); + let location = match &parsed.tool_calls[0].arguments { + ToolCallArguments::ValidJson(value) => value + .get("location") + .and_then(|v| v.as_str()) + .map(str::to_owned), + ToolCallArguments::InvalidJson(raw) => { + bail!("expected ValidJson, got InvalidJson: {raw}"); + } + }; + assert_eq!(location.as_deref(), Some("Paris")); + + Ok(()) +} diff --git a/llama-cpp-bindings-tests/tests/deepseek_r1_8b_duck_types_qwen_xml.rs b/llama-cpp-bindings-tests/tests/deepseek_r1_8b_duck_types_qwen_xml.rs new file mode 100644 index 00000000..72f8bcfd --- /dev/null +++ b/llama-cpp-bindings-tests/tests/deepseek_r1_8b_duck_types_qwen_xml.rs @@ -0,0 +1,75 @@ +use anyhow::Result; +use anyhow::bail; +use llama_cpp_bindings::ChatMessageParseOutcome; +use llama_cpp_bindings::ToolCallArguments; +use llama_cpp_bindings::llama_backend::LlamaBackend; +use llama_cpp_bindings::model::LlamaModel; +use llama_cpp_bindings_tests::gpu_backend::inference_model_params; +use llama_cpp_bindings_tests::gpu_backend::require_compiled_backends_present; +use llama_cpp_bindings_tests::test_model::download_file_from; + +const DEEPSEEK_R1_8B_REPO: &str = "unsloth/DeepSeek-R1-Distill-Llama-8B-GGUF"; +const DEEPSEEK_R1_8B_FILE: &str = "DeepSeek-R1-Distill-Llama-8B-Q4_K_M.gguf"; + +const TOOLS_JSON: &str = r#"[ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather for a location", + "parameters": { + "type": "object", + "properties": { + "location": {"type": "string", "description": "The city name"} + }, + "required": ["location"] + } + } + } +]"#; + +const QWEN_XML_PAYLOAD: &str = "\n\ +\n\ +\n\ +Paris\n\ +\n\ +\n\ +"; + +#[test] +fn deepseek_r1_8b_duck_types_qwen_xml() -> Result<()> { + let backend = LlamaBackend::init()?; + require_compiled_backends_present()?; + + let path = download_file_from(DEEPSEEK_R1_8B_REPO, DEEPSEEK_R1_8B_FILE)?; + let params = inference_model_params(); + let model = LlamaModel::load_from_file(&backend, &path, ¶ms)?; + + let outcome = model.parse_chat_message(TOOLS_JSON, QWEN_XML_PAYLOAD, false)?; + + let ChatMessageParseOutcome::Recognized(parsed) = outcome else { + bail!( + "duck-type pass must recognise Qwen XML on a model with no registered template; \ + got Unrecognized" + ); + }; + assert_eq!( + parsed.tool_calls.len(), + 1, + "expected one tool call; got {:?}", + parsed.tool_calls + ); + assert_eq!(parsed.tool_calls[0].name, "get_weather"); + let location = match &parsed.tool_calls[0].arguments { + ToolCallArguments::ValidJson(value) => value + .get("location") + .and_then(|v| v.as_str()) + .map(str::to_owned), + ToolCallArguments::InvalidJson(raw) => { + bail!("expected ValidJson, got InvalidJson: {raw}"); + } + }; + assert_eq!(location.as_deref(), Some("Paris")); + + Ok(()) +} diff --git a/llama-cpp-bindings-tests/tests/deepseek_r1_8b_recognizes_empty_tool_calls_when_input_is_plain_content_with_tools_requested.rs b/llama-cpp-bindings-tests/tests/deepseek_r1_8b_recognizes_empty_tool_calls_when_input_is_plain_content_with_tools_requested.rs new file mode 100644 index 00000000..60828698 --- /dev/null +++ b/llama-cpp-bindings-tests/tests/deepseek_r1_8b_recognizes_empty_tool_calls_when_input_is_plain_content_with_tools_requested.rs @@ -0,0 +1,57 @@ +use anyhow::Result; +use anyhow::bail; +use llama_cpp_bindings::ChatMessageParseOutcome; +use llama_cpp_bindings::llama_backend::LlamaBackend; +use llama_cpp_bindings::model::LlamaModel; +use llama_cpp_bindings_tests::gpu_backend::inference_model_params; +use llama_cpp_bindings_tests::gpu_backend::require_compiled_backends_present; +use llama_cpp_bindings_tests::test_model::download_file_from; + +const DEEPSEEK_R1_8B_REPO: &str = "unsloth/DeepSeek-R1-Distill-Llama-8B-GGUF"; +const DEEPSEEK_R1_8B_FILE: &str = "DeepSeek-R1-Distill-Llama-8B-Q4_K_M.gguf"; + +const TOOLS_JSON: &str = r#"[ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather for a location", + "parameters": { + "type": "object", + "properties": { + "location": {"type": "string", "description": "The city name"} + }, + "required": ["location"] + } + } + } +]"#; + +const PLAIN_CONTENT: &str = "Sorry, I cannot help with that."; + +#[test] +fn deepseek_r1_8b_recognizes_empty_tool_calls_when_input_is_plain_content_with_tools_requested() +-> Result<()> { + let backend = LlamaBackend::init()?; + require_compiled_backends_present()?; + + let path = download_file_from(DEEPSEEK_R1_8B_REPO, DEEPSEEK_R1_8B_FILE)?; + let params = inference_model_params(); + let model = LlamaModel::load_from_file(&backend, &path, ¶ms)?; + + let outcome = model.parse_chat_message(TOOLS_JSON, PLAIN_CONTENT, false)?; + + let ChatMessageParseOutcome::Recognized(parsed) = outcome else { + bail!( + "plain content with tools requested must produce Recognized (with empty tool_calls); \ + got Unrecognized" + ); + }; + assert!( + parsed.tool_calls.is_empty(), + "expected no tool calls; got {:?}", + parsed.tool_calls + ); + + Ok(()) +} diff --git a/llama-cpp-bindings-tests/tests/deepseek_r1_8b_recognizes_empty_tool_calls_when_tools_not_requested.rs b/llama-cpp-bindings-tests/tests/deepseek_r1_8b_recognizes_empty_tool_calls_when_tools_not_requested.rs new file mode 100644 index 00000000..931a9b1c --- /dev/null +++ b/llama-cpp-bindings-tests/tests/deepseek_r1_8b_recognizes_empty_tool_calls_when_tools_not_requested.rs @@ -0,0 +1,36 @@ +use anyhow::Result; +use anyhow::bail; +use llama_cpp_bindings::ChatMessageParseOutcome; +use llama_cpp_bindings::llama_backend::LlamaBackend; +use llama_cpp_bindings::model::LlamaModel; +use llama_cpp_bindings_tests::gpu_backend::inference_model_params; +use llama_cpp_bindings_tests::gpu_backend::require_compiled_backends_present; +use llama_cpp_bindings_tests::test_model::download_file_from; + +const DEEPSEEK_R1_8B_REPO: &str = "unsloth/DeepSeek-R1-Distill-Llama-8B-GGUF"; +const DEEPSEEK_R1_8B_FILE: &str = "DeepSeek-R1-Distill-Llama-8B-Q4_K_M.gguf"; + +const PLAIN_CONTENT: &str = "Hello there."; + +#[test] +fn deepseek_r1_8b_recognizes_empty_tool_calls_when_tools_not_requested() -> Result<()> { + let backend = LlamaBackend::init()?; + require_compiled_backends_present()?; + + let path = download_file_from(DEEPSEEK_R1_8B_REPO, DEEPSEEK_R1_8B_FILE)?; + let params = inference_model_params(); + let model = LlamaModel::load_from_file(&backend, &path, ¶ms)?; + + let outcome = model.parse_chat_message("[]", PLAIN_CONTENT, false)?; + + let ChatMessageParseOutcome::Recognized(parsed) = outcome else { + bail!("plain content with empty tools array must produce Recognized; got Unrecognized"); + }; + assert!( + parsed.tool_calls.is_empty(), + "expected no tool calls; got {:?}", + parsed.tool_calls + ); + + Ok(()) +} diff --git a/llama-cpp-bindings-tests/tests/embeddings.rs b/llama-cpp-bindings-tests/tests/embeddings.rs index 83fc008f..840dff79 100644 --- a/llama-cpp-bindings-tests/tests/embeddings.rs +++ b/llama-cpp-bindings-tests/tests/embeddings.rs @@ -1,11 +1,12 @@ use std::time::Duration; use anyhow::{Context, Result}; +use llama_cpp_bindings::context::LlamaContext; use llama_cpp_bindings::context::params::LlamaContextParams; use llama_cpp_bindings::ggml_time_us; use llama_cpp_bindings::llama_batch::LlamaBatch; use llama_cpp_bindings::model::AddBos; -use llama_cpp_bindings_tests::TestFixture; +use llama_cpp_bindings_tests::FixtureSession; fn normalize(input: &[f32]) -> Vec { let magnitude = input @@ -18,15 +19,14 @@ fn normalize(input: &[f32]) -> Vec { #[test] fn embedding_generation_produces_vectors() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.embedding_model()?; let ctx_params = LlamaContextParams::default() .with_n_threads_batch(std::thread::available_parallelism()?.get().try_into()?) .with_embeddings(true); - let mut ctx = model - .new_context(backend, ctx_params) + let mut ctx = LlamaContext::from_model(model, backend, ctx_params) .with_context(|| "unable to create context")?; let prompt = "Hello my name is"; @@ -40,12 +40,12 @@ fn embedding_generation_produces_vectors() -> Result<()> { let t_main_start = ggml_time_us(); - let mut classifier = model.reasoning_token_classifier()?; + let mut classifier = model.sampled_token_classifier(); let mut batch = LlamaBatch::new(n_ctx, 1)?; classifier.feed_prompt_sequence_to_batch(&mut batch, &tokens, 0, false)?; assert_eq!(classifier.pending_prompt_tokens(), prompt_token_count); - assert_eq!(classifier.usage().prompt_tokens(), 0); + assert_eq!(classifier.usage().prompt_tokens, 0); ctx.clear_kv_cache(); ctx.decode(&mut batch) @@ -84,7 +84,7 @@ fn embedding_generation_produces_vectors() -> Result<()> { ); let usage = classifier.into_usage(); - assert_eq!(usage.prompt_tokens(), prompt_token_count); + assert_eq!(usage.prompt_tokens, prompt_token_count); assert_eq!(usage.completion_tokens(), 0); Ok(()) diff --git a/llama-cpp-bindings-tests/tests/eval_multimodal_chunks_records_exact_token_counts.rs b/llama-cpp-bindings-tests/tests/eval_multimodal_chunks_records_exact_token_counts.rs new file mode 100644 index 00000000..53cdbb53 --- /dev/null +++ b/llama-cpp-bindings-tests/tests/eval_multimodal_chunks_records_exact_token_counts.rs @@ -0,0 +1,148 @@ +#![cfg(feature = "multimodal_capable")] + +use std::num::NonZeroU32; + +use anyhow::Result; +use llama_cpp_bindings::TokenUsage; +use llama_cpp_bindings::context::LlamaContext; +use llama_cpp_bindings::context::params::LlamaContextParams; +use llama_cpp_bindings::mtmd::MtmdBitmap; +use llama_cpp_bindings::mtmd::MtmdInputChunkType; +use llama_cpp_bindings::mtmd::MtmdInputChunks; +use llama_cpp_bindings::mtmd::MtmdInputText; +use llama_cpp_bindings::mtmd::mtmd_default_marker; +use llama_cpp_bindings_tests::FixtureSession; +use llama_cpp_bindings_tests::test_model::fixtures_dir; + +const PROMPT_QUESTION: &str = "What animals do you see in this image?"; + +struct ExpectedChunkTotals { + text: u64, + image: u64, + audio: u64, +} + +fn sum_chunk_token_counts_by_type(chunks: &MtmdInputChunks) -> Result { + let mut totals = ExpectedChunkTotals { + text: 0, + image: 0, + audio: 0, + }; + for index in 0..chunks.len() { + let chunk = chunks + .get(index) + .ok_or_else(|| anyhow::anyhow!("chunk index {index} should exist"))?; + let n_tokens = u64::try_from(chunk.n_tokens())?; + match chunk.chunk_type()? { + MtmdInputChunkType::Text => { + totals.text = totals.text.saturating_add(n_tokens); + } + MtmdInputChunkType::Image => { + totals.image = totals.image.saturating_add(n_tokens); + } + MtmdInputChunkType::Audio => { + totals.audio = totals.audio.saturating_add(n_tokens); + } + } + } + Ok(totals) +} + +fn build_multimodal_chunks_and_eval_into_usage() -> Result<(TokenUsage, ExpectedChunkTotals)> { + let fixture = FixtureSession::open()?; + let backend = fixture.backend(); + let model = fixture.default_model(); + let mtmd_ctx = fixture.mtmd_context()?; + + let image_path = fixtures_dir().join("llamas.jpg"); + let image_path_str = image_path + .to_str() + .ok_or_else(|| anyhow::anyhow!("image path is not valid UTF-8"))?; + let bitmap = MtmdBitmap::from_file(mtmd_ctx, image_path_str)?; + + let marker = mtmd_default_marker(); + let prompt = format!("{marker}{PROMPT_QUESTION}"); + + let input_text = MtmdInputText { + text: prompt, + add_special: false, + parse_special: true, + }; + + let chunks = mtmd_ctx.tokenize(input_text, &[&bitmap])?; + let expected = sum_chunk_token_counts_by_type(&chunks)?; + + let context_params = LlamaContextParams::default() + .with_n_ctx(NonZeroU32::new(4096)) + .with_n_batch(512); + let context = LlamaContext::from_model(model, backend, context_params)?; + + let mut classifier = model.sampled_token_classifier(); + classifier.eval_multimodal_chunks(&chunks, mtmd_ctx, &context, 0, 0, 512, true)?; + + Ok((classifier.into_usage(), expected)) +} + +#[test] +fn prompt_tokens_match_text_chunk_total() -> Result<()> { + let (usage, expected) = build_multimodal_chunks_and_eval_into_usage()?; + + if usage.prompt_tokens != expected.text { + anyhow::bail!( + "prompt_tokens must equal sum of text-chunk n_tokens; expected {}, got {}", + expected.text, + usage.prompt_tokens + ); + } + + Ok(()) +} + +#[test] +fn input_image_tokens_match_image_chunk_total() -> Result<()> { + let (usage, expected) = build_multimodal_chunks_and_eval_into_usage()?; + + if usage.input_image_tokens != expected.image { + anyhow::bail!( + "input_image_tokens must equal sum of image-chunk n_tokens; expected {}, got {}", + expected.image, + usage.input_image_tokens + ); + } + + Ok(()) +} + +#[test] +fn input_audio_tokens_are_zero_for_image_only_input() -> Result<()> { + let (usage, expected) = build_multimodal_chunks_and_eval_into_usage()?; + + if expected.audio != 0 { + anyhow::bail!( + "fixture invariant: image-only multimodal input should produce zero audio chunk tokens, got {}", + expected.audio + ); + } + if usage.input_audio_tokens != 0 { + anyhow::bail!( + "input_audio_tokens must be zero when no audio chunks are evaluated; got {}", + usage.input_audio_tokens + ); + } + + Ok(()) +} + +#[test] +fn completion_tokens_are_zero_after_eval_before_generation() -> Result<()> { + let (usage, _expected) = build_multimodal_chunks_and_eval_into_usage()?; + + if usage.completion_tokens() != 0 { + anyhow::bail!( + "completion_tokens must be zero immediately after eval (no generation has occurred); got {}", + usage.completion_tokens() + ); + } + + Ok(()) +} diff --git a/llama-cpp-bindings-tests/tests/gemma4_classifier_does_not_emit_reasoning_for_thinking_disabled_prompt.rs b/llama-cpp-bindings-tests/tests/gemma4_classifier_does_not_emit_reasoning_for_thinking_disabled_prompt.rs new file mode 100644 index 00000000..71b2a1ef --- /dev/null +++ b/llama-cpp-bindings-tests/tests/gemma4_classifier_does_not_emit_reasoning_for_thinking_disabled_prompt.rs @@ -0,0 +1,115 @@ +use std::num::NonZeroU32; + +use anyhow::Result; +use llama_cpp_bindings::context::LlamaContext; +use llama_cpp_bindings::context::params::LlamaContextParams; +use llama_cpp_bindings::llama_backend::LlamaBackend; +use llama_cpp_bindings::llama_batch::LlamaBatch; +use llama_cpp_bindings::model::AddBos; +use llama_cpp_bindings::model::LlamaModel; +use llama_cpp_bindings::sampling::LlamaSampler; +use llama_cpp_bindings_tests::classify_sample_loop::ClassifySampleLoop; +use llama_cpp_bindings_tests::gpu_backend::inference_model_params; +use llama_cpp_bindings_tests::gpu_backend::require_compiled_backends_present; +use llama_cpp_bindings_tests::test_model::download_file_from; + +const GEMMA4_REPO: &str = "unsloth/gemma-4-E4B-it-GGUF"; +const GEMMA4_FILE: &str = "gemma-4-E4B-it-Q4_K_M.gguf"; + +const MAX_GENERATED_TOKENS: i32 = 200; + +// Mirrors what Gemma 4's chat template renders when the caller asks for +// `enable_thinking=false`: the model turn opens with a closed empty +// `<|channel>thought\n\n` block, so generation begins in CONTENT. +const GEMMA4_THINKING_DISABLED_PROMPT: &str = "\ +user\nReply with the single word: four. Do not explain.\n\ +model\n<|channel>thought\n\n"; + +const FORBIDDEN_MARKERS: &[&str] = &["<|channel>thought", ""]; + +#[test] +fn gemma4_classifier_does_not_emit_reasoning_for_thinking_disabled_prompt() -> Result<()> { + let backend = LlamaBackend::init()?; + require_compiled_backends_present()?; + + let path = download_file_from(GEMMA4_REPO, GEMMA4_FILE)?; + let params = inference_model_params(); + let model = LlamaModel::load_from_file(&backend, &path, ¶ms)?; + + let mut classifier = model.sampled_token_classifier(); + let prompt_tokens = model.str_to_token(GEMMA4_THINKING_DISABLED_PROMPT, AddBos::Never)?; + let prompt_token_count = u64::try_from(prompt_tokens.len())?; + + let mut batch = LlamaBatch::new(2048, 1)?; + classifier.feed_prompt_sequence_to_batch(&mut batch, &prompt_tokens, 0, false)?; + + let context_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(8192)); + let mut context = LlamaContext::from_model(&model, &backend, context_params)?; + + context.decode(&mut batch)?; + + let promoted = classifier.commit_prompt_tokens(); + assert_eq!(promoted, prompt_token_count); + + let mut sampler = LlamaSampler::greedy(); + let initial_position = batch.n_tokens(); + let outcome = ClassifySampleLoop { + model: &model, + classifier: &mut classifier, + sampler: &mut sampler, + context: &mut context, + batch: &mut batch, + initial_position, + max_generated_tokens: MAX_GENERATED_TOKENS, + } + .run()?; + + let usage = classifier.usage(); + + assert!( + !outcome.generated_raw.is_empty(), + "Gemma 4 must generate at least one token" + ); + assert_eq!( + outcome.observed_reasoning, 0, + "Gemma 4 thinking-disabled: classifier must not emit any Reasoning token \ + when the prompt closes the thought channel before generation begins; \ + generated={:?}", + outcome.generated_raw + ); + assert_eq!( + outcome.observed_undeterminable, 0, + "Gemma 4 thinking-disabled: prompt-token replay must move section to Content \ + before generation, so no Undeterminable tokens may be emitted; \ + generated={:?}", + outcome.generated_raw + ); + assert_eq!( + usage.reasoning_tokens, 0, + "Gemma 4 thinking-disabled: usage.reasoning_tokens must be zero; usage={usage:?}" + ); + assert_eq!( + usage.undeterminable_tokens, 0, + "Gemma 4 thinking-disabled: usage.undeterminable_tokens must be zero; usage={usage:?}" + ); + assert!( + outcome.observed_content > 0, + "Gemma 4 thinking-disabled: classifier must emit at least one Content token" + ); + assert_eq!( + usage.completion_tokens(), + outcome.observed_content, + "Gemma 4 thinking-disabled: completion tokens must equal observed Content tokens" + ); + + for forbidden in FORBIDDEN_MARKERS { + assert!( + !outcome.content_stream.contains(forbidden), + "Gemma 4 thinking-disabled: content_stream leaked marker {forbidden:?}; \ + content_stream={:?}", + outcome.content_stream + ); + } + + Ok(()) +} diff --git a/llama-cpp-bindings-tests/tests/gemma4_classifier_emits_reasoning.rs b/llama-cpp-bindings-tests/tests/gemma4_classifier_emits_reasoning.rs new file mode 100644 index 00000000..0ad59240 --- /dev/null +++ b/llama-cpp-bindings-tests/tests/gemma4_classifier_emits_reasoning.rs @@ -0,0 +1,135 @@ +use std::num::NonZeroU32; + +use anyhow::Result; +use anyhow::bail; +use llama_cpp_bindings::ChatMessageParseOutcome; +use llama_cpp_bindings::context::LlamaContext; +use llama_cpp_bindings::context::params::LlamaContextParams; +use llama_cpp_bindings::llama_backend::LlamaBackend; +use llama_cpp_bindings::llama_batch::LlamaBatch; +use llama_cpp_bindings::model::AddBos; +use llama_cpp_bindings::model::LlamaModel; +use llama_cpp_bindings::sampling::LlamaSampler; +use llama_cpp_bindings_tests::classify_sample_loop::ClassifySampleLoop; +use llama_cpp_bindings_tests::gpu_backend::inference_model_params; +use llama_cpp_bindings_tests::gpu_backend::require_compiled_backends_present; +use llama_cpp_bindings_tests::test_model::download_file_from; + +const GEMMA4_REPO: &str = "unsloth/gemma-4-E4B-it-GGUF"; +const GEMMA4_FILE: &str = "gemma-4-E4B-it-Q4_K_M.gguf"; + +const MAX_GENERATED_TOKENS: i32 = 1500; + +// Gemma 4 uses asymmetric reasoning markers: `<|channel>thought` opens +// the thinking block and `` closes it. We pre-inject the +// `<|channel>thought\n` opener at the model turn so the classifier sees +// the marker via prompt-token replay and starts generation in `Reasoning`, +// matching the behaviour of Qwen3.5/3.6's auto-injected `\n`. +const GEMMA4_THINKING_PROMPT: &str = "\ +user\nReply with the single word: four. Do not explain.\n\ +model\n<|channel>thought\n"; + +const FORBIDDEN_MARKERS: &[&str] = &["<|channel>thought", ""]; + +#[test] +fn gemma4_classifier_emits_reasoning_for_thinking_prompt() -> Result<()> { + let backend = LlamaBackend::init()?; + require_compiled_backends_present()?; + + let path = download_file_from(GEMMA4_REPO, GEMMA4_FILE)?; + let params = inference_model_params(); + let model = LlamaModel::load_from_file(&backend, &path, ¶ms)?; + + let mut classifier = model.sampled_token_classifier(); + let prompt_tokens = model.str_to_token(GEMMA4_THINKING_PROMPT, AddBos::Never)?; + let prompt_token_count = u64::try_from(prompt_tokens.len())?; + + let mut batch = LlamaBatch::new(2048, 1)?; + classifier.feed_prompt_sequence_to_batch(&mut batch, &prompt_tokens, 0, false)?; + + let context_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(8192)); + let mut context = LlamaContext::from_model(&model, &backend, context_params)?; + + context.decode(&mut batch)?; + + let promoted = classifier.commit_prompt_tokens(); + assert_eq!(promoted, prompt_token_count); + + let mut sampler = LlamaSampler::greedy(); + let initial_position = batch.n_tokens(); + let outcome = ClassifySampleLoop { + model: &model, + classifier: &mut classifier, + sampler: &mut sampler, + context: &mut context, + batch: &mut batch, + initial_position, + max_generated_tokens: MAX_GENERATED_TOKENS, + } + .run()?; + + let usage = classifier.usage(); + let parse_outcome = model.parse_chat_message("[]", &outcome.generated_raw, false)?; + let ChatMessageParseOutcome::Recognized(parsed) = parse_outcome else { + bail!("Gemma 4 chat template must be recognised by the parser; got Unrecognized"); + }; + + assert!( + !outcome.generated_raw.is_empty(), + "Gemma 4 must generate at least one token" + ); + assert!( + outcome.observed_reasoning > 0, + "Gemma 4 classifier must emit at least one Reasoning token when the model \ + emits a `<|channel>thought` block; outcome={outcome:?}", + ); + assert!( + usage.reasoning_tokens > 0, + "Gemma 4 usage.reasoning_tokens must be non-zero when the model emits a \ + reasoning block; usage was {usage:?}" + ); + assert_eq!( + outcome.observed_undeterminable, 0, + "Gemma 4: classifier must not emit Undeterminable when the model emits a \ + detected `<|channel>thought` marker; outcome={outcome:?}" + ); + assert_eq!( + usage.undeterminable_tokens, 0, + "Gemma 4: usage.undeterminable_tokens must be zero; usage={usage:?}" + ); + assert_eq!( + usage.completion_tokens(), + outcome.observed_content + outcome.observed_reasoning, + "Gemma 4: completion tokens must equal observed Content + Reasoning" + ); + assert!( + !parsed.reasoning_content.is_empty(), + "Gemma 4 must close its reasoning block within {MAX_GENERATED_TOKENS} tokens; \ + increase the budget or pick a more direct prompt. generated={:?}", + outcome.generated_raw, + ); + + // Gemma 4 goes through llama.cpp's specialized-template path, which leaves the + // raw `<|channel>thought` prefix in `parsed.reasoning_content` rather than + // stripping it like the differential autoparser does for Qwen3-family. So the + // parser-equality cross-check would require a per-template carve-out — instead, + // rely on the FORBIDDEN_MARKERS substring check below: the streams the user + // actually sees must not contain marker text, regardless of what the parser + // chose to keep. + for forbidden in FORBIDDEN_MARKERS { + assert!( + !outcome.reasoning_stream.contains(forbidden), + "Gemma 4: reasoning_stream leaked marker {forbidden:?}; \ + reasoning_stream={:?}", + outcome.reasoning_stream + ); + assert!( + !outcome.content_stream.contains(forbidden), + "Gemma 4: content_stream leaked marker {forbidden:?}; \ + content_stream={:?}", + outcome.content_stream + ); + } + + Ok(()) +} diff --git a/llama-cpp-bindings-tests/tests/gemma4_classifier_emits_reasoning_for_multimodal_thinking_prompt.rs b/llama-cpp-bindings-tests/tests/gemma4_classifier_emits_reasoning_for_multimodal_thinking_prompt.rs new file mode 100644 index 00000000..b64b89a6 --- /dev/null +++ b/llama-cpp-bindings-tests/tests/gemma4_classifier_emits_reasoning_for_multimodal_thinking_prompt.rs @@ -0,0 +1,109 @@ +#![cfg(feature = "multimodal_capable")] + +use std::num::NonZeroU32; + +use anyhow::Result; +use llama_cpp_bindings::context::LlamaContext; +use llama_cpp_bindings::context::params::LlamaContextParams; +use llama_cpp_bindings::llama_backend::LlamaBackend; +use llama_cpp_bindings::llama_batch::LlamaBatch; +use llama_cpp_bindings::model::LlamaModel; +use llama_cpp_bindings::mtmd::MtmdBitmap; +use llama_cpp_bindings::mtmd::MtmdContext; +use llama_cpp_bindings::mtmd::MtmdContextParams; +use llama_cpp_bindings::mtmd::MtmdInputText; +use llama_cpp_bindings::mtmd::mtmd_default_marker; +use llama_cpp_bindings::sampling::LlamaSampler; +use llama_cpp_bindings_tests::classify_sample_loop::ClassifySampleLoop; +use llama_cpp_bindings_tests::gpu_backend::inference_model_params; +use llama_cpp_bindings_tests::gpu_backend::require_compiled_backends_present; +use llama_cpp_bindings_tests::test_model::download_file_from; +use llama_cpp_bindings_tests::test_model::fixtures_dir; + +const GEMMA4_REPO: &str = "unsloth/gemma-4-E4B-it-GGUF"; +const GEMMA4_FILE: &str = "gemma-4-E4B-it-Q4_K_M.gguf"; +const GEMMA4_MMPROJ_FILE: &str = "mmproj-F16.gguf"; + +const MAX_GENERATED_TOKENS: i32 = 200; + +#[test] +fn gemma4_classifier_emits_reasoning_for_multimodal_thinking_prompt() -> Result<()> { + let backend = LlamaBackend::init()?; + require_compiled_backends_present()?; + + let model_path = download_file_from(GEMMA4_REPO, GEMMA4_FILE)?; + let mmproj_path = download_file_from(GEMMA4_REPO, GEMMA4_MMPROJ_FILE)?; + let params = inference_model_params(); + let model = LlamaModel::load_from_file(&backend, &model_path, ¶ms)?; + + let mtmd_params = MtmdContextParams::default(); + let mmproj_str = mmproj_path + .to_str() + .ok_or_else(|| anyhow::anyhow!("mmproj path is not valid UTF-8"))?; + let mtmd_ctx = MtmdContext::init_from_file(mmproj_str, &model, &mtmd_params)?; + + let context_params = LlamaContextParams::default() + .with_n_ctx(NonZeroU32::new(8192)) + .with_n_batch(512); + let mut context = LlamaContext::from_model(&model, &backend, context_params)?; + + let image_path = fixtures_dir().join("llamas.jpg"); + let image_path_str = image_path + .to_str() + .ok_or_else(|| anyhow::anyhow!("image path is not valid UTF-8"))?; + let bitmap = MtmdBitmap::from_file(&mtmd_ctx, image_path_str)?; + + let marker = mtmd_default_marker(); + let prompt = format!( + "user\n{marker}What animals do you see in this image?\nmodel\n<|channel>thought\n" + ); + + let input_text = MtmdInputText { + text: prompt, + add_special: false, + parse_special: true, + }; + + let chunks = mtmd_ctx.tokenize(input_text, &[&bitmap])?; + + let mut classifier = model.sampled_token_classifier(); + let n_past = + classifier.eval_multimodal_chunks(&chunks, &mtmd_ctx, &context, 0, 0, 512, true)?; + + let mut sampler = LlamaSampler::chain_simple([ + LlamaSampler::penalties(64, 1.1, 0.0, 0.0), + LlamaSampler::top_k(40), + LlamaSampler::top_p(0.9, 1), + LlamaSampler::min_p(0.05, 1), + LlamaSampler::temp(0.7), + LlamaSampler::dist(0x00C0_FFEE), + ]); + + let mut batch = LlamaBatch::new(2048, 1)?; + let outcome = ClassifySampleLoop { + model: &model, + classifier: &mut classifier, + sampler: &mut sampler, + context: &mut context, + batch: &mut batch, + initial_position: n_past, + max_generated_tokens: MAX_GENERATED_TOKENS, + } + .run()?; + + let usage = classifier.usage(); + + if outcome.observed_reasoning == 0 { + anyhow::bail!( + "Gemma 4 multimodal + thinking: classifier must emit at least one Reasoning token \ + when the prompt opens a `<|channel>thought` block; outcome={outcome:?}" + ); + } + if usage.reasoning_tokens == 0 { + anyhow::bail!( + "Gemma 4 multimodal + thinking: usage.reasoning_tokens must be non-zero; usage={usage:?}" + ); + } + + Ok(()) +} diff --git a/llama-cpp-bindings-tests/tests/gemma4_parses_tool_call_payload.rs b/llama-cpp-bindings-tests/tests/gemma4_parses_tool_call_payload.rs new file mode 100644 index 00000000..87204774 --- /dev/null +++ b/llama-cpp-bindings-tests/tests/gemma4_parses_tool_call_payload.rs @@ -0,0 +1,67 @@ +use anyhow::Result; +use anyhow::bail; +use llama_cpp_bindings::ChatMessageParseOutcome; +use llama_cpp_bindings::ToolCallArguments; +use llama_cpp_bindings::llama_backend::LlamaBackend; +use llama_cpp_bindings::model::LlamaModel; +use llama_cpp_bindings_tests::gpu_backend::inference_model_params; +use llama_cpp_bindings_tests::gpu_backend::require_compiled_backends_present; +use llama_cpp_bindings_tests::test_model::download_file_from; + +const GEMMA4_REPO: &str = "unsloth/gemma-4-E4B-it-GGUF"; +const GEMMA4_FILE: &str = "gemma-4-E4B-it-Q4_K_M.gguf"; + +const TOOLS_JSON: &str = r#"[ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather for a location", + "parameters": { + "type": "object", + "properties": { + "location": {"type": "string", "description": "The city name"} + }, + "required": ["location"] + } + } + } +]"#; + +const GEMMA4_PAIRED_QUOTE_PAYLOAD: &str = + "<|tool_call>call:get_weather{location:<|\"|>Paris<|\"|>}"; + +#[test] +fn gemma4_parses_tool_call_payload() -> Result<()> { + let backend = LlamaBackend::init()?; + require_compiled_backends_present()?; + + let path = download_file_from(GEMMA4_REPO, GEMMA4_FILE)?; + let params = inference_model_params(); + let model = LlamaModel::load_from_file(&backend, &path, ¶ms)?; + + let outcome = model.parse_chat_message(TOOLS_JSON, GEMMA4_PAIRED_QUOTE_PAYLOAD, false)?; + + let ChatMessageParseOutcome::Recognized(parsed) = outcome else { + bail!("expected Recognized for Gemma 4 PairedQuote on a Gemma-4 model; got Unrecognized"); + }; + assert_eq!( + parsed.tool_calls.len(), + 1, + "expected one tool call; got {:?}", + parsed.tool_calls + ); + assert_eq!(parsed.tool_calls[0].name, "get_weather"); + let location = match &parsed.tool_calls[0].arguments { + ToolCallArguments::ValidJson(value) => value + .get("location") + .and_then(|v| v.as_str()) + .map(str::to_owned), + ToolCallArguments::InvalidJson(raw) => { + bail!("expected ValidJson, got InvalidJson: {raw}"); + } + }; + assert_eq!(location.as_deref(), Some("Paris")); + + Ok(()) +} diff --git a/llama-cpp-bindings-tests/tests/gemma4_template_override_returns_full_markers.rs b/llama-cpp-bindings-tests/tests/gemma4_template_override_returns_full_markers.rs new file mode 100644 index 00000000..8acea37b --- /dev/null +++ b/llama-cpp-bindings-tests/tests/gemma4_template_override_returns_full_markers.rs @@ -0,0 +1,46 @@ +use anyhow::Result; +use llama_cpp_bindings::ToolCallArgsShape; +use llama_cpp_bindings::llama_backend::LlamaBackend; +use llama_cpp_bindings::model::LlamaModel; +use llama_cpp_bindings_tests::gpu_backend::inference_model_params; +use llama_cpp_bindings_tests::gpu_backend::require_compiled_backends_present; +use llama_cpp_bindings_tests::test_model::download_file_from; + +const GEMMA4_REPO: &str = "unsloth/gemma-4-E4B-it-GGUF"; +const GEMMA4_FILE: &str = "gemma-4-E4B-it-Q4_K_M.gguf"; + +#[test] +fn gemma4_template_override_returns_full_markers() -> Result<()> { + let backend = LlamaBackend::init()?; + require_compiled_backends_present()?; + + let path = download_file_from(GEMMA4_REPO, GEMMA4_FILE)?; + let params = inference_model_params(); + let model = LlamaModel::load_from_file(&backend, &path, ¶ms)?; + + let template = model + .chat_template(None) + .expect("Gemma 4 chat template must be present"); + let template_str = template.to_str().expect("template must be valid UTF-8"); + assert!( + template_str.contains("<|tool_call>call:"), + "Gemma 4 chat template must contain '<|tool_call>call:' fingerprint; \ + template starts with: {:?}", + &template_str[..template_str.len().min(200)], + ); + + let markers = model + .tool_call_markers() + .expect("Gemma 4 must produce ToolCallMarkers via override registry"); + + assert_eq!(markers.open, "<|tool_call>call:"); + assert_eq!(markers.close, "}"); + let ToolCallArgsShape::PairedQuote(shape) = markers.args_shape else { + panic!("expected PairedQuote variant, got {:?}", markers.args_shape); + }; + assert_eq!(shape.name_args_separator, "{"); + assert_eq!(shape.value_quote.open, "<|\"|>"); + assert_eq!(shape.value_quote.close, "<|\"|>"); + + Ok(()) +} diff --git a/llama-cpp-bindings-tests/tests/glm47_classifier_does_not_emit_reasoning_for_thinking_disabled_prompt.rs b/llama-cpp-bindings-tests/tests/glm47_classifier_does_not_emit_reasoning_for_thinking_disabled_prompt.rs new file mode 100644 index 00000000..cea184bf --- /dev/null +++ b/llama-cpp-bindings-tests/tests/glm47_classifier_does_not_emit_reasoning_for_thinking_disabled_prompt.rs @@ -0,0 +1,127 @@ +use std::num::NonZeroU32; + +use anyhow::Result; +use llama_cpp_bindings::context::LlamaContext; +use llama_cpp_bindings::context::params::LlamaContextParams; +use llama_cpp_bindings::llama_backend::LlamaBackend; +use llama_cpp_bindings::llama_batch::LlamaBatch; +use llama_cpp_bindings::model::AddBos; +use llama_cpp_bindings::model::LlamaModel; +use llama_cpp_bindings::sampling::LlamaSampler; +use llama_cpp_bindings_tests::classify_sample_loop::ClassifySampleLoop; +use llama_cpp_bindings_tests::gpu_backend::inference_model_params; +use llama_cpp_bindings_tests::gpu_backend::require_compiled_backends_present; +use llama_cpp_bindings_tests::test_model::download_file_from; + +const GLM47_REPO: &str = "unsloth/GLM-4.7-Flash-GGUF"; +const GLM47_FILE: &str = "GLM-4.7-Flash-Q4_K_M.gguf"; + +const MAX_GENERATED_TOKENS: i32 = 200; + +// GLM-4.7-Flash with reasoning disabled: the chat template renders a closed +// `` immediately after `<|assistant|>\n`, leaving the model outside +// the reasoning section before generation begins. No reasoning tokens should +// ever be classified. +const GLM47_THINKING_DISABLED_PROMPT: &str = "\ +<|user|> +What is 2 + 2? +<|assistant|> + + +"; + +const FORBIDDEN_MARKERS: &[&str] = &["", ""]; + +#[test] +fn glm47_classifier_does_not_emit_reasoning_for_thinking_disabled_prompt() -> Result<()> { + let backend = LlamaBackend::init()?; + require_compiled_backends_present()?; + + let path = download_file_from(GLM47_REPO, GLM47_FILE)?; + let params = inference_model_params(); + let model = LlamaModel::load_from_file(&backend, &path, ¶ms)?; + + let mut classifier = model.sampled_token_classifier(); + let prompt_tokens = model.str_to_token(GLM47_THINKING_DISABLED_PROMPT, AddBos::Never)?; + let prompt_token_count = u64::try_from(prompt_tokens.len())?; + + let mut batch = LlamaBatch::new(2048, 1)?; + classifier.feed_prompt_sequence_to_batch(&mut batch, &prompt_tokens, 0, false)?; + + let context_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(8192)); + let mut context = LlamaContext::from_model(&model, &backend, context_params)?; + + context.decode(&mut batch)?; + + let promoted = classifier.commit_prompt_tokens(); + assert_eq!(promoted, prompt_token_count); + + let mut sampler = LlamaSampler::chain_simple([ + LlamaSampler::penalties(64, 1.1, 0.0, 0.0), + LlamaSampler::top_k(40), + LlamaSampler::top_p(0.9, 1), + LlamaSampler::min_p(0.05, 1), + LlamaSampler::temp(0.7), + LlamaSampler::dist(0x00C0_FFEE), + ]); + let initial_position = batch.n_tokens(); + let outcome = ClassifySampleLoop { + model: &model, + classifier: &mut classifier, + sampler: &mut sampler, + context: &mut context, + batch: &mut batch, + initial_position, + max_generated_tokens: MAX_GENERATED_TOKENS, + } + .run()?; + + let usage = classifier.usage(); + + assert!( + !outcome.generated_raw.is_empty(), + "GLM-4.7: must generate at least one token" + ); + assert_eq!( + outcome.observed_reasoning, 0, + "GLM-4.7 thinking-disabled: classifier must not emit any Reasoning token \ + when the prompt closes the think block before generation begins; \ + generated={:?}", + outcome.generated_raw + ); + assert_eq!( + outcome.observed_undeterminable, 0, + "GLM-4.7 thinking-disabled: prompt-token replay must move section to Content \ + before generation, so no Undeterminable tokens may be emitted; \ + generated={:?}", + outcome.generated_raw + ); + assert_eq!( + usage.reasoning_tokens, 0, + "GLM-4.7 thinking-disabled: usage.reasoning_tokens must be zero; usage={usage:?}" + ); + assert_eq!( + usage.undeterminable_tokens, 0, + "GLM-4.7 thinking-disabled: usage.undeterminable_tokens must be zero; usage={usage:?}" + ); + assert!( + outcome.observed_content > 0, + "GLM-4.7 thinking-disabled: classifier must emit at least one Content token" + ); + assert_eq!( + usage.completion_tokens(), + outcome.observed_content, + "GLM-4.7 thinking-disabled: completion tokens must equal observed Content tokens" + ); + + for forbidden in FORBIDDEN_MARKERS { + assert!( + !outcome.content_stream.contains(forbidden), + "GLM-4.7 thinking-disabled: content_stream leaked marker {forbidden:?}; \ + content_stream={:?}", + outcome.content_stream + ); + } + + Ok(()) +} diff --git a/llama-cpp-bindings-tests/tests/glm47_classifier_emits_reasoning.rs b/llama-cpp-bindings-tests/tests/glm47_classifier_emits_reasoning.rs new file mode 100644 index 00000000..d4fec908 --- /dev/null +++ b/llama-cpp-bindings-tests/tests/glm47_classifier_emits_reasoning.rs @@ -0,0 +1,152 @@ +use std::num::NonZeroU32; + +use anyhow::Result; +use anyhow::bail; +use llama_cpp_bindings::ChatMessageParseOutcome; +use llama_cpp_bindings::context::LlamaContext; +use llama_cpp_bindings::context::params::LlamaContextParams; +use llama_cpp_bindings::llama_backend::LlamaBackend; +use llama_cpp_bindings::llama_batch::LlamaBatch; +use llama_cpp_bindings::model::AddBos; +use llama_cpp_bindings::model::LlamaModel; +use llama_cpp_bindings::sampling::LlamaSampler; +use llama_cpp_bindings_tests::classify_sample_loop::ClassifySampleLoop; +use llama_cpp_bindings_tests::gpu_backend::inference_model_params; +use llama_cpp_bindings_tests::gpu_backend::require_compiled_backends_present; +use llama_cpp_bindings_tests::test_model::download_file_from; + +const GLM47_REPO: &str = "unsloth/GLM-4.7-Flash-GGUF"; +const GLM47_FILE: &str = "GLM-4.7-Flash-Q4_K_M.gguf"; + +// Budget tuned so the close marker reliably emits — enough thinking space for a +// concise question. The companion prompt is intentionally direct so the model +// finishes thinking quickly and emits . +const MAX_GENERATED_TOKENS: i32 = 1500; + +// GLM-4.7-Flash uses `...` reasoning markers (same lexical form +// as Qwen3.5/3.6) and `<|user|>` / `<|assistant|>` role tokens. The prompt +// ends inside an open `` block so generation resumes in the reasoning +// section, mirroring how the chat template renders when reasoning is enabled. +const GLM47_THINKING_PROMPT: &str = "\ +<|user|> +What is 2 + 2? +<|assistant|> + +"; + +const FORBIDDEN_MARKERS: &[&str] = &["", ""]; + +#[test] +fn glm47_classifier_emits_reasoning_for_thinking_enabled_prompt() -> Result<()> { + let backend = LlamaBackend::init()?; + require_compiled_backends_present()?; + + let path = download_file_from(GLM47_REPO, GLM47_FILE)?; + let params = inference_model_params(); + let model = LlamaModel::load_from_file(&backend, &path, ¶ms)?; + + let mut classifier = model.sampled_token_classifier(); + let prompt_tokens = model.str_to_token(GLM47_THINKING_PROMPT, AddBos::Never)?; + let prompt_token_count = u64::try_from(prompt_tokens.len())?; + + let mut batch = LlamaBatch::new(2048, 1)?; + classifier.feed_prompt_sequence_to_batch(&mut batch, &prompt_tokens, 0, false)?; + + let context_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(8192)); + let mut context = LlamaContext::from_model(&model, &backend, context_params)?; + + context.decode(&mut batch)?; + + let promoted = classifier.commit_prompt_tokens(); + assert_eq!(promoted, prompt_token_count); + + let mut sampler = LlamaSampler::chain_simple([ + LlamaSampler::penalties(64, 1.1, 0.0, 0.0), + LlamaSampler::top_k(40), + LlamaSampler::top_p(0.9, 1), + LlamaSampler::min_p(0.05, 1), + LlamaSampler::temp(0.7), + LlamaSampler::dist(0x00C0_FFEE), + ]); + let initial_position = batch.n_tokens(); + let outcome = ClassifySampleLoop { + model: &model, + classifier: &mut classifier, + sampler: &mut sampler, + context: &mut context, + batch: &mut batch, + initial_position, + max_generated_tokens: MAX_GENERATED_TOKENS, + } + .run()?; + + let usage = classifier.usage(); + let parse_outcome = model.parse_chat_message("[]", &outcome.generated_raw, false)?; + let ChatMessageParseOutcome::Recognized(parsed) = parse_outcome else { + bail!("GLM-4.7 chat template must be recognised by the parser; got Unrecognized"); + }; + + assert!( + !outcome.generated_raw.is_empty(), + "GLM-4.7: must generate at least one token" + ); + assert!( + outcome.observed_reasoning > 0, + "GLM-4.7: classifier must emit at least one Reasoning token when the prompt \ + opens a block; outcome={outcome:?}", + ); + assert!( + usage.reasoning_tokens > 0, + "GLM-4.7: usage.reasoning_tokens must be non-zero when the prompt opens a \ + block; usage was {usage:?}" + ); + assert_eq!( + outcome.observed_undeterminable, 0, + "GLM-4.7: prompt-token replay must move section to Reasoning before generation, \ + so no Undeterminable tokens may be emitted; outcome={outcome:?}" + ); + assert_eq!( + usage.undeterminable_tokens, 0, + "GLM-4.7: usage.undeterminable_tokens must be zero; usage={usage:?}" + ); + assert_eq!( + usage.completion_tokens(), + outcome.observed_content + outcome.observed_reasoning, + "GLM-4.7: completion tokens must equal observed Content + Reasoning" + ); + + if parsed.reasoning_content.is_empty() { + eprintln!( + "GLM-4.7 didn't close its reasoning block within {MAX_GENERATED_TOKENS} tokens — \ + skipping strict parser-equality assertions" + ); + } else { + assert_eq!( + outcome.reasoning_stream, parsed.reasoning_content, + "GLM-4.7: per-token reasoning stream must equal parser-side reasoning_content \ + (any difference means a marker leaked into the user-visible stream)", + ); + assert_eq!( + outcome.content_stream, parsed.content, + "GLM-4.7: per-token content stream must equal parser-side content \ + (any difference means a marker leaked into the user-visible stream)", + ); + } + + for forbidden in FORBIDDEN_MARKERS { + assert!( + !outcome.reasoning_stream.contains(forbidden), + "GLM-4.7: reasoning_stream leaked marker {forbidden:?}; \ + reasoning_stream={:?}", + outcome.reasoning_stream + ); + assert!( + !outcome.content_stream.contains(forbidden), + "GLM-4.7: content_stream leaked marker {forbidden:?}; \ + content_stream={:?}", + outcome.content_stream + ); + } + + Ok(()) +} diff --git a/llama-cpp-bindings-tests/tests/glm47_parses_tool_call_payload.rs b/llama-cpp-bindings-tests/tests/glm47_parses_tool_call_payload.rs new file mode 100644 index 00000000..f3b076ec --- /dev/null +++ b/llama-cpp-bindings-tests/tests/glm47_parses_tool_call_payload.rs @@ -0,0 +1,71 @@ +use anyhow::Result; +use anyhow::bail; +use llama_cpp_bindings::ChatMessageParseOutcome; +use llama_cpp_bindings::ToolCallArguments; +use llama_cpp_bindings::llama_backend::LlamaBackend; +use llama_cpp_bindings::model::LlamaModel; +use llama_cpp_bindings_tests::gpu_backend::inference_model_params; +use llama_cpp_bindings_tests::gpu_backend::require_compiled_backends_present; +use llama_cpp_bindings_tests::test_model::download_file_from; + +const GLM47_REPO: &str = "unsloth/GLM-4.7-Flash-GGUF"; +const GLM47_FILE: &str = "GLM-4.7-Flash-Q4_K_M.gguf"; + +const TOOLS_JSON: &str = r#"[ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather for a location", + "parameters": { + "type": "object", + "properties": { + "location": {"type": "string", "description": "The city name"} + }, + "required": ["location"] + } + } + } +]"#; + +const GLM47_KEY_VALUE_PAYLOAD: &str = "get_weather\ +location\ +Paris\ +"; + +#[test] +fn glm47_parses_tool_call_payload() -> Result<()> { + let backend = LlamaBackend::init()?; + require_compiled_backends_present()?; + + let path = download_file_from(GLM47_REPO, GLM47_FILE)?; + let params = inference_model_params(); + let model = LlamaModel::load_from_file(&backend, &path, ¶ms)?; + + let outcome = model.parse_chat_message(TOOLS_JSON, GLM47_KEY_VALUE_PAYLOAD, false)?; + + let ChatMessageParseOutcome::Recognized(parsed) = outcome else { + bail!( + "expected Recognized for GLM-4.7 key-value tags on a GLM-4.7-Flash model; got Unrecognized" + ); + }; + assert_eq!( + parsed.tool_calls.len(), + 1, + "expected one tool call; got {:?}", + parsed.tool_calls + ); + assert_eq!(parsed.tool_calls[0].name, "get_weather"); + let location = match &parsed.tool_calls[0].arguments { + ToolCallArguments::ValidJson(value) => value + .get("location") + .and_then(|v| v.as_str()) + .map(str::to_owned), + ToolCallArguments::InvalidJson(raw) => { + bail!("expected ValidJson, got InvalidJson: {raw}"); + } + }; + assert_eq!(location.as_deref(), Some("Paris")); + + Ok(()) +} diff --git a/llama-cpp-bindings-tests/tests/glm47_template_override_returns_full_markers.rs b/llama-cpp-bindings-tests/tests/glm47_template_override_returns_full_markers.rs new file mode 100644 index 00000000..72ac1edb --- /dev/null +++ b/llama-cpp-bindings-tests/tests/glm47_template_override_returns_full_markers.rs @@ -0,0 +1,50 @@ +use anyhow::Result; +use llama_cpp_bindings::ToolCallArgsShape; +use llama_cpp_bindings::llama_backend::LlamaBackend; +use llama_cpp_bindings::model::LlamaModel; +use llama_cpp_bindings_tests::gpu_backend::inference_model_params; +use llama_cpp_bindings_tests::gpu_backend::require_compiled_backends_present; +use llama_cpp_bindings_tests::test_model::download_file_from; + +const GLM47_REPO: &str = "unsloth/GLM-4.7-Flash-GGUF"; +const GLM47_FILE: &str = "GLM-4.7-Flash-Q4_K_M.gguf"; + +#[test] +fn glm47_template_override_returns_full_markers() -> Result<()> { + let backend = LlamaBackend::init()?; + require_compiled_backends_present()?; + + let path = download_file_from(GLM47_REPO, GLM47_FILE)?; + let params = inference_model_params(); + let model = LlamaModel::load_from_file(&backend, &path, ¶ms)?; + + let template = model + .chat_template(None) + .expect("GLM-4.7 chat template must be present"); + let template_str = template.to_str().expect("template must be valid UTF-8"); + assert!( + template_str.contains(""), + "GLM-4.7 chat template must contain '' fingerprint; \ + template starts with: {:?}", + &template_str[..template_str.len().min(200)], + ); + + let markers = model + .tool_call_markers() + .expect("GLM-4.7 must produce ToolCallMarkers via override registry"); + + assert_eq!(markers.open, ""); + assert_eq!(markers.close, ""); + let ToolCallArgsShape::KeyValueXmlTags(shape) = markers.args_shape else { + panic!( + "expected KeyValueXmlTags variant, got {:?}", + markers.args_shape + ); + }; + assert_eq!(shape.key_open, ""); + assert_eq!(shape.key_close, ""); + assert_eq!(shape.value_open, ""); + assert_eq!(shape.value_close, ""); + + Ok(()) +} diff --git a/llama-cpp-bindings-tests/tests/ingest_prompt_chunk.rs b/llama-cpp-bindings-tests/tests/ingest_prompt_chunk.rs new file mode 100644 index 00000000..df1af8b6 --- /dev/null +++ b/llama-cpp-bindings-tests/tests/ingest_prompt_chunk.rs @@ -0,0 +1,149 @@ +#![cfg(feature = "multimodal_capable")] + +use anyhow::Result; +use llama_cpp_bindings::ingest_prompt_chunk::ingest_prompt_chunk; +use llama_cpp_bindings::mtmd::MtmdBitmap; +use llama_cpp_bindings::mtmd::MtmdInputChunkType; +use llama_cpp_bindings::mtmd::MtmdInputText; +use llama_cpp_bindings::mtmd::mtmd_default_marker; +use llama_cpp_bindings_tests::FixtureSession; +use llama_cpp_bindings_tests::test_model::fixtures_dir; + +#[test] +fn text_chunk_records_prompt_tokens() -> Result<()> { + let fixture = FixtureSession::open()?; + let model = fixture.default_model(); + let mtmd_ctx = fixture.mtmd_context()?; + + let input_text = MtmdInputText { + text: "hello world".to_owned(), + add_special: false, + parse_special: false, + }; + let chunks = mtmd_ctx.tokenize(input_text, &[])?; + + let text_chunk = (0..chunks.len()) + .filter_map(|index| chunks.get(index)) + .find(|chunk| chunk.chunk_type() == Ok(MtmdInputChunkType::Text)) + .ok_or_else(|| { + anyhow::anyhow!("text-only tokenization should produce at least one text chunk") + })?; + + let n_tokens = text_chunk.n_tokens() as u64; + + let mut classifier = model.sampled_token_classifier(); + + ingest_prompt_chunk(&mut classifier, &text_chunk)?; + + let usage = classifier.usage(); + if usage.prompt_tokens != n_tokens { + anyhow::bail!( + "text chunk must record n_tokens as prompt_tokens; expected {n_tokens}, got {}", + usage.prompt_tokens + ); + } + if usage.input_image_tokens != 0 { + anyhow::bail!( + "text chunk must not bump input_image_tokens; got {}", + usage.input_image_tokens + ); + } + if usage.input_audio_tokens != 0 { + anyhow::bail!( + "text chunk must not bump input_audio_tokens; got {}", + usage.input_audio_tokens + ); + } + + Ok(()) +} + +#[test] +fn image_chunk_records_input_image_tokens_only() -> Result<()> { + let fixture = FixtureSession::open()?; + let model = fixture.default_model(); + let mtmd_ctx = fixture.mtmd_context()?; + + let image_path = fixtures_dir().join("llamas.jpg"); + let image_path_str = image_path + .to_str() + .ok_or_else(|| anyhow::anyhow!("image path is not valid UTF-8"))?; + let bitmap = MtmdBitmap::from_file(mtmd_ctx, image_path_str)?; + + let marker = mtmd_default_marker(); + let input_text = MtmdInputText { + text: marker.to_owned(), + add_special: false, + parse_special: true, + }; + let chunks = mtmd_ctx.tokenize(input_text, &[&bitmap])?; + + let image_chunk = (0..chunks.len()) + .filter_map(|index| chunks.get(index)) + .find(|chunk| chunk.chunk_type() == Ok(MtmdInputChunkType::Image)) + .ok_or_else(|| anyhow::anyhow!("multimodal tokenization should produce an image chunk"))?; + + let n_tokens = image_chunk.n_tokens() as u64; + if n_tokens == 0 { + anyhow::bail!("image chunk should report at least one token"); + } + + let mut classifier = model.sampled_token_classifier(); + + ingest_prompt_chunk(&mut classifier, &image_chunk)?; + + let usage = classifier.usage(); + if usage.input_image_tokens != n_tokens { + anyhow::bail!( + "image chunk must record n_tokens as input_image_tokens; expected {n_tokens}, got {}", + usage.input_image_tokens + ); + } + if usage.prompt_tokens != 0 { + anyhow::bail!( + "image chunk must not bump prompt_tokens; got {}", + usage.prompt_tokens + ); + } + if usage.input_audio_tokens != 0 { + anyhow::bail!( + "image chunk must not bump input_audio_tokens; got {}", + usage.input_audio_tokens + ); + } + + Ok(()) +} + +#[test] +fn text_chunk_drives_marker_state_machine_to_reasoning() -> Result<()> { + let fixture = FixtureSession::open()?; + let model = fixture.default_model(); + let mtmd_ctx = fixture.mtmd_context()?; + + let input_text = MtmdInputText { + text: "<|im_start|>user\nHi<|im_end|>\n<|im_start|>assistant\n\n".to_owned(), + add_special: false, + parse_special: true, + }; + let chunks = mtmd_ctx.tokenize(input_text, &[])?; + + let mut classifier = model.sampled_token_classifier(); + + for index in 0..chunks.len() { + let chunk = chunks + .get(index) + .ok_or_else(|| anyhow::anyhow!("chunk index {index} must exist"))?; + ingest_prompt_chunk(&mut classifier, &chunk)?; + } + + if classifier.current_section() != llama_cpp_bindings::SampledTokenSection::Reasoning { + anyhow::bail!( + "text chunk replay must transition the classifier section to Reasoning when the \ + prompt opens a `` block; got {:?}", + classifier.current_section() + ); + } + + Ok(()) +} diff --git a/llama-cpp-bindings-tests/tests/llama_backend.rs b/llama-cpp-bindings-tests/tests/llama_backend.rs index 6e3a19ec..aec05c41 100644 --- a/llama-cpp-bindings-tests/tests/llama_backend.rs +++ b/llama-cpp-bindings-tests/tests/llama_backend.rs @@ -1,7 +1,7 @@ use anyhow::Result; use llama_cpp_bindings::llama_backend::LlamaBackend; use llama_cpp_bindings::model::LlamaModel; -use llama_cpp_bindings::model::params::LlamaModelParams; +use llama_cpp_bindings_tests::gpu_backend::inference_model_params; use llama_cpp_bindings_tests::test_model; use serial_test::serial; @@ -11,7 +11,7 @@ fn void_logs_suppresses_output() -> Result<()> { let mut backend = LlamaBackend::init()?; backend.void_logs(); let model_path = test_model::download_model()?; - let model_params = LlamaModelParams::default(); + let model_params = inference_model_params(); let _model = LlamaModel::load_from_file(&backend, model_path, &model_params)?; Ok(()) diff --git a/llama-cpp-bindings-tests/tests/llguidance.rs b/llama-cpp-bindings-tests/tests/llguidance.rs index 88e8e711..06427e36 100644 --- a/llama-cpp-bindings-tests/tests/llguidance.rs +++ b/llama-cpp-bindings-tests/tests/llguidance.rs @@ -1,14 +1,16 @@ use std::ffi::CStr; use std::num::NonZeroU32; +use std::sync::Arc; use anyhow::Result; +use llama_cpp_bindings::context::LlamaContext; use llama_cpp_bindings::context::params::LlamaContextParams; use llama_cpp_bindings::llama_batch::LlamaBatch; use llama_cpp_bindings::llguidance_sampler::create_llg_sampler; use llama_cpp_bindings::model::AddBos; use llama_cpp_bindings::sampling::LlamaSampler; use llama_cpp_bindings::token::LlamaToken; -use llama_cpp_bindings_tests::TestFixture; +use llama_cpp_bindings_tests::FixtureSession; use serial_test::serial; const JSON_SCHEMA: &str = @@ -19,7 +21,7 @@ const LARK_GRAMMAR: &str = r#"start: "yes" | "no""#; #[test] #[serial] fn creates_sampler_with_valid_json_schema() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let model = fixture.default_model(); let sampler = create_llg_sampler(model, "json", JSON_SCHEMA)?; @@ -31,7 +33,7 @@ fn creates_sampler_with_valid_json_schema() -> Result<()> { #[test] #[serial] fn creates_sampler_with_valid_regex_grammar() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let model = fixture.default_model(); let sampler = create_llg_sampler(model, "regex", REGEX_GRAMMAR)?; @@ -43,7 +45,7 @@ fn creates_sampler_with_valid_regex_grammar() -> Result<()> { #[test] #[serial] fn creates_sampler_with_valid_lark_grammar() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let model = fixture.default_model(); let sampler = create_llg_sampler(model, "lark", LARK_GRAMMAR)?; @@ -55,7 +57,7 @@ fn creates_sampler_with_valid_lark_grammar() -> Result<()> { #[test] #[serial] fn returns_error_for_unknown_grammar_kind() { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open().expect("open fixture"); let model = fixture.default_model(); let result = create_llg_sampler(model, "not_a_real_kind", "anything"); @@ -65,7 +67,7 @@ fn returns_error_for_unknown_grammar_kind() { #[test] #[serial] fn returns_error_for_malformed_json_schema() { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open().expect("open fixture"); let model = fixture.default_model(); let result = create_llg_sampler(model, "json", "{this is not valid json"); @@ -75,7 +77,7 @@ fn returns_error_for_malformed_json_schema() { #[test] #[serial] fn returns_error_for_malformed_regex() { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open().expect("open fixture"); let model = fixture.default_model(); let result = create_llg_sampler(model, "regex", "[invalid"); @@ -85,7 +87,7 @@ fn returns_error_for_malformed_regex() { #[test] #[serial] fn name_callback_returns_llguidance() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let model = fixture.default_model(); let sampler = create_llg_sampler(model, "regex", REGEX_GRAMMAR)?; @@ -101,7 +103,7 @@ fn name_callback_returns_llguidance() -> Result<()> { #[test] #[serial] fn reset_clears_sampler_state() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let model = fixture.default_model(); let mut sampler = create_llg_sampler(model, "regex", REGEX_GRAMMAR)?; @@ -113,7 +115,7 @@ fn reset_clears_sampler_state() -> Result<()> { #[test] #[serial] fn clone_via_ffi_creates_independent_sampler() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let model = fixture.default_model(); let sampler = create_llg_sampler(model, "regex", REGEX_GRAMMAR)?; @@ -129,11 +131,11 @@ fn clone_via_ffi_creates_independent_sampler() -> Result<()> { #[test] #[serial] fn samples_token_constrained_by_grammar() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let prompt = "Answer yes or no:"; let tokens = model.str_to_token(prompt, AddBos::Always)?; @@ -153,7 +155,7 @@ fn samples_token_constrained_by_grammar() -> Result<()> { #[test] #[serial] fn accept_invalid_token_id_does_not_panic() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let model = fixture.default_model(); let mut sampler = create_llg_sampler(model, "regex", REGEX_GRAMMAR)?; @@ -163,14 +165,43 @@ fn accept_invalid_token_id_does_not_panic() -> Result<()> { Ok(()) } +#[test] +#[serial] +fn approximate_tok_env_returns_same_arc_across_calls() -> Result<()> { + let fixture = FixtureSession::open()?; + let model = fixture.default_model(); + + let first = model.approximate_tok_env(); + let second = model.approximate_tok_env(); + + assert!(Arc::ptr_eq(&first, &second)); + + Ok(()) +} + +#[test] +#[serial] +fn approximate_tok_env_drives_consistent_grammar_constraint() -> Result<()> { + let fixture = FixtureSession::open()?; + let model = fixture.default_model(); + + let first = create_llg_sampler(model, "regex", REGEX_GRAMMAR)?; + let second = create_llg_sampler(model, "regex", REGEX_GRAMMAR)?; + + assert!(!first.sampler.is_null()); + assert!(!second.sampler.is_null()); + + Ok(()) +} + #[test] #[serial] fn apply_through_chain_during_sample_does_not_panic() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let tokens = model.str_to_token("Answer:", AddBos::Always)?; let mut batch = LlamaBatch::new(512, 1)?; diff --git a/llama-cpp-bindings-tests/tests/mistral3_classifier_does_not_emit_reasoning_for_thinking_disabled_prompt.rs b/llama-cpp-bindings-tests/tests/mistral3_classifier_does_not_emit_reasoning_for_thinking_disabled_prompt.rs new file mode 100644 index 00000000..08708097 --- /dev/null +++ b/llama-cpp-bindings-tests/tests/mistral3_classifier_does_not_emit_reasoning_for_thinking_disabled_prompt.rs @@ -0,0 +1,114 @@ +use std::num::NonZeroU32; + +use anyhow::Result; +use llama_cpp_bindings::context::LlamaContext; +use llama_cpp_bindings::context::params::LlamaContextParams; +use llama_cpp_bindings::llama_backend::LlamaBackend; +use llama_cpp_bindings::llama_batch::LlamaBatch; +use llama_cpp_bindings::model::AddBos; +use llama_cpp_bindings::model::LlamaModel; +use llama_cpp_bindings::sampling::LlamaSampler; +use llama_cpp_bindings_tests::classify_sample_loop::ClassifySampleLoop; +use llama_cpp_bindings_tests::gpu_backend::inference_model_params; +use llama_cpp_bindings_tests::gpu_backend::require_compiled_backends_present; +use llama_cpp_bindings_tests::test_model::download_file_from; + +const MISTRAL3_REPO: &str = "unsloth/Ministral-3-14B-Reasoning-2512-GGUF"; +const MISTRAL3_FILE: &str = "Ministral-3-14B-Reasoning-2512-Q4_K_M.gguf"; + +const MAX_GENERATED_TOKENS: i32 = 200; + +// Mirrors what Mistral 3 Reasoning's chat template renders when the caller +// asks for `enable_thinking=false`: the user turn is followed by a closed +// empty `[THINK][/THINK]` block, so generation begins in CONTENT. +const MISTRAL3_THINKING_DISABLED_PROMPT: &str = "\ +[INST]Reply with the single word: four. Do not explain.[/INST][THINK][/THINK]"; + +const FORBIDDEN_MARKERS: &[&str] = &["[THINK]", "[/THINK]"]; + +#[test] +fn mistral3_classifier_does_not_emit_reasoning_for_thinking_disabled_prompt() -> Result<()> { + let backend = LlamaBackend::init()?; + require_compiled_backends_present()?; + + let path = download_file_from(MISTRAL3_REPO, MISTRAL3_FILE)?; + let params = inference_model_params(); + let model = LlamaModel::load_from_file(&backend, &path, ¶ms)?; + + let mut classifier = model.sampled_token_classifier(); + let prompt_tokens = model.str_to_token(MISTRAL3_THINKING_DISABLED_PROMPT, AddBos::Always)?; + let prompt_token_count = u64::try_from(prompt_tokens.len())?; + + let mut batch = LlamaBatch::new(2048, 1)?; + classifier.feed_prompt_sequence_to_batch(&mut batch, &prompt_tokens, 0, false)?; + + let context_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(8192)); + let mut context = LlamaContext::from_model(&model, &backend, context_params)?; + + context.decode(&mut batch)?; + + let promoted = classifier.commit_prompt_tokens(); + assert_eq!(promoted, prompt_token_count); + + let mut sampler = LlamaSampler::greedy(); + let initial_position = batch.n_tokens(); + let outcome = ClassifySampleLoop { + model: &model, + classifier: &mut classifier, + sampler: &mut sampler, + context: &mut context, + batch: &mut batch, + initial_position, + max_generated_tokens: MAX_GENERATED_TOKENS, + } + .run()?; + + let usage = classifier.usage(); + + assert!( + !outcome.generated_raw.is_empty(), + "Mistral 3 must generate at least one token" + ); + assert_eq!( + outcome.observed_reasoning, 0, + "Mistral 3 thinking-disabled: classifier must not emit any Reasoning token \ + when the prompt closes the [THINK] block before generation begins; \ + generated={:?}", + outcome.generated_raw + ); + assert_eq!( + outcome.observed_undeterminable, 0, + "Mistral 3 thinking-disabled: prompt-token replay must move section to Content \ + before generation, so no Undeterminable tokens may be emitted; \ + generated={:?}", + outcome.generated_raw + ); + assert_eq!( + usage.reasoning_tokens, 0, + "Mistral 3 thinking-disabled: usage.reasoning_tokens must be zero; usage={usage:?}" + ); + assert_eq!( + usage.undeterminable_tokens, 0, + "Mistral 3 thinking-disabled: usage.undeterminable_tokens must be zero; usage={usage:?}" + ); + assert!( + outcome.observed_content > 0, + "Mistral 3 thinking-disabled: classifier must emit at least one Content token" + ); + assert_eq!( + usage.completion_tokens(), + outcome.observed_content, + "Mistral 3 thinking-disabled: completion tokens must equal observed Content tokens" + ); + + for forbidden in FORBIDDEN_MARKERS { + assert!( + !outcome.content_stream.contains(forbidden), + "Mistral 3 thinking-disabled: content_stream leaked marker {forbidden:?}; \ + content_stream={:?}", + outcome.content_stream + ); + } + + Ok(()) +} diff --git a/llama-cpp-bindings-tests/tests/mistral3_classifier_emits_reasoning.rs b/llama-cpp-bindings-tests/tests/mistral3_classifier_emits_reasoning.rs new file mode 100644 index 00000000..83e39cb5 --- /dev/null +++ b/llama-cpp-bindings-tests/tests/mistral3_classifier_emits_reasoning.rs @@ -0,0 +1,146 @@ +use std::num::NonZeroU32; + +use anyhow::Result; +use anyhow::bail; +use llama_cpp_bindings::ChatMessageParseOutcome; +use llama_cpp_bindings::context::LlamaContext; +use llama_cpp_bindings::context::params::LlamaContextParams; +use llama_cpp_bindings::llama_backend::LlamaBackend; +use llama_cpp_bindings::llama_batch::LlamaBatch; +use llama_cpp_bindings::model::AddBos; +use llama_cpp_bindings::model::LlamaModel; +use llama_cpp_bindings::sampling::LlamaSampler; +use llama_cpp_bindings_tests::classify_sample_loop::ClassifySampleLoop; +use llama_cpp_bindings_tests::gpu_backend::inference_model_params; +use llama_cpp_bindings_tests::gpu_backend::require_compiled_backends_present; +use llama_cpp_bindings_tests::test_model::download_file_from; + +const MISTRAL3_REPO: &str = "unsloth/Ministral-3-14B-Reasoning-2512-GGUF"; +const MISTRAL3_FILE: &str = "Ministral-3-14B-Reasoning-2512-Q4_K_M.gguf"; + +const MAX_GENERATED_TOKENS: i32 = 768; + +// Mistral 3 Reasoning's chat template wraps thoughts in `[THINK]...[/THINK]` and +// relies on a fine-tuned default system prompt to make the model emit them. +// Unlike Qwen3.5/3.6, Mistral does not pre-inject `[THINK]` into the generation +// prompt — the model itself emits the open marker as its first generated token. +// We craft the prompt manually rather than going through the legacy chat-template +// engine to keep the test independent of jinja-engine quirks. +const MISTRAL3_THINKING_PROMPT: &str = "\ +[SYSTEM_PROMPT]# HOW YOU SHOULD THINK AND ANSWER\n\n\ +First draft your thinking process (inner monologue) until you arrive at a response. \ +Format your response using Markdown, and use LaTeX for any mathematical equations. \ +Write both your thoughts and the response in the same language as the input.\n\n\ +Your thinking process must follow the template below:\ +[THINK]Your thoughts or/and draft, like working through an exercise on scratch paper. \ +Be as casual and as long as you want until you are confident to generate the response \ +to the user.[/THINK]Here, provide a self-contained response.[/SYSTEM_PROMPT]\ +[INST]Reply with the single word: four. Do not explain.[/INST]"; + +const FORBIDDEN_MARKERS: &[&str] = &["[THINK]", "[/THINK]"]; + +#[test] +fn mistral3_classifier_emits_reasoning_for_thinking_prompt() -> Result<()> { + let backend = LlamaBackend::init()?; + require_compiled_backends_present()?; + + let path = download_file_from(MISTRAL3_REPO, MISTRAL3_FILE)?; + let params = inference_model_params(); + let model = LlamaModel::load_from_file(&backend, &path, ¶ms)?; + + let mut classifier = model.sampled_token_classifier(); + let prompt_tokens = model.str_to_token(MISTRAL3_THINKING_PROMPT, AddBos::Always)?; + let prompt_token_count = u64::try_from(prompt_tokens.len())?; + + let mut batch = LlamaBatch::new(2048, 1)?; + classifier.feed_prompt_sequence_to_batch(&mut batch, &prompt_tokens, 0, false)?; + + let context_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(8192)); + let mut context = LlamaContext::from_model(&model, &backend, context_params)?; + + context.decode(&mut batch)?; + + let promoted = classifier.commit_prompt_tokens(); + assert_eq!(promoted, prompt_token_count); + + let mut sampler = LlamaSampler::greedy(); + let initial_position = batch.n_tokens(); + let outcome = ClassifySampleLoop { + model: &model, + classifier: &mut classifier, + sampler: &mut sampler, + context: &mut context, + batch: &mut batch, + initial_position, + max_generated_tokens: MAX_GENERATED_TOKENS, + } + .run()?; + + let usage = classifier.usage(); + let parse_outcome = model.parse_chat_message("[]", &outcome.generated_raw, false)?; + let ChatMessageParseOutcome::Recognized(parsed) = parse_outcome else { + bail!("Mistral 3 chat template must be recognised by the parser; got Unrecognized"); + }; + + assert!( + !outcome.generated_raw.is_empty(), + "Mistral 3 must generate at least one token" + ); + assert!( + outcome.observed_reasoning > 0, + "Mistral 3 classifier must emit at least one Reasoning token when the model \ + opens a [THINK] block; outcome={outcome:?}", + ); + assert!( + usage.reasoning_tokens > 0, + "Mistral 3 usage.reasoning_tokens must be non-zero when the model emits a \ + [THINK] block; usage was {usage:?}" + ); + assert_eq!( + outcome.observed_undeterminable, 0, + "Mistral 3: prompt-token replay must transition the section before generation, \ + so no Undeterminable tokens may be emitted; outcome={outcome:?}" + ); + assert_eq!( + usage.undeterminable_tokens, 0, + "Mistral 3: usage.undeterminable_tokens must be zero; usage={usage:?}" + ); + assert_eq!( + usage.completion_tokens(), + outcome.observed_content + outcome.observed_reasoning, + "Mistral 3: completion tokens must equal observed Content + Reasoning" + ); + assert!( + !parsed.reasoning_content.is_empty(), + "Mistral 3 must close its reasoning block within {MAX_GENERATED_TOKENS} tokens; \ + increase the budget or pick a more direct prompt. generated={:?}", + outcome.generated_raw, + ); + assert_eq!( + outcome.reasoning_stream, parsed.reasoning_content, + "Mistral 3: per-token reasoning stream must equal parser-side reasoning_content \ + (any difference means a marker leaked into the user-visible stream)", + ); + assert_eq!( + outcome.content_stream, parsed.content, + "Mistral 3: per-token content stream must equal parser-side content \ + (any difference means a marker leaked into the user-visible stream)", + ); + + for forbidden in FORBIDDEN_MARKERS { + assert!( + !outcome.reasoning_stream.contains(forbidden), + "Mistral 3: reasoning_stream leaked marker {forbidden:?}; \ + reasoning_stream={:?}", + outcome.reasoning_stream + ); + assert!( + !outcome.content_stream.contains(forbidden), + "Mistral 3: content_stream leaked marker {forbidden:?}; \ + content_stream={:?}", + outcome.content_stream + ); + } + + Ok(()) +} diff --git a/llama-cpp-bindings-tests/tests/mistral3_classifier_emits_reasoning_for_multimodal_thinking_prompt.rs b/llama-cpp-bindings-tests/tests/mistral3_classifier_emits_reasoning_for_multimodal_thinking_prompt.rs new file mode 100644 index 00000000..53138078 --- /dev/null +++ b/llama-cpp-bindings-tests/tests/mistral3_classifier_emits_reasoning_for_multimodal_thinking_prompt.rs @@ -0,0 +1,109 @@ +#![cfg(feature = "multimodal_capable")] + +use std::num::NonZeroU32; + +use anyhow::Result; +use llama_cpp_bindings::context::LlamaContext; +use llama_cpp_bindings::context::params::LlamaContextParams; +use llama_cpp_bindings::llama_backend::LlamaBackend; +use llama_cpp_bindings::llama_batch::LlamaBatch; +use llama_cpp_bindings::model::LlamaModel; +use llama_cpp_bindings::mtmd::MtmdBitmap; +use llama_cpp_bindings::mtmd::MtmdContext; +use llama_cpp_bindings::mtmd::MtmdContextParams; +use llama_cpp_bindings::mtmd::MtmdInputText; +use llama_cpp_bindings::mtmd::mtmd_default_marker; +use llama_cpp_bindings::sampling::LlamaSampler; +use llama_cpp_bindings_tests::classify_sample_loop::ClassifySampleLoop; +use llama_cpp_bindings_tests::gpu_backend::inference_model_params; +use llama_cpp_bindings_tests::gpu_backend::require_compiled_backends_present; +use llama_cpp_bindings_tests::test_model::download_file_from; +use llama_cpp_bindings_tests::test_model::fixtures_dir; + +const MISTRAL3_REPO: &str = "unsloth/Ministral-3-14B-Reasoning-2512-GGUF"; +const MISTRAL3_FILE: &str = "Ministral-3-14B-Reasoning-2512-Q4_K_M.gguf"; +const MISTRAL3_MMPROJ_FILE: &str = "mmproj-F16.gguf"; + +const MAX_GENERATED_TOKENS: i32 = 768; + +#[test] +fn mistral3_classifier_emits_reasoning_for_multimodal_thinking_prompt() -> Result<()> { + let backend = LlamaBackend::init()?; + require_compiled_backends_present()?; + + let model_path = download_file_from(MISTRAL3_REPO, MISTRAL3_FILE)?; + let mmproj_path = download_file_from(MISTRAL3_REPO, MISTRAL3_MMPROJ_FILE)?; + let params = inference_model_params(); + let model = LlamaModel::load_from_file(&backend, &model_path, ¶ms)?; + + let mtmd_params = MtmdContextParams::default(); + let mmproj_str = mmproj_path + .to_str() + .ok_or_else(|| anyhow::anyhow!("mmproj path is not valid UTF-8"))?; + let mtmd_ctx = MtmdContext::init_from_file(mmproj_str, &model, &mtmd_params)?; + + let context_params = LlamaContextParams::default() + .with_n_ctx(NonZeroU32::new(8192)) + .with_n_batch(512); + let mut context = LlamaContext::from_model(&model, &backend, context_params)?; + + let image_path = fixtures_dir().join("llamas.jpg"); + let image_path_str = image_path + .to_str() + .ok_or_else(|| anyhow::anyhow!("image path is not valid UTF-8"))?; + let bitmap = MtmdBitmap::from_file(&mtmd_ctx, image_path_str)?; + + let marker = mtmd_default_marker(); + let prompt = format!( + "[SYSTEM_PROMPT]# HOW YOU SHOULD THINK AND ANSWER\n\n\ + First draft your thinking process (inner monologue) until you arrive at a response. \ + Format your response using Markdown, and use LaTeX for any mathematical equations. \ + Write both your thoughts and the response in the same language as the input.\n\n\ + Your thinking process must follow the template below:\ + [THINK]Your thoughts or/and draft, like working through an exercise on scratch paper. \ + Be as casual and as long as you want until you are confident to generate the response \ + to the user.[/THINK]Here, provide a self-contained response.[/SYSTEM_PROMPT]\ + [INST]{marker}What animals do you see in this image?[/INST]" + ); + + let input_text = MtmdInputText { + text: prompt, + add_special: true, + parse_special: true, + }; + + let chunks = mtmd_ctx.tokenize(input_text, &[&bitmap])?; + + let mut classifier = model.sampled_token_classifier(); + let n_past = + classifier.eval_multimodal_chunks(&chunks, &mtmd_ctx, &context, 0, 0, 512, true)?; + + let mut sampler = LlamaSampler::greedy(); + let mut batch = LlamaBatch::new(2048, 1)?; + let outcome = ClassifySampleLoop { + model: &model, + classifier: &mut classifier, + sampler: &mut sampler, + context: &mut context, + batch: &mut batch, + initial_position: n_past, + max_generated_tokens: MAX_GENERATED_TOKENS, + } + .run()?; + + let usage = classifier.usage(); + + if outcome.observed_reasoning == 0 { + anyhow::bail!( + "Mistral 3 multimodal + thinking: classifier must emit at least one Reasoning token \ + when the model opens a `[THINK]` block; outcome={outcome:?}" + ); + } + if usage.reasoning_tokens == 0 { + anyhow::bail!( + "Mistral 3 multimodal + thinking: usage.reasoning_tokens must be non-zero; usage={usage:?}" + ); + } + + Ok(()) +} diff --git a/llama-cpp-bindings-tests/tests/mistral3_parses_tool_call_payload.rs b/llama-cpp-bindings-tests/tests/mistral3_parses_tool_call_payload.rs new file mode 100644 index 00000000..e576de18 --- /dev/null +++ b/llama-cpp-bindings-tests/tests/mistral3_parses_tool_call_payload.rs @@ -0,0 +1,69 @@ +use anyhow::Result; +use anyhow::bail; +use llama_cpp_bindings::ChatMessageParseOutcome; +use llama_cpp_bindings::ToolCallArguments; +use llama_cpp_bindings::llama_backend::LlamaBackend; +use llama_cpp_bindings::model::LlamaModel; +use llama_cpp_bindings_tests::gpu_backend::inference_model_params; +use llama_cpp_bindings_tests::gpu_backend::require_compiled_backends_present; +use llama_cpp_bindings_tests::test_model::download_file_from; + +const MISTRAL3_REPO: &str = "unsloth/Ministral-3-14B-Reasoning-2512-GGUF"; +const MISTRAL3_FILE: &str = "Ministral-3-14B-Reasoning-2512-Q4_K_M.gguf"; + +const TOOLS_JSON: &str = r#"[ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather for a location", + "parameters": { + "type": "object", + "properties": { + "location": {"type": "string", "description": "The city name"} + }, + "required": ["location"] + } + } + } +]"#; + +const MISTRAL3_BRACKETED_JSON_PAYLOAD: &str = + r#"[TOOL_CALLS]get_weather[ARGS]{"location":"Paris"}"#; + +#[test] +fn mistral3_parses_tool_call_payload() -> Result<()> { + let backend = LlamaBackend::init()?; + require_compiled_backends_present()?; + + let path = download_file_from(MISTRAL3_REPO, MISTRAL3_FILE)?; + let params = inference_model_params(); + let model = LlamaModel::load_from_file(&backend, &path, ¶ms)?; + + let outcome = model.parse_chat_message(TOOLS_JSON, MISTRAL3_BRACKETED_JSON_PAYLOAD, false)?; + + let ChatMessageParseOutcome::Recognized(parsed) = outcome else { + bail!( + "expected Recognized for Mistral 3 BracketedJson on a Mistral-3 model; got Unrecognized" + ); + }; + assert_eq!( + parsed.tool_calls.len(), + 1, + "expected one tool call; got {:?}", + parsed.tool_calls + ); + assert_eq!(parsed.tool_calls[0].name, "get_weather"); + let location = match &parsed.tool_calls[0].arguments { + ToolCallArguments::ValidJson(value) => value + .get("location") + .and_then(|v| v.as_str()) + .map(str::to_owned), + ToolCallArguments::InvalidJson(raw) => { + bail!("expected ValidJson, got InvalidJson: {raw}"); + } + }; + assert_eq!(location.as_deref(), Some("Paris")); + + Ok(()) +} diff --git a/llama-cpp-bindings-tests/tests/model.rs b/llama-cpp-bindings-tests/tests/model.rs index 30b61532..b69f0bd9 100644 --- a/llama-cpp-bindings-tests/tests/model.rs +++ b/llama-cpp-bindings-tests/tests/model.rs @@ -7,6 +7,7 @@ use llama_cpp_bindings::ChatTemplateError; use llama_cpp_bindings::LlamaLoraAdapterInitError; use llama_cpp_bindings::LlamaModelLoadError; use llama_cpp_bindings::SampledToken; +use llama_cpp_bindings::context::LlamaContext; use llama_cpp_bindings::context::params::LlamaContextParams; use llama_cpp_bindings::json_schema_to_grammar; use llama_cpp_bindings::llama_batch::LlamaBatch; @@ -15,13 +16,14 @@ use llama_cpp_bindings::model::LlamaChatMessage; use llama_cpp_bindings::model::LlamaModel; use llama_cpp_bindings::model::params::LlamaModelParams; use llama_cpp_bindings::sampling::LlamaSampler; -use llama_cpp_bindings_tests::TestFixture; +use llama_cpp_bindings_tests::FixtureSession; +use llama_cpp_bindings_tests::classify_sample_loop::ClassifySampleLoop; use serial_test::serial; #[test] #[serial] fn model_loads_with_valid_metadata() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let model = fixture.default_model(); assert!(model.n_vocab() > 0); @@ -35,7 +37,7 @@ fn model_loads_with_valid_metadata() -> Result<()> { #[test] #[serial] fn special_tokens_exist() { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open().expect("open fixture"); let model = fixture.default_model(); let bos = model.token_bos(); let eos = model.token_eos(); @@ -47,7 +49,7 @@ fn special_tokens_exist() { #[test] #[serial] fn str_to_token_roundtrip() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let model = fixture.default_model(); let tokens = model.str_to_token("hello world", AddBos::Never)?; assert!(!tokens.is_empty()); @@ -63,7 +65,7 @@ fn str_to_token_roundtrip() -> Result<()> { #[test] #[serial] fn chat_template_returns_non_empty() { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open().expect("open fixture"); let model = fixture.default_model(); let template = model.chat_template(None); @@ -73,7 +75,7 @@ fn chat_template_returns_non_empty() { #[test] #[serial] fn apply_chat_template_produces_prompt() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let model = fixture.default_model(); let template = model.chat_template(None)?; let message = LlamaChatMessage::new("user".to_string(), "hello".to_string())?; @@ -88,7 +90,7 @@ fn apply_chat_template_produces_prompt() -> Result<()> { #[test] #[serial] fn meta_count_returns_positive() { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open().expect("open fixture"); let model = fixture.default_model(); assert!(model.meta_count() > 0); @@ -97,7 +99,7 @@ fn meta_count_returns_positive() { #[test] #[serial] fn tokens_iterator_produces_valid_entries() { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open().expect("open fixture"); let model = fixture.default_model(); let mut count = 0; @@ -116,7 +118,7 @@ fn tokens_iterator_produces_valid_entries() { #[test] #[serial] fn token_to_piece_bytes_returns_bytes_for_known_token() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let model = fixture.default_model(); let tokens = model.str_to_token("hello", AddBos::Never)?; let bytes = model.token_to_piece_bytes(tokens[0], 32, false, None)?; @@ -129,7 +131,7 @@ fn token_to_piece_bytes_returns_bytes_for_known_token() -> Result<()> { #[test] #[serial] fn n_layer_returns_positive() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let model = fixture.default_model(); assert!(model.n_layer()? > 0); @@ -140,7 +142,7 @@ fn n_layer_returns_positive() -> Result<()> { #[test] #[serial] fn n_head_returns_positive() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let model = fixture.default_model(); assert!(model.n_head()? > 0); @@ -151,7 +153,7 @@ fn n_head_returns_positive() -> Result<()> { #[test] #[serial] fn n_head_kv_returns_positive() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let model = fixture.default_model(); assert!(model.n_head_kv()? > 0); @@ -162,7 +164,7 @@ fn n_head_kv_returns_positive() -> Result<()> { #[test] #[serial] fn is_hybrid_returns_bool_for_test_model() { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open().expect("open fixture"); let model = fixture.default_model(); let _ = model.is_hybrid(); @@ -171,7 +173,7 @@ fn is_hybrid_returns_bool_for_test_model() { #[test] #[serial] fn meta_key_by_index_returns_valid_key() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let model = fixture.default_model(); let key = model.meta_key_by_index(0)?; @@ -183,7 +185,7 @@ fn meta_key_by_index_returns_valid_key() -> Result<()> { #[test] #[serial] fn meta_val_str_by_index_returns_valid_value() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let model = fixture.default_model(); let value = model.meta_val_str_by_index(0)?; @@ -195,7 +197,7 @@ fn meta_val_str_by_index_returns_valid_value() -> Result<()> { #[test] #[serial] fn meta_key_by_index_out_of_range_returns_error() { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open().expect("open fixture"); let model = fixture.default_model(); let result = model.meta_key_by_index(999_999); @@ -205,7 +207,7 @@ fn meta_key_by_index_out_of_range_returns_error() { #[test] #[serial] fn meta_val_str_by_index_out_of_range_returns_error() { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open().expect("open fixture"); let model = fixture.default_model(); let result = model.meta_val_str_by_index(999_999); @@ -215,7 +217,7 @@ fn meta_val_str_by_index_out_of_range_returns_error() { #[test] #[serial] fn meta_val_str_returns_value_for_known_key() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let model = fixture.default_model(); let first_key = model.meta_key_by_index(0)?; let value = model.meta_val_str(&first_key)?; @@ -228,7 +230,7 @@ fn meta_val_str_returns_value_for_known_key() -> Result<()> { #[test] #[serial] fn model_size_returns_nonzero() { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open().expect("open fixture"); let model = fixture.default_model(); assert!(model.size() > 0); @@ -237,7 +239,7 @@ fn model_size_returns_nonzero() { #[test] #[serial] fn is_recurrent_returns_false_for_transformer() { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open().expect("open fixture"); let model = fixture.default_model(); assert!(!model.is_recurrent()); @@ -246,7 +248,7 @@ fn is_recurrent_returns_false_for_transformer() { #[test] #[serial] fn rope_type_does_not_panic() { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open().expect("open fixture"); let model = fixture.default_model(); let _rope_type = model.rope_type(); } @@ -254,7 +256,7 @@ fn rope_type_does_not_panic() { #[test] #[serial] fn load_model_with_invalid_path_returns_error() { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open().expect("open fixture"); let backend = fixture.backend(); let model_params = LlamaModelParams::default(); let result = LlamaModel::load_from_file(backend, "/nonexistent/model.gguf", &model_params); @@ -268,7 +270,7 @@ fn load_model_with_invalid_path_returns_error() { #[test] #[serial] fn load_model_with_invalid_file_content_returns_null_result() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model_params = LlamaModelParams::default(); let dummy_path = std::env::temp_dir().join("llama_test_invalid_model.gguf"); @@ -289,7 +291,7 @@ fn load_model_with_non_utf8_path_returns_path_to_str_error() { use std::ffi::OsStr; use std::os::unix::ffi::OsStrExt; - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open().expect("open fixture"); let backend = fixture.backend(); let model_params = LlamaModelParams::default(); let non_utf8_path = std::path::Path::new(OsStr::from_bytes(b"/tmp/\xff\xfe.gguf")); @@ -309,7 +311,7 @@ fn lora_adapter_init_with_non_utf8_path_returns_error() { use std::ffi::OsStr; use std::os::unix::ffi::OsStrExt; - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open().expect("open fixture"); let model = fixture.default_model(); let non_utf8_path = std::path::Path::new(OsStr::from_bytes(b"/tmp/\xff\xfe.gguf")); @@ -324,7 +326,7 @@ fn lora_adapter_init_with_non_utf8_path_returns_error() { #[test] #[serial] fn lora_adapter_init_with_invalid_path_returns_error() { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open().expect("open fixture"); let model = fixture.default_model(); let result = model.lora_adapter_init("/nonexistent/path/lora.gguf"); @@ -337,11 +339,11 @@ fn lora_adapter_init_with_invalid_path_returns_error() { #[test] #[serial] fn new_context_returns_valid_context() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(256)); - let context = model.new_context(backend, ctx_params)?; + let context = LlamaContext::from_model(model, backend, ctx_params)?; assert!(context.n_ctx() > 0); @@ -351,7 +353,7 @@ fn new_context_returns_valid_context() -> Result<()> { #[test] #[serial] fn token_nl_returns_valid_token() { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open().expect("open fixture"); let model = fixture.default_model(); let nl_token = model.token_nl(); @@ -361,7 +363,7 @@ fn token_nl_returns_valid_token() { #[test] #[serial] fn decode_start_token_returns_valid_token() { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open().expect("open fixture"); let model = fixture.default_model(); let _decode_start = model.decode_start_token(); } @@ -369,7 +371,7 @@ fn decode_start_token_returns_valid_token() { #[test] #[serial] fn token_sep_returns_valid_token() { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open().expect("open fixture"); let model = fixture.default_model(); let _sep_token = model.token_sep(); } @@ -377,7 +379,7 @@ fn token_sep_returns_valid_token() { #[test] #[serial] fn token_to_piece_handles_large_token_requiring_buffer_resize() { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open().expect("open fixture"); let model = fixture.default_model(); let mut decoder = encoding_rs::UTF_8.new_decoder(); @@ -390,7 +392,7 @@ fn token_to_piece_handles_large_token_requiring_buffer_resize() { #[test] #[serial] fn token_to_piece_bytes_insufficient_buffer_returns_error() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let model = fixture.default_model(); let tokens = model.str_to_token("hello", AddBos::Never)?; let result = model.token_to_piece_bytes(tokens[0], 1, false, None); @@ -408,7 +410,7 @@ fn token_to_piece_bytes_insufficient_buffer_returns_error() -> Result<()> { #[test] #[serial] fn token_to_piece_with_lstrip() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let model = fixture.default_model(); let mut decoder = encoding_rs::UTF_8.new_decoder(); let tokens = model.str_to_token("hello", AddBos::Never)?; @@ -424,10 +426,118 @@ fn token_to_piece_with_lstrip() -> Result<()> { Ok(()) } +#[test] +#[serial] +fn is_eog_token_classifies_reasoning_variant() { + let fixture = FixtureSession::open().expect("open fixture"); + let model = fixture.default_model(); + let eos = model.token_eos(); + + assert!(model.is_eog_token(&SampledToken::Reasoning(eos))); +} + +#[test] +#[serial] +fn is_eog_token_classifies_tool_call_variant() { + let fixture = FixtureSession::open().expect("open fixture"); + let model = fixture.default_model(); + let eos = model.token_eos(); + + assert!(model.is_eog_token(&SampledToken::ToolCall(eos))); +} + +#[test] +#[serial] +fn is_eog_token_classifies_undeterminable_variant() { + let fixture = FixtureSession::open().expect("open fixture"); + let model = fixture.default_model(); + let eos = model.token_eos(); + + assert!(model.is_eog_token(&SampledToken::Undeterminable(eos))); +} + +#[test] +#[serial] +fn token_to_piece_decodes_reasoning_variant() -> Result<()> { + let fixture = FixtureSession::open()?; + let model = fixture.default_model(); + let mut decoder = encoding_rs::UTF_8.new_decoder(); + let tokens = model.str_to_token("hi", AddBos::Never)?; + + let piece = model.token_to_piece( + &SampledToken::Reasoning(tokens[0]), + &mut decoder, + true, + None, + )?; + + assert!(!piece.is_empty()); + + Ok(()) +} + +#[test] +#[serial] +fn token_to_piece_decodes_tool_call_variant() -> Result<()> { + let fixture = FixtureSession::open()?; + let model = fixture.default_model(); + let mut decoder = encoding_rs::UTF_8.new_decoder(); + let tokens = model.str_to_token("hi", AddBos::Never)?; + + let piece = + model.token_to_piece(&SampledToken::ToolCall(tokens[0]), &mut decoder, true, None)?; + + assert!(!piece.is_empty()); + + Ok(()) +} + +#[test] +#[serial] +fn token_to_piece_decodes_undeterminable_variant() -> Result<()> { + let fixture = FixtureSession::open()?; + let model = fixture.default_model(); + let mut decoder = encoding_rs::UTF_8.new_decoder(); + let tokens = model.str_to_token("hi", AddBos::Never)?; + + let piece = model.token_to_piece( + &SampledToken::Undeterminable(tokens[0]), + &mut decoder, + true, + None, + )?; + + assert!(!piece.is_empty()); + + Ok(()) +} + +#[test] +#[serial] +fn str_to_token_grows_buffer_when_initial_estimation_too_small() -> Result<()> { + let fixture = FixtureSession::open()?; + let model = fixture.default_model(); + + // A short input that tokenises to many small tokens. The initial + // capacity is `max(8, str.len()/2 + 1)` so a string with len < 16 may + // tokenise to >8 tokens, forcing the second `llama_tokenize` call along + // the buffer-grow path. + let many_short_chars = "a b c d e f g h i j k l"; + let tokens = model.str_to_token(many_short_chars, AddBos::Always)?; + + assert!( + tokens.len() > 8, + "expected regrow; got {} tokens", + tokens.len() + ); + + Ok(()) +} + #[test] #[serial] fn n_vocab_matches_tokens_iterator_count() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let model = fixture.default_model(); let n_vocab = model.n_vocab(); let count = model.tokens(false).count(); @@ -440,7 +550,7 @@ fn n_vocab_matches_tokens_iterator_count() -> Result<()> { #[test] #[serial] fn token_attr_returns_valid_attr() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let model = fixture.default_model(); let bos = model.token_bos(); let _attr = model.token_attr(bos)?; @@ -451,7 +561,7 @@ fn token_attr_returns_valid_attr() -> Result<()> { #[test] #[serial] fn vocab_type_returns_valid_type() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let model = fixture.default_model(); let _vocab_type = model.vocab_type()?; @@ -461,7 +571,7 @@ fn vocab_type_returns_valid_type() -> Result<()> { #[test] #[serial] fn apply_chat_template_buffer_resize_with_long_messages() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let model = fixture.default_model(); let template = model.chat_template(None)?; let long_content = "a".repeat(2000); @@ -477,7 +587,7 @@ fn apply_chat_template_buffer_resize_with_long_messages() -> Result<()> { #[test] #[serial] fn meta_val_str_with_long_value_triggers_buffer_resize() { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open().expect("open fixture"); let model = fixture.default_model(); let count = model.meta_count(); @@ -492,7 +602,7 @@ fn meta_val_str_with_long_value_triggers_buffer_resize() { #[test] #[serial] fn str_to_token_with_add_bos_never() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let model = fixture.default_model(); let tokens_with_bos = model.str_to_token("hello", AddBos::Always)?; let tokens_without_bos = model.str_to_token("hello", AddBos::Never)?; @@ -505,7 +615,7 @@ fn str_to_token_with_add_bos_never() -> Result<()> { #[test] #[serial] fn chat_template_with_nonexistent_name_returns_error() { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open().expect("open fixture"); let model = fixture.default_model(); let result = model.chat_template(Some("nonexistent_template_name_xyz")); @@ -516,7 +626,7 @@ fn chat_template_with_nonexistent_name_returns_error() { #[test] #[serial] fn lora_adapter_init_with_invalid_gguf_returns_null_result() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let model = fixture.default_model(); let dummy_path = std::env::temp_dir().join("llama_test_dummy_lora.gguf"); std::fs::write(&dummy_path, b"not a valid gguf")?; @@ -534,7 +644,7 @@ fn lora_adapter_init_with_invalid_gguf_returns_null_result() -> Result<()> { fn str_to_token_with_many_tokens_triggers_buffer_resize() -> Result<()> { use std::fmt::Write; - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let model = fixture.default_model(); let many_numbers = (0..2000).fold(String::new(), |mut accumulator, number| { let _ = write!(accumulator, "{number} "); @@ -551,7 +661,7 @@ fn str_to_token_with_many_tokens_triggers_buffer_resize() -> Result<()> { #[test] #[serial] fn rope_type_returns_valid_result_for_test_model() { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open().expect("open fixture"); let model = fixture.default_model(); let _rope_type = model.rope_type(); @@ -560,7 +670,7 @@ fn rope_type_returns_valid_result_for_test_model() { #[test] #[serial] fn meta_val_str_with_null_byte_in_key_returns_error() { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open().expect("open fixture"); let model = fixture.default_model(); let result = model.meta_val_str("key\0with_null"); @@ -570,12 +680,12 @@ fn meta_val_str_with_null_byte_in_key_returns_error() { #[test] #[serial] fn new_context_with_huge_ctx_returns_null_error() { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open().expect("open fixture"); let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(u32::MAX)); - let result = model.new_context(backend, ctx_params); + let result = LlamaContext::from_model(model, backend, ctx_params); assert!(result.is_err()); } @@ -583,11 +693,11 @@ fn new_context_with_huge_ctx_returns_null_error() { #[test] #[serial] fn sample_returns_result_and_succeeds_with_valid_index() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(256)); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let tokens = model.str_to_token("Hello", AddBos::Always)?; let mut batch = LlamaBatch::new(512, 1)?; @@ -608,12 +718,12 @@ fn sample_returns_result_and_succeeds_with_valid_index() -> Result<()> { #[test] #[serial] fn grammar_sampler_constrains_output_to_yes_or_no() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let prompt = "<|im_start|>user\nIs the sky blue? Answer yes or no.<|im_end|>\n<|im_start|>assistant\n\n\n\n\n"; let tokens = model.str_to_token(prompt, AddBos::Always)?; @@ -629,16 +739,25 @@ fn grammar_sampler_constrains_output_to_yes_or_no() -> Result<()> { LlamaSampler::greedy(), ]); - let mut classifier = model.reasoning_token_classifier()?; - let token = classifier.sample(&mut sampler, &context, batch.n_tokens() - 1)?; + let mut classifier = model.sampled_token_classifier(); + let (raw_token, mut outcomes) = + classifier.sample(&mut sampler, &context, batch.n_tokens() - 1)?; + outcomes.extend(classifier.flush()); + + assert_eq!( + outcomes.len(), + 1, + "expected one finalised outcome after flush" + ); + let outcome = &outcomes[0]; + let raw_as_sampled = SampledToken::Content(raw_token); assert!( - !model.is_eog_token(&token), + !model.is_eog_token(&raw_as_sampled), "Grammar sampler should not allow EOS as first token" ); - let mut decoder = encoding_rs::UTF_8.new_decoder(); - let piece = model.token_to_piece(&token, &mut decoder, true, None)?; + let piece = &outcome.raw_piece; let first_char = piece .chars() .next() @@ -663,12 +782,12 @@ fn grammar_sampler_constrains_output_to_yes_or_no() -> Result<()> { #[test] #[serial] fn json_schema_grammar_sampler_constrains_output_to_json() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let prompt = "<|im_start|>user\nWhat is 2+2? Respond with a JSON object.<|im_end|>\n<|im_start|>assistant\n\n\n\n\n"; let tokens = model.str_to_token(prompt, AddBos::Always)?; @@ -688,16 +807,25 @@ fn json_schema_grammar_sampler_constrains_output_to_json() -> Result<()> { LlamaSampler::greedy(), ]); - let mut classifier = model.reasoning_token_classifier()?; - let token = classifier.sample(&mut sampler, &context, batch.n_tokens() - 1)?; + let mut classifier = model.sampled_token_classifier(); + let (raw_token, mut outcomes) = + classifier.sample(&mut sampler, &context, batch.n_tokens() - 1)?; + outcomes.extend(classifier.flush()); + + assert_eq!( + outcomes.len(), + 1, + "expected one finalised outcome after flush" + ); + let outcome = &outcomes[0]; + let raw_as_sampled = SampledToken::Content(raw_token); assert!( - !model.is_eog_token(&token), + !model.is_eog_token(&raw_as_sampled), "Grammar sampler should not allow EOS as first token" ); - let mut decoder = encoding_rs::UTF_8.new_decoder(); - let piece = model.token_to_piece(&token, &mut decoder, true, None)?; + let piece = &outcome.raw_piece; assert!( piece.starts_with('{'), @@ -715,20 +843,22 @@ fn json_schema_grammar_sampler_constrains_output_to_json() -> Result<()> { #[test] #[serial] fn sample_with_grammar_produces_constrained_output_in_loop() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let prompt = "<|im_start|>user\nIs the sky blue? yes or no<|im_end|>\n<|im_start|>assistant\n\n\n\n\n"; let tokens = model.str_to_token(prompt, AddBos::Always)?; let mut batch = LlamaBatch::new(512, 1)?; - batch.add_sequence(&tokens, 0, false)?; + let mut classifier = model.sampled_token_classifier(); + classifier.feed_prompt_sequence_to_batch(&mut batch, &tokens, 0, false)?; context.decode(&mut batch)?; + classifier.commit_prompt_tokens(); let mut sampler = LlamaSampler::chain_simple([ LlamaSampler::grammar(model, r#"root ::= "yes" | "no""#, "root")?, @@ -736,73 +866,60 @@ fn sample_with_grammar_produces_constrained_output_in_loop() -> Result<()> { LlamaSampler::greedy(), ]); - let mut classifier = model.reasoning_token_classifier()?; - let mut generated = String::new(); - let mut decoder = encoding_rs::UTF_8.new_decoder(); - let mut position = batch.n_tokens(); - let mut observed_content: u64 = 0; - let mut observed_reasoning: u64 = 0; - - for iteration in 0..10 { - let token = classifier.sample(&mut sampler, &context, -1)?; - let is_eog = model.is_eog_token(&token); - - match token { - SampledToken::Content(raw) => { - eprintln!( - " iteration={iteration} token={} eog={is_eog} content", - raw.0 - ); - observed_content += 1; - } - SampledToken::Reasoning(raw) => { - eprintln!( - " iteration={iteration} token={} eog={is_eog} reasoning", - raw.0 - ); - observed_reasoning += 1; - } - SampledToken::Undeterminable(raw) => { - eprintln!( - " iteration={iteration} token={} eog={is_eog} undeterminable", - raw.0 - ); - } - } - - if is_eog { - break; - } - - let piece = model.token_to_piece(&token, &mut decoder, true, None)?; - - eprintln!(" piece='{piece}'"); - - generated.push_str(&piece); - - batch.clear(); - batch.add(&token, position, &[0], true)?; - position += 1; - - context.decode(&mut batch)?; + let initial_position = batch.n_tokens(); + let outcome = ClassifySampleLoop { + model, + classifier: &mut classifier, + sampler: &mut sampler, + context: &mut context, + batch: &mut batch, + initial_position, + max_generated_tokens: 10, } + .run()?; - let lowercase = generated.to_lowercase(); - + let lowercase = outcome.generated_raw.to_lowercase(); assert!( lowercase == "yes" || lowercase == "no", - "Grammar loop should produce 'yes' or 'no', got: '{generated}'" + "Grammar loop should produce 'yes' or 'no', got: '{}'", + outcome.generated_raw + ); + assert!( + outcome.eog_seen, + "loop must terminate via EOG once grammar accepts, not by exhausting the budget; \ + outcome={outcome:?}" + ); + assert_eq!( + outcome.observed_reasoning, 0, + "closed-think prompt must not produce Reasoning tokens; outcome={outcome:?}" + ); + assert_eq!( + outcome.observed_undeterminable, 0, + "prompt-token replay closes the think block before generation, so the section \ + must be Content and no Undeterminable tokens may be emitted; outcome={outcome:?}" + ); + assert_eq!( + outcome.observed_tool_call, 0, + "prompt without tool definitions must not produce ToolCall tokens; outcome={outcome:?}" ); - - let usage = classifier.into_usage(); assert!( - usage.completion_tokens() > 0, - "loop should record at least one completion token" + outcome.observed_content > 0, + "grammar must yield at least one Content token (the answer); outcome={outcome:?}" ); + + let usage = classifier.into_usage(); assert_eq!( usage.completion_tokens(), - observed_content + observed_reasoning, - "completion_tokens must equal observed content + reasoning" + outcome.observed_content, + "for the closed-think grammar prompt, completion_tokens equals observed Content" + ); + assert_eq!( + usage.reasoning_tokens, 0, + "usage.reasoning_tokens must be zero; usage={usage:?}" + ); + assert_eq!( + usage.undeterminable_tokens, 0, + "usage.undeterminable_tokens must be zero; usage={usage:?}" ); Ok(()) @@ -811,12 +928,12 @@ fn sample_with_grammar_produces_constrained_output_in_loop() -> Result<()> { #[test] #[serial] fn sample_without_grammar_produces_multiple_tokens() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let prompt = "<|im_start|>user\nSay hello<|im_end|>\n<|im_start|>assistant\n\n\n\n\n"; @@ -829,34 +946,37 @@ fn sample_without_grammar_produces_multiple_tokens() -> Result<()> { let mut sampler = LlamaSampler::chain_simple([LlamaSampler::temp(0.8), LlamaSampler::greedy()]); - let mut classifier = model.reasoning_token_classifier()?; - let mut token_count: u64 = 0; + let mut classifier = model.sampled_token_classifier(); + let mut sampled_count: u64 = 0; let mut position = batch.n_tokens(); for _ in 0..5 { - let token = classifier.sample(&mut sampler, &context, -1)?; + let (raw_token, _outcomes) = classifier.sample(&mut sampler, &context, -1)?; + let raw_as_sampled = SampledToken::Content(raw_token); - if model.is_eog_token(&token) { + if model.is_eog_token(&raw_as_sampled) { break; } - token_count += 1; + sampled_count += 1; batch.clear(); - batch.add(&token, position, &[0], true)?; + batch.add(&raw_as_sampled, position, &[0], true)?; position += 1; context.decode(&mut batch)?; } + let _ = classifier.flush(); + assert!( - token_count > 0, + sampled_count > 0, "Should produce at least one token without grammar" ); let usage = classifier.into_usage(); assert!( - usage.completion_tokens() >= token_count, - "completion_tokens ({}) must include the {token_count} non-EOG samples", + usage.completion_tokens() >= sampled_count, + "completion_tokens ({}) must include the {sampled_count} non-EOG samples", usage.completion_tokens() ); diff --git a/llama-cpp-bindings-tests/tests/model_helpers.rs b/llama-cpp-bindings-tests/tests/model_helpers.rs new file mode 100644 index 00000000..7605521c --- /dev/null +++ b/llama-cpp-bindings-tests/tests/model_helpers.rs @@ -0,0 +1,62 @@ +use anyhow::Result; +use llama_cpp_bindings_tests::FixtureSession; + +#[test] +fn debug_format_includes_struct_name_and_model_field() -> Result<()> { + let fixture = FixtureSession::open()?; + let model = fixture.default_model(); + + let formatted = format!("{model:?}"); + + assert!(formatted.contains("LlamaModel")); + assert!(formatted.contains("model")); + + Ok(()) +} + +#[test] +fn embedding_model_tool_call_markers_call_does_not_panic() -> Result<()> { + let fixture = FixtureSession::open()?; + let embedding_model = fixture.embedding_model()?; + + let _markers = embedding_model.tool_call_markers(); + + Ok(()) +} + +#[test] +fn embedding_model_streaming_markers_returns_ok_for_a_model_without_tool_calls() -> Result<()> { + let fixture = FixtureSession::open()?; + let embedding_model = fixture.embedding_model()?; + + // The exact set of detected markers depends on the embedding model's chat template; + // assertion is just that the call returns Ok without panicking, exercising the + // streaming_markers + autoparser-fallthrough + override-detect paths even on a model + // that lacks tool calls. + let _markers = embedding_model.streaming_markers()?; + + Ok(()) +} + +#[test] +fn approximate_tok_env_is_cached_across_calls() -> Result<()> { + let fixture = FixtureSession::open()?; + let model = fixture.default_model(); + + let first = model.approximate_tok_env(); + let second = model.approximate_tok_env(); + + assert!(std::sync::Arc::ptr_eq(&first, &second)); + + Ok(()) +} + +#[test] +fn approximate_tok_env_falls_back_to_eos_when_eot_unavailable() -> Result<()> { + let fixture = FixtureSession::open()?; + let embedding_model = fixture.embedding_model()?; + + let _env = embedding_model.approximate_tok_env(); + + Ok(()) +} diff --git a/llama-cpp-bindings-tests/tests/model_params.rs b/llama-cpp-bindings-tests/tests/model_params.rs index ff27e70d..59bd7d51 100644 --- a/llama-cpp-bindings-tests/tests/model_params.rs +++ b/llama-cpp-bindings-tests/tests/model_params.rs @@ -5,14 +5,14 @@ use anyhow::Result; use llama_cpp_bindings::context::params::LlamaContextParams; use llama_cpp_bindings::max_devices; use llama_cpp_bindings::model::params::LlamaModelParams; -use llama_cpp_bindings_tests::TestFixture; +use llama_cpp_bindings_tests::FixtureSession; use llama_cpp_bindings_tests::test_model; use serial_test::serial; #[test] #[serial] fn fit_params_succeeds_with_test_model() -> Result<()> { - let _fixture = TestFixture::shared(); + let _fixture = FixtureSession::open()?; let model_path = test_model::download_model()?; let model_path_str = model_path diff --git a/llama-cpp-bindings-tests/tests/mtmd.rs b/llama-cpp-bindings-tests/tests/mtmd.rs index 71620b71..cd0057bf 100644 --- a/llama-cpp-bindings-tests/tests/mtmd.rs +++ b/llama-cpp-bindings-tests/tests/mtmd.rs @@ -1,6 +1,9 @@ +#![cfg(feature = "multimodal_capable")] + use std::num::NonZeroU32; use anyhow::Result; +use llama_cpp_bindings::context::LlamaContext; use llama_cpp_bindings::context::params::LlamaContextParams; use llama_cpp_bindings::llama_backend::LlamaBackend; use llama_cpp_bindings::model::LlamaModel; @@ -11,7 +14,7 @@ use llama_cpp_bindings::mtmd::MtmdEvalError; use llama_cpp_bindings::mtmd::MtmdInputChunkType; use llama_cpp_bindings::mtmd::MtmdInputChunks; use llama_cpp_bindings::mtmd::MtmdInputText; -use llama_cpp_bindings_tests::TestFixture; +use llama_cpp_bindings_tests::FixtureSession; use llama_cpp_bindings_tests::test_model; use serial_test::serial; @@ -33,7 +36,7 @@ fn eval_synthetic_bitmap( let n_positions = chunks.total_positions(); let context_size = u32::try_from(n_positions + 256).unwrap_or(8192); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(context_size)); - let llama_ctx = model.new_context(backend, ctx_params)?; + let llama_ctx = LlamaContext::from_model(model, backend, ctx_params)?; let n_batch = i32::try_from(llama_ctx.n_batch())?; chunks.eval_chunks(mtmd_ctx, &llama_ctx, 0, 0, n_batch, false)?; @@ -43,13 +46,13 @@ fn eval_synthetic_bitmap( #[test] #[serial] fn eval_chunks_returns_batch_size_exceeds_context_limit_for_huge_batch() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let mtmd_ctx = fixture.mtmd_context()?; let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(64)); - let llama_ctx = model.new_context(backend, ctx_params)?; + let llama_ctx = LlamaContext::from_model(model, backend, ctx_params)?; let chunks = MtmdInputChunks::new()?; let huge_batch = i32::try_from(llama_ctx.n_batch() + 1)?; @@ -67,7 +70,7 @@ fn eval_chunks_returns_batch_size_exceeds_context_limit_for_huge_batch() -> Resu #[test] #[serial] fn from_buffer_creates_bitmap_from_image_bytes() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let mtmd_ctx = fixture.mtmd_context()?; let fixtures = test_model::fixtures_dir(); @@ -85,7 +88,7 @@ fn from_buffer_creates_bitmap_from_image_bytes() -> Result<()> { #[test] #[serial] fn from_file_with_null_byte_in_path_returns_error() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let mtmd_ctx = fixture.mtmd_context()?; let result = MtmdBitmap::from_file(mtmd_ctx, "path\0null"); @@ -97,7 +100,7 @@ fn from_file_with_null_byte_in_path_returns_error() -> Result<()> { #[test] #[serial] fn text_chunk_has_text_type() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let mtmd_ctx = fixture.mtmd_context()?; let image_data = vec![128u8; 64 * 64 * 3]; @@ -120,7 +123,7 @@ fn text_chunk_has_text_type() -> Result<()> { #[test] #[serial] fn text_chunk_returns_text_tokens() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let mtmd_ctx = fixture.mtmd_context()?; let image_data = vec![128u8; 64 * 64 * 3]; @@ -145,7 +148,7 @@ fn text_chunk_returns_text_tokens() -> Result<()> { #[test] #[serial] fn chunk_n_tokens_is_positive() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let mtmd_ctx = fixture.mtmd_context()?; let image_data = vec![128u8; 64 * 64 * 3]; @@ -168,7 +171,7 @@ fn chunk_n_tokens_is_positive() -> Result<()> { #[test] #[serial] fn chunk_n_positions_is_positive() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let mtmd_ctx = fixture.mtmd_context()?; let image_data = vec![128u8; 64 * 64 * 3]; @@ -191,7 +194,7 @@ fn chunk_n_positions_is_positive() -> Result<()> { #[test] #[serial] fn copy_creates_owned_duplicate() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let mtmd_ctx = fixture.mtmd_context()?; let image_data = vec![128u8; 64 * 64 * 3]; @@ -216,7 +219,7 @@ fn copy_creates_owned_duplicate() -> Result<()> { #[test] #[serial] fn text_chunk_id_returns_none() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let mtmd_ctx = fixture.mtmd_context()?; let image_data = vec![128u8; 64 * 64 * 3]; @@ -240,7 +243,7 @@ fn text_chunk_id_returns_none() -> Result<()> { #[test] #[serial] fn image_chunk_returns_none_for_text_tokens() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let mtmd_ctx = fixture.mtmd_context()?; let image_data = vec![128u8; 64 * 64 * 3]; @@ -269,7 +272,7 @@ fn image_chunk_returns_none_for_text_tokens() -> Result<()> { #[test] #[serial] fn image_chunk_id_returns_some() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let mtmd_ctx = fixture.mtmd_context()?; let image_data = vec![128u8; 64 * 64 * 3]; @@ -298,7 +301,7 @@ fn image_chunk_id_returns_some() -> Result<()> { #[test] #[serial] fn init_and_supports_vision() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let mtmd_ctx = fixture.mtmd_context()?; assert!(mtmd_ctx.support_vision()); @@ -309,7 +312,7 @@ fn init_and_supports_vision() -> Result<()> { #[test] #[serial] fn tokenize_text_with_image() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let mtmd_ctx = fixture.mtmd_context()?; let image_data = vec![128u8; 64 * 64 * 3]; @@ -330,7 +333,7 @@ fn tokenize_text_with_image() -> Result<()> { #[test] #[serial] fn eval_chunks_with_standard_image() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let mtmd_ctx = fixture.mtmd_context()?; @@ -350,7 +353,7 @@ fn eval_chunks_with_standard_image() -> Result<()> { let n_positions = chunks.total_positions(); let context_size = u32::try_from(n_positions + 256).unwrap_or(2048); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(context_size)); - let llama_ctx = model.new_context(backend, ctx_params)?; + let llama_ctx = LlamaContext::from_model(model, backend, ctx_params)?; let n_batch = i32::try_from(llama_ctx.n_batch())?; let result = chunks.eval_chunks(mtmd_ctx, &llama_ctx, 0, 0, n_batch, false); @@ -362,7 +365,7 @@ fn eval_chunks_with_standard_image() -> Result<()> { #[test] #[serial] fn eval_chunks_with_varied_dimensions() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let mtmd_ctx = fixture.mtmd_context()?; @@ -384,7 +387,7 @@ fn eval_chunks_with_varied_dimensions() -> Result<()> { #[test] #[serial] fn decode_use_non_causal_returns_bool() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let mtmd_ctx = fixture.mtmd_context()?; let image_data = vec![128u8; 64 * 64 * 3]; @@ -406,7 +409,7 @@ fn decode_use_non_causal_returns_bool() -> Result<()> { #[test] #[serial] fn decode_use_mrope_returns_bool() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let mtmd_ctx = fixture.mtmd_context()?; let _mrope = mtmd_ctx.decode_use_mrope(); @@ -417,7 +420,7 @@ fn decode_use_mrope_returns_bool() -> Result<()> { #[test] #[serial] fn support_audio_returns_bool() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let mtmd_ctx = fixture.mtmd_context()?; let _audio = mtmd_ctx.support_audio(); @@ -428,7 +431,7 @@ fn support_audio_returns_bool() -> Result<()> { #[test] #[serial] fn get_audio_sample_rate_returns_option() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let mtmd_ctx = fixture.mtmd_context()?; let _rate = mtmd_ctx.get_audio_sample_rate(); @@ -439,7 +442,7 @@ fn get_audio_sample_rate_returns_option() -> Result<()> { #[test] #[serial] fn encode_chunk_succeeds_for_image_chunk() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let mtmd_ctx = fixture.mtmd_context()?; let image_data = vec![128u8; 64 * 64 * 3]; @@ -470,7 +473,7 @@ fn encode_chunk_succeeds_for_image_chunk() -> Result<()> { #[test] #[serial] fn tokenize_bitmap_count_mismatch_returns_error() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let mtmd_ctx = fixture.mtmd_context()?; let input_text = MtmdInputText { @@ -490,7 +493,7 @@ fn tokenize_bitmap_count_mismatch_returns_error() -> Result<()> { #[test] #[serial] fn eval_chunks_with_extreme_dimensions_does_not_crash() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let mtmd_ctx = fixture.mtmd_context()?; @@ -524,7 +527,7 @@ fn eval_chunks_with_extreme_dimensions_does_not_crash() -> Result<()> { #[test] #[serial] fn init_from_file_with_null_byte_in_path_returns_error() { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open().expect("open fixture"); let model = fixture.default_model(); let mtmd_params = MtmdContextParams::default(); let result = MtmdContext::init_from_file("path\0null", model, &mtmd_params); @@ -535,7 +538,7 @@ fn init_from_file_with_null_byte_in_path_returns_error() { #[test] #[serial] fn tokenize_with_null_byte_in_text_returns_error() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let mtmd_ctx = fixture.mtmd_context()?; let input_text = MtmdInputText { diff --git a/llama-cpp-bindings-tests/tests/multimodal.rs b/llama-cpp-bindings-tests/tests/multimodal.rs index 335cdf06..b87f93c6 100644 --- a/llama-cpp-bindings-tests/tests/multimodal.rs +++ b/llama-cpp-bindings-tests/tests/multimodal.rs @@ -1,16 +1,18 @@ +#![cfg(feature = "multimodal_capable")] + use std::num::NonZeroU32; use anyhow::{Context, Result}; +use llama_cpp_bindings::SampledTokenClassifier; use llama_cpp_bindings::context::LlamaContext; use llama_cpp_bindings::context::params::LlamaContextParams; use llama_cpp_bindings::llama_batch::LlamaBatch; use llama_cpp_bindings::model::{LlamaChatMessage, LlamaModel}; use llama_cpp_bindings::mtmd::{MtmdBitmap, MtmdInputChunkType, MtmdInputChunks, MtmdInputText}; -use llama_cpp_bindings::reasoning_token_classifier::ReasoningTokenClassifier; use llama_cpp_bindings::sampled_token::SampledToken; use llama_cpp_bindings::sampling::LlamaSampler; use llama_cpp_bindings_sys::llama_pos; -use llama_cpp_bindings_tests::{TestFixture, test_model}; +use llama_cpp_bindings_tests::{FixtureSession, test_model}; struct ChunkTokenBreakdown { text: u64, @@ -55,7 +57,7 @@ struct SamplingTotals { } fn drive_sampling_loop( - classifier: &mut ReasoningTokenClassifier, + classifier: &mut SampledTokenClassifier, model: &LlamaModel, ctx: &mut LlamaContext, starting_position: llama_pos, @@ -67,41 +69,48 @@ fn drive_sampling_loop( observed_content: 0, observed_reasoning: 0, }; - let mut decoder = encoding_rs::UTF_8.new_decoder(); let mut batch = LlamaBatch::new(512, 1)?; let mut current_position = starting_position; for _ in 0..max_tokens { - let token = classifier.sample(&mut sampler, ctx, -1)?; - match token { - SampledToken::Content(_) => totals.observed_content += 1, - SampledToken::Reasoning(_) => totals.observed_reasoning += 1, - SampledToken::Undeterminable(_) => {} + let (raw_token, outcomes) = classifier.sample(&mut sampler, ctx, -1)?; + for outcome in &outcomes { + totals.generated.push_str(&outcome.raw_piece); + match outcome.sampled_token { + SampledToken::Content(_) => totals.observed_content += 1, + SampledToken::Reasoning(_) => totals.observed_reasoning += 1, + SampledToken::ToolCall(_) | SampledToken::Undeterminable(_) => {} + } } - if model.is_eog_token(&token) { + let raw_as_sampled = SampledToken::Content(raw_token); + if model.is_eog_token(&raw_as_sampled) { break; } - let piece = model - .token_to_piece(&token, &mut decoder, false, None) - .with_context(|| "failed to convert token to piece")?; - totals.generated.push_str(&piece); - batch.clear(); - batch.add(&token, current_position, &[0], true)?; + batch.add(&raw_as_sampled, current_position, &[0], true)?; current_position += 1; ctx.decode(&mut batch) .with_context(|| "failed to decode generated token")?; } + for outcome in classifier.flush() { + totals.generated.push_str(&outcome.raw_piece); + match outcome.sampled_token { + SampledToken::Content(_) => totals.observed_content += 1, + SampledToken::Reasoning(_) => totals.observed_reasoning += 1, + SampledToken::ToolCall(_) | SampledToken::Undeterminable(_) => {} + } + } + Ok(totals) } #[test] fn multimodal_vision_inference_produces_output() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let mtmd_ctx = fixture.mtmd_context()?; @@ -110,8 +119,7 @@ fn multimodal_vision_inference_produces_output() -> Result<()> { let ctx_params = LlamaContextParams::default() .with_n_ctx(n_ctx) .with_n_batch(512); - let mut ctx = model - .new_context(backend, ctx_params) + let mut ctx = LlamaContext::from_model(model, backend, ctx_params) .with_context(|| "unable to create llama context")?; assert!( @@ -159,7 +167,7 @@ fn multimodal_vision_inference_produces_output() -> Result<()> { "vision input must produce at least one image chunk" ); - let mut classifier = model.reasoning_token_classifier()?; + let mut classifier = model.sampled_token_classifier(); let n_past = classifier .eval_multimodal_chunks(&chunks, mtmd_ctx, &ctx, 0, 0, 512, true) .with_context(|| "failed to evaluate chunks")?; @@ -168,9 +176,9 @@ fn multimodal_vision_inference_produces_output() -> Result<()> { { let usage = classifier.usage(); - assert_eq!(usage.prompt_tokens(), expected.text); - assert_eq!(usage.input_image_tokens(), expected.image); - assert_eq!(usage.input_audio_tokens(), expected.audio); + assert_eq!(usage.prompt_tokens, expected.text); + assert_eq!(usage.input_image_tokens, expected.image); + assert_eq!(usage.input_audio_tokens, expected.audio); } let totals = drive_sampling_loop(&mut classifier, model, &mut ctx, n_past, 512)?; @@ -183,11 +191,11 @@ fn multimodal_vision_inference_produces_output() -> Result<()> { ); let usage = classifier.into_usage(); - assert_eq!(usage.prompt_tokens(), expected.text); - assert_eq!(usage.input_image_tokens(), expected.image); - assert_eq!(usage.input_audio_tokens(), expected.audio); - assert_eq!(usage.content_tokens(), totals.observed_content); - assert_eq!(usage.reasoning_tokens(), totals.observed_reasoning); + assert_eq!(usage.prompt_tokens, expected.text); + assert_eq!(usage.input_image_tokens, expected.image); + assert_eq!(usage.input_audio_tokens, expected.audio); + assert_eq!(usage.content_tokens, totals.observed_content); + assert_eq!(usage.reasoning_tokens, totals.observed_reasoning); assert_eq!( usage.completion_tokens(), totals.observed_content + totals.observed_reasoning diff --git a/llama-cpp-bindings-tests/tests/parse_chat_message.rs b/llama-cpp-bindings-tests/tests/parse_chat_message.rs new file mode 100644 index 00000000..05b64269 --- /dev/null +++ b/llama-cpp-bindings-tests/tests/parse_chat_message.rs @@ -0,0 +1,113 @@ +use anyhow::Result; +use anyhow::bail; +use llama_cpp_bindings::ChatMessageParseOutcome; +use llama_cpp_bindings_tests::FixtureSession; + +#[test] +fn parses_pure_content_response() -> Result<()> { + let fixture = FixtureSession::open()?; + let model = fixture.default_model(); + + let outcome = model.parse_chat_message("[]", "hello world", false)?; + + let ChatMessageParseOutcome::Recognized(parsed) = outcome else { + bail!("expected Recognized for plain content; got Unrecognized"); + }; + assert!(parsed.tool_calls.is_empty()); + assert!(!parsed.is_empty()); + assert!(parsed.content.contains("hello world")); + + Ok(()) +} + +#[test] +fn parses_reasoning_section_into_reasoning_content() -> Result<()> { + let fixture = FixtureSession::open()?; + let model = fixture.default_model(); + + let input = "step one, step two\n\nactual response"; + let outcome = model.parse_chat_message("[]", input, false)?; + + let ChatMessageParseOutcome::Recognized(parsed) = outcome else { + bail!("expected Recognized for reasoning section; got Unrecognized"); + }; + assert!( + parsed.reasoning_content.contains("step") || parsed.content.contains("step"), + "neither content nor reasoning contains 'step'; content={:?} reasoning={:?}", + parsed.content, + parsed.reasoning_content + ); + + Ok(()) +} + +#[test] +fn parses_empty_input_yields_empty_message() -> Result<()> { + let fixture = FixtureSession::open()?; + let model = fixture.default_model(); + + let outcome = model.parse_chat_message("[]", "", false)?; + + let ChatMessageParseOutcome::Recognized(parsed) = outcome else { + bail!("expected Recognized for empty input; got Unrecognized"); + }; + assert!(parsed.tool_calls.is_empty()); + + Ok(()) +} + +#[test] +fn parses_malformed_tools_json_returns_tools_json_invalid_error() { + let fixture = FixtureSession::open().expect("open fixture"); + let model = fixture.default_model(); + + let result = model.parse_chat_message("not_a_json[}", "hello", false); + + assert!(matches!( + result, + Err(llama_cpp_bindings::ParseChatMessageError::ToolsJsonInvalid( + _ + )) + )); +} + +#[test] +fn parses_non_array_tools_json_returns_tools_json_not_array_error() { + let fixture = FixtureSession::open().expect("open fixture"); + let model = fixture.default_model(); + + let result = model.parse_chat_message("{\"foo\": 1}", "hello", false); + + assert!(matches!( + result, + Err(llama_cpp_bindings::ParseChatMessageError::ToolsJsonNotArray) + )); +} + +#[test] +fn parses_with_tools_null_byte_returns_tools_json_invalid_error() { + let fixture = FixtureSession::open().expect("open fixture"); + let model = fixture.default_model(); + + let result = model.parse_chat_message("[]\0extra", "hello", false); + + assert!(matches!( + result, + Err(llama_cpp_bindings::ParseChatMessageError::ToolsJsonInvalid( + _ + )) + )); +} + +#[test] +fn parses_with_input_null_byte_returns_tools_serialization_error() { + let fixture = FixtureSession::open().expect("open fixture"); + let model = fixture.default_model(); + + let result = model.parse_chat_message("[]", "hello\0world", false); + + assert!(matches!( + result, + Err(llama_cpp_bindings::ParseChatMessageError::ToolsSerialization(_)) + )); +} diff --git a/llama-cpp-bindings-tests/tests/qwen35_chat_inference_emits_reasoning_when_template_auto_opens.rs b/llama-cpp-bindings-tests/tests/qwen35_chat_inference_emits_reasoning_when_template_auto_opens.rs new file mode 100644 index 00000000..88d40f95 --- /dev/null +++ b/llama-cpp-bindings-tests/tests/qwen35_chat_inference_emits_reasoning_when_template_auto_opens.rs @@ -0,0 +1,112 @@ +use anyhow::Result; +use anyhow::bail; +use llama_cpp_bindings::ChatMessageParseOutcome; +use llama_cpp_bindings::context::LlamaContext; +use llama_cpp_bindings::context::params::LlamaContextParams; +use llama_cpp_bindings::llama_backend::LlamaBackend; +use llama_cpp_bindings::llama_batch::LlamaBatch; +use llama_cpp_bindings::model::AddBos; +use llama_cpp_bindings::model::LlamaChatMessage; +use llama_cpp_bindings::model::LlamaModel; +use llama_cpp_bindings::sampling::LlamaSampler; +use llama_cpp_bindings_tests::classify_sample_loop::ClassifySampleLoop; +use llama_cpp_bindings_tests::gpu_backend::inference_model_params; +use llama_cpp_bindings_tests::gpu_backend::require_compiled_backends_present; +use llama_cpp_bindings_tests::test_model::download_file_from; + +const QWEN35_REPO: &str = "unsloth/Qwen3.5-0.8B-GGUF"; +const QWEN35_FILE: &str = "Qwen3.5-0.8B-Q4_K_M.gguf"; + +#[test] +fn qwen35_chat_inference_emits_reasoning_when_template_auto_opens() -> Result<()> { + let backend = LlamaBackend::init()?; + require_compiled_backends_present()?; + + let path = download_file_from(QWEN35_REPO, QWEN35_FILE)?; + let params = inference_model_params(); + let model = LlamaModel::load_from_file(&backend, &path, ¶ms)?; + + let context_params = LlamaContextParams::default(); + let mut context = LlamaContext::from_model(&model, &backend, context_params)?; + + let chat_template = model.chat_template(None)?; + let messages = vec![LlamaChatMessage::new( + "user".to_owned(), + "Hello! How are you?".to_owned(), + )?]; + let prompt = model.apply_chat_template(&chat_template, &messages, true)?; + + let mut classifier = model.sampled_token_classifier(); + let tokens = model.str_to_token(&prompt, AddBos::Always)?; + let prompt_token_count = u64::try_from(tokens.len())?; + + let mut batch = LlamaBatch::new(512, 1)?; + classifier.feed_prompt_sequence_to_batch(&mut batch, &tokens, 0, false)?; + + context.decode(&mut batch)?; + + let promoted = classifier.commit_prompt_tokens(); + assert_eq!(promoted, prompt_token_count); + + let mut sampler = LlamaSampler::greedy(); + let initial_position = batch.n_tokens(); + let outcome = ClassifySampleLoop { + model: &model, + classifier: &mut classifier, + sampler: &mut sampler, + context: &mut context, + batch: &mut batch, + initial_position, + max_generated_tokens: 1024, + } + .run()?; + + assert!( + !outcome.generated_raw.is_empty(), + "Qwen3.5 must generate at least one token" + ); + assert!( + outcome.observed_reasoning > 0, + "Qwen3.5 chat template auto-opens reasoning, so the classifier must emit at \ + least one Reasoning token; outcome={outcome:?}" + ); + assert!( + outcome.observed_content > 0, + "Qwen3.5 must emit at least one Content token after ; outcome={outcome:?}" + ); + assert_eq!( + outcome.observed_undeterminable, 0, + "Qwen3.5 chat template auto-opens reasoning, so the classifier must never emit \ + Undeterminable; outcome={outcome:?}" + ); + assert_eq!( + outcome.observed_tool_call, 0, + "chat without tool definitions must not produce ToolCall tokens; outcome={outcome:?}" + ); + + let parse_outcome = model.parse_chat_message("[]", &outcome.generated_raw, false)?; + let ChatMessageParseOutcome::Recognized(parsed) = parse_outcome else { + bail!("Qwen3.5 chat template must be recognised by the parser; got Unrecognized"); + }; + assert!( + !parsed.content.is_empty(), + "parser must see post- content in generated text; generated={:?}", + outcome.generated_raw + ); + + let usage = classifier.into_usage(); + assert_eq!( + usage.prompt_tokens, prompt_token_count, + "prompt_tokens must equal the tokenizer's prompt length" + ); + assert_eq!( + usage.reasoning_tokens, outcome.observed_reasoning, + "reasoning_tokens must equal observed Reasoning variants" + ); + assert_eq!( + usage.undeterminable_tokens, 0, + "Qwen3.5 with auto-opening reasoning must never produce Undeterminable" + ); + + Ok(()) +} diff --git a/llama-cpp-bindings-tests/tests/qwen35_classifier_does_not_emit_reasoning_for_thinking_disabled_prompt.rs b/llama-cpp-bindings-tests/tests/qwen35_classifier_does_not_emit_reasoning_for_thinking_disabled_prompt.rs new file mode 100644 index 00000000..075ea34b --- /dev/null +++ b/llama-cpp-bindings-tests/tests/qwen35_classifier_does_not_emit_reasoning_for_thinking_disabled_prompt.rs @@ -0,0 +1,129 @@ +use std::num::NonZeroU32; + +use anyhow::Result; +use llama_cpp_bindings::context::LlamaContext; +use llama_cpp_bindings::context::params::LlamaContextParams; +use llama_cpp_bindings::llama_backend::LlamaBackend; +use llama_cpp_bindings::llama_batch::LlamaBatch; +use llama_cpp_bindings::model::AddBos; +use llama_cpp_bindings::model::LlamaModel; +use llama_cpp_bindings::sampling::LlamaSampler; +use llama_cpp_bindings_tests::classify_sample_loop::ClassifySampleLoop; +use llama_cpp_bindings_tests::gpu_backend::inference_model_params; +use llama_cpp_bindings_tests::gpu_backend::require_compiled_backends_present; +use llama_cpp_bindings_tests::test_model::download_file_from; + +const QWEN35_REPO: &str = "unsloth/Qwen3.5-0.8B-GGUF"; +const QWEN35_FILE: &str = "Qwen3.5-0.8B-Q4_K_M.gguf"; + +const MAX_GENERATED_TOKENS: i32 = 200; + +// Mirrors what Qwen3.5's chat template renders when `enable_thinking=false`: +// the assistant header is followed by a closed empty `...` +// block, so generation begins in CONTENT — no reasoning tokens should ever be +// classified. +const QWEN35_THINKING_DISABLED_PROMPT: &str = "\ +<|im_start|>user +What is 2 + 2?<|im_end|> +<|im_start|>assistant + + + + +"; + +const FORBIDDEN_MARKERS: &[&str] = &["", ""]; + +#[test] +fn qwen35_classifier_does_not_emit_reasoning_for_thinking_disabled_prompt() -> Result<()> { + let backend = LlamaBackend::init()?; + require_compiled_backends_present()?; + + let path = download_file_from(QWEN35_REPO, QWEN35_FILE)?; + let params = inference_model_params(); + let model = LlamaModel::load_from_file(&backend, &path, ¶ms)?; + + let mut classifier = model.sampled_token_classifier(); + let prompt_tokens = model.str_to_token(QWEN35_THINKING_DISABLED_PROMPT, AddBos::Never)?; + let prompt_token_count = u64::try_from(prompt_tokens.len())?; + + let mut batch = LlamaBatch::new(2048, 1)?; + classifier.feed_prompt_sequence_to_batch(&mut batch, &prompt_tokens, 0, false)?; + + let context_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(8192)); + let mut context = LlamaContext::from_model(&model, &backend, context_params)?; + + context.decode(&mut batch)?; + + let promoted = classifier.commit_prompt_tokens(); + assert_eq!(promoted, prompt_token_count); + + let mut sampler = LlamaSampler::chain_simple([ + LlamaSampler::penalties(64, 1.1, 0.0, 0.0), + LlamaSampler::top_k(40), + LlamaSampler::top_p(0.9, 1), + LlamaSampler::min_p(0.05, 1), + LlamaSampler::temp(0.7), + LlamaSampler::dist(0x00C0_FFEE), + ]); + let initial_position = batch.n_tokens(); + let outcome = ClassifySampleLoop { + model: &model, + classifier: &mut classifier, + sampler: &mut sampler, + context: &mut context, + batch: &mut batch, + initial_position, + max_generated_tokens: MAX_GENERATED_TOKENS, + } + .run()?; + + let usage = classifier.usage(); + + assert!( + !outcome.generated_raw.is_empty(), + "Qwen3.5 must generate at least one token" + ); + assert_eq!( + outcome.observed_reasoning, 0, + "Qwen3.5 thinking-disabled: classifier must not emit any Reasoning token \ + when the prompt closes the think block before generation begins; \ + generated={:?}", + outcome.generated_raw + ); + assert_eq!( + outcome.observed_undeterminable, 0, + "Qwen3.5 thinking-disabled: prompt-token replay must move section to Content \ + before generation, so no Undeterminable tokens may be emitted; \ + generated={:?}", + outcome.generated_raw + ); + assert_eq!( + usage.reasoning_tokens, 0, + "Qwen3.5 thinking-disabled: usage.reasoning_tokens must be zero; usage={usage:?}" + ); + assert_eq!( + usage.undeterminable_tokens, 0, + "Qwen3.5 thinking-disabled: usage.undeterminable_tokens must be zero; usage={usage:?}" + ); + assert!( + outcome.observed_content > 0, + "Qwen3.5 thinking-disabled: classifier must emit at least one Content token" + ); + assert_eq!( + usage.completion_tokens(), + outcome.observed_content, + "Qwen3.5 thinking-disabled: completion tokens must equal observed Content tokens" + ); + + for forbidden in FORBIDDEN_MARKERS { + assert!( + !outcome.content_stream.contains(forbidden), + "Qwen3.5 thinking-disabled: content_stream leaked marker {forbidden:?}; \ + content_stream={:?}", + outcome.content_stream + ); + } + + Ok(()) +} diff --git a/llama-cpp-bindings-tests/tests/qwen35_classifier_emits_reasoning.rs b/llama-cpp-bindings-tests/tests/qwen35_classifier_emits_reasoning.rs new file mode 100644 index 00000000..76671c96 --- /dev/null +++ b/llama-cpp-bindings-tests/tests/qwen35_classifier_emits_reasoning.rs @@ -0,0 +1,160 @@ +use std::num::NonZeroU32; + +use anyhow::Result; +use anyhow::bail; +use llama_cpp_bindings::ChatMessageParseOutcome; +use llama_cpp_bindings::context::LlamaContext; +use llama_cpp_bindings::context::params::LlamaContextParams; +use llama_cpp_bindings::llama_backend::LlamaBackend; +use llama_cpp_bindings::llama_batch::LlamaBatch; +use llama_cpp_bindings::model::AddBos; +use llama_cpp_bindings::model::LlamaModel; +use llama_cpp_bindings::sampling::LlamaSampler; +use llama_cpp_bindings_tests::classify_sample_loop::ClassifySampleLoop; +use llama_cpp_bindings_tests::gpu_backend::inference_model_params; +use llama_cpp_bindings_tests::gpu_backend::require_compiled_backends_present; +use llama_cpp_bindings_tests::test_model::download_file_from; + +const QWEN35_REPO: &str = "unsloth/Qwen3.5-0.8B-GGUF"; +const QWEN35_FILE: &str = "Qwen3.5-0.8B-Q4_K_M.gguf"; + +// Budget tuned so the close marker reliably emits — enough thinking space for a +// concise question. The companion prompt is intentionally direct so the model +// finishes thinking quickly and emits . +const MAX_GENERATED_TOKENS: i32 = 1500; + +// Qwen3.5's chat template injects `\n` directly into the generation prompt +// when `enable_thinking=true` (the default). The legacy `llama_chat_apply_template` +// path bypasses that jinja branch, so we craft the prompt manually to faithfully +// reproduce the production case where the model resumes generation already inside +// the reasoning block. +const QWEN35_THINKING_PROMPT: &str = "\ +<|im_start|>user +What is 2 + 2?<|im_end|> +<|im_start|>assistant + +"; + +const FORBIDDEN_MARKERS: &[&str] = &["", ""]; + +#[test] +fn qwen35_classifier_emits_reasoning_for_thinking_enabled_prompt() -> Result<()> { + let backend = LlamaBackend::init()?; + require_compiled_backends_present()?; + + let path = download_file_from(QWEN35_REPO, QWEN35_FILE)?; + let params = inference_model_params(); + let model = LlamaModel::load_from_file(&backend, &path, ¶ms)?; + + let mut classifier = model.sampled_token_classifier(); + let prompt_tokens = model.str_to_token(QWEN35_THINKING_PROMPT, AddBos::Never)?; + let prompt_token_count = u64::try_from(prompt_tokens.len())?; + + let mut batch = LlamaBatch::new(2048, 1)?; + classifier.feed_prompt_sequence_to_batch(&mut batch, &prompt_tokens, 0, false)?; + + let context_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(8192)); + let mut context = LlamaContext::from_model(&model, &backend, context_params)?; + + context.decode(&mut batch)?; + + let promoted = classifier.commit_prompt_tokens(); + assert_eq!(promoted, prompt_token_count); + + // Mirrors paddler's production sampler chain: rep penalty + top_k/top_p/min_p + + // temp + dist. The 0.8B model loops on plain greedy; this chain breaks the + // loop and lets the model emit `` reliably. + let mut sampler = LlamaSampler::chain_simple([ + LlamaSampler::penalties(64, 1.1, 0.0, 0.0), + LlamaSampler::top_k(40), + LlamaSampler::top_p(0.9, 1), + LlamaSampler::min_p(0.05, 1), + LlamaSampler::temp(0.7), + LlamaSampler::dist(0x00C0_FFEE), + ]); + let initial_position = batch.n_tokens(); + let outcome = ClassifySampleLoop { + model: &model, + classifier: &mut classifier, + sampler: &mut sampler, + context: &mut context, + batch: &mut batch, + initial_position, + max_generated_tokens: MAX_GENERATED_TOKENS, + } + .run()?; + + let usage = classifier.usage(); + let parse_outcome = model.parse_chat_message("[]", &outcome.generated_raw, false)?; + let ChatMessageParseOutcome::Recognized(parsed) = parse_outcome else { + bail!("Qwen3.5 chat template must be recognised by the parser; got Unrecognized"); + }; + + assert!( + !outcome.generated_raw.is_empty(), + "Qwen3.5 must generate at least one token" + ); + assert!( + outcome.observed_reasoning > 0, + "Qwen3.5: classifier must emit at least one Reasoning token when the prompt \ + opens a block; outcome={outcome:?}", + ); + assert!( + usage.reasoning_tokens > 0, + "Qwen3.5: usage.reasoning_tokens must be non-zero when the prompt opens a \ + block; usage was {usage:?}" + ); + assert_eq!( + outcome.observed_undeterminable, 0, + "Qwen3.5: prompt-token replay must move section to Reasoning before generation, \ + so no Undeterminable tokens may be emitted; outcome={outcome:?}" + ); + assert_eq!( + usage.undeterminable_tokens, 0, + "Qwen3.5: usage.undeterminable_tokens must be zero; usage={usage:?}" + ); + assert_eq!( + usage.completion_tokens(), + outcome.observed_content + outcome.observed_reasoning, + "Qwen3.5: completion tokens must equal observed Content + Reasoning" + ); + + // Qwen3.5-0.8B genuinely loops on simple prompts even with rep penalty + + // sampling — it cannot reliably close the reasoning block within a tight + // budget. Skip the strict leak assertions when the model never emitted + // ; the parser-equality check is meaningless then. + if parsed.reasoning_content.is_empty() { + eprintln!( + "Qwen3.5 didn't close its reasoning block within {MAX_GENERATED_TOKENS} tokens — \ + skipping strict parser-equality assertions" + ); + } else { + assert_eq!( + outcome.reasoning_stream, parsed.reasoning_content, + "Qwen3.5: per-token reasoning stream must equal parser-side reasoning_content \ + (any difference means a marker leaked into the user-visible stream)", + ); + assert_eq!( + outcome.content_stream, parsed.content, + "Qwen3.5: per-token content stream must equal parser-side content \ + (any difference means a marker leaked into the user-visible stream)", + ); + } + + for forbidden in FORBIDDEN_MARKERS { + assert!( + !outcome.reasoning_stream.contains(forbidden), + "Qwen3.5: reasoning_stream leaked marker {forbidden:?}; \ + reasoning_stream={:?}", + outcome.reasoning_stream + ); + assert!( + !outcome.content_stream.contains(forbidden), + "Qwen3.5: content_stream leaked marker {forbidden:?}; \ + content_stream={:?}", + outcome.content_stream + ); + } + + Ok(()) +} diff --git a/llama-cpp-bindings-tests/tests/qwen35_classifier_emits_reasoning_for_multimodal_thinking_prompt.rs b/llama-cpp-bindings-tests/tests/qwen35_classifier_emits_reasoning_for_multimodal_thinking_prompt.rs new file mode 100644 index 00000000..be1578f8 --- /dev/null +++ b/llama-cpp-bindings-tests/tests/qwen35_classifier_emits_reasoning_for_multimodal_thinking_prompt.rs @@ -0,0 +1,89 @@ +#![cfg(feature = "multimodal_capable")] + +use std::num::NonZeroU32; + +use anyhow::Result; +use llama_cpp_bindings::context::LlamaContext; +use llama_cpp_bindings::context::params::LlamaContextParams; +use llama_cpp_bindings::llama_batch::LlamaBatch; +use llama_cpp_bindings::mtmd::MtmdBitmap; +use llama_cpp_bindings::mtmd::MtmdInputText; +use llama_cpp_bindings::mtmd::mtmd_default_marker; +use llama_cpp_bindings::sampling::LlamaSampler; +use llama_cpp_bindings_tests::FixtureSession; +use llama_cpp_bindings_tests::classify_sample_loop::ClassifySampleLoop; +use llama_cpp_bindings_tests::test_model::fixtures_dir; + +const MAX_GENERATED_TOKENS: i32 = 200; + +#[test] +fn qwen35_classifier_emits_reasoning_for_multimodal_thinking_prompt() -> Result<()> { + let fixture = FixtureSession::open()?; + let backend = fixture.backend(); + let model = fixture.default_model(); + let mtmd_ctx = fixture.mtmd_context()?; + + let context_params = LlamaContextParams::default() + .with_n_ctx(NonZeroU32::new(4096)) + .with_n_batch(512); + let mut context = LlamaContext::from_model(model, backend, context_params)?; + + let image_path = fixtures_dir().join("llamas.jpg"); + let image_path_str = image_path + .to_str() + .ok_or_else(|| anyhow::anyhow!("image path is not valid UTF-8"))?; + let bitmap = MtmdBitmap::from_file(mtmd_ctx, image_path_str)?; + + let marker = mtmd_default_marker(); + let prompt = format!( + "<|im_start|>user\n{marker}What animals do you see in this image?<|im_end|>\n<|im_start|>assistant\n\n" + ); + + let input_text = MtmdInputText { + text: prompt, + add_special: false, + parse_special: true, + }; + + let chunks = mtmd_ctx.tokenize(input_text, &[&bitmap])?; + + let mut classifier = model.sampled_token_classifier(); + let n_past = classifier.eval_multimodal_chunks(&chunks, mtmd_ctx, &context, 0, 0, 512, true)?; + + let mut sampler = LlamaSampler::chain_simple([ + LlamaSampler::penalties(64, 1.1, 0.0, 0.0), + LlamaSampler::top_k(40), + LlamaSampler::top_p(0.9, 1), + LlamaSampler::min_p(0.05, 1), + LlamaSampler::temp(0.7), + LlamaSampler::dist(0x00C0_FFEE), + ]); + + let mut batch = LlamaBatch::new(2048, 1)?; + let outcome = ClassifySampleLoop { + model, + classifier: &mut classifier, + sampler: &mut sampler, + context: &mut context, + batch: &mut batch, + initial_position: n_past, + max_generated_tokens: MAX_GENERATED_TOKENS, + } + .run()?; + + let usage = classifier.usage(); + + if outcome.observed_reasoning == 0 { + anyhow::bail!( + "Qwen 3.5 multimodal + thinking: classifier must emit at least one Reasoning token \ + when the prompt opens a `` block; outcome={outcome:?}" + ); + } + if usage.reasoning_tokens == 0 { + anyhow::bail!( + "Qwen 3.5 multimodal + thinking: usage.reasoning_tokens must be non-zero; usage={usage:?}" + ); + } + + Ok(()) +} diff --git a/llama-cpp-bindings-tests/tests/qwen35_parses_constrained_schema_payload.rs b/llama-cpp-bindings-tests/tests/qwen35_parses_constrained_schema_payload.rs new file mode 100644 index 00000000..712f09d3 --- /dev/null +++ b/llama-cpp-bindings-tests/tests/qwen35_parses_constrained_schema_payload.rs @@ -0,0 +1,111 @@ +use anyhow::Result; +use anyhow::bail; +use llama_cpp_bindings::ChatMessageParseOutcome; +use llama_cpp_bindings::ToolCallArguments; +use llama_cpp_bindings::llama_backend::LlamaBackend; +use llama_cpp_bindings::model::LlamaModel; +use llama_cpp_bindings_tests::gpu_backend::inference_model_params; +use llama_cpp_bindings_tests::gpu_backend::require_compiled_backends_present; +use llama_cpp_bindings_tests::test_model::download_file_from; +use serde_json::Value; +use serde_json::json; + +const QWEN35_REPO: &str = "unsloth/Qwen3.5-0.8B-GGUF"; +const QWEN35_FILE: &str = "Qwen3.5-0.8B-Q4_K_M.gguf"; + +const NEGOTIATE_WITH_CAT_TOOLS_JSON: &str = r#"[ + { + "type": "function", + "function": { + "name": "negotiate_with_cat", + "description": "Attempt to negotiate with a cat. Outcomes are not guaranteed and may include the silent treatment.", + "parameters": { + "type": "object", + "properties": { + "topic": { + "type": "string", + "description": "What you are trying to negotiate, e.g. 'get off the keyboard' or 'stop knocking things off the table'" + }, + "bribe": { + "type": "string", + "enum": ["tuna", "salmon", "treats", "ear_scritches", "cardboard_box", "none"], + "description": "What you are offering in exchange" + }, + "desperation_level": { + "type": "integer", + "description": "How desperate you are, on a scale from 1 (mildly annoyed human) to 10 (it is 3am)", + "minimum": 1, + "maximum": 10 + } + }, + "required": ["topic"], + "additionalProperties": false + } + } + } +]"#; + +const NEGOTIATE_WITH_CAT_INPUT: &str = "\n\ +\n\ +\n\ +tuna\n\ +\n\ +\n\ +8\n\ +\n\ +\n\ +get off the keyboard\n\ +\n\ +\n\ +"; + +fn arguments_as_json(arguments: &ToolCallArguments) -> Result<&Value> { + match arguments { + ToolCallArguments::ValidJson(value) => Ok(value), + ToolCallArguments::InvalidJson(raw) => { + bail!("expected ValidJson arguments, got InvalidJson: {raw}") + } + } +} + +#[test] +fn qwen35_parses_constrained_schema_payload() -> Result<()> { + let backend = LlamaBackend::init()?; + require_compiled_backends_present()?; + + let path = download_file_from(QWEN35_REPO, QWEN35_FILE)?; + let params = inference_model_params(); + let model = LlamaModel::load_from_file(&backend, &path, ¶ms)?; + + let outcome = model.parse_chat_message( + NEGOTIATE_WITH_CAT_TOOLS_JSON, + NEGOTIATE_WITH_CAT_INPUT, + false, + )?; + + let ChatMessageParseOutcome::Recognized(parsed) = outcome else { + bail!( + "Qwen 3.5's tool-call payload must be parsed by the wrapper-side duck-type pass; \ + got Unrecognized" + ); + }; + + assert_eq!( + parsed.tool_calls.len(), + 1, + "expected exactly one parsed tool call; got {:?}", + parsed.tool_calls + ); + assert_eq!(parsed.tool_calls[0].name, "negotiate_with_cat"); + assert_eq!(parsed.tool_calls[0].id, "call_0"); + assert_eq!( + arguments_as_json(&parsed.tool_calls[0].arguments)?, + &json!({ + "bribe": "tuna", + "desperation_level": 8, + "topic": "get off the keyboard", + }), + ); + + Ok(()) +} diff --git a/llama-cpp-bindings-tests/tests/qwen35_parses_tool_call_payload.rs b/llama-cpp-bindings-tests/tests/qwen35_parses_tool_call_payload.rs new file mode 100644 index 00000000..28efc3fc --- /dev/null +++ b/llama-cpp-bindings-tests/tests/qwen35_parses_tool_call_payload.rs @@ -0,0 +1,128 @@ +use anyhow::Result; +use anyhow::bail; +use llama_cpp_bindings::ChatMessageParseOutcome; +use llama_cpp_bindings::ToolCallArguments; +use llama_cpp_bindings::llama_backend::LlamaBackend; +use llama_cpp_bindings::model::LlamaModel; +use llama_cpp_bindings_tests::gpu_backend::inference_model_params; +use llama_cpp_bindings_tests::gpu_backend::require_compiled_backends_present; +use llama_cpp_bindings_tests::test_model::download_file_from; + +const QWEN35_REPO: &str = "unsloth/Qwen3.5-0.8B-GGUF"; +const QWEN35_FILE: &str = "Qwen3.5-0.8B-Q4_K_M.gguf"; + +const TOOLS_JSON: &str = r#"[ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather for a location", + "parameters": { + "type": "object", + "properties": { + "location": {"type": "string", "description": "The city name"} + }, + "required": ["location"] + } + } + } +]"#; + +const QWEN_XML_PAYLOAD: &str = "\n\ +\n\ +\n\ +Paris\n\ +\n\ +\n\ +"; + +const PARTIAL_QWEN_XML_PAYLOAD: &str = "\n\n Result<(LlamaBackend, LlamaModel)> { + let backend = LlamaBackend::init()?; + require_compiled_backends_present()?; + let path = download_file_from(QWEN35_REPO, QWEN35_FILE)?; + let params = inference_model_params(); + let model = LlamaModel::load_from_file(&backend, &path, ¶ms)?; + + Ok((backend, model)) +} + +#[test] +fn qwen35_parses_tool_call_payload() -> Result<()> { + let (_backend, model) = load_qwen35()?; + + let outcome = model.parse_chat_message(TOOLS_JSON, QWEN_XML_PAYLOAD, false)?; + + let ChatMessageParseOutcome::Recognized(parsed) = outcome else { + bail!("expected Recognized for Qwen XML on a Qwen-3.5 model; got Unrecognized"); + }; + assert_eq!( + parsed.tool_calls.len(), + 1, + "expected one tool call; got {:?}", + parsed.tool_calls + ); + assert_eq!(parsed.tool_calls[0].name, "get_weather"); + let location = match &parsed.tool_calls[0].arguments { + ToolCallArguments::ValidJson(value) => value + .get("location") + .and_then(|v| v.as_str()) + .map(str::to_owned), + ToolCallArguments::InvalidJson(raw) => { + bail!("expected ValidJson, got InvalidJson: {raw}"); + } + }; + assert_eq!(location.as_deref(), Some("Paris")); + + Ok(()) +} + +#[test] +fn qwen35_parses_partial_tool_call_returns_pending_state() -> Result<()> { + let (_backend, model) = load_qwen35()?; + + let outcome = model.parse_chat_message(TOOLS_JSON, PARTIAL_QWEN_XML_PAYLOAD, true)?; + + let ChatMessageParseOutcome::Recognized(parsed) = outcome else { + bail!("expected Recognized for partial Qwen XML on a Qwen-3.5 model; got Unrecognized"); + }; + assert!(parsed.tool_calls.is_empty() || parsed.tool_calls.len() == 1); + + Ok(()) +} + +#[test] +fn qwen35_parses_multiple_tool_calls() -> Result<()> { + let (_backend, model) = load_qwen35()?; + + let outcome = model.parse_chat_message(TOOLS_JSON, TWO_QWEN_XML_PAYLOADS, false)?; + + let ChatMessageParseOutcome::Recognized(parsed) = outcome else { + bail!( + "expected Recognized for two Qwen XML payloads on a Qwen-3.5 model; got Unrecognized" + ); + }; + assert!( + !parsed.tool_calls.is_empty(), + "expected at least one tool call; got {:?}", + parsed.tool_calls + ); + + Ok(()) +} diff --git a/llama-cpp-bindings-tests/tests/qwen35_recognizes_empty_tool_calls_when_input_is_plain_content_with_tools_requested.rs b/llama-cpp-bindings-tests/tests/qwen35_recognizes_empty_tool_calls_when_input_is_plain_content_with_tools_requested.rs new file mode 100644 index 00000000..b4ea9692 --- /dev/null +++ b/llama-cpp-bindings-tests/tests/qwen35_recognizes_empty_tool_calls_when_input_is_plain_content_with_tools_requested.rs @@ -0,0 +1,57 @@ +use anyhow::Result; +use anyhow::bail; +use llama_cpp_bindings::ChatMessageParseOutcome; +use llama_cpp_bindings::llama_backend::LlamaBackend; +use llama_cpp_bindings::model::LlamaModel; +use llama_cpp_bindings_tests::gpu_backend::inference_model_params; +use llama_cpp_bindings_tests::gpu_backend::require_compiled_backends_present; +use llama_cpp_bindings_tests::test_model::download_file_from; + +const QWEN35_REPO: &str = "unsloth/Qwen3.5-0.8B-GGUF"; +const QWEN35_FILE: &str = "Qwen3.5-0.8B-Q4_K_M.gguf"; + +const TOOLS_JSON: &str = r#"[ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather for a location", + "parameters": { + "type": "object", + "properties": { + "location": {"type": "string", "description": "The city name"} + }, + "required": ["location"] + } + } + } +]"#; + +const PLAIN_CONTENT: &str = "Sorry, I cannot help with that."; + +#[test] +fn qwen35_recognizes_empty_tool_calls_when_input_is_plain_content_with_tools_requested() +-> Result<()> { + let backend = LlamaBackend::init()?; + require_compiled_backends_present()?; + + let path = download_file_from(QWEN35_REPO, QWEN35_FILE)?; + let params = inference_model_params(); + let model = LlamaModel::load_from_file(&backend, &path, ¶ms)?; + + let outcome = model.parse_chat_message(TOOLS_JSON, PLAIN_CONTENT, false)?; + + let ChatMessageParseOutcome::Recognized(parsed) = outcome else { + bail!( + "Qwen 3.5 with tools requested + plain content must produce Recognized (with empty \ + tool_calls); got Unrecognized" + ); + }; + assert!( + parsed.tool_calls.is_empty(), + "expected no tool calls; got {:?}", + parsed.tool_calls + ); + + Ok(()) +} diff --git a/llama-cpp-bindings-tests/tests/qwen36_chat_inference_emits_reasoning_when_template_auto_opens.rs b/llama-cpp-bindings-tests/tests/qwen36_chat_inference_emits_reasoning_when_template_auto_opens.rs new file mode 100644 index 00000000..f402f0be --- /dev/null +++ b/llama-cpp-bindings-tests/tests/qwen36_chat_inference_emits_reasoning_when_template_auto_opens.rs @@ -0,0 +1,112 @@ +use anyhow::Result; +use anyhow::bail; +use llama_cpp_bindings::ChatMessageParseOutcome; +use llama_cpp_bindings::context::LlamaContext; +use llama_cpp_bindings::context::params::LlamaContextParams; +use llama_cpp_bindings::llama_backend::LlamaBackend; +use llama_cpp_bindings::llama_batch::LlamaBatch; +use llama_cpp_bindings::model::AddBos; +use llama_cpp_bindings::model::LlamaChatMessage; +use llama_cpp_bindings::model::LlamaModel; +use llama_cpp_bindings::sampling::LlamaSampler; +use llama_cpp_bindings_tests::classify_sample_loop::ClassifySampleLoop; +use llama_cpp_bindings_tests::gpu_backend::inference_model_params; +use llama_cpp_bindings_tests::gpu_backend::require_compiled_backends_present; +use llama_cpp_bindings_tests::test_model::download_file_from; + +const QWEN36_REPO: &str = "unsloth/Qwen3.6-35B-A3B-GGUF"; +const QWEN36_FILE: &str = "Qwen3.6-35B-A3B-UD-Q4_K_M.gguf"; + +#[test] +fn qwen36_chat_inference_emits_reasoning_when_template_auto_opens() -> Result<()> { + let backend = LlamaBackend::init()?; + require_compiled_backends_present()?; + + let path = download_file_from(QWEN36_REPO, QWEN36_FILE)?; + let params = inference_model_params(); + let model = LlamaModel::load_from_file(&backend, &path, ¶ms)?; + + let context_params = LlamaContextParams::default(); + let mut context = LlamaContext::from_model(&model, &backend, context_params)?; + + let chat_template = model.chat_template(None)?; + let messages = vec![LlamaChatMessage::new( + "user".to_owned(), + "Hello! How are you?".to_owned(), + )?]; + let prompt = model.apply_chat_template(&chat_template, &messages, true)?; + + let mut classifier = model.sampled_token_classifier(); + let tokens = model.str_to_token(&prompt, AddBos::Always)?; + let prompt_token_count = u64::try_from(tokens.len())?; + + let mut batch = LlamaBatch::new(512, 1)?; + classifier.feed_prompt_sequence_to_batch(&mut batch, &tokens, 0, false)?; + + context.decode(&mut batch)?; + + let promoted = classifier.commit_prompt_tokens(); + assert_eq!(promoted, prompt_token_count); + + let mut sampler = LlamaSampler::greedy(); + let initial_position = batch.n_tokens(); + let outcome = ClassifySampleLoop { + model: &model, + classifier: &mut classifier, + sampler: &mut sampler, + context: &mut context, + batch: &mut batch, + initial_position, + max_generated_tokens: 1024, + } + .run()?; + + assert!( + !outcome.generated_raw.is_empty(), + "Qwen3.6 must generate at least one token" + ); + assert!( + outcome.observed_reasoning > 0, + "Qwen3.6 chat template auto-opens reasoning, so the classifier must emit at \ + least one Reasoning token; outcome={outcome:?}" + ); + assert!( + outcome.observed_content > 0, + "Qwen3.6 must emit at least one Content token after ; outcome={outcome:?}" + ); + assert_eq!( + outcome.observed_undeterminable, 0, + "Qwen3.6 chat template auto-opens reasoning, so the classifier must never emit \ + Undeterminable; outcome={outcome:?}" + ); + assert_eq!( + outcome.observed_tool_call, 0, + "chat without tool definitions must not produce ToolCall tokens; outcome={outcome:?}" + ); + + let parse_outcome = model.parse_chat_message("[]", &outcome.generated_raw, false)?; + let ChatMessageParseOutcome::Recognized(parsed) = parse_outcome else { + bail!("Qwen3.6 chat template must be recognised by the parser; got Unrecognized"); + }; + assert!( + !parsed.content.is_empty(), + "parser must see post- content in generated text; generated={:?}", + outcome.generated_raw + ); + + let usage = classifier.into_usage(); + assert_eq!( + usage.prompt_tokens, prompt_token_count, + "prompt_tokens must equal the tokenizer's prompt length" + ); + assert_eq!( + usage.reasoning_tokens, outcome.observed_reasoning, + "reasoning_tokens must equal observed Reasoning variants" + ); + assert_eq!( + usage.undeterminable_tokens, 0, + "Qwen3.6 with auto-opening reasoning must never produce Undeterminable" + ); + + Ok(()) +} diff --git a/llama-cpp-bindings-tests/tests/qwen36_classifier_does_not_emit_reasoning_for_thinking_disabled_prompt.rs b/llama-cpp-bindings-tests/tests/qwen36_classifier_does_not_emit_reasoning_for_thinking_disabled_prompt.rs new file mode 100644 index 00000000..aee03a2a --- /dev/null +++ b/llama-cpp-bindings-tests/tests/qwen36_classifier_does_not_emit_reasoning_for_thinking_disabled_prompt.rs @@ -0,0 +1,128 @@ +use std::num::NonZeroU32; + +use anyhow::Result; +use llama_cpp_bindings::context::LlamaContext; +use llama_cpp_bindings::context::params::LlamaContextParams; +use llama_cpp_bindings::llama_backend::LlamaBackend; +use llama_cpp_bindings::llama_batch::LlamaBatch; +use llama_cpp_bindings::model::AddBos; +use llama_cpp_bindings::model::LlamaModel; +use llama_cpp_bindings::sampling::LlamaSampler; +use llama_cpp_bindings_tests::classify_sample_loop::ClassifySampleLoop; +use llama_cpp_bindings_tests::gpu_backend::inference_model_params; +use llama_cpp_bindings_tests::gpu_backend::require_compiled_backends_present; +use llama_cpp_bindings_tests::test_model::download_file_from; + +const QWEN36_REPO: &str = "unsloth/Qwen3.6-35B-A3B-GGUF"; +const QWEN36_FILE: &str = "Qwen3.6-35B-A3B-UD-Q4_K_M.gguf"; + +const MAX_GENERATED_TOKENS: i32 = 200; + +// Mirrors what Qwen3.6's chat template renders when `enable_thinking=false`: +// the assistant header is followed by a closed empty `...` +// block, so generation begins in CONTENT. +const QWEN36_THINKING_DISABLED_PROMPT: &str = "\ +<|im_start|>user +What is 2 + 2?<|im_end|> +<|im_start|>assistant + + + + +"; + +const FORBIDDEN_MARKERS: &[&str] = &["", ""]; + +#[test] +fn qwen36_classifier_does_not_emit_reasoning_for_thinking_disabled_prompt() -> Result<()> { + let backend = LlamaBackend::init()?; + require_compiled_backends_present()?; + + let path = download_file_from(QWEN36_REPO, QWEN36_FILE)?; + let params = inference_model_params(); + let model = LlamaModel::load_from_file(&backend, &path, ¶ms)?; + + let mut classifier = model.sampled_token_classifier(); + let prompt_tokens = model.str_to_token(QWEN36_THINKING_DISABLED_PROMPT, AddBos::Never)?; + let prompt_token_count = u64::try_from(prompt_tokens.len())?; + + let mut batch = LlamaBatch::new(2048, 1)?; + classifier.feed_prompt_sequence_to_batch(&mut batch, &prompt_tokens, 0, false)?; + + let context_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(8192)); + let mut context = LlamaContext::from_model(&model, &backend, context_params)?; + + context.decode(&mut batch)?; + + let promoted = classifier.commit_prompt_tokens(); + assert_eq!(promoted, prompt_token_count); + + let mut sampler = LlamaSampler::chain_simple([ + LlamaSampler::penalties(64, 1.1, 0.0, 0.0), + LlamaSampler::top_k(40), + LlamaSampler::top_p(0.9, 1), + LlamaSampler::min_p(0.05, 1), + LlamaSampler::temp(0.7), + LlamaSampler::dist(0x00C0_FFEE), + ]); + let initial_position = batch.n_tokens(); + let outcome = ClassifySampleLoop { + model: &model, + classifier: &mut classifier, + sampler: &mut sampler, + context: &mut context, + batch: &mut batch, + initial_position, + max_generated_tokens: MAX_GENERATED_TOKENS, + } + .run()?; + + let usage = classifier.usage(); + + assert!( + !outcome.generated_raw.is_empty(), + "Qwen3.6 must generate at least one token" + ); + assert_eq!( + outcome.observed_reasoning, 0, + "Qwen3.6 thinking-disabled: classifier must not emit any Reasoning token \ + when the prompt closes the think block before generation begins; \ + generated={:?}", + outcome.generated_raw + ); + assert_eq!( + outcome.observed_undeterminable, 0, + "Qwen3.6 thinking-disabled: prompt-token replay must move section to Content \ + before generation, so no Undeterminable tokens may be emitted; \ + generated={:?}", + outcome.generated_raw + ); + assert_eq!( + usage.reasoning_tokens, 0, + "Qwen3.6 thinking-disabled: usage.reasoning_tokens must be zero; usage={usage:?}" + ); + assert_eq!( + usage.undeterminable_tokens, 0, + "Qwen3.6 thinking-disabled: usage.undeterminable_tokens must be zero; usage={usage:?}" + ); + assert!( + outcome.observed_content > 0, + "Qwen3.6 thinking-disabled: classifier must emit at least one Content token" + ); + assert_eq!( + usage.completion_tokens(), + outcome.observed_content, + "Qwen3.6 thinking-disabled: completion tokens must equal observed Content tokens" + ); + + for forbidden in FORBIDDEN_MARKERS { + assert!( + !outcome.content_stream.contains(forbidden), + "Qwen3.6 thinking-disabled: content_stream leaked marker {forbidden:?}; \ + content_stream={:?}", + outcome.content_stream + ); + } + + Ok(()) +} diff --git a/llama-cpp-bindings-tests/tests/qwen36_classifier_emits_reasoning.rs b/llama-cpp-bindings-tests/tests/qwen36_classifier_emits_reasoning.rs new file mode 100644 index 00000000..19596fa6 --- /dev/null +++ b/llama-cpp-bindings-tests/tests/qwen36_classifier_emits_reasoning.rs @@ -0,0 +1,151 @@ +use std::num::NonZeroU32; + +use anyhow::Result; +use anyhow::bail; +use llama_cpp_bindings::ChatMessageParseOutcome; +use llama_cpp_bindings::context::LlamaContext; +use llama_cpp_bindings::context::params::LlamaContextParams; +use llama_cpp_bindings::llama_backend::LlamaBackend; +use llama_cpp_bindings::llama_batch::LlamaBatch; +use llama_cpp_bindings::model::AddBos; +use llama_cpp_bindings::model::LlamaModel; +use llama_cpp_bindings::sampling::LlamaSampler; +use llama_cpp_bindings_tests::classify_sample_loop::ClassifySampleLoop; +use llama_cpp_bindings_tests::gpu_backend::inference_model_params; +use llama_cpp_bindings_tests::gpu_backend::require_compiled_backends_present; +use llama_cpp_bindings_tests::test_model::download_file_from; + +const QWEN36_REPO: &str = "unsloth/Qwen3.6-35B-A3B-GGUF"; +const QWEN36_FILE: &str = "Qwen3.6-35B-A3B-UD-Q4_K_M.gguf"; + +const MAX_GENERATED_TOKENS: i32 = 1500; + +// Qwen3.6's chat template injects `\n` directly into the generation prompt +// when `enable_thinking=true` (the default). The legacy `llama_chat_apply_template` +// path bypasses that jinja branch, so we craft the prompt manually to faithfully +// reproduce the production case where the model resumes generation already inside +// the reasoning block. +const QWEN36_THINKING_PROMPT: &str = "\ +<|im_start|>user +What is 2 + 2?<|im_end|> +<|im_start|>assistant + +"; + +const FORBIDDEN_MARKERS: &[&str] = &["", ""]; + +#[test] +fn qwen36_classifier_emits_reasoning_for_thinking_enabled_prompt() -> Result<()> { + let backend = LlamaBackend::init()?; + require_compiled_backends_present()?; + + let path = download_file_from(QWEN36_REPO, QWEN36_FILE)?; + let params = inference_model_params(); + let model = LlamaModel::load_from_file(&backend, &path, ¶ms)?; + + let mut classifier = model.sampled_token_classifier(); + let prompt_tokens = model.str_to_token(QWEN36_THINKING_PROMPT, AddBos::Never)?; + let prompt_token_count = u64::try_from(prompt_tokens.len())?; + + let mut batch = LlamaBatch::new(2048, 1)?; + classifier.feed_prompt_sequence_to_batch(&mut batch, &prompt_tokens, 0, false)?; + + let context_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(8192)); + let mut context = LlamaContext::from_model(&model, &backend, context_params)?; + + context.decode(&mut batch)?; + + let promoted = classifier.commit_prompt_tokens(); + assert_eq!(promoted, prompt_token_count); + + let mut sampler = LlamaSampler::chain_simple([ + LlamaSampler::penalties(64, 1.1, 0.0, 0.0), + LlamaSampler::top_k(40), + LlamaSampler::top_p(0.9, 1), + LlamaSampler::min_p(0.05, 1), + LlamaSampler::temp(0.7), + LlamaSampler::dist(0x00C0_FFEE), + ]); + let initial_position = batch.n_tokens(); + let outcome = ClassifySampleLoop { + model: &model, + classifier: &mut classifier, + sampler: &mut sampler, + context: &mut context, + batch: &mut batch, + initial_position, + max_generated_tokens: MAX_GENERATED_TOKENS, + } + .run()?; + + let usage = classifier.usage(); + let parse_outcome = model.parse_chat_message("[]", &outcome.generated_raw, true)?; + let ChatMessageParseOutcome::Recognized(parsed) = parse_outcome else { + bail!("Qwen3.6 chat template must be recognised by the parser; got Unrecognized"); + }; + + assert!( + !outcome.generated_raw.is_empty(), + "Qwen3.6 must generate at least one token" + ); + assert!( + outcome.observed_reasoning > 0, + "Qwen3.6: classifier must emit at least one Reasoning token when the prompt \ + opens a block; outcome={outcome:?}", + ); + assert!( + usage.reasoning_tokens > 0, + "Qwen3.6: usage.reasoning_tokens must be non-zero when the prompt opens a \ + block; usage was {usage:?}" + ); + assert_eq!( + outcome.observed_undeterminable, 0, + "Qwen3.6: prompt-token replay must move section to Reasoning before generation, \ + so no Undeterminable tokens may be emitted; outcome={outcome:?}" + ); + assert_eq!( + usage.undeterminable_tokens, 0, + "Qwen3.6: usage.undeterminable_tokens must be zero; usage={usage:?}" + ); + assert_eq!( + usage.completion_tokens(), + outcome.observed_content + outcome.observed_reasoning, + "Qwen3.6: completion tokens must equal observed Content + Reasoning" + ); + + if parsed.reasoning_content.is_empty() { + eprintln!( + "Qwen3.6 parser returned empty reasoning_content (likely a partial parse \ + over `<|im_end|>`-truncated output) — relying on the FORBIDDEN_MARKERS \ + substring check below for leak detection." + ); + } else { + assert_eq!( + outcome.reasoning_stream, parsed.reasoning_content, + "Qwen3.6: per-token reasoning stream must equal parser-side reasoning_content \ + (any difference means a marker leaked into the user-visible stream)", + ); + assert_eq!( + outcome.content_stream, parsed.content, + "Qwen3.6: per-token content stream must equal parser-side content \ + (any difference means a marker leaked into the user-visible stream)", + ); + } + + for forbidden in FORBIDDEN_MARKERS { + assert!( + !outcome.reasoning_stream.contains(forbidden), + "Qwen3.6: reasoning_stream leaked marker {forbidden:?}; \ + reasoning_stream={:?}", + outcome.reasoning_stream + ); + assert!( + !outcome.content_stream.contains(forbidden), + "Qwen3.6: content_stream leaked marker {forbidden:?}; \ + content_stream={:?}", + outcome.content_stream + ); + } + + Ok(()) +} diff --git a/llama-cpp-bindings-tests/tests/qwen36_classifier_emits_reasoning_for_multimodal_thinking_prompt.rs b/llama-cpp-bindings-tests/tests/qwen36_classifier_emits_reasoning_for_multimodal_thinking_prompt.rs new file mode 100644 index 00000000..1d9c1621 --- /dev/null +++ b/llama-cpp-bindings-tests/tests/qwen36_classifier_emits_reasoning_for_multimodal_thinking_prompt.rs @@ -0,0 +1,109 @@ +#![cfg(feature = "multimodal_capable")] + +use std::num::NonZeroU32; + +use anyhow::Result; +use llama_cpp_bindings::context::LlamaContext; +use llama_cpp_bindings::context::params::LlamaContextParams; +use llama_cpp_bindings::llama_backend::LlamaBackend; +use llama_cpp_bindings::llama_batch::LlamaBatch; +use llama_cpp_bindings::model::LlamaModel; +use llama_cpp_bindings::mtmd::MtmdBitmap; +use llama_cpp_bindings::mtmd::MtmdContext; +use llama_cpp_bindings::mtmd::MtmdContextParams; +use llama_cpp_bindings::mtmd::MtmdInputText; +use llama_cpp_bindings::mtmd::mtmd_default_marker; +use llama_cpp_bindings::sampling::LlamaSampler; +use llama_cpp_bindings_tests::classify_sample_loop::ClassifySampleLoop; +use llama_cpp_bindings_tests::gpu_backend::inference_model_params; +use llama_cpp_bindings_tests::gpu_backend::require_compiled_backends_present; +use llama_cpp_bindings_tests::test_model::download_file_from; +use llama_cpp_bindings_tests::test_model::fixtures_dir; + +const QWEN36_REPO: &str = "unsloth/Qwen3.6-35B-A3B-GGUF"; +const QWEN36_FILE: &str = "Qwen3.6-35B-A3B-UD-Q4_K_M.gguf"; +const QWEN36_MMPROJ_FILE: &str = "mmproj-F16.gguf"; + +const MAX_GENERATED_TOKENS: i32 = 200; + +#[test] +fn qwen36_classifier_emits_reasoning_for_multimodal_thinking_prompt() -> Result<()> { + let backend = LlamaBackend::init()?; + require_compiled_backends_present()?; + + let model_path = download_file_from(QWEN36_REPO, QWEN36_FILE)?; + let mmproj_path = download_file_from(QWEN36_REPO, QWEN36_MMPROJ_FILE)?; + let params = inference_model_params(); + let model = LlamaModel::load_from_file(&backend, &model_path, ¶ms)?; + + let mtmd_params = MtmdContextParams::default(); + let mmproj_str = mmproj_path + .to_str() + .ok_or_else(|| anyhow::anyhow!("mmproj path is not valid UTF-8"))?; + let mtmd_ctx = MtmdContext::init_from_file(mmproj_str, &model, &mtmd_params)?; + + let context_params = LlamaContextParams::default() + .with_n_ctx(NonZeroU32::new(8192)) + .with_n_batch(512); + let mut context = LlamaContext::from_model(&model, &backend, context_params)?; + + let image_path = fixtures_dir().join("llamas.jpg"); + let image_path_str = image_path + .to_str() + .ok_or_else(|| anyhow::anyhow!("image path is not valid UTF-8"))?; + let bitmap = MtmdBitmap::from_file(&mtmd_ctx, image_path_str)?; + + let marker = mtmd_default_marker(); + let prompt = format!( + "<|im_start|>user\n{marker}What animals do you see in this image?<|im_end|>\n<|im_start|>assistant\n\n" + ); + + let input_text = MtmdInputText { + text: prompt, + add_special: false, + parse_special: true, + }; + + let chunks = mtmd_ctx.tokenize(input_text, &[&bitmap])?; + + let mut classifier = model.sampled_token_classifier(); + let n_past = + classifier.eval_multimodal_chunks(&chunks, &mtmd_ctx, &context, 0, 0, 512, true)?; + + let mut sampler = LlamaSampler::chain_simple([ + LlamaSampler::penalties(64, 1.1, 0.0, 0.0), + LlamaSampler::top_k(40), + LlamaSampler::top_p(0.9, 1), + LlamaSampler::min_p(0.05, 1), + LlamaSampler::temp(0.7), + LlamaSampler::dist(0x00C0_FFEE), + ]); + + let mut batch = LlamaBatch::new(2048, 1)?; + let outcome = ClassifySampleLoop { + model: &model, + classifier: &mut classifier, + sampler: &mut sampler, + context: &mut context, + batch: &mut batch, + initial_position: n_past, + max_generated_tokens: MAX_GENERATED_TOKENS, + } + .run()?; + + let usage = classifier.usage(); + + if outcome.observed_reasoning == 0 { + anyhow::bail!( + "Qwen 3.6 multimodal + thinking: classifier must emit at least one Reasoning token \ + when the prompt opens a `` block; outcome={outcome:?}" + ); + } + if usage.reasoning_tokens == 0 { + anyhow::bail!( + "Qwen 3.6 multimodal + thinking: usage.reasoning_tokens must be non-zero; usage={usage:?}" + ); + } + + Ok(()) +} diff --git a/llama-cpp-bindings-tests/tests/reranker.rs b/llama-cpp-bindings-tests/tests/reranker.rs index 79a2332e..08f0de6a 100644 --- a/llama-cpp-bindings-tests/tests/reranker.rs +++ b/llama-cpp-bindings-tests/tests/reranker.rs @@ -1,11 +1,12 @@ use std::time::Duration; use anyhow::{Context, Result, bail}; +use llama_cpp_bindings::context::LlamaContext; use llama_cpp_bindings::context::params::LlamaContextParams; use llama_cpp_bindings::ggml_time_us; use llama_cpp_bindings::llama_batch::LlamaBatch; use llama_cpp_bindings::model::AddBos; -use llama_cpp_bindings_tests::TestFixture; +use llama_cpp_bindings_tests::FixtureSession; fn normalize(input: &[f32]) -> Vec { let magnitude = input @@ -26,7 +27,7 @@ fn cosine_similarity(vec_a: &[f32], vec_b: &[f32]) -> f32 { #[test] fn reranking_produces_scores() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.embedding_model()?; @@ -42,8 +43,7 @@ fn reranking_produces_scores() -> Result<()> { .with_n_threads_batch(std::thread::available_parallelism()?.get().try_into()?) .with_n_seq_max(u32::try_from(document_count)?) .with_embeddings(true); - let mut ctx = model - .new_context(backend, ctx_params) + let mut ctx = LlamaContext::from_model(model, backend, ctx_params) .with_context(|| "unable to create context")?; let prompt_lines: Vec = documents @@ -63,7 +63,7 @@ fn reranking_produces_scores() -> Result<()> { bail!("one of the provided prompts exceeds the size of the context window"); } - let mut classifier = model.reasoning_token_classifier()?; + let mut classifier = model.sampled_token_classifier(); let mut batch = LlamaBatch::new(2048, i32::try_from(document_count)?)?; let t_main_start = ggml_time_us(); @@ -80,7 +80,7 @@ fn reranking_produces_scores() -> Result<()> { let total_token_count = u64::try_from(total_tokens)?; assert_eq!(classifier.pending_prompt_tokens(), total_token_count); - assert_eq!(classifier.usage().prompt_tokens(), 0); + assert_eq!(classifier.usage().prompt_tokens, 0); ctx.clear_kv_cache(); ctx.decode(&mut batch) @@ -101,7 +101,10 @@ fn reranking_produces_scores() -> Result<()> { let t_main_end = ggml_time_us(); let duration = Duration::from_micros(u64::try_from(t_main_end - t_main_start)?); - #[allow(clippy::cast_precision_loss)] + #[expect( + clippy::cast_precision_loss, + reason = "logged throughput tolerates f32 precision" + )] let tokens_per_second = total_tokens as f32 / duration.as_secs_f32(); eprintln!( @@ -131,7 +134,7 @@ fn reranking_produces_scores() -> Result<()> { ); let usage = classifier.into_usage(); - assert_eq!(usage.prompt_tokens(), total_token_count); + assert_eq!(usage.prompt_tokens, total_token_count); assert_eq!(usage.completion_tokens(), 0); Ok(()) diff --git a/llama-cpp-bindings-tests/tests/sampled_token_classifier_markers.rs b/llama-cpp-bindings-tests/tests/sampled_token_classifier_markers.rs new file mode 100644 index 00000000..ee747c61 --- /dev/null +++ b/llama-cpp-bindings-tests/tests/sampled_token_classifier_markers.rs @@ -0,0 +1,159 @@ +use anyhow::Result; +use llama_cpp_bindings::SampledToken; +use llama_cpp_bindings::llama_batch::LlamaBatch; +use llama_cpp_bindings::sampled_token_classifier::SampledTokenClassifier; +use llama_cpp_bindings::sampled_token_classifier::SampledTokenSection; +use llama_cpp_bindings::sampled_token_classifier::StreamingMarkers; +use llama_cpp_bindings_tests::FixtureSession; + +#[test] +fn classifier_starts_in_pending_section_for_default_fixture() { + let fixture = FixtureSession::open().expect("open fixture"); + let model = fixture.default_model(); + + let classifier = model.sampled_token_classifier(); + + assert_eq!(classifier.current_section(), SampledTokenSection::Pending); +} + +#[test] +fn classifier_construction_is_idempotent_across_calls() { + let fixture = FixtureSession::open().expect("open fixture"); + let model = fixture.default_model(); + + let first = model.sampled_token_classifier(); + let second = model.sampled_token_classifier(); + + assert_eq!(first.current_section(), second.current_section()); + assert_eq!(first.usage(), second.usage()); +} + +#[test] +fn diagnose_tool_call_synthetic_renders_runs_without_panic() -> Result<()> { + let fixture = FixtureSession::open()?; + let model = fixture.default_model(); + + let _ = model.diagnose_tool_call_synthetic_renders()?; + + Ok(()) +} + +#[test] +fn ingest_with_no_markers_emits_undeterminable_with_visible_and_raw_piece() { + let fixture = FixtureSession::open().expect("open fixture"); + let model = fixture.default_model(); + + let mut classifier = SampledTokenClassifier::new(model, StreamingMarkers::default()); + + let outcomes = classifier.ingest(model.token_bos()); + + assert_eq!(outcomes.len(), 1); + let outcome = &outcomes[0]; + assert!(matches!( + outcome.sampled_token, + SampledToken::Undeterminable(_) + )); + assert_eq!(outcome.visible_piece, outcome.raw_piece); + assert_eq!(classifier.usage().undeterminable_tokens, 1); +} + +#[test] +fn ingest_with_no_markers_decodes_each_token_independently() { + let fixture = FixtureSession::open().expect("open fixture"); + let model = fixture.default_model(); + + let mut classifier = SampledTokenClassifier::new(model, StreamingMarkers::default()); + + let _ = classifier.ingest(model.token_bos()); + let _ = classifier.ingest(model.token_eos()); + + assert_eq!(classifier.usage().undeterminable_tokens, 2); +} + +#[test] +fn ingest_prompt_token_with_no_markers_is_a_noop() { + let fixture = FixtureSession::open().expect("open fixture"); + let model = fixture.default_model(); + + let mut classifier = SampledTokenClassifier::new(model, StreamingMarkers::default()); + let usage_before = *classifier.usage(); + + classifier.ingest_prompt_token(model.token_bos()); + classifier.ingest_prompt_tokens(&[model.token_eos(), model.token_nl()]); + + assert_eq!(*classifier.usage(), usage_before); + assert_eq!(classifier.current_section(), SampledTokenSection::Pending); +} + +#[test] +fn feed_prompt_to_batch_increments_pending_prompt_tokens() -> Result<()> { + let fixture = FixtureSession::open()?; + let model = fixture.default_model(); + + let mut classifier = SampledTokenClassifier::new(model, StreamingMarkers::default()); + let mut batch = LlamaBatch::new(8, 1)?; + + classifier.feed_prompt_to_batch(&mut batch, model.token_bos(), 0, &[0], false)?; + classifier.feed_prompt_to_batch(&mut batch, model.token_eos(), 1, &[0], false)?; + + assert_eq!(classifier.pending_prompt_tokens(), 2); + assert_eq!(batch.n_tokens(), 2); + + Ok(()) +} + +#[test] +fn feed_prompt_sequence_to_batch_stages_all_tokens() -> Result<()> { + let fixture = FixtureSession::open()?; + let model = fixture.default_model(); + + let mut classifier = SampledTokenClassifier::new(model, StreamingMarkers::default()); + let mut batch = LlamaBatch::new(8, 1)?; + + let tokens = vec![model.token_bos(), model.token_eos(), model.token_nl()]; + classifier.feed_prompt_sequence_to_batch(&mut batch, &tokens, 0, false)?; + + assert_eq!(classifier.pending_prompt_tokens(), 3); + assert_eq!(batch.n_tokens(), 3); + + Ok(()) +} + +#[test] +fn commit_prompt_tokens_promotes_pending_count_to_usage_and_clears() -> Result<()> { + let fixture = FixtureSession::open()?; + let model = fixture.default_model(); + + let mut classifier = SampledTokenClassifier::new(model, StreamingMarkers::default()); + let mut batch = LlamaBatch::new(8, 1)?; + + classifier.feed_prompt_to_batch(&mut batch, model.token_bos(), 0, &[0], false)?; + classifier.feed_prompt_to_batch(&mut batch, model.token_eos(), 1, &[0], false)?; + + let promoted = classifier.commit_prompt_tokens(); + + assert_eq!(promoted, 2); + assert_eq!(classifier.pending_prompt_tokens(), 0); + assert_eq!(classifier.usage().prompt_tokens, 2); + + Ok(()) +} + +#[test] +fn discard_pending_prompt_tokens_clears_count_without_recording_usage() -> Result<()> { + let fixture = FixtureSession::open()?; + let model = fixture.default_model(); + + let mut classifier = SampledTokenClassifier::new(model, StreamingMarkers::default()); + let mut batch = LlamaBatch::new(8, 1)?; + + classifier.feed_prompt_to_batch(&mut batch, model.token_bos(), 0, &[0], false)?; + + let discarded = classifier.discard_pending_prompt_tokens(); + + assert_eq!(discarded, 1); + assert_eq!(classifier.pending_prompt_tokens(), 0); + assert_eq!(classifier.usage().prompt_tokens, 0); + + Ok(()) +} diff --git a/llama-cpp-bindings-tests/tests/sampling.rs b/llama-cpp-bindings-tests/tests/sampling.rs index 3b906f4c..8033ccfc 100644 --- a/llama-cpp-bindings-tests/tests/sampling.rs +++ b/llama-cpp-bindings-tests/tests/sampling.rs @@ -2,17 +2,19 @@ use std::num::NonZeroU32; use anyhow::Result; use llama_cpp_bindings::GrammarError; +use llama_cpp_bindings::context::LlamaContext; use llama_cpp_bindings::context::params::LlamaContextParams; use llama_cpp_bindings::llama_batch::LlamaBatch; use llama_cpp_bindings::model::AddBos; use llama_cpp_bindings::sampling::LlamaSampler; -use llama_cpp_bindings_tests::TestFixture; +use llama_cpp_bindings::token::LlamaToken; +use llama_cpp_bindings_tests::FixtureSession; use serial_test::serial; #[test] #[serial] fn dry_sampler_with_model() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let model = fixture.default_model(); let breakers: Vec<&[u8]> = vec![b"\n", b"\t"]; let _sampler = LlamaSampler::dry(model, 1.5, 2.0, 128, 2, &breakers); @@ -23,7 +25,7 @@ fn dry_sampler_with_model() -> Result<()> { #[test] #[serial] fn dry_sampler_with_null_byte_in_seq_breakers_returns_error() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let model = fixture.default_model(); let breakers: Vec<&[u8]> = vec![b"hello\0world"]; let result = LlamaSampler::dry(model, 1.5, 2.0, 128, 2, breakers); @@ -36,7 +38,7 @@ fn dry_sampler_with_null_byte_in_seq_breakers_returns_error() -> Result<()> { #[test] #[serial] fn grammar_returns_sampler_for_valid_grammar() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let model = fixture.default_model(); let sampler = LlamaSampler::grammar(model, "root ::= \"hello\"", "root"); @@ -48,7 +50,7 @@ fn grammar_returns_sampler_for_valid_grammar() -> Result<()> { #[test] #[serial] fn grammar_lazy_returns_sampler_for_valid_grammar_with_triggers() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let model = fixture.default_model(); let trigger_words: Vec<&[u8]> = vec![b"function"]; let sampler = @@ -62,7 +64,7 @@ fn grammar_lazy_returns_sampler_for_valid_grammar_with_triggers() -> Result<()> #[test] #[serial] fn grammar_lazy_patterns_returns_sampler_for_valid_grammar_with_patterns() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let model = fixture.default_model(); let patterns = vec!["\\{.*".to_string()]; let sampler = @@ -76,7 +78,7 @@ fn grammar_lazy_patterns_returns_sampler_for_valid_grammar_with_patterns() -> Re #[test] #[serial] fn grammar_lazy_with_root_not_found_returns_error() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let model = fixture.default_model(); let trigger_words: Vec<&[u8]> = vec![b"function"]; let result = @@ -90,7 +92,7 @@ fn grammar_lazy_with_root_not_found_returns_error() -> Result<()> { #[test] #[serial] fn grammar_lazy_with_null_byte_in_trigger_word_returns_error() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let model = fixture.default_model(); let trigger_words: Vec<&[u8]> = vec![b"hel\0lo"]; let result = @@ -104,7 +106,7 @@ fn grammar_lazy_with_null_byte_in_trigger_word_returns_error() -> Result<()> { #[test] #[serial] fn grammar_lazy_patterns_with_root_not_found_returns_error() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let model = fixture.default_model(); let patterns = vec!["\\{.*".to_string()]; let result = @@ -118,7 +120,7 @@ fn grammar_lazy_patterns_with_root_not_found_returns_error() -> Result<()> { #[test] #[serial] fn grammar_lazy_patterns_with_null_byte_in_pattern_returns_error() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let model = fixture.default_model(); let patterns = vec!["hel\0lo".to_string()]; let result = @@ -132,7 +134,7 @@ fn grammar_lazy_patterns_with_null_byte_in_pattern_returns_error() -> Result<()> #[test] #[serial] fn llguidance_method_creates_sampler() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let model = fixture.default_model(); let result = LlamaSampler::llguidance(model, "regex", r"yes|no"); @@ -152,7 +154,7 @@ fn logit_bias_with_empty_biases_succeeds() { #[test] #[serial] fn dry_sampler_with_root_not_found_grammar_does_not_apply() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let model = fixture.default_model(); let breakers: Vec<&[u8]> = vec![b"\n"]; let _sampler = LlamaSampler::dry(model, 1.5, 2.0, 128, 2, &breakers); @@ -160,14 +162,89 @@ fn dry_sampler_with_root_not_found_grammar_does_not_apply() -> Result<()> { Ok(()) } +#[test] +#[serial] +fn accept_many_iterates_over_borrowed_tokens() -> Result<()> { + let fixture = FixtureSession::open()?; + let model = fixture.default_model(); + let mut sampler = LlamaSampler::chain_simple([LlamaSampler::greedy()]); + let tokens = vec![model.token_bos(), model.token_eos()]; + + sampler.accept_many(&tokens)?; + + Ok(()) +} + +#[test] +#[serial] +fn with_tokens_returns_self_after_accepting_each_token() -> Result<()> { + let fixture = FixtureSession::open()?; + let model = fixture.default_model(); + let sampler = LlamaSampler::chain_simple([LlamaSampler::greedy()]); + let tokens = [model.token_bos(), model.token_eos()]; + + let _consumed = sampler.with_tokens(tokens.iter().copied())?; + + Ok(()) +} + +#[test] +#[serial] +fn accept_consumes_a_single_token() -> Result<()> { + let fixture = FixtureSession::open()?; + let model = fixture.default_model(); + let mut sampler = LlamaSampler::chain_simple([LlamaSampler::greedy()]); + + sampler.accept(model.token_bos())?; + + Ok(()) +} + +#[test] +#[serial] +fn try_accept_returns_ok_for_a_valid_token() -> Result<()> { + let _fixture = FixtureSession::open()?; + let mut sampler = LlamaSampler::chain_simple([LlamaSampler::greedy()]); + + sampler.try_accept(LlamaToken::new(0))?; + + Ok(()) +} + +#[test] +#[serial] +fn apply_runs_sampler_over_token_data_array() -> Result<()> { + use std::num::NonZeroU32; + + use llama_cpp_bindings::context::params::LlamaContextParams; + use llama_cpp_bindings::llama_batch::LlamaBatch; + use llama_cpp_bindings::model::AddBos; + + let fixture = FixtureSession::open()?; + let backend = fixture.backend(); + let model = fixture.default_model(); + let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; + let tokens = model.str_to_token("Hi", AddBos::Always)?; + let mut batch = LlamaBatch::new(512, 1)?; + batch.add_sequence(&tokens, 0, false)?; + context.decode(&mut batch)?; + + let mut data_array = context.token_data_array_ith(batch.n_tokens() - 1)?; + let sampler = LlamaSampler::greedy(); + sampler.apply(&mut data_array); + + Ok(()) +} + #[test] #[serial] fn sample_returns_token_after_decode() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(512)); - let mut context = model.new_context(backend, ctx_params)?; + let mut context = LlamaContext::from_model(model, backend, ctx_params)?; let tokens = model.str_to_token("Hello", AddBos::Always)?; let mut batch = LlamaBatch::new(512, 1)?; batch.add_sequence(&tokens, 0, false)?; diff --git a/llama-cpp-bindings-tests/tests/text_generation.rs b/llama-cpp-bindings-tests/tests/text_generation.rs index f053b701..ad59463b 100644 --- a/llama-cpp-bindings-tests/tests/text_generation.rs +++ b/llama-cpp-bindings-tests/tests/text_generation.rs @@ -1,30 +1,33 @@ use std::io::Write; use std::time::Duration; -use anyhow::{Context, Result}; +use anyhow::Context as _; +use anyhow::Result; +use llama_cpp_bindings::context::LlamaContext; use llama_cpp_bindings::context::params::LlamaContextParams; use llama_cpp_bindings::ggml_time_us; use llama_cpp_bindings::llama_batch::LlamaBatch; -use llama_cpp_bindings::model::{AddBos, LlamaChatMessage}; +use llama_cpp_bindings::model::AddBos; +use llama_cpp_bindings::model::LlamaChatMessage; use llama_cpp_bindings::sampled_token::SampledToken; use llama_cpp_bindings::sampling::LlamaSampler; -use llama_cpp_bindings_tests::TestFixture; +use llama_cpp_bindings_tests::FixtureSession; +use llama_cpp_bindings_tests::classify_sample_loop::ClassifySampleLoop; #[test] fn raw_prompt_completion_with_timing() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let ctx_params = LlamaContextParams::default(); - let mut ctx = model - .new_context(backend, ctx_params) + let mut ctx = LlamaContext::from_model(model, backend, ctx_params) .with_context(|| "unable to create context")?; let prompt = "Hello my name is"; - let n_len: i32 = 64; + let max_generated_tokens: i32 = 64; - let mut classifier = model.reasoning_token_classifier()?; + let mut classifier = model.sampled_token_classifier(); let tokens_list = model .str_to_token(prompt, AddBos::Always) .with_context(|| format!("failed to tokenize {prompt}"))?; @@ -44,87 +47,84 @@ fn raw_prompt_completion_with_timing() -> Result<()> { classifier.feed_prompt_sequence_to_batch(&mut batch, &tokens_list, 0, false)?; assert_eq!(classifier.pending_prompt_tokens(), prompt_token_count); - assert_eq!(classifier.usage().prompt_tokens(), 0); + assert_eq!(classifier.usage().prompt_tokens, 0); ctx.decode(&mut batch) .with_context(|| "llama_decode() failed")?; let promoted = classifier.commit_prompt_tokens(); assert_eq!(promoted, prompt_token_count); - assert_eq!(classifier.usage().prompt_tokens(), prompt_token_count); - - let mut n_cur = batch.n_tokens(); - let mut n_decode: i32 = 0; - let mut observed_content: u64 = 0; - let mut observed_reasoning: u64 = 0; - let t_main_start = ggml_time_us(); + assert_eq!(classifier.usage().prompt_tokens, prompt_token_count); let mut sampler = LlamaSampler::chain_simple([LlamaSampler::dist(1234), LlamaSampler::greedy()]); - - let mut generated = String::new(); - - while n_cur <= n_len { - let token = classifier.sample(&mut sampler, &ctx, batch.n_tokens() - 1)?; - - match token { - SampledToken::Content(_) => observed_content += 1, - SampledToken::Reasoning(_) => observed_reasoning += 1, - SampledToken::Undeterminable(_) => {} - } - - if model.is_eog_token(&token) { - break; - } - - let output_string = model.token_to_piece(&token, &mut decoder, true, None)?; - generated.push_str(&output_string); - print!("{output_string}"); - std::io::stdout().flush()?; - - batch.clear(); - batch.add(&token, n_cur, &[0], true)?; - n_cur += 1; - - ctx.decode(&mut batch).with_context(|| "failed to eval")?; - n_decode += 1; + let initial_position = batch.n_tokens(); + let t_main_start = ggml_time_us(); + let outcome = ClassifySampleLoop { + model, + classifier: &mut classifier, + sampler: &mut sampler, + context: &mut ctx, + batch: &mut batch, + initial_position, + max_generated_tokens, } - + .run()?; let t_main_end = ggml_time_us(); let duration = Duration::from_micros(u64::try_from(t_main_end - t_main_start)?); + let total_observed = + outcome.observed_content + outcome.observed_reasoning + outcome.observed_undeterminable; - #[allow(clippy::cast_precision_loss)] - let tokens_per_second = n_decode as f32 / duration.as_secs_f32(); + #[expect( + clippy::cast_precision_loss, + reason = "logged throughput tolerates f32 precision" + )] + let tokens_per_second = total_observed as f32 / duration.as_secs_f32(); eprintln!( - "\ndecoded {n_decode} tokens in {:.2} s, speed {tokens_per_second:.2} t/s", + "\ndecoded {total_observed} tokens in {:.2} s, speed {tokens_per_second:.2} t/s", duration.as_secs_f32(), ); assert!( - !generated.is_empty(), + !outcome.generated_raw.is_empty(), "model should generate at least one token" ); + assert_eq!( + outcome.observed_tool_call, 0, + "raw prompt without tool-call markers must not produce ToolCall tokens; \ + outcome={outcome:?}" + ); + assert!( + total_observed > 0, + "model must produce at least one classified token; outcome={outcome:?}" + ); let usage = classifier.into_usage(); assert_eq!( - usage.prompt_tokens(), - prompt_token_count, + usage.prompt_tokens, prompt_token_count, "prompt_tokens must equal the tokenizer's prompt length" ); assert_eq!( - usage.content_tokens(), - observed_content, + usage.content_tokens, outcome.observed_content, "content_tokens must equal observed Content variants" ); assert_eq!( - usage.reasoning_tokens(), - observed_reasoning, + usage.reasoning_tokens, outcome.observed_reasoning, "reasoning_tokens must equal observed Reasoning variants" ); + assert_eq!( + usage.undeterminable_tokens, outcome.observed_undeterminable, + "undeterminable_tokens must equal observed Undeterminable variants" + ); + assert_eq!( + usage.tool_call_tokens, outcome.observed_tool_call, + "tool_call_tokens must equal observed ToolCall variants" + ); assert_eq!( usage.completion_tokens(), - observed_content + observed_reasoning + total_observed, + "completion_tokens must equal Content + Reasoning + Undeterminable" ); Ok(()) @@ -132,12 +132,12 @@ fn raw_prompt_completion_with_timing() -> Result<()> { #[test] fn chat_inference_produces_coherent_output() -> Result<()> { - let fixture = TestFixture::shared(); + let fixture = FixtureSession::open()?; let backend = fixture.backend(); let model = fixture.default_model(); let context_params = LlamaContextParams::default(); - let mut context = model.new_context(backend, context_params)?; + let mut context = LlamaContext::from_model(model, backend, context_params)?; let chat_template = model.chat_template(None)?; let messages = vec![LlamaChatMessage::new( @@ -146,7 +146,7 @@ fn chat_inference_produces_coherent_output() -> Result<()> { )?]; let prompt = model.apply_chat_template(&chat_template, &messages, true)?; - let mut classifier = model.reasoning_token_classifier()?; + let mut classifier = model.sampled_token_classifier(); let tokens = model.str_to_token(&prompt, AddBos::Always)?; let prompt_token_count = u64::try_from(tokens.len())?; @@ -154,90 +154,69 @@ fn chat_inference_produces_coherent_output() -> Result<()> { classifier.feed_prompt_sequence_to_batch(&mut batch, &tokens, 0, false)?; assert_eq!(classifier.pending_prompt_tokens(), prompt_token_count); - assert_eq!(classifier.usage().prompt_tokens(), 0); + assert_eq!(classifier.usage().prompt_tokens, 0); context.decode(&mut batch)?; let promoted = classifier.commit_prompt_tokens(); assert_eq!(promoted, prompt_token_count); - let mut decoder = encoding_rs::UTF_8.new_decoder(); let mut sampler = LlamaSampler::greedy(); - let mut position = batch.n_tokens(); - let max_tokens = 1024; - let mut generated = String::new(); - let mut observed_content: u64 = 0; - let mut observed_reasoning: u64 = 0; - - while position <= max_tokens { - let token = classifier.sample(&mut sampler, &context, batch.n_tokens() - 1)?; - - match token { - SampledToken::Content(_) => observed_content += 1, - SampledToken::Reasoning(_) => observed_reasoning += 1, - SampledToken::Undeterminable(_) => { - unreachable!( - "Qwen3 chat template uses detected reasoning markers; classifier must not emit Undeterminable" - ) - } - } - - if model.is_eog_token(&token) { - break; - } - - let piece = model.token_to_piece(&token, &mut decoder, true, None)?; - generated.push_str(&piece); - print!("{piece}"); - std::io::stdout().flush()?; - - batch.clear(); - batch.add(&token, position, &[0], true)?; - position += 1; - - context.decode(&mut batch)?; + let initial_position = batch.n_tokens(); + let outcome = ClassifySampleLoop { + model, + classifier: &mut classifier, + sampler: &mut sampler, + context: &mut context, + batch: &mut batch, + initial_position, + max_generated_tokens: 1024, } + .run()?; println!(); assert!( - !generated.is_empty(), + !outcome.generated_raw.is_empty(), "model should generate at least one token" ); + let total_observed = + outcome.observed_content + outcome.observed_reasoning + outcome.observed_undeterminable; assert!( - observed_reasoning > 0, - "reasoning model should emit at least one Reasoning token" + total_observed > 0, + "model must produce at least one classified token; outcome={outcome:?}" ); - assert!( - observed_content > 0, - "reasoning model should emit at least one Content token after " + assert_eq!( + outcome.observed_tool_call, 0, + "chat without tool definitions must not produce ToolCall tokens; outcome={outcome:?}" ); let usage = classifier.into_usage(); assert_eq!( - usage.prompt_tokens(), - prompt_token_count, + usage.prompt_tokens, prompt_token_count, "prompt_tokens must equal the tokenizer's prompt length" ); assert_eq!( - usage.content_tokens(), - observed_content, + usage.content_tokens, outcome.observed_content, "content_tokens must equal observed Content variants" ); assert_eq!( - usage.reasoning_tokens(), - observed_reasoning, + usage.reasoning_tokens, outcome.observed_reasoning, "reasoning_tokens must equal observed Reasoning variants" ); + assert_eq!( + usage.undeterminable_tokens, outcome.observed_undeterminable, + "undeterminable_tokens must equal observed Undeterminable variants" + ); assert_eq!( usage.completion_tokens(), - observed_content + observed_reasoning + total_observed, + "completion_tokens must equal Content + Reasoning + Undeterminable" ); assert_eq!( - usage.undeterminable_tokens(), - 0, - "model with detected markers should never produce Undeterminable" + usage.tool_call_tokens, outcome.observed_tool_call, + "tool_call_tokens must equal observed ToolCall variants" ); Ok(()) diff --git a/llama-cpp-bindings-types/Cargo.toml b/llama-cpp-bindings-types/Cargo.toml new file mode 100644 index 00000000..6cba8e9f --- /dev/null +++ b/llama-cpp-bindings-types/Cargo.toml @@ -0,0 +1,22 @@ +[package] +name = "llama-cpp-bindings-types" +description = "Shared value types for llama-cpp-bindings, free of llama.cpp/FFI dependencies" +version = "0.5.0" +edition.workspace = true +license = "Apache-2.0" +repository = "https://github.com/intentee/llama-cpp-bindings" + +[dependencies] +serde = { workspace = true } +serde_json = { workspace = true } +thiserror = { workspace = true } + +[lints.rust] +unsafe_op_in_unsafe_fn = "warn" +unused_qualifications = "warn" + +[lints.clippy] +all = { level = "deny", priority = -1 } +pedantic = { level = "warn", priority = -1 } +nursery = { level = "warn", priority = -1 } +module_name_repetitions = "allow" diff --git a/llama-cpp-bindings-types/src/bracketed_json_shape.rs b/llama-cpp-bindings-types/src/bracketed_json_shape.rs new file mode 100644 index 00000000..51b18f4b --- /dev/null +++ b/llama-cpp-bindings-types/src/bracketed_json_shape.rs @@ -0,0 +1,4 @@ +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct BracketedJsonShape { + pub name_args_separator: String, +} diff --git a/llama-cpp-bindings-types/src/json_object_shape.rs b/llama-cpp-bindings-types/src/json_object_shape.rs new file mode 100644 index 00000000..b20a5e20 --- /dev/null +++ b/llama-cpp-bindings-types/src/json_object_shape.rs @@ -0,0 +1,5 @@ +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct JsonObjectShape { + pub name_field: String, + pub arguments_field: String, +} diff --git a/llama-cpp-bindings-types/src/key_value_xml_tags_shape.rs b/llama-cpp-bindings-types/src/key_value_xml_tags_shape.rs new file mode 100644 index 00000000..220e94b6 --- /dev/null +++ b/llama-cpp-bindings-types/src/key_value_xml_tags_shape.rs @@ -0,0 +1,7 @@ +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct KeyValueXmlTagsShape { + pub key_open: String, + pub key_close: String, + pub value_open: String, + pub value_close: String, +} diff --git a/llama-cpp-bindings-types/src/lib.rs b/llama-cpp-bindings-types/src/lib.rs new file mode 100644 index 00000000..f3db5990 --- /dev/null +++ b/llama-cpp-bindings-types/src/lib.rs @@ -0,0 +1,29 @@ +pub mod bracketed_json_shape; +pub mod json_object_shape; +pub mod key_value_xml_tags_shape; +pub mod paired_quote_shape; +pub mod parsed_chat_message; +pub mod parsed_tool_call; +pub mod reasoning_markers; +pub mod token_usage; +pub mod token_usage_error; +pub mod tool_call_args_shape; +pub mod tool_call_arguments; +pub mod tool_call_markers; +pub mod tool_call_value_quote; +pub mod xml_tags_shape; + +pub use bracketed_json_shape::BracketedJsonShape; +pub use json_object_shape::JsonObjectShape; +pub use key_value_xml_tags_shape::KeyValueXmlTagsShape; +pub use paired_quote_shape::PairedQuoteShape; +pub use parsed_chat_message::ParsedChatMessage; +pub use parsed_tool_call::ParsedToolCall; +pub use reasoning_markers::ReasoningMarkers; +pub use token_usage::TokenUsage; +pub use token_usage_error::TokenUsageError; +pub use tool_call_args_shape::ToolCallArgsShape; +pub use tool_call_arguments::ToolCallArguments; +pub use tool_call_markers::ToolCallMarkers; +pub use tool_call_value_quote::ToolCallValueQuote; +pub use xml_tags_shape::XmlTagsShape; diff --git a/llama-cpp-bindings-types/src/paired_quote_shape.rs b/llama-cpp-bindings-types/src/paired_quote_shape.rs new file mode 100644 index 00000000..1126d3ae --- /dev/null +++ b/llama-cpp-bindings-types/src/paired_quote_shape.rs @@ -0,0 +1,7 @@ +use crate::tool_call_value_quote::ToolCallValueQuote; + +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct PairedQuoteShape { + pub name_args_separator: String, + pub value_quote: ToolCallValueQuote, +} diff --git a/llama-cpp-bindings-types/src/parsed_chat_message.rs b/llama-cpp-bindings-types/src/parsed_chat_message.rs new file mode 100644 index 00000000..df7bef7c --- /dev/null +++ b/llama-cpp-bindings-types/src/parsed_chat_message.rs @@ -0,0 +1,91 @@ +use serde::Deserialize; +use serde::Serialize; + +use crate::parsed_tool_call::ParsedToolCall; + +#[derive(Clone, Debug, Default, Deserialize, Eq, PartialEq, Serialize)] +#[serde(deny_unknown_fields)] +pub struct ParsedChatMessage { + pub content: String, + pub reasoning_content: String, + pub tool_calls: Vec, +} + +impl ParsedChatMessage { + #[must_use] + pub const fn new( + content: String, + reasoning_content: String, + tool_calls: Vec, + ) -> Self { + Self { + content, + reasoning_content, + tool_calls, + } + } + + #[must_use] + pub const fn is_empty(&self) -> bool { + self.content.is_empty() && self.reasoning_content.is_empty() && self.tool_calls.is_empty() + } +} + +#[cfg(test)] +mod tests { + use super::ParsedChatMessage; + use super::ParsedToolCall; + use crate::tool_call_arguments::ToolCallArguments; + + #[test] + fn empty_message_reports_empty() { + assert!(ParsedChatMessage::default().is_empty()); + } + + #[test] + fn message_with_content_is_not_empty() { + let parsed = ParsedChatMessage::new("hello".to_owned(), String::new(), Vec::new()); + + assert!(!parsed.is_empty()); + } + + #[test] + fn message_with_reasoning_is_not_empty() { + let parsed = ParsedChatMessage::new(String::new(), "thinking".to_owned(), Vec::new()); + + assert!(!parsed.is_empty()); + } + + #[test] + fn message_with_tool_call_is_not_empty() { + let parsed = ParsedChatMessage::new( + String::new(), + String::new(), + vec![ParsedToolCall::new( + String::new(), + "tool".to_owned(), + ToolCallArguments::default(), + )], + ); + + assert!(!parsed.is_empty()); + } + + #[test] + fn message_with_all_three_fields_populated_is_not_empty() { + let parsed = ParsedChatMessage::new( + "hello".to_owned(), + "thinking".to_owned(), + vec![ParsedToolCall::new( + "id-1".to_owned(), + "tool".to_owned(), + ToolCallArguments::default(), + )], + ); + + assert!(!parsed.is_empty()); + assert_eq!(parsed.content, "hello"); + assert_eq!(parsed.reasoning_content, "thinking"); + assert_eq!(parsed.tool_calls.len(), 1); + } +} diff --git a/llama-cpp-bindings-types/src/parsed_tool_call.rs b/llama-cpp-bindings-types/src/parsed_tool_call.rs new file mode 100644 index 00000000..27f69370 --- /dev/null +++ b/llama-cpp-bindings-types/src/parsed_tool_call.rs @@ -0,0 +1,56 @@ +use serde::Deserialize; +use serde::Serialize; + +use crate::tool_call_arguments::ToolCallArguments; + +#[derive(Clone, Debug, Default, Deserialize, Eq, PartialEq, Serialize)] +#[serde(deny_unknown_fields)] +pub struct ParsedToolCall { + pub id: String, + pub name: String, + pub arguments: ToolCallArguments, +} + +impl ParsedToolCall { + #[must_use] + pub const fn new(id: String, name: String, arguments: ToolCallArguments) -> Self { + Self { + id, + name, + arguments, + } + } +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::ParsedToolCall; + use crate::tool_call_arguments::ToolCallArguments; + + #[test] + fn new_assigns_fields_in_order() { + let parsed = ParsedToolCall::new( + "id-1".to_owned(), + "tool".to_owned(), + ToolCallArguments::ValidJson(json!({})), + ); + + assert_eq!(parsed.id, "id-1"); + assert_eq!(parsed.name, "tool"); + assert_eq!(parsed.arguments, ToolCallArguments::ValidJson(json!({}))); + } + + #[test] + fn default_is_empty_strings_and_invalid_arguments() { + let parsed = ParsedToolCall::default(); + + assert!(parsed.id.is_empty()); + assert!(parsed.name.is_empty()); + assert_eq!( + parsed.arguments, + ToolCallArguments::InvalidJson(String::new()) + ); + } +} diff --git a/llama-cpp-bindings-types/src/reasoning_markers.rs b/llama-cpp-bindings-types/src/reasoning_markers.rs new file mode 100644 index 00000000..02d7586a --- /dev/null +++ b/llama-cpp-bindings-types/src/reasoning_markers.rs @@ -0,0 +1,5 @@ +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct ReasoningMarkers { + pub open: String, + pub close: String, +} diff --git a/llama-cpp-bindings/src/token_usage.rs b/llama-cpp-bindings-types/src/token_usage.rs similarity index 60% rename from llama-cpp-bindings/src/token_usage.rs rename to llama-cpp-bindings-types/src/token_usage.rs index 3502cb27..7bf67448 100644 --- a/llama-cpp-bindings/src/token_usage.rs +++ b/llama-cpp-bindings-types/src/token_usage.rs @@ -2,22 +2,22 @@ use std::iter::Sum; use std::ops::Add; use std::ops::AddAssign; -use crate::TokenUsageError; -use crate::sampled_token::SampledToken; - -#[expect( - clippy::struct_field_names, - reason = "every field counts a kind of token" -)] -#[derive(Copy, Clone, Debug, Default, Eq, PartialEq)] +use serde::Deserialize; +use serde::Serialize; + +use crate::token_usage_error::TokenUsageError; + +#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, PartialEq, Serialize)] +#[serde(deny_unknown_fields)] pub struct TokenUsage { - prompt_tokens: u64, - cached_prompt_tokens: u64, - input_image_tokens: u64, - input_audio_tokens: u64, - content_tokens: u64, - reasoning_tokens: u64, - undeterminable_tokens: u64, + pub prompt_tokens: u64, + pub cached_prompt_tokens: u64, + pub input_image_tokens: u64, + pub input_audio_tokens: u64, + pub content_tokens: u64, + pub reasoning_tokens: u64, + pub tool_call_tokens: u64, + pub undeterminable_tokens: u64, } impl TokenUsage { @@ -30,6 +30,7 @@ impl TokenUsage { input_audio_tokens: 0, content_tokens: 0, reasoning_tokens: 0, + tool_call_tokens: 0, undeterminable_tokens: 0, } } @@ -39,8 +40,8 @@ impl TokenUsage { } /// # Errors - /// Returns [`TokenUsageError::CachedExceedsPrompt`] when the running cached total would - /// exceed [`Self::prompt_tokens`]. + /// Returns [`TokenUsageError::CachedExceedsPrompt`] when the running cached + /// total would exceed [`Self::prompt_tokens`]. pub const fn record_cached_prompt_tokens(&mut self, count: u64) -> Result<(), TokenUsageError> { let next = self.cached_prompt_tokens.saturating_add(count); @@ -72,56 +73,25 @@ impl TokenUsage { self.reasoning_tokens = self.reasoning_tokens.saturating_add(1); } - pub const fn record_undeterminable_token(&mut self) { - self.undeterminable_tokens = self.undeterminable_tokens.saturating_add(1); + pub const fn record_tool_call_token(&mut self) { + self.tool_call_tokens = self.tool_call_tokens.saturating_add(1); } - pub const fn record_sampled(&mut self, token: &SampledToken) { - match token { - SampledToken::Content(_) => self.record_content_token(), - SampledToken::Reasoning(_) => self.record_reasoning_token(), - SampledToken::Undeterminable(_) => self.record_undeterminable_token(), - } - } - - #[must_use] - pub const fn prompt_tokens(&self) -> u64 { - self.prompt_tokens - } - - #[must_use] - pub const fn cached_prompt_tokens(&self) -> u64 { - self.cached_prompt_tokens - } - - #[must_use] - pub const fn input_image_tokens(&self) -> u64 { - self.input_image_tokens - } - - #[must_use] - pub const fn input_audio_tokens(&self) -> u64 { - self.input_audio_tokens + pub const fn record_undeterminable_token(&mut self) { + self.undeterminable_tokens = self.undeterminable_tokens.saturating_add(1); } #[must_use] - pub const fn content_tokens(&self) -> u64 { + pub const fn completion_tokens(&self) -> u64 { self.content_tokens + .saturating_add(self.reasoning_tokens) + .saturating_add(self.tool_call_tokens) + .saturating_add(self.undeterminable_tokens) } #[must_use] - pub const fn reasoning_tokens(&self) -> u64 { - self.reasoning_tokens - } - - #[must_use] - pub const fn undeterminable_tokens(&self) -> u64 { - self.undeterminable_tokens - } - - #[must_use] - pub const fn completion_tokens(&self) -> u64 { - self.content_tokens.saturating_add(self.reasoning_tokens) + pub const fn total_tokens(&self) -> u64 { + self.prompt_tokens.saturating_add(self.completion_tokens()) } } @@ -130,7 +100,6 @@ impl Add for TokenUsage { fn add(mut self, other: Self) -> Self { self += other; - self } } @@ -140,7 +109,6 @@ impl Add<&Self> for TokenUsage { fn add(mut self, other: &Self) -> Self { self += other; - self } } @@ -165,6 +133,7 @@ impl AddAssign<&Self> for TokenUsage { .saturating_add(other.input_audio_tokens); self.content_tokens = self.content_tokens.saturating_add(other.content_tokens); self.reasoning_tokens = self.reasoning_tokens.saturating_add(other.reasoning_tokens); + self.tool_call_tokens = self.tool_call_tokens.saturating_add(other.tool_call_tokens); self.undeterminable_tokens = self .undeterminable_tokens .saturating_add(other.undeterminable_tokens); @@ -186,23 +155,20 @@ impl<'usage> Sum<&'usage Self> for TokenUsage { #[cfg(test)] mod tests { use super::TokenUsage; - use crate::TokenUsageError; - use crate::sampled_token::SampledToken; - use crate::token::LlamaToken; - - const TOKEN: LlamaToken = LlamaToken::new(7); + use super::TokenUsageError; #[test] fn new_starts_with_all_counters_at_zero() { let usage = TokenUsage::new(); - assert_eq!(usage.prompt_tokens(), 0); - assert_eq!(usage.cached_prompt_tokens(), 0); - assert_eq!(usage.input_image_tokens(), 0); - assert_eq!(usage.input_audio_tokens(), 0); - assert_eq!(usage.content_tokens(), 0); - assert_eq!(usage.reasoning_tokens(), 0); - assert_eq!(usage.undeterminable_tokens(), 0); + assert_eq!(usage.prompt_tokens, 0); + assert_eq!(usage.cached_prompt_tokens, 0); + assert_eq!(usage.input_image_tokens, 0); + assert_eq!(usage.input_audio_tokens, 0); + assert_eq!(usage.content_tokens, 0); + assert_eq!(usage.reasoning_tokens, 0); + assert_eq!(usage.tool_call_tokens, 0); + assert_eq!(usage.undeterminable_tokens, 0); } #[test] @@ -212,9 +178,17 @@ mod tests { #[test] fn completion_is_zero_when_no_events_recorded() { - let usage = TokenUsage::new(); + assert_eq!(TokenUsage::new().completion_tokens(), 0); + } - assert_eq!(usage.completion_tokens(), 0); + #[test] + fn total_equals_prompt_plus_completion() { + let mut usage = TokenUsage::new(); + usage.record_prompt_tokens(3); + usage.record_content_token(); + usage.record_reasoning_token(); + + assert_eq!(usage.total_tokens(), 5); } #[test] @@ -223,26 +197,30 @@ mod tests { usage.record_prompt_tokens(3); usage.record_prompt_tokens(4); - assert_eq!(usage.prompt_tokens(), 7); + assert_eq!(usage.prompt_tokens, 7); } #[test] - fn record_cached_below_prompt_succeeds_and_accumulates() { + fn record_cached_below_prompt_succeeds_and_accumulates() -> Result<(), TokenUsageError> { let mut usage = TokenUsage::new(); usage.record_prompt_tokens(10); - usage.record_cached_prompt_tokens(3).unwrap(); - usage.record_cached_prompt_tokens(4).unwrap(); + usage.record_cached_prompt_tokens(3)?; + usage.record_cached_prompt_tokens(4)?; + + assert_eq!(usage.cached_prompt_tokens, 7); - assert_eq!(usage.cached_prompt_tokens(), 7); + Ok(()) } #[test] - fn record_cached_equal_to_prompt_succeeds() { + fn record_cached_equal_to_prompt_succeeds() -> Result<(), TokenUsageError> { let mut usage = TokenUsage::new(); usage.record_prompt_tokens(5); - usage.record_cached_prompt_tokens(5).unwrap(); + usage.record_cached_prompt_tokens(5)?; + + assert_eq!(usage.cached_prompt_tokens, 5); - assert_eq!(usage.cached_prompt_tokens(), 5); + Ok(()) } #[test] @@ -259,21 +237,7 @@ mod tests { prompt: 2, }) ); - assert_eq!(usage.cached_prompt_tokens(), 0); - } - - #[test] - fn record_cached_can_be_recorded_after_more_prompt_tokens_arrive() { - let mut usage = TokenUsage::new(); - usage.record_prompt_tokens(2); - - let first = usage.record_cached_prompt_tokens(3); - assert!(first.is_err()); - - usage.record_prompt_tokens(5); - usage.record_cached_prompt_tokens(3).unwrap(); - - assert_eq!(usage.cached_prompt_tokens(), 3); + assert_eq!(usage.cached_prompt_tokens, 0); } #[test] @@ -282,7 +246,7 @@ mod tests { usage.record_input_image_tokens(5); usage.record_input_image_tokens(3); - assert_eq!(usage.input_image_tokens(), 8); + assert_eq!(usage.input_image_tokens, 8); } #[test] @@ -291,7 +255,7 @@ mod tests { usage.record_input_audio_tokens(2); usage.record_input_audio_tokens(9); - assert_eq!(usage.input_audio_tokens(), 11); + assert_eq!(usage.input_audio_tokens, 11); } #[test] @@ -299,7 +263,7 @@ mod tests { let mut usage = TokenUsage::new(); usage.record_input_image_tokens(40); - assert_eq!(usage.prompt_tokens(), 0); + assert_eq!(usage.prompt_tokens, 0); assert_eq!(usage.completion_tokens(), 0); } @@ -308,100 +272,92 @@ mod tests { let mut usage = TokenUsage::new(); usage.record_input_audio_tokens(40); - assert_eq!(usage.prompt_tokens(), 0); + assert_eq!(usage.prompt_tokens, 0); assert_eq!(usage.completion_tokens(), 0); } #[test] - fn record_sampled_content_increments_only_content() { + fn record_content_token_increments_only_content() { let mut usage = TokenUsage::new(); - usage.record_sampled(&SampledToken::Content(TOKEN)); + usage.record_content_token(); - assert_eq!(usage.content_tokens(), 1); - assert_eq!(usage.reasoning_tokens(), 0); - assert_eq!(usage.undeterminable_tokens(), 0); + assert_eq!(usage.content_tokens, 1); + assert_eq!(usage.reasoning_tokens, 0); + assert_eq!(usage.tool_call_tokens, 0); + assert_eq!(usage.undeterminable_tokens, 0); } #[test] - fn record_sampled_reasoning_increments_only_reasoning() { + fn record_reasoning_token_increments_only_reasoning() { let mut usage = TokenUsage::new(); - usage.record_sampled(&SampledToken::Reasoning(TOKEN)); + usage.record_reasoning_token(); - assert_eq!(usage.content_tokens(), 0); - assert_eq!(usage.reasoning_tokens(), 1); - assert_eq!(usage.undeterminable_tokens(), 0); + assert_eq!(usage.content_tokens, 0); + assert_eq!(usage.reasoning_tokens, 1); + assert_eq!(usage.tool_call_tokens, 0); + assert_eq!(usage.undeterminable_tokens, 0); } #[test] - fn record_sampled_undeterminable_increments_only_undeterminable() { + fn record_tool_call_token_increments_only_tool_call() { let mut usage = TokenUsage::new(); - usage.record_sampled(&SampledToken::Undeterminable(TOKEN)); + usage.record_tool_call_token(); - assert_eq!(usage.content_tokens(), 0); - assert_eq!(usage.reasoning_tokens(), 0); - assert_eq!(usage.undeterminable_tokens(), 1); + assert_eq!(usage.content_tokens, 0); + assert_eq!(usage.reasoning_tokens, 0); + assert_eq!(usage.tool_call_tokens, 1); + assert_eq!(usage.undeterminable_tokens, 0); } #[test] - fn undeterminable_tokens_do_not_contribute_to_completion_tokens() { + fn record_undeterminable_token_increments_only_undeterminable() { let mut usage = TokenUsage::new(); usage.record_undeterminable_token(); - usage.record_undeterminable_token(); - assert_eq!(usage.undeterminable_tokens(), 2); - assert_eq!(usage.completion_tokens(), 0); + assert_eq!(usage.content_tokens, 0); + assert_eq!(usage.reasoning_tokens, 0); + assert_eq!(usage.tool_call_tokens, 0); + assert_eq!(usage.undeterminable_tokens, 1); } #[test] - fn completion_tokens_sums_only_content_and_reasoning() { + fn completion_tokens_sums_every_output_kind() { let mut usage = TokenUsage::new(); usage.record_content_token(); usage.record_content_token(); usage.record_reasoning_token(); + usage.record_tool_call_token(); + usage.record_undeterminable_token(); - assert_eq!(usage.completion_tokens(), 3); - } - - #[test] - fn independent_instances_do_not_share_counts() { - let mut first = TokenUsage::new(); - let mut second = TokenUsage::new(); - - first.record_prompt_tokens(11); - first.record_content_token(); - - second.record_reasoning_token(); - - assert_eq!(first.prompt_tokens(), 11); - assert_eq!(first.content_tokens(), 1); - assert_eq!(first.reasoning_tokens(), 0); - - assert_eq!(second.prompt_tokens(), 0); - assert_eq!(second.content_tokens(), 0); - assert_eq!(second.reasoning_tokens(), 1); + assert_eq!(usage.completion_tokens(), 5); } #[test] - fn add_combines_field_by_field() { + fn add_combines_field_by_field() -> Result<(), TokenUsageError> { let mut left = TokenUsage::new(); left.record_prompt_tokens(2); - left.record_cached_prompt_tokens(1).unwrap(); + left.record_cached_prompt_tokens(1)?; left.record_content_token(); left.record_reasoning_token(); + left.record_tool_call_token(); left.record_undeterminable_token(); let mut right = TokenUsage::new(); right.record_prompt_tokens(5); - right.record_cached_prompt_tokens(2).unwrap(); + right.record_cached_prompt_tokens(2)?; right.record_content_token(); + right.record_tool_call_token(); let combined = left + right; - assert_eq!(combined.prompt_tokens(), 7); - assert_eq!(combined.cached_prompt_tokens(), 3); - assert_eq!(combined.content_tokens(), 2); - assert_eq!(combined.reasoning_tokens(), 1); - assert_eq!(combined.undeterminable_tokens(), 1); + assert_eq!(combined.prompt_tokens, 7); + assert_eq!(combined.cached_prompt_tokens, 3); + assert_eq!(combined.content_tokens, 2); + assert_eq!(combined.reasoning_tokens, 1); + assert_eq!(combined.tool_call_tokens, 2); + assert_eq!(combined.undeterminable_tokens, 1); + + Ok(()) } #[test] @@ -416,8 +372,8 @@ mod tests { let combined = left + right; - assert_eq!(combined.input_image_tokens(), 7); - assert_eq!(combined.input_audio_tokens(), 8); + assert_eq!(combined.input_image_tokens, 7); + assert_eq!(combined.input_audio_tokens, 8); } #[test] @@ -449,7 +405,7 @@ mod tests { let combined = left + right_ref; - assert_eq!(combined.prompt_tokens(), 7); + assert_eq!(combined.prompt_tokens, 7); } #[test] diff --git a/llama-cpp-bindings-types/src/token_usage_error.rs b/llama-cpp-bindings-types/src/token_usage_error.rs new file mode 100644 index 00000000..b3de4fef --- /dev/null +++ b/llama-cpp-bindings-types/src/token_usage_error.rs @@ -0,0 +1,7 @@ +#[derive(Debug, Eq, PartialEq, thiserror::Error)] +pub enum TokenUsageError { + #[error( + "cached prompt tokens would reach {cached_after} but only {prompt} prompt tokens were recorded" + )] + CachedExceedsPrompt { cached_after: u64, prompt: u64 }, +} diff --git a/llama-cpp-bindings-types/src/tool_call_args_shape.rs b/llama-cpp-bindings-types/src/tool_call_args_shape.rs new file mode 100644 index 00000000..10f3b1fb --- /dev/null +++ b/llama-cpp-bindings-types/src/tool_call_args_shape.rs @@ -0,0 +1,14 @@ +use crate::bracketed_json_shape::BracketedJsonShape; +use crate::json_object_shape::JsonObjectShape; +use crate::key_value_xml_tags_shape::KeyValueXmlTagsShape; +use crate::paired_quote_shape::PairedQuoteShape; +use crate::xml_tags_shape::XmlTagsShape; + +#[derive(Clone, Debug, Eq, PartialEq)] +pub enum ToolCallArgsShape { + BracketedJson(BracketedJsonShape), + JsonObject(JsonObjectShape), + KeyValueXmlTags(KeyValueXmlTagsShape), + PairedQuote(PairedQuoteShape), + XmlTags(XmlTagsShape), +} diff --git a/llama-cpp-bindings-types/src/tool_call_arguments.rs b/llama-cpp-bindings-types/src/tool_call_arguments.rs new file mode 100644 index 00000000..05c77e20 --- /dev/null +++ b/llama-cpp-bindings-types/src/tool_call_arguments.rs @@ -0,0 +1,76 @@ +use serde::Deserialize; +use serde::Serialize; +use serde_json::Value; + +#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)] +pub enum ToolCallArguments { + ValidJson(Value), + InvalidJson(String), +} + +impl ToolCallArguments { + #[must_use] + pub fn from_string(raw: String) -> Self { + serde_json::from_str::(&raw).map_or_else(|_| Self::InvalidJson(raw), Self::ValidJson) + } +} + +impl Default for ToolCallArguments { + fn default() -> Self { + Self::InvalidJson(String::new()) + } +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::ToolCallArguments; + + #[test] + fn from_string_object_returns_valid() { + let result = ToolCallArguments::from_string(r#"{"location":"Paris"}"#.to_owned()); + + assert_eq!( + result, + ToolCallArguments::ValidJson(json!({"location": "Paris"})) + ); + } + + #[test] + fn from_string_array_returns_valid() { + let result = ToolCallArguments::from_string("[1,2,3]".to_owned()); + + assert_eq!(result, ToolCallArguments::ValidJson(json!([1, 2, 3]))); + } + + #[test] + fn from_string_scalar_returns_valid() { + let result = ToolCallArguments::from_string("42".to_owned()); + + assert_eq!(result, ToolCallArguments::ValidJson(json!(42))); + } + + #[test] + fn from_string_unparseable_returns_invalid() { + let raw = "{not really json".to_owned(); + let result = ToolCallArguments::from_string(raw.clone()); + + assert_eq!(result, ToolCallArguments::InvalidJson(raw)); + } + + #[test] + fn from_string_empty_returns_invalid() { + let result = ToolCallArguments::from_string(String::new()); + + assert_eq!(result, ToolCallArguments::InvalidJson(String::new())); + } + + #[test] + fn default_is_empty_invalid() { + assert_eq!( + ToolCallArguments::default(), + ToolCallArguments::InvalidJson(String::new()) + ); + } +} diff --git a/llama-cpp-bindings-types/src/tool_call_markers.rs b/llama-cpp-bindings-types/src/tool_call_markers.rs new file mode 100644 index 00000000..1f6610cd --- /dev/null +++ b/llama-cpp-bindings-types/src/tool_call_markers.rs @@ -0,0 +1,8 @@ +use crate::tool_call_args_shape::ToolCallArgsShape; + +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct ToolCallMarkers { + pub open: String, + pub close: String, + pub args_shape: ToolCallArgsShape, +} diff --git a/llama-cpp-bindings-types/src/tool_call_value_quote.rs b/llama-cpp-bindings-types/src/tool_call_value_quote.rs new file mode 100644 index 00000000..aca34cbf --- /dev/null +++ b/llama-cpp-bindings-types/src/tool_call_value_quote.rs @@ -0,0 +1,5 @@ +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct ToolCallValueQuote { + pub open: String, + pub close: String, +} diff --git a/llama-cpp-bindings-types/src/xml_tags_shape.rs b/llama-cpp-bindings-types/src/xml_tags_shape.rs new file mode 100644 index 00000000..c09634be --- /dev/null +++ b/llama-cpp-bindings-types/src/xml_tags_shape.rs @@ -0,0 +1,7 @@ +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct XmlTagsShape { + pub function_open_prefix: String, + pub function_close: String, + pub parameter_open_prefix: String, + pub parameter_close: String, +} diff --git a/llama-cpp-bindings/Cargo.toml b/llama-cpp-bindings/Cargo.toml index e1cfcbb7..0f592a1f 100644 --- a/llama-cpp-bindings/Cargo.toml +++ b/llama-cpp-bindings/Cargo.toml @@ -8,17 +8,20 @@ repository = "https://github.com/intentee/llama-cpp-bindings" [dependencies] encoding_rs = { workspace = true } -enumflags2 = "0.7.12" +enumflags2 = { workspace = true } llama-cpp-bindings-sys = { workspace = true } -thiserror = "2" +llama-cpp-bindings-types = { workspace = true } +llguidance = { workspace = true } +nom = { workspace = true } +serde_json = { workspace = true } +thiserror = { workspace = true } +toktrie = { workspace = true } tracing = { workspace = true } -tracing-core = "0.1" -llguidance = { version = "1.7.0", optional = true } -toktrie = { version = "1.7.0", optional = true } +tracing-core = { workspace = true } [dev-dependencies] -serial_test = "3" -tracing-subscriber = { version = "0.3", features = ["json"] } +serial_test = { workspace = true } +tracing-subscriber = { workspace = true } [features] default = ["openmp", "android-shared-stdcxx"] @@ -36,7 +39,6 @@ android-shared-stdcxx = ["llama-cpp-bindings-sys/shared-stdcxx"] android-static-stdcxx = ["llama-cpp-bindings-sys/static-stdcxx"] system-ggml = ["llama-cpp-bindings-sys/system-ggml"] system-ggml-static = ["system-ggml", "llama-cpp-bindings-sys/system-ggml-static"] -llguidance = ["dep:llguidance", "dep:toktrie"] [target.'cfg(all(target_os = "macos", any(target_arch = "aarch64", target_arch = "arm64")))'.dependencies] llama-cpp-bindings-sys = { workspace = true, features = ["metal"] } diff --git a/llama-cpp-bindings/src/batch_add_error.rs b/llama-cpp-bindings/src/batch_add_error.rs new file mode 100644 index 00000000..ea4cb154 --- /dev/null +++ b/llama-cpp-bindings/src/batch_add_error.rs @@ -0,0 +1,13 @@ +/// Errors that can occur when adding a token to a batch. +#[derive(thiserror::Error, Debug, PartialEq, Eq)] +pub enum BatchAddError { + /// There was not enough space in the batch to add the token. + #[error("Insufficient Space of {0}")] + InsufficientSpace(usize), + /// Empty buffer is provided for [`crate::llama_batch::LlamaBatch::get_one`] + #[error("Empty buffer")] + EmptyBuffer, + /// An integer value exceeded the allowed range. + #[error("Integer overflow: {0}")] + IntegerOverflow(String), +} diff --git a/llama-cpp-bindings/src/chat_message_parse_outcome.rs b/llama-cpp-bindings/src/chat_message_parse_outcome.rs new file mode 100644 index 00000000..12550664 --- /dev/null +++ b/llama-cpp-bindings/src/chat_message_parse_outcome.rs @@ -0,0 +1,56 @@ +use llama_cpp_bindings_types::ParsedChatMessage; + +use crate::raw_chat_message::RawChatMessage; + +pub enum ChatMessageParseOutcome { + Recognized(ParsedChatMessage), + Unrecognized(RawChatMessage), +} + +#[cfg(test)] +mod tests { + use llama_cpp_bindings_types::ParsedChatMessage; + + use super::ChatMessageParseOutcome; + use crate::raw_chat_message::RawChatMessage; + + #[test] + fn recognized_variant_exposes_parsed_chat_message() { + let parsed = + ParsedChatMessage::new("content".to_owned(), "reasoning".to_owned(), Vec::new()); + let outcome = ChatMessageParseOutcome::Recognized(parsed); + + match outcome { + ChatMessageParseOutcome::Recognized(parsed) => { + assert_eq!(parsed.content, "content"); + assert_eq!(parsed.reasoning_content, "reasoning"); + assert!(parsed.tool_calls.is_empty()); + } + ChatMessageParseOutcome::Unrecognized(_) => { + panic!("expected Recognized variant"); + } + } + } + + #[test] + fn unrecognized_variant_exposes_raw_chat_message() { + let outcome = ChatMessageParseOutcome::Unrecognized(RawChatMessage { + tools_json: "[]".to_owned(), + text: "raw input".to_owned(), + is_partial: false, + ffi_error_message: "parser bailed".to_owned(), + }); + + match outcome { + ChatMessageParseOutcome::Unrecognized(raw) => { + assert_eq!(raw.tools_json, "[]"); + assert_eq!(raw.text, "raw input"); + assert!(!raw.is_partial); + assert_eq!(raw.ffi_error_message, "parser bailed"); + } + ChatMessageParseOutcome::Recognized(_) => { + panic!("expected Unrecognized variant"); + } + } + } +} diff --git a/llama-cpp-bindings/src/context.rs b/llama-cpp-bindings/src/context.rs index cafbcfb5..09d6560d 100644 --- a/llama-cpp-bindings/src/context.rs +++ b/llama-cpp-bindings/src/context.rs @@ -9,6 +9,8 @@ use std::sync::Arc; use std::sync::atomic::AtomicBool; use std::sync::atomic::Ordering; +use crate::context::params::LlamaContextParams; +use crate::llama_backend::LlamaBackend; use crate::llama_batch::LlamaBatch; use crate::model::{LlamaLoraAdapter, LlamaModel}; use crate::timing::LlamaTimings; @@ -16,7 +18,7 @@ use crate::token::LlamaToken; use crate::token::data::LlamaTokenData; use crate::token::data_array::LlamaTokenDataArray; use crate::{ - DecodeError, EmbeddingsError, EncodeError, LlamaLoraAdapterRemoveError, + DecodeError, EmbeddingsError, EncodeError, LlamaContextLoadError, LlamaLoraAdapterRemoveError, LlamaLoraAdapterSetError, LogitsError, }; @@ -87,6 +89,35 @@ impl<'model> LlamaContext<'model> { } } + /// Create a new context bound to `model`. + /// + /// `_backend` is unused in the body but serves as a compile-time witness that + /// the global llama.cpp backend has been initialised before context creation. + /// + /// # Errors + /// + /// Returns [`LlamaContextLoadError`] when llama.cpp fails to allocate the context. + #[expect( + clippy::needless_pass_by_value, + reason = "LlamaContextParams may become non-trivially copyable upstream" + )] + pub fn from_model( + model: &'model LlamaModel, + _backend: &LlamaBackend, + params: LlamaContextParams, + ) -> Result { + let context_params = params.context_params; + let context = unsafe { + llama_cpp_bindings_sys::llama_new_context_with_model( + model.model.as_ptr(), + context_params, + ) + }; + let context = NonNull::new(context).ok_or(LlamaContextLoadError::NullReturn)?; + + Ok(Self::new(model, context, params.embeddings())) + } + /// Gets the max number of logical tokens that can be submitted to decode. Must be greater than or equal to [`Self::n_ubatch`]. #[must_use] pub fn n_batch(&self) -> u32 { diff --git a/llama-cpp-bindings/src/context/params.rs b/llama-cpp-bindings/src/context/params.rs index e85ddeda..bcea5898 100644 --- a/llama-cpp-bindings/src/context/params.rs +++ b/llama-cpp-bindings/src/context/params.rs @@ -121,7 +121,16 @@ impl From for i32 { } /// A rusty wrapper around `ggml_type` for KV cache types. -#[allow(non_camel_case_types, missing_docs)] +#[expect( + non_camel_case_types, + reason = "variant names mirror llama.cpp's `enum ggml_type` symbol names verbatim so they can \ + be matched 1:1 against the C ABI without a translation table" +)] +#[expect( + missing_docs, + reason = "each variant denotes a quantisation flavour whose semantics are defined upstream in \ + ggml; restating the upstream spec inline would risk drifting from the source of truth" +)] #[derive(Copy, Clone, Debug, PartialEq, Eq)] pub enum KvCacheType { /// Represents an unknown or not-yet-mapped `ggml_type` and carries the raw value. @@ -260,10 +269,16 @@ impl From for KvCacheType { /// assert_eq!(ctx_params.n_ctx(), NonZeroU32::new(2048)); /// ``` #[derive(Debug, Clone)] -#[allow( +#[expect( missing_docs, - clippy::struct_excessive_bools, - clippy::module_name_repetitions + reason = "field meanings mirror llama.cpp's `llama_context_params` C struct; restating each \ + one inline would risk drift from the upstream spec — the doc-comment on the struct \ + points at the canonical reference" +)] +#[expect( + clippy::module_name_repetitions, + reason = "`LlamaContextParams` is the canonical Rust name in the public API; renaming it to \ + `Params` would force `params::Params` at every call site" )] pub struct LlamaContextParams { pub context_params: llama_cpp_bindings_sys::llama_context_params, diff --git a/llama-cpp-bindings/src/error.rs b/llama-cpp-bindings/src/error.rs index cb22ca2c..d48e2596 100644 --- a/llama-cpp-bindings/src/error.rs +++ b/llama-cpp-bindings/src/error.rs @@ -4,7 +4,7 @@ use std::os::raw::c_int; use std::path::PathBuf; use std::string::FromUtf8Error; -use crate::llama_batch::BatchAddError; +use crate::batch_add_error::BatchAddError; use crate::mtmd::MtmdEvalError; use crate::mtmd::mtmd_input_chunk_type::MtmdInputChunkTypeError; @@ -333,9 +333,9 @@ pub enum ApplyChatTemplateError { IntConversionError(#[from] std::num::TryFromIntError), } -/// Failed to build a [`crate::reasoning_token_classifier::ReasoningTokenClassifier`] for a model. +/// Failed to detect tool-call diagnostic markers for a model. #[derive(Debug, thiserror::Error)] -pub enum ReasoningClassifierError { +pub enum MarkerDetectionError { /// llama.cpp returned an error code from the marker detection FFI call. #[error("ffi error {0}")] FfiError(i32), @@ -345,39 +345,143 @@ pub enum ReasoningClassifierError { /// llama.cpp returned a marker string but its bytes were not valid UTF-8. #[error("ffi returned non-utf8 marker bytes: {0}")] MarkerUtf8Error(#[from] FromUtf8Error), - /// Tokenizing a detected marker string failed. - #[error("marker tokenization failed: {0}")] - MarkerTokenization(#[from] StringToTokenError), - /// Reading token attributes for a resolved marker token failed. - #[error("token attribute lookup failed: {0}")] - TokenAttr(#[from] crate::token_type::LlamaTokenTypeFromIntError), - /// The detected open-marker string did not tokenize to exactly one token. - #[error("open marker {marker:?} tokenized to {token_count} tokens, expected 1")] - OpenMarkerNotSingleToken { - /// The marker string returned by llama.cpp. - marker: String, - /// The number of tokens the marker tokenized to. - token_count: usize, +} + +/// Failed to parse a chat message via [`crate::Model::parse_chat_message`]. +#[derive(Debug, thiserror::Error)] +pub enum ParseChatMessageError { + /// llama.cpp returned an error code from the parse FFI call. + #[error("ffi error {0}")] + FfiError(i32), + /// The C++ side threw an exception while parsing. + #[error("c++ exception during chat parse: {0}")] + ParseException(String), + /// An accessor returned bytes that were not valid UTF-8. + #[error("ffi returned non-utf8 string: {0}")] + StringUtf8Error(#[from] FromUtf8Error), + /// The caller passed a `tools_json` argument that is not valid JSON. + #[error("tools_json is not valid JSON: {0}")] + ToolsJsonInvalid(#[source] serde_json::Error), + /// The caller passed a `tools_json` argument that parses as JSON but is not an array. + #[error("tools_json must be a JSON array")] + ToolsJsonNotArray, + /// Failed to serialize the tools array for the FFI call. + #[error("could not serialize tools to JSON: {0}")] + ToolsSerialization(String), + /// The model has no usable chat template, so the parser cannot be built. + #[error("model has no chat template")] + NoChatTemplate, + /// The wrapper-side fallback parser detected a structural issue while parsing the body. + #[error("template-override fallback parser failed: {0}")] + TemplateOverrideFailed(#[from] ToolCallFormatFailure), +} + +/// Top-level failure for the wrapper-side template-override parsers (one variant per supported shape). +#[derive(Debug, thiserror::Error)] +pub enum ToolCallFormatFailure { + #[error("bracketed-args fallback parser: {0}")] + BracketedArgs(#[from] BracketedArgsFailure), + #[error("json-object fallback parser: {0}")] + JsonObject(#[from] JsonObjectFailure), + #[error("key-value-xml-tags fallback parser: {0}")] + KeyValueXmlTags(#[from] KeyValueXmlTagsFailure), + #[error("paired-quote fallback parser: {0}")] + PairedQuote(#[from] PairedQuoteFailure), + #[error("xml-function-tags fallback parser: {0}")] + XmlFunctionTags(#[from] XmlFunctionTagsFailure), +} + +/// Failures specific to the JSON-object args parser (Qwen 3 `{"name":..., "arguments":...}`). +#[derive(Debug, thiserror::Error)] +pub enum JsonObjectFailure { + #[error("tool call body has malformed JSON: {message}")] + InvalidJson { message: String }, +} + +/// Failures specific to the bracketed-JSON args parser (Mistral 3 `[TOOL_CALLS]name[ARGS]{...}`). +#[derive(Debug, thiserror::Error)] +pub enum BracketedArgsFailure { + #[error("tool call '{tool_name}' arguments are not valid JSON: {message}")] + InvalidJsonArguments { tool_name: String, message: String }, + #[error("tool call '{tool_name}' arguments truncated before JSON value completed")] + UnterminatedArguments { tool_name: String }, +} + +/// Failures specific to the paired-quote args parser (Gemma 4 `<|tool_call>call:name{key:<|"|>val<|"|>}`). +#[derive(Debug, thiserror::Error)] +pub enum PairedQuoteFailure { + #[error("empty key in tool call '{tool_name}' arguments")] + EmptyKey { tool_name: String }, + #[error("tool call '{tool_name}' translated arguments are not valid JSON: {message}")] + InvalidJsonArguments { tool_name: String, message: String }, + #[error("tool call '{tool_name}' has unclosed quoted value for key '{key}'")] + UnclosedQuotedValue { tool_name: String, key: String }, + #[error("tool call '{tool_name}' arguments ended without close marker (state: {state})")] + UnclosedArgumentBlock { + tool_name: String, + state: &'static str, + }, + #[error( + "tool call '{tool_name}' has unexpected character '{character}' after value for key '{key}'" + )] + UnexpectedCharAfterValue { + tool_name: String, + key: String, + character: char, }, - /// The detected close-marker string did not tokenize to exactly one token. - #[error("close marker {marker:?} tokenized to {token_count} tokens, expected 1")] - CloseMarkerNotSingleToken { - /// The marker string returned by llama.cpp. - marker: String, - /// The number of tokens the marker tokenized to. - token_count: usize, +} + +/// Failures specific to the key-value XML-tags parser (GLM-4.7 `{name}{k}{v}...`). +#[derive(Debug, thiserror::Error)] +pub enum KeyValueXmlTagsFailure { + #[error("tool call function tag has empty name")] + EmptyFunctionName, + #[error("tool call function block is missing close tag '{expected_close}'")] + UnclosedFunctionBlock { expected_close: String }, + #[error("tool call function '{function_name}' has key tag with empty content")] + EmptyKey { function_name: String }, + #[error("tool call function '{function_name}' is missing key close tag '{expected_close}'")] + UnclosedKeyTag { + function_name: String, + expected_close: String, + }, + #[error( + "tool call function '{function_name}' key '{key}' is missing value open tag '{expected_open}'" + )] + MissingValueTag { + function_name: String, + key: String, + expected_open: String, }, - /// The detected open-marker token is not registered as a special token (Control or `UserDefined`). - #[error("open marker {marker:?} is not a registered special token")] - OpenMarkerNotSpecial { - /// The marker string returned by llama.cpp. - marker: String, + #[error( + "tool call function '{function_name}' key '{key}' is missing value close tag '{expected_close}'" + )] + UnclosedValueTag { + function_name: String, + key: String, + expected_close: String, + }, +} + +/// Failures specific to the XML function-tags parser (Qwen 3.5+ `val`). +#[derive(Debug, thiserror::Error)] +pub enum XmlFunctionTagsFailure { + #[error("tool call function tag has empty name")] + EmptyFunctionName, + #[error("tool call function '{function_name}' is missing close tag '{expected_close}'")] + UnclosedFunctionBlock { + function_name: String, + expected_close: String, }, - /// The detected close-marker token is not registered as a special token (Control or `UserDefined`). - #[error("close marker {marker:?} is not a registered special token")] - CloseMarkerNotSpecial { - /// The marker string returned by llama.cpp. - marker: String, + #[error("tool call function '{function_name}' has parameter with empty name")] + EmptyParameterName { function_name: String }, + #[error( + "tool call function '{function_name}' parameter '{parameter_name}' is missing close tag '{expected_close}'" + )] + UnclosedParameterBlock { + function_name: String, + parameter_name: String, + expected_close: String, }, } @@ -395,21 +499,6 @@ pub enum EvalMultimodalChunksError { ChunkOutOfBounds(usize), } -/// Token-usage accounting violations. -#[derive(Debug, Eq, PartialEq, thiserror::Error)] -pub enum TokenUsageError { - /// Cached prompt tokens cannot exceed the recorded prompt total. - #[error( - "cached prompt tokens would reach {cached_after} but only {prompt} prompt tokens were recorded" - )] - CachedExceedsPrompt { - /// Running cached total after this would-be call. - cached_after: u64, - /// Currently recorded prompt-token total. - prompt: u64, - }, -} - /// Failed to accept a token in a sampler. #[derive(Debug, thiserror::Error)] pub enum SamplerAcceptError { diff --git a/llama-cpp-bindings/src/extract_tool_call_markers_from_haystack.rs b/llama-cpp-bindings/src/extract_tool_call_markers_from_haystack.rs new file mode 100644 index 00000000..bbbbdb4e --- /dev/null +++ b/llama-cpp-bindings/src/extract_tool_call_markers_from_haystack.rs @@ -0,0 +1,143 @@ +use crate::tool_call_marker_pair::ToolCallMarkerPair; + +#[must_use] +pub fn extract_tool_call_markers_from_haystack(haystack: &str) -> Option { + if haystack.is_empty() { + return None; + } + + let json_start = haystack.find('{')?; + let json_end = haystack.rfind('}')?; + if json_end < json_start { + return None; + } + + let json_slice = &haystack[json_start..=json_end]; + serde_json::from_str::(json_slice).ok()?; + + let open = haystack[..json_start].trim().to_owned(); + let close = haystack[json_end + 1..].trim().to_owned(); + + if open.is_empty() || close.is_empty() { + return None; + } + + Some(ToolCallMarkerPair { open, close }) +} + +#[cfg(test)] +mod tests { + use super::ToolCallMarkerPair; + use super::extract_tool_call_markers_from_haystack; + + #[test] + fn extracts_open_and_close_around_a_simple_json_payload() { + let pair = extract_tool_call_markers_from_haystack( + "{\"name\":\"x\",\"arguments\":{}}", + ); + + assert_eq!( + pair, + Some(ToolCallMarkerPair { + open: "".to_owned(), + close: "".to_owned(), + }), + ); + } + + #[test] + fn trims_surrounding_whitespace_from_each_marker() { + let pair = extract_tool_call_markers_from_haystack( + " \n {\"k\": 1}\n ", + ); + + assert_eq!( + pair, + Some(ToolCallMarkerPair { + open: "".to_owned(), + close: "".to_owned(), + }), + ); + } + + #[test] + fn returns_none_when_haystack_is_empty() { + assert_eq!(extract_tool_call_markers_from_haystack(""), None); + } + + #[test] + fn returns_none_when_haystack_has_no_open_brace() { + assert_eq!( + extract_tool_call_markers_from_haystack("plain assistant text"), + None + ); + } + + #[test] + fn returns_none_when_haystack_has_open_brace_but_no_close() { + assert_eq!( + extract_tool_call_markers_from_haystack("{ unclosed"), + None + ); + } + + #[test] + fn returns_none_when_close_brace_precedes_open_brace() { + assert_eq!( + extract_tool_call_markers_from_haystack("}{"), + None + ); + } + + #[test] + fn returns_none_when_brace_payload_is_not_valid_json() { + assert_eq!( + extract_tool_call_markers_from_haystack("{not valid json}"), + None + ); + } + + #[test] + fn returns_none_when_open_marker_resolves_to_empty_after_trim() { + assert_eq!( + extract_tool_call_markers_from_haystack(" {\"x\":1}"), + None + ); + } + + #[test] + fn returns_none_when_close_marker_resolves_to_empty_after_trim() { + assert_eq!( + extract_tool_call_markers_from_haystack("{\"x\":1} "), + None + ); + } + + #[test] + fn extracts_around_an_object_that_contains_nested_braces() { + let pair = extract_tool_call_markers_from_haystack( + "{\"args\":{\"k\":[1,2,{\"deep\":true}]}}", + ); + + assert_eq!( + pair, + Some(ToolCallMarkerPair { + open: "".to_owned(), + close: "".to_owned(), + }), + ); + } + + #[test] + fn extracts_when_open_marker_contains_multibyte_utf8() { + let pair = extract_tool_call_markers_from_haystack("<|tool→call|>{\"k\":1}<|/tool→call|>"); + + assert_eq!( + pair, + Some(ToolCallMarkerPair { + open: "<|tool→call|>".to_owned(), + close: "<|/tool→call|>".to_owned(), + }), + ); + } +} diff --git a/llama-cpp-bindings/src/ingest_prompt_chunk.rs b/llama-cpp-bindings/src/ingest_prompt_chunk.rs new file mode 100644 index 00000000..c17b0993 --- /dev/null +++ b/llama-cpp-bindings/src/ingest_prompt_chunk.rs @@ -0,0 +1,37 @@ +use crate::mtmd::MtmdInputChunk; +use crate::mtmd::MtmdInputChunkType; +use crate::mtmd::MtmdInputChunkTypeError; +use crate::sampled_token_classifier::SampledTokenClassifier; + +/// Dispatches a single multimodal chunk into the classifier: +/// - Text chunks bump `prompt_tokens` and replay every text token through the +/// marker state machine, so prompt-end markers like `` reach the +/// classifier and the section transitions before generation begins. +/// - Image / Audio chunks bump only their own usage counters; they have no +/// text token IDs to replay. +/// +/// This is the single canonical per-chunk ingest path for the multimodal +/// driver. Any future per-chunk invariant (e.g. cached prefix replay) lives +/// here so it cannot diverge between consumers. +/// +/// # Errors +/// Returns [`MtmdInputChunkTypeError`] when the chunk reports a type unknown +/// to this binding. Counters are not updated on error. +pub fn ingest_prompt_chunk( + classifier: &mut SampledTokenClassifier<'_>, + chunk: &MtmdInputChunk, +) -> Result<(), MtmdInputChunkTypeError> { + let n_tokens = chunk.n_tokens() as u64; + match chunk.chunk_type()? { + MtmdInputChunkType::Text => { + classifier.record_prompt_tokens(n_tokens); + if let Some(tokens) = chunk.text_tokens() { + classifier.ingest_prompt_tokens(tokens); + } + } + MtmdInputChunkType::Image => classifier.record_input_image_tokens(n_tokens), + MtmdInputChunkType::Audio => classifier.record_input_audio_tokens(n_tokens), + } + + Ok(()) +} diff --git a/llama-cpp-bindings/src/lib.rs b/llama-cpp-bindings/src/lib.rs index 5c1288c3..4ee62c7e 100644 --- a/llama-cpp-bindings/src/lib.rs +++ b/llama-cpp-bindings/src/lib.rs @@ -10,8 +10,11 @@ //! - `cuda` enables CUDA gpu support. //! - `sampler` adds the [`context::sample::sampler`] struct for a more rusty way of sampling. +pub mod batch_add_error; +pub mod chat_message_parse_outcome; pub mod context; pub mod error; +pub mod extract_tool_call_markers_from_haystack; pub mod ffi_error_reader; pub mod ffi_status_is_ok; pub mod ffi_status_to_i32; @@ -19,13 +22,16 @@ pub mod ggml_time_us; pub mod gguf_context; pub mod gguf_context_error; pub mod gguf_type; +pub mod ingest_prompt_chunk; pub mod json_schema_to_grammar; pub mod llama_backend; pub mod llama_backend_device; pub mod llama_backend_numa_strategy; pub mod llama_batch; pub mod llama_time_us; -#[cfg(feature = "llguidance")] +pub mod llama_token_attr; +pub mod llama_token_attrs; +pub mod llama_token_attrs_from_int_error; pub mod llguidance_sampler; #[cfg(feature = "dynamic-backends")] pub mod load_backends; @@ -40,33 +46,45 @@ pub mod mlock_supported; pub mod mmap_supported; pub mod model; pub mod mtmd; -pub mod reasoning_token_classifier; +pub mod raw_chat_message; +pub mod resolved_tool_call_markers; pub mod sampled_token; +pub mod sampled_token_classifier; pub mod sampling; +pub mod streaming_json_probe; pub mod timing; pub mod token; -pub mod token_type; -pub mod token_usage; +pub mod tool_call_format; +pub mod tool_call_marker_pair; +pub mod tool_call_template_overrides; pub use error::{ ApplyChatTemplateError, ChatTemplateError, DecodeError, EmbeddingsError, EncodeError, EvalMultimodalChunksError, GrammarError, LlamaContextLoadError, LlamaCppError, LlamaLoraAdapterInitError, LlamaLoraAdapterRemoveError, LlamaLoraAdapterSetError, - LlamaModelLoadError, LogitsError, MetaValError, ModelParamsError, NewLlamaChatMessageError, - ReasoningClassifierError, Result, SampleError, SamplerAcceptError, SamplingError, - StringToTokenError, TokenSamplingError, TokenToStringError, TokenUsageError, + LlamaModelLoadError, LogitsError, MarkerDetectionError, MetaValError, ModelParamsError, + NewLlamaChatMessageError, ParseChatMessageError, Result, SampleError, SamplerAcceptError, + SamplingError, StringToTokenError, TokenSamplingError, TokenToStringError, }; +pub use chat_message_parse_outcome::ChatMessageParseOutcome; pub use llama_backend_device::{ LlamaBackendDevice, LlamaBackendDeviceType, list_llama_ggml_backend_devices, }; -pub use reasoning_token_classifier::ReasoningTokenClassifier; +pub use llama_cpp_bindings_types::{ + BracketedJsonShape, KeyValueXmlTagsShape, PairedQuoteShape, ParsedChatMessage, ParsedToolCall, + ReasoningMarkers, TokenUsage, TokenUsageError, ToolCallArgsShape, ToolCallArguments, + ToolCallMarkers, ToolCallValueQuote, XmlTagsShape, +}; +pub use raw_chat_message::RawChatMessage; pub use sampled_token::SampledToken; -pub use token_usage::TokenUsage; +pub use sampled_token_classifier::SampledTokenClassifier; +pub use sampled_token_classifier::SampledTokenSection; pub use ffi_status_is_ok::status_is_ok; pub use ffi_status_to_i32::status_to_i32; pub use ggml_time_us::ggml_time_us; +pub use ingest_prompt_chunk::ingest_prompt_chunk; pub use json_schema_to_grammar::json_schema_to_grammar; pub use llama_time_us::llama_time_us; pub use max_devices::max_devices; diff --git a/llama-cpp-bindings/src/llama_backend.rs b/llama-cpp-bindings/src/llama_backend.rs index 223775e2..803c27a2 100644 --- a/llama-cpp-bindings/src/llama_backend.rs +++ b/llama-cpp-bindings/src/llama_backend.rs @@ -19,8 +19,8 @@ impl LlamaBackend { /// Mark the llama backend as initialized fn mark_init() -> crate::Result<()> { match LLAMA_BACKEND_INITIALIZED.compare_exchange(false, true, SeqCst, SeqCst) { - Ok(_) => Ok(()), - Err(_) => Err(LlamaCppError::BackendAlreadyInitialized), + Ok(_was_uninitialized) => Ok(()), + Err(_was_already_initialized) => Err(LlamaCppError::BackendAlreadyInitialized), } } diff --git a/llama-cpp-bindings/src/llama_batch.rs b/llama-cpp-bindings/src/llama_batch.rs index 9c412fde..b6b8b189 100644 --- a/llama-cpp-bindings/src/llama_batch.rs +++ b/llama-cpp-bindings/src/llama_batch.rs @@ -1,5 +1,6 @@ //! Safe wrapper around `llama_batch`. +use crate::batch_add_error::BatchAddError; use crate::sampled_token::SampledToken; use crate::token::LlamaToken; use llama_cpp_bindings_sys::{ @@ -67,20 +68,6 @@ pub struct LlamaBatch<'tokens> { phantom: PhantomData<&'tokens [LlamaToken]>, } -/// Errors that can occur when adding a token to a batch. -#[derive(thiserror::Error, Debug, PartialEq, Eq)] -pub enum BatchAddError { - /// There was not enough space in the batch to add the token. - #[error("Insufficient Space of {0}")] - InsufficientSpace(usize), - /// Empty buffer is provided for [`LlamaBatch::get_one`] - #[error("Empty buffer")] - EmptyBuffer, - /// An integer value exceeded the allowed range. - #[error("Integer overflow: {0}")] - IntegerOverflow(String), -} - impl<'tokens> LlamaBatch<'tokens> { /// Clear the batch. This does not free the memory associated with the batch, but it does reset /// the number of tokens to 0. @@ -104,6 +91,7 @@ impl<'tokens> LlamaBatch<'tokens> { ) -> Result<(), BatchAddError> { let (SampledToken::Content(LlamaToken(id)) | SampledToken::Reasoning(LlamaToken(id)) + | SampledToken::ToolCall(LlamaToken(id)) | SampledToken::Undeterminable(LlamaToken(id))) = *sampled_token; let required = checked_n_tokens_plus_one_as_usize(self.n_tokens())?; @@ -324,6 +312,44 @@ mod tests { assert_eq!(result, Err(BatchAddError::InsufficientSpace(1))); } + #[test] + fn add_accepts_reasoning_sampled_token_variant() { + let mut batch = LlamaBatch::new(4, 1).unwrap(); + + batch + .add(&SampledToken::Reasoning(LlamaToken::new(11)), 0, &[0], true) + .unwrap(); + + assert_eq!(batch.n_tokens(), 1); + } + + #[test] + fn add_accepts_tool_call_sampled_token_variant() { + let mut batch = LlamaBatch::new(4, 1).unwrap(); + + batch + .add(&SampledToken::ToolCall(LlamaToken::new(22)), 0, &[0], true) + .unwrap(); + + assert_eq!(batch.n_tokens(), 1); + } + + #[test] + fn add_accepts_undeterminable_sampled_token_variant() { + let mut batch = LlamaBatch::new(4, 1).unwrap(); + + batch + .add( + &SampledToken::Undeterminable(LlamaToken::new(33)), + 0, + &[0], + false, + ) + .unwrap(); + + assert_eq!(batch.n_tokens(), 1); + } + #[test] fn add_sequence_adds_all_tokens() { let mut batch = LlamaBatch::new(16, 1).unwrap(); diff --git a/llama-cpp-bindings/src/llama_token_attr.rs b/llama-cpp-bindings/src/llama_token_attr.rs new file mode 100644 index 00000000..fb9de83c --- /dev/null +++ b/llama-cpp-bindings/src/llama_token_attr.rs @@ -0,0 +1,28 @@ +use enumflags2::bitflags; + +/// A rust flavored equivalent of `llama_token_type`. +#[derive(Eq, PartialEq, Debug, Clone, Copy)] +#[bitflags] +#[repr(u32)] +pub enum LlamaTokenAttr { + /// Unknown token attribute. + Unknown = llama_cpp_bindings_sys::LLAMA_TOKEN_ATTR_UNKNOWN as _, + /// Unused token attribute. + Unused = llama_cpp_bindings_sys::LLAMA_TOKEN_ATTR_UNUSED as _, + /// Normal text token. + Normal = llama_cpp_bindings_sys::LLAMA_TOKEN_ATTR_NORMAL as _, + /// Control token (e.g. BOS, EOS). + Control = llama_cpp_bindings_sys::LLAMA_TOKEN_ATTR_CONTROL as _, + /// User-defined token. + UserDefined = llama_cpp_bindings_sys::LLAMA_TOKEN_ATTR_USER_DEFINED as _, + /// Byte-level fallback token. + Byte = llama_cpp_bindings_sys::LLAMA_TOKEN_ATTR_BYTE as _, + /// Token with normalized text. + Normalized = llama_cpp_bindings_sys::LLAMA_TOKEN_ATTR_NORMALIZED as _, + /// Token with left-stripped whitespace. + LStrip = llama_cpp_bindings_sys::LLAMA_TOKEN_ATTR_LSTRIP as _, + /// Token with right-stripped whitespace. + RStrip = llama_cpp_bindings_sys::LLAMA_TOKEN_ATTR_RSTRIP as _, + /// Token representing a single word. + SingleWord = llama_cpp_bindings_sys::LLAMA_TOKEN_ATTR_SINGLE_WORD as _, +} diff --git a/llama-cpp-bindings/src/token_type.rs b/llama-cpp-bindings/src/llama_token_attrs.rs similarity index 50% rename from llama-cpp-bindings/src/token_type.rs rename to llama-cpp-bindings/src/llama_token_attrs.rs index 4405582b..37d46651 100644 --- a/llama-cpp-bindings/src/token_type.rs +++ b/llama-cpp-bindings/src/llama_token_attrs.rs @@ -1,35 +1,11 @@ -//! Utilities for working with `llama_token_type` values. -use enumflags2::{BitFlags, bitflags}; use std::ops::{Deref, DerefMut}; -/// A rust flavored equivalent of `llama_token_type`. -#[derive(Eq, PartialEq, Debug, Clone, Copy)] -#[bitflags] -#[repr(u32)] -pub enum LlamaTokenAttr { - /// Unknown token attribute. - Unknown = llama_cpp_bindings_sys::LLAMA_TOKEN_ATTR_UNKNOWN as _, - /// Unused token attribute. - Unused = llama_cpp_bindings_sys::LLAMA_TOKEN_ATTR_UNUSED as _, - /// Normal text token. - Normal = llama_cpp_bindings_sys::LLAMA_TOKEN_ATTR_NORMAL as _, - /// Control token (e.g. BOS, EOS). - Control = llama_cpp_bindings_sys::LLAMA_TOKEN_ATTR_CONTROL as _, - /// User-defined token. - UserDefined = llama_cpp_bindings_sys::LLAMA_TOKEN_ATTR_USER_DEFINED as _, - /// Byte-level fallback token. - Byte = llama_cpp_bindings_sys::LLAMA_TOKEN_ATTR_BYTE as _, - /// Token with normalized text. - Normalized = llama_cpp_bindings_sys::LLAMA_TOKEN_ATTR_NORMALIZED as _, - /// Token with left-stripped whitespace. - LStrip = llama_cpp_bindings_sys::LLAMA_TOKEN_ATTR_LSTRIP as _, - /// Token with right-stripped whitespace. - RStrip = llama_cpp_bindings_sys::LLAMA_TOKEN_ATTR_RSTRIP as _, - /// Token representing a single word. - SingleWord = llama_cpp_bindings_sys::LLAMA_TOKEN_ATTR_SINGLE_WORD as _, -} +use enumflags2::BitFlags; + +use crate::llama_token_attr::LlamaTokenAttr; +use crate::llama_token_attrs_from_int_error::LlamaTokenAttrsFromIntError; -/// A set of `LlamaTokenAttrs` +/// A set of [`LlamaTokenAttr`] flags. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct LlamaTokenAttrs(pub BitFlags); @@ -48,28 +24,22 @@ impl DerefMut for LlamaTokenAttrs { } impl TryFrom for LlamaTokenAttrs { - type Error = LlamaTokenTypeFromIntError; + type Error = LlamaTokenAttrsFromIntError; fn try_from(value: llama_cpp_bindings_sys::llama_vocab_type) -> Result { - Ok(Self(BitFlags::from_bits(value as _).map_err(|e| { - LlamaTokenTypeFromIntError::UnknownValue(e.invalid_bits()) - })?)) + Ok(Self(BitFlags::from_bits(value as _).map_err( + |bit_flag_error| { + LlamaTokenAttrsFromIntError::UnknownValue(bit_flag_error.invalid_bits()) + }, + )?)) } } -/// An error type for `LlamaTokenType::try_from`. -#[derive(thiserror::Error, Debug, Eq, PartialEq)] -pub enum LlamaTokenTypeFromIntError { - /// The value is not a valid `llama_token_type`. - #[error("Unknown Value {0}")] - UnknownValue(std::ffi::c_uint), -} - #[cfg(test)] mod tests { use enumflags2::BitFlags; - use super::{LlamaTokenAttr, LlamaTokenAttrs, LlamaTokenTypeFromIntError}; + use super::{LlamaTokenAttr, LlamaTokenAttrs, LlamaTokenAttrsFromIntError}; #[test] fn try_from_valid_single_attribute() { @@ -99,7 +69,7 @@ mod tests { assert!(result.is_err()); matches!( result.expect_err("should fail"), - LlamaTokenTypeFromIntError::UnknownValue(_) + LlamaTokenAttrsFromIntError::UnknownValue(_) ); } diff --git a/llama-cpp-bindings/src/llama_token_attrs_from_int_error.rs b/llama-cpp-bindings/src/llama_token_attrs_from_int_error.rs new file mode 100644 index 00000000..df1ad6c2 --- /dev/null +++ b/llama-cpp-bindings/src/llama_token_attrs_from_int_error.rs @@ -0,0 +1,9 @@ +/// Returned by [`crate::llama_token_attrs::LlamaTokenAttrs::try_from`] when the +/// integer bit pattern contains bits not defined by +/// [`crate::llama_token_attr::LlamaTokenAttr`]. +#[derive(thiserror::Error, Debug, Eq, PartialEq)] +pub enum LlamaTokenAttrsFromIntError { + /// The value is not a valid `llama_token_type`. + #[error("Unknown Value {0}")] + UnknownValue(std::ffi::c_uint), +} diff --git a/llama-cpp-bindings/src/llguidance_sampler.rs b/llama-cpp-bindings/src/llguidance_sampler.rs index b4ab2288..67da9f09 100644 --- a/llama-cpp-bindings/src/llguidance_sampler.rs +++ b/llama-cpp-bindings/src/llguidance_sampler.rs @@ -7,12 +7,11 @@ use std::ffi::c_void; use std::sync::Arc; use llguidance::Matcher; -use toktrie::{ApproximateTokEnv, TokRxInfo, TokTrie}; +use toktrie::ApproximateTokEnv; use crate::GrammarError; use crate::model::LlamaModel; use crate::sampling::LlamaSampler; -use crate::token::LlamaToken; /// Internal state for the llguidance sampler. struct LlgContext { @@ -22,51 +21,6 @@ struct LlgContext { grammar_data: String, } -/// Build a [`toktrie::TokEnv`] from a [`LlamaModel`]'s vocabulary. -/// -/// This mirrors the logic in upstream `llguidance.cpp` — for each token: -/// - Try normal detokenize (special=false) -/// - If empty, detokenize with special=true and prefix with 0xFF marker byte -fn build_tok_env(model: &LlamaModel) -> Arc { - let n_vocab = model.n_vocab().cast_unsigned(); - let tok_eos = { - let eot = unsafe { llama_cpp_bindings_sys::llama_vocab_eot(model.vocab_ptr()) }; - if eot == -1 { - model.token_eos().0.cast_unsigned() - } else { - eot.cast_unsigned() - } - }; - let info = TokRxInfo::new(n_vocab, tok_eos); - - let mut words = Vec::with_capacity(n_vocab as usize); - - for token_id in 0..n_vocab.cast_signed() { - let token = LlamaToken(token_id); - let bytes = model - .token_to_piece_bytes(token, 32, false, None) - .unwrap_or_default(); - if bytes.is_empty() { - let special_bytes = model - .token_to_piece_bytes(token, 32, true, None) - .unwrap_or_default(); - if special_bytes.is_empty() { - words.push(vec![]); - } else { - let mut marked = Vec::with_capacity(special_bytes.len() + 1); - marked.push(0xFF); - marked.extend(special_bytes); - words.push(marked); - } - } else { - words.push(bytes); - } - } - - let trie = TokTrie::from(&info, &words); - Arc::new(ApproximateTokEnv::new(trie)) -} - const unsafe extern "C" fn llg_name( _smpl: *const llama_cpp_bindings_sys::llama_sampler, ) -> *const std::os::raw::c_char { @@ -175,7 +129,7 @@ pub fn create_llg_sampler( grammar_kind: &str, grammar_data: &str, ) -> Result { - let tok_env = build_tok_env(model); + let tok_env = model.approximate_tok_env(); let tok_env_dyn: Arc = tok_env.clone(); let factory = llguidance::ParserFactory::new_simple(&tok_env_dyn) diff --git a/llama-cpp-bindings/src/log.rs b/llama-cpp-bindings/src/log.rs index 38404e03..639cc5f0 100644 --- a/llama-cpp-bindings/src/log.rs +++ b/llama-cpp-bindings/src/log.rs @@ -500,7 +500,11 @@ mod tests { } struct Logger { - #[allow(unused)] + #[expect( + unused, + reason = "guard must outlive the test body so the tracing subscriber stays installed; \ + dropping it un-installs the subscriber and tests would silently miss log lines" + )] guard: tracing::subscriber::DefaultGuard, logs: Arc>>, } diff --git a/llama-cpp-bindings/src/model.rs b/llama-cpp-bindings/src/model.rs index e976c949..e8d5ac01 100644 --- a/llama-cpp-bindings/src/model.rs +++ b/llama-cpp-bindings/src/model.rs @@ -3,6 +3,12 @@ use std::ffi::{CStr, CString, c_char}; use std::num::NonZeroU16; use std::os::raw::c_int; use std::path::Path; +use std::sync::Arc; +use std::sync::OnceLock; + +use toktrie::ApproximateTokEnv; +use toktrie::TokRxInfo; +use toktrie::TokTrie; fn truncated_buffer_to_string( mut buffer: Vec, @@ -24,19 +30,31 @@ fn cstring_with_validated_len(str: &str) -> Result<(CString, c_int), StringToTok } use std::ptr::{self, NonNull}; -use crate::context::LlamaContext; -use crate::context::params::LlamaContextParams; +use crate::chat_message_parse_outcome::ChatMessageParseOutcome; use crate::ffi_status_to_i32::status_to_i32; use crate::llama_backend::LlamaBackend; -use crate::reasoning_token_classifier::ReasoningTokenClassifier; +use crate::llama_token_attrs::LlamaTokenAttrs; +use crate::llama_token_attrs_from_int_error::LlamaTokenAttrsFromIntError; +use crate::raw_chat_message::RawChatMessage; +use crate::resolved_tool_call_markers::ResolvedToolCallMarkers; use crate::sampled_token::SampledToken; +use crate::sampled_token_classifier::SampledTokenClassifier; +use crate::sampled_token_classifier::StreamingMarkers; use crate::token::LlamaToken; -use crate::token_type::{LlamaTokenAttr, LlamaTokenAttrs}; use crate::{ - ApplyChatTemplateError, ChatTemplateError, LlamaContextLoadError, LlamaLoraAdapterInitError, - LlamaModelLoadError, MetaValError, ReasoningClassifierError, StringToTokenError, + ApplyChatTemplateError, ChatTemplateError, LlamaLoraAdapterInitError, LlamaModelLoadError, + MarkerDetectionError, MetaValError, ParseChatMessageError, StringToTokenError, TokenToStringError, }; +use llama_cpp_bindings_types::ParsedChatMessage; +use llama_cpp_bindings_types::ParsedToolCall; +use llama_cpp_bindings_types::ReasoningMarkers; +use llama_cpp_bindings_types::ToolCallArguments; +use llama_cpp_bindings_types::ToolCallMarkers; + +use crate::tool_call_format; +use crate::tool_call_format::ToolCallFormatOutcome; +use crate::tool_call_template_overrides; pub mod add_bos; pub mod llama_chat_message; @@ -46,22 +64,31 @@ pub mod params; pub mod rope_type; pub mod split_mode; pub mod vocab_type; +pub mod vocab_type_from_int_error; pub use add_bos::AddBos; pub use llama_chat_message::LlamaChatMessage; pub use llama_chat_template::LlamaChatTemplate; pub use llama_lora_adapter::LlamaLoraAdapter; pub use rope_type::RopeType; -pub use vocab_type::{LlamaTokenTypeFromIntError, VocabType}; +pub use vocab_type::VocabType; +pub use vocab_type_from_int_error::VocabTypeFromIntError; use params::LlamaModelParams; /// A safe wrapper around `llama_model`. -#[derive(Debug)] -#[repr(transparent)] pub struct LlamaModel { /// Raw pointer to the underlying `llama_model`. pub model: NonNull, + tok_env: OnceLock>, +} + +impl std::fmt::Debug for LlamaModel { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("LlamaModel") + .field("model", &self.model) + .finish_non_exhaustive() + } } unsafe impl Send for LlamaModel {} @@ -133,6 +160,7 @@ impl LlamaModel { pub fn is_eog_token(&self, token: &SampledToken) -> bool { let (SampledToken::Content(LlamaToken(id)) | SampledToken::Reasoning(LlamaToken(id)) + | SampledToken::ToolCall(LlamaToken(id)) | SampledToken::Undeterminable(LlamaToken(id))) = *token; unsafe { llama_cpp_bindings_sys::llama_token_is_eog(self.vocab_ptr(), id) } @@ -237,7 +265,7 @@ impl LlamaModel { pub fn token_attr( &self, LlamaToken(id): LlamaToken, - ) -> Result { + ) -> Result { let token_type = unsafe { llama_cpp_bindings_sys::llama_token_get_attr(self.vocab_ptr(), id) }; @@ -268,6 +296,7 @@ impl LlamaModel { ) -> Result { let (SampledToken::Content(inner) | SampledToken::Reasoning(inner) + | SampledToken::ToolCall(inner) | SampledToken::Undeterminable(inner)) = *token; let bytes = match self.token_to_piece_bytes(inner, 8, special, lstrip) { Err(TokenToStringError::InsufficientBufferSpace(required_size)) => { @@ -296,7 +325,6 @@ impl LlamaModel { /// - if the token type is unknown /// - the resultant token is larger than `buffer_size`. /// - if an integer conversion fails - #[allow(clippy::missing_panics_doc)] pub fn token_to_piece_bytes( &self, token: LlamaToken, @@ -304,18 +332,15 @@ impl LlamaModel { special: bool, lstrip: Option, ) -> Result, TokenToStringError> { - // SAFETY: `*` (0x2A) is never `\0`, so CString::new cannot fail here - let string = CString::new(vec![b'*'; buffer_size]).expect("no null"); - let len = string.as_bytes().len(); - let len = c_int::try_from(len)?; - let buf = string.into_raw(); + let mut buffer: Vec = vec![0u8; buffer_size]; + let buffer_len = c_int::try_from(buffer.len())?; let lstrip = lstrip.map_or(0, |strip_count| i32::from(strip_count.get())); let size = unsafe { llama_cpp_bindings_sys::llama_token_to_piece( self.vocab_ptr(), token.0, - buf, - len, + buffer.as_mut_ptr().cast::(), + buffer_len, lstrip, special, ) @@ -327,12 +352,10 @@ impl LlamaModel { Err(TokenToStringError::InsufficientBufferSpace(error_code)) } size => { - let string = unsafe { CString::from_raw(buf) }; - let mut bytes = string.into_bytes(); - let len = usize::try_from(size)?; - bytes.truncate(len); + let written = usize::try_from(size)?; + buffer.truncate(written); - Ok(bytes) + Ok(buffer) } } } @@ -351,7 +374,7 @@ impl LlamaModel { /// # Errors /// /// Returns an error if llama.cpp emits a vocab type that is not known to this library. - pub fn vocab_type(&self) -> Result { + pub fn vocab_type(&self) -> Result { let vocab_type = unsafe { llama_cpp_bindings_sys::llama_vocab_type(self.vocab_ptr()) }; VocabType::try_from(vocab_type) @@ -522,10 +545,8 @@ impl LlamaModel { Err(ChatTemplateError::MissingTemplate) } else { let chat_template_cstr = unsafe { CStr::from_ptr(result) }; - let chat_template = CString::new(chat_template_cstr.to_bytes()) - .expect("CStr bytes cannot contain interior null bytes"); - Ok(LlamaChatTemplate(chat_template)) + Ok(LlamaChatTemplate(chat_template_cstr.to_owned())) } } @@ -567,7 +588,10 @@ impl LlamaModel { None => return Err(LlamaModelLoadError::NullResult), }; - Ok(Self { model }) + Ok(Self { + model, + tok_env: OnceLock::new(), + }) } /// Initializes a lora adapter from a file. @@ -603,32 +627,6 @@ impl LlamaModel { }) } - /// Create a new context from this model. - /// - /// # Errors - /// - /// There is many ways this can fail. See [`LlamaContextLoadError`] for more information. - #[expect( - clippy::needless_pass_by_value, - reason = "LlamaContextParams may become non-trivially copyable upstream" - )] - pub fn new_context<'model>( - &'model self, - _: &LlamaBackend, - params: LlamaContextParams, - ) -> Result, LlamaContextLoadError> { - let context_params = params.context_params; - let context = unsafe { - llama_cpp_bindings_sys::llama_new_context_with_model( - self.model.as_ptr(), - context_params, - ) - }; - let context = NonNull::new(context).ok_or(LlamaContextLoadError::NullReturn)?; - - Ok(LlamaContext::new(self, context, params.embeddings())) - } - /// Apply the models chat template to some messages. /// See /// @@ -707,126 +705,542 @@ impl LlamaModel { truncated_buffer_to_string(buff, final_size) } - /// Build a [`ReasoningTokenClassifier`] for this model by detecting the model's - /// reasoning markers via llama.cpp's chat-template analyzer and resolving them - /// to single Control-attribute token ids. + /// Build a streaming [`SampledTokenClassifier`] for this model. + /// + /// At construction the bindings detect reasoning markers (via the + /// autoparser, with a chunked-thinking fallback for templates that consume + /// thoughts via content blocks), tool-call markers, and the trailing + /// generation-prompt slice. The classifier then runs a state machine over + /// the decoded token stream — no per-model branches. + /// + /// If the model has no usable chat template the classifier is built in a + /// blind mode that classifies every token as + /// [`SampledToken::Undeterminable`]. + pub fn sampled_token_classifier(&self) -> SampledTokenClassifier<'_> { + let markers = match self.streaming_markers() { + Ok(markers) => markers, + Err(detection_error) => { + tracing::warn!( + "streaming markers detection failed; classifier will run blind: {detection_error}" + ); + StreamingMarkers::default() + } + }; + + SampledTokenClassifier::new(self, markers) + } + + /// Detect reasoning / tool-call markers (as token-ID sequences) and the + /// trailing generation-prompt slice for this model's chat template. The + /// returned `StreamingMarkers` carry tokenised markers — never raw strings + /// — so the classifier matches by `LlamaToken` equality rather than text + /// scanning. + /// + /// # Errors + /// Returns [`MarkerDetectionError`] when any underlying FFI call fails. + pub fn streaming_markers(&self) -> Result { + let (reasoning_open_str, reasoning_close_str) = + invoke_ffi_string_pair_detector(|first, second, error| unsafe { + llama_cpp_bindings_sys::llama_rs_detect_reasoning_markers( + self.model.as_ptr(), + first, + second, + error, + ) + })?; + + let tool_call_haystack = invoke_ffi_single_string_detector(|haystack, error| unsafe { + llama_cpp_bindings_sys::llama_rs_compute_tool_call_haystack( + self.model.as_ptr(), + haystack, + error, + ) + })?; + + let autoparser_pair = tool_call_haystack.as_deref().and_then( + crate::extract_tool_call_markers_from_haystack::extract_tool_call_markers_from_haystack, + ); + + let (autoparser_open, autoparser_close) = match autoparser_pair { + Some(crate::tool_call_marker_pair::ToolCallMarkerPair { open, close }) => { + (Some(open), Some(close)) + } + None => (None, None), + }; + + let resolved_tool_call_markers = + self.resolve_tool_call_marker_strings(autoparser_open, autoparser_close); + + Ok(StreamingMarkers { + reasoning_open: self.tokenize_marker(reasoning_open_str.as_deref()), + reasoning_close: self.tokenize_marker(reasoning_close_str.as_deref()), + tool_call_open: self.tokenize_marker(resolved_tool_call_markers.open.as_deref()), + tool_call_close: self.tokenize_marker(resolved_tool_call_markers.close.as_deref()), + }) + } + + /// When the autoparser-driven FFI returned no tool-call markers, consult the + /// per-template override registry so wrapper-known templates (Gemma 4, + /// Mistral 3, ...) still drive the classifier. + fn resolve_tool_call_marker_strings( + &self, + autoparser_open: Option, + autoparser_close: Option, + ) -> ResolvedToolCallMarkers { + if autoparser_open + .as_deref() + .is_some_and(|raw| !raw.trim().is_empty()) + { + return ResolvedToolCallMarkers { + open: autoparser_open, + close: autoparser_close, + }; + } + let Some(markers) = self.tool_call_markers() else { + return ResolvedToolCallMarkers { + open: autoparser_open, + close: autoparser_close, + }; + }; + let close = if markers.close.is_empty() { + None + } else { + Some(markers.close) + }; + ResolvedToolCallMarkers { + open: Some(markers.open), + close, + } + } + + /// # Errors + /// Returns [`MarkerDetectionError`] when the underlying FFI call fails. + pub fn reasoning_markers(&self) -> Result, MarkerDetectionError> { + let (open, close) = invoke_ffi_string_pair_detector(|first, second, error| unsafe { + llama_cpp_bindings_sys::llama_rs_detect_reasoning_markers( + self.model.as_ptr(), + first, + second, + error, + ) + })?; + + match (open, close) { + (Some(open), Some(close)) if !open.is_empty() && !close.is_empty() => { + Ok(Some(ReasoningMarkers { open, close })) + } + _ => Ok(None), + } + } + + /// Returns the rich tool-call marker bundle (open / separator / close / + /// optional value-quote pair) for this model's chat template, sourced from + /// the wrapper's per-template override registry. Returns `None` when no + /// registered override matches — callers in that case fall back to + /// llama.cpp's autoparser via [`Self::parse_chat_message`]. + #[must_use] + pub fn tool_call_markers(&self) -> Option { + let template = match self.chat_template(None) { + Ok(template) => template, + Err(error) => { + tracing::debug!( + "tool-call markers unavailable: chat template missing or invalid: {error}" + ); + return None; + } + }; + let template_str = match template.to_str() { + Ok(template_str) => template_str, + Err(error) => { + tracing::debug!( + "tool-call markers unavailable: chat template is not valid UTF-8: {error}" + ); + return None; + } + }; + tool_call_template_overrides::detect(template_str) + } + + fn tokenize_marker(&self, marker: Option<&str>) -> Option> { + let marker = marker?.trim(); + if marker.is_empty() { + return None; + } + match self.str_to_token(marker, AddBos::Never) { + Ok(tokens) if !tokens.is_empty() => Some(tokens), + Ok(_) => None, + Err(tokenize_error) => { + tracing::debug!( + "marker {marker:?} failed to tokenise; classifier will ignore it: {tokenize_error}" + ); + None + } + } + } + + /// Parse the assistant's output text into structured content, reasoning, + /// and tool calls. + /// + /// Two passes, in order: + /// 1. Duck-type the wrapper-side parsers across every known shape + /// (Qwen XML, GLM key-value, Gemma paired-quote, Mistral bracketed-JSON). + /// First match wins. The shapes are ordered so that more restrictive + /// shapes run first, which keeps the duck-type pass safe for inputs + /// that share an open marker but differ in inner structure. + /// 2. Delegate to llama.cpp's `common_chat_parse`. If it succeeds the + /// result is `Recognized`; if it throws `ParseException` the result is + /// `Unrecognized` with the raw input plus the FFI's diagnostic, so the + /// caller can pass the unstructured tokens to the client. + /// + /// Empty tool-call `id` fields are filled with `call_{index}` before + /// returning, so callers always see well-formed identifiers. /// - /// Returns an `Ok(undetermined)` classifier when the model exposes no detectable - /// reasoning markers — that is the canonical "this model has no reasoning" signal. + /// `tools_json` is a JSON-array string of OpenAI-style tool definitions + /// (use `"[]"` when no tools are in scope). `is_partial` switches between + /// mid-stream (lenient) and final (strict) parses for the FFI step. /// /// # Errors /// - /// Returns [`ReasoningClassifierError`] when the C++ analyzer throws, when a - /// detected marker does not tokenize to exactly one token, or when the resolved - /// token does not have the [`LlamaTokenAttr::Control`] attribute. - pub fn reasoning_token_classifier( + /// Returns [`ParseChatMessageError`] when `tools_json` is not valid JSON, + /// the FFI returns a non-OK status other than `ParseException`, or + /// accessor strings are not valid UTF-8. + pub fn parse_chat_message( &self, - ) -> Result { - let mut out_open: *mut c_char = ptr::null_mut(); - let mut out_close: *mut c_char = ptr::null_mut(); + tools_json: &str, + input: &str, + is_partial: bool, + ) -> Result { + let tools_value: serde_json::Value = + serde_json::from_str(tools_json).map_err(ParseChatMessageError::ToolsJsonInvalid)?; + if !tools_value.is_array() { + return Err(ParseChatMessageError::ToolsJsonNotArray); + } + + let reasoning_markers = self.reasoning_markers().ok().flatten(); + + for candidate in tool_call_template_overrides::known_marker_candidates() { + if let ToolCallFormatOutcome::Parsed(calls) = + tool_call_format::try_parse(input, &candidate) + { + let split = + split_reasoning_prefix(input, reasoning_markers.as_ref(), &candidate.open); + let mut parsed = ParsedChatMessage::new(split.content, split.reasoning, calls); + synthesize_missing_tool_call_ids(&mut parsed.tool_calls); + return Ok(ChatMessageParseOutcome::Recognized(parsed)); + } + } + + match self.parse_chat_message_via_ffi(tools_json, input, is_partial) { + Ok(mut parsed) => { + synthesize_missing_tool_call_ids(&mut parsed.tool_calls); + Ok(ChatMessageParseOutcome::Recognized(parsed)) + } + Err(ParseChatMessageError::ParseException(ffi_error_message)) => { + Ok(ChatMessageParseOutcome::Unrecognized(RawChatMessage { + tools_json: tools_json.to_owned(), + text: input.to_owned(), + is_partial, + ffi_error_message, + })) + } + Err(other) => Err(other), + } + } + + fn parse_chat_message_via_ffi( + &self, + tools_json: &str, + input: &str, + is_partial: bool, + ) -> Result { + let tools_cstring = CString::new(tools_json) + .map_err(|err| ParseChatMessageError::ToolsSerialization(err.to_string()))?; + let input_cstring = CString::new(input) + .map_err(|err| ParseChatMessageError::ToolsSerialization(err.to_string()))?; + + let mut handle: *mut llama_cpp_bindings_sys::llama_rs_parsed_chat = ptr::null_mut(); let mut out_error: *mut c_char = ptr::null_mut(); let status = unsafe { - llama_cpp_bindings_sys::llama_rs_detect_reasoning_markers( + llama_cpp_bindings_sys::llama_rs_parse_chat_message( self.model.as_ptr(), - &raw mut out_open, - &raw mut out_close, + tools_cstring.as_ptr(), + input_cstring.as_ptr(), + i32::from(is_partial), + &raw mut handle, &raw mut out_error, ) }; - let parsed = (|| match status { - llama_cpp_bindings_sys::LLAMA_RS_STATUS_OK => { - let open_string = read_optional_owned_cstr(out_open)?; - let close_string = read_optional_owned_cstr(out_close)?; - - Ok((open_string, close_string)) - } + let parsed = match status { + llama_cpp_bindings_sys::LLAMA_RS_STATUS_OK => collect_parsed_chat_message(handle), llama_cpp_bindings_sys::LLAMA_RS_STATUS_EXCEPTION => { let message = read_optional_owned_cstr_lossy(out_error); - - Err(ReasoningClassifierError::AnalyzeException(message)) + Err(ParseChatMessageError::ParseException(message)) } - other => Err(ReasoningClassifierError::FfiError(status_to_i32(other))), - })(); + other => Err(ParseChatMessageError::FfiError(status_to_i32(other))), + }; - unsafe { llama_cpp_bindings_sys::llama_rs_string_free(out_open) }; - unsafe { llama_cpp_bindings_sys::llama_rs_string_free(out_close) }; + unsafe { llama_cpp_bindings_sys::llama_rs_parsed_chat_free(handle) }; unsafe { llama_cpp_bindings_sys::llama_rs_string_free(out_error) }; - let (open_string, close_string) = parsed?; - - let (Some(open_marker), Some(close_marker)) = (open_string, close_string) else { - return Ok(ReasoningTokenClassifier::undetermined()); - }; + parsed + } - let open_marker = open_marker.trim(); - let close_marker = close_marker.trim(); + /// Render the model's chat template with the autoparser's synthetic + /// no-tools and with-tools inputs. Returns `(output_no_tools, + /// output_with_tools)`. Either side can be empty when the template throws + /// during rendering. Useful for debugging tool-call marker detection. + /// + /// # Errors + /// + /// Returns [`MarkerDetectionError`] when the C++ analyzer throws or the FFI + /// returns a non-OK status. + pub fn diagnose_tool_call_synthetic_renders( + &self, + ) -> Result<(String, String), MarkerDetectionError> { + let (no_tools, with_tools) = + invoke_ffi_string_pair_detector(|first, second, error| unsafe { + llama_cpp_bindings_sys::llama_rs_diagnose_tool_call_synthetic_renders( + self.model.as_ptr(), + first, + second, + error, + ) + })?; - let open_token = self.resolve_open_reasoning_marker(open_marker)?; - let close_token = self.resolve_close_reasoning_marker(close_marker)?; + Ok((no_tools.unwrap_or_default(), with_tools.unwrap_or_default())) + } +} - Ok(ReasoningTokenClassifier::new(open_token, close_token)) +impl LlamaModel { + /// Returns a process-cached, approximate token environment built from this model's vocabulary. + /// + /// The first call iterates the full vocabulary and constructs the trie; subsequent calls + /// return the cached `Arc` without further FFI work. + pub fn approximate_tok_env(&self) -> Arc { + Arc::clone(self.tok_env.get_or_init(|| build_approximate_tok_env(self))) } +} - fn resolve_open_reasoning_marker( - &self, - marker: &str, - ) -> Result { - let tokens = self.str_to_token(marker, AddBos::Never)?; - - if tokens.len() != 1 { - return Err(ReasoningClassifierError::OpenMarkerNotSingleToken { - marker: marker.to_string(), - token_count: tokens.len(), - }); +fn build_approximate_tok_env(model: &LlamaModel) -> Arc { + let n_vocab = model.n_vocab().cast_unsigned(); + let tok_eos = { + let eot = unsafe { llama_cpp_bindings_sys::llama_vocab_eot(model.vocab_ptr()) }; + if eot == -1 { + model.token_eos().0.cast_unsigned() + } else { + eot.cast_unsigned() + } + }; + let info = TokRxInfo::new(n_vocab, tok_eos); + + let mut words = Vec::with_capacity(n_vocab as usize); + + for token_id in 0..n_vocab.cast_signed() { + let token = LlamaToken(token_id); + let bytes = model + .token_to_piece_bytes(token, 32, false, None) + .unwrap_or_default(); + if bytes.is_empty() { + let special_bytes = model + .token_to_piece_bytes(token, 32, true, None) + .unwrap_or_default(); + if special_bytes.is_empty() { + words.push(vec![]); + } else { + let mut marked = Vec::with_capacity(special_bytes.len() + 1); + marked.push(0xFF); + marked.extend(special_bytes); + words.push(marked); + } + } else { + words.push(bytes); } + } - let token = tokens[0]; - let attrs = self.token_attr(token)?; + let trie = TokTrie::from(&info, &words); + Arc::new(ApproximateTokEnv::new(trie)) +} - if !is_special_marker_attr(attrs) { - return Err(ReasoningClassifierError::OpenMarkerNotSpecial { - marker: marker.to_string(), - }); +fn collect_parsed_chat_message( + handle: *mut llama_cpp_bindings_sys::llama_rs_parsed_chat, +) -> Result { + if handle.is_null() { + return Ok(ParsedChatMessage::default()); + } + + let content = read_owned_cstr_for_parse(unsafe { + llama_cpp_bindings_sys::llama_rs_parsed_chat_content(handle) + })?; + let reasoning_content = read_owned_cstr_for_parse(unsafe { + llama_cpp_bindings_sys::llama_rs_parsed_chat_reasoning_content(handle) + })?; + + let count = unsafe { llama_cpp_bindings_sys::llama_rs_parsed_chat_tool_call_count(handle) }; + + let mut tool_calls = Vec::with_capacity(count); + for index in 0..count { + let id = read_owned_cstr_for_parse(unsafe { + llama_cpp_bindings_sys::llama_rs_parsed_chat_tool_call_id(handle, index) + })?; + let name = read_owned_cstr_for_parse(unsafe { + llama_cpp_bindings_sys::llama_rs_parsed_chat_tool_call_name(handle, index) + })?; + let arguments_json = read_owned_cstr_for_parse(unsafe { + llama_cpp_bindings_sys::llama_rs_parsed_chat_tool_call_arguments(handle, index) + })?; + + let arguments = ToolCallArguments::from_string(arguments_json); + tool_calls.push(ParsedToolCall::new(id, name, arguments)); + } + + Ok(ParsedChatMessage::new( + content, + reasoning_content, + tool_calls, + )) +} + +struct ReasoningSplit { + reasoning: String, + content: String, +} + +fn split_reasoning_prefix( + input: &str, + reasoning_markers: Option<&ReasoningMarkers>, + tool_call_open: &str, +) -> ReasoningSplit { + let content_only = || ReasoningSplit { + reasoning: String::new(), + content: prefix_before(input, tool_call_open), + }; + + let Some(reasoning_markers) = reasoning_markers else { + return content_only(); + }; + let Some(open_pos) = input.find(&reasoning_markers.open) else { + return content_only(); + }; + + let after_open = &input[open_pos + reasoning_markers.open.len()..]; + let Some(close_offset) = after_open.find(&reasoning_markers.close) else { + return content_only(); + }; + + let reasoning = after_open[..close_offset].to_owned(); + let after_close = &after_open[close_offset + reasoning_markers.close.len()..]; + + ReasoningSplit { + reasoning, + content: prefix_before(after_close, tool_call_open), + } +} + +fn prefix_before(text: &str, marker: &str) -> String { + text.find(marker) + .map_or_else(|| text.to_owned(), |pos| text[..pos].to_owned()) +} + +fn synthesize_missing_tool_call_ids(tool_calls: &mut [ParsedToolCall]) { + for (index, call) in tool_calls.iter_mut().enumerate() { + if call.id.is_empty() { + call.id = format!("call_{index}"); } + } +} - Ok(token) +fn parse_single_string_status( + status: llama_cpp_bindings_sys::llama_rs_status, + out_value: *mut c_char, + out_error: *mut c_char, +) -> Result, MarkerDetectionError> { + match status { + llama_cpp_bindings_sys::LLAMA_RS_STATUS_OK => read_optional_owned_cstr(out_value), + llama_cpp_bindings_sys::LLAMA_RS_STATUS_EXCEPTION => { + let message = read_optional_owned_cstr_lossy(out_error); + + Err(MarkerDetectionError::AnalyzeException(message)) + } + other => Err(MarkerDetectionError::FfiError(status_to_i32(other))), } +} - fn resolve_close_reasoning_marker( - &self, - marker: &str, - ) -> Result { - let tokens = self.str_to_token(marker, AddBos::Never)?; - - if tokens.len() != 1 { - return Err(ReasoningClassifierError::CloseMarkerNotSingleToken { - marker: marker.to_string(), - token_count: tokens.len(), - }); +fn invoke_ffi_single_string_detector( + invoke: TInvoke, +) -> Result, MarkerDetectionError> +where + TInvoke: FnOnce(*mut *mut c_char, *mut *mut c_char) -> llama_cpp_bindings_sys::llama_rs_status, +{ + let mut out_value: *mut c_char = ptr::null_mut(); + let mut out_error: *mut c_char = ptr::null_mut(); + + let status = invoke(&raw mut out_value, &raw mut out_error); + let parsed = parse_single_string_status(status, out_value, out_error); + + unsafe { + if !out_value.is_null() { + llama_cpp_bindings_sys::llama_rs_string_free(out_value); + } + if !out_error.is_null() { + llama_cpp_bindings_sys::llama_rs_string_free(out_error); } + } - let token = tokens[0]; - let attrs = self.token_attr(token)?; + parsed +} + +fn invoke_ffi_string_pair_detector( + invoke: TInvoke, +) -> Result<(Option, Option), MarkerDetectionError> +where + TInvoke: FnOnce( + *mut *mut c_char, + *mut *mut c_char, + *mut *mut c_char, + ) -> llama_cpp_bindings_sys::llama_rs_status, +{ + let mut out_first: *mut c_char = ptr::null_mut(); + let mut out_second: *mut c_char = ptr::null_mut(); + let mut out_error: *mut c_char = ptr::null_mut(); - if !is_special_marker_attr(attrs) { - return Err(ReasoningClassifierError::CloseMarkerNotSpecial { - marker: marker.to_string(), - }); + let status = invoke(&raw mut out_first, &raw mut out_second, &raw mut out_error); + + let parsed = (|| match status { + llama_cpp_bindings_sys::LLAMA_RS_STATUS_OK => { + let first = read_optional_owned_cstr(out_first)?; + let second = read_optional_owned_cstr(out_second)?; + + Ok((first, second)) } + llama_cpp_bindings_sys::LLAMA_RS_STATUS_EXCEPTION => { + let message = read_optional_owned_cstr_lossy(out_error); - Ok(token) - } + Err(MarkerDetectionError::AnalyzeException(message)) + } + other => Err(MarkerDetectionError::FfiError(status_to_i32(other))), + })(); + + unsafe { llama_cpp_bindings_sys::llama_rs_string_free(out_first) }; + unsafe { llama_cpp_bindings_sys::llama_rs_string_free(out_second) }; + unsafe { llama_cpp_bindings_sys::llama_rs_string_free(out_error) }; + + parsed } -fn is_special_marker_attr(attrs: LlamaTokenAttrs) -> bool { - attrs.contains(LlamaTokenAttr::Control) || attrs.contains(LlamaTokenAttr::UserDefined) +fn read_owned_cstr_for_parse(ptr: *mut c_char) -> Result { + if ptr.is_null() { + return Ok(String::new()); + } + + let bytes = unsafe { CStr::from_ptr(ptr) }.to_bytes().to_vec(); + unsafe { llama_cpp_bindings_sys::llama_rs_string_free(ptr) }; + + Ok(String::from_utf8(bytes)?) } -fn read_optional_owned_cstr( - ptr: *const c_char, -) -> Result, ReasoningClassifierError> { +fn read_optional_owned_cstr(ptr: *const c_char) -> Result, MarkerDetectionError> { if ptr.is_null() { return Ok(None); } @@ -977,3 +1391,152 @@ mod extract_meta_string_tests { assert!(result.is_err()); } } + +#[cfg(test)] +mod ffi_helper_tests { + use std::ffi::CString; + use std::ptr; + + use super::invoke_ffi_single_string_detector; + use super::invoke_ffi_string_pair_detector; + use super::parse_single_string_status; + use super::read_optional_owned_cstr_lossy; + use crate::MarkerDetectionError; + + #[test] + fn read_optional_owned_cstr_lossy_returns_empty_for_null() { + let result = read_optional_owned_cstr_lossy(ptr::null()); + + assert!(result.is_empty()); + } + + #[test] + fn read_optional_owned_cstr_lossy_returns_string_for_valid_pointer() { + let owned = CString::new("hello").expect("static literal has no nuls"); + let result = read_optional_owned_cstr_lossy(owned.as_ptr()); + + assert_eq!(result, "hello"); + } + + #[test] + fn read_optional_owned_cstr_lossy_handles_invalid_utf8_via_replacement() { + let owned = CString::new(vec![b'a', 0xFF, b'b']).expect("no interior nul"); + let result = read_optional_owned_cstr_lossy(owned.as_ptr()); + + assert!(result.starts_with('a')); + assert!(result.ends_with('b')); + } + + #[test] + fn parse_single_string_status_returns_none_for_ok_with_null() { + let result = parse_single_string_status( + llama_cpp_bindings_sys::LLAMA_RS_STATUS_OK, + ptr::null_mut(), + ptr::null_mut(), + ); + + assert_eq!(result.expect("OK + null returns Ok(None)"), None); + } + + #[test] + fn parse_single_string_status_returns_some_for_ok_with_value() { + let owned = CString::new("present").expect("no nul"); + let result = parse_single_string_status( + llama_cpp_bindings_sys::LLAMA_RS_STATUS_OK, + owned.as_ptr().cast_mut(), + ptr::null_mut(), + ); + + assert_eq!( + result.expect("OK + value returns Ok(Some)"), + Some("present".to_owned()) + ); + } + + #[test] + fn parse_single_string_status_returns_analyze_exception() { + let owned = CString::new("boom").expect("no nul"); + let result = parse_single_string_status( + llama_cpp_bindings_sys::LLAMA_RS_STATUS_EXCEPTION, + ptr::null_mut(), + owned.as_ptr().cast_mut(), + ); + + match result.expect_err("EXCEPTION must yield Err") { + MarkerDetectionError::AnalyzeException(message) => assert_eq!(message, "boom"), + other => panic!("expected AnalyzeException, got {other:?}"), + } + } + + #[test] + fn parse_single_string_status_returns_ffi_error_for_other_status() { + let result = parse_single_string_status( + llama_cpp_bindings_sys::LLAMA_RS_STATUS_INVALID_ARGUMENT, + ptr::null_mut(), + ptr::null_mut(), + ); + + match result.expect_err("invalid status must yield Err") { + MarkerDetectionError::FfiError(_) => {} + other => panic!("expected FfiError, got {other:?}"), + } + } + + #[test] + fn invoke_ffi_single_string_detector_propagates_invalid_argument_status() { + let result = invoke_ffi_single_string_detector(|_value, _error| { + llama_cpp_bindings_sys::LLAMA_RS_STATUS_INVALID_ARGUMENT + }); + + assert!(matches!(result, Err(MarkerDetectionError::FfiError(_)))); + } + + #[test] + fn invoke_ffi_single_string_detector_returns_none_for_ok_with_null() { + let result = invoke_ffi_single_string_detector(|value, _error| { + unsafe { + *value = ptr::null_mut(); + } + llama_cpp_bindings_sys::LLAMA_RS_STATUS_OK + }); + + assert_eq!(result.expect("OK + null returns Ok(None)"), None); + } + + #[test] + fn invoke_ffi_string_pair_detector_propagates_invalid_argument_status() { + let result = invoke_ffi_string_pair_detector(|_first, _second, _error| { + llama_cpp_bindings_sys::LLAMA_RS_STATUS_INVALID_ARGUMENT + }); + + assert!(matches!(result, Err(MarkerDetectionError::FfiError(_)))); + } + + #[test] + fn invoke_ffi_string_pair_detector_returns_pair_of_none_for_ok_with_nulls() { + let result = invoke_ffi_string_pair_detector(|first, second, _error| { + unsafe { + *first = ptr::null_mut(); + *second = ptr::null_mut(); + } + llama_cpp_bindings_sys::LLAMA_RS_STATUS_OK + }); + + assert_eq!( + result.expect("OK with both null returns Ok((None, None))"), + (None, None) + ); + } + + #[test] + fn invoke_ffi_string_pair_detector_propagates_invalid_status_codes() { + let result = invoke_ffi_string_pair_detector(|_first, _second, _error| { + llama_cpp_bindings_sys::LLAMA_RS_STATUS_ALLOCATION_FAILED + }); + + match result.expect_err("non-OK status yields Err") { + MarkerDetectionError::FfiError(code) => assert!(code != 0), + other => panic!("expected FfiError, got {other:?}"), + } + } +} diff --git a/llama-cpp-bindings/src/model/vocab_type.rs b/llama-cpp-bindings/src/model/vocab_type.rs index c5a6f819..4c790755 100644 --- a/llama-cpp-bindings/src/model/vocab_type.rs +++ b/llama-cpp-bindings/src/model/vocab_type.rs @@ -1,3 +1,5 @@ +use crate::model::vocab_type_from_int_error::VocabTypeFromIntError; + /// a rusty equivalent of `llama_vocab_type` #[repr(u32)] #[derive(Debug, Eq, Copy, Clone, PartialEq)] @@ -8,29 +10,21 @@ pub enum VocabType { SPM = llama_cpp_bindings_sys::LLAMA_VOCAB_TYPE_SPM as _, } -/// There was an error converting a `llama_vocab_type` to a `VocabType`. -#[derive(thiserror::Error, Debug, Eq, PartialEq)] -pub enum LlamaTokenTypeFromIntError { - /// The value is not a valid `llama_token_type`. Contains the int value that was invalid. - #[error("Unknown Value {0}")] - UnknownValue(llama_cpp_bindings_sys::llama_vocab_type), -} - impl TryFrom for VocabType { - type Error = LlamaTokenTypeFromIntError; + type Error = VocabTypeFromIntError; fn try_from(value: llama_cpp_bindings_sys::llama_vocab_type) -> Result { match value { llama_cpp_bindings_sys::LLAMA_VOCAB_TYPE_BPE => Ok(Self::BPE), llama_cpp_bindings_sys::LLAMA_VOCAB_TYPE_SPM => Ok(Self::SPM), - unknown => Err(LlamaTokenTypeFromIntError::UnknownValue(unknown)), + unknown => Err(VocabTypeFromIntError::UnknownValue(unknown)), } } } #[cfg(test)] mod tests { - use super::{LlamaTokenTypeFromIntError, VocabType}; + use super::{VocabType, VocabTypeFromIntError}; #[test] fn try_from_bpe() { @@ -50,6 +44,6 @@ mod tests { fn try_from_unknown_value() { let result = VocabType::try_from(99999); - assert_eq!(result, Err(LlamaTokenTypeFromIntError::UnknownValue(99999))); + assert_eq!(result, Err(VocabTypeFromIntError::UnknownValue(99999))); } } diff --git a/llama-cpp-bindings/src/model/vocab_type_from_int_error.rs b/llama-cpp-bindings/src/model/vocab_type_from_int_error.rs new file mode 100644 index 00000000..3e7bcf8e --- /dev/null +++ b/llama-cpp-bindings/src/model/vocab_type_from_int_error.rs @@ -0,0 +1,8 @@ +/// Returned by [`crate::model::vocab_type::VocabType::try_from`] when the +/// integer value does not match a known `llama_vocab_type` discriminant. +#[derive(thiserror::Error, Debug, Eq, PartialEq)] +pub enum VocabTypeFromIntError { + /// The value is not a valid `llama_vocab_type`. Contains the int value that was invalid. + #[error("Unknown Value {0}")] + UnknownValue(llama_cpp_bindings_sys::llama_vocab_type), +} diff --git a/llama-cpp-bindings/src/mtmd.rs b/llama-cpp-bindings/src/mtmd.rs index 90cbaf80..e5787a83 100644 --- a/llama-cpp-bindings/src/mtmd.rs +++ b/llama-cpp-bindings/src/mtmd.rs @@ -6,6 +6,7 @@ //! # Warning //! This API is experimental and subject to breaking changes. +pub mod image_chunk_batch_size_mismatch; pub mod mtmd_bitmap; pub mod mtmd_context; pub mod mtmd_context_params; @@ -16,6 +17,7 @@ pub mod mtmd_input_chunk_type; pub mod mtmd_input_chunks; pub mod mtmd_input_text; +pub use image_chunk_batch_size_mismatch::ImageChunkBatchSizeMismatch; pub use mtmd_bitmap::MtmdBitmap; pub use mtmd_context::MtmdContext; pub use mtmd_context_params::MtmdContextParams; diff --git a/llama-cpp-bindings/src/mtmd/image_chunk_batch_size_mismatch.rs b/llama-cpp-bindings/src/mtmd/image_chunk_batch_size_mismatch.rs new file mode 100644 index 00000000..992b0eec --- /dev/null +++ b/llama-cpp-bindings/src/mtmd/image_chunk_batch_size_mismatch.rs @@ -0,0 +1,12 @@ +/// Carried by [`super::mtmd_error::MtmdEvalError::ImageChunkExceedsBatchSize`]. +/// +/// `n_batch` is the per-decode batch budget enforced by `cparams.n_batch` in +/// llama.cpp; `image_tokens` is the number of tokens this image chunk would +/// hand to `llama_decode`. When `image_tokens > n_batch` the C-side +/// `GGML_ASSERT(n_tokens_all <= cparams.n_batch)` would abort the process — +/// the binding refuses the call instead. +#[derive(Debug)] +pub struct ImageChunkBatchSizeMismatch { + pub image_tokens: u32, + pub n_batch: u32, +} diff --git a/llama-cpp-bindings/src/mtmd/mtmd_bitmap.rs b/llama-cpp-bindings/src/mtmd/mtmd_bitmap.rs index b068998d..8076d6e6 100644 --- a/llama-cpp-bindings/src/mtmd/mtmd_bitmap.rs +++ b/llama-cpp-bindings/src/mtmd/mtmd_bitmap.rs @@ -80,7 +80,7 @@ impl MtmdBitmap { /// /// // Create a simple sine wave audio sample /// let audio_data: Vec = (0..100) - /// .map(|i| (i as f32 * 0.1).sin()) + /// .map(|sample_index| (sample_index as f32 * 0.1).sin()) /// .collect(); /// /// let bitmap = MtmdBitmap::from_audio_data(&audio_data); @@ -283,7 +283,11 @@ mod tests { #[test] fn from_audio_data_creates_valid_bitmap() { - #[allow(clippy::cast_precision_loss)] + #[expect( + clippy::cast_precision_loss, + reason = "test fixture casts a small i32 (0..100) to f32 to synthesise a sine wave; \ + the values are well within f32's exact-representation range" + )] let audio_samples: Vec = (0..100).map(|index| (index as f32 * 0.1).sin()).collect(); let bitmap = MtmdBitmap::from_audio_data(&audio_samples).unwrap(); diff --git a/llama-cpp-bindings/src/mtmd/mtmd_error.rs b/llama-cpp-bindings/src/mtmd/mtmd_error.rs index 09048ab8..687b7243 100644 --- a/llama-cpp-bindings/src/mtmd/mtmd_error.rs +++ b/llama-cpp-bindings/src/mtmd/mtmd_error.rs @@ -70,6 +70,8 @@ pub enum MtmdEncodeError { EncodeFailure(i32), } +use crate::mtmd::image_chunk_batch_size_mismatch::ImageChunkBatchSizeMismatch; + /// Errors that can occur during evaluation #[derive(thiserror::Error, Debug)] pub enum MtmdEvalError { @@ -81,6 +83,14 @@ pub enum MtmdEvalError { /// The maximum batch size configured on the context context_max: u32, }, + /// An image chunk's token count exceeds the per-decode `n_batch` budget, + /// so handing it to `llama_decode` would trip the `GGML_ASSERT`. + #[error( + "image chunk has {} tokens but n_batch is {}", + .0.image_tokens, + .0.n_batch, + )] + ImageChunkExceedsBatchSize(ImageChunkBatchSizeMismatch), /// Evaluation operation failed #[error("Eval failed with code: {0}")] EvalFailure(i32), diff --git a/llama-cpp-bindings/src/mtmd/mtmd_input_chunk.rs b/llama-cpp-bindings/src/mtmd/mtmd_input_chunk.rs index a1987e4e..4bfa1110 100644 --- a/llama-cpp-bindings/src/mtmd/mtmd_input_chunk.rs +++ b/llama-cpp-bindings/src/mtmd/mtmd_input_chunk.rs @@ -2,8 +2,12 @@ use std::ffi::CStr; use std::ptr::NonNull; use std::slice; +use crate::context::LlamaContext; use crate::token::LlamaToken; +use super::image_chunk_batch_size_mismatch::ImageChunkBatchSizeMismatch; +use super::mtmd_context::MtmdContext; +use super::mtmd_error::MtmdEvalError; use super::mtmd_error::MtmdInputChunkError; use super::mtmd_input_chunk_type::{MtmdInputChunkType, MtmdInputChunkTypeError}; @@ -109,6 +113,74 @@ impl MtmdInputChunk { Ok(Self { chunk, owned: true }) } + + /// Evaluate this single chunk through the multimodal helper. + /// + /// Mirrors `MtmdInputChunks::eval_chunks` but for one chunk at a time, so + /// callers can interleave per-chunk decode with per-chunk bookkeeping + /// (token counting, marker state-machine replay) inside one loop instead + /// of running the helper-level all-chunks eval and a separate ingest pass. + /// + /// Image chunks are decoded as one `llama_decode` call inside the helper, + /// so their token count must fit in `n_batch`. When it would not, the + /// binding refuses the call up front because the C-side + /// `GGML_ASSERT(n_tokens_all <= cparams.n_batch)` would otherwise abort + /// the process. + /// + /// # Errors + /// + /// Returns [`MtmdEvalError::ImageChunkExceedsBatchSize`] when this is an + /// image chunk whose token count exceeds `n_batch`. Returns + /// [`MtmdEvalError::EvalFailure`] if the underlying encode or decode step + /// fails. + pub fn eval_single( + &self, + mtmd_ctx: &MtmdContext, + llama_ctx: &LlamaContext, + start_position: llama_cpp_bindings_sys::llama_pos, + seq_id: llama_cpp_bindings_sys::llama_seq_id, + n_batch: i32, + logits_last: bool, + ) -> Result { + let chunk_token_count = self.n_tokens(); + + if matches!(self.chunk_type(), Ok(MtmdInputChunkType::Image)) + && i64::try_from(chunk_token_count).is_ok_and(|tokens| tokens > i64::from(n_batch)) + { + #[expect( + clippy::cast_possible_truncation, + clippy::cast_sign_loss, + reason = "image token counts and n_batch are model-bounded and fit in u32" + )] + return Err(MtmdEvalError::ImageChunkExceedsBatchSize( + ImageChunkBatchSizeMismatch { + image_tokens: chunk_token_count as u32, + n_batch: n_batch as u32, + }, + )); + } + + let mut final_position: llama_cpp_bindings_sys::llama_pos = start_position; + + let result = unsafe { + llama_cpp_bindings_sys::mtmd_helper_eval_chunk_single( + mtmd_ctx.context.as_ptr(), + llama_ctx.context.as_ptr(), + self.chunk.as_ptr(), + start_position, + seq_id, + n_batch, + logits_last, + &raw mut final_position, + ) + }; + + if result == 0 { + Ok(final_position) + } else { + Err(MtmdEvalError::EvalFailure(result)) + } + } } impl Drop for MtmdInputChunk { diff --git a/llama-cpp-bindings/src/mtmd/mtmd_input_chunks.rs b/llama-cpp-bindings/src/mtmd/mtmd_input_chunks.rs index cc564a39..d9b3a9d8 100644 --- a/llama-cpp-bindings/src/mtmd/mtmd_input_chunks.rs +++ b/llama-cpp-bindings/src/mtmd/mtmd_input_chunks.rs @@ -99,7 +99,7 @@ impl MtmdInputChunks { &self, mtmd_ctx: &MtmdContext, llama_ctx: &LlamaContext, - n_past: llama_cpp_bindings_sys::llama_pos, + start_position: llama_cpp_bindings_sys::llama_pos, seq_id: llama_cpp_bindings_sys::llama_seq_id, n_batch: i32, logits_last: bool, @@ -113,24 +113,29 @@ impl MtmdInputChunks { }); } - let mut new_n_past: llama_cpp_bindings_sys::llama_pos = 0; + // mtmd_helper_eval_chunks overwrites `*new_n_past` at the end of its + // chunk loop (mtmd-helper.cpp:413), so any seed would be fine — but + // we mirror the per-chunk wrapper's `start_position` / `final_position` + // shape here for parity, keeping the read-only input and write-only + // output strictly separated. + let mut final_position: llama_cpp_bindings_sys::llama_pos = start_position; let result = unsafe { llama_cpp_bindings_sys::mtmd_helper_eval_chunks( mtmd_ctx.context.as_ptr(), llama_ctx.context.as_ptr(), self.chunks.as_ptr(), - n_past, + start_position, seq_id, n_batch, logits_last, - &raw mut new_n_past, + &raw mut final_position, ) }; check_eval_result(result)?; - Ok(new_n_past) + Ok(final_position) } } diff --git a/llama-cpp-bindings/src/raw_chat_message.rs b/llama-cpp-bindings/src/raw_chat_message.rs new file mode 100644 index 00000000..ad3cc4a5 --- /dev/null +++ b/llama-cpp-bindings/src/raw_chat_message.rs @@ -0,0 +1,26 @@ +pub struct RawChatMessage { + pub tools_json: String, + pub text: String, + pub is_partial: bool, + pub ffi_error_message: String, +} + +#[cfg(test)] +mod tests { + use super::RawChatMessage; + + #[test] + fn carries_tools_json_text_partial_flag_and_ffi_error_message() { + let raw = RawChatMessage { + tools_json: "[]".to_owned(), + text: "hello".to_owned(), + is_partial: true, + ffi_error_message: "parser bailed".to_owned(), + }; + + assert_eq!(raw.tools_json, "[]"); + assert_eq!(raw.text, "hello"); + assert!(raw.is_partial); + assert_eq!(raw.ffi_error_message, "parser bailed"); + } +} diff --git a/llama-cpp-bindings/src/reasoning_token_classifier.rs b/llama-cpp-bindings/src/reasoning_token_classifier.rs deleted file mode 100644 index 3c14741a..00000000 --- a/llama-cpp-bindings/src/reasoning_token_classifier.rs +++ /dev/null @@ -1,647 +0,0 @@ -use llama_cpp_bindings_sys::llama_pos; -use llama_cpp_bindings_sys::llama_seq_id; - -use crate::context::LlamaContext; -use crate::error::EvalMultimodalChunksError; -use crate::error::SampleError; -use crate::error::TokenUsageError; -use crate::llama_batch::BatchAddError; -use crate::llama_batch::LlamaBatch; -use crate::mtmd::MtmdContext; -use crate::mtmd::MtmdInputChunkType; -use crate::mtmd::MtmdInputChunks; -use crate::sampled_token::SampledToken; -use crate::sampling::LlamaSampler; -use crate::token::LlamaToken; -use crate::token_usage::TokenUsage; - -#[derive(Copy, Clone, Debug, Eq, PartialEq)] -struct ReasoningBoundary { - open: LlamaToken, - close: LlamaToken, -} - -#[derive(Copy, Clone, Debug, Eq, PartialEq)] -pub struct ReasoningTokenClassifier { - boundary: Option, - in_reasoning: bool, - pending_prompt_tokens: u64, - usage: TokenUsage, -} - -impl ReasoningTokenClassifier { - #[must_use] - pub const fn new(open_token: LlamaToken, close_token: LlamaToken) -> Self { - Self { - boundary: Some(ReasoningBoundary { - open: open_token, - close: close_token, - }), - in_reasoning: false, - pending_prompt_tokens: 0, - usage: TokenUsage::new(), - } - } - - #[must_use] - pub const fn undetermined() -> Self { - Self { - boundary: None, - in_reasoning: false, - pending_prompt_tokens: 0, - usage: TokenUsage::new(), - } - } - - pub fn ingest(&mut self, token: LlamaToken) -> SampledToken { - let Some(boundary) = self.boundary else { - self.usage.record_undeterminable_token(); - - return SampledToken::Undeterminable(token); - }; - - if self.in_reasoning { - if token == boundary.close { - self.in_reasoning = false; - } - self.usage.record_reasoning_token(); - - SampledToken::Reasoning(token) - } else if token == boundary.open { - self.in_reasoning = true; - self.usage.record_reasoning_token(); - - SampledToken::Reasoning(token) - } else { - self.usage.record_content_token(); - - SampledToken::Content(token) - } - } - - /// # Errors - /// Forwards [`LlamaSampler::sample`] errors verbatim. Nothing is recorded on failure. - pub fn sample( - &mut self, - sampler: &mut LlamaSampler, - context: &LlamaContext, - idx: i32, - ) -> Result { - let raw = sampler.sample(context, idx)?; - - Ok(self.ingest(raw)) - } - - /// # Errors - /// Forwards [`LlamaBatch::add`] errors verbatim. Nothing is staged on failure. - pub fn feed_prompt_to_batch( - &mut self, - batch: &mut LlamaBatch, - token: LlamaToken, - position: llama_pos, - seq_ids: &[llama_seq_id], - logits: bool, - ) -> Result<(), BatchAddError> { - batch.add(&SampledToken::Content(token), position, seq_ids, logits)?; - self.pending_prompt_tokens = self.pending_prompt_tokens.saturating_add(1); - - Ok(()) - } - - /// # Errors - /// Forwards [`LlamaBatch::add_sequence`] errors verbatim. Nothing is staged on failure. - pub fn feed_prompt_sequence_to_batch( - &mut self, - batch: &mut LlamaBatch, - tokens: &[LlamaToken], - seq_id: llama_seq_id, - logits_all: bool, - ) -> Result<(), BatchAddError> { - batch.add_sequence(tokens, seq_id, logits_all)?; - self.pending_prompt_tokens = self - .pending_prompt_tokens - .saturating_add(tokens.len() as u64); - - Ok(()) - } - - pub const fn commit_prompt_tokens(&mut self) -> u64 { - let promoted = self.pending_prompt_tokens; - self.usage.record_prompt_tokens(promoted); - self.pending_prompt_tokens = 0; - - promoted - } - - pub const fn discard_pending_prompt_tokens(&mut self) -> u64 { - let discarded = self.pending_prompt_tokens; - self.pending_prompt_tokens = 0; - - discarded - } - - #[must_use] - pub const fn pending_prompt_tokens(&self) -> u64 { - self.pending_prompt_tokens - } - - /// # Errors - /// Returns [`EvalMultimodalChunksError::EvalFailed`] when the underlying - /// `eval_chunks` call fails (no counters move), - /// [`EvalMultimodalChunksError::UnknownChunkType`] when a chunk reports a - /// type unknown to this binding, or - /// [`EvalMultimodalChunksError::ChunkOutOfBounds`] when a valid index returns - /// `None` from `chunks.get`. - #[expect( - clippy::too_many_arguments, - reason = "thin wrapper over MtmdInputChunks::eval_chunks; parameter shape mirrors the underlying API" - )] - pub fn eval_multimodal_chunks( - &mut self, - chunks: &MtmdInputChunks, - mtmd_ctx: &MtmdContext, - llama_ctx: &LlamaContext, - n_past: llama_pos, - seq_id: llama_seq_id, - n_batch: i32, - logits_last: bool, - ) -> Result { - let n_past_after = - chunks.eval_chunks(mtmd_ctx, llama_ctx, n_past, seq_id, n_batch, logits_last)?; - - for index in 0..chunks.len() { - let chunk = chunks - .get(index) - .ok_or(EvalMultimodalChunksError::ChunkOutOfBounds(index))?; - let n_tokens = chunk.n_tokens() as u64; - match chunk.chunk_type()? { - MtmdInputChunkType::Text => self.usage.record_prompt_tokens(n_tokens), - MtmdInputChunkType::Image => self.usage.record_input_image_tokens(n_tokens), - MtmdInputChunkType::Audio => self.usage.record_input_audio_tokens(n_tokens), - } - } - - Ok(n_past_after) - } - - pub const fn record_prompt_tokens(&mut self, count: u64) { - self.usage.record_prompt_tokens(count); - } - - /// # Errors - /// Forwards [`TokenUsageError::CachedExceedsPrompt`] when the running cached total would - /// exceed the prompt total. - pub const fn record_cached_prompt_tokens(&mut self, count: u64) -> Result<(), TokenUsageError> { - self.usage.record_cached_prompt_tokens(count) - } - - #[must_use] - pub const fn usage(&self) -> &TokenUsage { - &self.usage - } - - #[must_use] - pub const fn into_usage(self) -> TokenUsage { - self.usage - } -} - -#[cfg(test)] -mod tests { - use super::ReasoningTokenClassifier; - use crate::error::TokenUsageError; - use crate::llama_batch::LlamaBatch; - use crate::sampled_token::SampledToken; - use crate::token::LlamaToken; - use crate::token_usage::TokenUsage; - - const OPEN: LlamaToken = LlamaToken::new(100); - const CLOSE: LlamaToken = LlamaToken::new(200); - - fn fresh_classifier() -> ReasoningTokenClassifier { - ReasoningTokenClassifier::new(OPEN, CLOSE) - } - - #[test] - fn content_token_outside_reasoning_classified_as_content() { - let mut classifier = fresh_classifier(); - let token = LlamaToken::new(1); - - assert_eq!(classifier.ingest(token), SampledToken::Content(token)); - } - - #[test] - fn open_token_emits_reasoning_and_enters_reasoning_state() { - let mut classifier = fresh_classifier(); - - assert_eq!(classifier.ingest(OPEN), SampledToken::Reasoning(OPEN)); - let after_open = LlamaToken::new(1); - assert_eq!( - classifier.ingest(after_open), - SampledToken::Reasoning(after_open) - ); - } - - #[test] - fn token_inside_reasoning_classified_as_reasoning() { - let mut classifier = fresh_classifier(); - classifier.ingest(OPEN); - let inner = LlamaToken::new(42); - - assert_eq!(classifier.ingest(inner), SampledToken::Reasoning(inner)); - } - - #[test] - fn close_token_emits_reasoning_and_exits_reasoning_state() { - let mut classifier = fresh_classifier(); - classifier.ingest(OPEN); - - assert_eq!(classifier.ingest(CLOSE), SampledToken::Reasoning(CLOSE)); - let after_close = LlamaToken::new(7); - assert_eq!( - classifier.ingest(after_close), - SampledToken::Content(after_close) - ); - } - - #[test] - fn token_after_close_classified_as_content() { - let mut classifier = fresh_classifier(); - classifier.ingest(OPEN); - classifier.ingest(LlamaToken::new(5)); - classifier.ingest(CLOSE); - let after = LlamaToken::new(9); - - assert_eq!(classifier.ingest(after), SampledToken::Content(after)); - } - - #[test] - fn multiple_reasoning_blocks_alternate_correctly() { - let mut classifier = fresh_classifier(); - let regular = LlamaToken::new(1); - let inner = LlamaToken::new(2); - - assert_eq!(classifier.ingest(regular), SampledToken::Content(regular)); - assert_eq!(classifier.ingest(OPEN), SampledToken::Reasoning(OPEN)); - assert_eq!(classifier.ingest(inner), SampledToken::Reasoning(inner)); - assert_eq!(classifier.ingest(CLOSE), SampledToken::Reasoning(CLOSE)); - assert_eq!(classifier.ingest(regular), SampledToken::Content(regular)); - assert_eq!(classifier.ingest(OPEN), SampledToken::Reasoning(OPEN)); - assert_eq!(classifier.ingest(CLOSE), SampledToken::Reasoning(CLOSE)); - assert_eq!(classifier.ingest(regular), SampledToken::Content(regular)); - } - - #[test] - fn close_token_outside_reasoning_classified_as_content() { - let mut classifier = fresh_classifier(); - - assert_eq!(classifier.ingest(CLOSE), SampledToken::Content(CLOSE)); - let next = LlamaToken::new(3); - assert_eq!(classifier.ingest(next), SampledToken::Content(next)); - } - - #[test] - fn open_token_while_already_in_reasoning_stays_in_reasoning() { - let mut classifier = fresh_classifier(); - classifier.ingest(OPEN); - - assert_eq!(classifier.ingest(OPEN), SampledToken::Reasoning(OPEN)); - let inner = LlamaToken::new(4); - assert_eq!(classifier.ingest(inner), SampledToken::Reasoning(inner)); - } - - #[test] - fn undetermined_classifier_emits_undeterminable_for_every_input() { - let mut classifier = ReasoningTokenClassifier::undetermined(); - - assert_eq!(classifier.ingest(OPEN), SampledToken::Undeterminable(OPEN)); - assert_eq!( - classifier.ingest(CLOSE), - SampledToken::Undeterminable(CLOSE) - ); - let other = LlamaToken::new(7); - assert_eq!( - classifier.ingest(other), - SampledToken::Undeterminable(other) - ); - } - - #[test] - fn usage_starts_at_default_for_fresh_classifier() { - assert_eq!(*fresh_classifier().usage(), TokenUsage::default()); - assert_eq!( - *ReasoningTokenClassifier::undetermined().usage(), - TokenUsage::default() - ); - } - - #[test] - fn ingest_records_content_in_usage() { - let mut classifier = fresh_classifier(); - classifier.ingest(LlamaToken::new(1)); - classifier.ingest(LlamaToken::new(2)); - - assert_eq!(classifier.usage().content_tokens(), 2); - assert_eq!(classifier.usage().reasoning_tokens(), 0); - assert_eq!(classifier.usage().undeterminable_tokens(), 0); - } - - #[test] - fn ingest_records_reasoning_in_usage_for_open_token_and_inner() { - let mut classifier = fresh_classifier(); - classifier.ingest(OPEN); - classifier.ingest(LlamaToken::new(5)); - classifier.ingest(LlamaToken::new(6)); - classifier.ingest(CLOSE); - - assert_eq!(classifier.usage().reasoning_tokens(), 4); - assert_eq!(classifier.usage().content_tokens(), 0); - } - - #[test] - fn ingest_records_undeterminable_in_usage_when_no_boundary() { - let mut classifier = ReasoningTokenClassifier::undetermined(); - classifier.ingest(LlamaToken::new(1)); - classifier.ingest(LlamaToken::new(2)); - classifier.ingest(LlamaToken::new(3)); - - assert_eq!(classifier.usage().undeterminable_tokens(), 3); - assert_eq!(classifier.usage().content_tokens(), 0); - assert_eq!(classifier.usage().reasoning_tokens(), 0); - assert_eq!(classifier.usage().completion_tokens(), 0); - } - - #[test] - fn record_prompt_tokens_updates_usage() { - let mut classifier = fresh_classifier(); - classifier.record_prompt_tokens(11); - classifier.record_prompt_tokens(2); - - assert_eq!(classifier.usage().prompt_tokens(), 13); - } - - #[test] - fn record_cached_prompt_tokens_updates_usage() { - let mut classifier = fresh_classifier(); - classifier.record_prompt_tokens(10); - classifier.record_cached_prompt_tokens(4).unwrap(); - - assert_eq!(classifier.usage().cached_prompt_tokens(), 4); - } - - #[test] - fn record_cached_above_prompt_returns_error_in_classifier_too() { - let mut classifier = fresh_classifier(); - classifier.record_prompt_tokens(2); - - let result = classifier.record_cached_prompt_tokens(3); - - assert_eq!( - result, - Err(TokenUsageError::CachedExceedsPrompt { - cached_after: 3, - prompt: 2, - }) - ); - assert_eq!(classifier.usage().cached_prompt_tokens(), 0); - } - - #[test] - fn into_usage_returns_accumulated_counters_and_consumes_classifier() { - let mut classifier = fresh_classifier(); - classifier.record_prompt_tokens(5); - classifier.ingest(LlamaToken::new(1)); - classifier.ingest(OPEN); - classifier.ingest(CLOSE); - - let usage = classifier.into_usage(); - - assert_eq!(usage.prompt_tokens(), 5); - assert_eq!(usage.content_tokens(), 1); - assert_eq!(usage.reasoning_tokens(), 2); - assert_eq!(usage.completion_tokens(), 3); - } - - #[test] - fn feed_prompt_to_batch_stages_one_pending_on_success() { - let mut classifier = fresh_classifier(); - let mut batch = LlamaBatch::new(4, 1).unwrap(); - - classifier - .feed_prompt_to_batch(&mut batch, LlamaToken::new(1), 0, &[0], false) - .unwrap(); - - assert_eq!(classifier.pending_prompt_tokens(), 1); - assert_eq!(classifier.usage().prompt_tokens(), 0); - } - - #[test] - fn feed_prompt_to_batch_does_not_stage_when_batch_rejects() { - let mut classifier = fresh_classifier(); - let mut batch = LlamaBatch::new(1, 1).unwrap(); - classifier - .feed_prompt_to_batch(&mut batch, LlamaToken::new(1), 0, &[0], false) - .unwrap(); - - let rejection = - classifier.feed_prompt_to_batch(&mut batch, LlamaToken::new(2), 1, &[0], false); - - assert!(rejection.is_err()); - assert_eq!(classifier.pending_prompt_tokens(), 1); - } - - #[test] - fn feed_prompt_sequence_to_batch_stages_count_on_success() { - let mut classifier = fresh_classifier(); - let mut batch = LlamaBatch::new(8, 1).unwrap(); - let tokens = [LlamaToken::new(1), LlamaToken::new(2), LlamaToken::new(3)]; - - classifier - .feed_prompt_sequence_to_batch(&mut batch, &tokens, 0, false) - .unwrap(); - - assert_eq!(classifier.pending_prompt_tokens(), 3); - assert_eq!(classifier.usage().prompt_tokens(), 0); - } - - #[test] - fn feed_prompt_sequence_to_batch_does_not_stage_full_count_when_batch_rejects() { - let mut classifier = fresh_classifier(); - let mut batch = LlamaBatch::new(2, 1).unwrap(); - let tokens = [LlamaToken::new(1), LlamaToken::new(2), LlamaToken::new(3)]; - - let rejection = classifier.feed_prompt_sequence_to_batch(&mut batch, &tokens, 0, false); - - assert!(rejection.is_err()); - assert_eq!(classifier.pending_prompt_tokens(), 0); - } - - #[test] - fn pending_prompt_tokens_does_not_contribute_to_prompt_or_completion() { - let mut classifier = fresh_classifier(); - let mut batch = LlamaBatch::new(8, 1).unwrap(); - let tokens = [LlamaToken::new(1), LlamaToken::new(2)]; - classifier - .feed_prompt_sequence_to_batch(&mut batch, &tokens, 0, false) - .unwrap(); - - assert_eq!(classifier.usage().prompt_tokens(), 0); - assert_eq!(classifier.usage().completion_tokens(), 0); - } - - #[test] - fn commit_prompt_tokens_moves_pending_into_committed_prompt_tokens_and_resets_pending() { - let mut classifier = fresh_classifier(); - let mut batch = LlamaBatch::new(8, 1).unwrap(); - let tokens = [LlamaToken::new(1), LlamaToken::new(2), LlamaToken::new(3)]; - classifier - .feed_prompt_sequence_to_batch(&mut batch, &tokens, 0, false) - .unwrap(); - - let promoted = classifier.commit_prompt_tokens(); - - assert_eq!(promoted, 3); - assert_eq!(classifier.pending_prompt_tokens(), 0); - assert_eq!(classifier.usage().prompt_tokens(), 3); - } - - #[test] - fn commit_prompt_tokens_with_no_pending_returns_zero_and_changes_nothing() { - let mut classifier = fresh_classifier(); - - let promoted = classifier.commit_prompt_tokens(); - - assert_eq!(promoted, 0); - assert_eq!(classifier.pending_prompt_tokens(), 0); - assert_eq!(classifier.usage().prompt_tokens(), 0); - } - - #[test] - fn discard_pending_prompt_tokens_resets_pending_without_touching_usage() { - let mut classifier = fresh_classifier(); - let mut batch = LlamaBatch::new(8, 1).unwrap(); - let tokens = [LlamaToken::new(1), LlamaToken::new(2)]; - classifier - .feed_prompt_sequence_to_batch(&mut batch, &tokens, 0, false) - .unwrap(); - - let discarded = classifier.discard_pending_prompt_tokens(); - - assert_eq!(discarded, 2); - assert_eq!(classifier.pending_prompt_tokens(), 0); - assert_eq!(classifier.usage().prompt_tokens(), 0); - } - - #[test] - fn multiple_feed_then_commit_aggregates_correctly() { - let mut classifier = fresh_classifier(); - let mut batch = LlamaBatch::new(16, 1).unwrap(); - classifier - .feed_prompt_to_batch(&mut batch, LlamaToken::new(1), 0, &[0], false) - .unwrap(); - classifier - .feed_prompt_sequence_to_batch( - &mut batch, - &[LlamaToken::new(2), LlamaToken::new(3)], - 1, - false, - ) - .unwrap(); - - let promoted = classifier.commit_prompt_tokens(); - - assert_eq!(promoted, 3); - assert_eq!(classifier.usage().prompt_tokens(), 3); - } - - #[test] - fn multiple_feed_then_discard_drops_everything() { - let mut classifier = fresh_classifier(); - let mut batch = LlamaBatch::new(16, 1).unwrap(); - classifier - .feed_prompt_to_batch(&mut batch, LlamaToken::new(1), 0, &[0], false) - .unwrap(); - classifier - .feed_prompt_sequence_to_batch( - &mut batch, - &[LlamaToken::new(2), LlamaToken::new(3)], - 1, - false, - ) - .unwrap(); - - let discarded = classifier.discard_pending_prompt_tokens(); - - assert_eq!(discarded, 3); - assert_eq!(classifier.usage().prompt_tokens(), 0); - } - - #[test] - fn two_classifiers_sharing_a_batch_track_their_own_pending_and_committed_counts() { - let mut request_a = fresh_classifier(); - let mut request_b = fresh_classifier(); - let mut batch = LlamaBatch::new(16, 2).unwrap(); - - let tokens_a = [LlamaToken::new(1), LlamaToken::new(2), LlamaToken::new(3)]; - let tokens_b = [LlamaToken::new(4), LlamaToken::new(5)]; - - request_a - .feed_prompt_sequence_to_batch(&mut batch, &tokens_a, 0, false) - .unwrap(); - request_b - .feed_prompt_sequence_to_batch(&mut batch, &tokens_b, 1, false) - .unwrap(); - - assert_eq!(request_a.pending_prompt_tokens(), 3); - assert_eq!(request_b.pending_prompt_tokens(), 2); - assert_eq!(request_a.usage().prompt_tokens(), 0); - assert_eq!(request_b.usage().prompt_tokens(), 0); - - request_a.ingest(LlamaToken::new(99)); - - assert_eq!(request_a.usage().content_tokens(), 1); - assert_eq!(request_b.usage().content_tokens(), 0); - - let promoted_a = request_a.commit_prompt_tokens(); - let promoted_b = request_b.commit_prompt_tokens(); - - assert_eq!(promoted_a, 3); - assert_eq!(promoted_b, 2); - assert_eq!(request_a.usage().prompt_tokens(), 3); - assert_eq!(request_b.usage().prompt_tokens(), 2); - assert_eq!(request_a.pending_prompt_tokens(), 0); - assert_eq!(request_b.pending_prompt_tokens(), 0); - } - - #[test] - fn discarding_one_classifier_does_not_affect_another_sharing_the_batch() { - let mut request_a = fresh_classifier(); - let mut request_b = fresh_classifier(); - let mut batch = LlamaBatch::new(16, 2).unwrap(); - - request_a - .feed_prompt_sequence_to_batch( - &mut batch, - &[LlamaToken::new(1), LlamaToken::new(2)], - 0, - false, - ) - .unwrap(); - request_b - .feed_prompt_sequence_to_batch( - &mut batch, - &[LlamaToken::new(3), LlamaToken::new(4), LlamaToken::new(5)], - 1, - false, - ) - .unwrap(); - - let discarded_a = request_a.discard_pending_prompt_tokens(); - let promoted_b = request_b.commit_prompt_tokens(); - - assert_eq!(discarded_a, 2); - assert_eq!(promoted_b, 3); - assert_eq!(request_a.usage().prompt_tokens(), 0); - assert_eq!(request_b.usage().prompt_tokens(), 3); - } -} diff --git a/llama-cpp-bindings/src/resolved_tool_call_markers.rs b/llama-cpp-bindings/src/resolved_tool_call_markers.rs new file mode 100644 index 00000000..ced6510c --- /dev/null +++ b/llama-cpp-bindings/src/resolved_tool_call_markers.rs @@ -0,0 +1,11 @@ +/// Effective tool-call marker strings resolved from either the autoparser +/// output or the per-template override registry. +/// +/// Each side is independently optional because the autoparser may report only +/// one of the two strings, and the override registry may not match the +/// template at all. +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct ResolvedToolCallMarkers { + pub open: Option, + pub close: Option, +} diff --git a/llama-cpp-bindings/src/sampled_token.rs b/llama-cpp-bindings/src/sampled_token.rs index a7afa83e..776ead80 100644 --- a/llama-cpp-bindings/src/sampled_token.rs +++ b/llama-cpp-bindings/src/sampled_token.rs @@ -4,5 +4,6 @@ use crate::token::LlamaToken; pub enum SampledToken { Content(LlamaToken), Reasoning(LlamaToken), + ToolCall(LlamaToken), Undeterminable(LlamaToken), } diff --git a/llama-cpp-bindings/src/sampled_token_classifier.rs b/llama-cpp-bindings/src/sampled_token_classifier.rs new file mode 100644 index 00000000..89c034f2 --- /dev/null +++ b/llama-cpp-bindings/src/sampled_token_classifier.rs @@ -0,0 +1,1638 @@ +use std::collections::VecDeque; + +use llama_cpp_bindings_sys::llama_pos; +use llama_cpp_bindings_sys::llama_seq_id; + +use llama_cpp_bindings_types::TokenUsage; +use llama_cpp_bindings_types::TokenUsageError; + +use crate::batch_add_error::BatchAddError; +use crate::context::LlamaContext; +use crate::error::EvalMultimodalChunksError; +use crate::error::SampleError; +use crate::llama_batch::LlamaBatch; +use crate::model::LlamaModel; +use crate::mtmd::MtmdContext; +use crate::mtmd::MtmdInputChunks; +use crate::sampled_token::SampledToken; +use crate::sampling::LlamaSampler; +use crate::streaming_json_probe::JsonProbeOutcome; +use crate::token::LlamaToken; + +#[derive(Copy, Clone, Debug, Eq, PartialEq)] +pub enum SampledTokenSection { + Pending, + Content, + Reasoning, + ToolCall, +} + +#[derive(Copy, Clone, Debug, Eq, PartialEq)] +enum MarkerKind { + ReasoningOpen, + ReasoningClose, + ToolCallOpen, + ToolCallClose, +} + +/// Tokenized marker sequences (token IDs, not strings). +/// +/// Each marker is a `Vec` of length `>= 1`; absent markers are +/// `None`. Sequence matching at every `ingest()` is by token-ID equality, +/// never by substring scanning of decoded text. +#[derive(Clone, Debug, Default, Eq, PartialEq)] +pub struct StreamingMarkers { + pub reasoning_open: Option>, + pub reasoning_close: Option>, + pub tool_call_open: Option>, + pub tool_call_close: Option>, +} + +impl StreamingMarkers { + const fn has_any(&self) -> bool { + self.reasoning_open.is_some() + || self.reasoning_close.is_some() + || self.tool_call_open.is_some() + || self.tool_call_close.is_some() + } + + fn max_token_len(&self) -> usize { + [ + self.reasoning_open.as_deref(), + self.reasoning_close.as_deref(), + self.tool_call_open.as_deref(), + self.tool_call_close.as_deref(), + ] + .into_iter() + .flatten() + .map(<[LlamaToken]>::len) + .max() + .unwrap_or(0) + } + + fn lookup(&self, kind: MarkerKind) -> Option<&[LlamaToken]> { + match kind { + MarkerKind::ReasoningOpen => self.reasoning_open.as_deref(), + MarkerKind::ReasoningClose => self.reasoning_close.as_deref(), + MarkerKind::ToolCallOpen => self.tool_call_open.as_deref(), + MarkerKind::ToolCallClose => self.tool_call_close.as_deref(), + } + } +} + +#[derive(Clone, Debug)] +pub struct IngestOutcome { + pub sampled_token: SampledToken, + /// Empty when the token is part of a recognised marker boundary; otherwise + /// the decoded UTF-8 piece. Callers should stream `visible_piece` and skip + /// emission when it is empty. + pub visible_piece: String, + /// Always the decoded UTF-8 piece, even for marker-boundary tokens. Useful + /// for accumulating the full raw model output (e.g. for downstream parser + /// cross-checks) without losing marker bytes. + pub raw_piece: String, +} + +#[derive(Clone, Debug)] +struct PendingToken { + token: LlamaToken, + decoded: String, + section: SampledTokenSection, + is_boundary: bool, + is_from_prompt: bool, + is_held_for_probe: bool, +} + +#[derive(Clone, Debug)] +struct JsonProbeState { + held_text: String, +} + +#[derive(Clone, Debug)] +enum ProbeMode { + Idle, + Active(JsonProbeState), +} + +pub struct SampledTokenClassifier<'model> { + model: &'model LlamaModel, + markers: StreamingMarkers, + decoder: encoding_rs::Decoder, + pending: VecDeque, + section: SampledTokenSection, + pending_prompt_tokens: u64, + usage: TokenUsage, + probe_mode: ProbeMode, +} + +impl<'model> SampledTokenClassifier<'model> { + #[must_use] + pub fn new(model: &'model LlamaModel, markers: StreamingMarkers) -> Self { + Self { + model, + markers, + decoder: encoding_rs::UTF_8.new_decoder(), + pending: VecDeque::new(), + section: SampledTokenSection::Pending, + pending_prompt_tokens: 0, + usage: TokenUsage::new(), + probe_mode: ProbeMode::Idle, + } + } + + /// Ingest one sampled token. Returns the outcomes that have finalised this + /// turn — typically a single outcome, occasionally zero (the classifier is + /// holding back tokens that may yet form a marker), or several when a + /// buffered marker prefix diverges and the held-back tokens flush. + /// + /// Each [`IngestOutcome`] carries both the [`SampledToken`] variant for + /// classification and the decoded `visible_piece` for streaming. Marker + /// boundaries get an empty `visible_piece` so their text never reaches + /// user-visible streams. + pub fn ingest(&mut self, token: LlamaToken) -> Vec { + if !self.markers.has_any() { + self.usage.record_undeterminable_token(); + let piece = self.decode(token); + return vec![IngestOutcome { + sampled_token: SampledToken::Undeterminable(token), + visible_piece: piece.clone(), + raw_piece: piece, + }]; + } + + let decoded = self.decode(token); + self.pending.push_back(PendingToken { + token, + decoded: decoded.clone(), + section: self.section, + is_boundary: false, + is_from_prompt: false, + is_held_for_probe: false, + }); + + self.try_consume_marker_at_tail(); + + let probe_was_active = matches!(self.probe_mode, ProbeMode::Active(_)); + let mut outcomes = if probe_was_active && self.section_disengages_probe() { + self.abandon_probe() + } else { + self.update_probe(&decoded) + }; + + outcomes.extend(self.drain_overflow()); + outcomes + } + + const fn section_disengages_probe(&self) -> bool { + matches!( + self.section, + SampledTokenSection::ToolCall | SampledTokenSection::Reasoning + ) + } + + /// Replay one prompt token through the marker state machine so that the + /// section at end-of-prompt reflects the chat template's rendered tail + /// (e.g. for Qwen3.5/3.6 with `enable_thinking=false` the prompt ends with + /// a closed empty `...` block, leaving the section in + /// `Content`; with `enable_thinking=true` it ends inside an open ``, + /// leaving the section in `Reasoning`). + /// + /// Prompt tokens never produce [`IngestOutcome`]s and never increment usage + /// counters — they are not generated content. + pub fn ingest_prompt_token(&mut self, token: LlamaToken) { + if !self.markers.has_any() { + return; + } + + self.pending.push_back(PendingToken { + token, + decoded: String::new(), + section: self.section, + is_boundary: false, + is_from_prompt: true, + is_held_for_probe: false, + }); + + self.try_consume_marker_at_tail(); + self.drain_overflow(); + } + + pub fn ingest_prompt_tokens(&mut self, tokens: &[LlamaToken]) { + if !self.markers.has_any() { + return; + } + for &token in tokens { + self.ingest_prompt_token(token); + } + } + + /// Drain every still-buffered token. Call once at end of generation (EOG) + /// to make sure no decoded text is silently dropped. After `flush()` the + /// classifier behaves as if freshly constructed in terms of buffer state. + pub fn flush(&mut self) -> Vec { + self.probe_mode = ProbeMode::Idle; + let mut outcomes = Vec::with_capacity(self.pending.len()); + while let Some(entry) = self.pending.pop_front() { + if entry.is_from_prompt { + continue; + } + outcomes.push(self.finalize_entry(entry)); + } + outcomes + } + + fn decode(&mut self, token: LlamaToken) -> String { + match self.model.token_to_piece( + &SampledToken::Content(token), + &mut self.decoder, + true, + None, + ) { + Ok(piece) => piece, + Err(detokenize_error) => { + tracing::debug!( + "token_to_piece failed during classification, dropping piece: {detokenize_error}" + ); + String::new() + } + } + } + + fn try_consume_marker_at_tail(&mut self) { + // Probe every marker in every section so the user-visible streams stay + // free of marker text even when the model misbehaves: a stray + // `` / `` / `[/THINK]` while in `Content` is + // suppressed (close markers transition to Content — a no-op when + // already there); a nested `` while in `Reasoning` is also + // suppressed (open markers keep the section in Reasoning). Without + // this, models like Gemma 4 E4B that emit close markers without ever + // opening leak the literal marker text into `content_stream`. + const PROBE_KINDS: &[MarkerKind] = &[ + MarkerKind::ReasoningOpen, + MarkerKind::ReasoningClose, + MarkerKind::ToolCallOpen, + MarkerKind::ToolCallClose, + ]; + + for &kind in PROBE_KINDS { + let Some(marker) = self.markers.lookup(kind) else { + continue; + }; + if marker.is_empty() || self.pending.len() < marker.len() { + continue; + } + let span_start = self.pending.len() - marker.len(); + let matches = self + .pending + .iter() + .skip(span_start) + .zip(marker) + .all(|(entry, marker_token)| entry.token == *marker_token); + if matches { + self.mark_marker_span(span_start, kind); + return; + } + } + } + + fn mark_marker_span(&mut self, span_start: usize, kind: MarkerKind) { + let next_section = match kind { + MarkerKind::ReasoningOpen => SampledTokenSection::Reasoning, + MarkerKind::ReasoningClose | MarkerKind::ToolCallClose => SampledTokenSection::Content, + MarkerKind::ToolCallOpen => SampledTokenSection::ToolCall, + }; + // For open markers, the boundary tokens are classified as the destination + // section — they are the marker itself (`` is part of reasoning, + // `` is part of the tool-call protocol). For close markers, + // the boundary tokens are classified as the section the model was in: + // a normal `` while in `Reasoning` is still reasoning, but a + // spurious `` while in `Content` (e.g. some Gemma variants + // re-emit close markers without ever opening) is just noise in the + // content section — counting it as `Reasoning` would inflate + // `observed_reasoning` and falsely indicate the model thought. + let span_section = match kind { + MarkerKind::ReasoningOpen => SampledTokenSection::Reasoning, + MarkerKind::ToolCallOpen => SampledTokenSection::ToolCall, + MarkerKind::ReasoningClose => { + if self.section == SampledTokenSection::Reasoning { + SampledTokenSection::Reasoning + } else { + SampledTokenSection::Content + } + } + MarkerKind::ToolCallClose => { + if self.section == SampledTokenSection::ToolCall { + SampledTokenSection::ToolCall + } else { + SampledTokenSection::Content + } + } + }; + + for entry in self.pending.iter_mut().skip(span_start) { + entry.is_boundary = true; + entry.section = span_section; + } + + self.section = next_section; + } + + fn drain_overflow(&mut self) -> Vec { + let lookback = self.markers.max_token_len().saturating_sub(1); + let mut outcomes = Vec::new(); + + loop { + let Some(front) = self.pending.front() else { + break; + }; + if front.is_held_for_probe { + break; + } + let probe_held = self + .pending + .iter() + .filter(|entry| entry.is_held_for_probe) + .count(); + let drainable = self.pending.len().saturating_sub(probe_held); + let beyond_lookback = drainable > lookback; + if !front.is_boundary && !beyond_lookback { + break; + } + let Some(entry) = self.pending.pop_front() else { + break; + }; + if entry.is_from_prompt { + continue; + } + outcomes.push(self.finalize_entry(entry)); + } + + outcomes + } + + fn update_probe(&mut self, piece: &str) -> Vec { + let probe_active = matches!(self.probe_mode, ProbeMode::Active(_)); + if !probe_active { + if !self.section_allows_probe_engagement() { + return Vec::new(); + } + if !piece.trim_start().starts_with('{') { + return Vec::new(); + } + if let Some(entry) = self.pending.back_mut() { + entry.is_held_for_probe = true; + } + self.probe_mode = ProbeMode::Active(JsonProbeState { + held_text: piece.to_owned(), + }); + return self.evaluate_probe(); + } + + if let Some(entry) = self.pending.back_mut() { + entry.is_held_for_probe = true; + } + if let ProbeMode::Active(state) = &mut self.probe_mode { + state.held_text.push_str(piece); + } + self.evaluate_probe() + } + + const fn section_allows_probe_engagement(&self) -> bool { + matches!( + self.section, + SampledTokenSection::Content | SampledTokenSection::Pending + ) + } + + fn evaluate_probe(&mut self) -> Vec { + let outcome = match &self.probe_mode { + ProbeMode::Active(state) => JsonProbeOutcome::validate_prefix(&state.held_text), + ProbeMode::Idle => return Vec::new(), + }; + match outcome { + JsonProbeOutcome::StillPossiblyValid => Vec::new(), + JsonProbeOutcome::CompletedValid => self.commit_probe_as_tool_call(), + JsonProbeOutcome::Failed => self.abandon_probe(), + } + } + + fn commit_probe_as_tool_call(&mut self) -> Vec { + if !matches!(self.probe_mode, ProbeMode::Active(_)) { + return Vec::new(); + } + self.probe_mode = ProbeMode::Idle; + self.section = SampledTokenSection::Content; + + let drained: Vec<_> = self.pending.drain(..).collect(); + let mut outcomes = Vec::new(); + for mut entry in drained { + if entry.is_held_for_probe { + entry.section = SampledTokenSection::ToolCall; + entry.is_held_for_probe = false; + if !entry.is_from_prompt { + outcomes.push(self.finalize_entry(entry)); + } + } else { + self.pending.push_back(entry); + } + } + outcomes + } + + fn abandon_probe(&mut self) -> Vec { + if !matches!(self.probe_mode, ProbeMode::Active(_)) { + return Vec::new(); + } + self.probe_mode = ProbeMode::Idle; + + let drained: Vec<_> = self.pending.drain(..).collect(); + let mut outcomes = Vec::new(); + for mut entry in drained { + if entry.is_held_for_probe { + entry.is_held_for_probe = false; + if !entry.is_from_prompt { + outcomes.push(self.finalize_entry(entry)); + } + } else { + self.pending.push_back(entry); + } + } + outcomes + } + + fn finalize_entry(&mut self, entry: PendingToken) -> IngestOutcome { + let section = entry.section; + match section { + SampledTokenSection::Reasoning => self.usage.record_reasoning_token(), + SampledTokenSection::Content => self.usage.record_content_token(), + SampledTokenSection::ToolCall => self.usage.record_tool_call_token(), + SampledTokenSection::Pending => self.usage.record_undeterminable_token(), + } + + let sampled_token = match section { + SampledTokenSection::Reasoning => SampledToken::Reasoning(entry.token), + SampledTokenSection::Content => SampledToken::Content(entry.token), + SampledTokenSection::ToolCall => SampledToken::ToolCall(entry.token), + SampledTokenSection::Pending => SampledToken::Undeterminable(entry.token), + }; + + let visible_piece = if entry.is_boundary { + String::new() + } else { + entry.decoded.clone() + }; + + IngestOutcome { + sampled_token, + visible_piece, + raw_piece: entry.decoded, + } + } + + /// # Errors + /// Forwards [`LlamaSampler::sample`] errors verbatim. Nothing is recorded on failure. + /// + /// Returns the raw sampled token (for downstream `batch.add` / `is_eog_token` + /// calls) alongside the outcomes that finalised this turn — see + /// [`Self::ingest`] for buffering semantics. + pub fn sample( + &mut self, + sampler: &mut LlamaSampler, + context: &LlamaContext, + idx: i32, + ) -> Result<(LlamaToken, Vec), SampleError> { + let raw = sampler.sample(context, idx)?; + let outcomes = self.ingest(raw); + + Ok((raw, outcomes)) + } + + /// # Errors + /// Forwards [`LlamaBatch::add`] errors verbatim. Nothing is staged on failure. + pub fn feed_prompt_to_batch( + &mut self, + batch: &mut LlamaBatch, + token: LlamaToken, + position: llama_pos, + seq_ids: &[llama_seq_id], + logits: bool, + ) -> Result<(), BatchAddError> { + batch.add(&SampledToken::Content(token), position, seq_ids, logits)?; + self.ingest_prompt_token(token); + self.pending_prompt_tokens = self.pending_prompt_tokens.saturating_add(1); + + Ok(()) + } + + /// # Errors + /// Forwards [`LlamaBatch::add_sequence`] errors verbatim. Nothing is staged on failure. + pub fn feed_prompt_sequence_to_batch( + &mut self, + batch: &mut LlamaBatch, + tokens: &[LlamaToken], + seq_id: llama_seq_id, + logits_all: bool, + ) -> Result<(), BatchAddError> { + batch.add_sequence(tokens, seq_id, logits_all)?; + self.ingest_prompt_tokens(tokens); + self.pending_prompt_tokens = self + .pending_prompt_tokens + .saturating_add(tokens.len() as u64); + + Ok(()) + } + + pub const fn commit_prompt_tokens(&mut self) -> u64 { + let promoted = self.pending_prompt_tokens; + self.usage.record_prompt_tokens(promoted); + self.pending_prompt_tokens = 0; + + promoted + } + + pub const fn discard_pending_prompt_tokens(&mut self) -> u64 { + let discarded = self.pending_prompt_tokens; + self.pending_prompt_tokens = 0; + + discarded + } + + #[must_use] + pub const fn pending_prompt_tokens(&self) -> u64 { + self.pending_prompt_tokens + } + + /// # Errors + /// Returns [`EvalMultimodalChunksError::EvalFailed`] when the underlying + /// `eval_chunks` call fails (no counters move), + /// [`EvalMultimodalChunksError::UnknownChunkType`] when a chunk reports a + /// type unknown to this binding, or + /// [`EvalMultimodalChunksError::ChunkOutOfBounds`] when a valid index returns + /// `None` from `chunks.get`. + #[expect( + clippy::too_many_arguments, + reason = "thin wrapper over MtmdInputChunks::eval_chunks; parameter shape mirrors the underlying API" + )] + pub fn eval_multimodal_chunks( + &mut self, + chunks: &MtmdInputChunks, + mtmd_ctx: &MtmdContext, + llama_ctx: &LlamaContext, + start_position: llama_pos, + seq_id: llama_seq_id, + n_batch: i32, + logits_last: bool, + ) -> Result { + let chunk_count = chunks.len(); + // `start_position` stays read-only; `next_position` is the loop + // accumulator that walks forward chunk-by-chunk and is the function's + // return value. Two locals, single responsibility each. + let mut next_position = start_position; + + for index in 0..chunk_count { + let chunk = chunks + .get(index) + .ok_or(EvalMultimodalChunksError::ChunkOutOfBounds(index))?; + let logits_for_this_chunk = logits_last && index + 1 == chunk_count; + + next_position = chunk.eval_single( + mtmd_ctx, + llama_ctx, + next_position, + seq_id, + n_batch, + logits_for_this_chunk, + )?; + crate::ingest_prompt_chunk::ingest_prompt_chunk(self, &chunk)?; + } + + Ok(next_position) + } + + pub const fn record_prompt_tokens(&mut self, count: u64) { + self.usage.record_prompt_tokens(count); + } + + pub const fn record_input_image_tokens(&mut self, count: u64) { + self.usage.record_input_image_tokens(count); + } + + pub const fn record_input_audio_tokens(&mut self, count: u64) { + self.usage.record_input_audio_tokens(count); + } + + /// # Errors + /// Forwards [`TokenUsageError::CachedExceedsPrompt`] when the running cached total would + /// exceed the prompt total. + pub const fn record_cached_prompt_tokens(&mut self, count: u64) -> Result<(), TokenUsageError> { + self.usage.record_cached_prompt_tokens(count) + } + + #[must_use] + pub const fn usage(&self) -> &TokenUsage { + &self.usage + } + + #[must_use] + pub fn into_usage(self) -> TokenUsage { + self.usage + } + + #[must_use] + pub const fn current_section(&self) -> SampledTokenSection { + self.section + } + + #[must_use] + pub const fn markers(&self) -> &StreamingMarkers { + &self.markers + } +} + +#[cfg(test)] +mod tests { + use super::IngestOutcome; + use super::PendingToken; + use super::ProbeMode; + use super::SampledTokenClassifier; + use super::SampledTokenSection; + use super::StreamingMarkers; + use crate::sampled_token::SampledToken; + use crate::token::LlamaToken; + + fn token(id: i32) -> LlamaToken { + LlamaToken::new(id) + } + + fn markers_with( + reasoning_open: Option>, + reasoning_close: Option>, + ) -> StreamingMarkers { + StreamingMarkers { + reasoning_open, + reasoning_close, + tool_call_open: None, + tool_call_close: None, + } + } + + /// Builds a classifier without a real model — only safe for tests that go + /// through `try_consume_marker_at_tail` / `drain_overflow` directly, never + /// through `ingest()` (which calls `model.token_to_piece`). + fn synthetic_classifier(markers: StreamingMarkers) -> SampledTokenClassifier<'static> { + SampledTokenClassifier { + model: unsafe { &*std::ptr::NonNull::::dangling().as_ptr() }, + markers, + decoder: encoding_rs::UTF_8.new_decoder(), + pending: std::collections::VecDeque::new(), + section: SampledTokenSection::Pending, + pending_prompt_tokens: 0, + usage: llama_cpp_bindings_types::TokenUsage::new(), + probe_mode: ProbeMode::Idle, + } + } + + fn push_pending(classifier: &mut SampledTokenClassifier<'_>, token_id: i32, decoded: &str) { + classifier.pending.push_back(PendingToken { + token: token(token_id), + decoded: decoded.to_owned(), + section: classifier.section, + is_boundary: false, + is_from_prompt: false, + is_held_for_probe: false, + }); + } + + fn push_pending_from_prompt(classifier: &mut SampledTokenClassifier<'_>, token_id: i32) { + classifier.pending.push_back(PendingToken { + token: token(token_id), + decoded: String::new(), + section: classifier.section, + is_boundary: false, + is_from_prompt: true, + is_held_for_probe: false, + }); + } + + fn push_and_probe( + classifier: &mut SampledTokenClassifier<'_>, + token_id: i32, + decoded: &str, + ) -> Vec { + push_pending(classifier, token_id, decoded); + classifier.try_consume_marker_at_tail(); + let probe_was_active = matches!(classifier.probe_mode, ProbeMode::Active(_)); + let mut outcomes = if probe_was_active && classifier.section_disengages_probe() { + classifier.abandon_probe() + } else { + classifier.update_probe(decoded) + }; + outcomes.extend(classifier.drain_overflow()); + outcomes + } + + fn outcome_pieces(outcomes: &[IngestOutcome]) -> Vec<&str> { + outcomes + .iter() + .map(|outcome| outcome.visible_piece.as_str()) + .collect() + } + + fn outcome_sections(outcomes: &[IngestOutcome]) -> Vec { + outcomes + .iter() + .map(|outcome| match outcome.sampled_token { + SampledToken::Reasoning(_) => SampledTokenSection::Reasoning, + SampledToken::Content(_) => SampledTokenSection::Content, + SampledToken::ToolCall(_) => SampledTokenSection::ToolCall, + SampledToken::Undeterminable(_) => SampledTokenSection::Pending, + }) + .collect() + } + + #[test] + fn streaming_markers_with_no_markers_reports_none() { + let markers = StreamingMarkers::default(); + assert!(!markers.has_any()); + assert_eq!(markers.max_token_len(), 0); + } + + #[test] + fn streaming_markers_max_token_len_takes_longest() { + let markers = StreamingMarkers { + reasoning_open: Some(vec![token(1)]), + reasoning_close: Some(vec![token(2), token(3), token(4)]), + tool_call_open: Some(vec![token(5), token(6)]), + tool_call_close: None, + }; + assert_eq!(markers.max_token_len(), 3); + } + + #[test] + fn single_token_close_marker_when_already_in_reasoning_emits_empty_piece_for_marker() { + let markers = markers_with(Some(vec![token(100)]), Some(vec![token(200)])); + let mut classifier = synthetic_classifier(markers); + classifier.section = SampledTokenSection::Reasoning; + + push_pending(&mut classifier, 7, "step"); + classifier.try_consume_marker_at_tail(); + let mut outcomes = classifier.drain_overflow(); + + push_pending(&mut classifier, 200, ""); + classifier.try_consume_marker_at_tail(); + outcomes.extend(classifier.drain_overflow()); + + push_pending(&mut classifier, 9, "Hi"); + classifier.try_consume_marker_at_tail(); + outcomes.extend(classifier.drain_overflow()); + + outcomes.extend(classifier.flush()); + + assert_eq!( + outcome_sections(&outcomes), + vec![ + SampledTokenSection::Reasoning, + SampledTokenSection::Reasoning, + SampledTokenSection::Content, + ], + ); + assert_eq!(outcome_pieces(&outcomes), vec!["step", "", "Hi"]); + assert_eq!(classifier.section, SampledTokenSection::Content); + } + + #[test] + fn multi_token_close_marker_suppresses_every_marker_token() { + let markers = markers_with( + Some(vec![token(100)]), + Some(vec![token(200), token(201), token(202)]), + ); + let mut classifier = synthetic_classifier(markers); + classifier.section = SampledTokenSection::Reasoning; + + let mut outcomes = Vec::new(); + for (id, decoded) in [(7, "r"), (200, ""), (9, "OK")] { + push_pending(&mut classifier, id, decoded); + classifier.try_consume_marker_at_tail(); + outcomes.extend(classifier.drain_overflow()); + } + outcomes.extend(classifier.flush()); + + assert_eq!(outcome_pieces(&outcomes), vec!["r", "", "", "", "OK"]); + assert_eq!(classifier.section, SampledTokenSection::Content); + } + + #[test] + fn marker_prefix_that_diverges_does_not_suppress_buffered_tokens() { + let markers = markers_with( + Some(vec![token(100)]), + Some(vec![token(200), token(201), token(202)]), + ); + let mut classifier = synthetic_classifier(markers); + classifier.section = SampledTokenSection::Reasoning; + + let mut outcomes = Vec::new(); + for (id, decoded) in [(7, "r"), (200, "a"), (201, "b"), (300, "x")] { + push_pending(&mut classifier, id, decoded); + classifier.try_consume_marker_at_tail(); + outcomes.extend(classifier.drain_overflow()); + } + outcomes.extend(classifier.flush()); + + assert_eq!(outcome_pieces(&outcomes), vec!["r", "a", "b", "x"]); + assert!( + outcomes + .iter() + .all(|outcome| matches!(outcome.sampled_token, SampledToken::Reasoning(_))) + ); + assert_eq!(classifier.section, SampledTokenSection::Reasoning); + } + + #[test] + fn open_then_close_back_to_back_emits_two_empty_pieces_around_zero_content() { + let markers = markers_with(Some(vec![token(100)]), Some(vec![token(200)])); + let mut classifier = synthetic_classifier(markers); + classifier.section = SampledTokenSection::Content; + + let mut outcomes = Vec::new(); + for (id, decoded) in [(100, ""), (200, ""), (9, "Hi")] { + push_pending(&mut classifier, id, decoded); + classifier.try_consume_marker_at_tail(); + outcomes.extend(classifier.drain_overflow()); + } + outcomes.extend(classifier.flush()); + + assert_eq!( + outcome_sections(&outcomes), + vec![ + SampledTokenSection::Reasoning, + SampledTokenSection::Reasoning, + SampledTokenSection::Content, + ], + ); + assert_eq!(outcome_pieces(&outcomes), vec!["", "", "Hi"]); + assert_eq!(classifier.section, SampledTokenSection::Content); + } + + #[test] + fn spurious_reasoning_close_in_content_section_classifies_as_content() { + let markers = markers_with(Some(vec![token(100)]), Some(vec![token(200)])); + let mut classifier = synthetic_classifier(markers); + classifier.section = SampledTokenSection::Content; + + push_pending(&mut classifier, 200, ""); + classifier.try_consume_marker_at_tail(); + let outcomes = classifier.drain_overflow(); + + assert_eq!( + outcome_sections(&outcomes), + vec![SampledTokenSection::Content], + ); + assert_eq!(classifier.section, SampledTokenSection::Content); + } + + #[test] + fn spurious_tool_call_close_in_reasoning_section_classifies_as_tool_call() { + let markers = StreamingMarkers { + reasoning_open: Some(vec![token(100)]), + reasoning_close: Some(vec![token(200)]), + tool_call_open: Some(vec![token(300)]), + tool_call_close: Some(vec![token(400)]), + }; + let mut classifier = synthetic_classifier(markers); + classifier.section = SampledTokenSection::ToolCall; + + push_pending(&mut classifier, 400, ""); + classifier.try_consume_marker_at_tail(); + let outcomes = classifier.drain_overflow(); + + assert_eq!( + outcome_sections(&outcomes), + vec![SampledTokenSection::ToolCall], + ); + assert_eq!(classifier.section, SampledTokenSection::Content); + } + + #[test] + fn flush_drains_remaining_pending_at_eog() { + let markers = markers_with( + Some(vec![token(100)]), + Some(vec![token(200), token(201), token(202)]), + ); + let mut classifier = synthetic_classifier(markers); + classifier.section = SampledTokenSection::Reasoning; + + push_pending(&mut classifier, 7, "abc"); + push_pending(&mut classifier, 200, "".to_owned(), + section: classifier.section, + is_boundary: false, + is_from_prompt: false, + is_held_for_probe: false, + }); + classifier.try_consume_marker_at_tail(); + let outcomes = classifier.drain_overflow(); + + assert_eq!(outcomes.len(), 1); + assert!(matches!( + outcomes[0].sampled_token, + SampledToken::Reasoning(_) + )); + assert_eq!(outcomes[0].visible_piece, ""); + assert_eq!(outcomes[0].raw_piece, "k>"); + + assert_eq!(classifier.section, SampledTokenSection::Content); + assert_eq!(classifier.usage().reasoning_tokens, 1); + assert_eq!(classifier.usage().content_tokens, 0); + } + + #[test] + fn ingest_prompt_tokens_with_multiple_round_trips_ends_in_content() { + let markers = markers_with(Some(vec![token(100)]), Some(vec![token(200)])); + let mut classifier = synthetic_classifier(markers); + + // body body + for token_id in [100, 7, 200, 100, 8, 200] { + push_pending_from_prompt(&mut classifier, token_id); + classifier.try_consume_marker_at_tail(); + classifier.drain_overflow(); + } + + assert_eq!(classifier.section, SampledTokenSection::Content); + assert_eq!(classifier.usage().reasoning_tokens, 0); + assert_eq!(classifier.usage().content_tokens, 0); + assert_eq!(classifier.usage().tool_call_tokens, 0); + assert_eq!(classifier.usage().undeterminable_tokens, 0); + } + + #[test] + fn ingest_prompt_tokens_initial_section_is_always_pending() { + let markers = markers_with(Some(vec![token(100)]), Some(vec![token(200)])); + let classifier = synthetic_classifier(markers); + + assert_eq!(classifier.section, SampledTokenSection::Pending); + } + + #[test] + fn ingest_prompt_tokens_then_drain_for_generated_token_classifies_correctly() { + let markers = markers_with(Some(vec![token(100)]), Some(vec![token(200)])); + let mut classifier = synthetic_classifier(markers); + + // Closed-think prompt: body + for token_id in [100, 7, 200] { + push_pending_from_prompt(&mut classifier, token_id); + classifier.try_consume_marker_at_tail(); + classifier.drain_overflow(); + } + + assert_eq!(classifier.section, SampledTokenSection::Content); + assert_eq!(classifier.usage().reasoning_tokens, 0); + assert_eq!(classifier.usage().content_tokens, 0); + + // Generated content token (not from prompt): pushed with section=Content, + // is_from_prompt=false. drain_overflow finalises it as SampledToken::Content + // and increments usage.content_tokens. + classifier.pending.push_back(PendingToken { + token: token(50), + decoded: "hi".to_owned(), + section: classifier.section, + is_boundary: false, + is_from_prompt: false, + is_held_for_probe: false, + }); + classifier.try_consume_marker_at_tail(); + let outcomes = classifier.drain_overflow(); + + assert_eq!(outcomes.len(), 1); + assert!(matches!( + outcomes[0].sampled_token, + SampledToken::Content(_) + )); + assert_eq!(outcomes[0].visible_piece, "hi"); + assert_eq!(classifier.usage().content_tokens, 1); + assert_eq!(classifier.usage().reasoning_tokens, 0); + assert_eq!(classifier.usage().undeterminable_tokens, 0); + } + + #[test] + fn close_marker_in_content_section_is_suppressed_as_boundary() { + // When a misbehaving model emits a close marker (e.g. ``) while + // already in the Content section, the classifier must treat it as a + // boundary so the marker text never reaches the user-visible content + // stream. The boundary token is classified as Content (not Reasoning): + // there is no reasoning to close, the close marker is just noise in + // the content section. This is the architectural backstop against + // models that re-emit close markers without a preceding open. + let markers = markers_with(Some(vec![token(100)]), Some(vec![token(200)])); + let mut classifier = synthetic_classifier(markers); + classifier.section = SampledTokenSection::Content; + + let mut outcomes = Vec::new(); + for (id, decoded) in [(7, "hi"), (200, ""), (8, "ok")] { + push_pending(&mut classifier, id, decoded); + classifier.try_consume_marker_at_tail(); + outcomes.extend(classifier.drain_overflow()); + } + outcomes.extend(classifier.flush()); + + assert_eq!( + outcome_sections(&outcomes), + vec![ + SampledTokenSection::Content, + SampledTokenSection::Content, + SampledTokenSection::Content, + ], + ); + // The close marker's `visible_piece` is empty (boundary), so the + // user-visible content stream is "hi" + "" + "ok" = "hiok". + assert_eq!(outcome_pieces(&outcomes), vec!["hi", "", "ok"]); + assert_eq!(classifier.section, SampledTokenSection::Content); + } + + #[test] + fn open_marker_in_reasoning_section_is_suppressed_as_boundary() { + // A nested `` while already in Reasoning is suppressed (so the + // user never sees the marker text in the reasoning stream) and the + // section stays Reasoning. + let markers = markers_with(Some(vec![token(100)]), Some(vec![token(200)])); + let mut classifier = synthetic_classifier(markers); + classifier.section = SampledTokenSection::Reasoning; + + let mut outcomes = Vec::new(); + for (id, decoded) in [(7, "step1"), (100, ""), (8, "step2")] { + push_pending(&mut classifier, id, decoded); + classifier.try_consume_marker_at_tail(); + outcomes.extend(classifier.drain_overflow()); + } + outcomes.extend(classifier.flush()); + + assert_eq!(outcome_pieces(&outcomes), vec!["step1", "", "step2"]); + assert_eq!(classifier.section, SampledTokenSection::Reasoning); + } + + #[test] + fn record_prompt_tokens_updates_usage() { + let markers = markers_with(None, None); + let mut classifier = synthetic_classifier(markers); + + classifier.record_prompt_tokens(7); + + assert_eq!(classifier.usage().prompt_tokens, 7); + } + + #[test] + fn record_cached_prompt_tokens_updates_usage_when_under_limit() { + let markers = markers_with(None, None); + let mut classifier = synthetic_classifier(markers); + classifier.record_prompt_tokens(10); + + classifier.record_cached_prompt_tokens(3).unwrap(); + + assert_eq!(classifier.usage().cached_prompt_tokens, 3); + } + + #[test] + fn record_cached_prompt_tokens_returns_error_when_over_prompt_total() { + let markers = markers_with(None, None); + let mut classifier = synthetic_classifier(markers); + classifier.record_prompt_tokens(2); + + let result = classifier.record_cached_prompt_tokens(5); + + assert!(result.is_err()); + } + + #[test] + fn markers_accessor_returns_configured_markers() { + let configured = markers_with(Some(vec![token(1)]), Some(vec![token(2)])); + let classifier = synthetic_classifier(configured); + + let returned = classifier.markers(); + + assert_eq!(returned.reasoning_open.as_deref(), Some(&[token(1)][..])); + assert_eq!(returned.reasoning_close.as_deref(), Some(&[token(2)][..])); + } + + #[test] + fn into_usage_consumes_classifier_and_yields_usage_snapshot() { + let markers = markers_with(None, None); + let mut classifier = synthetic_classifier(markers); + classifier.record_prompt_tokens(11); + + let usage = classifier.into_usage(); + + assert_eq!(usage.prompt_tokens, 11); + } + + #[test] + fn spurious_tool_call_close_in_content_section_classifies_as_content() { + // A `` while in Content (model misbehaves) is classified as + // Content (not ToolCall) so observed_tool_calls isn't inflated. + let mut markers = markers_with(None, None); + markers.tool_call_close = Some(vec![token(300)]); + let mut classifier = synthetic_classifier(markers); + classifier.section = SampledTokenSection::Content; + + push_pending(&mut classifier, 300, ""); + classifier.try_consume_marker_at_tail(); + let outcomes = classifier.drain_overflow(); + + assert_eq!( + outcome_sections(&outcomes), + vec![SampledTokenSection::Content], + ); + assert_eq!(classifier.section, SampledTokenSection::Content); + } + + fn markers_with_tool_call_open(tool_call_open: Vec) -> StreamingMarkers { + StreamingMarkers { + reasoning_open: None, + reasoning_close: None, + tool_call_open: Some(tool_call_open), + tool_call_close: None, + } + } + + fn feed_json_string( + classifier: &mut SampledTokenClassifier<'_>, + text: &str, + starting_token_id: i32, + ) -> Vec { + let mut outcomes = Vec::new(); + for (offset, ch) in text.char_indices() { + let token_id = starting_token_id + i32::try_from(offset).unwrap_or(i32::MAX); + let mut buffer = [0_u8; 4]; + let chunk = ch.encode_utf8(&mut buffer); + outcomes.extend(push_and_probe(classifier, token_id, chunk)); + } + outcomes + } + + #[test] + fn json_probe_engages_when_first_non_whitespace_is_open_brace_in_content() { + let markers = markers_with_tool_call_open(vec![token(900)]); + let mut classifier = synthetic_classifier(markers); + classifier.section = SampledTokenSection::Content; + + push_and_probe(&mut classifier, 1, "{"); + + assert!(matches!(classifier.probe_mode, ProbeMode::Active(_))); + } + + #[test] + fn json_probe_releases_tokens_as_tool_call_when_signature_matches() { + let markers = markers_with_tool_call_open(vec![token(900)]); + let mut classifier = synthetic_classifier(markers); + classifier.section = SampledTokenSection::Content; + + let outcomes = feed_json_string(&mut classifier, r#"{"name":"f","arguments":{}}"#, 100); + + assert!(!outcomes.is_empty()); + assert!( + outcomes + .iter() + .all(|outcome| matches!(outcome.sampled_token, SampledToken::ToolCall(_))), + "every emitted outcome should be ToolCall, got {:?}", + outcome_sections(&outcomes), + ); + assert!(matches!(classifier.probe_mode, ProbeMode::Idle)); + } + + #[test] + fn json_probe_releases_tokens_as_content_when_signature_does_not_match() { + let markers = markers_with_tool_call_open(vec![token(900)]); + let mut classifier = synthetic_classifier(markers); + classifier.section = SampledTokenSection::Content; + + let outcomes = feed_json_string(&mut classifier, r#"{"foo":"bar"}"#, 100); + + assert!( + outcomes + .iter() + .all(|outcome| matches!(outcome.sampled_token, SampledToken::Content(_))), + "every emitted outcome should be Content, got {:?}", + outcome_sections(&outcomes), + ); + assert!(matches!(classifier.probe_mode, ProbeMode::Idle)); + } + + #[test] + fn json_probe_releases_tokens_as_content_when_extra_top_level_key() { + let markers = markers_with_tool_call_open(vec![token(900)]); + let mut classifier = synthetic_classifier(markers); + classifier.section = SampledTokenSection::Content; + + let outcomes = feed_json_string( + &mut classifier, + r#"{"name":"f","arguments":{},"extra":1}"#, + 100, + ); + + assert!( + outcomes + .iter() + .all(|outcome| matches!(outcome.sampled_token, SampledToken::Content(_))), + ); + } + + #[test] + fn json_probe_releases_tokens_as_content_when_arguments_is_not_object() { + let markers = markers_with_tool_call_open(vec![token(900)]); + let mut classifier = synthetic_classifier(markers); + classifier.section = SampledTokenSection::Content; + + let outcomes = feed_json_string(&mut classifier, r#"{"name":"f","arguments":"hi"}"#, 100); + + assert!( + outcomes + .iter() + .all(|outcome| matches!(outcome.sampled_token, SampledToken::Content(_))), + ); + } + + #[test] + fn json_probe_handles_strings_with_quoted_braces_in_arguments() { + let markers = markers_with_tool_call_open(vec![token(900)]); + let mut classifier = synthetic_classifier(markers); + classifier.section = SampledTokenSection::Content; + + let outcomes = feed_json_string( + &mut classifier, + r#"{"name":"f","arguments":{"q":"a } b"}}"#, + 100, + ); + + assert!( + outcomes + .iter() + .all(|outcome| matches!(outcome.sampled_token, SampledToken::ToolCall(_))), + ); + } + + #[test] + fn json_probe_handles_escaped_quotes_in_string_values() { + let markers = markers_with_tool_call_open(vec![token(900)]); + let mut classifier = synthetic_classifier(markers); + classifier.section = SampledTokenSection::Content; + + let outcomes = feed_json_string( + &mut classifier, + r#"{"name":"f","arguments":{"q":"he said \"hi\""}}"#, + 100, + ); + + assert!( + outcomes + .iter() + .all(|outcome| matches!(outcome.sampled_token, SampledToken::ToolCall(_))), + ); + } + + #[test] + fn json_probe_handles_unicode_letters_in_strings() { + let markers = markers_with_tool_call_open(vec![token(900)]); + let mut classifier = synthetic_classifier(markers); + classifier.section = SampledTokenSection::Content; + + let outcomes = feed_json_string( + &mut classifier, + r#"{"name":"日本語","arguments":{"city":"パリ"}}"#, + 100, + ); + + assert!( + outcomes + .iter() + .all(|outcome| matches!(outcome.sampled_token, SampledToken::ToolCall(_))), + ); + } + + #[test] + fn json_probe_handles_nested_objects() { + let markers = markers_with_tool_call_open(vec![token(900)]); + let mut classifier = synthetic_classifier(markers); + classifier.section = SampledTokenSection::Content; + + let outcomes = feed_json_string( + &mut classifier, + r#"{"name":"f","arguments":{"a":{"b":{"c":1}}}}"#, + 100, + ); + + assert!( + outcomes + .iter() + .all(|outcome| matches!(outcome.sampled_token, SampledToken::ToolCall(_))), + ); + } + + #[test] + fn json_probe_handles_arrays_inside_arguments() { + let markers = markers_with_tool_call_open(vec![token(900)]); + let mut classifier = synthetic_classifier(markers); + classifier.section = SampledTokenSection::Content; + + let outcomes = feed_json_string( + &mut classifier, + r#"{"name":"f","arguments":{"items":[1,2,3]}}"#, + 100, + ); + + assert!( + outcomes + .iter() + .all(|outcome| matches!(outcome.sampled_token, SampledToken::ToolCall(_))), + ); + } + + #[test] + fn json_probe_does_not_engage_when_first_byte_is_close_brace() { + let markers = markers_with_tool_call_open(vec![token(900)]); + let mut classifier = synthetic_classifier(markers); + classifier.section = SampledTokenSection::Content; + + let outcomes = feed_json_string(&mut classifier, "}}", 100); + + assert!(matches!(classifier.probe_mode, ProbeMode::Idle)); + assert!( + outcomes + .iter() + .all(|outcome| matches!(outcome.sampled_token, SampledToken::Content(_))), + ); + } + + #[test] + fn json_probe_does_not_engage_in_reasoning_section() { + let markers = StreamingMarkers { + reasoning_open: Some(vec![token(800)]), + reasoning_close: Some(vec![token(801)]), + tool_call_open: Some(vec![token(900)]), + tool_call_close: None, + }; + let mut classifier = synthetic_classifier(markers); + classifier.section = SampledTokenSection::Reasoning; + + push_and_probe(&mut classifier, 1, "{"); + + assert!(matches!(classifier.probe_mode, ProbeMode::Idle)); + } + + #[test] + fn json_probe_does_not_engage_in_tool_call_section() { + let markers = markers_with_tool_call_open(vec![token(900)]); + let mut classifier = synthetic_classifier(markers); + classifier.section = SampledTokenSection::ToolCall; + + push_and_probe(&mut classifier, 1, "{"); + + assert!(matches!(classifier.probe_mode, ProbeMode::Idle)); + } + + #[test] + fn marker_probe_takes_precedence_when_both_could_match() { + // Marker is a single token whose decoded text starts with `"` (a JSON + // signature-valid byte). The JSON probe holds the leading `{`, the + // marker matches at the next token, the section transitions to ToolCall, + // the JSON probe abandons. The leading `{` releases as Content; the + // marker token releases as a ToolCall boundary (suppressed). + let markers = markers_with_tool_call_open(vec![token(900)]); + let mut classifier = synthetic_classifier(markers); + classifier.section = SampledTokenSection::Content; + + let mut outcomes = Vec::new(); + outcomes.extend(push_and_probe(&mut classifier, 1, "{")); + outcomes.extend(push_and_probe(&mut classifier, 900, r#"""#)); + + assert_eq!(classifier.section, SampledTokenSection::ToolCall); + assert_eq!(outcome_pieces(&outcomes), vec!["{", ""]); + assert_eq!( + outcome_sections(&outcomes), + vec![SampledTokenSection::Content, SampledTokenSection::ToolCall], + ); + } + + #[test] + fn json_probe_consumes_two_consecutive_objects_separately() { + let markers = markers_with_tool_call_open(vec![token(900)]); + let mut classifier = synthetic_classifier(markers); + classifier.section = SampledTokenSection::Content; + + let mut outcomes = Vec::new(); + outcomes.extend(feed_json_string( + &mut classifier, + r#"{"name":"a","arguments":{}}"#, + 100, + )); + outcomes.extend(feed_json_string( + &mut classifier, + r#"{"name":"b","arguments":{"x":1}}"#, + 200, + )); + + assert!( + outcomes + .iter() + .all(|outcome| matches!(outcome.sampled_token, SampledToken::ToolCall(_))), + "two consecutive markerless tool calls must both classify as ToolCall, got {:?}", + outcome_sections(&outcomes), + ); + } + + #[test] + fn json_probe_with_leading_whitespace_then_open_brace_classifies_whitespace_as_content_and_json_as_tool_call() + { + let markers = markers_with_tool_call_open(vec![token(900)]); + let mut classifier = synthetic_classifier(markers); + classifier.section = SampledTokenSection::Content; + + let outcomes = feed_json_string( + &mut classifier, + "\n {\"name\":\"f\",\"arguments\":{}}", + 100, + ); + + let tool_call_count = outcomes + .iter() + .filter(|outcome| matches!(outcome.sampled_token, SampledToken::ToolCall(_))) + .count(); + let content_count = outcomes + .iter() + .filter(|outcome| matches!(outcome.sampled_token, SampledToken::Content(_))) + .count(); + assert_eq!( + content_count, 3, + "leading `\\n ` should classify as content" + ); + assert!( + tool_call_count > 0, + "the JSON object should classify as ToolCall", + ); + assert_eq!(content_count + tool_call_count, outcomes.len()); + } + + #[test] + fn json_probe_records_tool_call_token_usage_on_commit() { + let markers = markers_with_tool_call_open(vec![token(900)]); + let mut classifier = synthetic_classifier(markers); + classifier.section = SampledTokenSection::Content; + + let json = r#"{"name":"f","arguments":{}}"#; + let outcomes = feed_json_string(&mut classifier, json, 100); + + let emitted = outcomes.len(); + let usage = classifier.usage(); + assert_eq!(usage.tool_call_tokens, emitted as u64); + assert_eq!(usage.content_tokens, 0); + } + + #[test] + fn json_probe_records_content_token_usage_on_abandon() { + let markers = markers_with_tool_call_open(vec![token(900)]); + let mut classifier = synthetic_classifier(markers); + classifier.section = SampledTokenSection::Content; + + let json = r#"{"foo":"bar"}"#; + let outcomes = feed_json_string(&mut classifier, json, 100); + + let emitted = outcomes.len(); + let usage = classifier.usage(); + assert_eq!(usage.content_tokens, emitted as u64); + assert_eq!(usage.tool_call_tokens, 0); + } + + #[test] + fn flush_during_active_json_probe_releases_held_tokens_as_content() { + let markers = markers_with_tool_call_open(vec![token(900)]); + let mut classifier = synthetic_classifier(markers); + classifier.section = SampledTokenSection::Content; + + push_and_probe(&mut classifier, 1, "{"); + push_and_probe(&mut classifier, 2, r#""name""#); + assert!(matches!(classifier.probe_mode, ProbeMode::Active(_))); + + let outcomes = classifier.flush(); + + assert!( + outcomes + .iter() + .all(|outcome| matches!(outcome.sampled_token, SampledToken::Content(_))), + "mid-probe flush must release held tokens as Content, got {:?}", + outcome_sections(&outcomes), + ); + assert!(matches!(classifier.probe_mode, ProbeMode::Idle)); + } +} diff --git a/llama-cpp-bindings/src/sampling.rs b/llama-cpp-bindings/src/sampling.rs index 81b829b1..e9aadb21 100644 --- a/llama-cpp-bindings/src/sampling.rs +++ b/llama-cpp-bindings/src/sampling.rs @@ -475,7 +475,6 @@ impl LlamaSampler { /// # Errors /// /// Returns [`GrammarError`] if the grammar is invalid or the sampler cannot be initialized. - #[cfg(feature = "llguidance")] pub fn llguidance( model: &LlamaModel, grammar_kind: &str, @@ -522,7 +521,6 @@ impl LlamaSampler { /// /// # Errors /// Returns an error if any string in `seq_breakers` contains null bytes. - #[allow(missing_docs)] pub fn dry( model: &LlamaModel, multiplier: f32, @@ -533,10 +531,12 @@ impl LlamaSampler { ) -> Result { let seq_breakers: Vec = seq_breakers .into_iter() - .map(|s| CString::new(s.as_ref())) + .map(|seq_breaker| CString::new(seq_breaker.as_ref())) .collect::, _>>()?; - let mut seq_breaker_pointers: Vec<*const c_char> = - seq_breakers.iter().map(|s| s.as_ptr()).collect(); + let mut seq_breaker_pointers: Vec<*const c_char> = seq_breakers + .iter() + .map(|seq_breaker| seq_breaker.as_ptr()) + .collect(); let n_ctx_train_value = model.n_ctx_train().map_err(|convert_error| { GrammarError::IntegerOverflow(format!( diff --git a/llama-cpp-bindings/src/streaming_json_probe.rs b/llama-cpp-bindings/src/streaming_json_probe.rs new file mode 100644 index 00000000..388b06fb --- /dev/null +++ b/llama-cpp-bindings/src/streaming_json_probe.rs @@ -0,0 +1,449 @@ +use serde_json::Value; +use serde_json::error::Category; + +const NAME_FIELD: &str = "name"; +const ARGUMENTS_FIELD: &str = "arguments"; + +#[derive(Copy, Clone, Debug, Eq, PartialEq)] +pub enum JsonProbeOutcome { + StillPossiblyValid, + CompletedValid, + Failed, +} + +impl JsonProbeOutcome { + #[must_use] + pub fn validate_prefix(buffer: &str) -> Self { + let trimmed = buffer.trim_start(); + if trimmed.is_empty() { + return Self::StillPossiblyValid; + } + if !trimmed.starts_with('{') { + return Self::Failed; + } + + let mut stream = serde_json::Deserializer::from_str(trimmed).into_iter::(); + match stream.next() { + Some(Ok(value)) => evaluate_completed_value(&value, &trimmed[stream.byte_offset()..]), + Some(Err(parse_error)) => match parse_error.classify() { + Category::Eof => Self::StillPossiblyValid, + Category::Io | Category::Syntax | Category::Data => Self::Failed, + }, + None => Self::StillPossiblyValid, + } + } +} + +fn evaluate_completed_value(value: &Value, trailing: &str) -> JsonProbeOutcome { + let Value::Object(map) = value else { + return JsonProbeOutcome::Failed; + }; + + let Some(Value::String(name)) = map.get(NAME_FIELD) else { + return JsonProbeOutcome::Failed; + }; + if name.is_empty() { + return JsonProbeOutcome::Failed; + } + + if let Some(arguments) = map.get(ARGUMENTS_FIELD) + && !matches!(arguments, Value::Object(_)) + { + return JsonProbeOutcome::Failed; + } + + for key in map.keys() { + if key != NAME_FIELD && key != ARGUMENTS_FIELD { + return JsonProbeOutcome::Failed; + } + } + + if trailing.trim().is_empty() { + JsonProbeOutcome::CompletedValid + } else { + JsonProbeOutcome::Failed + } +} + +#[cfg(test)] +mod tests { + use super::JsonProbeOutcome; + + #[test] + fn empty_buffer_is_still_possibly_valid() { + assert_eq!( + JsonProbeOutcome::validate_prefix(""), + JsonProbeOutcome::StillPossiblyValid, + ); + } + + #[test] + fn whitespace_only_buffer_is_still_possibly_valid() { + assert_eq!( + JsonProbeOutcome::validate_prefix(" \n "), + JsonProbeOutcome::StillPossiblyValid, + ); + } + + #[test] + fn single_open_brace_is_still_possibly_valid() { + assert_eq!( + JsonProbeOutcome::validate_prefix("{"), + JsonProbeOutcome::StillPossiblyValid, + ); + } + + #[test] + fn open_brace_with_trailing_space_is_still_possibly_valid() { + assert_eq!( + JsonProbeOutcome::validate_prefix("{ "), + JsonProbeOutcome::StillPossiblyValid, + ); + } + + #[test] + fn open_brace_with_quote_starting_key_is_still_possibly_valid() { + assert_eq!( + JsonProbeOutcome::validate_prefix(r#"{ ""#), + JsonProbeOutcome::StillPossiblyValid, + ); + } + + #[test] + fn partial_name_key_is_still_possibly_valid() { + assert_eq!( + JsonProbeOutcome::validate_prefix(r#"{ "name""#), + JsonProbeOutcome::StillPossiblyValid, + ); + } + + #[test] + fn partial_name_value_quote_is_still_possibly_valid() { + assert_eq!( + JsonProbeOutcome::validate_prefix(r#"{ "name": ""#), + JsonProbeOutcome::StillPossiblyValid, + ); + } + + #[test] + fn partial_name_value_letters_is_still_possibly_valid() { + assert_eq!( + JsonProbeOutcome::validate_prefix(r#"{ "name": "ge"#), + JsonProbeOutcome::StillPossiblyValid, + ); + } + + #[test] + fn complete_name_string_no_comma_is_still_possibly_valid() { + assert_eq!( + JsonProbeOutcome::validate_prefix(r#"{ "name": "get_weather""#), + JsonProbeOutcome::StillPossiblyValid, + ); + } + + #[test] + fn name_then_comma_is_still_possibly_valid() { + assert_eq!( + JsonProbeOutcome::validate_prefix(r#"{ "name": "get_weather","#), + JsonProbeOutcome::StillPossiblyValid, + ); + } + + #[test] + fn name_then_partial_arguments_key_is_still_possibly_valid() { + assert_eq!( + JsonProbeOutcome::validate_prefix(r#"{ "name": "get_weather", "argum"#), + JsonProbeOutcome::StillPossiblyValid, + ); + } + + #[test] + fn name_then_arguments_key_is_still_possibly_valid() { + assert_eq!( + JsonProbeOutcome::validate_prefix(r#"{ "name": "get_weather", "arguments""#), + JsonProbeOutcome::StillPossiblyValid, + ); + } + + #[test] + fn name_then_arguments_open_brace_is_still_possibly_valid() { + assert_eq!( + JsonProbeOutcome::validate_prefix(r#"{ "name": "get_weather", "arguments": {"#), + JsonProbeOutcome::StillPossiblyValid, + ); + } + + #[test] + fn arguments_with_partial_inner_key_value_is_still_possibly_valid() { + assert_eq!( + JsonProbeOutcome::validate_prefix( + r#"{ "name": "get_weather", "arguments": {"location":"# + ), + JsonProbeOutcome::StillPossiblyValid, + ); + } + + #[test] + fn arguments_with_partial_inner_string_value_is_still_possibly_valid() { + assert_eq!( + JsonProbeOutcome::validate_prefix( + r#"{ "name": "get_weather", "arguments": {"location": "Pa"# + ), + JsonProbeOutcome::StillPossiblyValid, + ); + } + + #[test] + fn complete_simple_tool_call_is_completed_valid() { + assert_eq!( + JsonProbeOutcome::validate_prefix(r#"{"name":"f","arguments":{}}"#), + JsonProbeOutcome::CompletedValid, + ); + } + + #[test] + fn complete_tool_call_with_internal_whitespace_is_completed_valid() { + assert_eq!( + JsonProbeOutcome::validate_prefix(r#"{"name": "f", "arguments": {}}"#), + JsonProbeOutcome::CompletedValid, + ); + } + + #[test] + fn complete_tool_call_with_string_argument_is_completed_valid() { + assert_eq!( + JsonProbeOutcome::validate_prefix( + r#"{"name":"get_weather","arguments":{"location":"Paris"}}"# + ), + JsonProbeOutcome::CompletedValid, + ); + } + + #[test] + fn complete_tool_call_with_multiple_arguments_is_completed_valid() { + assert_eq!( + JsonProbeOutcome::validate_prefix( + r#"{"name":"book_flight","arguments":{"from":"NYC","to":"PAR","passengers":2}}"# + ), + JsonProbeOutcome::CompletedValid, + ); + } + + #[test] + fn complete_tool_call_with_nested_arguments_is_completed_valid() { + assert_eq!( + JsonProbeOutcome::validate_prefix(r#"{"name":"f","arguments":{"a":{"b":[1,2,3]}}}"#), + JsonProbeOutcome::CompletedValid, + ); + } + + #[test] + fn complete_tool_call_with_close_brace_inside_string_is_completed_valid() { + assert_eq!( + JsonProbeOutcome::validate_prefix(r#"{"name":"f","arguments":{"q":"a } b"}}"#), + JsonProbeOutcome::CompletedValid, + ); + } + + #[test] + fn complete_tool_call_with_escaped_quotes_in_string_is_completed_valid() { + assert_eq!( + JsonProbeOutcome::validate_prefix(r#"{"name":"f","arguments":{"q":"he said \"hi\""}}"#), + JsonProbeOutcome::CompletedValid, + ); + } + + #[test] + fn complete_tool_call_with_unicode_strings_is_completed_valid() { + assert_eq!( + JsonProbeOutcome::validate_prefix(r#"{"name":"日本語","arguments":{"city":"パリ"}}"#), + JsonProbeOutcome::CompletedValid, + ); + } + + #[test] + fn complete_tool_call_with_trailing_whitespace_is_completed_valid() { + assert_eq!( + JsonProbeOutcome::validate_prefix("{\"name\":\"f\",\"arguments\":{}}\n"), + JsonProbeOutcome::CompletedValid, + ); + } + + #[test] + fn complete_tool_call_with_array_inside_arguments_is_completed_valid() { + assert_eq!( + JsonProbeOutcome::validate_prefix(r#"{"name":"f","arguments":{"items":[1,2,3]}}"#), + JsonProbeOutcome::CompletedValid, + ); + } + + #[test] + fn complete_tool_call_without_arguments_field_is_completed_valid() { + assert_eq!( + JsonProbeOutcome::validate_prefix(r#"{"name":"ping"}"#), + JsonProbeOutcome::CompletedValid, + ); + } + + #[test] + fn top_level_array_is_failed() { + assert_eq!( + JsonProbeOutcome::validate_prefix("["), + JsonProbeOutcome::Failed + ); + } + + #[test] + fn top_level_scalar_number_is_failed() { + assert_eq!( + JsonProbeOutcome::validate_prefix("123"), + JsonProbeOutcome::Failed + ); + } + + #[test] + fn top_level_string_is_failed() { + assert_eq!( + JsonProbeOutcome::validate_prefix(r#""hi""#), + JsonProbeOutcome::Failed + ); + } + + #[test] + fn complete_object_with_wrong_first_key_is_failed() { + assert_eq!( + JsonProbeOutcome::validate_prefix(r#"{"foo":"bar"}"#), + JsonProbeOutcome::Failed, + ); + } + + #[test] + fn complete_object_with_non_string_name_is_failed() { + assert_eq!( + JsonProbeOutcome::validate_prefix(r#"{"name":123,"arguments":{}}"#), + JsonProbeOutcome::Failed, + ); + } + + #[test] + fn complete_object_with_null_name_is_failed() { + assert_eq!( + JsonProbeOutcome::validate_prefix(r#"{"name":null,"arguments":{}}"#), + JsonProbeOutcome::Failed, + ); + } + + #[test] + fn complete_object_with_arguments_as_array_is_failed() { + assert_eq!( + JsonProbeOutcome::validate_prefix(r#"{"name":"f","arguments":[]}"#), + JsonProbeOutcome::Failed, + ); + } + + #[test] + fn complete_object_with_arguments_as_string_is_failed() { + assert_eq!( + JsonProbeOutcome::validate_prefix(r#"{"name":"f","arguments":"hi"}"#), + JsonProbeOutcome::Failed, + ); + } + + #[test] + fn complete_object_with_third_top_level_key_is_failed() { + assert_eq!( + JsonProbeOutcome::validate_prefix(r#"{"name":"f","arguments":{},"extra":1}"#), + JsonProbeOutcome::Failed, + ); + } + + #[test] + fn complete_object_with_empty_name_is_failed() { + assert_eq!( + JsonProbeOutcome::validate_prefix(r#"{"name":"","arguments":{}}"#), + JsonProbeOutcome::Failed, + ); + } + + #[test] + fn complete_object_with_trailing_garbage_is_failed() { + assert_eq!( + JsonProbeOutcome::validate_prefix(r#"{"name":"f","arguments":{}}garbage"#), + JsonProbeOutcome::Failed, + ); + } + + #[test] + fn empty_object_is_failed_due_to_missing_required_name() { + assert_eq!( + JsonProbeOutcome::validate_prefix("{}"), + JsonProbeOutcome::Failed + ); + } + + #[test] + fn complete_object_with_arguments_only_no_name_is_failed() { + assert_eq!( + JsonProbeOutcome::validate_prefix(r#"{"arguments":{}}"#), + JsonProbeOutcome::Failed, + ); + } + + #[test] + fn leading_whitespace_then_open_brace_is_still_possibly_valid() { + assert_eq!( + JsonProbeOutcome::validate_prefix("\n \n{"), + JsonProbeOutcome::StillPossiblyValid, + ); + } + + #[test] + fn leading_whitespace_then_complete_tool_call_is_completed_valid() { + assert_eq!( + JsonProbeOutcome::validate_prefix("\n {\"name\":\"f\",\"arguments\":{}}"), + JsonProbeOutcome::CompletedValid, + ); + } + + #[test] + fn complete_tool_call_followed_by_second_object_is_failed() { + assert_eq!( + JsonProbeOutcome::validate_prefix( + r#"{"name":"a","arguments":{}}{"name":"b","arguments":{}}"# + ), + JsonProbeOutcome::Failed, + ); + } + + #[test] + fn buffer_with_only_open_quote_is_still_possibly_valid() { + assert_eq!( + JsonProbeOutcome::validate_prefix(r#"{ "n"#), + JsonProbeOutcome::StillPossiblyValid, + ); + } + + #[test] + fn buffer_with_complete_first_field_unknown_second_key_is_failed() { + assert_eq!( + JsonProbeOutcome::validate_prefix(r#"{ "name": "f", "foo": 1}"#), + JsonProbeOutcome::Failed, + ); + } + + #[test] + fn unicode_letter_inside_name_value_completes_validly() { + assert_eq!( + JsonProbeOutcome::validate_prefix(r#"{"name":"éclair","arguments":{}}"#), + JsonProbeOutcome::CompletedValid, + ); + } + + #[test] + fn arguments_field_with_explicit_null_is_failed() { + assert_eq!( + JsonProbeOutcome::validate_prefix(r#"{"name":"f","arguments":null}"#), + JsonProbeOutcome::Failed, + ); + } +} diff --git a/llama-cpp-bindings/src/token/data_array.rs b/llama-cpp-bindings/src/token/data_array.rs index d7dc28c8..af2134df 100644 --- a/llama-cpp-bindings/src/token/data_array.rs +++ b/llama-cpp-bindings/src/token/data_array.rs @@ -93,7 +93,10 @@ impl LlamaTokenDataArray { let mut c_llama_token_data_array = llama_cpp_bindings_sys::llama_token_data_array { data, size, - selected: self.selected.and_then(|s| s.try_into().ok()).unwrap_or(-1), + selected: self + .selected + .and_then(|selected_index| selected_index.try_into().ok()) + .unwrap_or(-1), sorted: self.sorted, }; diff --git a/llama-cpp-bindings/src/tool_call_format/bracketed_args.rs b/llama-cpp-bindings/src/tool_call_format/bracketed_args.rs new file mode 100644 index 00000000..0020c90a --- /dev/null +++ b/llama-cpp-bindings/src/tool_call_format/bracketed_args.rs @@ -0,0 +1,229 @@ +use llama_cpp_bindings_types::BracketedJsonShape; +use llama_cpp_bindings_types::ParsedToolCall; +use llama_cpp_bindings_types::ToolCallArguments; +use llama_cpp_bindings_types::ToolCallMarkers; + +use crate::error::BracketedArgsFailure; + +enum ParseStep<'body> { + Done, + Call(ParsedToolCall, &'body str), +} + +fn consume_optional_prefix<'body>(input: &'body str, literal: &str) -> &'body str { + input.strip_prefix(literal).unwrap_or(input) +} + +fn split_at_separator<'body>( + input: &'body str, + separator: &str, +) -> Option<(&'body str, &'body str)> { + let (name_raw, after_separator) = input.split_once(separator)?; + Some((name_raw, after_separator)) +} + +fn consume_one_json_value<'body>( + input: &'body str, + tool_name: &str, +) -> Result<(serde_json::Value, &'body str), BracketedArgsFailure> { + let mut stream = serde_json::Deserializer::from_str(input).into_iter::(); + let value = stream + .next() + .ok_or_else(|| BracketedArgsFailure::UnterminatedArguments { + tool_name: tool_name.to_owned(), + })? + .map_err(|err| BracketedArgsFailure::InvalidJsonArguments { + tool_name: tool_name.to_owned(), + message: err.to_string(), + })?; + let consumed = stream.byte_offset(); + + Ok((value, &input[consumed..])) +} + +fn parse_one_call<'body>( + input: &'body str, + markers: &ToolCallMarkers, + shape: &BracketedJsonShape, +) -> Result, BracketedArgsFailure> { + if input.is_empty() { + return Ok(ParseStep::Done); + } + + let after_open = consume_optional_prefix(input, markers.open.as_str()); + + let Some((name_raw, after_separator)) = + split_at_separator(after_open, shape.name_args_separator.as_str()) + else { + return Ok(ParseStep::Done); + }; + + let name = name_raw.trim().to_owned(); + if name.is_empty() { + return Ok(ParseStep::Done); + } + + let (arguments_value, after_arguments) = consume_one_json_value(after_separator, &name)?; + + let after_close = consume_optional_prefix(after_arguments, markers.close.as_str()); + + Ok(ParseStep::Call( + ParsedToolCall::new( + String::new(), + name, + ToolCallArguments::ValidJson(arguments_value), + ), + after_close, + )) +} + +/// # Errors +/// +/// Returns [`BracketedArgsFailure`] when the body looks like a bracketed-JSON +/// tool-call block (matches the name/args separator) but contains a structural +/// issue: invalid JSON arguments or a JSON value truncated mid-stream. +pub fn parse( + body: &str, + markers: &ToolCallMarkers, + shape: &BracketedJsonShape, +) -> Result, BracketedArgsFailure> { + if shape.name_args_separator.is_empty() { + return Ok(Vec::new()); + } + + let mut parsed = Vec::new(); + let mut remaining = body.trim_start(); + + loop { + match parse_one_call(remaining, markers, shape)? { + ParseStep::Done => break, + ParseStep::Call(call, rest) => { + parsed.push(call); + remaining = rest.trim_start(); + } + } + } + + Ok(parsed) +} + +#[cfg(test)] +mod tests { + use llama_cpp_bindings_types::BracketedJsonShape; + use llama_cpp_bindings_types::ToolCallArgsShape; + use llama_cpp_bindings_types::ToolCallArguments; + use llama_cpp_bindings_types::ToolCallMarkers; + use serde_json::json; + + use super::parse; + use crate::error::BracketedArgsFailure; + + fn mistral3_markers() -> ToolCallMarkers { + ToolCallMarkers { + open: "[TOOL_CALLS]".to_owned(), + close: String::new(), + args_shape: ToolCallArgsShape::BracketedJson(BracketedJsonShape { + name_args_separator: "[ARGS]".to_owned(), + }), + } + } + + fn mistral3_shape() -> BracketedJsonShape { + BracketedJsonShape { + name_args_separator: "[ARGS]".to_owned(), + } + } + + #[test] + fn parses_single_tool_call_with_open_marker_present() { + let parsed = parse( + "[TOOL_CALLS]get_weather[ARGS]{\"location\":\"Paris\"}", + &mistral3_markers(), + &mistral3_shape(), + ) + .expect("must parse"); + + assert_eq!(parsed.len(), 1); + assert_eq!(parsed[0].name, "get_weather"); + assert_eq!( + parsed[0].arguments, + ToolCallArguments::ValidJson(json!({"location": "Paris"})), + ); + } + + #[test] + fn parses_single_tool_call_when_classifier_stripped_open_marker() { + let parsed = parse( + "get_weather[ARGS]{\"location\":\"Paris\"}", + &mistral3_markers(), + &mistral3_shape(), + ) + .expect("must parse"); + + assert_eq!(parsed.len(), 1); + assert_eq!(parsed[0].name, "get_weather"); + assert_eq!( + parsed[0].arguments, + ToolCallArguments::ValidJson(json!({"location": "Paris"})), + ); + } + + #[test] + fn parses_two_consecutive_tool_calls_with_repeated_open_marker() { + let parsed = parse( + "[TOOL_CALLS]a[ARGS]{\"x\":1}[TOOL_CALLS]b[ARGS]{\"y\":2}", + &mistral3_markers(), + &mistral3_shape(), + ) + .expect("must parse"); + + assert_eq!(parsed.len(), 2); + assert_eq!(parsed[0].name, "a"); + assert_eq!( + parsed[0].arguments, + ToolCallArguments::ValidJson(json!({"x": 1})) + ); + assert_eq!(parsed[1].name, "b"); + assert_eq!( + parsed[1].arguments, + ToolCallArguments::ValidJson(json!({"y": 2})) + ); + } + + #[test] + fn rejects_malformed_json_arguments_with_typed_failure() { + let result = parse( + "[TOOL_CALLS]get_weather[ARGS]{\"location\":}", + &mistral3_markers(), + &mistral3_shape(), + ); + + let failure = result.expect_err("malformed JSON must produce a typed failure"); + match failure { + BracketedArgsFailure::InvalidJsonArguments { tool_name, .. } => { + assert_eq!(tool_name, "get_weather"); + } + other @ BracketedArgsFailure::UnterminatedArguments { .. } => { + panic!("expected InvalidJsonArguments, got {other:?}") + } + } + } + + #[test] + fn returns_empty_vec_for_empty_body() { + let parsed = + parse("", &mistral3_markers(), &mistral3_shape()).expect("empty body must parse"); + assert!(parsed.is_empty()); + } + + #[test] + fn returns_empty_vec_when_body_lacks_separator() { + let parsed = parse( + "plain text without separator", + &mistral3_markers(), + &mistral3_shape(), + ) + .expect("body without separator must parse"); + assert!(parsed.is_empty()); + } +} diff --git a/llama-cpp-bindings/src/tool_call_format/json_object.rs b/llama-cpp-bindings/src/tool_call_format/json_object.rs new file mode 100644 index 00000000..08633d72 --- /dev/null +++ b/llama-cpp-bindings/src/tool_call_format/json_object.rs @@ -0,0 +1,199 @@ +use llama_cpp_bindings_types::JsonObjectShape; +use llama_cpp_bindings_types::ParsedToolCall; +use llama_cpp_bindings_types::ToolCallArguments; + +use crate::error::JsonObjectFailure; + +fn try_parse_one_object( + input: &str, + shape: &JsonObjectShape, +) -> Result, JsonObjectFailure> { + let trimmed_start = input.find('{'); + let Some(start) = trimmed_start else { + return Ok(None); + }; + + let mut stream = + serde_json::Deserializer::from_str(&input[start..]).into_iter::(); + let value = match stream.next() { + Some(Ok(value)) => value, + Some(Err(err)) => { + return Err(JsonObjectFailure::InvalidJson { + message: err.to_string(), + }); + } + None => return Ok(None), + }; + let consumed = stream.byte_offset(); + + let serde_json::Value::Object(map) = value else { + return Ok(None); + }; + + let Some(name_value) = map.get(&shape.name_field) else { + return Ok(None); + }; + let serde_json::Value::String(name) = name_value else { + return Ok(None); + }; + + let arguments_value = map + .get(&shape.arguments_field) + .cloned() + .unwrap_or(serde_json::Value::Object(serde_json::Map::new())); + let arguments = ToolCallArguments::from_string(arguments_value.to_string()); + + let trailing_extras = map + .keys() + .any(|key| key != &shape.name_field && key != &shape.arguments_field); + if trailing_extras { + return Ok(None); + } + + Ok(Some(( + ParsedToolCall::new(String::new(), name.clone(), arguments), + start + consumed, + ))) +} + +/// # Errors +/// +/// Returns [`JsonObjectFailure`] when the body contains a JSON object that +/// looks like a tool call (matches the open brace at start) but the JSON itself +/// is malformed. +pub fn parse( + body: &str, + shape: &JsonObjectShape, +) -> Result, JsonObjectFailure> { + if shape.name_field.is_empty() || shape.arguments_field.is_empty() { + return Ok(Vec::new()); + } + + let mut parsed = Vec::new(); + let mut remaining = body; + + while let Some((call, consumed)) = try_parse_one_object(remaining, shape)? { + parsed.push(call); + remaining = &remaining[consumed..]; + } + + Ok(parsed) +} + +#[cfg(test)] +mod tests { + use llama_cpp_bindings_types::JsonObjectShape; + use llama_cpp_bindings_types::ToolCallArguments; + use serde_json::json; + + use super::parse; + use crate::error::JsonObjectFailure; + + fn qwen3_shape() -> JsonObjectShape { + JsonObjectShape { + name_field: "name".to_owned(), + arguments_field: "arguments".to_owned(), + } + } + + #[test] + fn parses_single_json_object_with_name_and_arguments() { + let parsed = parse( + r#"{"name": "get_weather", "arguments": {"location": "Paris"}}"#, + &qwen3_shape(), + ) + .expect("must parse"); + + assert_eq!(parsed.len(), 1); + assert_eq!(parsed[0].name, "get_weather"); + assert_eq!( + parsed[0].arguments, + ToolCallArguments::ValidJson(json!({"location": "Paris"})), + ); + } + + #[test] + fn parses_json_object_after_leading_whitespace_and_newlines() { + let parsed = parse( + "\n {\"name\": \"f\", \"arguments\": {\"a\": 1}}\n", + &qwen3_shape(), + ) + .expect("must parse"); + + assert_eq!(parsed.len(), 1); + assert_eq!(parsed[0].name, "f"); + } + + #[test] + fn parses_two_consecutive_json_objects() { + let parsed = parse( + r#"{"name": "a", "arguments": {}}{"name": "b", "arguments": {"x": 2}}"#, + &qwen3_shape(), + ) + .expect("must parse"); + + assert_eq!(parsed.len(), 2); + assert_eq!(parsed[0].name, "a"); + assert_eq!(parsed[1].name, "b"); + } + + #[test] + fn parses_object_with_arguments_field_missing_yields_empty_arguments() { + let parsed = parse(r#"{"name": "ping"}"#, &qwen3_shape()).expect("must parse"); + + assert_eq!(parsed.len(), 1); + assert_eq!(parsed[0].name, "ping"); + assert_eq!(parsed[0].arguments, ToolCallArguments::ValidJson(json!({})),); + } + + #[test] + fn rejects_json_object_with_extra_unexpected_top_level_keys() { + let parsed = parse( + r#"{"name": "f", "arguments": {}, "extra": 1}"#, + &qwen3_shape(), + ) + .expect("must parse"); + + assert!(parsed.is_empty(), "extra top-level key must reject"); + } + + #[test] + fn rejects_json_object_with_non_string_name() { + let parsed = + parse(r#"{"name": 123, "arguments": {}}"#, &qwen3_shape()).expect("must parse"); + + assert!(parsed.is_empty(), "non-string name must reject"); + } + + #[test] + fn rejects_input_without_open_brace() { + let parsed = parse("plain content", &qwen3_shape()).expect("must parse"); + assert!(parsed.is_empty()); + } + + #[test] + fn rejects_array_instead_of_object() { + let parsed = parse("[1, 2, 3]", &qwen3_shape()).expect("must parse"); + assert!(parsed.is_empty()); + } + + #[test] + fn returns_failure_for_malformed_json() { + let result = parse(r#"{"name": "f", "arguments": {"a": }"#, &qwen3_shape()); + + match result { + Err(JsonObjectFailure::InvalidJson { message }) => { + assert!(!message.is_empty()); + } + other => panic!("expected InvalidJson, got {other:?}"), + } + } + + #[test] + fn returns_empty_when_shape_has_empty_required_field() { + let mut shape = qwen3_shape(); + shape.name_field.clear(); + let parsed = parse(r#"{"name": "x", "arguments": {}}"#, &shape).expect("must parse"); + assert!(parsed.is_empty()); + } +} diff --git a/llama-cpp-bindings/src/tool_call_format/key_value_xml_tags.rs b/llama-cpp-bindings/src/tool_call_format/key_value_xml_tags.rs new file mode 100644 index 00000000..0ea21787 --- /dev/null +++ b/llama-cpp-bindings/src/tool_call_format/key_value_xml_tags.rs @@ -0,0 +1,330 @@ +use llama_cpp_bindings_types::KeyValueXmlTagsShape; +use llama_cpp_bindings_types::ParsedToolCall; +use llama_cpp_bindings_types::ToolCallArguments; +use llama_cpp_bindings_types::ToolCallMarkers; +use nom::IResult; +use nom::Parser; +use nom::bytes::complete::tag; +use nom::bytes::complete::take_until; + +use crate::error::KeyValueXmlTagsFailure; + +enum ParseStep<'body> { + Done, + Call(ParsedToolCall, &'body str), +} + +const fn shape_is_complete(shape: &KeyValueXmlTagsShape) -> bool { + !shape.key_open.is_empty() + && !shape.key_close.is_empty() + && !shape.value_open.is_empty() + && !shape.value_close.is_empty() +} + +fn skip_to_next_open<'body>(input: &'body str, open: &str) -> Option<&'body str> { + let take_result: IResult<&'body str, &'body str> = take_until(open).parse(input); + let (after_prefix_inclusive, _) = take_result.ok()?; + let consume_result: IResult<&'body str, &'body str> = tag(open).parse(after_prefix_inclusive); + let (after_open, _) = consume_result.ok()?; + + Some(after_open) +} + +fn parameter_value_to_json(raw: &str) -> serde_json::Value { + serde_json::from_str::(raw) + .ok() + .unwrap_or_else(|| serde_json::Value::String(raw.to_owned())) +} + +fn parse_one_parameter<'body>( + input: &'body str, + shape: &KeyValueXmlTagsShape, + function_name: &str, +) -> Result, KeyValueXmlTagsFailure> { + let take_result: IResult<&'body str, &'body str> = + take_until(shape.key_open.as_str()).parse(input); + let Ok((after_key_open_inclusive, _)) = take_result else { + return Ok(None); + }; + let consume_result: IResult<&'body str, &'body str> = + tag(shape.key_open.as_str()).parse(after_key_open_inclusive); + let Ok((after_key_open, _)) = consume_result else { + return Ok(None); + }; + + let key_close_position = after_key_open + .find(shape.key_close.as_str()) + .ok_or_else(|| KeyValueXmlTagsFailure::UnclosedKeyTag { + function_name: function_name.to_owned(), + expected_close: shape.key_close.clone(), + })?; + let key = after_key_open[..key_close_position].trim().to_owned(); + if key.is_empty() { + return Err(KeyValueXmlTagsFailure::EmptyKey { + function_name: function_name.to_owned(), + }); + } + let after_key_close = &after_key_open[key_close_position + shape.key_close.len()..]; + + let value_open_take: IResult<&str, &str> = + take_until(shape.value_open.as_str()).parse(after_key_close); + let Ok((after_value_open_inclusive, _)) = value_open_take else { + return Err(KeyValueXmlTagsFailure::MissingValueTag { + function_name: function_name.to_owned(), + key, + expected_open: shape.value_open.clone(), + }); + }; + let value_open_consume: IResult<&str, &str> = + tag(shape.value_open.as_str()).parse(after_value_open_inclusive); + let Ok((after_value_open, _)) = value_open_consume else { + return Err(KeyValueXmlTagsFailure::MissingValueTag { + function_name: function_name.to_owned(), + key, + expected_open: shape.value_open.clone(), + }); + }; + + let value_close_position = after_value_open + .find(shape.value_close.as_str()) + .ok_or_else(|| KeyValueXmlTagsFailure::UnclosedValueTag { + function_name: function_name.to_owned(), + key: key.clone(), + expected_close: shape.value_close.clone(), + })?; + let raw_value = &after_value_open[..value_close_position]; + let value = parameter_value_to_json(raw_value); + let after_value_close = &after_value_open[value_close_position + shape.value_close.len()..]; + + Ok(Some((key, value, after_value_close))) +} + +fn collect_parameters( + function_body: &str, + shape: &KeyValueXmlTagsShape, + function_name: &str, +) -> Result, KeyValueXmlTagsFailure> { + let mut parameters = serde_json::Map::new(); + let mut remaining = function_body; + + while let Some((key, value, rest)) = parse_one_parameter(remaining, shape, function_name)? { + parameters.insert(key, value); + remaining = rest; + } + + Ok(parameters) +} + +fn parse_one_call<'body>( + input: &'body str, + markers: &ToolCallMarkers, + shape: &KeyValueXmlTagsShape, +) -> Result, KeyValueXmlTagsFailure> { + let Some(after_open) = skip_to_next_open(input, &markers.open) else { + return Ok(ParseStep::Done); + }; + + let Some(close_position) = after_open.find(markers.close.as_str()) else { + return Err(KeyValueXmlTagsFailure::UnclosedFunctionBlock { + expected_close: markers.close.clone(), + }); + }; + let function_block = &after_open[..close_position]; + let after_function_close = &after_open[close_position + markers.close.len()..]; + + let (name_end, has_args) = function_block + .find(shape.key_open.as_str()) + .map_or((function_block.len(), false), |position| (position, true)); + let function_name = function_block[..name_end].trim().to_owned(); + if function_name.is_empty() { + return Err(KeyValueXmlTagsFailure::EmptyFunctionName); + } + + let args_section = if has_args { + &function_block[name_end..] + } else { + "" + }; + let arguments_object = collect_parameters(args_section, shape, &function_name)?; + let arguments_value = serde_json::Value::Object(arguments_object); + let arguments = ToolCallArguments::from_string(arguments_value.to_string()); + + Ok(ParseStep::Call( + ParsedToolCall::new(String::new(), function_name, arguments), + after_function_close, + )) +} + +/// # Errors +/// +/// Returns [`KeyValueXmlTagsFailure`] when the body looks like a key-value-XML +/// tool-call block (matches the open marker) but contains a structural issue: +/// empty function/key name, missing key/value tag, or unclosed function block. +pub fn parse( + body: &str, + markers: &ToolCallMarkers, + shape: &KeyValueXmlTagsShape, +) -> Result, KeyValueXmlTagsFailure> { + if !shape_is_complete(shape) || markers.open.is_empty() || markers.close.is_empty() { + return Ok(Vec::new()); + } + + let mut parsed = Vec::new(); + let mut remaining = body; + + loop { + match parse_one_call(remaining, markers, shape)? { + ParseStep::Done => break, + ParseStep::Call(call, rest) => { + parsed.push(call); + remaining = rest; + } + } + } + + Ok(parsed) +} + +#[cfg(test)] +mod tests { + use llama_cpp_bindings_types::KeyValueXmlTagsShape; + use llama_cpp_bindings_types::ToolCallArgsShape; + use llama_cpp_bindings_types::ToolCallArguments; + use llama_cpp_bindings_types::ToolCallMarkers; + use serde_json::json; + + use super::parse; + use crate::error::KeyValueXmlTagsFailure; + + fn glm47_markers() -> ToolCallMarkers { + ToolCallMarkers { + open: "".to_owned(), + close: "".to_owned(), + args_shape: ToolCallArgsShape::KeyValueXmlTags(glm47_shape()), + } + } + + fn glm47_shape() -> KeyValueXmlTagsShape { + KeyValueXmlTagsShape { + key_open: "".to_owned(), + key_close: "".to_owned(), + value_open: "".to_owned(), + value_close: "".to_owned(), + } + } + + #[test] + fn parses_single_call_with_one_argument() { + let body = "get_weatherlocationParis"; + let parsed = parse(body, &glm47_markers(), &glm47_shape()).expect("must parse"); + + assert_eq!(parsed.len(), 1); + assert_eq!(parsed[0].name, "get_weather"); + assert_eq!( + parsed[0].arguments, + ToolCallArguments::ValidJson(json!({"location": "Paris"})), + ); + } + + #[test] + fn parses_call_with_multiple_arguments() { + let body = "set_thermostatroomkitchencelsius21"; + let parsed = parse(body, &glm47_markers(), &glm47_shape()).expect("must parse"); + + assert_eq!(parsed.len(), 1); + assert_eq!(parsed[0].name, "set_thermostat"); + assert_eq!( + parsed[0].arguments, + ToolCallArguments::ValidJson(json!({"room": "kitchen", "celsius": 21})), + ); + } + + #[test] + fn parses_two_calls_in_one_body() { + let body = "ax1by2"; + let parsed = parse(body, &glm47_markers(), &glm47_shape()).expect("must parse"); + + assert_eq!(parsed.len(), 2); + assert_eq!(parsed[0].name, "a"); + assert_eq!(parsed[1].name, "b"); + } + + #[test] + fn parses_call_with_no_arguments() { + let body = "ping"; + let parsed = parse(body, &glm47_markers(), &glm47_shape()).expect("must parse"); + + assert_eq!(parsed.len(), 1); + assert_eq!(parsed[0].name, "ping"); + assert_eq!(parsed[0].arguments, ToolCallArguments::ValidJson(json!({})),); + } + + #[test] + fn rejects_unclosed_function_block_with_typed_failure() { + let body = "get_weatherlocationParis"; + let result = parse(body, &glm47_markers(), &glm47_shape()); + + match result.expect_err("must error") { + KeyValueXmlTagsFailure::UnclosedFunctionBlock { expected_close } => { + assert_eq!(expected_close, ""); + } + other => panic!("expected UnclosedFunctionBlock, got {other:?}"), + } + } + + #[test] + fn rejects_empty_function_name_with_typed_failure() { + let body = "kv"; + let result = parse(body, &glm47_markers(), &glm47_shape()); + + match result.expect_err("must error") { + KeyValueXmlTagsFailure::EmptyFunctionName => {} + other => panic!("expected EmptyFunctionName, got {other:?}"), + } + } + + #[test] + fn rejects_unclosed_key_tag_with_typed_failure() { + let body = "flocation"; + let result = parse(body, &glm47_markers(), &glm47_shape()); + + match result.expect_err("must error") { + KeyValueXmlTagsFailure::UnclosedKeyTag { function_name, .. } => { + assert_eq!(function_name, "f"); + } + other => panic!("expected UnclosedKeyTag, got {other:?}"), + } + } + + #[test] + fn rejects_missing_value_tag_with_typed_failure() { + let body = "flocationParis"; + let result = parse(body, &glm47_markers(), &glm47_shape()); + + match result.expect_err("must error") { + KeyValueXmlTagsFailure::MissingValueTag { + function_name, key, .. + } => { + assert_eq!(function_name, "f"); + assert_eq!(key, "location"); + } + other => panic!("expected MissingValueTag, got {other:?}"), + } + } + + #[test] + fn returns_empty_for_body_without_open_marker() { + let parsed = + parse("plain text", &glm47_markers(), &glm47_shape()).expect("must parse empty"); + assert!(parsed.is_empty()); + } + + #[test] + fn returns_empty_when_shape_is_incomplete() { + let mut shape = glm47_shape(); + shape.value_close.clear(); + let body = "fkv"; + let parsed = parse(body, &glm47_markers(), &shape).expect("must parse empty"); + assert!(parsed.is_empty()); + } +} diff --git a/llama-cpp-bindings/src/tool_call_format/mod.rs b/llama-cpp-bindings/src/tool_call_format/mod.rs new file mode 100644 index 00000000..134b9e8e --- /dev/null +++ b/llama-cpp-bindings/src/tool_call_format/mod.rs @@ -0,0 +1,384 @@ +pub mod bracketed_args; +pub mod json_object; +pub mod key_value_xml_tags; +pub mod paired_quote_args; +pub mod tool_call_format_outcome; +pub mod xml_function_tags; + +pub use self::tool_call_format_outcome::ToolCallFormatOutcome; + +use llama_cpp_bindings_types::ToolCallArgsShape; +use llama_cpp_bindings_types::ToolCallMarkers; + +use crate::error::ToolCallFormatFailure; + +#[must_use] +pub fn try_parse(body: &str, markers: &ToolCallMarkers) -> ToolCallFormatOutcome { + if markers.open.is_empty() { + return ToolCallFormatOutcome::NoMatch; + } + + let parsed: Result, ToolCallFormatFailure> = match &markers.args_shape { + ToolCallArgsShape::BracketedJson(shape) => { + bracketed_args::parse(body, markers, shape).map_err(Into::into) + } + ToolCallArgsShape::JsonObject(shape) => json_object::parse(body, shape).map_err(Into::into), + ToolCallArgsShape::KeyValueXmlTags(shape) => { + key_value_xml_tags::parse(body, markers, shape).map_err(Into::into) + } + ToolCallArgsShape::PairedQuote(shape) => { + paired_quote_args::parse(body, markers, shape).map_err(Into::into) + } + ToolCallArgsShape::XmlTags(shape) => { + xml_function_tags::parse(body, shape).map_err(Into::into) + } + }; + + match parsed { + Ok(parsed) if parsed.is_empty() => ToolCallFormatOutcome::NoMatch, + Ok(parsed) => ToolCallFormatOutcome::Parsed(parsed), + Err(failure) => ToolCallFormatOutcome::Failed(failure), + } +} + +#[cfg(test)] +mod tests { + use llama_cpp_bindings_types::BracketedJsonShape; + use llama_cpp_bindings_types::KeyValueXmlTagsShape; + use llama_cpp_bindings_types::PairedQuoteShape; + use llama_cpp_bindings_types::ToolCallArgsShape; + use llama_cpp_bindings_types::ToolCallArguments; + use llama_cpp_bindings_types::ToolCallMarkers; + use llama_cpp_bindings_types::ToolCallValueQuote; + use llama_cpp_bindings_types::XmlTagsShape; + use serde_json::json; + + use super::ToolCallFormatOutcome; + use super::try_parse; + + fn mistral3_markers() -> ToolCallMarkers { + ToolCallMarkers { + open: "[TOOL_CALLS]".to_owned(), + close: String::new(), + args_shape: ToolCallArgsShape::BracketedJson(BracketedJsonShape { + name_args_separator: "[ARGS]".to_owned(), + }), + } + } + + fn gemma4_markers() -> ToolCallMarkers { + ToolCallMarkers { + open: "<|tool_call>call:".to_owned(), + close: "}".to_owned(), + args_shape: ToolCallArgsShape::PairedQuote(PairedQuoteShape { + name_args_separator: "{".to_owned(), + value_quote: ToolCallValueQuote { + open: "<|\"|>".to_owned(), + close: "<|\"|>".to_owned(), + }, + }), + } + } + + fn qwen35_markers() -> ToolCallMarkers { + ToolCallMarkers { + open: "".to_owned(), + close: "".to_owned(), + args_shape: ToolCallArgsShape::XmlTags(XmlTagsShape { + function_open_prefix: "".to_owned(), + parameter_open_prefix: "".to_owned(), + }), + } + } + + fn glm47_markers() -> ToolCallMarkers { + ToolCallMarkers { + open: "".to_owned(), + close: "".to_owned(), + args_shape: ToolCallArgsShape::KeyValueXmlTags(KeyValueXmlTagsShape { + key_open: "".to_owned(), + key_close: "".to_owned(), + value_open: "".to_owned(), + value_close: "".to_owned(), + }), + } + } + + #[test] + fn dispatches_to_bracketed_args_for_mistral3_shape() { + let outcome = try_parse( + "[TOOL_CALLS]get_weather[ARGS]{\"location\":\"Paris\"}", + &mistral3_markers(), + ); + + match outcome { + ToolCallFormatOutcome::Parsed(calls) => { + assert_eq!(calls.len(), 1); + assert_eq!(calls[0].name, "get_weather"); + assert_eq!( + calls[0].arguments, + ToolCallArguments::ValidJson(json!({"location": "Paris"})), + ); + } + other => panic!("expected Parsed, got {other:?}"), + } + } + + #[test] + fn dispatches_to_paired_quote_args_for_gemma4_shape() { + let outcome = try_parse( + "<|tool_call>call:get_weather{location:<|\"|>Paris<|\"|>}", + &gemma4_markers(), + ); + + match outcome { + ToolCallFormatOutcome::Parsed(calls) => { + assert_eq!(calls.len(), 1); + assert_eq!(calls[0].name, "get_weather"); + assert_eq!( + calls[0].arguments, + ToolCallArguments::ValidJson(json!({"location": "Paris"})), + ); + } + other => panic!("expected Parsed, got {other:?}"), + } + } + + #[test] + fn dispatches_to_key_value_xml_tags_for_glm47_shape() { + let outcome = try_parse( + "get_weatherlocationParis", + &glm47_markers(), + ); + + match outcome { + ToolCallFormatOutcome::Parsed(calls) => { + assert_eq!(calls.len(), 1); + assert_eq!(calls[0].name, "get_weather"); + assert_eq!( + calls[0].arguments, + ToolCallArguments::ValidJson(json!({"location": "Paris"})), + ); + } + other => panic!("expected Parsed, got {other:?}"), + } + } + + #[test] + fn dispatches_to_xml_function_tags_for_qwen35_shape() { + let outcome = try_parse( + "Paris", + &qwen35_markers(), + ); + + match outcome { + ToolCallFormatOutcome::Parsed(calls) => { + assert_eq!(calls.len(), 1); + assert_eq!(calls[0].name, "get_weather"); + assert_eq!( + calls[0].arguments, + ToolCallArguments::ValidJson(json!({"location": "Paris"})), + ); + } + other => panic!("expected Parsed, got {other:?}"), + } + } + + #[test] + fn no_match_when_open_marker_is_empty() { + let markers = ToolCallMarkers { + open: String::new(), + close: String::new(), + args_shape: ToolCallArgsShape::BracketedJson(BracketedJsonShape { + name_args_separator: "[ARGS]".to_owned(), + }), + }; + + match try_parse("[TOOL_CALLS]get_weather[ARGS]{}", &markers) { + ToolCallFormatOutcome::NoMatch => {} + other => panic!("expected NoMatch, got {other:?}"), + } + } + + #[test] + fn no_match_when_body_lacks_markers() { + match try_parse("plain text without tool calls", &mistral3_markers()) { + ToolCallFormatOutcome::NoMatch => {} + other => panic!("expected NoMatch, got {other:?}"), + } + } + + #[test] + fn failed_when_inner_parser_returns_typed_failure() { + match try_parse( + "[TOOL_CALLS]get_weather[ARGS]{\"location\":}", + &mistral3_markers(), + ) { + ToolCallFormatOutcome::Failed(_) => {} + other => panic!("expected Failed, got {other:?}"), + } + } + + #[test] + fn try_parse_returns_no_match_for_glm_input_under_qwen_markers() { + let glm_input = "get_weather\ + location\ + Paris\ + "; + + match try_parse(glm_input, &qwen35_markers()) { + ToolCallFormatOutcome::NoMatch => {} + other => panic!("expected NoMatch for GLM input under Qwen markers, got {other:?}"), + } + } + + #[test] + fn try_parse_returns_no_match_for_plain_content_under_every_known_shape() { + use crate::tool_call_template_overrides::known_marker_candidates; + + let plain_content = "Sorry, I cannot help with that request."; + + for candidate in known_marker_candidates() { + match try_parse(plain_content, &candidate) { + ToolCallFormatOutcome::NoMatch => {} + other => panic!( + "expected NoMatch for plain content under candidate {candidate:?}, got {other:?}" + ), + } + } + } + + #[test] + fn duck_type_resolves_qwen_xml_input_via_xml_tags_shape_first() { + use llama_cpp_bindings_types::ToolCallArguments; + + use crate::tool_call_template_overrides::known_marker_candidates; + + let qwen_input = "\n\ + \n\ + \n\ + Paris\n\ + \n\ + \n\ + "; + + let mut resolved = None; + for candidate in known_marker_candidates() { + if let ToolCallFormatOutcome::Parsed(calls) = try_parse(qwen_input, &candidate) { + resolved = Some((candidate.args_shape, calls)); + break; + } + } + + let (args_shape, calls) = + resolved.expect("Qwen XML input must resolve via at least one duck-type candidate"); + assert!( + matches!(args_shape, ToolCallArgsShape::XmlTags(_)), + "duck-type ordering must resolve Qwen XML via the XmlTags shape (most restrictive \ + shape that requires `Paris<|\"|>}"; + + let mut resolved = None; + for candidate in known_marker_candidates() { + if let ToolCallFormatOutcome::Parsed(calls) = try_parse(gemma_input, &candidate) { + resolved = Some((candidate.args_shape, calls)); + break; + } + } + + let (args_shape, calls) = + resolved.expect("Gemma input must resolve via at least one duck-type candidate"); + assert!( + matches!(args_shape, ToolCallArgsShape::PairedQuote(_)), + "Gemma input must resolve via the PairedQuote shape, got {args_shape:?}" + ); + assert_eq!(calls.len(), 1); + assert_eq!(calls[0].name, "get_weather"); + assert_eq!( + calls[0].arguments, + ToolCallArguments::ValidJson(json!({"location": "Paris"})), + ); + } +} diff --git a/llama-cpp-bindings/src/tool_call_format/paired_quote_args.rs b/llama-cpp-bindings/src/tool_call_format/paired_quote_args.rs new file mode 100644 index 00000000..eba1b87e --- /dev/null +++ b/llama-cpp-bindings/src/tool_call_format/paired_quote_args.rs @@ -0,0 +1,433 @@ +use llama_cpp_bindings_types::PairedQuoteShape; +use llama_cpp_bindings_types::ParsedToolCall; +use llama_cpp_bindings_types::ToolCallArguments; +use llama_cpp_bindings_types::ToolCallMarkers; +use llama_cpp_bindings_types::ToolCallValueQuote; + +use crate::error::PairedQuoteFailure; + +enum ParseStep<'body> { + Done, + Call(ParsedToolCall, &'body str), +} + +fn consume_optional_prefix<'body>(input: &'body str, literal: &str) -> &'body str { + input.strip_prefix(literal).unwrap_or(input) +} + +fn split_at_separator<'body>( + input: &'body str, + separator: &str, +) -> Option<(&'body str, &'body str)> { + let (name_raw, after_separator) = input.split_once(separator)?; + Some((name_raw, after_separator)) +} + +fn bare_value_to_json(text: &str) -> serde_json::Value { + if text.is_empty() { + return serde_json::Value::Null; + } + serde_json::from_str::(text) + .ok() + .unwrap_or_else(|| serde_json::Value::String(text.to_owned())) +} + +fn find_bare_value_end(input: &str, close_marker: &str) -> usize { + for (byte_index, character) in input.char_indices() { + if character == ',' { + return byte_index; + } + if !close_marker.is_empty() && input[byte_index..].starts_with(close_marker) { + return byte_index; + } + } + + input.len() +} + +fn parse_one_key<'body>( + input: &'body str, + tool_name: &str, +) -> Result<(String, &'body str), PairedQuoteFailure> { + let Some((key_raw, after_colon)) = input.split_once(':') else { + return Err(PairedQuoteFailure::UnclosedArgumentBlock { + tool_name: tool_name.to_owned(), + state: "key", + }); + }; + let key = key_raw.trim().to_owned(); + if key.is_empty() { + return Err(PairedQuoteFailure::EmptyKey { + tool_name: tool_name.to_owned(), + }); + } + + Ok((key, after_colon)) +} + +fn parse_one_value<'body>( + input: &'body str, + value_quote: &ToolCallValueQuote, + close_marker: &str, + tool_name: &str, + key: &str, +) -> Result<(serde_json::Value, &'body str), PairedQuoteFailure> { + let trimmed = input.trim_start(); + + if !value_quote.open.is_empty() + && !value_quote.close.is_empty() + && let Some(after_open) = trimmed.strip_prefix(value_quote.open.as_str()) + { + let Some(close_position) = after_open.find(value_quote.close.as_str()) else { + return Err(PairedQuoteFailure::UnclosedQuotedValue { + tool_name: tool_name.to_owned(), + key: key.to_owned(), + }); + }; + let value_text = after_open[..close_position].to_owned(); + let after_close = &after_open[close_position + value_quote.close.len()..]; + + return Ok((serde_json::Value::String(value_text), after_close)); + } + + let bare_end = find_bare_value_end(trimmed, close_marker); + let bare_text = trimmed[..bare_end].trim(); + let value = bare_value_to_json(bare_text); + + Ok((value, &trimmed[bare_end..])) +} + +fn parse_args_body<'body>( + input: &'body str, + value_quote: &ToolCallValueQuote, + close_marker: &str, + tool_name: &str, +) -> Result<(serde_json::Map, &'body str), PairedQuoteFailure> { + let mut map = serde_json::Map::new(); + let mut remaining = input.trim_start(); + + loop { + if remaining.is_empty() { + return Ok((map, remaining)); + } + if !close_marker.is_empty() + && let Some(after_close) = remaining.strip_prefix(close_marker) + { + return Ok((map, after_close)); + } + + let (key, after_key) = parse_one_key(remaining, tool_name)?; + let (value, after_value) = + parse_one_value(after_key, value_quote, close_marker, tool_name, &key)?; + map.insert(key.clone(), value); + + remaining = after_value.trim_start(); + if remaining.is_empty() { + return Ok((map, remaining)); + } + if !close_marker.is_empty() + && let Some(after_close) = remaining.strip_prefix(close_marker) + { + return Ok((map, after_close)); + } + if let Some(after_comma) = remaining.strip_prefix(',') { + remaining = after_comma.trim_start(); + continue; + } + + let Some(character) = remaining.chars().next() else { + return Ok((map, remaining)); + }; + + return Err(PairedQuoteFailure::UnexpectedCharAfterValue { + tool_name: tool_name.to_owned(), + key, + character, + }); + } +} + +fn parse_one_call<'body>( + input: &'body str, + markers: &ToolCallMarkers, + shape: &PairedQuoteShape, +) -> Result, PairedQuoteFailure> { + if input.is_empty() { + return Ok(ParseStep::Done); + } + + let after_open = consume_optional_prefix(input, markers.open.as_str()); + + let Some((name_raw, after_separator)) = + split_at_separator(after_open, shape.name_args_separator.as_str()) + else { + return Ok(ParseStep::Done); + }; + + let name = name_raw.trim().to_owned(); + if name.is_empty() { + return Ok(ParseStep::Done); + } + + let (args_object, after_args) = parse_args_body( + after_separator, + &shape.value_quote, + markers.close.as_str(), + &name, + )?; + let arguments_value = serde_json::Value::Object(args_object); + + Ok(ParseStep::Call( + ParsedToolCall::new( + String::new(), + name, + ToolCallArguments::ValidJson(arguments_value), + ), + after_args, + )) +} + +/// # Errors +/// +/// Returns [`PairedQuoteFailure`] when the body looks like a paired-quote +/// tool-call block (matches the open marker and separator) but contains a +/// structural issue: empty key, unclosed quoted value, unexpected character +/// after a value, or an unfinished argument block. +pub fn parse( + body: &str, + markers: &ToolCallMarkers, + shape: &PairedQuoteShape, +) -> Result, PairedQuoteFailure> { + if shape.name_args_separator.is_empty() { + return Ok(Vec::new()); + } + + let mut parsed = Vec::new(); + let mut remaining = body.trim_start(); + + loop { + match parse_one_call(remaining, markers, shape)? { + ParseStep::Done => break, + ParseStep::Call(call, rest) => { + parsed.push(call); + remaining = rest.trim_start(); + } + } + } + + Ok(parsed) +} + +#[cfg(test)] +mod tests { + #![expect( + clippy::literal_string_with_formatting_args, + reason = "Gemma tool-call format literals contain braces that resemble format args" + )] + + use llama_cpp_bindings_types::PairedQuoteShape; + use llama_cpp_bindings_types::ToolCallArgsShape; + use llama_cpp_bindings_types::ToolCallArguments; + use llama_cpp_bindings_types::ToolCallMarkers; + use llama_cpp_bindings_types::ToolCallValueQuote; + use serde_json::json; + + use super::parse; + use crate::error::PairedQuoteFailure; + + fn gemma4_markers() -> ToolCallMarkers { + ToolCallMarkers { + open: "<|tool_call>call:".to_owned(), + close: "}".to_owned(), + args_shape: ToolCallArgsShape::PairedQuote(gemma4_shape()), + } + } + + fn gemma4_shape() -> PairedQuoteShape { + PairedQuoteShape { + name_args_separator: "{".to_owned(), + value_quote: ToolCallValueQuote { + open: "<|\"|>".to_owned(), + close: "<|\"|>".to_owned(), + }, + } + } + + #[test] + fn parses_single_quoted_string_argument_with_full_markers() { + let parsed = parse( + "<|tool_call>call:get_weather{location:<|\"|>Paris<|\"|>}", + &gemma4_markers(), + &gemma4_shape(), + ) + .expect("must parse"); + + assert_eq!(parsed.len(), 1); + assert_eq!(parsed[0].name, "get_weather"); + assert_eq!( + parsed[0].arguments, + ToolCallArguments::ValidJson(json!({"location": "Paris"})), + ); + } + + #[test] + fn parses_classifier_stripped_body_without_open_or_close() { + let parsed = parse( + "get_weather{location:<|\"|>Paris<|\"|>", + &gemma4_markers(), + &gemma4_shape(), + ) + .expect("must parse"); + + assert_eq!(parsed.len(), 1); + assert_eq!(parsed[0].name, "get_weather"); + assert_eq!( + parsed[0].arguments, + ToolCallArguments::ValidJson(json!({"location": "Paris"})), + ); + } + + #[test] + fn parses_multiple_quoted_string_arguments() { + let parsed = parse( + "<|tool_call>call:f{a:<|\"|>1<|\"|>,b:<|\"|>2<|\"|>}", + &gemma4_markers(), + &gemma4_shape(), + ) + .expect("must parse"); + + assert_eq!(parsed.len(), 1); + assert_eq!( + parsed[0].arguments, + ToolCallArguments::ValidJson(json!({"a": "1", "b": "2"})), + ); + } + + #[test] + fn parses_bare_numeric_value() { + let parsed = parse( + "<|tool_call>call:f{a:42}", + &gemma4_markers(), + &gemma4_shape(), + ) + .expect("must parse"); + + assert_eq!( + parsed[0].arguments, + ToolCallArguments::ValidJson(json!({"a": 42})), + ); + } + + #[test] + fn parses_bare_boolean_value() { + let parsed = parse( + "<|tool_call>call:f{a:true}", + &gemma4_markers(), + &gemma4_shape(), + ) + .expect("must parse"); + + assert_eq!( + parsed[0].arguments, + ToolCallArguments::ValidJson(json!({"a": true})), + ); + } + + #[test] + fn rejects_unclosed_quoted_value_with_typed_failure() { + let result = parse( + "<|tool_call>call:f{a:<|\"|>oops", + &gemma4_markers(), + &gemma4_shape(), + ); + + match result.expect_err("unclosed quote must produce a typed failure") { + PairedQuoteFailure::UnclosedQuotedValue { tool_name, key } => { + assert_eq!(tool_name, "f"); + assert_eq!(key, "a"); + } + other => panic!("expected UnclosedQuotedValue, got {other:?}"), + } + } + + #[test] + fn rejects_unexpected_char_after_value_with_typed_failure() { + let result = parse( + "<|tool_call>call:f{a:<|\"|>v<|\"|>$bad}", + &gemma4_markers(), + &gemma4_shape(), + ); + + match result.expect_err("garbage after value must produce a typed failure") { + PairedQuoteFailure::UnexpectedCharAfterValue { + tool_name, + key, + character, + } => { + assert_eq!(tool_name, "f"); + assert_eq!(key, "a"); + assert_eq!(character, '$'); + } + other => panic!("expected UnexpectedCharAfterValue, got {other:?}"), + } + } + + #[test] + fn returns_empty_vec_for_empty_body() { + let parsed = parse("", &gemma4_markers(), &gemma4_shape()).expect("empty body must parse"); + assert!(parsed.is_empty()); + } + + #[test] + fn returns_empty_vec_when_body_lacks_separator() { + let parsed = parse("no separator anywhere", &gemma4_markers(), &gemma4_shape()) + .expect("body without separator must parse"); + assert!(parsed.is_empty()); + } + + #[test] + fn parses_args_body_terminated_by_end_of_input_after_quoted_value() { + let parsed = parse( + "<|tool_call>call:f{x:<|\"|>v<|\"|>", + &gemma4_markers(), + &gemma4_shape(), + ) + .expect("end-of-input after quoted value must parse"); + + assert_eq!( + parsed[0].arguments, + ToolCallArguments::ValidJson(json!({"x": "v"})), + ); + } + + #[test] + fn parses_args_body_terminated_by_end_of_input_after_bare_value() { + let parsed = parse( + "<|tool_call>call:f{n:42", + &gemma4_markers(), + &gemma4_shape(), + ) + .expect("end-of-input after bare value must parse"); + + assert_eq!( + parsed[0].arguments, + ToolCallArguments::ValidJson(json!({"n": 42})), + ); + } + + #[test] + fn rejects_empty_key_with_typed_failure() { + let result = parse( + "<|tool_call>call:f{:42}", + &gemma4_markers(), + &gemma4_shape(), + ); + + match result.expect_err("empty key must produce a typed failure") { + PairedQuoteFailure::EmptyKey { tool_name } => { + assert_eq!(tool_name, "f"); + } + other => panic!("expected EmptyKey, got {other:?}"), + } + } +} diff --git a/llama-cpp-bindings/src/tool_call_format/tool_call_format_outcome.rs b/llama-cpp-bindings/src/tool_call_format/tool_call_format_outcome.rs new file mode 100644 index 00000000..fa5e1368 --- /dev/null +++ b/llama-cpp-bindings/src/tool_call_format/tool_call_format_outcome.rs @@ -0,0 +1,10 @@ +use llama_cpp_bindings_types::ParsedToolCall; + +use crate::error::ToolCallFormatFailure; + +#[derive(Debug)] +pub enum ToolCallFormatOutcome { + Parsed(Vec), + NoMatch, + Failed(ToolCallFormatFailure), +} diff --git a/llama-cpp-bindings/src/tool_call_format/xml_function_tags.rs b/llama-cpp-bindings/src/tool_call_format/xml_function_tags.rs new file mode 100644 index 00000000..0e1cb0af --- /dev/null +++ b/llama-cpp-bindings/src/tool_call_format/xml_function_tags.rs @@ -0,0 +1,369 @@ +use llama_cpp_bindings_types::ParsedToolCall; +use llama_cpp_bindings_types::ToolCallArguments; +use llama_cpp_bindings_types::XmlTagsShape; +use nom::IResult; +use nom::Parser; +use nom::bytes::complete::tag; +use nom::bytes::complete::take_until; + +use crate::error::XmlFunctionTagsFailure; + +const fn shape_is_complete(shape: &XmlTagsShape) -> bool { + !shape.function_open_prefix.is_empty() + && !shape.function_close.is_empty() + && !shape.parameter_open_prefix.is_empty() + && !shape.parameter_close.is_empty() +} + +fn trim_surrounding_newlines(input: &str) -> &str { + input.trim_start_matches('\n').trim_end_matches('\n') +} + +fn parameter_value_to_json(raw: &str) -> serde_json::Value { + serde_json::from_str::(raw) + .ok() + .unwrap_or_else(|| serde_json::Value::String(raw.to_owned())) +} + +fn locate_tag_name_end(after_prefix: &str) -> Option { + let close_position = after_prefix.find('>'); + let next_open_position = after_prefix.find('<'); + + match (close_position, next_open_position) { + (Some(close), Some(open)) if open < close => None, + (Some(close), _) => Some(close), + (None, _) => None, + } +} + +fn skip_to_next_function_open<'body>( + input: &'body str, + function_open_prefix: &str, +) -> Option<&'body str> { + let take_result: IResult<&'body str, &'body str> = + take_until(function_open_prefix).parse(input); + let (after_prefix_inclusive, _) = take_result.ok()?; + let consume_result: IResult<&'body str, &'body str> = + tag(function_open_prefix).parse(after_prefix_inclusive); + let (after_prefix, _) = consume_result.ok()?; + + Some(after_prefix) +} + +fn parse_one_parameter<'body>( + input: &'body str, + shape: &XmlTagsShape, + function_name: &str, +) -> Result, XmlFunctionTagsFailure> { + let take_result: IResult<&'body str, &'body str> = + take_until(shape.parameter_open_prefix.as_str()).parse(input); + let Ok((after_prefix_inclusive, _)) = take_result else { + return Ok(None); + }; + let consume_result: IResult<&'body str, &'body str> = + tag(shape.parameter_open_prefix.as_str()).parse(after_prefix_inclusive); + let Ok((after_prefix, _)) = consume_result else { + return Ok(None); + }; + + let Some(name_end) = locate_tag_name_end(after_prefix) else { + return Err(XmlFunctionTagsFailure::UnclosedParameterBlock { + function_name: function_name.to_owned(), + parameter_name: String::new(), + expected_close: shape.parameter_close.clone(), + }); + }; + let parameter_name = after_prefix[..name_end].trim().to_owned(); + if parameter_name.is_empty() { + return Err(XmlFunctionTagsFailure::EmptyParameterName { + function_name: function_name.to_owned(), + }); + } + let value_start = &after_prefix[name_end + 1..]; + + let Some(value_end_position) = value_start.find(shape.parameter_close.as_str()) else { + return Err(XmlFunctionTagsFailure::UnclosedParameterBlock { + function_name: function_name.to_owned(), + parameter_name, + expected_close: shape.parameter_close.clone(), + }); + }; + let raw_value = trim_surrounding_newlines(&value_start[..value_end_position]); + let after_close = &value_start[value_end_position + shape.parameter_close.len()..]; + let parameter_value = parameter_value_to_json(raw_value); + + Ok(Some((parameter_name, parameter_value, after_close))) +} + +fn collect_parameters( + function_body: &str, + shape: &XmlTagsShape, + function_name: &str, +) -> Result, XmlFunctionTagsFailure> { + let mut parameters = serde_json::Map::new(); + let mut remaining = function_body; + + while let Some((parameter_name, parameter_value, rest)) = + parse_one_parameter(remaining, shape, function_name)? + { + parameters.insert(parameter_name, parameter_value); + remaining = rest; + } + + Ok(parameters) +} + +fn parse_one_function<'body>( + input: &'body str, + shape: &XmlTagsShape, +) -> Result, XmlFunctionTagsFailure> { + let Some(after_function_prefix) = + skip_to_next_function_open(input, &shape.function_open_prefix) + else { + return Ok(None); + }; + + let Some(name_end) = locate_tag_name_end(after_function_prefix) else { + return Err(XmlFunctionTagsFailure::UnclosedFunctionBlock { + function_name: String::new(), + expected_close: shape.function_close.clone(), + }); + }; + let function_name = after_function_prefix[..name_end].trim().to_owned(); + if function_name.is_empty() { + return Err(XmlFunctionTagsFailure::EmptyFunctionName); + } + let function_body_start = &after_function_prefix[name_end + 1..]; + + let Some(function_body_end) = function_body_start.find(shape.function_close.as_str()) else { + return Err(XmlFunctionTagsFailure::UnclosedFunctionBlock { + function_name, + expected_close: shape.function_close.clone(), + }); + }; + let function_body = &function_body_start[..function_body_end]; + let after_function_close = + &function_body_start[function_body_end + shape.function_close.len()..]; + + let arguments_object = collect_parameters(function_body, shape, &function_name)?; + let arguments_value = serde_json::Value::Object(arguments_object); + let arguments = ToolCallArguments::from_string(arguments_value.to_string()); + + Ok(Some(( + ParsedToolCall::new(String::new(), function_name, arguments), + after_function_close, + ))) +} + +/// # Errors +/// +/// Returns [`XmlFunctionTagsFailure`] when the body looks like an XML +/// function-tag tool-call block (matches the function open prefix) but +/// contains a structural issue: empty function/parameter name or an +/// unclosed function/parameter block. +pub fn parse( + body: &str, + shape: &XmlTagsShape, +) -> Result, XmlFunctionTagsFailure> { + if !shape_is_complete(shape) { + return Ok(Vec::new()); + } + + let mut parsed = Vec::new(); + let mut remaining = body; + + while let Some((call, rest)) = parse_one_function(remaining, shape)? { + parsed.push(call); + remaining = rest; + } + + Ok(parsed) +} + +#[cfg(test)] +mod tests { + use llama_cpp_bindings_types::ToolCallArguments; + use llama_cpp_bindings_types::XmlTagsShape; + use serde_json::json; + + use super::parse; + use crate::error::XmlFunctionTagsFailure; + + fn xml_shape() -> XmlTagsShape { + XmlTagsShape { + function_open_prefix: "".to_owned(), + parameter_open_prefix: "".to_owned(), + } + } + + #[test] + fn parses_single_function_with_one_parameter() { + let body = + "\n\n\nParis\n\n\n"; + let parsed = parse(body, &xml_shape()).expect("must parse"); + + assert_eq!(parsed.len(), 1); + assert_eq!(parsed[0].name, "get_weather"); + assert_eq!( + parsed[0].arguments, + ToolCallArguments::ValidJson(json!({"location": "Paris"})), + ); + } + + #[test] + fn parses_function_with_multiple_parameters() { + let body = "1two"; + let parsed = parse(body, &xml_shape()).expect("must parse"); + + assert_eq!(parsed.len(), 1); + assert_eq!( + parsed[0].arguments, + ToolCallArguments::ValidJson(json!({"a": 1, "b": "two"})), + ); + } + + #[test] + fn parses_two_function_blocks_in_one_body() { + let body = "12"; + let parsed = parse(body, &xml_shape()).expect("must parse"); + + assert_eq!(parsed.len(), 2); + assert_eq!(parsed[0].name, "a"); + assert_eq!(parsed[1].name, "b"); + } + + #[test] + fn preserves_multi_line_parameter_value() { + let body = "\n\nline one\nline two\n\n"; + let parsed = parse(body, &xml_shape()).expect("must parse"); + + assert_eq!( + parsed[0].arguments, + ToolCallArguments::ValidJson(json!({"msg": "line one\nline two"})), + ); + } + + #[test] + fn rejects_function_tag_missing_closing_angle_with_typed_failure() { + let body = "Paris"; + let result = parse(body, &xml_shape()); + + match result.expect_err("must error") { + XmlFunctionTagsFailure::UnclosedFunctionBlock { .. } => {} + other => panic!("expected UnclosedFunctionBlock, got {other:?}"), + } + } + + #[test] + fn rejects_function_block_missing_close_tag_with_typed_failure() { + let body = "Paris"; + let result = parse(body, &xml_shape()); + + match result.expect_err("must error") { + XmlFunctionTagsFailure::UnclosedFunctionBlock { + function_name, + expected_close, + } => { + assert_eq!(function_name, "get_weather"); + assert_eq!(expected_close, ""); + } + other => panic!("expected UnclosedFunctionBlock, got {other:?}"), + } + } + + #[test] + fn rejects_parameter_block_missing_close_tag_with_typed_failure() { + let body = "Paris"; + let result = parse(body, &xml_shape()); + + match result.expect_err("must error") { + XmlFunctionTagsFailure::UnclosedParameterBlock { + function_name, + parameter_name, + expected_close, + } => { + assert_eq!(function_name, "get_weather"); + assert_eq!(parameter_name, "location"); + assert_eq!(expected_close, ""); + } + other => panic!("expected UnclosedParameterBlock, got {other:?}"), + } + } + + #[test] + fn rejects_empty_function_name_with_typed_failure() { + let body = "1"; + let result = parse(body, &xml_shape()); + + match result.expect_err("must error") { + XmlFunctionTagsFailure::EmptyFunctionName => {} + other => panic!("expected EmptyFunctionName, got {other:?}"), + } + } + + #[test] + fn rejects_empty_parameter_name_with_typed_failure() { + let body = "1"; + let result = parse(body, &xml_shape()); + + match result.expect_err("must error") { + XmlFunctionTagsFailure::EmptyParameterName { function_name } => { + assert_eq!(function_name, "f"); + } + other => panic!("expected EmptyParameterName, got {other:?}"), + } + } + + #[test] + fn returns_empty_when_body_has_no_function_tag() { + let parsed = + parse("plain text without function tags", &xml_shape()).expect("must parse empty"); + assert!(parsed.is_empty()); + } + + #[test] + fn returns_empty_for_empty_body() { + let parsed = parse("", &xml_shape()).expect("must parse empty"); + assert!(parsed.is_empty()); + } + + #[test] + fn returns_empty_when_shape_has_empty_required_field() { + let mut shape = xml_shape(); + shape.function_close.clear(); + let body = "1"; + let parsed = parse(body, &shape).expect("must parse empty"); + assert!(parsed.is_empty()); + } + + #[test] + fn parses_negotiate_with_cat_reproducer_payload() { + let body = "\n\ +\n\ +\n\ +tuna\n\ +\n\ +\n\ +8\n\ +\n\ +\n\ +get off the keyboard\n\ +\n\ +\n\ +"; + let parsed = parse(body, &xml_shape()).expect("must parse"); + + assert_eq!(parsed.len(), 1); + assert_eq!(parsed[0].name, "negotiate_with_cat"); + assert_eq!( + parsed[0].arguments, + ToolCallArguments::ValidJson(json!({ + "bribe": "tuna", + "desperation_level": 8, + "topic": "get off the keyboard", + })), + ); + } +} diff --git a/llama-cpp-bindings/src/tool_call_marker_pair.rs b/llama-cpp-bindings/src/tool_call_marker_pair.rs new file mode 100644 index 00000000..3ee5fd42 --- /dev/null +++ b/llama-cpp-bindings/src/tool_call_marker_pair.rs @@ -0,0 +1,5 @@ +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct ToolCallMarkerPair { + pub open: String, + pub close: String, +} diff --git a/llama-cpp-bindings/src/tool_call_template_overrides/detect.rs b/llama-cpp-bindings/src/tool_call_template_overrides/detect.rs new file mode 100644 index 00000000..9dab2cdc --- /dev/null +++ b/llama-cpp-bindings/src/tool_call_template_overrides/detect.rs @@ -0,0 +1,66 @@ +use llama_cpp_bindings_types::ToolCallMarkers; + +use crate::tool_call_template_overrides::gemma4_call_block::Gemma4CallBlockOverride; +use crate::tool_call_template_overrides::glm47_key_value_tags::Glm47KeyValueTagsOverride; +use crate::tool_call_template_overrides::mistral3_arrow_args::Mistral3ArrowArgsOverride; +use crate::tool_call_template_overrides::qwen_xml_tags::QwenXmlTagsOverride; +use crate::tool_call_template_overrides::qwen3_json_inside_tool_call::Qwen3JsonInsideToolCallOverride; + +#[must_use] +pub fn detect(template: &str) -> Option { + let detectors: [fn(&str) -> Option; 5] = [ + Gemma4CallBlockOverride::detect, + Glm47KeyValueTagsOverride::detect, + Mistral3ArrowArgsOverride::detect, + Qwen3JsonInsideToolCallOverride::detect, + QwenXmlTagsOverride::detect, + ]; + detectors + .into_iter() + .find_map(|detector| detector(template)) +} + +#[cfg(test)] +mod tests { + use llama_cpp_bindings_types::ToolCallArgsShape; + + use super::detect; + + #[test] + fn dispatches_to_gemma4_override() { + let template = "{{- '<|tool_call>call:' + function['name'] + '{' -}}"; + let markers = detect(template).expect("must dispatch to Gemma 4"); + + assert_eq!(markers.open, "<|tool_call>call:"); + assert!(matches!( + markers.args_shape, + ToolCallArgsShape::PairedQuote(_) + )); + } + + #[test] + fn dispatches_to_mistral3_override() { + let template = "{{- name + '[ARGS]' + arguments }}"; + let markers = detect(template).expect("must dispatch to Mistral 3"); + + assert_eq!(markers.open, "[TOOL_CALLS]"); + assert!(matches!( + markers.args_shape, + ToolCallArgsShape::BracketedJson(_) + )); + } + + #[test] + fn dispatches_to_qwen_xml_tags_override() { + let template = "{{- '\\n\\n' }}"; + let markers = detect(template).expect("must dispatch to Qwen XML tags"); + + assert_eq!(markers.open, ""); + assert!(matches!(markers.args_shape, ToolCallArgsShape::XmlTags(_))); + } + + #[test] + fn returns_none_when_no_override_matches() { + assert!(detect("plain unrelated template").is_none()); + } +} diff --git a/llama-cpp-bindings/src/tool_call_template_overrides/gemma4_call_block.rs b/llama-cpp-bindings/src/tool_call_template_overrides/gemma4_call_block.rs new file mode 100644 index 00000000..f09a7b42 --- /dev/null +++ b/llama-cpp-bindings/src/tool_call_template_overrides/gemma4_call_block.rs @@ -0,0 +1,72 @@ +use llama_cpp_bindings_types::PairedQuoteShape; +use llama_cpp_bindings_types::ToolCallArgsShape; +use llama_cpp_bindings_types::ToolCallMarkers; +use llama_cpp_bindings_types::ToolCallValueQuote; + +pub struct Gemma4CallBlockOverride; + +impl Gemma4CallBlockOverride { + const TEMPLATE_FINGERPRINT: &'static str = "'<|tool_call>call:'"; + + #[must_use] + pub fn markers() -> ToolCallMarkers { + ToolCallMarkers { + open: "<|tool_call>call:".to_owned(), + close: "}".to_owned(), + args_shape: ToolCallArgsShape::PairedQuote(PairedQuoteShape { + name_args_separator: "{".to_owned(), + value_quote: ToolCallValueQuote { + open: "<|\"|>".to_owned(), + close: "<|\"|>".to_owned(), + }, + }), + } + } + + #[must_use] + pub fn detect(template: &str) -> Option { + if !template.contains(Self::TEMPLATE_FINGERPRINT) { + return None; + } + Some(Self::markers()) + } +} + +#[cfg(test)] +mod tests { + use llama_cpp_bindings_types::ToolCallArgsShape; + + use super::Gemma4CallBlockOverride; + + #[test] + fn detects_gemma4_template_with_tool_call_call_literal() { + let template = "...{{- '<|tool_call>call:' + function['name'] + '{' -}}..."; + let markers = + Gemma4CallBlockOverride::detect(template).expect("Gemma 4 template must be detected"); + + assert_eq!(markers.open, "<|tool_call>call:"); + assert_eq!(markers.close, "}"); + let ToolCallArgsShape::PairedQuote(shape) = markers.args_shape else { + panic!("expected PairedQuote variant, got {:?}", markers.args_shape); + }; + assert_eq!(shape.name_args_separator, "{"); + assert_eq!(shape.value_quote.open, "<|\"|>"); + assert_eq!(shape.value_quote.close, "<|\"|>"); + } + + #[test] + fn returns_none_for_template_without_fingerprint() { + assert!(Gemma4CallBlockOverride::detect("just some plain template body").is_none()); + } + + #[test] + fn returns_none_for_empty_template() { + assert!(Gemma4CallBlockOverride::detect("").is_none()); + } + + #[test] + fn returns_none_when_fingerprint_substring_appears_without_jinja_apostrophes() { + let template = "doc explaining the <|tool_call>call: format in prose, not as a literal"; + assert!(Gemma4CallBlockOverride::detect(template).is_none()); + } +} diff --git a/llama-cpp-bindings/src/tool_call_template_overrides/glm47_key_value_tags.rs b/llama-cpp-bindings/src/tool_call_template_overrides/glm47_key_value_tags.rs new file mode 100644 index 00000000..73373472 --- /dev/null +++ b/llama-cpp-bindings/src/tool_call_template_overrides/glm47_key_value_tags.rs @@ -0,0 +1,68 @@ +use llama_cpp_bindings_types::KeyValueXmlTagsShape; +use llama_cpp_bindings_types::ToolCallArgsShape; +use llama_cpp_bindings_types::ToolCallMarkers; + +pub struct Glm47KeyValueTagsOverride; + +impl Glm47KeyValueTagsOverride { + const TEMPLATE_FINGERPRINT: &'static str = ""; + + #[must_use] + pub fn markers() -> ToolCallMarkers { + ToolCallMarkers { + open: "".to_owned(), + close: "".to_owned(), + args_shape: ToolCallArgsShape::KeyValueXmlTags(KeyValueXmlTagsShape { + key_open: "".to_owned(), + key_close: "".to_owned(), + value_open: "".to_owned(), + value_close: "".to_owned(), + }), + } + } + + #[must_use] + pub fn detect(template: &str) -> Option { + if !template.contains(Self::TEMPLATE_FINGERPRINT) { + return None; + } + Some(Self::markers()) + } +} + +#[cfg(test)] +mod tests { + use llama_cpp_bindings_types::ToolCallArgsShape; + + use super::Glm47KeyValueTagsOverride; + + #[test] + fn detects_glm47_template_with_arg_key_literal() { + let template = "{{- '' + tool_call.name }}{% for k, v in args.items() %}{{ k }}{{ v }}{% endfor %}"; + let markers = + Glm47KeyValueTagsOverride::detect(template).expect("GLM-4.7 template must be detected"); + + assert_eq!(markers.open, ""); + assert_eq!(markers.close, ""); + let ToolCallArgsShape::KeyValueXmlTags(shape) = markers.args_shape else { + panic!( + "expected KeyValueXmlTags variant, got {:?}", + markers.args_shape + ); + }; + assert_eq!(shape.key_open, ""); + assert_eq!(shape.key_close, ""); + assert_eq!(shape.value_open, ""); + assert_eq!(shape.value_close, ""); + } + + #[test] + fn returns_none_for_template_without_fingerprint() { + assert!(Glm47KeyValueTagsOverride::detect("just some plain template body").is_none()); + } + + #[test] + fn returns_none_for_empty_template() { + assert!(Glm47KeyValueTagsOverride::detect("").is_none()); + } +} diff --git a/llama-cpp-bindings/src/tool_call_template_overrides/known_marker_candidates.rs b/llama-cpp-bindings/src/tool_call_template_overrides/known_marker_candidates.rs new file mode 100644 index 00000000..9448c866 --- /dev/null +++ b/llama-cpp-bindings/src/tool_call_template_overrides/known_marker_candidates.rs @@ -0,0 +1,54 @@ +use llama_cpp_bindings_types::ToolCallMarkers; + +use crate::tool_call_template_overrides::gemma4_call_block::Gemma4CallBlockOverride; +use crate::tool_call_template_overrides::glm47_key_value_tags::Glm47KeyValueTagsOverride; +use crate::tool_call_template_overrides::mistral3_arrow_args::Mistral3ArrowArgsOverride; +use crate::tool_call_template_overrides::qwen_xml_tags::QwenXmlTagsOverride; +use crate::tool_call_template_overrides::qwen3_json_inside_tool_call::Qwen3JsonInsideToolCallOverride; + +#[must_use] +pub fn known_marker_candidates() -> Vec { + vec![ + Qwen3JsonInsideToolCallOverride::markers(), + QwenXmlTagsOverride::markers(), + Glm47KeyValueTagsOverride::markers(), + Mistral3ArrowArgsOverride::markers(), + Gemma4CallBlockOverride::markers(), + ] +} + +#[cfg(test)] +mod tests { + use std::collections::HashSet; + + use llama_cpp_bindings_types::ToolCallArgsShape; + + use super::known_marker_candidates; + + #[test] + fn known_marker_candidates_returns_one_per_registered_shape() { + let candidates = known_marker_candidates(); + assert_eq!( + candidates.len(), + 5, + "expected exactly five registered shapes, got {}", + candidates.len() + ); + + let shape_discriminants: HashSet<&'static str> = candidates + .iter() + .map(|markers| match &markers.args_shape { + ToolCallArgsShape::BracketedJson(_) => "BracketedJson", + ToolCallArgsShape::JsonObject(_) => "JsonObject", + ToolCallArgsShape::KeyValueXmlTags(_) => "KeyValueXmlTags", + ToolCallArgsShape::PairedQuote(_) => "PairedQuote", + ToolCallArgsShape::XmlTags(_) => "XmlTags", + }) + .collect(); + assert_eq!( + shape_discriminants.len(), + 5, + "duplicate shape discriminants in known_marker_candidates: {shape_discriminants:?}" + ); + } +} diff --git a/llama-cpp-bindings/src/tool_call_template_overrides/mistral3_arrow_args.rs b/llama-cpp-bindings/src/tool_call_template_overrides/mistral3_arrow_args.rs new file mode 100644 index 00000000..3337a120 --- /dev/null +++ b/llama-cpp-bindings/src/tool_call_template_overrides/mistral3_arrow_args.rs @@ -0,0 +1,68 @@ +use llama_cpp_bindings_types::BracketedJsonShape; +use llama_cpp_bindings_types::ToolCallArgsShape; +use llama_cpp_bindings_types::ToolCallMarkers; + +pub struct Mistral3ArrowArgsOverride; + +impl Mistral3ArrowArgsOverride { + const TEMPLATE_FINGERPRINT: &'static str = "'[ARGS]'"; + + #[must_use] + pub fn markers() -> ToolCallMarkers { + ToolCallMarkers { + open: "[TOOL_CALLS]".to_owned(), + close: String::new(), + args_shape: ToolCallArgsShape::BracketedJson(BracketedJsonShape { + name_args_separator: "[ARGS]".to_owned(), + }), + } + } + + #[must_use] + pub fn detect(template: &str) -> Option { + if !template.contains(Self::TEMPLATE_FINGERPRINT) { + return None; + } + Some(Self::markers()) + } +} + +#[cfg(test)] +mod tests { + use llama_cpp_bindings_types::ToolCallArgsShape; + + use super::Mistral3ArrowArgsOverride; + + #[test] + fn detects_mistral3_template_with_args_literal() { + let template = "...{{- name + '[ARGS]' + arguments }}..."; + let markers = Mistral3ArrowArgsOverride::detect(template) + .expect("Mistral 3 template must be detected"); + + assert_eq!(markers.open, "[TOOL_CALLS]"); + assert!(markers.close.is_empty()); + let ToolCallArgsShape::BracketedJson(shape) = markers.args_shape else { + panic!( + "expected BracketedJson variant, got {:?}", + markers.args_shape + ); + }; + assert_eq!(shape.name_args_separator, "[ARGS]"); + } + + #[test] + fn returns_none_for_template_without_fingerprint() { + assert!(Mistral3ArrowArgsOverride::detect("just some plain template body").is_none()); + } + + #[test] + fn returns_none_for_empty_template() { + assert!(Mistral3ArrowArgsOverride::detect("").is_none()); + } + + #[test] + fn returns_none_when_fingerprint_substring_appears_without_jinja_apostrophes() { + let template = "doc text mentioning the [ARGS] tag without quoting it as a literal"; + assert!(Mistral3ArrowArgsOverride::detect(template).is_none()); + } +} diff --git a/llama-cpp-bindings/src/tool_call_template_overrides/mod.rs b/llama-cpp-bindings/src/tool_call_template_overrides/mod.rs new file mode 100644 index 00000000..b8717ad5 --- /dev/null +++ b/llama-cpp-bindings/src/tool_call_template_overrides/mod.rs @@ -0,0 +1,10 @@ +pub mod detect; +pub mod gemma4_call_block; +pub mod glm47_key_value_tags; +pub mod known_marker_candidates; +pub mod mistral3_arrow_args; +pub mod qwen3_json_inside_tool_call; +pub mod qwen_xml_tags; + +pub use detect::detect; +pub use known_marker_candidates::known_marker_candidates; diff --git a/llama-cpp-bindings/src/tool_call_template_overrides/qwen3_json_inside_tool_call.rs b/llama-cpp-bindings/src/tool_call_template_overrides/qwen3_json_inside_tool_call.rs new file mode 100644 index 00000000..7ac4bda6 --- /dev/null +++ b/llama-cpp-bindings/src/tool_call_template_overrides/qwen3_json_inside_tool_call.rs @@ -0,0 +1,80 @@ +use llama_cpp_bindings_types::JsonObjectShape; +use llama_cpp_bindings_types::ToolCallArgsShape; +use llama_cpp_bindings_types::ToolCallMarkers; + +pub struct Qwen3JsonInsideToolCallOverride; + +impl Qwen3JsonInsideToolCallOverride { + const TEMPLATE_FINGERPRINT_OPEN: &'static str = "'\\n{\"name\": \"'"; + const TEMPLATE_FINGERPRINT_ARGS_JOIN: &'static str = "'\", \"arguments\": '"; + + #[must_use] + pub fn markers() -> ToolCallMarkers { + ToolCallMarkers { + open: "".to_owned(), + close: "".to_owned(), + args_shape: ToolCallArgsShape::JsonObject(JsonObjectShape { + name_field: "name".to_owned(), + arguments_field: "arguments".to_owned(), + }), + } + } + + #[must_use] + pub fn detect(template: &str) -> Option { + if !template.contains(Self::TEMPLATE_FINGERPRINT_OPEN) { + return None; + } + if !template.contains(Self::TEMPLATE_FINGERPRINT_ARGS_JOIN) { + return None; + } + Some(Self::markers()) + } +} + +#[cfg(test)] +mod tests { + use llama_cpp_bindings_types::ToolCallArgsShape; + + use super::Qwen3JsonInsideToolCallOverride; + + #[test] + fn detects_qwen3_json_inside_tool_call_template() { + let template = "{{- '\\n{\"name\": \"' + tool_call.name + '\", \"arguments\": ' + (tool_call.arguments | tojson) + '}\\n' -}}"; + let markers = Qwen3JsonInsideToolCallOverride::detect(template) + .expect("Qwen 3 template must be detected"); + + assert_eq!(markers.open, ""); + assert_eq!(markers.close, ""); + let ToolCallArgsShape::JsonObject(shape) = markers.args_shape else { + panic!("expected JsonObject variant, got {:?}", markers.args_shape); + }; + assert_eq!(shape.name_field, "name"); + assert_eq!(shape.arguments_field, "arguments"); + } + + #[test] + fn returns_none_for_template_without_fingerprint() { + assert!(Qwen3JsonInsideToolCallOverride::detect("just some plain template body").is_none()); + } + + #[test] + fn returns_none_for_empty_template() { + assert!(Qwen3JsonInsideToolCallOverride::detect("").is_none()); + } + + #[test] + fn returns_none_when_only_open_fingerprint_present() { + let template = "{{- '\\n{\"name\": \"' + tool_call.name + ..."; + assert!( + Qwen3JsonInsideToolCallOverride::detect(template).is_none(), + "open fingerprint alone must not match (Qwen3-Embedding-style false positive)", + ); + } + + #[test] + fn returns_none_when_only_args_join_fingerprint_present() { + let template = "some text '\", \"arguments\": ' more text"; + assert!(Qwen3JsonInsideToolCallOverride::detect(template).is_none()); + } +} diff --git a/llama-cpp-bindings/src/tool_call_template_overrides/qwen_xml_tags.rs b/llama-cpp-bindings/src/tool_call_template_overrides/qwen_xml_tags.rs new file mode 100644 index 00000000..b0d013fe --- /dev/null +++ b/llama-cpp-bindings/src/tool_call_template_overrides/qwen_xml_tags.rs @@ -0,0 +1,71 @@ +use llama_cpp_bindings_types::ToolCallArgsShape; +use llama_cpp_bindings_types::ToolCallMarkers; +use llama_cpp_bindings_types::XmlTagsShape; + +pub struct QwenXmlTagsOverride; + +impl QwenXmlTagsOverride { + const TEMPLATE_FINGERPRINT: &'static str = " ToolCallMarkers { + ToolCallMarkers { + open: "".to_owned(), + close: "".to_owned(), + args_shape: ToolCallArgsShape::XmlTags(XmlTagsShape { + function_open_prefix: "".to_owned(), + parameter_open_prefix: "".to_owned(), + }), + } + } + + #[must_use] + pub fn detect(template: &str) -> Option { + if !template.contains(Self::TEMPLATE_FINGERPRINT) { + return None; + } + Some(Self::markers()) + } +} + +#[cfg(test)] +mod tests { + use llama_cpp_bindings_types::ToolCallArgsShape; + + use super::QwenXmlTagsOverride; + + #[test] + fn detects_qwen_xml_template_with_function_tag_literal() { + let template = "{{- '\\n\\n' }}"; + let markers = + QwenXmlTagsOverride::detect(template).expect("Qwen XML template must be detected"); + + assert_eq!(markers.open, ""); + assert_eq!(markers.close, ""); + let ToolCallArgsShape::XmlTags(shape) = markers.args_shape else { + panic!("expected XmlTags variant, got {:?}", markers.args_shape); + }; + assert_eq!(shape.function_open_prefix, ""); + assert_eq!(shape.parameter_open_prefix, ""); + } + + #[test] + fn returns_none_for_template_without_fingerprint() { + assert!(QwenXmlTagsOverride::detect("just some plain template body").is_none()); + } + + #[test] + fn returns_none_for_empty_template() { + assert!(QwenXmlTagsOverride::detect("").is_none()); + } + + #[test] + fn detects_qwen_xml_template_with_concatenated_string_literal() { + let template = "{{- '\\n\\n\\n\\n' }}"; + assert!(QwenXmlTagsOverride::detect(template).is_some()); + } +}