From 15d1fbf5e28808e015eaff1eabbf2c03aed117a5 Mon Sep 17 00:00:00 2001 From: Mateusz Charytoniuk Date: Fri, 28 Aug 2026 00:27:26 +0200 Subject: [PATCH 1/2] let the tool-call opener win over a reasoning close and treat an absent KV cache as already cleared --- .../tests/kv_cache_without_memory_module.rs | 35 +++++++++++++ llama-cpp-bindings-tests/tests/main.rs | 1 + llama-cpp-bindings/src/context/kv_cache.rs | 34 ++++++++----- .../src/sampled_token_classifier.rs | 51 ++++++++++++++++--- 4 files changed, 103 insertions(+), 18 deletions(-) create mode 100644 llama-cpp-bindings-tests/tests/kv_cache_without_memory_module.rs diff --git a/llama-cpp-bindings-tests/tests/kv_cache_without_memory_module.rs b/llama-cpp-bindings-tests/tests/kv_cache_without_memory_module.rs new file mode 100644 index 00000000..da64e4cb --- /dev/null +++ b/llama-cpp-bindings-tests/tests/kv_cache_without_memory_module.rs @@ -0,0 +1,35 @@ +use anyhow::Context; +use anyhow::Result; +use llama_cpp_bindings::context::LlamaContext; +use llama_cpp_test_harness::LlamaFixture; +use llama_cpp_test_harness::llama_test; + +#[llama_test( + model_source = HuggingFace("nomic-ai/nomic-embed-text-v1.5-GGUF", "nomic-embed-text-v1.5.Q4_K_M.gguf"), + n_gpu_layers = 999, + load_mode = Mmap, + n_ctx = 512, + n_batch = 2048, + n_ubatch = 512, + n_threads_batch = 8, + embeddings = true, +)] +fn kv_cache_mutations_succeed_when_the_context_has_no_memory_module( + fixture: &LlamaFixture<'_>, +) -> Result<()> { + let mut ctx = LlamaContext::from_model( + fixture.model, + fixture.backend, + (*fixture.context_params).into_llama_context_params(), + ) + .context("unable to create context")?; + + ctx.clear_kv_cache() + .context("clearing an absent KV cache must succeed")?; + ctx.clear_kv_cache_seq(Some(0), None, None) + .context("removing a sequence from an absent KV cache must succeed")?; + ctx.kv_cache_seq_keep(0) + .context("keeping a sequence in an absent KV cache must succeed")?; + + Ok(()) +} diff --git a/llama-cpp-bindings-tests/tests/main.rs b/llama-cpp-bindings-tests/tests/main.rs index 6b07d941..29305e20 100644 --- a/llama-cpp-bindings-tests/tests/main.rs +++ b/llama-cpp-bindings-tests/tests/main.rs @@ -7,6 +7,7 @@ mod chat_protocol; mod context_state; mod embedding_models; mod generation_control; +mod kv_cache_without_memory_module; mod model_introspection; mod model_loading_errors; mod multimodal_audio; diff --git a/llama-cpp-bindings/src/context/kv_cache.rs b/llama-cpp-bindings/src/context/kv_cache.rs index 85dfce32..4a66fe60 100644 --- a/llama-cpp-bindings/src/context/kv_cache.rs +++ b/llama-cpp-bindings/src/context/kv_cache.rs @@ -174,13 +174,14 @@ impl LlamaContext<'_> { fn memory_handle( &self, ) -> Result { - let mem = unsafe { llama_cpp_bindings_sys::llama_get_memory(self.context.as_ptr()) }; + self.optional_memory_handle() + .ok_or(KvCacheConversionError::MemoryHandleUnavailable) + } - if mem.is_null() { - return Err(KvCacheConversionError::MemoryHandleUnavailable); - } + fn optional_memory_handle(&self) -> Option { + let mem = unsafe { llama_cpp_bindings_sys::llama_get_memory(self.context.as_ptr()) }; - Ok(mem) + if mem.is_null() { None } else { Some(mem) } } /// # Errors @@ -205,8 +206,9 @@ impl LlamaContext<'_> { } /// # Errors - /// If the sequence id or either position exceeds [`i32::MAX`], the context has no - /// memory module, or llama.cpp reports that the partial sequence could not be removed. + /// If the sequence id or either position exceeds [`i32::MAX`], or llama.cpp reports + /// that the partial sequence could not be removed. A context without a memory module + /// holds no KV cache, so the removal trivially succeeds. pub fn clear_kv_cache_seq( &mut self, src: Option, @@ -222,7 +224,9 @@ impl LlamaContext<'_> { let p1 = p1 .map_or(Ok(-1), i32::try_from) .map_err(KvCacheConversionError::P1TooLarge)?; - let mem = self.memory_handle()?; + let Some(mem) = self.optional_memory_handle() else { + return Ok(()); + }; if unsafe { llama_cpp_bindings_sys::llama_memory_seq_rm(mem, src, p0, p1) } { return Ok(()); @@ -236,9 +240,12 @@ impl LlamaContext<'_> { } /// # Errors - /// If the context has no memory module. + /// Never returns an error: a context without a memory module holds no KV cache, so + /// there is nothing to clear. pub fn clear_kv_cache(&mut self) -> Result<(), KvCacheConversionError> { - let mem = self.memory_handle()?; + let Some(mem) = self.optional_memory_handle() else { + return Ok(()); + }; let clear_data_buffers = true; unsafe { llama_cpp_bindings_sys::llama_memory_clear(mem, clear_data_buffers) }; @@ -246,9 +253,12 @@ impl LlamaContext<'_> { } /// # Errors - /// If the context has no memory module. + /// Never returns an error: a context without a memory module holds no KV cache, so + /// every sequence is already absent from it. pub fn kv_cache_seq_keep(&mut self, seq_id: i32) -> Result<(), KvCacheConversionError> { - let mem = self.memory_handle()?; + let Some(mem) = self.optional_memory_handle() else { + return Ok(()); + }; unsafe { llama_cpp_bindings_sys::llama_memory_seq_keep(mem, seq_id) }; Ok(()) diff --git a/llama-cpp-bindings/src/sampled_token_classifier.rs b/llama-cpp-bindings/src/sampled_token_classifier.rs index ceb273ca..8061fa36 100644 --- a/llama-cpp-bindings/src/sampled_token_classifier.rs +++ b/llama-cpp-bindings/src/sampled_token_classifier.rs @@ -174,12 +174,6 @@ impl<'model> SampledTokenClassifier<'model> { .as_deref() .and_then(|marker| self.marker_span_start(marker)) .map(|span_start| (span_start, MarkerKind::ReasoningOpen)) - .or_else(|| { - self.markers.reasoning_closes.iter().find_map(|marker| { - self.marker_span_start(marker) - .map(|span_start| (span_start, MarkerKind::ReasoningClose)) - }) - }) .or_else(|| { self.markers .tool_call_open @@ -187,6 +181,12 @@ impl<'model> SampledTokenClassifier<'model> { .and_then(|marker| self.marker_span_start(marker)) .map(|span_start| (span_start, MarkerKind::ToolCallOpen)) }) + .or_else(|| { + self.markers.reasoning_closes.iter().find_map(|marker| { + self.marker_span_start(marker) + .map(|span_start| (span_start, MarkerKind::ReasoningClose)) + }) + }) .or_else(|| { self.markers .tool_call_close @@ -639,6 +639,45 @@ mod tests { .collect() } + fn markers_sharing_tool_call_open_with_a_reasoning_close( + shared: Vec, + ) -> StreamingMarkers { + StreamingMarkers { + reasoning_open: Some(vec![token(100)]), + reasoning_closes: vec![vec![token(200)], shared.clone()], + tool_call_open: Some(shared), + tool_call_close: Some(vec![token(201)]), + } + } + + #[test] + fn a_token_that_is_both_a_reasoning_close_and_the_tool_call_open_opens_the_tool_call() { + let mut classifier = + synthetic_classifier(markers_sharing_tool_call_open_with_a_reasoning_close(vec![ + token(300), + ])); + classifier.section = SampledTokenSection::Content; + + push_pending(&mut classifier, 300, ""); + classifier.try_consume_marker_at_tail(); + + assert_eq!(classifier.section, SampledTokenSection::ToolCall); + } + + #[test] + fn a_shared_marker_ends_reasoning_by_opening_the_tool_call() { + let mut classifier = + synthetic_classifier(markers_sharing_tool_call_open_with_a_reasoning_close(vec![ + token(300), + ])); + classifier.section = SampledTokenSection::Reasoning; + + push_pending(&mut classifier, 300, ""); + classifier.try_consume_marker_at_tail(); + + assert_eq!(classifier.section, SampledTokenSection::ToolCall); + } + #[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)])); From a4765218fb86d778935bfc7aeb1808da7899b091 Mon Sep 17 00:00:00 2001 From: Mateusz Charytoniuk Date: Fri, 28 Aug 2026 13:51:37 +0200 Subject: [PATCH 2/2] Redesign streaming markers and optional KV cache APIs --- .../tests/context_state.rs | 21 +- .../tests/embedding_models.rs | 57 ++++- .../tests/generation_control.rs | 28 +-- .../tests/kv_cache_without_memory_module.rs | 35 --- llama-cpp-bindings-tests/tests/main.rs | 1 - .../tests/structured_chat_output.rs | 40 ++- llama-cpp-bindings/src/context/kv_cache.rs | 77 +++--- llama-cpp-bindings/src/error.rs | 6 +- ...n_error.rs => clear_kv_cache_seq_error.rs} | 10 +- .../src/error/copy_kv_cache_seq_error.rs | 11 + .../src/error/marker_detection_error.rs | 5 + llama-cpp-bindings/src/lib.rs | 20 +- llama-cpp-bindings/src/marker_kind.rs | 7 - llama-cpp-bindings/src/marker_role.rs | 28 +++ llama-cpp-bindings/src/model.rs | 22 +- .../src/sampled_token_classifier.rs | 237 +++++++++--------- llama-cpp-bindings/src/streaming_marker.rs | 67 +++++ llama-cpp-bindings/src/streaming_markers.rs | 179 ++++++++++--- 18 files changed, 553 insertions(+), 298 deletions(-) delete mode 100644 llama-cpp-bindings-tests/tests/kv_cache_without_memory_module.rs rename llama-cpp-bindings/src/error/{kv_cache_conversion_error.rs => clear_kv_cache_seq_error.rs} (56%) create mode 100644 llama-cpp-bindings/src/error/copy_kv_cache_seq_error.rs delete mode 100644 llama-cpp-bindings/src/marker_kind.rs create mode 100644 llama-cpp-bindings/src/marker_role.rs create mode 100644 llama-cpp-bindings/src/streaming_marker.rs diff --git a/llama-cpp-bindings-tests/tests/context_state.rs b/llama-cpp-bindings-tests/tests/context_state.rs index 6a52a4da..b9efaac7 100644 --- a/llama-cpp-bindings-tests/tests/context_state.rs +++ b/llama-cpp-bindings-tests/tests/context_state.rs @@ -6,7 +6,8 @@ use anyhow::Result; use llama_cpp_bindings::DecodeError; use llama_cpp_bindings::LogitsError; use llama_cpp_bindings::context::LlamaContext; -use llama_cpp_bindings::error::KvCacheConversionError; +use llama_cpp_bindings::error::ClearKvCacheSeqError; +use llama_cpp_bindings::error::CopyKvCacheSeqError; use llama_cpp_bindings::error::KvCacheSeqAddError; use llama_cpp_bindings::error::KvCacheSeqDivError; use llama_cpp_bindings::llama_batch::LlamaBatch; @@ -726,7 +727,7 @@ fn clear_kv_cache_resets_positions(fixture: &LlamaFixture<'_>) -> Result<()> { prime_kv_cache(fixture, &mut context)?; - context.clear_kv_cache()?; + context.clear_kv_cache(); assert_eq!(context.kv_cache_seq_pos_max(0)?, -1); Ok(()) @@ -1000,7 +1001,7 @@ fn kv_cache_seq_keep_retains_specified_sequence(fixture: &LlamaFixture<'_>) -> R prime_kv_cache(fixture, &mut context)?; - context.kv_cache_seq_keep(0)?; + context.kv_cache_seq_keep(0); assert!(context.kv_cache_seq_pos_max(0)? >= 0); @@ -1133,7 +1134,7 @@ fn copy_kv_cache_seq_rejects_p0_exceeding_i32_max(fixture: &LlamaFixture<'_>) -> assert_eq!( result.unwrap_err(), - KvCacheConversionError::P0TooLarge( + CopyKvCacheSeqError::P0TooLarge( i32::try_from(u32::MAX).expect_err("u32::MAX does not fit into i32") ) ); @@ -1180,7 +1181,7 @@ fn copy_kv_cache_seq_rejects_p1_exceeding_i32_max(fixture: &LlamaFixture<'_>) -> assert_eq!( result.unwrap_err(), - KvCacheConversionError::P1TooLarge( + CopyKvCacheSeqError::P1TooLarge( i32::try_from(u32::MAX).expect_err("u32::MAX does not fit into i32") ) ); @@ -1220,14 +1221,16 @@ fn copy_kv_cache_seq_rejects_p1_exceeding_i32_max(fixture: &LlamaFixture<'_>) -> n_batch = 512, n_ubatch = 128, )] -fn clear_kv_cache_seq_rejects_src_exceeding_i32_max(fixture: &LlamaFixture<'_>) -> Result<()> { +fn clear_kv_cache_seq_rejects_sequence_id_exceeding_i32_max( + fixture: &LlamaFixture<'_>, +) -> Result<()> { let mut context = fixture.build_context()?; let result = context.clear_kv_cache_seq(Some(u32::MAX), None, None); assert_eq!( result.unwrap_err(), - KvCacheConversionError::SeqIdTooLarge( + ClearKvCacheSeqError::SeqIdTooLarge( i32::try_from(u32::MAX).expect_err("u32::MAX does not fit into i32") ) ); @@ -1274,7 +1277,7 @@ fn clear_kv_cache_seq_rejects_p0_exceeding_i32_max(fixture: &LlamaFixture<'_>) - assert_eq!( result.unwrap_err(), - KvCacheConversionError::P0TooLarge( + ClearKvCacheSeqError::P0TooLarge( i32::try_from(u32::MAX).expect_err("u32::MAX does not fit into i32") ) ); @@ -1321,7 +1324,7 @@ fn clear_kv_cache_seq_rejects_p1_exceeding_i32_max(fixture: &LlamaFixture<'_>) - assert_eq!( result.unwrap_err(), - KvCacheConversionError::P1TooLarge( + ClearKvCacheSeqError::P1TooLarge( i32::try_from(u32::MAX).expect_err("u32::MAX does not fit into i32") ) ); diff --git a/llama-cpp-bindings-tests/tests/embedding_models.rs b/llama-cpp-bindings-tests/tests/embedding_models.rs index 98f9ec9a..12a8d074 100644 --- a/llama-cpp-bindings-tests/tests/embedding_models.rs +++ b/llama-cpp-bindings-tests/tests/embedding_models.rs @@ -4,6 +4,11 @@ use std::time::Duration; use anyhow::Context; use anyhow::Result; use anyhow::bail; +use llama_cpp_bindings::ClearKvCacheSeqError; +use llama_cpp_bindings::CopyKvCacheSeqError; +use llama_cpp_bindings::KvCacheSeqAddError; +use llama_cpp_bindings::KvCacheSeqDivError; +use llama_cpp_bindings::KvCacheSeqPosMaxError; use llama_cpp_bindings::context::LlamaContext; use llama_cpp_bindings::ggml_time_us; use llama_cpp_bindings::llama_batch::LlamaBatch; @@ -67,7 +72,7 @@ fn embedding_generation_produces_vectors(fixture: &LlamaFixture<'_>) -> Result<( assert_eq!(classifier.pending_prompt_tokens(), prompt_token_count); assert_eq!(classifier.usage().prompt_tokens, 0); - ctx.clear_kv_cache()?; + ctx.clear_kv_cache(); ctx.decode(&mut batch) .with_context(|| "llama_decode() failed")?; @@ -180,7 +185,7 @@ fn reranking_produces_scores(fixture: &LlamaFixture<'_>) -> Result<()> { assert_eq!(classifier.pending_prompt_tokens(), total_token_count); assert_eq!(classifier.usage().prompt_tokens, 0); - ctx.clear_kv_cache()?; + ctx.clear_kv_cache(); ctx.decode(&mut batch) .with_context(|| "llama_decode() failed")?; @@ -393,7 +398,7 @@ fn embeddings_returns_distinct_values_when_reused_batch_has_extra_capacity( batch.add_sequence(&tokens, sequence_id, true)?; } - context.clear_kv_cache()?; + context.clear_kv_cache(); context.decode(&mut batch)?; for sequence_index in 0..iteration_inputs.len() { @@ -581,7 +586,51 @@ fn embedding_model_exposes_tool_call_markers(fixture: &LlamaFixture<'_>) -> Resu fn embedding_model_exposes_streaming_markers(fixture: &LlamaFixture<'_>) -> Result<()> { let markers = fixture.model.streaming_markers()?; - assert!(markers.has_any()); + assert!(!markers.is_empty()); + + Ok(()) +} + +#[llama_test( + model_source = HuggingFace("nomic-ai/nomic-embed-text-v1.5-GGUF", "nomic-embed-text-v1.5.Q4_K_M.gguf"), + n_gpu_layers = 999, + load_mode = Mmap, + n_ctx = 512, + n_batch = 2048, + n_ubatch = 512, + n_threads_batch = 8, + embeddings = true, +)] +fn kv_cache_operations_respect_an_embedding_context_without_memory( + fixture: &LlamaFixture<'_>, +) -> Result<()> { + let mut context = fixture.build_context()?; + + context.clear_kv_cache(); + assert_eq!(context.clear_kv_cache_seq(Some(0), None, None), Ok(())); + context.kv_cache_seq_keep(0); + + assert_eq!( + context.copy_kv_cache_seq(0, 1, None, None), + Err(CopyKvCacheSeqError::MemoryHandleUnavailable) + ); + assert_eq!( + context.kv_cache_seq_add(0, None, None, 1), + Err(KvCacheSeqAddError::MemoryHandleUnavailable) + ); + let divisor = NonZeroU8::new(2).ok_or_else(|| anyhow::anyhow!("2 is non-zero"))?; + assert_eq!( + context.kv_cache_seq_div(0, None, None, divisor), + Err(KvCacheSeqDivError::MemoryHandleUnavailable) + ); + assert_eq!( + context.kv_cache_seq_pos_max(0), + Err(KvCacheSeqPosMaxError::MemoryHandleUnavailable) + ); + assert!(matches!( + context.clear_kv_cache_seq(Some(u32::MAX), None, None), + Err(ClearKvCacheSeqError::SeqIdTooLarge(_)) + )); Ok(()) } diff --git a/llama-cpp-bindings-tests/tests/generation_control.rs b/llama-cpp-bindings-tests/tests/generation_control.rs index ccd7342e..2fe17873 100644 --- a/llama-cpp-bindings-tests/tests/generation_control.rs +++ b/llama-cpp-bindings-tests/tests/generation_control.rs @@ -15,10 +15,8 @@ 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::model::LlamaChatMessage; -use llama_cpp_bindings::sampled_token_classifier::SampledTokenClassifier; use llama_cpp_bindings::sampled_token_section::SampledTokenSection; use llama_cpp_bindings::sampling::LlamaSampler; -use llama_cpp_bindings::streaming_markers::StreamingMarkers; use llama_cpp_bindings::token::LlamaToken; use llama_cpp_bindings_tests::classify_sample_loop::ClassifySampleLoop; use llama_cpp_test_harness::LlamaFixture; @@ -1903,13 +1901,14 @@ fn classifier_construction_is_idempotent_across_calls(fixture: &LlamaFixture<'_> n_batch = 128, n_ubatch = 64, )] -fn ingest_with_no_markers_emits_undeterminable_with_visible_and_raw_piece( +fn ingest_flushes_an_unmatched_token_with_its_visible_and_raw_piece( fixture: &LlamaFixture<'_>, ) -> Result<()> { let model = fixture.model; - let mut classifier = SampledTokenClassifier::new(model, StreamingMarkers::default()); + let mut classifier = model.sampled_token_classifier()?; - let outcomes = classifier.ingest(model.token_bos())?; + let mut outcomes = classifier.ingest(model.token_bos())?; + outcomes.extend(classifier.flush()); assert_eq!(outcomes.len(), 1); let outcome = &outcomes[0]; @@ -1954,14 +1953,13 @@ fn ingest_with_no_markers_emits_undeterminable_with_visible_and_raw_piece( n_batch = 128, n_ubatch = 64, )] -fn ingest_with_no_markers_decodes_each_token_independently( - fixture: &LlamaFixture<'_>, -) -> Result<()> { +fn ingest_accounts_for_each_unmatched_token_after_flush(fixture: &LlamaFixture<'_>) -> Result<()> { let model = fixture.model; - let mut classifier = SampledTokenClassifier::new(model, StreamingMarkers::default()); + let mut classifier = model.sampled_token_classifier()?; classifier.ingest(model.token_bos())?; classifier.ingest(model.token_eos())?; + classifier.flush(); assert_eq!(classifier.usage().undeterminable_tokens, 2); Ok(()) @@ -1999,9 +1997,9 @@ fn ingest_with_no_markers_decodes_each_token_independently( n_batch = 128, n_ubatch = 64, )] -fn ingest_prompt_token_with_no_markers_is_a_noop(fixture: &LlamaFixture<'_>) -> Result<()> { +fn ingest_unmatched_prompt_tokens_does_not_record_usage(fixture: &LlamaFixture<'_>) -> Result<()> { let model = fixture.model; - let mut classifier = SampledTokenClassifier::new(model, StreamingMarkers::default()); + let mut classifier = model.sampled_token_classifier()?; let usage_before = *classifier.usage(); classifier.ingest_prompt_token(model.token_bos()); @@ -2046,7 +2044,7 @@ fn ingest_prompt_token_with_no_markers_is_a_noop(fixture: &LlamaFixture<'_>) -> )] fn feed_prompt_to_batch_increments_pending_prompt_tokens(fixture: &LlamaFixture<'_>) -> Result<()> { let model = fixture.model; - let mut classifier = SampledTokenClassifier::new(model, StreamingMarkers::default()); + let mut classifier = model.sampled_token_classifier()?; let mut batch = LlamaBatch::new(8, 1)?; classifier.feed_prompt_to_batch(&mut batch, model.token_bos(), 0, &[0], false)?; @@ -2092,7 +2090,7 @@ fn feed_prompt_to_batch_increments_pending_prompt_tokens(fixture: &LlamaFixture< )] fn feed_prompt_sequence_to_batch_stages_all_tokens(fixture: &LlamaFixture<'_>) -> Result<()> { let model = fixture.model; - let mut classifier = SampledTokenClassifier::new(model, StreamingMarkers::default()); + let mut classifier = model.sampled_token_classifier()?; let mut batch = LlamaBatch::new(8, 1)?; let tokens = vec![model.token_bos(), model.token_eos(), model.token_nl()]; @@ -2140,7 +2138,7 @@ fn commit_prompt_tokens_promotes_pending_count_to_usage_and_clears( fixture: &LlamaFixture<'_>, ) -> Result<()> { let model = fixture.model; - let mut classifier = SampledTokenClassifier::new(model, StreamingMarkers::default()); + let mut classifier = model.sampled_token_classifier()?; let mut batch = LlamaBatch::new(8, 1)?; classifier.feed_prompt_to_batch(&mut batch, model.token_bos(), 0, &[0], false)?; @@ -2191,7 +2189,7 @@ fn discard_pending_prompt_tokens_clears_count_without_recording_usage( fixture: &LlamaFixture<'_>, ) -> Result<()> { let model = fixture.model; - let mut classifier = SampledTokenClassifier::new(model, StreamingMarkers::default()); + let mut classifier = model.sampled_token_classifier()?; let mut batch = LlamaBatch::new(8, 1)?; classifier.feed_prompt_to_batch(&mut batch, model.token_bos(), 0, &[0], false)?; diff --git a/llama-cpp-bindings-tests/tests/kv_cache_without_memory_module.rs b/llama-cpp-bindings-tests/tests/kv_cache_without_memory_module.rs deleted file mode 100644 index da64e4cb..00000000 --- a/llama-cpp-bindings-tests/tests/kv_cache_without_memory_module.rs +++ /dev/null @@ -1,35 +0,0 @@ -use anyhow::Context; -use anyhow::Result; -use llama_cpp_bindings::context::LlamaContext; -use llama_cpp_test_harness::LlamaFixture; -use llama_cpp_test_harness::llama_test; - -#[llama_test( - model_source = HuggingFace("nomic-ai/nomic-embed-text-v1.5-GGUF", "nomic-embed-text-v1.5.Q4_K_M.gguf"), - n_gpu_layers = 999, - load_mode = Mmap, - n_ctx = 512, - n_batch = 2048, - n_ubatch = 512, - n_threads_batch = 8, - embeddings = true, -)] -fn kv_cache_mutations_succeed_when_the_context_has_no_memory_module( - fixture: &LlamaFixture<'_>, -) -> Result<()> { - let mut ctx = LlamaContext::from_model( - fixture.model, - fixture.backend, - (*fixture.context_params).into_llama_context_params(), - ) - .context("unable to create context")?; - - ctx.clear_kv_cache() - .context("clearing an absent KV cache must succeed")?; - ctx.clear_kv_cache_seq(Some(0), None, None) - .context("removing a sequence from an absent KV cache must succeed")?; - ctx.kv_cache_seq_keep(0) - .context("keeping a sequence in an absent KV cache must succeed")?; - - Ok(()) -} diff --git a/llama-cpp-bindings-tests/tests/main.rs b/llama-cpp-bindings-tests/tests/main.rs index 29305e20..6b07d941 100644 --- a/llama-cpp-bindings-tests/tests/main.rs +++ b/llama-cpp-bindings-tests/tests/main.rs @@ -7,7 +7,6 @@ mod chat_protocol; mod context_state; mod embedding_models; mod generation_control; -mod kv_cache_without_memory_module; mod model_introspection; mod model_loading_errors; mod multimodal_audio; diff --git a/llama-cpp-bindings-tests/tests/structured_chat_output.rs b/llama-cpp-bindings-tests/tests/structured_chat_output.rs index f7754d8f..5ffb3362 100644 --- a/llama-cpp-bindings-tests/tests/structured_chat_output.rs +++ b/llama-cpp-bindings-tests/tests/structured_chat_output.rs @@ -1,7 +1,9 @@ use anyhow::Result; use anyhow::bail; use llama_cpp_bindings::ChatMessageParseOutcome; +use llama_cpp_bindings::MarkerRole; use llama_cpp_bindings::ParsedChatMessage; +use llama_cpp_bindings::SampledTokenSection; use llama_cpp_bindings::TokenUsage; use llama_cpp_bindings::ToolCallArgsShape; use llama_cpp_bindings::ToolCallArguments; @@ -1465,7 +1467,7 @@ fn qwen35_chat_inference_emits_reasoning_when_template_auto_opens( n_batch = 512, n_ubatch = 128, )] -fn qwen35_streaming_markers_tokenize_every_reasoning_boundary( +fn qwen35_shared_reasoning_close_and_tool_call_open_is_one_transition( fixture: &LlamaFixture<'_>, ) -> Result<()> { let reasoning_markers = fixture @@ -1474,18 +1476,40 @@ fn qwen35_streaming_markers_tokenize_every_reasoning_boundary( .expect("Qwen3.5 must expose reasoning markers"); let streaming_markers = fixture.model.streaming_markers()?; - assert!(streaming_markers.reasoning_open.is_some()); + assert!(streaming_markers.iter().any(|marker| { + marker.roles().contains(&MarkerRole::ReasoningOpen) && !marker.tokens().is_empty() + })); assert_eq!( - streaming_markers.reasoning_closes.len(), - reasoning_markers.closes.len() - ); - assert!( streaming_markers - .reasoning_closes .iter() - .all(|tokens| !tokens.is_empty()) + .filter(|marker| marker.roles().contains(&MarkerRole::ReasoningClose)) + .count(), + reasoning_markers.closes.len() ); + let reasoning_open = streaming_markers + .iter() + .find(|marker| marker.roles().contains(&MarkerRole::ReasoningOpen)) + .expect("Qwen3.5 must expose a reasoning opener") + .tokens() + .to_vec(); + let shared_boundary = streaming_markers + .iter() + .find(|marker| { + marker.roles().contains(&MarkerRole::ReasoningClose) + && marker.roles().contains(&MarkerRole::ToolCallOpen) + }) + .expect("Qwen3.5 must share its tool-call opener with a reasoning close") + .tokens() + .to_vec(); + + let mut classifier = fixture.model.sampled_token_classifier()?; + classifier.ingest_prompt_tokens(&reasoning_open); + assert_eq!(classifier.current_section(), SampledTokenSection::Reasoning); + + classifier.ingest_prompt_tokens(&shared_boundary); + assert_eq!(classifier.current_section(), SampledTokenSection::ToolCall); + Ok(()) } diff --git a/llama-cpp-bindings/src/context/kv_cache.rs b/llama-cpp-bindings/src/context/kv_cache.rs index 4a66fe60..bfebba7d 100644 --- a/llama-cpp-bindings/src/context/kv_cache.rs +++ b/llama-cpp-bindings/src/context/kv_cache.rs @@ -4,8 +4,10 @@ use std::os::raw::c_char; use std::ptr; use crate::context::LlamaContext; -use crate::error::kv_cache_conversion_error::KvCacheConversionError; -use crate::error::{KvCacheSeqAddError, KvCacheSeqDivError, KvCacheSeqPosMaxError}; +use crate::error::{ + ClearKvCacheSeqError, CopyKvCacheSeqError, KvCacheSeqAddError, KvCacheSeqDivError, + KvCacheSeqPosMaxError, +}; use llama_cpp_ffi_status::read_and_free_cpp_string; fn kv_cache_seq_add_status_to_result( @@ -168,22 +170,19 @@ fn kv_cache_seq_pos_max_status_to_result( } impl LlamaContext<'_> { - /// # Errors - /// Returns [`KvCacheConversionError::MemoryHandleUnavailable`] when the context was - /// built without a memory module, so a null handle is never handed to llama.cpp. - fn memory_handle( - &self, - ) -> Result { - self.optional_memory_handle() - .ok_or(KvCacheConversionError::MemoryHandleUnavailable) - } - - fn optional_memory_handle(&self) -> Option { + fn memory_handle(&self) -> Option { let mem = unsafe { llama_cpp_bindings_sys::llama_get_memory(self.context.as_ptr()) }; if mem.is_null() { None } else { Some(mem) } } + fn required_memory_handle( + &self, + ) -> Result { + self.memory_handle() + .ok_or(CopyKvCacheSeqError::MemoryHandleUnavailable) + } + /// # Errors /// If either position exceeds [`i32::MAX`], or the context has no memory module. pub fn copy_kv_cache_seq( @@ -192,14 +191,14 @@ impl LlamaContext<'_> { dest: i32, p0: Option, p1: Option, - ) -> Result<(), KvCacheConversionError> { + ) -> Result<(), CopyKvCacheSeqError> { let p0 = p0 .map_or(Ok(-1), i32::try_from) - .map_err(KvCacheConversionError::P0TooLarge)?; + .map_err(CopyKvCacheSeqError::P0TooLarge)?; let p1 = p1 .map_or(Ok(-1), i32::try_from) - .map_err(KvCacheConversionError::P1TooLarge)?; - let mem = self.memory_handle()?; + .map_err(CopyKvCacheSeqError::P1TooLarge)?; + let mem = self.required_memory_handle()?; unsafe { llama_cpp_bindings_sys::llama_memory_seq_cp(mem, src, dest, p0, p1) }; Ok(()) @@ -211,57 +210,43 @@ impl LlamaContext<'_> { /// holds no KV cache, so the removal trivially succeeds. pub fn clear_kv_cache_seq( &mut self, - src: Option, + seq_id: Option, p0: Option, p1: Option, - ) -> Result<(), KvCacheConversionError> { - let src = src + ) -> Result<(), ClearKvCacheSeqError> { + let seq_id = seq_id .map_or(Ok(-1), i32::try_from) - .map_err(KvCacheConversionError::SeqIdTooLarge)?; + .map_err(ClearKvCacheSeqError::SeqIdTooLarge)?; let p0 = p0 .map_or(Ok(-1), i32::try_from) - .map_err(KvCacheConversionError::P0TooLarge)?; + .map_err(ClearKvCacheSeqError::P0TooLarge)?; let p1 = p1 .map_or(Ok(-1), i32::try_from) - .map_err(KvCacheConversionError::P1TooLarge)?; - let Some(mem) = self.optional_memory_handle() else { + .map_err(ClearKvCacheSeqError::P1TooLarge)?; + let Some(mem) = self.memory_handle() else { return Ok(()); }; - if unsafe { llama_cpp_bindings_sys::llama_memory_seq_rm(mem, src, p0, p1) } { + if unsafe { llama_cpp_bindings_sys::llama_memory_seq_rm(mem, seq_id, p0, p1) } { return Ok(()); } - Err(KvCacheConversionError::PartialSequenceNotRemoved { - seq_id: src, - p0, - p1, - }) + Err(ClearKvCacheSeqError::PartialSequenceNotRemoved { seq_id, p0, p1 }) } - /// # Errors - /// Never returns an error: a context without a memory module holds no KV cache, so - /// there is nothing to clear. - pub fn clear_kv_cache(&mut self) -> Result<(), KvCacheConversionError> { - let Some(mem) = self.optional_memory_handle() else { - return Ok(()); + pub fn clear_kv_cache(&mut self) { + let Some(mem) = self.memory_handle() else { + return; }; let clear_data_buffers = true; unsafe { llama_cpp_bindings_sys::llama_memory_clear(mem, clear_data_buffers) }; - - Ok(()) } - /// # Errors - /// Never returns an error: a context without a memory module holds no KV cache, so - /// every sequence is already absent from it. - pub fn kv_cache_seq_keep(&mut self, seq_id: i32) -> Result<(), KvCacheConversionError> { - let Some(mem) = self.optional_memory_handle() else { - return Ok(()); + pub fn kv_cache_seq_keep(&mut self, seq_id: i32) { + let Some(mem) = self.memory_handle() else { + return; }; unsafe { llama_cpp_bindings_sys::llama_memory_seq_keep(mem, seq_id) }; - - Ok(()) } /// # Errors diff --git a/llama-cpp-bindings/src/error.rs b/llama-cpp-bindings/src/error.rs index 9896295e..04deed8f 100644 --- a/llama-cpp-bindings/src/error.rs +++ b/llama-cpp-bindings/src/error.rs @@ -1,6 +1,8 @@ pub mod apply_chat_template_error; pub mod bracketed_args_failure; pub mod chat_template_error; +pub mod clear_kv_cache_seq_error; +pub mod copy_kv_cache_seq_error; pub mod decode_error; pub mod embeddings_error; pub mod encode_error; @@ -11,7 +13,6 @@ pub mod grammar_runtime_error; pub mod json_object_failure; pub mod json_schema_to_grammar_error; pub mod key_value_xml_tags_failure; -pub mod kv_cache_conversion_error; pub mod kv_cache_seq_add_error; pub mod kv_cache_seq_div_error; pub mod kv_cache_seq_pos_max_error; @@ -43,6 +44,8 @@ pub use llama_cpp_ffi_status::FfiStatusError; pub use apply_chat_template_error::ApplyChatTemplateError; pub use bracketed_args_failure::BracketedArgsFailure; pub use chat_template_error::ChatTemplateError; +pub use clear_kv_cache_seq_error::ClearKvCacheSeqError; +pub use copy_kv_cache_seq_error::CopyKvCacheSeqError; pub use decode_error::DecodeError; pub use embeddings_error::EmbeddingsError; pub use encode_error::EncodeError; @@ -53,7 +56,6 @@ pub use grammar_runtime_error::GrammarRuntimeError; pub use json_object_failure::JsonObjectFailure; pub use json_schema_to_grammar_error::JsonSchemaToGrammarError; pub use key_value_xml_tags_failure::KeyValueXmlTagsFailure; -pub use kv_cache_conversion_error::KvCacheConversionError; pub use kv_cache_seq_add_error::KvCacheSeqAddError; pub use kv_cache_seq_div_error::KvCacheSeqDivError; pub use kv_cache_seq_pos_max_error::KvCacheSeqPosMaxError; diff --git a/llama-cpp-bindings/src/error/kv_cache_conversion_error.rs b/llama-cpp-bindings/src/error/clear_kv_cache_seq_error.rs similarity index 56% rename from llama-cpp-bindings/src/error/kv_cache_conversion_error.rs rename to llama-cpp-bindings/src/error/clear_kv_cache_seq_error.rs index f545744f..45c01bcd 100644 --- a/llama-cpp-bindings/src/error/kv_cache_conversion_error.rs +++ b/llama-cpp-bindings/src/error/clear_kv_cache_seq_error.rs @@ -2,15 +2,13 @@ use std::ffi::c_int; use std::num::TryFromIntError; #[derive(Debug, Eq, PartialEq, thiserror::Error)] -pub enum KvCacheConversionError { - #[error("Provided sequence id is too large for a i32")] +pub enum ClearKvCacheSeqError { + #[error("provided sequence id is too large for an i32")] SeqIdTooLarge(#[source] TryFromIntError), - #[error("Provided start position is too large for a i32")] + #[error("provided start position is too large for an i32")] P0TooLarge(#[source] TryFromIntError), - #[error("Provided end position is too large for a i32")] + #[error("provided end position is too large for an i32")] P1TooLarge(#[source] TryFromIntError), - #[error("the context has no memory module attached")] - MemoryHandleUnavailable, #[error("sequence {seq_id} could not be partially removed over positions [{p0}, {p1})")] PartialSequenceNotRemoved { seq_id: c_int, p0: c_int, p1: c_int }, } diff --git a/llama-cpp-bindings/src/error/copy_kv_cache_seq_error.rs b/llama-cpp-bindings/src/error/copy_kv_cache_seq_error.rs new file mode 100644 index 00000000..011e7074 --- /dev/null +++ b/llama-cpp-bindings/src/error/copy_kv_cache_seq_error.rs @@ -0,0 +1,11 @@ +use std::num::TryFromIntError; + +#[derive(Debug, Eq, PartialEq, thiserror::Error)] +pub enum CopyKvCacheSeqError { + #[error("provided start position is too large for an i32")] + P0TooLarge(#[source] TryFromIntError), + #[error("provided end position is too large for an i32")] + P1TooLarge(#[source] TryFromIntError), + #[error("the context has no memory module attached")] + MemoryHandleUnavailable, +} diff --git a/llama-cpp-bindings/src/error/marker_detection_error.rs b/llama-cpp-bindings/src/error/marker_detection_error.rs index 0d6abc58..6ea6a8d8 100644 --- a/llama-cpp-bindings/src/error/marker_detection_error.rs +++ b/llama-cpp-bindings/src/error/marker_detection_error.rs @@ -3,6 +3,7 @@ use std::string::FromUtf8Error; use crate::error::chat_template_error::ChatTemplateError; use crate::error::string_to_token_error::StringToTokenError; +use crate::token::LlamaToken; #[derive(Debug, PartialEq, Eq, thiserror::Error)] pub enum MarkerDetectionError { @@ -30,6 +31,10 @@ pub enum MarkerDetectionError { ReasoningMarkersFreeFailed { message: String }, #[error("a detected marker string could not be tokenised: {0}")] MarkerTokenizationFailed(#[from] StringToTokenError), + #[error("marker detection produced an empty token sequence")] + EmptyMarker, + #[error("marker token sequence {tokens:?} opens more than one section")] + AmbiguousMarkerOpeners { tokens: Vec }, #[error("the chat template is not valid UTF-8: {0}")] ToolCallTemplateNotUtf8(#[from] Utf8Error), #[error("the chat template could not be retrieved for tool-call marker detection: {0}")] diff --git a/llama-cpp-bindings/src/lib.rs b/llama-cpp-bindings/src/lib.rs index ba5f1bcc..c9760c1a 100644 --- a/llama-cpp-bindings/src/lib.rs +++ b/llama-cpp-bindings/src/lib.rs @@ -38,7 +38,7 @@ pub mod load_backends_error; pub mod load_backends_from_path; pub mod log_options; pub mod log_record; -pub mod marker_kind; +pub mod marker_role; pub mod mask_outcome; pub mod max_devices; pub mod mlock_supported; @@ -54,6 +54,7 @@ pub mod sampling; pub mod sanitized_grammar; pub mod send_logs_to_log; pub mod streaming_json_probe; +pub mod streaming_marker; pub mod streaming_markers; pub mod synthetic_tool_call_renders; pub mod timing; @@ -62,13 +63,13 @@ pub mod tool_call_format; pub mod tool_call_marker_pair; pub use error::{ - ApplyChatTemplateError, ChatTemplateError, DecodeError, EmbeddingsError, EncodeError, - EvalMultimodalChunksError, FfiContractError, FfiStatusError, GrammarError, - JsonSchemaToGrammarError, KvCacheSeqAddError, KvCacheSeqDivError, KvCacheSeqPosMaxError, - LlamaContextLoadError, LlamaCppError, LlamaLoraAdapterInitError, LlamaLoraAdaptersError, - LlamaModelLoadError, LogitsError, MarkerDetectionError, MetaValError, ModelParamsError, - NewLlamaChatMessageError, ParseChatMessageError, Result, SampleError, SamplerAcceptError, - SamplingError, StringToTokenError, TokenSamplingError, TokenToStringError, + ApplyChatTemplateError, ChatTemplateError, ClearKvCacheSeqError, CopyKvCacheSeqError, + DecodeError, EmbeddingsError, EncodeError, EvalMultimodalChunksError, FfiContractError, + FfiStatusError, GrammarError, JsonSchemaToGrammarError, KvCacheSeqAddError, KvCacheSeqDivError, + KvCacheSeqPosMaxError, LlamaContextLoadError, LlamaCppError, LlamaLoraAdapterInitError, + LlamaLoraAdaptersError, LlamaModelLoadError, LogitsError, MarkerDetectionError, MetaValError, + ModelParamsError, NewLlamaChatMessageError, ParseChatMessageError, Result, SampleError, + SamplerAcceptError, SamplingError, StringToTokenError, TokenSamplingError, TokenToStringError, }; pub use chat_message_parse_outcome::ChatMessageParseOutcome; @@ -81,10 +82,13 @@ pub use llama_cpp_bindings_types::{ ReasoningMarkers, TokenUsage, TokenUsageError, ToolCallArgsShape, ToolCallArguments, ToolCallMarkers, ToolCallValueQuote, XmlTagsShape, }; +pub use marker_role::MarkerRole; pub use raw_chat_message::RawChatMessage; pub use sampled_token::SampledToken; pub use sampled_token_classifier::SampledTokenClassifier; pub use sampled_token_section::SampledTokenSection; +pub use streaming_marker::StreamingMarker; +pub use streaming_markers::StreamingMarkers; pub use synthetic_tool_call_renders::SyntheticToolCallRenders; pub use ggml_time_us::ggml_time_us; diff --git a/llama-cpp-bindings/src/marker_kind.rs b/llama-cpp-bindings/src/marker_kind.rs deleted file mode 100644 index fe027e7a..00000000 --- a/llama-cpp-bindings/src/marker_kind.rs +++ /dev/null @@ -1,7 +0,0 @@ -#[derive(Copy, Clone, Debug, Eq, PartialEq)] -pub enum MarkerKind { - ReasoningOpen, - ReasoningClose, - ToolCallOpen, - ToolCallClose, -} diff --git a/llama-cpp-bindings/src/marker_role.rs b/llama-cpp-bindings/src/marker_role.rs new file mode 100644 index 00000000..f2f93aaf --- /dev/null +++ b/llama-cpp-bindings/src/marker_role.rs @@ -0,0 +1,28 @@ +use crate::sampled_token_section::SampledTokenSection; + +#[derive(Copy, Clone, Debug, Eq, PartialEq)] +/// A semantic transition performed when a streaming marker is consumed. +pub enum MarkerRole { + ReasoningOpen, + ReasoningClose, + ToolCallOpen, + ToolCallClose, +} + +impl MarkerRole { + pub(crate) const fn opened_section(self) -> Option { + match self { + Self::ReasoningOpen => Some(SampledTokenSection::Reasoning), + Self::ToolCallOpen => Some(SampledTokenSection::ToolCall), + Self::ReasoningClose | Self::ToolCallClose => None, + } + } + + pub(crate) const fn closed_section(self) -> Option { + match self { + Self::ReasoningClose => Some(SampledTokenSection::Reasoning), + Self::ToolCallClose => Some(SampledTokenSection::ToolCall), + Self::ReasoningOpen | Self::ToolCallOpen => None, + } + } +} diff --git a/llama-cpp-bindings/src/model.rs b/llama-cpp-bindings/src/model.rs index 8ce33b3b..149cc0ff 100644 --- a/llama-cpp-bindings/src/model.rs +++ b/llama-cpp-bindings/src/model.rs @@ -37,6 +37,7 @@ use crate::chat_template_tool_calls; use crate::llama_backend::LlamaBackend; use crate::llama_token_attrs::LlamaTokenAttrs; use crate::llama_token_attrs_from_int_error::LlamaTokenAttrsFromIntError; +use crate::marker_role::MarkerRole; use crate::model::tokenizer_input::TokenizerInput; use crate::raw_chat_message::RawChatMessage; use crate::resolved_tool_call_markers::ResolvedToolCallMarkers; @@ -1017,11 +1018,11 @@ impl LlamaModel { let resolved_tool_call_markers = self.resolve_tool_call_marker_strings(autoparser_open, autoparser_close)?; - let mut reasoning_closes = Vec::new(); + let mut candidates = Vec::new(); if let Some(markers) = &reasoning_markers { for marker in &markers.closes { if let Some(tokens) = self.tokenize_marker(Some(marker))? { - reasoning_closes.push(tokens); + candidates.push((tokens, MarkerRole::ReasoningClose)); } } } @@ -1029,14 +1030,17 @@ impl LlamaModel { let reasoning_open = reasoning_markers .as_ref() .map(|markers| markers.open.as_str()); - let reasoning_open = self.tokenize_marker(reasoning_open)?; + if let Some(tokens) = self.tokenize_marker(reasoning_open)? { + candidates.push((tokens, MarkerRole::ReasoningOpen)); + } + if let Some(tokens) = self.tokenize_marker(resolved_tool_call_markers.open.as_deref())? { + candidates.push((tokens, MarkerRole::ToolCallOpen)); + } + if let Some(tokens) = self.tokenize_marker(resolved_tool_call_markers.close.as_deref())? { + candidates.push((tokens, MarkerRole::ToolCallClose)); + } - Ok(StreamingMarkers { - reasoning_open, - reasoning_closes, - 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())?, - }) + StreamingMarkers::from_candidates(candidates) } fn resolve_tool_call_marker_strings( diff --git a/llama-cpp-bindings/src/sampled_token_classifier.rs b/llama-cpp-bindings/src/sampled_token_classifier.rs index 8061fa36..782b81d6 100644 --- a/llama-cpp-bindings/src/sampled_token_classifier.rs +++ b/llama-cpp-bindings/src/sampled_token_classifier.rs @@ -13,7 +13,6 @@ use crate::error::SampleError; use crate::error::TokenToStringError; use crate::eval_multimodal_chunks_params::EvalMultimodalChunksParams; use crate::llama_batch::LlamaBatch; -use crate::marker_kind::MarkerKind; use crate::model::LlamaModel; use crate::mtmd::MtmdContext; use crate::mtmd::MtmdInputChunks; @@ -27,12 +26,19 @@ pub use crate::classified_sample::ClassifiedSample; use crate::ingest_outcome::IngestOutcome; pub use crate::sampled_token_section::SampledTokenSection; +#[derive(Copy, Clone, Debug)] +enum PendingMarkerStatus { + Unmatched, + ResolvedBoundary, + AmbiguousBoundary, +} + #[derive(Clone, Debug)] struct PendingToken { token: LlamaToken, decoded: String, section: SampledTokenSection, - is_boundary: bool, + marker_status: PendingMarkerStatus, is_from_prompt: bool, is_held_for_probe: bool, } @@ -61,7 +67,7 @@ pub struct SampledTokenClassifier<'model> { impl<'model> SampledTokenClassifier<'model> { #[must_use] - pub fn new(model: &'model LlamaModel, markers: StreamingMarkers) -> Self { + pub(crate) fn new(model: &'model LlamaModel, markers: StreamingMarkers) -> Self { Self { model, markers, @@ -79,7 +85,7 @@ impl<'model> SampledTokenClassifier<'model> { /// detokenised. The failure is surfaced rather than substituting an empty /// piece, so classification never silently drops generated text. pub fn ingest(&mut self, token: LlamaToken) -> Result, TokenToStringError> { - if !self.markers.has_any() { + if self.markers.is_empty() { self.usage.record_undeterminable_token(); let piece = self.decode(token)?; return Ok(vec![IngestOutcome { @@ -94,7 +100,7 @@ impl<'model> SampledTokenClassifier<'model> { token, decoded: decoded.clone(), section: self.section, - is_boundary: false, + marker_status: PendingMarkerStatus::Unmatched, is_from_prompt: false, is_held_for_probe: false, }); @@ -124,7 +130,7 @@ impl<'model> SampledTokenClassifier<'model> { } pub fn ingest_prompt_token(&mut self, token: LlamaToken) { - if !self.markers.has_any() { + if self.markers.is_empty() { return; } @@ -132,7 +138,7 @@ impl<'model> SampledTokenClassifier<'model> { token, decoded: String::new(), section: self.section, - is_boundary: false, + marker_status: PendingMarkerStatus::Unmatched, is_from_prompt: true, is_held_for_probe: false, }); @@ -142,7 +148,7 @@ impl<'model> SampledTokenClassifier<'model> { } pub fn ingest_prompt_tokens(&mut self, tokens: &[LlamaToken]) { - if !self.markers.has_any() { + if self.markers.is_empty() { return; } for &token in tokens { @@ -168,78 +174,32 @@ impl<'model> SampledTokenClassifier<'model> { } fn try_consume_marker_at_tail(&mut self) { - let marker_match = self - .markers - .reasoning_open - .as_deref() - .and_then(|marker| self.marker_span_start(marker)) - .map(|span_start| (span_start, MarkerKind::ReasoningOpen)) - .or_else(|| { - self.markers - .tool_call_open - .as_deref() - .and_then(|marker| self.marker_span_start(marker)) - .map(|span_start| (span_start, MarkerKind::ToolCallOpen)) - }) - .or_else(|| { - self.markers.reasoning_closes.iter().find_map(|marker| { - self.marker_span_start(marker) - .map(|span_start| (span_start, MarkerKind::ReasoningClose)) - }) - }) - .or_else(|| { - self.markers - .tool_call_close - .as_deref() - .and_then(|marker| self.marker_span_start(marker)) - .map(|span_start| (span_start, MarkerKind::ToolCallClose)) - }); - - if let Some((span_start, marker_kind)) = marker_match { - self.mark_marker_span(span_start, marker_kind); - } - } + let pending_tokens: Vec<_> = self.pending.iter().map(|entry| entry.token).collect(); + let Some(marker) = self.markers.longest_matching_suffix(&pending_tokens) else { + return; + }; + let span_start = self.pending.len() - marker.tokens().len(); + let span_section = marker.span_section(self.section); + let next_section = marker.next_section(); + let is_ambiguous_prefix = self.markers.is_prefix_of_longer_marker(marker.tokens()); - fn marker_span_start(&self, marker: &[LlamaToken]) -> Option { - if marker.is_empty() || self.pending.len() < marker.len() { - return None; - } - let span_start = self.pending.len() - marker.len(); - self.pending - .iter() - .skip(span_start) - .zip(marker) - .all(|(entry, marker_token)| entry.token == *marker_token) - .then_some(span_start) + self.mark_marker_span(span_start, span_section, next_section, is_ambiguous_prefix); } - 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, - }; - 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 - } - } + fn mark_marker_span( + &mut self, + span_start: usize, + span_section: SampledTokenSection, + next_section: SampledTokenSection, + is_ambiguous_prefix: bool, + ) { + let marker_status = if is_ambiguous_prefix { + PendingMarkerStatus::AmbiguousBoundary + } else { + PendingMarkerStatus::ResolvedBoundary }; - for entry in self.pending.iter_mut().skip(span_start) { - entry.is_boundary = true; + entry.marker_status = marker_status; entry.section = span_section; } @@ -261,7 +221,9 @@ impl<'model> SampledTokenClassifier<'model> { .count(); let drainable = self.pending.len().saturating_sub(probe_held); let beyond_lookback = drainable > lookback; - if !front.is_boundary && !beyond_lookback { + let resolved_boundary = + matches!(front.marker_status, PendingMarkerStatus::ResolvedBoundary); + if !resolved_boundary && !beyond_lookback { break; } let Some(entry) = self.pending.pop_front() else { @@ -382,10 +344,10 @@ impl<'model> SampledTokenClassifier<'model> { SampledTokenSection::Pending => SampledToken::Undeterminable(entry.token), }; - let visible_piece = if entry.is_boundary { - String::new() - } else { + let visible_piece = if matches!(entry.marker_status, PendingMarkerStatus::Unmatched) { entry.decoded.clone() + } else { + String::new() }; IngestOutcome { @@ -548,10 +510,12 @@ impl<'model> SampledTokenClassifier<'model> { #[cfg(test)] mod tests { use super::JsonProbeState; + use super::PendingMarkerStatus; use super::PendingToken; use super::ProbeMode; use super::SampledTokenClassifier; use crate::ingest_outcome::IngestOutcome; + use crate::marker_role::MarkerRole; use crate::sampled_token::SampledToken; use crate::sampled_token_section::SampledTokenSection; use crate::streaming_markers::StreamingMarkers; @@ -565,12 +529,16 @@ mod tests { reasoning_open: Option>, reasoning_close: Option>, ) -> StreamingMarkers { - StreamingMarkers { - reasoning_open, - reasoning_closes: reasoning_close.into_iter().collect(), - tool_call_open: None, - tool_call_close: None, - } + let candidates = reasoning_open + .into_iter() + .map(|tokens| (tokens, MarkerRole::ReasoningOpen)) + .chain( + reasoning_close + .into_iter() + .map(|tokens| (tokens, MarkerRole::ReasoningClose)), + ); + + StreamingMarkers::from_candidates(candidates).expect("synthetic markers must be valid") } fn synthetic_classifier(markers: StreamingMarkers) -> SampledTokenClassifier<'static> { @@ -591,7 +559,7 @@ mod tests { token: token(token_id), decoded: decoded.to_owned(), section: classifier.section, - is_boundary: false, + marker_status: PendingMarkerStatus::Unmatched, is_from_prompt: false, is_held_for_probe: false, }); @@ -602,7 +570,7 @@ mod tests { token: token(token_id), decoded: String::new(), section: classifier.section, - is_boundary: false, + marker_status: PendingMarkerStatus::Unmatched, is_from_prompt: true, is_held_for_probe: false, }); @@ -642,12 +610,14 @@ mod tests { fn markers_sharing_tool_call_open_with_a_reasoning_close( shared: Vec, ) -> StreamingMarkers { - StreamingMarkers { - reasoning_open: Some(vec![token(100)]), - reasoning_closes: vec![vec![token(200)], shared.clone()], - tool_call_open: Some(shared), - tool_call_close: Some(vec![token(201)]), - } + StreamingMarkers::from_candidates([ + (vec![token(100)], MarkerRole::ReasoningOpen), + (vec![token(200)], MarkerRole::ReasoningClose), + (shared.clone(), MarkerRole::ReasoningClose), + (shared, MarkerRole::ToolCallOpen), + (vec![token(201)], MarkerRole::ToolCallClose), + ]) + .expect("synthetic markers must be valid") } #[test] @@ -678,6 +648,33 @@ mod tests { assert_eq!(classifier.section, SampledTokenSection::ToolCall); } + #[test] + fn longer_marker_reclassifies_a_completed_marker_that_was_its_prefix() { + let markers = StreamingMarkers::from_candidates([ + (vec![token(300)], MarkerRole::ReasoningClose), + (vec![token(300), token(301)], MarkerRole::ToolCallOpen), + ]) + .expect("synthetic markers must be valid"); + let mut classifier = synthetic_classifier(markers); + classifier.section = SampledTokenSection::Reasoning; + + push_pending(&mut classifier, 300, ""); + classifier.try_consume_marker_at_tail(); + let outcomes = classifier.drain_overflow(); + + assert_eq!(classifier.section, SampledTokenSection::ToolCall); + assert_eq!( + outcome_sections(&outcomes), + vec![SampledTokenSection::ToolCall, SampledTokenSection::ToolCall] + ); + assert_eq!(outcome_pieces(&outcomes), vec!["", ""]); + } + #[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)])); @@ -801,12 +798,13 @@ mod tests { #[test] fn spurious_tool_call_close_in_reasoning_section_classifies_as_tool_call() { - let markers = StreamingMarkers { - reasoning_open: Some(vec![token(100)]), - reasoning_closes: vec![vec![token(200)]], - tool_call_open: Some(vec![token(300)]), - tool_call_close: Some(vec![token(400)]), - }; + let markers = StreamingMarkers::from_candidates([ + (vec![token(100)], MarkerRole::ReasoningOpen), + (vec![token(200)], MarkerRole::ReasoningClose), + (vec![token(300)], MarkerRole::ToolCallOpen), + (vec![token(400)], MarkerRole::ToolCallClose), + ]) + .expect("synthetic markers must be valid"); let mut classifier = synthetic_classifier(markers); classifier.section = SampledTokenSection::ToolCall; @@ -949,7 +947,7 @@ mod tests { token: token(202), decoded: "k>".to_owned(), section: classifier.section, - is_boundary: false, + marker_status: PendingMarkerStatus::Unmatched, is_from_prompt: false, is_held_for_probe: false, }); @@ -1014,7 +1012,7 @@ mod tests { token: token(50), decoded: "hi".to_owned(), section: classifier.section, - is_boundary: false, + marker_status: PendingMarkerStatus::Unmatched, is_from_prompt: false, is_held_for_probe: false, }); @@ -1115,8 +1113,12 @@ mod tests { let returned = classifier.markers(); - assert_eq!(returned.reasoning_open.as_deref(), Some(&[token(1)][..])); - assert_eq!(returned.reasoning_closes, vec![vec![token(2)]]); + assert!(returned.iter().any(|marker| { + marker.tokens() == [token(1)] && marker.roles() == [MarkerRole::ReasoningOpen] + })); + assert!(returned.iter().any(|marker| { + marker.tokens() == [token(2)] && marker.roles() == [MarkerRole::ReasoningClose] + })); } #[test] @@ -1132,8 +1134,9 @@ mod tests { #[test] fn spurious_tool_call_close_in_content_section_classifies_as_content() { - let mut markers = markers_with(None, None); - markers.tool_call_close = Some(vec![token(300)]); + let markers = + StreamingMarkers::from_candidates([(vec![token(300)], MarkerRole::ToolCallClose)]) + .expect("synthetic markers must be valid"); let mut classifier = synthetic_classifier(markers); classifier.section = SampledTokenSection::Content; @@ -1149,12 +1152,8 @@ mod tests { } fn markers_with_tool_call_open(tool_call_open: Vec) -> StreamingMarkers { - StreamingMarkers { - reasoning_open: None, - reasoning_closes: Vec::new(), - tool_call_open: Some(tool_call_open), - tool_call_close: None, - } + StreamingMarkers::from_candidates([(tool_call_open, MarkerRole::ToolCallOpen)]) + .expect("synthetic markers must be valid") } fn feed_json_string( @@ -1359,12 +1358,12 @@ mod tests { #[test] fn json_probe_does_not_engage_in_reasoning_section() { - let markers = StreamingMarkers { - reasoning_open: Some(vec![token(800)]), - reasoning_closes: vec![vec![token(801)]], - tool_call_open: Some(vec![token(900)]), - tool_call_close: None, - }; + let markers = StreamingMarkers::from_candidates([ + (vec![token(800)], MarkerRole::ReasoningOpen), + (vec![token(801)], MarkerRole::ReasoningClose), + (vec![token(900)], MarkerRole::ToolCallOpen), + ]) + .expect("synthetic markers must be valid"); let mut classifier = synthetic_classifier(markers); classifier.section = SampledTokenSection::Reasoning; @@ -1559,7 +1558,7 @@ mod tests { token: token(1), decoded: "before".to_owned(), section: SampledTokenSection::Content, - is_boundary: false, + marker_status: PendingMarkerStatus::Unmatched, is_from_prompt: false, is_held_for_probe: false, }); @@ -1567,7 +1566,7 @@ mod tests { token: token(2), decoded: "{}".to_owned(), section: SampledTokenSection::Content, - is_boundary: false, + marker_status: PendingMarkerStatus::Unmatched, is_from_prompt: false, is_held_for_probe: true, }); diff --git a/llama-cpp-bindings/src/streaming_marker.rs b/llama-cpp-bindings/src/streaming_marker.rs new file mode 100644 index 00000000..347f5ed0 --- /dev/null +++ b/llama-cpp-bindings/src/streaming_marker.rs @@ -0,0 +1,67 @@ +use crate::marker_role::MarkerRole; +use crate::sampled_token_section::SampledTokenSection; +use crate::token::LlamaToken; + +#[derive(Clone, Debug, Eq, PartialEq)] +/// A normalized token sequence and every semantic role attached to it. +pub struct StreamingMarker { + tokens: Vec, + roles: Vec, +} + +impl StreamingMarker { + pub(crate) fn new(tokens: Vec, role: MarkerRole) -> Self { + Self { + tokens, + roles: vec![role], + } + } + + pub(crate) fn add_role(&mut self, role: MarkerRole) { + if !self.roles.contains(&role) { + self.roles.push(role); + } + } + + #[must_use] + /// Returns the tokens that form this marker. + pub fn tokens(&self) -> &[LlamaToken] { + &self.tokens + } + + #[must_use] + /// Returns the transitions associated with this marker. + pub fn roles(&self) -> &[MarkerRole] { + &self.roles + } + + pub(crate) fn opener_count(&self) -> usize { + self.roles + .iter() + .filter(|role| role.opened_section().is_some()) + .count() + } + + fn opened_section(&self) -> Option { + self.roles.iter().find_map(|role| role.opened_section()) + } + + pub(crate) fn span_section(&self, current: SampledTokenSection) -> SampledTokenSection { + self.opened_section().unwrap_or_else(|| { + if self + .roles + .iter() + .any(|role| role.closed_section() == Some(current)) + { + current + } else { + SampledTokenSection::Content + } + }) + } + + pub(crate) fn next_section(&self) -> SampledTokenSection { + self.opened_section() + .unwrap_or(SampledTokenSection::Content) + } +} diff --git a/llama-cpp-bindings/src/streaming_markers.rs b/llama-cpp-bindings/src/streaming_markers.rs index 87400eb9..cb8b2028 100644 --- a/llama-cpp-bindings/src/streaming_markers.rs +++ b/llama-cpp-bindings/src/streaming_markers.rs @@ -1,41 +1,89 @@ +use crate::error::MarkerDetectionError; +use crate::marker_role::MarkerRole; +use crate::streaming_marker::StreamingMarker; use crate::token::LlamaToken; #[derive(Clone, Debug, Default, Eq, PartialEq)] +/// The normalized streaming markers detected for a model. pub struct StreamingMarkers { - pub reasoning_open: Option>, - pub reasoning_closes: Vec>, - pub tool_call_open: Option>, - pub tool_call_close: Option>, + markers: Vec, } impl StreamingMarkers { + pub(crate) fn from_candidates( + candidates: impl IntoIterator, MarkerRole)>, + ) -> Result { + let mut markers: Vec = Vec::new(); + + for (tokens, role) in candidates { + if tokens.is_empty() { + return Err(MarkerDetectionError::EmptyMarker); + } + + if let Some(marker) = markers.iter_mut().find(|marker| marker.tokens() == tokens) { + marker.add_role(role); + } else { + markers.push(StreamingMarker::new(tokens, role)); + } + } + + if let Some(marker) = markers.iter().find(|marker| marker.opener_count() > 1) { + return Err(MarkerDetectionError::AmbiguousMarkerOpeners { + tokens: marker.tokens().to_vec(), + }); + } + + Ok(Self { markers }) + } + #[must_use] - pub const fn has_any(&self) -> bool { - self.reasoning_open.is_some() - || !self.reasoning_closes.is_empty() - || self.tool_call_open.is_some() - || self.tool_call_close.is_some() + /// Returns whether the model exposes no streaming markers. + pub const fn is_empty(&self) -> bool { + self.markers.is_empty() + } + + #[must_use] + /// Returns the number of distinct marker token sequences. + pub const fn len(&self) -> usize { + self.markers.len() + } + + /// Iterates over the distinct marker token sequences. + pub fn iter(&self) -> impl Iterator { + self.markers.iter() } #[must_use] pub fn max_token_len(&self) -> usize { - [ - self.reasoning_open.as_deref(), - self.tool_call_open.as_deref(), - self.tool_call_close.as_deref(), - ] - .into_iter() - .flatten() - .map(<[LlamaToken]>::len) - .chain(self.reasoning_closes.iter().map(Vec::len)) - .max() - .unwrap_or(0) + self.markers + .iter() + .map(|marker| marker.tokens().len()) + .max() + .unwrap_or(0) + } + + pub(crate) fn longest_matching_suffix( + &self, + tokens: &[LlamaToken], + ) -> Option<&StreamingMarker> { + self.markers + .iter() + .filter(|marker| tokens.ends_with(marker.tokens())) + .max_by_key(|marker| marker.tokens().len()) + } + + pub(crate) fn is_prefix_of_longer_marker(&self, tokens: &[LlamaToken]) -> bool { + self.markers.iter().any(|marker| { + marker.tokens().len() > tokens.len() && marker.tokens().starts_with(tokens) + }) } } #[cfg(test)] mod tests { use super::StreamingMarkers; + use crate::error::MarkerDetectionError; + use crate::marker_role::MarkerRole; use crate::token::LlamaToken; fn token(id: i32) -> LlamaToken { @@ -43,20 +91,93 @@ mod tests { } #[test] - fn streaming_markers_with_no_markers_reports_none() { + fn empty_collection_reports_no_markers() { let markers = StreamingMarkers::default(); - assert!(!markers.has_any()); + + assert!(markers.is_empty()); + assert_eq!(markers.len(), 0); 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_closes: vec![vec![token(2), token(3), token(4)]], - tool_call_open: Some(vec![token(5), token(6)]), - tool_call_close: None, - }; + fn candidates_with_the_same_tokens_are_one_marker_with_multiple_roles() { + let markers = StreamingMarkers::from_candidates([ + (vec![token(1)], MarkerRole::ReasoningClose), + (vec![token(1)], MarkerRole::ToolCallOpen), + ]) + .expect("a close and an opener compose into one transition"); + + let marker = markers.iter().next().expect("one marker must remain"); + assert_eq!(marker.tokens(), &[token(1)]); + assert_eq!( + marker.roles(), + &[MarkerRole::ReasoningClose, MarkerRole::ToolCallOpen] + ); + } + + #[test] + fn empty_marker_is_rejected() { + assert_eq!( + StreamingMarkers::from_candidates([(Vec::new(), MarkerRole::ReasoningOpen)]), + Err(MarkerDetectionError::EmptyMarker) + ); + } + + #[test] + fn two_openers_for_the_same_tokens_are_rejected() { + let marker_tokens = vec![token(1), token(2)]; + + assert_eq!( + StreamingMarkers::from_candidates([ + (marker_tokens.clone(), MarkerRole::ReasoningOpen), + (marker_tokens.clone(), MarkerRole::ToolCallOpen), + ]), + Err(MarkerDetectionError::AmbiguousMarkerOpeners { + tokens: marker_tokens + }) + ); + } + + #[test] + fn longest_matching_suffix_wins() { + let markers = StreamingMarkers::from_candidates([ + (vec![token(2)], MarkerRole::ReasoningClose), + (vec![token(1), token(2)], MarkerRole::ToolCallOpen), + ]) + .expect("markers are valid"); + + let matched = markers + .longest_matching_suffix(&[token(1), token(2)]) + .expect("a suffix must match"); + + assert_eq!(matched.tokens(), &[token(1), token(2)]); + assert_eq!(matched.roles(), &[MarkerRole::ToolCallOpen]); + } + + #[test] + fn shorter_complete_marker_reports_when_it_is_still_an_ambiguous_prefix() { + let markers = StreamingMarkers::from_candidates([ + (vec![token(1)], MarkerRole::ReasoningClose), + (vec![token(1), token(2)], MarkerRole::ToolCallOpen), + ]) + .expect("markers are valid"); + + assert!(markers.is_prefix_of_longer_marker(&[token(1)])); + assert!(!markers.is_prefix_of_longer_marker(&[token(1), token(2)])); + } + + #[test] + fn max_token_len_uses_the_longest_normalized_marker() { + let markers = StreamingMarkers::from_candidates([ + (vec![token(1)], MarkerRole::ReasoningOpen), + ( + vec![token(2), token(3), token(4)], + MarkerRole::ReasoningClose, + ), + (vec![token(5), token(6)], MarkerRole::ToolCallOpen), + ]) + .expect("markers are valid"); + assert_eq!(markers.max_token_len(), 3); } }