From ba55a1c8bc64a4e76d2bb1ee6315b29ef4fb89b5 Mon Sep 17 00:00:00 2001 From: Andrea Donetti Date: Fri, 28 Aug 2026 17:25:55 +0200 Subject: [PATCH 1/3] fix(chat): release the cursor when llm_chat() xOpen fails llm_chat_cursor_open() allocates the ai_cursor before validating that the connection has a usable context, then returns SQLITE_ERROR without freeing it. sqlite does not call xClose for a cursor whose xOpen failed, so the allocation is lost for the lifetime of the connection - 48 bytes per rejected statement. Reachable today by querying the vtab with no context created: SELECT llm_model_load('model.gguf'); SELECT reply FROM llm_chat('hi'); -- errors, and leaks Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01Ks738YbgnnJBio2wLtpbAJ --- src/sqlite-ai.c | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/src/sqlite-ai.c b/src/sqlite-ai.c index 2d1adbd..3036c38 100644 --- a/src/sqlite-ai.c +++ b/src/sqlite-ai.c @@ -2139,7 +2139,12 @@ static int llm_chat_cursor_open (sqlite3_vtab *pVtab, sqlite3_vtab_cursor **ppCu c->ai = vtab->ai; ai_context *ai = c->ai; - if (llm_chat_check_context(ai) == false) return SQLITE_ERROR; + // sqlite never calls xClose for a cursor whose xOpen failed, so the cursor has to be + // released here or it leaks for the lifetime of the connection + if (llm_chat_check_context(ai) == false) { + sqlite3_free(c); + return SQLITE_ERROR; + } *ppCursor = (sqlite3_vtab_cursor *)c; return SQLITE_OK; From 46e68beed61b6e79f5ef0d390bd53db44b6bc17f Mon Sep 17 00:00:00 2001 From: Andrea Donetti Date: Fri, 28 Aug 2026 17:27:04 +0200 Subject: [PATCH 2/3] fix(error): report failures to the caller being served, not to stale sinks ai->context and ai->vtab are the error sink for the whole chat and sampler subsystem. Roughly twenty sites report through sqlite_common_set_error(ai->context, ai->vtab, ...) - llm_chat_run(), llm_chat_generate_response(), llm_chat_tokenize_input(), llm_chat_save_response(), llm_chat_check_context(), llm_sampler_check() - and none of them receives the caller directly. Only three places ever assigned those fields, so at nearly every report site the sink was whatever the previous statement happened to leave behind. Two consequences, both reachable from ordinary SQL: SELECT llm_context_create_chat('context_size=512'); SELECT llm_chat_respond('hi'); -- publishes this statement's context SELECT llm_context_free(); SELECT llm_chat_restore('x'); -- reports into the finalized statement That last line called sqlite3_result_error() on freed memory: SIGSEGV. Without the llm_chat_respond() line the sink is still NULL and the same call returned NULL with no error at all. After a vtab scan the sink is the vtab instead, so the message went into vtab->zErrMsg, which sqlite only imports immediately after a vtab method - silently discarded, and the sqlite3_vmprintf buffer leaked. llm_context_create_with_options() had the same bug from the other direction: it reported a llama_init_from_model() failure through ai->context/ai->vtab while every other error in that function used the context parameter it was handed. On a fresh connection both fields are NULL and the failure vanished entirely, so the function returned NULL as though it had succeeded. The rule now: every entry point publishes its own sink before running anything that can fail. Scalars own a sqlite3_context, vtab methods own the sqlite3_vtab. Applied at all of them: - chat scalars: create, free, save, restore, system_prompt, respond - vtab: xConnect, xOpen, xFilter, xNext, xClose. xFilter and xNext matter as much as xOpen - both reach llm_chat_run() and llm_chat_generate_response(), and an interleaved scalar statement can repoint the sink mid-scan. - the 16 llm_sampler_* scalars, which reach llm_sampler_check(). Its allocation failure is the one report site with no caller of its own, so without this an OOM right after a vtab scan wrote into a vtab->zErrMsg that sqlite no longer imports: message lost, buffer leaked. llm_chat_disconnect() additionally left ai->vtab dangling at freed memory after the vtab was released, which is what turned the second case above from a lost message into a use-after-free. The ai_context outlives the vtab, so the back-pointer has to be cleared. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01Ks738YbgnnJBio2wLtpbAJ --- src/sqlite-ai.c | 68 +++++++++++++++++++++++++++++++++++++++++----- tests/c/unittest.c | 40 +++++++++++++++++++++++++++ 2 files changed, 101 insertions(+), 7 deletions(-) diff --git a/src/sqlite-ai.c b/src/sqlite-ai.c index 3036c38..371dd18 100644 --- a/src/sqlite-ai.c +++ b/src/sqlite-ai.c @@ -946,6 +946,28 @@ static bool llm_check_context (sqlite3_context *context) { return true; } +// ai->context / ai->vtab are the error sink for the whole chat and sampler subsystem: +// roughly twenty sites report through sqlite_common_set_error(ai->context, ai->vtab, ...), +// including llm_chat_run(), llm_chat_generate_response(), llm_chat_tokenize_input(), +// llm_chat_save_response(), llm_chat_check_context() and llm_sampler_check(). None of +// them receives the caller directly, so both fields must name something that is alive +// *right now* - a stale sqlite3_context belongs to a finalized statement, and reporting +// into it calls sqlite3_result_error() on freed memory. +// +// The rule: every entry point publishes its own sink before running anything that can +// fail. Scalars own a sqlite3_context; vtab methods own the sqlite3_vtab. +static inline void llm_error_sink_scalar (ai_context *ai, sqlite3_context *context) { + if (!ai) return; + ai->context = context; + ai->vtab = NULL; +} + +static inline void llm_error_sink_vtab (ai_context *ai, sqlite3_vtab *vtab) { + if (!ai) return; + ai->context = NULL; + ai->vtab = vtab; +} + // MARK: - Chat Messages - bool llm_messages_append (ai_messages *list, const char *role, const char *content) { @@ -2098,8 +2120,7 @@ static int llm_chat_connect (sqlite3 *db, void *pAux, int argc, const char *cons vtab->ai = ai; ai->db = db; - ai->context = NULL; - ai->vtab = (sqlite3_vtab *)vtab; + llm_error_sink_vtab(ai, (sqlite3_vtab *)vtab); *ppVtab = (sqlite3_vtab *)vtab; return SQLITE_OK; @@ -2107,6 +2128,12 @@ static int llm_chat_connect (sqlite3 *db, void *pAux, int argc, const char *cons static int llm_chat_disconnect (sqlite3_vtab *pVtab) { ai_vtab *vtab = (ai_vtab *)pVtab; + + // the ai_context outlives the vtab: leaving ai->vtab pointing at freed memory turns + // the next sqlite_common_set_error() that routes through it into a use-after-free. + ai_context *ai = vtab->ai; + if (ai && ai->vtab == pVtab) ai->vtab = NULL; + sqlite3_free(vtab); return SQLITE_OK; } @@ -2139,6 +2166,8 @@ static int llm_chat_cursor_open (sqlite3_vtab *pVtab, sqlite3_vtab_cursor **ppCu c->ai = vtab->ai; ai_context *ai = c->ai; + llm_error_sink_vtab(ai, (sqlite3_vtab *)vtab); + // sqlite never calls xClose for a cursor whose xOpen failed, so the cursor has to be // released here or it leaks for the lifetime of the connection if (llm_chat_check_context(ai) == false) { @@ -2153,6 +2182,7 @@ static int llm_chat_cursor_open (sqlite3_vtab *pVtab, sqlite3_vtab_cursor **ppCu static int llm_chat_cursor_close (sqlite3_vtab_cursor *cur) { ai_cursor *c = (ai_cursor *)cur; ai_context *ai = c->ai; + llm_error_sink_vtab(ai, (sqlite3_vtab *)c->vtab); // save response before freeing the cursor ai_messages *messages = &ai->chat.messages; @@ -2166,6 +2196,7 @@ static int llm_chat_cursor_close (sqlite3_vtab_cursor *cur) { static int llm_chat_cursor_next (sqlite3_vtab_cursor *cur) { ai_cursor *c = (ai_cursor *)cur; + llm_error_sink_vtab(c->ai, (sqlite3_vtab *)c->vtab); if (!llm_chat_generate_response (c->ai, c, NULL)) return SQLITE_ERROR; c->rowid++; return SQLITE_OK; @@ -2194,6 +2225,7 @@ static int llm_chat_cursor_filter (sqlite3_vtab_cursor *cur, int idxNum, const c ai_cursor *c = (ai_cursor *)cur; ai_context *ai = c->ai; ai_vtab *vtab = c->vtab; + llm_error_sink_vtab(ai, (sqlite3_vtab *)vtab); // sanity check arguments if (argc != 1) { @@ -2242,6 +2274,7 @@ static sqlite3_module llm_chat = { // MARK: - static void llm_chat_free (sqlite3_context *context, int argc, sqlite3_value **argv) { + llm_error_sink_scalar((ai_context *)sqlite3_user_data(context), context); ai_chat_release((ai_context *)sqlite3_user_data(context)); } @@ -2249,6 +2282,7 @@ static void llm_chat_create (sqlite3_context *context, int argc, sqlite3_value * if (llm_check_context(context) == false) return; ai_context *ai = (ai_context *)sqlite3_user_data(context); + llm_error_sink_scalar(ai, context); // clean-up old chat (if any) llm_chat_free(context, argc, argv); @@ -2274,6 +2308,7 @@ static bool llm_chat_check_tables (sqlite3_context *context) { static void llm_chat_save (sqlite3_context *context, int argc, sqlite3_value **argv) { ai_context *ai = (ai_context *)sqlite3_user_data(context); + llm_error_sink_scalar(ai, context); if (llm_chat_check_tables(context) == false) return; // sanity check if there is something to save @@ -2355,11 +2390,12 @@ static void llm_chat_restore (sqlite3_context *context, int argc, sqlite3_value int types[] = {SQLITE_TEXT}; if (sqlite_sanity_function(context, "llm_chat_restore", argc, argv, 1, types, false, false) == false) return; + ai_context *ai = (ai_context *)sqlite3_user_data(context); + llm_error_sink_scalar(ai, context); + // free old chat (if any) llm_chat_free(context, 0, NULL); - ai_context *ai = (ai_context *)sqlite3_user_data(context); - // re-initialize chat state (UUID, buffers, tokens) if (llm_chat_check_context(ai) == false) return; @@ -2411,6 +2447,7 @@ static void llm_chat_respond (sqlite3_context *context, int argc, sqlite3_value } ai_context *ai = (ai_context *)sqlite3_user_data(context); + llm_error_sink_scalar(ai, context); if (!ai->model) { sqlite_context_result_error(context, SQLITE_ERROR, "No model loaded"); return; @@ -2418,8 +2455,6 @@ static void llm_chat_respond (sqlite3_context *context, int argc, sqlite3_value if (llm_chat_check_context(ai) == false) return; const char *user_prompt = (const char *)sqlite3_value_text(argv[0]); - ai->context = context; - ai->vtab = NULL; ai->chat.token_count = 0; buffer_reset(&ai->chat.response); @@ -2449,6 +2484,7 @@ static void llm_chat_system_prompt(sqlite3_context *context, int argc, sqlite3_v return; ai_context *ai = (ai_context *)sqlite3_user_data(context); + llm_error_sink_scalar(ai, context); if (llm_chat_check_context(ai) == false) return; @@ -2494,6 +2530,7 @@ static void llm_chat_system_prompt(sqlite3_context *context, int argc, sqlite3_v static void llm_sampler_init_greedy (sqlite3_context *context, int argc, sqlite3_value **argv) { ai_context *ai = (ai_context *)sqlite3_user_data(context); + llm_error_sink_scalar(ai, context); llm_sampler_check(ai); if (ai->sampler) llama_sampler_chain_add(ai->sampler, llama_sampler_init_greedy()); } @@ -2505,6 +2542,7 @@ static void llm_sampler_init_dist (sqlite3_context *context, int argc, sqlite3_v } ai_context *ai = (ai_context *)sqlite3_user_data(context); + llm_error_sink_scalar(ai, context); llm_sampler_check(ai); if (ai->sampler) { int32_t seed = (argc == 1) ? (int32_t)sqlite3_value_int64(argv[0]) : (int32_t)LLAMA_DEFAULT_SEED; @@ -2519,6 +2557,7 @@ static void llm_sampler_init_top_k (sqlite3_context *context, int argc, sqlite3_ if (sqlite_sanity_function(context, "llm_sampler_init_top_k", argc, argv, 1, types, true, false) == false) return; ai_context *ai = (ai_context *)sqlite3_user_data(context); + llm_error_sink_scalar(ai, context); llm_sampler_check(ai); if (ai->sampler) { int32_t k = (int32_t)sqlite3_value_int64(argv[0]); @@ -2533,6 +2572,7 @@ static void llm_sampler_init_top_p (sqlite3_context *context, int argc, sqlite3_ if (sqlite_sanity_function(context, "llm_sampler_init_top_p", argc, argv, 2, types, true, false) == false) return; ai_context *ai = (ai_context *)sqlite3_user_data(context); + llm_error_sink_scalar(ai, context); llm_sampler_check(ai); if (ai->sampler) { float p = (float)sqlite3_value_double(argv[0]); @@ -2548,6 +2588,7 @@ static void llm_sampler_init_min_p (sqlite3_context *context, int argc, sqlite3_ if (sqlite_sanity_function(context, "llm_sampler_init_min_p", argc, argv, 2, types, true, false) == false) return; ai_context *ai = (ai_context *)sqlite3_user_data(context); + llm_error_sink_scalar(ai, context); llm_sampler_check(ai); if (ai->sampler) { float p = (float)sqlite3_value_double(argv[0]); @@ -2563,6 +2604,7 @@ static void llm_sampler_init_typical (sqlite3_context *context, int argc, sqlite if (sqlite_sanity_function(context, "llm_sampler_init_typical", argc, argv, 2, types, true, false) == false) return; ai_context *ai = (ai_context *)sqlite3_user_data(context); + llm_error_sink_scalar(ai, context); llm_sampler_check(ai); if (ai->sampler) { float p = (float)sqlite3_value_double(argv[0]); @@ -2576,6 +2618,7 @@ static void llm_sampler_init_temp (sqlite3_context *context, int argc, sqlite3_v if (sqlite_sanity_function(context, "llm_sampler_init_temp", argc, argv, 1, types, true, false) == false) return; ai_context *ai = (ai_context *)sqlite3_user_data(context); + llm_error_sink_scalar(ai, context); llm_sampler_check(ai); if (ai->sampler) { float t = (float)sqlite3_value_double(argv[0]); @@ -2590,6 +2633,7 @@ static void llm_sampler_init_temp_ext (sqlite3_context *context, int argc, sqlit if (sqlite_sanity_function(context, "llm_sampler_init_temp_ext", argc, argv, 3, types, true, false) == false) return; ai_context *ai = (ai_context *)sqlite3_user_data(context); + llm_error_sink_scalar(ai, context); llm_sampler_check(ai); if (ai->sampler) { float t = (float)sqlite3_value_double(argv[0]); @@ -2606,6 +2650,7 @@ static void llm_sampler_init_xtc (sqlite3_context *context, int argc, sqlite3_va if (sqlite_sanity_function(context, "llm_sampler_init_xtc", argc, argv, 4, types, true, false) == false) return; ai_context *ai = (ai_context *)sqlite3_user_data(context); + llm_error_sink_scalar(ai, context); llm_sampler_check(ai); if (ai->sampler) { float p = (float)sqlite3_value_double(argv[0]); @@ -2623,6 +2668,7 @@ static void llm_sampler_init_top_n_sigma (sqlite3_context *context, int argc, sq if (sqlite_sanity_function(context, "llm_sampler_init_top_n_sigma", argc, argv, 1, types, true, false) == false) return; ai_context *ai = (ai_context *)sqlite3_user_data(context); + llm_error_sink_scalar(ai, context); llm_sampler_check(ai); if (ai->sampler) { float n = (float)sqlite3_value_double(argv[0]); @@ -2643,6 +2689,7 @@ static void llm_sampler_init_mirostat (sqlite3_context *context, int argc, sqlit return; } + llm_error_sink_scalar(ai, context); llm_sampler_check(ai); if (ai->sampler) { uint32_t seed = (uint32_t)sqlite3_value_int64(argv[0]); @@ -2660,6 +2707,7 @@ static void llm_sampler_init_mirostat_v2 (sqlite3_context *context, int argc, sq if (sqlite_sanity_function(context, "llm_sampler_init_mirostat_v2", argc, argv, 3, types, true, false) == false) return; ai_context *ai = (ai_context *)sqlite3_user_data(context); + llm_error_sink_scalar(ai, context); llm_sampler_check(ai); if (ai->sampler) { uint32_t seed = (uint32_t)sqlite3_value_int64(argv[0]); @@ -2680,6 +2728,7 @@ static void llm_sampler_init_grammar (sqlite3_context *context, int argc, sqlite return; } + llm_error_sink_scalar(ai, context); llm_sampler_check(ai); if (ai->sampler) { const char *grammar_str = (const char *)sqlite3_value_text(argv[0]); @@ -2696,6 +2745,7 @@ static void llm_sampler_init_infill (sqlite3_context *context, int argc, sqlite3 return; } + llm_error_sink_scalar(ai, context); llm_sampler_check(ai); if (ai->sampler) { llama_sampler_chain_add(ai->sampler, llama_sampler_init_infill(vocab)); @@ -2707,6 +2757,7 @@ static void llm_sampler_init_penalties (sqlite3_context *context, int argc, sqli if (sqlite_sanity_function(context, "llm_sampler_init_penalties", argc, argv, 4, types, true, false) == false) return; ai_context *ai = (ai_context *)sqlite3_user_data(context); + llm_error_sink_scalar(ai, context); llm_sampler_check(ai); if (ai->sampler) { int32_t penalty_last_n = (int32_t)sqlite3_value_int64(argv[0]); @@ -2780,6 +2831,7 @@ static void llm_sampler_create (sqlite3_context *context, int argc, sqlite3_valu ai_context *ai = (ai_context *)sqlite3_user_data(context); if (ai->sampler) llama_sampler_free(ai->sampler); ai->sampler = NULL; + llm_error_sink_scalar(ai, context); llm_sampler_check(ai); } @@ -2831,7 +2883,9 @@ static bool llm_context_create_with_options (sqlite3_context *context, ai_contex struct llama_context *ctx = llama_init_from_model(ai->model, ctx_params); if (!ctx) { - sqlite_common_set_error(ai->context, ai->vtab, SQLITE_ERROR, "Unable to create context from model"); + // report to the context we were handed, not ai->context/ai->vtab: those are + // leftovers from an earlier statement and may already be freed. + sqlite_context_result_error(context, SQLITE_ERROR, "Unable to create context from model"); return false; } diff --git a/tests/c/unittest.c b/tests/c/unittest.c index bea5523..f983ab2 100644 --- a/tests/c/unittest.c +++ b/tests/c/unittest.c @@ -489,6 +489,45 @@ static int test_llm_chat_vtab(const test_env *env) { return 1; } +// ai->context / ai->vtab are the error sink for the whole chat subsystem, and every +// entry point has to publish its own before running anything that can fail. When it was +// left holding a finalized statement's sqlite3_context, llm_chat_restore() reported into +// freed memory and took the process down with it. +static int test_chat_error_sink_after_statement(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 (install_seeded_chat_sampler(env, db) != 0) goto fail; + + // (a) after a scalar chat call, whose sqlite3_context dies with its statement + if (exec_expect_ok(env, db, "SELECT llm_context_create_chat('context_size=512,n_predict=4');") != 0) goto fail; + if (exec_expect_ok(env, db, "SELECT llm_chat_respond('hi');") != 0) goto fail; + if (exec_expect_ok(env, db, "SELECT llm_context_free();") != 0) goto fail; + if (exec_expect_error(env, db, "SELECT llm_chat_restore('deadbeef');", "No context found") != 0) goto fail; + + // (b) after a vtab scan, which used to leave ai->vtab armed so the message went into + // vtab->zErrMsg and was discarded + if (exec_expect_ok(env, db, "SELECT llm_context_create_chat('context_size=512,n_predict=4');") != 0) goto fail; + if (exec_expect_ok(env, db, "SELECT count(*) FROM llm_chat('hi');") != 0) goto fail; + if (exec_expect_ok(env, db, "SELECT llm_context_free();") != 0) goto fail; + if (exec_expect_error(env, db, "SELECT llm_chat_restore('deadbeef');", "No context found") != 0) goto fail; + + // (c) with no context ever created, the sink is still unset + if (exec_expect_error(env, db, "SELECT llm_chat_system_prompt('hi');", "No context found") != 0) goto fail; + + if (exec_expect_ok(env, db, "SELECT llm_model_free();") != 0) goto fail; + sqlite3_close_v2(db); + return assert_sqlite_memory_clean("chat_error_sink_after_statement", env); + +fail: + if (db) sqlite3_close_v2(db); + return 1; +} + static int test_llm_embed_generate(const test_env *env) { sqlite3 *db = NULL; if (open_db_and_load(env, &db) != SQLITE_OK) { @@ -2584,6 +2623,7 @@ static const test_case TESTS[] = { {"issue15_llm_chat_without_context", test_issue15_chat_without_context}, {"llm_chat_respond_repeated", test_llm_chat_respond_repeated}, {"llm_chat_vtab", test_llm_chat_vtab}, + {"chat_error_sink_after_statement", test_chat_error_sink_after_statement}, {"test_llm_embed_generate", test_llm_embed_generate}, {"llm_embed_generate_basic", test_llm_embed_generate_basic}, {"llm_embedding_then_chat", test_llm_embedding_then_chat}, From 06554a21d95131c16fdfa80bcee8ce494d16abdd Mon Sep 17 00:00:00 2001 From: Andrea Donetti Date: Mon, 31 Aug 2026 11:04:57 +0200 Subject: [PATCH 3/3] fix(context): reject generation on a context that cannot produce tokens Closes #33. llm_text_generate() only checked that a context existed, so calling it against an embedding context ran on incompatible state. Pooled embeddings leave no per-token logits, so the first sampled token reads as EOG, the loop breaks immediately, and the empty buffer is returned as '' with no error. The chat paths are worse: llm_chat_respond() and the llm_chat() vtab reach llama_sampler_sample(), which dereferences that buffer. Track what the active llama context can actually do and reject the operations that need per-token logits with SQLITE_MISUSE. The kind is derived from how the context is configured, not from which llm_context_create_* wrapper was called, because two different things make a context an embedding context: - generate_embedding=1 was passed. llm_context_create() is a general-purpose constructor and API.md documents llm_context_create('generate_embedding=1,normalize_embedding=1,pooling_type=mean') as equivalent to llm_context_create_embedding(), so classifying by call site would leave #33 unfixed on the documented generic form. - the model pools by default. A BERT-family model (all-MiniLM, nomic-embed) carries its pooling type in its own GGUF, so llm_context_create('context_size=512,embedding_type=FLOAT32') on one yields a pooling context with no embedding settings passed at all - and that is the flow API.md recommends for embeddings, so #33 reproduced there too. Resolved pooling alone cannot decide this: forcing pooling_type=mean on a generative model still generates correctly. What identifies an embedding model is pooling the caller did not ask for, which llama makes visible by resolving LLAMA_POOLING_TYPE_UNSPECIFIED from the model's own hparams. So the test is ctx_params.embeddings, or an unrequested resolved pooling type. That makes "did the caller ask for pooling" load-bearing, which exposed a bug in the option parser: generate_embedding forced pooling_type = MEAN for any value, 0 included. An explicit generate_embedding=0 therefore both configured a pooling context while asking for embeddings to be off, and marked pooling as caller-requested - hiding embedding models from the check. The side effects now apply only when embeddings are actually enabled. Note llama_model_has_decoder() is no use here - it returns true for everything except T5ENCODER, BERT included. Text generation and chat contexts share one kind: llm_context_create_chat() and llm_context_create_textgen() both pass the same empty option string, so the contexts are byte-for-byte identical and must stay interchangeable - the README vision example creates a chat context and then calls llm_text_generate(). Embedding generation is deliberately NOT gated. What it needs is a context that pools, and llm_embed_generate_run() already checks the resolved pooling type - a better test than the declared kind, and one that keeps working on every configuration above. In the vtab the check is skipped when no context exists at all, so a missing context still reports the actionable "No context found" rather than a kind mismatch against LLM_CONTEXT_NONE. Known limitation, documented in API.md: passing pooling_type explicitly on an embedding model opts out of the second test and lets generation return '' again. That is the cost of not misreading a caller-forced pooling type on a generative model, which is the more common case. Tests cover both directions on both model families: the rejections, and that a generic-constructor embedding context, a caller-forced pooling type, and a chat context all keep working. Covering the embedding-model half needs an encoder-style model, so the suite now pulls all-MiniLM-L6-v2 (25MB) alongside the existing generative model; that test skips without --embed-model. CI never reaches the Makefile's download rule for the other models: the download-models job fetches them on the host and every build job restores them from cache. The new model is wired through that same path - URL hash, job output, host-side restore/download/verify, and a cache restore plus directory in the build jobs. Without it `make test` fell through to curl, which the Alpine container used by the linux-musl arm64 jobs does not have (Error 127). Bumps SQLITE_AI_VERSION to 1.0.8. The workflow reads it via `make version` and tags from it, so the bump is what makes the release job publish rather than warn that the version matches the latest release. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01Ks738YbgnnJBio2wLtpbAJ --- .github/workflows/main.yml | 30 ++++++ API.md | 37 ++++++- Makefile | 15 ++- src/sqlite-ai.c | 120 ++++++++++++++++++++-- src/sqlite-ai.h | 2 +- tests/c/unittest.c | 202 ++++++++++++++++++++++++++++++++++++- 6 files changed, 389 insertions(+), 17 deletions(-) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index ac8a252..e916a62 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -15,6 +15,9 @@ env: GGUF_MODEL_DIR: tests/models/unsloth/gemma-3-270m-it-GGUF GGUF_MODEL_NAME: gemma-3-270m-it-UD-IQ2_M.gguf GGUF_MODEL_URL: https://huggingface.co/unsloth/gemma-3-270m-it-GGUF/resolve/main/gemma-3-270m-it-UD-IQ2_M.gguf + EMBED_MODEL_DIR: tests/models/Mungert/all-MiniLM-L6-v2-GGUF + EMBED_MODEL_NAME: all-MiniLM-L6-v2-q8_0.gguf + EMBED_MODEL_URL: https://huggingface.co/Mungert/all-MiniLM-L6-v2-GGUF/resolve/main/all-MiniLM-L6-v2-q8_0.gguf WHISPER_MODEL_DIR: tests/models/ggerganov/whisper-tiny WHISPER_MODEL_NAME: ggml-tiny.bin WHISPER_MODEL_URL: https://huggingface.co/ggerganov/whisper.cpp/resolve/main/ggml-tiny.bin @@ -31,6 +34,8 @@ jobs: outputs: gguf-cache-key: gguf-${{ steps.meta.outputs.gguf-hash }} gguf-model-path: ${{ env.GGUF_MODEL_DIR }}/${{ env.GGUF_MODEL_NAME }} + embed-cache-key: embed-${{ steps.meta.outputs.embed-hash }} + embed-model-path: ${{ env.EMBED_MODEL_DIR }}/${{ env.EMBED_MODEL_NAME }} whisper-cache-key: whisper-${{ steps.meta.outputs.whisper-hash }} whisper-model-path: ${{ env.WHISPER_MODEL_DIR }}/${{ env.WHISPER_MODEL_NAME }} audio-cache-key: audio-${{ steps.meta.outputs.audio-hash }} @@ -43,20 +48,24 @@ jobs: run: | if command -v sha256sum >/dev/null 2>&1; then gguf_hash=$(echo -n "${{ env.GGUF_MODEL_URL }}" | sha256sum | cut -d' ' -f1) + embed_hash=$(echo -n "${{ env.EMBED_MODEL_URL }}" | sha256sum | cut -d' ' -f1) whisper_hash=$(echo -n "${{ env.WHISPER_MODEL_URL }}" | sha256sum | cut -d' ' -f1) audio_hash=$(echo -n "${{ env.AUDIO_TEST_WAV_URL }}" | sha256sum | cut -d' ' -f1) else gguf_hash=$(echo -n "${{ env.GGUF_MODEL_URL }}" | shasum -a 256 | cut -d' ' -f1) + embed_hash=$(echo -n "${{ env.EMBED_MODEL_URL }}" | shasum -a 256 | cut -d' ' -f1) whisper_hash=$(echo -n "${{ env.WHISPER_MODEL_URL }}" | shasum -a 256 | cut -d' ' -f1) audio_hash=$(echo -n "${{ env.AUDIO_TEST_WAV_URL }}" | shasum -a 256 | cut -d' ' -f1) fi echo "gguf-hash=$gguf_hash" >> "$GITHUB_OUTPUT" + echo "embed-hash=$embed_hash" >> "$GITHUB_OUTPUT" echo "whisper-hash=$whisper_hash" >> "$GITHUB_OUTPUT" echo "audio-hash=$audio_hash" >> "$GITHUB_OUTPUT" - name: Prepare directories run: | mkdir -p "${{ env.GGUF_MODEL_DIR }}" + mkdir -p "${{ env.EMBED_MODEL_DIR }}" mkdir -p "${{ env.WHISPER_MODEL_DIR }}" mkdir -p "${{ env.AUDIO_TEST_DIR }}" @@ -74,6 +83,20 @@ jobs: - name: Verify GGUF model run: test -f "${{ env.GGUF_MODEL_DIR }}/${{ env.GGUF_MODEL_NAME }}" + - name: Restore embedding model cache + id: cache-embed + uses: actions/cache@v4 + with: + path: ${{ env.EMBED_MODEL_DIR }}/${{ env.EMBED_MODEL_NAME }} + key: embed-${{ steps.meta.outputs.embed-hash }} + + - name: Download embedding model + if: steps.cache-embed.outputs.cache-hit != 'true' + run: curl -L --fail --retry 3 "${{ env.EMBED_MODEL_URL }}" -o "${{ env.EMBED_MODEL_DIR }}/${{ env.EMBED_MODEL_NAME }}" + + - name: Verify embedding model + run: test -f "${{ env.EMBED_MODEL_DIR }}/${{ env.EMBED_MODEL_NAME }}" + - name: Restore Whisper cache id: cache-whisper uses: actions/cache@v4 @@ -205,6 +228,7 @@ jobs: - name: Prepare test asset directories run: | mkdir -p "${{ env.GGUF_MODEL_DIR }}" + mkdir -p "${{ env.EMBED_MODEL_DIR }}" mkdir -p "${{ env.WHISPER_MODEL_DIR }}" mkdir -p "${{ env.AUDIO_TEST_DIR }}" @@ -214,6 +238,12 @@ jobs: path: ${{ needs.download-models.outputs.gguf-model-path }} key: ${{ needs.download-models.outputs.gguf-cache-key }} + - name: Restore embedding model cache + uses: actions/cache@v4 + with: + path: ${{ needs.download-models.outputs.embed-model-path }} + key: ${{ needs.download-models.outputs.embed-cache-key }} + - name: Restore Whisper cache uses: actions/cache@v4 with: diff --git a/API.md b/API.md index 5db77bd..33354a3 100644 --- a/API.md +++ b/API.md @@ -158,6 +158,21 @@ Creates a new inference context with comma separated key=value configuration. **Context must explicitly created before performing any AI operation!** +The context is classified from how it is actually configured, not from which constructor +you call. It is an *embedding* context when either `generate_embedding=1` was passed — so +`llm_context_create('generate_embedding=1,...')` behaves exactly like +`llm_context_create_embedding()` — or the model pools by default, which is how +BERT-family embedding models (`all-MiniLM`, `nomic-embed`) describe themselves in their +GGUF. Anything else is a *text generation* context. Forcing `pooling_type` yourself opts out of that second +test in both directions: on a generative model it does **not** make the context an +embedding one, and on an embedding model it suppresses the detection, so +`llm_text_generate()` is allowed again and returns an empty string. Only pooling you did +not ask for identifies the model as an embedding one — if you want the check, leave +`pooling_type` unset, or use `llm_context_create_embedding()`. Text generation and chat reject an +embedding context with `SQLITE_MISUSE`, since it produces no per-token logits to sample +from. Embedding generation is not restricted this way — it needs pooling rather than a +particular constructor, and checks for that directly. + ## context_settings The following keys are available in context_settings: @@ -165,7 +180,7 @@ The following keys are available in context_settings: | Key | Type | Meaning | | ------------------------| -------- | ---------------------------------------------------------------- | -| `generate_embedding` | `1 or 0` | Force the model to generate embeddings. | +| `generate_embedding` | `1 or 0` | Force the model to generate embeddings. This is what marks the context as an *embedding* context, which makes `llm_text_generate()`, `llm_chat_respond()` and the `llm_chat()` vtab reject it. Also forces `pooling_type` to `mean`, and `n_ubatch` is clamped to `n_batch` for embedding contexts. | | `normalize_embedding` | `1 or 0` | Force normalization during embedding generation (default to 1). | | `json_output` | `1 or 0` | Force JSON output in embedding generation (default to 0). | | `max_tokens` | `number` | Set a maximum number of tokens in input. If input is too large then an error is returned. | @@ -218,7 +233,6 @@ The following keys are available in context_settings: | Key | Type | Meaning | | -------------- | ------- | ------------------------------------------------------------------------------------------------------------------------------------------ | -| `embeddings` | `1 or 0` | If `1`, extract embeddings (with logits). Used by the embedding preset. | | `offload_kqv` | `1 or 0` | Offload KQV ops (incl. KV cache) to GPU. | | `no_perf` | `1 or 0` | Disable performance timing. | | `op_offload` | `1 or 0` | Offload host tensor ops to device. | @@ -685,6 +699,12 @@ Leave `json_output` off when storing embeddings for [sqlite-vector](https://github.com/sqliteai/sqlite-vector): the BLOB is already layout-compatible, so insert it directly rather than wrapping it. +Requires a context that pools token embeddings. `llm_context_create_embedding()` +guarantees that; a context built another way also works whenever pooling resolves to +something other than `none` — embedding models (BERT-family, such as `all-MiniLM` or +`nomic-embed-text`) inherit mean pooling from the GGUF, so they need no extra settings. A +context with no pooling fails with *"Embedding generation requires pooling"*. + **Example:** ```sql @@ -706,6 +726,13 @@ Generates a full-text completion based on input, with optional configuration pro When a vision model is loaded via `llm_vision_load()`, you can pass one or more images as additional arguments. Images can be file paths (TEXT) or raw image data (BLOB). Supported image formats: JPG, PNG, BMP, GIF. +Requires a text generation context. Contexts from `llm_context_create_textgen()`, +`llm_context_create_chat()` and a plain `llm_context_create()` are all accepted — they are +configured identically. Calling this against an embedding context fails with +`SQLITE_MISUSE` rather than returning an empty string, which also covers embedding +*models*: on a BERT-family model every context is an embedding context, so generation is +rejected there no matter which constructor was used. + **Examples:** ```sql @@ -750,6 +777,12 @@ Returns unique chat UUIDv7 value. If no chat is explicitly created, one will be created automatically when needed — but the UUID is needed for `llm_chat_save()` / `llm_chat_restore()`. +`llm_chat_respond()` and the `llm_chat()` virtual table require a text generation +context, because they decode and sample; running either against an embedding context +(`generate_embedding=1`) fails with `SQLITE_MISUSE`. `llm_chat_create()`, +`llm_chat_restore()` and `llm_chat_system_prompt()` only build up in-memory message state, +so they work regardless of the active context. + **Example:** ```sql diff --git a/Makefile b/Makefile index 14663ec..ea1de35 100644 --- a/Makefile +++ b/Makefile @@ -45,6 +45,13 @@ GGUF_MODEL_NAME ?= gemma-3-270m-it-UD-IQ2_M.gguf GGUF_MODEL_URL ?= https://huggingface.co/unsloth/gemma-3-270m-it-GGUF/resolve/main/gemma-3-270m-it-UD-IQ2_M.gguf GGUF_MODEL_PATH := $(GGUF_MODEL_DIR)/$(GGUF_MODEL_NAME) +# an embedding model (BERT-family) is needed to cover the context-kind checks: it pools +# by default from its own GGUF, which a generative model never does +EMBED_MODEL_DIR ?= tests/models/Mungert/all-MiniLM-L6-v2-GGUF +EMBED_MODEL_NAME ?= all-MiniLM-L6-v2-q8_0.gguf +EMBED_MODEL_URL ?= https://huggingface.co/Mungert/all-MiniLM-L6-v2-GGUF/resolve/main/all-MiniLM-L6-v2-q8_0.gguf +EMBED_MODEL_PATH := $(EMBED_MODEL_DIR)/$(EMBED_MODEL_NAME) + WHISPER_MODEL_DIR ?= tests/models/ggerganov/whisper-tiny WHISPER_MODEL_NAME ?= ggml-tiny.bin WHISPER_MODEL_URL ?= https://huggingface.co/ggerganov/whisper.cpp/resolve/main/ggml-tiny.bin @@ -245,6 +252,10 @@ $(GGUF_MODEL_PATH): @mkdir -p $(GGUF_MODEL_DIR) curl -L --fail --retry 3 -o $@ $(GGUF_MODEL_URL) +$(EMBED_MODEL_PATH): + @mkdir -p $(EMBED_MODEL_DIR) + curl -L --fail --retry 3 -o $@ $(EMBED_MODEL_URL) + $(WHISPER_MODEL_PATH): @mkdir -p $(WHISPER_MODEL_DIR) curl -L --fail --retry 3 -o $@ $(WHISPER_MODEL_URL) @@ -255,14 +266,14 @@ $(AUDIO_TEST_WAV): TEST_DEPS := $(TARGET) ifeq ($(SKIP_UNITTEST),0) -TEST_DEPS += $(CTEST_BIN) $(GGUF_MODEL_PATH) $(WHISPER_MODEL_PATH) $(AUDIO_TEST_WAV) +TEST_DEPS += $(CTEST_BIN) $(GGUF_MODEL_PATH) $(EMBED_MODEL_PATH) $(WHISPER_MODEL_PATH) $(AUDIO_TEST_WAV) endif test: $(TEST_DEPS) @echo "Running sqlite3 CLI smoke test (ensures .load works)..." $(SQLITE3) ":memory:" -cmd ".bail on" ".load ./dist/ai" "SELECT ai_version();" ifeq ($(SKIP_UNITTEST),0) - $(CTEST_BIN) --extension "$(TARGET)" --model "$(GGUF_MODEL_PATH)" --whisper-model "$(WHISPER_MODEL_PATH)" --audio "$(AUDIO_TEST_WAV)" + $(CTEST_BIN) --extension "$(TARGET)" --model "$(GGUF_MODEL_PATH)" --embed-model "$(EMBED_MODEL_PATH)" --whisper-model "$(WHISPER_MODEL_PATH)" --audio "$(AUDIO_TEST_WAV)" else @echo "Skipping C unit tests (SKIP_UNITTEST=$(SKIP_UNITTEST))." endif diff --git a/src/sqlite-ai.c b/src/sqlite-ai.c index 371dd18..67f51d6 100644 --- a/src/sqlite-ai.c +++ b/src/sqlite-ai.c @@ -140,6 +140,19 @@ typedef struct { } embedding; } llm_options; +// What the active llama context was configured for. A context created with +// generate_embedding=1 has embeddings enabled and pooling forced on, which makes it +// unusable for token-by-token generation (no per-token logits are produced, so +// sampling reads a NULL logits buffer). Contexts created for text generation and for +// chat are byte-for-byte identical - llm_context_create_chat() and +// llm_context_create_textgen() both pass the same empty option string - so they share +// a single kind and stay mutually interchangeable. +typedef enum { + LLM_CONTEXT_NONE = 0, + LLM_CONTEXT_GENERATIVE, + LLM_CONTEXT_EMBEDDING +} llm_context_kind; + typedef struct { llama_chat_message *items; size_t count; @@ -156,6 +169,7 @@ typedef struct { struct llama_model *model; struct llama_context *ctx; struct llama_sampler *sampler; + llm_context_kind context_kind; // what ctx was configured for (derived from the parsed options) struct llama_adapter_lora *lora[MAX_LORAS]; float lora_scale[MAX_LORAS]; @@ -500,11 +514,19 @@ static bool llm_context_options_callback (void *ctx, void *xdata, const char *ke // https://github.com/ggml-org/llama.cpp/discussions/15093 int value = (int)strtol(buffer, NULL, 0); options->embeddings = (value != 0); - options->pooling_type = LLAMA_POOLING_TYPE_MEAN; - - // for non-causal models, batch size must be equal to ubatch size - // when generating embeddings, always tie them together. - options->n_ubatch = options->n_batch; + + // Only when actually enabling embeddings. API.md documents this key as "1 or 0", + // and generate_embedding=0 used to force MEAN pooling anyway - which both made + // "embeddings off" configure a pooling context, and, because an explicit + // pooling_type is what tells the classifier the caller asked for pooling, hid + // embedding models from the check in llm_context_create_with_options(). + if (value != 0) { + options->pooling_type = LLAMA_POOLING_TYPE_MEAN; + + // for non-causal models, batch size must be equal to ubatch size + // when generating embeddings, always tie them together. + options->n_ubatch = options->n_batch; + } return true; } @@ -946,6 +968,36 @@ static bool llm_check_context (sqlite3_context *context) { return true; } +static const char *llm_context_kind_name (llm_context_kind kind) { + switch (kind) { + case LLM_CONTEXT_EMBEDDING: return "embedding generation"; + case LLM_CONTEXT_GENERATIVE: return "text generation"; + case LLM_CONTEXT_NONE: break; + } + // callers gate on a live context first, so LLM_CONTEXT_NONE should be unreachable + return "an unknown operation"; +} + +// Rejects an operation that needs per-token logits when the active context cannot +// produce them: generate_embedding=1 turns on pooled embeddings, and llama then returns +// no per-token logits to sample from. +// +// The error sink is passed explicitly rather than read from ai->context/ai->vtab so +// this check does not depend on those fields being current: scalar functions pass +// (context, NULL), the llm_chat() vtab passes (NULL, vtab). +static bool llm_check_generative (ai_context *ai, sqlite3_context *context, sqlite3_vtab *vtab, const char *function_name) { + if (!ai) return false; + + if (ai->context_kind != LLM_CONTEXT_GENERATIVE) { + sqlite_common_set_error(context, vtab, SQLITE_MISUSE, + "%s requires a text generation context, but the current context was created for %s", + function_name, llm_context_kind_name(ai->context_kind)); + return false; + } + + return true; +} + // ai->context / ai->vtab are the error sink for the whole chat and sampler subsystem: // roughly twenty sites report through sqlite_common_set_error(ai->context, ai->vtab, ...), // including llm_chat_run(), llm_chat_generate_response(), llm_chat_tokenize_input(), @@ -1039,6 +1091,7 @@ void llm_messages_free (ai_messages *list) { static void ai_context_release (ai_context *ai) { if (ai->ctx) llama_free(ai->ctx); ai->ctx = NULL; + ai->context_kind = LLM_CONTEXT_NONE; ai->chat.prev_len = 0; ai->chat.token_count = 0; @@ -1549,7 +1602,14 @@ static void llm_embed_generate_run (sqlite3_context *context, const char *text, static void llm_embed_generate (sqlite3_context *context, int argc, sqlite3_value **argv) { if (llm_check_context(context) == false) return; if (llm_common_args_check(context, "llm_embed_generate", argc, argv, true) == false) return; - + + // Deliberately NOT gated on context_kind. What embedding generation actually needs is + // a context that pools, and llm_embed_generate_run() already checks the *resolved* + // pooling type (llama_pooling_type()) further down. That check is strictly better than + // one based on how the context was declared: embedding models inherit mean pooling + // from the GGUF, so llm_context_create('context_size=512,embedding_type=FLOAT32') on + // e.g. all-MiniLM-L6-v2 produces working embeddings without generate_embedding=1. + const char *text = (const char *)sqlite3_value_text(argv[0]); int32_t text_len = (int32_t)sqlite3_value_bytes(argv[0]); const char *model_options = (argc == 2) ? (const char *)sqlite3_value_text(argv[1]) : NULL; @@ -1791,6 +1851,7 @@ static void llm_text_generate (sqlite3_context *context, int argc, sqlite3_value sqlite_context_result_error(context, SQLITE_ERROR, "No model loaded"); return; } + if (llm_check_generative(ai, context, NULL, "llm_text_generate") == false) return; const char *text = (const char *)sqlite3_value_text(argv[0]); int32_t text_len = (int32_t)sqlite3_value_bytes(argv[0]); @@ -1838,7 +1899,14 @@ static bool llm_chat_check_context (ai_context *ai) { sqlite_common_set_error(ai ? ai->context : NULL, ai ? ai->vtab : NULL, SQLITE_MISUSE, "No context found. Please call llm_context_create() before llm_chat_create()."); return false; } - + + // NOTE: no context-kind check here. Only the entry points that actually decode and + // sample - llm_chat_respond() and the llm_chat() vtab - need per-token logits, and + // they carry the check themselves. llm_chat_create/restore/system_prompt only build + // up in-memory message state, and llm_chat_restore() in particular has already + // dropped the previous chat by the time it gets here, so failing it would destroy + // state on behalf of an operation that would have worked. + // check sampler if (!ai->sampler) { llm_sampler_check(ai); @@ -2168,13 +2236,19 @@ static int llm_chat_cursor_open (sqlite3_vtab *pVtab, sqlite3_vtab_cursor **ppCu ai_context *ai = c->ai; llm_error_sink_vtab(ai, (sqlite3_vtab *)vtab); - // sqlite never calls xClose for a cursor whose xOpen failed, so the cursor has to be - // released here or it leaks for the lifetime of the connection + // Only meaningful once a context exists: with ai->ctx == NULL the kind is + // LLM_CONTEXT_NONE and this would report "created for an unknown operation" where + // llm_chat_check_context() below gives the caller the actionable "No context found". + // sqlite never calls xClose for a cursor whose xOpen failed, so free it on the way out. + if (ai->ctx && llm_check_generative(ai, NULL, (sqlite3_vtab *)vtab, "llm_chat") == false) { + sqlite3_free(c); + return SQLITE_ERROR; + } if (llm_chat_check_context(ai) == false) { sqlite3_free(c); return SQLITE_ERROR; } - + *ppCursor = (sqlite3_vtab_cursor *)c; return SQLITE_OK; } @@ -2452,6 +2526,7 @@ static void llm_chat_respond (sqlite3_context *context, int argc, sqlite3_value sqlite_context_result_error(context, SQLITE_ERROR, "No model loaded"); return; } + if (llm_check_generative(ai, context, NULL, "llm_chat_respond") == false) return; if (llm_chat_check_context(ai) == false) return; const char *user_prompt = (const char *)sqlite3_value_text(argv[0]); @@ -2861,6 +2936,13 @@ static bool llm_context_create_with_options (sqlite3_context *context, ai_contex } } + // Whether the caller asked for pooling at all. llama_context_default_params() leaves + // pooling_type UNSPECIFIED and the parser only writes it for an explicit + // pooling_type=... or generate_embedding=1, so anything else here came from the user. + // llama then resolves UNSPECIFIED from the model's own hparams, which is what lets us + // tell "this model pools by default" from "this caller asked for pooling". + const bool caller_set_pooling = (ctx_params.pooling_type != LLAMA_POOLING_TYPE_UNSPECIFIED); + // sanity check embedding_type if (ctx_params.embeddings && ai->options.embedding.type == 0) { sqlite_context_result_error(context, SQLITE_ERROR, "Embedding type (embedding_type) must be specified in the create context function"); @@ -2891,7 +2973,23 @@ static bool llm_context_create_with_options (sqlite3_context *context, ai_contex if (ai->ctx) llm_context_free(context, 0, NULL); ai->ctx = ctx; - + + // Classify from what was actually configured, not from which wrapper was called: + // llm_context_create('generate_embedding=1,...') is documented as equivalent to + // llm_context_create_embedding() and has to be recognised as one. ctx_params is + // rebuilt from llama_context_default_params() on every call, so it always describes + // this context. + // + // The second clause catches embedding *models*. A BERT-family model carries its own + // pooling type in the GGUF, so llm_context_create('context_size=512,...') on + // all-MiniLM yields a context that pools - and generation on it silently returns '' + // (GH #33) because the first sampled token reads as EOG. Resolved pooling alone is + // not enough to conclude that, though: a caller may force pooling_type=mean on a + // perfectly generative model, and that still generates. Only pooling the caller did + // not ask for tells us the model itself is an embedding model. + const bool model_pools_by_default = (!caller_set_pooling && llama_pooling_type(ctx) != LLAMA_POOLING_TYPE_NONE); + ai->context_kind = (ctx_params.embeddings || model_pools_by_default) ? LLM_CONTEXT_EMBEDDING : LLM_CONTEXT_GENERATIVE; + return true; } diff --git a/src/sqlite-ai.h b/src/sqlite-ai.h index 948be5b..b873b16 100644 --- a/src/sqlite-ai.h +++ b/src/sqlite-ai.h @@ -24,7 +24,7 @@ extern "C" { #endif -#define SQLITE_AI_VERSION "1.0.7" +#define SQLITE_AI_VERSION "1.0.8" SQLITE_AI_API int sqlite3_ai_init (sqlite3 *db, char **pzErrMsg, const sqlite3_api_routines *pApi); diff --git a/tests/c/unittest.c b/tests/c/unittest.c index f983ab2..d136184 100644 --- a/tests/c/unittest.c +++ b/tests/c/unittest.c @@ -15,6 +15,7 @@ typedef struct { const char *extension_path; const char *model_path; + const char *embed_model_path; const char *whisper_model_path; const char *audio_path; bool verbose; @@ -28,7 +29,7 @@ typedef struct { } test_case; static void usage(const char *prog) { - fprintf(stderr, "Usage: %s [--extension /path/to/ai] [--model /path/to/model] [--whisper-model /path/to/whisper] [--audio /path/to/audio.wav] [--verbose]\n", prog); + fprintf(stderr, "Usage: %s [--extension /path/to/ai] [--model /path/to/model] [--embed-model /path/to/embedding-model] [--whisper-model /path/to/whisper] [--audio /path/to/audio.wav] [--verbose]\n", prog); } static int expect_error_contains(const char *err_msg, const char *needle) { @@ -564,6 +565,195 @@ static int test_llm_embed_generate(const test_env *env) { return 1; } +// Regression (GH #33): an embedding context produces no per-token logits, so text +// generation on one used to sample from a NULL logits buffer and silently return ''. +// The entry points that decode and sample must reject it with an explicit error instead - +// including the chat paths, where reading that NULL buffer is a crash, not an empty +// string. The llm_chat() vtab row also covers the cursor it has to free on rejection: +// sqlite never calls xClose for a cursor whose xOpen failed, so a leak there would show +// up in the memory check below. +static int test_context_kind_mismatch(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; + + // an embedding context rejects every operation that decodes and samples + if (exec_expect_ok(env, db, "SELECT llm_context_create_embedding('context_size=512,embedding_type=UINT8');") != 0) goto fail; + if (exec_expect_error(env, db, "SELECT llm_text_generate('Say hi');", + "llm_text_generate requires a text generation context") != 0) goto fail; + if (exec_expect_error(env, db, "SELECT llm_chat_respond('Say hi');", + "llm_chat_respond requires a text generation context") != 0) goto fail; + if (exec_expect_error(env, db, "SELECT reply FROM llm_chat('Say hi');", + "llm_chat requires a text generation context") != 0) goto fail; + // repeat the vtab rejection: a cursor leaked per attempt would fail the memory check + if (exec_expect_error(env, db, "SELECT reply FROM llm_chat('Say hi again');", + "llm_chat requires a text generation context") != 0) goto fail; + + // the same rejection must apply to an embedding context built through the generic + // constructor, which API.md documents as equivalent to llm_context_create_embedding() + if (exec_expect_ok(env, db, "SELECT llm_context_create('generate_embedding=1,normalize_embedding=1,pooling_type=mean,context_size=512,embedding_type=UINT8');") != 0) goto fail; + if (exec_expect_error(env, db, "SELECT llm_text_generate('Say hi');", + "llm_text_generate requires a text generation context") != 0) goto fail; + + // the mirror direction is NOT gated on how the context was declared: what embeddings + // need is pooling, so this is rejected by the resolved-pooling check instead + if (exec_expect_ok(env, db, "SELECT llm_context_create_textgen('context_size=512');") != 0) goto fail; + if (exec_expect_error(env, db, "SELECT llm_embed_generate('Say hi');", + "Embedding generation requires pooling") != 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_v2(db); + return assert_sqlite_memory_clean("context_kind_mismatch", env); + +fail: + if (db) sqlite3_close_v2(db); + return 1; +} + +// An embedding *model* (BERT-family) carries its own pooling type in the GGUF, so a +// context built with no embedding settings at all still pools - and generation on it +// returns '' with no error, which is the GH #33 symptom on the model rather than the +// context. Requires an encoder-style model, so it is skipped without --embed-model. +static int test_context_kind_embedding_model(const test_env *env) { + if (!env->embed_model_path) { + printf(" [SKIP] no --embed-model provided\n"); + return 0; + } + + sqlite3 *db = NULL; + if (open_db_and_load(env, &db) != SQLITE_OK) return 1; + + char sqlbuf[512]; + snprintf(sqlbuf, sizeof(sqlbuf), "SELECT llm_model_load('%s');", env->embed_model_path); + if (exec_expect_ok(env, db, sqlbuf) != 0) goto fail; + + // the exact flow API.md recommends: a plain context, no generate_embedding. + // Embeddings must keep working... + if (exec_expect_ok(env, db, "SELECT llm_context_create('context_size=512,embedding_type=FLOAT32');") != 0) goto fail; + if (exec_expect_ok(env, db, "SELECT llm_embed_generate('hello world');") != 0) goto fail; + // ...and generation on it must say so rather than return '' + if (exec_expect_error(env, db, "SELECT llm_text_generate('Say hi');", + "llm_text_generate requires a text generation context") != 0) goto fail; + if (exec_expect_error(env, db, "SELECT llm_chat_respond('Say hi');", + "llm_chat_respond requires a text generation context") != 0) goto fail; + + // asking for a textgen context on this model does not make it generative either + if (exec_expect_ok(env, db, "SELECT llm_context_create_textgen('context_size=512');") != 0) goto fail; + if (exec_expect_error(env, db, "SELECT llm_text_generate('Say hi');", + "llm_text_generate requires a text generation context") != 0) goto fail; + + // nor does explicitly turning embeddings off: generate_embedding=0 must not be read + // as "the caller asked for pooling", which would hide the model from the check + if (exec_expect_ok(env, db, "SELECT llm_context_create('generate_embedding=0,context_size=512,embedding_type=FLOAT32');") != 0) goto fail; + if (exec_expect_error(env, db, "SELECT llm_text_generate('Say hi');", + "llm_text_generate requires a text generation context") != 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_v2(db); + return assert_sqlite_memory_clean("context_kind_embedding_model", env); + +fail: + if (db) sqlite3_close_v2(db); + return 1; +} + +// With no context at all the vtab must give the actionable "No context found", not the +// context-kind message: LLM_CONTEXT_NONE is not a kind the caller can act on. +static int test_chat_vtab_without_context(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_error(env, db, "SELECT reply FROM llm_chat('hi');", "No context found") != 0) goto fail; + // twice: a cursor leaked per rejected xOpen would fail the memory check below + if (exec_expect_error(env, db, "SELECT reply FROM llm_chat('hi again');", "No context found") != 0) goto fail; + + if (exec_expect_ok(env, db, "SELECT llm_model_free();") != 0) goto fail; + sqlite3_close_v2(db); + return assert_sqlite_memory_clean("chat_vtab_without_context", env); + +fail: + if (db) sqlite3_close_v2(db); + return 1; +} + +// The other half of the contract: the kind must come from the options that were actually +// parsed, not from which llm_context_create_* wrapper was called. A generic +// llm_context_create('generate_embedding=1,...') is an embedding context and must keep +// working with llm_embed_generate(), and chat/textgen contexts are byte-for-byte +// identical so llm_text_generate() has to accept either (see the README vision example, +// which calls llm_context_create_chat() and then llm_text_generate()). +static int test_context_kind_compatible(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]; + char result[4096]; + snprintf(sqlbuf, sizeof(sqlbuf), "SELECT llm_model_load('%s');", model); + if (exec_expect_ok(env, db, sqlbuf) != 0) goto fail; + + // embedding context built the generic way still generates embeddings + if (exec_expect_ok(env, db, "SELECT llm_context_create('generate_embedding=1,normalize_embedding=1,pooling_type=mean,context_size=512,embedding_type=UINT8');") != 0) goto fail; + if (exec_expect_ok(env, db, "SELECT llm_embed_generate('generic embedding context');") != 0) goto fail; + + // ...and so does a context that never passed generate_embedding but still resolves to + // a pooling type. Embedding models (BERT-family) inherit mean pooling from the GGUF, + // so gating llm_embed_generate() on how the context was declared would break them; + // here pooling_type is set explicitly to get the same shape from the test model. + if (exec_expect_ok(env, db, "SELECT llm_context_create('context_size=512,pooling_type=mean,embedding_type=UINT8');") != 0) goto fail; + if (exec_expect_ok(env, db, "SELECT llm_embed_generate('pooling without generate_embedding');") != 0) goto fail; + + // a chat context is a valid text generation context + if (exec_expect_ok(env, db, "SELECT llm_context_create_chat('context_size=1024,n_predict=32');") != 0) goto fail; + memset(result, 0, sizeof(result)); + 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, "[context_kind_compatible] llm_text_generate on a chat context returned empty\n"); + goto fail; + } + + // so is one from the generic constructor with no embedding options + if (exec_expect_ok(env, db, "SELECT llm_context_create('context_size=1024,n_predict=32');") != 0) goto fail; + memset(result, 0, sizeof(result)); + 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, "[context_kind_compatible] llm_text_generate on a generic context returned empty\n"); + goto fail; + } + + // chat calls that only build in-memory state must not be gated on the context kind: + // llm_chat_restore() in particular drops the previous chat before it could be checked. + // Kept last: llm_chat_create() installs a default dist sampler on ai->sampler that + // outlives the context, which would make any generation after it non-deterministic. + if (exec_expect_ok(env, db, "SELECT llm_context_create_embedding('context_size=512,embedding_type=UINT8');") != 0) goto fail; + if (exec_expect_ok(env, db, "SELECT llm_chat_create();") != 0) goto fail; + if (exec_expect_ok(env, db, "SELECT llm_chat_system_prompt('You are helpful.');") != 0) goto fail; + if (exec_expect_ok(env, db, "SELECT llm_chat_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_v2(db); + return assert_sqlite_memory_clean("context_kind_compatible", env); + +fail: + if (db) sqlite3_close_v2(db); + return 1; +} + static int test_llm_embed_generate_basic(const test_env *env) { sqlite3 *db = NULL; sqlite3_stmt *stmt = NULL; @@ -2625,6 +2815,10 @@ static const test_case TESTS[] = { {"llm_chat_vtab", test_llm_chat_vtab}, {"chat_error_sink_after_statement", test_chat_error_sink_after_statement}, {"test_llm_embed_generate", test_llm_embed_generate}, + {"context_kind_mismatch", test_context_kind_mismatch}, + {"context_kind_compatible", test_context_kind_compatible}, + {"context_kind_embedding_model", test_context_kind_embedding_model}, + {"chat_vtab_without_context", test_chat_vtab_without_context}, {"llm_embed_generate_basic", test_llm_embed_generate_basic}, {"llm_embedding_then_chat", test_llm_embedding_then_chat}, {"llm_context_size_errors", test_llm_context_size_errors}, @@ -2698,6 +2892,12 @@ int main(int argc, char **argv) { return EXIT_FAILURE; } env.model_path = argv[i]; + } else if (strcmp(argv[i], "--embed-model") == 0) { + if (++i >= argc) { + usage(argv[0]); + return EXIT_FAILURE; + } + env.embed_model_path = argv[i]; } else if (strcmp(argv[i], "--whisper-model") == 0) { if (++i >= argc) { usage(argv[0]);