Skip to content

Add an assistant prefill option to LlamaLanguageModel #198

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
james-333i wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Add an assistant prefill option to LlamaLanguageModel #198

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
3e168e3
Fix LlamaLanguageModel build against current llama.swift
james-333i Aug 25, 2026
9a2b94d
Add an assistant prefill option to LlamaLanguageModel
james-333i Aug 25, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 9 additions & 0 deletions Package.resolved
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

43 changes: 35 additions & 8 deletions Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -137,6 +137,14 @@ import Foundation
/// Mirostat sampling mode for adaptive perplexity control.
public var mirostat: MirostatMode?

/// Text appended after the assistant header of the rendered prompt.
///
/// The model continues generating from this text. Use it to steer the
/// start of the response, for example prefilling an empty
/// `<think></think>` block to suppress a model's default reasoning
/// output when its chat template offers no switch for it.
public var assistantPrefill: String?

/// Creates custom generation options for llama.cpp.
public init(
contextSize: UInt32? = nil,
Expand All @@ -150,7 +158,8 @@ import Foundation
repeatLastN: Int32? = nil,
frequencyPenalty: Float? = nil,
presencePenalty: Float? = nil,
mirostat: MirostatMode? = nil
mirostat: MirostatMode? = nil,
assistantPrefill: String? = nil
) {
self.contextSize = contextSize
self.batchSize = batchSize
Expand All @@ -164,6 +173,7 @@ import Foundation
self.frequencyPenalty = frequencyPenalty
self.presencePenalty = presencePenalty
self.mirostat = mirostat
self.assistantPrefill = assistantPrefill
}

/// Default llama.cpp options used when none are provided at runtime.
Expand Down Expand Up @@ -321,6 +331,7 @@ import Foundation
var frequencyPenalty: Float
var presencePenalty: Float
var mirostat: CustomGenerationOptions.MirostatMode?
var assistantPrefill: String?
var sampling: GenerationOptions.SamplingMode?
var maximumResponseTokens: Int?

Expand All @@ -337,6 +348,7 @@ import Foundation
frequencyPenalty: Float = 0.0,
presencePenalty: Float = 0.0,
mirostat: CustomGenerationOptions.MirostatMode? = nil,
assistantPrefill: String? = nil,
sampling: GenerationOptions.SamplingMode? = nil,
maximumResponseTokens: Int? = nil
) {
Expand All @@ -352,6 +364,7 @@ import Foundation
self.frequencyPenalty = frequencyPenalty
self.presencePenalty = presencePenalty
self.mirostat = mirostat
self.assistantPrefill = assistantPrefill
self.sampling = sampling
self.maximumResponseTokens = maximumResponseTokens
}
Expand Down Expand Up @@ -389,6 +402,7 @@ import Foundation
frequencyPenalty: base.frequencyPenalty,
presencePenalty: base.presencePenalty,
mirostat: base.mirostat,
assistantPrefill: base.assistantPrefill,
sampling: sampling ?? base.sampling,
maximumResponseTokens: maximumResponseTokens ?? base.maximumResponseTokens
)
Expand All @@ -407,6 +421,7 @@ import Foundation
self.frequencyPenalty = options.frequencyPenalty ?? base.frequencyPenalty
self.presencePenalty = options.presencePenalty ?? base.presencePenalty
self.mirostat = options.mirostat ?? base.mirostat
self.assistantPrefill = options.assistantPrefill ?? base.assistantPrefill
self.sampling = sampling ?? base.sampling
self.maximumResponseTokens = maximumResponseTokens ?? base.maximumResponseTokens
}
Expand Down Expand Up @@ -509,10 +524,11 @@ import Foundation
if includeSchemaInPrompt, type != String.self {
fullPrompt = try formatPrompt(
for: session,
extraSystemMessage: schemaPrompt(for: type.generationSchema)
extraSystemMessage: schemaPrompt(for: type.generationSchema),
assistantPrefill: runtimeOptions.assistantPrefill
)
} else {
fullPrompt = try formatPrompt(for: session)
fullPrompt = try formatPrompt(for: session, assistantPrefill: runtimeOptions.assistantPrefill)
}

if type == String.self {
Expand Down Expand Up @@ -600,7 +616,10 @@ import Foundation
llama_set_n_threads(context, runtimeOptions.threads, runtimeOptions.threads)

var accumulatedText = ""
let fullPrompt = try self.formatPrompt(for: session)
let fullPrompt = try self.formatPrompt(
for: session,
assistantPrefill: runtimeOptions.assistantPrefill
)

do {
for try await tokenText in generateTextStream(
Expand Down Expand Up @@ -673,8 +692,7 @@ import Foundation
params.n_gpu_layers = 0

// Try to reduce memory usage
params.use_mmap = true
params.use_mlock = false
params.load_mode = LLAMA_LOAD_MODE_MMAP
return params
}

Expand Down Expand Up @@ -829,6 +847,7 @@ import Foundation
llama_sampler_chain_add(
samplerPtr,
llama_sampler_init_penalties(
llama_vocab_n_tokens(vocab),
effectiveRepeatLastN,
effectiveRepeatPenalty,
effectiveFrequencyPenalty,
Expand Down Expand Up @@ -958,6 +977,7 @@ import Foundation
llama_sampler_chain_add(
samplerPointer,
llama_sampler_init_penalties(
llama_vocab_n_tokens(vocab),
options.repeatLastN,
options.repeatPenalty,
options.frequencyPenalty,
Expand Down Expand Up @@ -1197,6 +1217,7 @@ import Foundation
llama_sampler_chain_add(
samplerPtr,
llama_sampler_init_penalties(
llama_vocab_n_tokens(vocab),
effectiveRepeatLastN,
effectiveRepeatPenalty,
effectiveFrequencyPenalty,
Expand Down Expand Up @@ -1368,7 +1389,8 @@ import Foundation

private func formatPrompt(
for session: LanguageModelSession,
extraSystemMessage: String? = nil
extraSystemMessage: String? = nil,
assistantPrefill: String? = nil
) throws -> String {
guard let model = self.model else {
throw LlamaLanguageModelError.modelLoadFailed
Expand Down Expand Up @@ -1452,9 +1474,14 @@ import Foundation
throw LlamaLanguageModelError.encodingFailed
}

return buffer.withUnsafeBytes { rawBuffer in
let rendered = buffer.withUnsafeBytes { rawBuffer in
String(decoding: rawBuffer.prefix(Int(result)), as: UTF8.self)
}

if let assistantPrefill, !assistantPrefill.isEmpty {
return rendered + assistantPrefill
}
return rendered
}

private func extractText(from segments: [Transcript.Segment]) -> String {
Expand Down
4 changes: 3 additions & 1 deletion Tests/AnyLanguageModelTests/LlamaLanguageModelTests.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,8 @@ import Testing
repeatLastN: 48,
frequencyPenalty: 0.05,
presencePenalty: 0.05,
mirostat: .v2(tau: 5.0, eta: 0.2)
mirostat: .v2(tau: 5.0, eta: 0.2),
assistantPrefill: "<think></think>"
)
options[custom: LlamaLanguageModel.self] = custom

Expand All @@ -62,6 +63,7 @@ import Testing
#expect(retrieved?.frequencyPenalty == 0.05)
#expect(retrieved?.presencePenalty == 0.05)
#expect(retrieved?.mirostat == .v2(tau: 5.0, eta: 0.2))
#expect(retrieved?.assistantPrefill == "<think></think>")
}

@Test func customGenerationOptionsDefaults() {
Expand Down