# Spec — Ajouter le support embeddings (RAG) à kazeia-engine

> **Destinataire :** développeur de `kazeia-engine` (`/opt/Kazeia-engine`).
> **But :** exposer un entrypoint « texte → vecteur de phrase poolé » pour le RAG de Kazeia mobile.
> **Fichier cible :** `dist/jni/kazeia_engine_jni.cpp` (+ bindings Kotlin côté app).
> **Coût estimé :** ~40–60 lignes C++, aucune nouvelle dépendance.

## Contexte

Le RAG de Kazeia mobile doit transformer un texte en **vecteur de phrase poolé**. La machinerie
existe déjà dans le fork (`llama_pooling_type`, champ `pooling_type` dans les context params,
`llama_get_embeddings_seq`), mais **aucun entrypoint JNI texte→vecteur** n'est exposé. Il faut
l'ajouter **sans toucher au chemin Speaker/TTS existant**.

## Principe (important)

- L'embedder est un **modèle dédié** (GGUF type **e5 / bge**, arch BERT), chargé comme **handle
  séparé** — surtout **pas** le LLM Speaker (un décodeur causal est un mauvais embedder).
- **CPU pur, pas de HTP/HMX** (un embed coûte ~10–50 ms ; le NPU est inutile et déstabilisant ici).
- **Ne pas réutiliser** `prefillEmbeds`/`decodeEmbed`/`llama_get_embeddings_ith` : c'est l'I/O float
  du Talker TTS (hidden du *dernier token*), pas un embedding de phrase poolé.

---

## Tâche 1 — Nouvelle struct + `loadEmbedder`

À ajouter dans `kazeia_engine_jni.cpp` (section nouvelle, après le bloc LLM) :

```cpp
// ============================================================================
// EMBEDDINGS (RAG) : modèle dédié BERT-like (e5/bge), CPU, pooling de séquence.
// Indépendant du KEngine Speaker. Aucun HTP.
// ============================================================================
struct KEmbedder {
    llama_model*   m;
    llama_context* c;
    const llama_vocab* v;
    int n_embd;
};

// pooling: -1 = défaut du modèle (recommandé) ; 1 = MEAN (e5) ; 2 = CLS (bge)
extern "C" JNIEXPORT jlong JNICALL
Java_com_kazeia_llm_EngineJni_loadEmbedder(JNIEnv* e, jobject, jstring path, jint nThreads, jint pooling) {
    const char* pc = e->GetStringUTFChars(path, 0);
    std::string p(pc); e->ReleaseStringUTFChars(path, pc);

    llama_backend_init();                                   // idempotent
    auto mp = llama_model_default_params(); mp.n_gpu_layers = 0;   // CPU only
    llama_model* m = llama_model_load_from_file(p.c_str(), mp);
    if (!m) return 0;

    auto cp = llama_context_default_params();
    cp.n_threads = (nThreads > 0) ? nThreads : 4;
    cp.embeddings   = true;                                 // <-- clé
    cp.pooling_type = (pooling == 1) ? LLAMA_POOLING_TYPE_MEAN
                    : (pooling == 2) ? LLAMA_POOLING_TYPE_CLS
                    : LLAMA_POOLING_TYPE_UNSPECIFIED;        // défaut = métadonnées du modèle
    // Pooling MEAN/CLS exige que TOUTE la séquence tienne dans un seul ubatch :
    cp.n_ctx = 512; cp.n_batch = 512; cp.n_ubatch = 512;     // chunks <= 512 tokens
    // (ne PAS réutiliser make_ctx : pas de flash_attn ni de KV f16 ici)
    llama_context* c = llama_init_from_model(m, cp);
    if (!c) { llama_model_free(m); return 0; }

    auto* k = new KEmbedder{ m, c, llama_model_get_vocab(m), llama_model_n_embd(m) };
    fprintf(stderr, "kazeia-engine: EMBEDDER chargé (%s, n_embd=%d, pooling=%d)\n", p.c_str(), k->n_embd, pooling);
    return (jlong) k;
}
```

## Tâche 2 — `embedText` (texte → vecteur L2-normalisé)

```cpp
extern "C" JNIEXPORT jfloatArray JNICALL
Java_com_kazeia_llm_EngineJni_embedText(JNIEnv* e, jobject, jlong h, jstring text) {
    auto* k = (KEmbedder*) h;
    const char* tc = e->GetStringUTFChars(text, 0);
    std::string s(tc); e->ReleaseStringUTFChars(text, tc);

    // 1) tokenize (add_special=true -> CLS/BOS selon le modèle)
    int n = -llama_tokenize(k->v, s.c_str(), s.size(), nullptr, 0, true, true);
    if (n <= 0) return nullptr;
    if (n > 512) n = 512;                                    // garde-fou ubatch
    std::vector<llama_token> toks(n);
    llama_tokenize(k->v, s.c_str(), s.size(), toks.data(), n, true, true);

    // 2) batch : tous les tokens, seq 0, sortie activée (requis pour le pooling)
    llama_memory_clear(llama_get_memory(k->c), true);
    llama_batch b = llama_batch_init(n, 0, 1);
    for (int i = 0; i < n; ++i) {
        b.token[i] = toks[i]; b.pos[i] = i;
        b.n_seq_id[i] = 1; b.seq_id[i][0] = 0; b.logits[i] = 1;
    }
    b.n_tokens = n;
    if (llama_decode(k->c, b) != 0) { llama_batch_free(b); return nullptr; }

    // 3) embedding poolé de la séquence 0
    const float* emb = llama_get_embeddings_seq(k->c, 0);
    llama_batch_free(b);
    if (!emb) return nullptr;

    // 4) normalisation L2 (cosinus = produit scalaire côté Android)
    std::vector<float> out(k->n_embd);
    double norm = 0.0; for (int i=0;i<k->n_embd;++i) norm += (double)emb[i]*emb[i];
    norm = norm > 0 ? 1.0/std::sqrt(norm) : 0.0;
    for (int i=0;i<k->n_embd;++i) out[i] = (float)(emb[i]*norm);

    jfloatArray arr = e->NewFloatArray(k->n_embd);
    e->SetFloatArrayRegion(arr, 0, k->n_embd, out.data());
    return arr;
}

extern "C" JNIEXPORT void JNICALL
Java_com_kazeia_llm_EngineJni_freeEmbedder(JNIEnv*, jobject, jlong h) {
    auto* k = (KEmbedder*) h;
    if (k->c) llama_free(k->c);
    if (k->m) llama_model_free(k->m);
    delete k;
}
```

> ⚠️ **À cross-checker contre `tools/llama-embedding` (ou `examples/embedding`) de TON fork.**
> Selon la version, un modèle encoder-only (BERT) peut nécessiter `llama_encode()` au lieu de
> `llama_decode()`, et le drapeau `logits` du batch peut différer. **Reproduis exactement ce que
> fait ton outil d'embedding de référence** — c'est la source de vérité.

## Tâche 3 — Bindings Kotlin

Dans `EngineJni` (`kazeia-android/.../llm/EngineLlmEngine.kt`) :

```kotlin
external fun loadEmbedder(ggufPath: String, nThreads: Int, pooling: Int): Long
external fun embedText(handle: Long, text: String): FloatArray?
external fun freeEmbedder(handle: Long)
```

---

## Points de vigilance (impératifs)

1. **Modèle = e5/bge dédié, FR-capable** (ex. `multilingual-e5-small`, 384-dim). Le **même** modèle
   à l'ingestion et à la requête, sinon vecteurs incompatibles. **Versionne-le** (le n° de version
   du modèle d'embedding fait partie du contrat de distribution).
2. **Vérifier d'abord que le fork charge l'arch BERT** (`bert` / `nomic-bert`) : le fork a été modifié
   pour Qwen/TTS — **teste le GGUF d'embedding avec `llama-embedding` AVANT d'intégrer**. Si l'arch a
   été retirée du fork, c'est LE point de blocage à traiter en priorité (≈ 5 min à vérifier).
3. **`n_ubatch ≥ nombre de tokens`** : le pooling MEAN/CLS impose la séquence entière dans un seul
   ubatch. D'où `n_ctx = n_batch = n_ubatch = 512` et le cap à 512 tokens (chunks plus courts).
4. **`pooling = -1` (défaut modèle) recommandé** ; sinon e5→MEAN(1), bge→CLS(2).
5. **Préfixes e5 (`query:` / `passage:`) = côté appelant**, pas dans le moteur.
6. **CPU only** (`n_gpu_layers = 0`), pas de HTP/HMX.
7. **Déterminisme** : pas d'échantillonnage, l'embed est déterministe — indispensable pour
   l'invariant « même modèle ingestion = requête ».

## Test d'acceptation (sur PC d'abord)

- `loadEmbedder` retourne un handle ≠ 0 sur le GGUF choisi ; `embedText` renvoie un `float[]` de
  taille `n_embd` (384 pour e5-small).
- **Déterminisme** : même texte → vecteur **bit-identique** sur 2 appels.
- **Norme ≈ 1.0** (L2).
- **Cohérence sémantique** : `cos("j'ai du mal à dormir", "insomnie")` >
  `cos("j'ai du mal à dormir", "recette de tarte")`.
- **Parité référence** : le vecteur ≈ celui de `llama-embedding` sur le même texte (à epsilon près).

## Build

Rebuild `libkazeia_engine.so` via le script habituel (`dist/build_kazeia_*.sh` / CMake).
**Aucune nouvelle dépendance** (tout est dans libllama). Livrer la `.so` + un GGUF d'embedding e5/bge.

---

## Ce qui reste côté app mobile (hors de cette spec, pour info)

Une fois `embedText` livré, l'intégration RAG dans `kazeia-android` est :
schéma SQLite (`chunk_text`, `source`, `embedding BLOB`) dans la base SQLCipher existante →
ingestion côté app admin (chunk sémantique + `embedText`) → au runtime, charger la matrice en RAM
→ brute-force cosinus (produit scalaire NEON) top-k=3–5 + **seuil de similarité** → injection d'un
bloc « CONTEXTE » budgété en tokens dans le system prompt. **Pas de vector-DB, pas d'ANN, pas
d'ONNX** (cf. analyse d'architecture). Le RAG ne porte jamais la gestion de crise.