diff --git a/API.md b/API.md index 73a5d40..5db77bd 100644 --- a/API.md +++ b/API.md @@ -5,6 +5,66 @@ These functions enable loading and interacting with LLMs, configuring samplers, --- +## Lifecycle: what each `*_free()` releases + +A connection holds four things — a model, an inference context (which owns the KV +cache), a sampler, and a chat (the message history). They are released +independently, and releasing one does not necessarily release the others. + +| you call | model | context / KV cache | sampler | chat history | +| ---------------------------- | ------------ | ------------------ | ---------- | ------------ | +| `llm_model_load()` (again) | replaced | freed | **freed** | **kept** | +| `llm_model_free()` | freed | freed | **freed** | **kept** | +| `llm_context_free()` | kept | freed | kept | **kept** | +| `llm_chat_free()` | kept | **cache cleared** | kept | freed | +| closing the connection | freed | freed | freed | freed | + +Two consequences are worth knowing: + +**A conversation survives a model switch.** The history is not tied to a model, so +after loading a new one the next `llm_chat_respond()` replays the whole transcript +through the new model's chat template before answering. That first turn therefore +costs one full prompt evaluation, and it fails if the transcript no longer fits the +new context. To start clean instead, call `llm_chat_free()` first: + +```sql +SELECT llm_chat_save('before switching model'); -- optional: keep a copy +SELECT llm_chat_free(); -- drop history and clear the cache +SELECT llm_model_load('/path/to/other-model.gguf'); +SELECT llm_context_create_chat('context_size=4096'); +SELECT llm_sampler_create(); -- the model load released the old one +SELECT llm_sampler_init_temp(0.8); +SELECT llm_sampler_init_dist(42); +SELECT llm_chat_respond('Hello'); +``` + +**A sampler never survives a model load or free.** A sampler built with +`llm_sampler_init_grammar()` holds a reference to the model's vocabulary, so it +cannot outlive it. Rebuild the chain after loading a model; if you never configured +one, a default is created for you. + +### Recovering from a full context + +When a chat fills its context, `llm_chat_respond()` returns +`Context size exceeded`. The turn is not lost — whatever was generated is kept and +the conversation stays usable — but the context is full, so every later turn returns +the same error until you make room. Any of these works: + +- **Start over** — `llm_chat_free()` then carry on; the cache is cleared with it. +- **Enlarge the context** — `llm_context_free()` then + `llm_context_create_chat('context_size=...')` with a bigger window. The history is + kept and replayed into it. +- **Compact** — the history is an ordinary table, so trim it with SQL and restore: + + ```sql + SELECT llm_chat_save('before compaction'); + DELETE FROM ai_chat_messages + WHERE id NOT IN (SELECT id FROM ai_chat_messages ORDER BY id DESC LIMIT 8); + SELECT llm_chat_restore((SELECT uuid FROM ai_chat_history ORDER BY id DESC LIMIT 1)); + ``` + +--- + ## `ai_version()` **Returns:** `TEXT` @@ -44,6 +104,11 @@ SELECT ai_log_info(1); Loads a GGUF model from the specified file path with optional comma separated key=value configuration. If no options are provided the following default value is used: `gpu_layers=99` +Loading a model replaces whatever was loaded before: the previous model, its +inference context and the sampler are all released. The chat history is **kept**, so +an ongoing conversation carries over to the new model — call `llm_chat_free()` first +if you want it dropped. See [Lifecycle](#lifecycle-what-each-_free-releases). + The following keys are available: ``` gpu_layers=N (N is the number of layers to store in VRAM) @@ -70,6 +135,9 @@ SELECT llm_model_load('./models/llama.gguf', 'gpu_layers=99'); **Description:** Unloads the current model and frees associated memory. +Also releases the inference context and the sampler, since both are tied to the model. +The chat history is **kept** — load another model and the conversation continues. +See [Lifecycle](#lifecycle-what-each-_free-releases). **Example:** @@ -102,7 +170,7 @@ The following keys are available in context_settings: | `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. | | `n_predict` | `number` | Control the maximum number of tokens generated during text generation. | -| `embedding_type` | `FLOAT32, FLOAT16, BFLOAT16, UINT8, INT8` | Set the model native type, mandatory during embedding generation. | +| `embedding_type` | `FLOAT32, FLOAT16, FLOATB16, UINT8, INT8` | Set the model native type. **Required** for embedding contexts — omitting it, or passing an unrecognised name, fails with *"Embedding type (embedding_type) must be specified in the create context function"*. Note the spelling `FLOATB16`, not `BFLOAT16`. | ### Core sizing & threading @@ -243,7 +311,10 @@ SELECT llm_context_create_textgen(); **Returns:** `NULL` **Description:** -Frees the current inference context. +Frees the current inference context and its KV cache. +The chat history is **kept**: the next `llm_chat_respond()` replays it into whatever +context you create next, which is how a context is resized mid-conversation. +See [Lifecycle](#lifecycle-what-each-_free-releases). **Example:** @@ -305,6 +376,9 @@ SELECT llm_sampler_create(); **Description:** Frees resources associated with the current sampler. +Note that `llm_model_load()`, `llm_model_free()` and `llm_sampler_create()` already +release the previous sampler, so calling this is only necessary to drop a chain +without replacing it. **Example:** @@ -604,12 +678,21 @@ SELECT llm_token_count('Hello world!'); **Description:** Generates a text embedding as a BLOB vector, with optional configuration provided as a comma-separated list of key=value pairs. By default, the embedding is normalized unless `normalize_embedding=0` is specified. -If `json_output=1` is set, the function returns a JSON object instead of a BLOB. +If `json_output=1` is set, the function returns the vector as a JSON **array** of +numbers instead of a BLOB. + +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. **Example:** ```sql +SELECT llm_embed_generate('hello world'); +-- BLOB, ready to store in a sqlite-vector column + SELECT llm_embed_generate('hello world', 'json_output=1'); +-- '[-0.0355376,0.0334288,...]' ``` --- @@ -661,9 +744,11 @@ SELECT reply FROM llm_chat('Tell me a joke.'); **Returns:** `TEXT` **Description:** -Starts a new in-memory chat session. +Starts a new in-memory chat session, discarding any chat already in progress and +clearing it out of the KV cache. Returns unique chat UUIDv7 value. -If no chat is explicitly created, one will be created automatically when needed. +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()`. **Example:** @@ -678,7 +763,12 @@ SELECT llm_chat_create(); **Returns:** `NULL` **Description:** -Ends the current chat session. +Ends the current chat session: discards the in-memory history and clears the +conversation out of the KV cache. Anything already written by `llm_chat_save()` +survives in `ai_chat_history` / `ai_chat_messages` and can be brought back with +`llm_chat_restore()`. +Use this before `llm_model_load()` when you want the new model to start clean. +See [Lifecycle](#lifecycle-what-each-_free-releases). **Example:** @@ -705,10 +795,15 @@ SELECT llm_chat_save('Support Chat', '{"user": "Marco"}'); ## `llm_chat_restore(uuid TEXT)` -**Returns:** `NULL` +**Returns:** `INTEGER` **Description:** -Restores a previously saved chat session by UUID. +Restores a previously saved chat session by UUID, replacing whatever chat is current. +Returns the number of messages restored. +The KV cache is cleared, so the restored transcript is replayed on the next +`llm_chat_respond()`. Combined with a `DELETE` against `ai_chat_messages`, this is how +a long conversation is compacted — see +[Recovering from a full context](#recovering-from-a-full-context). **Example:** diff --git a/src/sqlite-ai.c b/src/sqlite-ai.c index 97d30b8..2d1adbd 100644 --- a/src/sqlite-ai.c +++ b/src/sqlite-ai.c @@ -835,6 +835,9 @@ void *ai_create (sqlite3 *db) { return ai; } +static void ai_context_release (ai_context *ai); +static void ai_chat_release (ai_context *ai); + static void ai_free (void *ctx, bool free_ai, bool free_llm, bool free_audio) { if (!ctx) return; ai_context *ai = (ai_context *)ctx; @@ -855,12 +858,19 @@ static void ai_free (void *ctx, bool free_ai, bool free_llm, bool free_audio) { 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); + + // A conversation outlives a model swap on purpose: prev_len is reset with the + // context below, so the next turn replays the whole transcript through the new + // model's template. It only becomes unreachable when the connection goes away, + // and nothing else frees it - llm_chat_free() is a SQL function the caller may + // never invoke. Must run before the context: clearing the KV cache needs it. + if (free_ai) ai_chat_release(ai); + ai_context_release(ai); + if (ai->model) llama_model_free(ai->model); llm_options_init(&ai->options); ai->model = NULL; - ai->ctx = NULL; ai->sampler = NULL; } @@ -1000,6 +1010,49 @@ void llm_messages_free (ai_messages *list) { list->capacity = 0; } +// Releases the llama context and everything that describes its contents. prev_len is +// how much of the rendered transcript is already in the KV cache, so it cannot outlive +// that cache; the message history is deliberately left alone - that belongs to the chat, +// and keeping it is what lets the next turn replay the conversation into a new context. +static void ai_context_release (ai_context *ai) { + if (ai->ctx) llama_free(ai->ctx); + ai->ctx = NULL; + + ai->chat.prev_len = 0; + ai->chat.token_count = 0; + memset(&ai->chat.batch, 0, sizeof(ai->chat.batch)); +} + +// Releases the conversation: its history, its buffers, and its tokens in the KV cache. +// A conversation lives in both places, so dropping one without the other leaves prev_len +// disagreeing with the cache. +static void ai_chat_release (ai_context *ai) { + if (ai->ctx) { + llama_memory_t memory = llama_get_memory(ai->ctx); + if (memory) llama_memory_clear(memory, true); + } + + memset(ai->chat.uuid, 0, UUID_STR_MAXLEN); + + buffer_destroy(&ai->chat.response); + buffer_destroy(&ai->chat.formatted); + llm_messages_free(&ai->chat.messages); + + if (ai->chat.tokens) sqlite3_free(ai->chat.tokens); + ai->chat.tokens = NULL; + ai->chat.ntokens = 0; + + if (ai->chat.prompt) sqlite3_free(ai->chat.prompt); + ai->chat.prompt = NULL; + + ai->chat.prev_len = 0; + ai->chat.template = NULL; + ai->chat.vocab = NULL; + ai->chat.token_count = 0; + // batch.token pointed into the tokens buffer just freed + memset(&ai->chat.batch, 0, sizeof(ai->chat.batch)); +} + // MARK: - Text Embedding and Normalization - static inline float llm_common_f32_sum (const float *src, int dim) { @@ -1833,10 +1886,24 @@ static bool llm_chat_generate_response (ai_context *ai, ai_cursor *c, bool *is_e return false; } - int32_t drc = llama_decode(ctx, batch); - if (drc != 0) { - sqlite_common_set_error (ai->context, ai->vtab, SQLITE_ERROR, "Failed to decode prompt batch (%d: %s)", drc, llm_decode_error_string(drc)); - return false; + // llama_decode() asserts n_tokens <= n_batch and aborts the process if it is not, + // so a prompt longer than n_batch has to be fed in chunks - llm_text_run() already + // does this. Only the last chunk's logits are sampled from, so the earlier ones can + // go through untouched. n_batch only tracks n_ctx when the caller passes + // context_size, so a big n_ctx with the default 2048 n_batch reaches here easily, + // and a replayed transcript is a single batch of whatever length the chat has. + const int32_t n_batch_max = (int32_t)llama_n_batch(ctx); + for (int32_t decoded = 0; decoded < batch.n_tokens; ) { + int32_t chunk = batch.n_tokens - decoded; + if (chunk > n_batch_max) chunk = n_batch_max; + + struct llama_batch sub = llama_batch_get_one(batch.token + decoded, chunk); + int32_t drc = llama_decode(ctx, sub); + if (drc != 0) { + sqlite_common_set_error (ai->context, ai->vtab, SQLITE_ERROR, "Failed to decode prompt batch (%d: %s)", drc, llm_decode_error_string(drc)); + return false; + } + decoded += chunk; } // sample next token @@ -2170,26 +2237,7 @@ static sqlite3_module llm_chat = { // MARK: - static void llm_chat_free (sqlite3_context *context, int argc, sqlite3_value **argv) { - ai_context *ai = (ai_context *)sqlite3_user_data(context); - - // reset UUID and cleanup chat related memory - memset(ai->chat.uuid, 0, UUID_STR_MAXLEN); - - buffer_destroy(&ai->chat.response); - buffer_destroy(&ai->chat.formatted); - llm_messages_free(&ai->chat.messages); - - if (ai->chat.tokens) sqlite3_free(ai->chat.tokens); - ai->chat.tokens = NULL; - ai->chat.ntokens = 0; - - if (ai->chat.prompt) sqlite3_free(ai->chat.prompt); - ai->chat.prompt = NULL; - ai->chat.prev_len = 0; - - ai->chat.template = NULL; - ai->chat.vocab = NULL; - ai->chat.token_count = 0; + ai_chat_release((ai_context *)sqlite3_user_data(context)); } static void llm_chat_create (sqlite3_context *context, int argc, sqlite3_value **argv) { @@ -2731,9 +2779,7 @@ static void llm_sampler_create (sqlite3_context *context, int argc, sqlite3_valu } static void llm_context_free (sqlite3_context *context, int argc, sqlite3_value **argv) { - ai_context *ai = (ai_context *)sqlite3_user_data(context); - if (ai->ctx) llama_free(ai->ctx); - ai->ctx = NULL; + ai_context_release((ai_context *)sqlite3_user_data(context)); } static bool llm_context_create_with_options (sqlite3_context *context, ai_context *ai, const char *options1, const char *options2) { @@ -2760,7 +2806,7 @@ static bool llm_context_create_with_options (sqlite3_context *context, ai_contex // 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 funtion"); + sqlite_context_result_error(context, SQLITE_ERROR, "Embedding type (embedding_type) must be specified in the create context function"); return false; } diff --git a/src/sqlite-ai.h b/src/sqlite-ai.h index 713d066..948be5b 100644 --- a/src/sqlite-ai.h +++ b/src/sqlite-ai.h @@ -24,7 +24,7 @@ extern "C" { #endif -#define SQLITE_AI_VERSION "1.0.6" +#define SQLITE_AI_VERSION "1.0.7" 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 a2362b5..bea5523 100644 --- a/tests/c/unittest.c +++ b/tests/c/unittest.c @@ -169,6 +169,23 @@ static int exec_capture_error(const test_env *env, sqlite3 *db, const char *sql, return (rc == SQLITE_OK) ? 0 : 1; } +// talks until the chat runs out of context; 0 once it has failed the way we expect +static int fill_chat_context(const test_env *env, sqlite3 *db, const char *label) { + const char *turn = "SELECT llm_chat_respond('Tell me about computers.');"; + for (int i = 0; i < 40; ++i) { + char err[256] = {0}; + if (exec_capture_error(env, db, turn, err, sizeof(err)) == 0) continue; + if (strstr(err, "Context size exceeded") == NULL) { + fprintf(stderr, "[%s] expected the context guard to stop us, got: %s\n", label, err); + return 1; + } + if (env->verbose) printf("[%s] context full after %d turns: %s\n", label, i, err); + return 0; + } + fprintf(stderr, "[%s] context never filled\n", label); + return 1; +} + static int exec_select_rows(const test_env *env, sqlite3 *db, const char *sql, int *rows_out) { if (env->verbose) { printf("[SQL] %s\n", sql); @@ -1524,6 +1541,306 @@ static int test_chat_context_full_is_recoverable(const test_env *env) { return 1; } +// Path A: llm_chat_free() must drop the conversation from the KV cache too, not just +// from ai->chat. It used to leave the cache full, so "start a new chat" - the obvious +// way to recover from a full context - failed with the very same error. +static int test_chat_free_clears_kv_cache(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; + if (exec_expect_ok(env, db, "SELECT llm_context_create_chat('context_size=256');") != 0) goto fail; + if (exec_expect_ok(env, db, "SELECT llm_chat_create();") != 0) goto fail; + + if (fill_chat_context(env, db, "chat_free_clears_kv_cache") != 0) goto fail; + + // a brand new chat must actually start from nothing + if (exec_expect_ok(env, db, "SELECT llm_chat_free();") != 0) goto fail; + if (exec_expect_ok(env, db, "SELECT llm_chat_create();") != 0) goto fail; + + int used = -1; + if (select_single_int(env, db, "SELECT llm_context_used();", &used) != 0) goto fail; + if (used > 0) { + fprintf(stderr, "[chat_free_clears_kv_cache] cache still holds %d tokens after llm_chat_free()\n", used); + goto fail; + } + if (exec_expect_ok(env, db, "SELECT llm_chat_respond('Hello');") != 0) goto fail; + + if (exec_expect_ok(env, db, "SELECT llm_chat_free();") != 0) goto fail; + 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_v2(db); + return assert_sqlite_memory_clean("chat_free_clears_kv_cache", env); + +fail: + if (db) sqlite3_close_v2(db); + return 1; +} + +// Path B: llm_context_free() must reset prev_len, which tracks how much of the transcript +// is in the KV cache. It used to survive the context it described, so the next turn sent +// only the newest message into an empty cache and the model silently lost the +// conversation while ai->chat.messages still claimed it. Reset, the transcript is +// re-primed - which is what makes resizing a context mid-conversation work. +static int test_context_free_replays_chat(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; + if (exec_expect_ok(env, db, "SELECT llm_context_create_chat('context_size=1024');") != 0) goto fail; + if (exec_expect_ok(env, db, "SELECT llm_chat_create();") != 0) goto fail; + + for (int i = 0; i < 3; ++i) { + if (exec_expect_ok(env, db, "SELECT llm_chat_respond('Tell me about computers.');") != 0) goto fail; + } + int used_before = 0; + if (select_single_int(env, db, "SELECT llm_context_used();", &used_before) != 0) goto fail; + if (used_before <= 0) { + fprintf(stderr, "[context_free_replays_chat] expected a populated cache, got %d\n", used_before); + goto fail; + } + + // resize the context mid-conversation + if (exec_expect_ok(env, db, "SELECT llm_context_free();") != 0) goto fail; + if (exec_expect_ok(env, db, "SELECT llm_context_create_chat('context_size=2048');") != 0) goto fail; + + int used_fresh = -1; + if (select_single_int(env, db, "SELECT llm_context_used();", &used_fresh) != 0) goto fail; + if (used_fresh > 0) { + fprintf(stderr, "[context_free_replays_chat] new context should start empty, holds %d\n", used_fresh); + goto fail; + } + + if (exec_expect_ok(env, db, "SELECT llm_chat_respond('And what else?');") != 0) goto fail; + int used_after = 0; + if (select_single_int(env, db, "SELECT llm_context_used();", &used_after) != 0) goto fail; + if (env->verbose) printf("[context_free_replays_chat] used %d -> 0 -> %d\n", used_before, used_after); + + // the whole transcript must be back, not just the newest message + if (used_after <= used_before) { + fprintf(stderr, "[context_free_replays_chat] transcript was not replayed: %d tokens before the " + "resize, only %d after the next turn\n", used_before, used_after); + goto fail; + } + + if (exec_expect_ok(env, db, "SELECT llm_chat_free();") != 0) goto fail; + 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_v2(db); + return assert_sqlite_memory_clean("context_free_replays_chat", env); + +fail: + if (db) sqlite3_close_v2(db); + return 1; +} + +// Path C: compaction. Save the conversation, drop the oldest turns with plain SQL, and +// restore - llm_chat_restore() goes through llm_chat_free(), so the cache is cleared and +// the trimmed transcript is replayed. No context juggling required, which is what this +// used to need. +static int test_chat_restore_compacts_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 (install_seeded_chat_sampler(env, db) != 0) goto fail; + if (exec_expect_ok(env, db, "SELECT llm_context_create_chat('context_size=256');") != 0) goto fail; + if (exec_expect_ok(env, db, "SELECT llm_chat_create();") != 0) goto fail; + + if (fill_chat_context(env, db, "chat_restore_compacts_context") != 0) goto fail; + + if (exec_expect_ok(env, db, "SELECT llm_chat_save('compaction');") != 0) goto fail; + int before = 0; + if (select_single_int(env, db, "SELECT count(*) FROM ai_chat_messages;", &before) != 0) goto fail; + + // keep only the last two exchanges + if (exec_expect_ok(env, db, "DELETE FROM ai_chat_messages WHERE id NOT IN " + "(SELECT id FROM ai_chat_messages ORDER BY id DESC LIMIT 4);") != 0) goto fail; + int after = 0; + if (select_single_int(env, db, "SELECT count(*) FROM ai_chat_messages;", &after) != 0) goto fail; + if (before <= after || after != 4) { + fprintf(stderr, "[chat_restore_compacts_context] prune did not take: %d -> %d\n", before, after); + goto fail; + } + + // note: no llm_context_free() here - restore alone must be enough now + if (exec_expect_ok(env, db, "SELECT llm_chat_restore((SELECT uuid FROM ai_chat_history " + "ORDER BY id DESC LIMIT 1));") != 0) goto fail; + if (exec_expect_ok(env, db, "SELECT llm_chat_respond('Carry on');") != 0) goto fail; + if (env->verbose) { + int used = 0; + if (select_single_int(env, db, "SELECT llm_context_used();", &used) == 0) + printf("[chat_restore_compacts_context] %d messages -> %d, cache now %d tokens\n", before, after, used); + } + + if (exec_expect_ok(env, db, "SELECT llm_chat_free();") != 0) goto fail; + 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_v2(db); + return assert_sqlite_memory_clean("chat_restore_compacts_context", env); + +fail: + if (db) sqlite3_close_v2(db); + return 1; +} + +// A model swap frees the context through ai_free(), which used to bypass +// llm_context_free() entirely - so prev_len survived the cache it described and the +// next turn sent only its tail into an empty one. The conversation is deliberately kept +// across the swap, so it must be replayed through the new model instead. +static int test_chat_survives_model_reload(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; + if (exec_expect_ok(env, db, "SELECT llm_context_create_chat('context_size=2048');") != 0) goto fail; + if (exec_expect_ok(env, db, "SELECT llm_chat_create();") != 0) goto fail; + + for (int i = 0; i < 2; ++i) { + if (exec_expect_ok(env, db, "SELECT llm_chat_respond('Tell me about computers.');") != 0) goto fail; + } + int used_before = 0; + if (select_single_int(env, db, "SELECT llm_context_used();", &used_before) != 0) goto fail; + if (used_before <= 0) { + fprintf(stderr, "[chat_survives_model_reload] expected a populated cache, got %d\n", used_before); + goto fail; + } + + // note: no llm_context_free() - ai_free() takes the context down from under us, + // which is exactly the path that used to leave prev_len behind + if (exec_expect_ok(env, db, "SELECT llm_model_free();") != 0) goto fail; + if (exec_expect_ok(env, db, sqlbuf) != 0) goto fail; + if (install_seeded_chat_sampler(env, db) != 0) goto fail; + if (exec_expect_ok(env, db, "SELECT llm_context_create_chat('context_size=2048');") != 0) goto fail; + + int used_fresh = -1; + if (select_single_int(env, db, "SELECT llm_context_used();", &used_fresh) != 0) goto fail; + if (used_fresh > 0) { + fprintf(stderr, "[chat_survives_model_reload] reloaded context should start empty, holds %d\n", used_fresh); + goto fail; + } + + if (exec_expect_ok(env, db, "SELECT llm_chat_respond('And what else?');") != 0) goto fail; + int used_after = 0; + if (select_single_int(env, db, "SELECT llm_context_used();", &used_after) != 0) goto fail; + if (env->verbose) printf("[chat_survives_model_reload] used %d -> 0 -> %d\n", used_before, used_after); + if (used_after <= used_before) { + fprintf(stderr, "[chat_survives_model_reload] transcript was not replayed: %d tokens before the " + "reload, only %d after the next turn\n", used_before, used_after); + goto fail; + } + + if (exec_expect_ok(env, db, "SELECT llm_chat_free();") != 0) goto fail; + 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_v2(db); + return assert_sqlite_memory_clean("chat_survives_model_reload", env); + +fail: + if (db) sqlite3_close_v2(db); + return 1; +} + +// llm_chat_free() is a SQL function a caller may simply never invoke, and it used to be +// the only thing that released the history, the buffers, the prompt and the token array. +// Closing the connection leaked all of it - 29920 bytes for this conversation. Every +// other test calls llm_chat_free() explicitly, which is why nothing caught it. +static int test_chat_released_on_connection_close(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; + if (exec_expect_ok(env, db, "SELECT llm_context_create_chat('context_size=1024');") != 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_respond('Hello');") != 0) goto fail; + + // deliberately no llm_chat_free(), no llm_context_free(), no llm_model_free(): + // closing the connection has to reclaim everything on its own + sqlite3_close_v2(db); + return assert_sqlite_memory_clean("chat_released_on_connection_close", env); + +fail: + if (db) sqlite3_close_v2(db); + return 1; +} + +// llama_decode() asserts n_tokens <= n_batch and aborts the process, and the chat path +// submitted the whole prompt as one batch. n_batch only follows n_ctx when the caller +// passes context_size, so n_ctx=4096 leaves n_batch at llama's default of 2048 and a +// long enough prompt kills the process. Replaying a transcript into a fresh context +// makes the same batch, which is why this matters more since prev_len is reset. +static int test_chat_prompt_larger_than_n_batch(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; + // n_ctx only: n_batch stays at 2048 + if (exec_expect_ok(env, db, "SELECT llm_context_create_chat('n_ctx=4096');") != 0) goto fail; + + int n_ctx = 0; + if (select_single_int(env, db, "SELECT llm_context_size();", &n_ctx) != 0) goto fail; + if (n_ctx != 4096) { + fprintf(stderr, "[chat_prompt_larger_than_n_batch] wanted n_ctx 4096, got %d\n", n_ctx); + goto fail; + } + if (exec_expect_ok(env, db, "SELECT llm_chat_create();") != 0) goto fail; + + // ~3000 tokens of prompt against a 2048 batch: this used to abort the process + if (exec_expect_ok(env, db, "SELECT llm_chat_respond(" + "replace(hex(zeroblob(1500)), '0', 'hello '));") != 0) goto fail; + int used = 0; + if (select_single_int(env, db, "SELECT llm_context_used();", &used) != 0) goto fail; + if (env->verbose) printf("[chat_prompt_larger_than_n_batch] cache holds %d tokens\n", used); + if (used <= 2048) { + fprintf(stderr, "[chat_prompt_larger_than_n_batch] expected more than one batch to be " + "decoded, cache holds only %d tokens\n", used); + goto fail; + } + + if (exec_expect_ok(env, db, "SELECT llm_chat_free();") != 0) goto fail; + 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_v2(db); + return assert_sqlite_memory_clean("chat_prompt_larger_than_n_batch", env); + +fail: + if (db) sqlite3_close_v2(db); + return 1; +} + // Test chat vtab streaming produces tokens and saves response correctly across turns static int test_chat_vtab_multi_turn(const test_env *env) { sqlite3 *db = NULL; @@ -2290,6 +2607,12 @@ static const test_case TESTS[] = { {"chat_create_free_cycle", test_chat_create_free_cycle}, {"chat_recreate_after_conversation", test_chat_recreate_after_conversation}, {"chat_context_full_is_recoverable", test_chat_context_full_is_recoverable}, + {"chat_free_clears_kv_cache", test_chat_free_clears_kv_cache}, + {"context_free_replays_chat", test_context_free_replays_chat}, + {"chat_restore_compacts_context", test_chat_restore_compacts_context}, + {"chat_survives_model_reload", test_chat_survives_model_reload}, + {"chat_released_on_connection_close", test_chat_released_on_connection_close}, + {"chat_prompt_larger_than_n_batch", test_chat_prompt_larger_than_n_batch}, {"chat_vtab_multi_turn", test_chat_vtab_multi_turn}, {"chat_save_restore_roundtrip", test_chat_save_restore_roundtrip}, {"chat_system_prompt_clear", test_chat_system_prompt_clear},