Skip to content

fix(sampler): allow llm_text_generate() to be called more than once #26

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

Merged
marcobambini merged 1 commit into from
Aug 22, 2026
+134 −37
Merged

fix(sampler): allow llm_text_generate() to be called more than once #26

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
90 changes: 54 additions & 36 deletions src/sqlite-ai.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -687,11 +687,17 @@ static const char *llm_decode_error_string (int32_t rc) {
}
}

// creates a new, unowned sampler chain; the caller owns the returned pointer
static struct llama_sampler *llm_sampler_new (void) {
struct llama_sampler_chain_params sampler_params = llama_sampler_chain_default_params();
return llama_sampler_chain_init(sampler_params);
}

// returns the connection-owned sampler chain, creating and publishing it if needed
struct llama_sampler *llm_sampler_check (ai_context *ai) {
if (ai->sampler) return ai->sampler;

struct llama_sampler_chain_params sampler_params = llama_sampler_chain_default_params();
struct llama_sampler *sampler = llama_sampler_chain_init(sampler_params);
struct llama_sampler *sampler = llm_sampler_new();
if (!sampler) {
sqlite_common_set_error(ai->context, ai->vtab, SQLITE_ERROR, "Unable to create sampler");
return NULL;
Expand Down Expand Up @@ -845,11 +851,12 @@ static void ai_free (void *ctx, bool free_ai, bool free_llm, bool free_audio) {
ai->vision = NULL;
memset(ai->lora, 0, sizeof(struct llama_adapter_lora *)*MAX_LORAS);
memset(ai->lora_scale, 0, sizeof(float)*MAX_LORAS);
// free before ctx/model: grammar/infill/mirostat samplers cache a vocab pointer
if (ai->sampler) llama_sampler_free(ai->sampler);
ai->sampler = NULL;
if (ai->ctx) llama_set_adapters_lora(ai->ctx, NULL, 0, NULL);
if (ai->ctx) llama_free(ai->ctx);
if (ai->model) llama_model_free(ai->model);
// sampler chain is freed explicitly via llm_sampler_free() or llm_sampler_create() SQL functions;
// freeing it here causes a double-free crash when ai_destroy runs after explicit cleanup
llm_options_init(&ai->options);

ai->model = NULL;
Expand Down Expand Up @@ -1515,6 +1522,7 @@ static void llm_text_run (sqlite3_context *context, const char *text, int32_t te
bool buffer_initialized = false;
buffer_t buffer = {0};
char *formatted_prompt = NULL;
struct llama_sampler *owned_sampler = NULL; // ephemeral chain, never published into ai->sampler

// sanity check vocab
const struct llama_vocab *vocab = llama_model_get_vocab(ai->model);
Expand Down Expand Up @@ -1608,20 +1616,28 @@ static void llm_text_run (sqlite3_context *context, const char *text, int32_t te
goto error;
}

// initialize the sampler
bool sampler_already_setup = (ai->sampler != NULL);
struct llama_sampler *sampler = llm_sampler_check(ai);
if (!sampler) goto error;
if (!sampler_already_setup) {
// no sampler was setup, so initialize it with some default values
// a user-configured chain belongs to the connection; otherwise build an
// ephemeral one that lives only for this call
struct llama_sampler *sampler = ai->sampler;
if (sampler == NULL || llama_sampler_chain_n(sampler) == 0) {
owned_sampler = llm_sampler_new();
if (!owned_sampler) {
sqlite_context_result_error(context, SQLITE_NOMEM, "Unable to create sampler");
goto error;
}
sampler = owned_sampler;
llama_sampler_chain_add(sampler, llama_sampler_init_penalties(64, 1.1, 0, 0));
llama_sampler_chain_add(sampler, llama_sampler_init_greedy());
}

// the KV cache was cleared above, so the sampler must start clean too:
// rebuild grammar state, re-seed dist, clear penalty history
llama_sampler_reset(sampler);

// allocate output buffer (starts small, grows dynamically via buffer_append)
if (!buffer_create(&buffer, 0)) {
sqlite_context_result_error(context, SQLITE_NOMEM, "Out of memory: failed to allocate buffer");
goto error_sampler;
goto error;
}
buffer_initialized = true;

Expand All @@ -1635,7 +1651,7 @@ static void llm_text_run (sqlite3_context *context, const char *text, int32_t te
int32_t drc = llama_decode(ctx, batch);
if (drc != 0) {
sqlite_context_result_error(context, SQLITE_ERROR, "Failed to execute the decoding function during prompt processing (%d: %s)", drc, llm_decode_error_string(drc));
goto error_sampler;
goto error;
}
prompt_pos += chunk;
}
Expand All @@ -1655,20 +1671,20 @@ static void llm_text_run (sqlite3_context *context, const char *text, int32_t te
int n = llama_token_to_piece(vocab, new_token_id, buf, sizeof(buf), 0, true);
if (n < 0) {
sqlite_context_result_error(context, SQLITE_ERROR, "Failed to convert token to piece (%d)", n);
goto error_sampler;
goto error;
}

if (buffer_append(&buffer, buf, n, true) == false) {
sqlite_context_result_error(context, SQLITE_NOMEM, "Out of memory: failed to append to buffer");
goto error_sampler;
goto error;
}

// decode the sampled token to advance the KV cache
struct llama_batch batch = llama_batch_get_one(&new_token_id, 1);
int32_t drc = llama_decode(ctx, batch);
if (drc != 0) {
sqlite_context_result_error(context, SQLITE_ERROR, "Failed to execute the decoding function during generation (%d: %s)", drc, llm_decode_error_string(drc));
goto error_sampler;
goto error;
}
}
}
Expand All @@ -1677,16 +1693,12 @@ static void llm_text_run (sqlite3_context *context, const char *text, int32_t te
sqlite3_result_text(context, buffer.data, buffer.length, sqlite3_free);
sqlite3_free(tokens);
sqlite3_free(formatted_prompt);
if (!sampler_already_setup) llama_sampler_free(sampler);
if (owned_sampler) llama_sampler_free(owned_sampler);
return;

error_sampler:
if (!sampler_already_setup && ai->sampler) {
llama_sampler_free(ai->sampler);
ai->sampler = NULL;
}
error:
if (buffer_initialized) buffer_destroy(&buffer);
if (owned_sampler) llama_sampler_free(owned_sampler);
sqlite3_free(tokens);
sqlite3_free(formatted_prompt);
}
Expand Down Expand Up @@ -3118,6 +3130,7 @@ static void llm_text_run_vision (sqlite3_context *context, const char *text, int
char *formatted_prompt = NULL;
mtmd_bitmap **bitmaps = NULL;
mtmd_input_chunks *chunks = NULL;
struct llama_sampler *owned_sampler = NULL; // ephemeral chain, never published into ai->sampler

struct llama_context *ctx = ai->ctx;
if (!ctx) {
Expand Down Expand Up @@ -3188,19 +3201,28 @@ static void llm_text_run_vision (sqlite3_context *context, const char *text, int
}
}

// initialize sampler
bool sampler_already_setup = (ai->sampler != NULL);
struct llama_sampler *sampler = llm_sampler_check(ai);
if (!sampler) goto error;
if (!sampler_already_setup) {
// a user-configured chain belongs to the connection; otherwise build an
// ephemeral one that lives only for this call
struct llama_sampler *sampler = ai->sampler;
if (sampler == NULL || llama_sampler_chain_n(sampler) == 0) {
owned_sampler = llm_sampler_new();
if (!owned_sampler) {
sqlite_context_result_error(context, SQLITE_NOMEM, "Unable to create sampler");
goto error;
}
sampler = owned_sampler;
llama_sampler_chain_add(sampler, llama_sampler_init_penalties(64, 1.1, 0, 0));
llama_sampler_chain_add(sampler, llama_sampler_init_greedy());
}

// the KV cache was cleared above, so the sampler must start clean too:
// rebuild grammar state, re-seed dist, clear penalty history
llama_sampler_reset(sampler);

// allocate output buffer
if (!buffer_create(&buffer, 0)) {
sqlite_context_result_error(context, SQLITE_NOMEM, "Out of memory");
goto error_sampler;
goto error;
}
buffer_initialized = true;

Expand All @@ -3215,19 +3237,19 @@ static void llm_text_run_vision (sqlite3_context *context, const char *text, int
int n = llama_token_to_piece(vocab, new_token_id, buf, sizeof(buf), 0, true);
if (n < 0) {
sqlite_context_result_error(context, SQLITE_ERROR, "Failed to convert token to piece");
goto error_sampler;
goto error;
}

if (!buffer_append(&buffer, buf, n, true)) {
sqlite_context_result_error(context, SQLITE_NOMEM, "Out of memory");
goto error_sampler;
goto error;
}

struct llama_batch batch = llama_batch_get_one(&new_token_id, 1);
int32_t drc = llama_decode(ctx, batch);
if (drc != 0) {
sqlite_context_result_error(context, SQLITE_ERROR, "Failed to decode during generation (%d: %s)", drc, llm_decode_error_string(drc));
goto error_sampler;
goto error;
}
}
}
Expand All @@ -3239,16 +3261,12 @@ static void llm_text_run_vision (sqlite3_context *context, const char *text, int
mtmd_input_chunks_free(chunks);
sqlite3_free(prompt_with_markers);
sqlite3_free(formatted_prompt);
if (!sampler_already_setup) llama_sampler_free(sampler);
if (owned_sampler) llama_sampler_free(owned_sampler);
return;

error_sampler:
if (!sampler_already_setup && ai->sampler) {
llama_sampler_free(ai->sampler);
ai->sampler = NULL;
}
error:
if (buffer_initialized) buffer_destroy(&buffer);
if (owned_sampler) llama_sampler_free(owned_sampler);
if (bitmaps) {
for (int i = 0; i < n_images; i++) if (bitmaps[i]) mtmd_bitmap_free(bitmaps[i]);
sqlite3_free(bitmaps);
Expand Down
2 changes: 1 addition & 1 deletion src/sqlite-ai.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@
extern "C" {
#endif

#define SQLITE_AI_VERSION "1.0.4"
#define SQLITE_AI_VERSION "1.0.5"

SQLITE_AI_API int sqlite3_ai_init (sqlite3 *db, char **pzErrMsg, const sqlite3_api_routines *pApi);

Expand Down
79 changes: 79 additions & 0 deletions tests/c/unittest.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1657,6 +1657,83 @@ static int test_text_generate_default_limit(const test_env *env) {
return 1;
}

// Regression: llm_text_generate() must be callable repeatedly on one connection.
// The implicit default sampler chain used to be published into ai->sampler and then
// freed on the success path, leaving a dangling pointer that crashed the second call.
static int test_text_generate_repeated(const test_env *env) {
sqlite3 *db = NULL;
if (open_db_and_load(env, &db) != SQLITE_OK) return 1;

const char *model = env->model_path ? env->model_path : DEFAULT_MODEL_PATH;
char sqlbuf[512];
snprintf(sqlbuf, sizeof(sqlbuf), "SELECT llm_model_load('%s');", model);
if (exec_expect_ok(env, db, sqlbuf) != 0) goto fail;
if (exec_expect_ok(env, db, "SELECT llm_context_create_textgen('context_size=1024,n_predict=32');") != 0) goto fail;

for (int i = 0; i < 3; ++i) {
char result[4096] = {0};
if (exec_query_text(env, db, "SELECT llm_text_generate('Say hello in one word.');", result, sizeof(result)) != 0) goto fail;
if (result[0] == '\0') {
fprintf(stderr, "[text_generate_repeated] call %d returned empty output\n", i + 1);
goto fail;
}
if (env->verbose) printf("[text_generate_repeated] call %d: %s\n", i + 1, result);
}

// deliberately no llm_sampler_free(): the implicit chain must not leak or dangle
if (exec_expect_ok(env, db, "SELECT llm_context_free();") != 0) goto fail;
if (exec_expect_ok(env, db, "SELECT llm_model_free();") != 0) goto fail;

sqlite3_close(db);
return assert_sqlite_memory_clean("text_generate_repeated", env);

fail:
if (db) sqlite3_close(db);
return 1;
}

// Regression: a user-configured grammar sampler must be reset between generations.
// Without llama_sampler_reset() the grammar stays in its terminal state after the
// first call and every later call samples EOG immediately, returning ''.
static int test_text_generate_repeated_grammar(const test_env *env) {
sqlite3 *db = NULL;
if (open_db_and_load(env, &db) != SQLITE_OK) return 1;

const char *model = env->model_path ? env->model_path : DEFAULT_MODEL_PATH;
char sqlbuf[512];
snprintf(sqlbuf, sizeof(sqlbuf), "SELECT llm_model_load('%s');", model);
if (exec_expect_ok(env, db, sqlbuf) != 0) goto fail;
if (exec_expect_ok(env, db, "SELECT llm_context_create_textgen('context_size=1024,n_predict=32');") != 0) goto fail;

// grammar needs the vocab of a loaded model, and a selecting sampler must
// terminate the chain, otherwise llama_sampler_sample() asserts
if (exec_expect_ok(env, db, "SELECT llm_sampler_create();") != 0) goto fail;
if (exec_expect_ok(env, db, "SELECT llm_sampler_init_grammar('root ::= \"yes\"', 'root');") != 0) goto fail;
if (exec_expect_ok(env, db, "SELECT llm_sampler_init_greedy();") != 0) goto fail;

for (int i = 0; i < 3; ++i) {
char result[4096] = {0};
if (exec_query_text(env, db, "SELECT llm_text_generate('Answer with yes.');", result, sizeof(result)) != 0) goto fail;
if (env->verbose) printf("[text_generate_repeated_grammar] call %d: '%s'\n", i + 1, result);
if (strcmp(result, "yes") != 0) {
fprintf(stderr, "[text_generate_repeated_grammar] call %d expected 'yes', got '%s'\n", i + 1, result);
goto fail;
}
}

// an explicit free followed by ai_free() at close must not double-free
if (exec_expect_ok(env, db, "SELECT llm_sampler_free();") != 0) goto fail;
if (exec_expect_ok(env, db, "SELECT llm_context_free();") != 0) goto fail;
if (exec_expect_ok(env, db, "SELECT llm_model_free();") != 0) goto fail;

sqlite3_close(db);
return assert_sqlite_memory_clean("text_generate_repeated_grammar", env);

fail:
if (db) sqlite3_close(db);
return 1;
}

// ---------------------------------------------------------------------
// Audio / Whisper tests
// ---------------------------------------------------------------------
Expand Down Expand Up @@ -2041,6 +2118,8 @@ static const test_case TESTS[] = {
{"chat_respond_auto_init", test_chat_respond_auto_init},
{"chat_save_with_metadata", test_chat_save_with_metadata},
{"text_generate_default_limit", test_text_generate_default_limit},
{"text_generate_repeated", test_text_generate_repeated},
{"text_generate_repeated_grammar", test_text_generate_repeated_grammar},
{"llm_chat_double_save", test_llm_chat_double_save},
// Audio / Whisper tests
{"audio_transcribe_no_model", test_audio_transcribe_no_model},
Expand Down
Loading